1 //===- LoopTiling.cpp --- Loop tiling pass ------------------------------*-===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // This file implements a pass to tile loop nests. 10 // 11 //===----------------------------------------------------------------------===// 12 13 #include "PassDetail.h" 14 #include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h" 15 #include "mlir/Dialect/Affine/Analysis/AffineStructures.h" 16 #include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h" 17 #include "mlir/Dialect/Affine/Analysis/Utils.h" 18 #include "mlir/Dialect/Affine/IR/AffineOps.h" 19 #include "mlir/Dialect/Affine/IR/AffineValueMap.h" 20 #include "mlir/Dialect/Affine/LoopUtils.h" 21 #include "mlir/Dialect/Affine/Passes.h" 22 #include "mlir/Dialect/Affine/Utils.h" 23 #include "mlir/IR/BlockAndValueMapping.h" 24 #include "mlir/IR/Builders.h" 25 #include "llvm/Support/CommandLine.h" 26 #include "llvm/Support/Debug.h" 27 using namespace mlir; 28 29 #define DEBUG_TYPE "affine-loop-tile" 30 31 namespace { 32 33 /// A pass to perform loop tiling on all suitable loop nests of a Function. 34 struct LoopTiling : public AffineLoopTilingBase<LoopTiling> { 35 LoopTiling() = default; 36 explicit LoopTiling(uint64_t cacheSizeBytes, bool avoidMaxMinBounds = true) 37 : avoidMaxMinBounds(avoidMaxMinBounds) { 38 this->cacheSizeInKiB = cacheSizeBytes / 1024; 39 } 40 41 void runOnOperation() override; 42 void getTileSizes(ArrayRef<AffineForOp> band, 43 SmallVectorImpl<unsigned> *tileSizes); 44 45 // Default tile size if nothing is provided. 46 constexpr static unsigned kDefaultTileSize = 4; 47 48 // If true, tile sizes are set to avoid max/min in bounds if possible. 49 bool avoidMaxMinBounds = true; 50 }; 51 52 } // namespace 53 54 /// Creates a pass to perform loop tiling on all suitable loop nests of a 55 /// Function. 56 std::unique_ptr<OperationPass<func::FuncOp>> 57 mlir::createLoopTilingPass(uint64_t cacheSizeBytes) { 58 return std::make_unique<LoopTiling>(cacheSizeBytes); 59 } 60 std::unique_ptr<OperationPass<func::FuncOp>> mlir::createLoopTilingPass() { 61 return std::make_unique<LoopTiling>(); 62 } 63 64 /// Reduces each tile size to the largest divisor of the corresponding trip 65 /// count (if the trip count is known). 66 static void adjustToDivisorsOfTripCounts(ArrayRef<AffineForOp> band, 67 SmallVectorImpl<unsigned> *tileSizes) { 68 assert(band.size() == tileSizes->size() && "invalid tile size count"); 69 for (unsigned i = 0, e = band.size(); i < e; i++) { 70 unsigned &tSizeAdjusted = (*tileSizes)[i]; 71 Optional<uint64_t> mayConst = getConstantTripCount(band[i]); 72 if (!mayConst) 73 continue; 74 // Adjust the tile size to largest factor of the trip count less than 75 // tSize. 76 uint64_t constTripCount = *mayConst; 77 if (constTripCount > 1 && tSizeAdjusted > constTripCount / 2) 78 tSizeAdjusted = constTripCount / 2; 79 while (constTripCount % tSizeAdjusted != 0) 80 tSizeAdjusted--; 81 } 82 } 83 84 // Returns tile sizes to use. Checks CL options; if none are specified, sets it 85 // based on a simple model that looks at the memory footprint and determines 86 // tile sizes assuming identity accesses / 1:1 tile size proportional footprint 87 // along each of the dimensions being tiled. 88 // TODO: evolve this model. Tile size determination is a large area 89 // to play with in general. 90 void LoopTiling::getTileSizes(ArrayRef<AffineForOp> band, 91 SmallVectorImpl<unsigned> *tileSizes) { 92 if (band.empty()) 93 return; 94 95 // Use command-line tileSize for all loops if specified. 96 if (tileSize) { 97 tileSizes->assign(band.size(), tileSize); 98 return; 99 } 100 101 // Use tileSizes and fill them with default tile size if it's short. 102 if (!this->tileSizes.empty()) { 103 tileSizes->assign(this->tileSizes.begin(), this->tileSizes.end()); 104 tileSizes->resize(band.size(), kDefaultTileSize); 105 return; 106 } 107 tileSizes->resize(band.size()); 108 109 // The first loop in the band. 110 AffineForOp rootForOp = band[0]; 111 (void)rootForOp; 112 113 // Obtain memory footprint and set tile sizes so that a tile fits in 114 // the cache size. This is an approximation with the assumption that the 115 // footprint increases with the tile size linearly in that dimension (i.e., 116 // assumes one-to-one access function). 117 Optional<int64_t> fp = getMemoryFootprintBytes(band[0], 0); 118 if (!fp) { 119 // Fill with default tile sizes if footprint is unknown. 120 std::fill(tileSizes->begin(), tileSizes->end(), 121 LoopTiling::kDefaultTileSize); 122 if (avoidMaxMinBounds) 123 adjustToDivisorsOfTripCounts(band, tileSizes); 124 LLVM_DEBUG( 125 rootForOp.emitWarning("memory footprint unknown: using default tile " 126 "sizes adjusted to trip count divisors")); 127 return; 128 } 129 130 // Check how many times larger the cache size is when compared to footprint. 131 uint64_t cacheSizeBytes = cacheSizeInKiB * 1024; 132 uint64_t excessFactor = llvm::divideCeil(*fp, cacheSizeBytes); 133 if (excessFactor <= 1) { 134 // No need of any tiling - set tile size to 1. 135 std::fill(tileSizes->begin(), tileSizes->end(), 1); 136 return; 137 } 138 139 // Divide all loops equally in an attempt to reduce footprint. 140 // TODO: this is approximate. Ideally, obtain reuse factor / 141 // profitability along each dimension and weight tile sizes based on that as 142 // one possible approach. Or compute a polynomial in tile sizes and solve for 143 // it. 144 145 // For an n-d tileable band, compute the n^th root of the excess. 146 unsigned tSize = 147 static_cast<unsigned>(floorl(std::pow(excessFactor, 1.0 / band.size()))); 148 // We'll keep a running product to determine the last tile size better. 149 unsigned cumulProductOfTileSizes = 1; 150 for (unsigned i = 0, e = band.size(); i < e; i++) { 151 if (i < e - 1) 152 (*tileSizes)[i] = tSize; 153 else 154 // Set last tile size to cover the balance. 155 (*tileSizes)[i] = std::max( 156 1U, static_cast<unsigned>(excessFactor / cumulProductOfTileSizes)); 157 cumulProductOfTileSizes *= (*tileSizes)[i]; 158 } 159 if (avoidMaxMinBounds) 160 adjustToDivisorsOfTripCounts(band, tileSizes); 161 } 162 163 void LoopTiling::runOnOperation() { 164 // Bands of loops to tile. 165 std::vector<SmallVector<AffineForOp, 6>> bands; 166 getTileableBands(getOperation(), &bands); 167 168 // Tile each band. 169 for (auto &band : bands) { 170 // Set up tile sizes; fill missing tile sizes at the end with default tile 171 // size or tileSize if one was provided. 172 SmallVector<unsigned, 6> tileSizes; 173 getTileSizes(band, &tileSizes); 174 if (llvm::DebugFlag) { 175 auto diag = band[0].emitRemark("using tile sizes ["); 176 for (unsigned tSize : tileSizes) 177 diag << tSize << ' '; 178 diag << "]\n"; 179 } 180 SmallVector<AffineForOp, 6> tiledNest; 181 if (failed(tilePerfectlyNested(band, tileSizes, &tiledNest))) { 182 // An empty band always succeeds. 183 assert(!band.empty() && "guaranteed to succeed on empty bands"); 184 LLVM_DEBUG(band.front()->emitRemark("loop tiling failed!\n")); 185 continue; 186 } 187 188 // Separate full and partial tiles. 189 if (separate) { 190 auto intraTileLoops = 191 MutableArrayRef<AffineForOp>(tiledNest).drop_front(band.size()); 192 if (failed(separateFullTiles(intraTileLoops))) { 193 assert(!intraTileLoops.empty() && 194 "guaranteed to succeed on empty bands"); 195 LLVM_DEBUG(intraTileLoops.front()->emitRemark( 196 "separation post tiling failed!\n")); 197 } 198 } 199 } 200 } 201 202 constexpr unsigned LoopTiling::kDefaultTileSize; 203