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 "mlir/Analysis/AffineAnalysis.h" 14 #include "mlir/Analysis/AffineStructures.h" 15 #include "mlir/Analysis/LoopAnalysis.h" 16 #include "mlir/Analysis/Utils.h" 17 #include "mlir/Dialect/Affine/IR/AffineOps.h" 18 #include "mlir/Dialect/Affine/IR/AffineValueMap.h" 19 #include "mlir/Dialect/Affine/Passes.h" 20 #include "mlir/IR/BlockAndValueMapping.h" 21 #include "mlir/IR/Builders.h" 22 #include "mlir/Pass/Pass.h" 23 #include "mlir/Transforms/LoopUtils.h" 24 #include "mlir/Transforms/Utils.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 static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options"); 32 33 static llvm::cl::opt<unsigned long long> 34 clCacheSizeKiB("affine-tile-cache-size", 35 llvm::cl::desc("Set size of cache to tile for in KiB"), 36 llvm::cl::cat(clOptionsCategory)); 37 38 // Separate full and partial tiles. 39 static llvm::cl::opt<bool> 40 clSeparate("affine-tile-separate", 41 llvm::cl::desc("Separate full and partial tiles"), 42 llvm::cl::cat(clOptionsCategory)); 43 44 // Tile size to use for all loops (overrides -tile-sizes if provided). 45 static llvm::cl::opt<unsigned> 46 clTileSize("affine-tile-size", 47 llvm::cl::desc("Use this tile size for all loops"), 48 llvm::cl::cat(clOptionsCategory)); 49 50 // List of tile sizes. If any of them aren't provided, they are filled with 51 // clTileSize / kDefaultTileSize. 52 static llvm::cl::list<unsigned> clTileSizes( 53 "affine-tile-sizes", 54 llvm::cl::desc( 55 "List of tile sizes for each perfect nest (overridden by -tile-size)"), 56 llvm::cl::ZeroOrMore, llvm::cl::cat(clOptionsCategory)); 57 58 namespace { 59 60 /// A pass to perform loop tiling on all suitable loop nests of a Function. 61 struct LoopTiling : public FunctionPass<LoopTiling> { 62 explicit LoopTiling(uint64_t cacheSizeBytes = kDefaultCacheMemCapacity, 63 bool avoidMaxMinBounds = true) 64 : cacheSizeBytes(cacheSizeBytes), avoidMaxMinBounds(avoidMaxMinBounds) {} 65 66 void runOnFunction() override; 67 void getTileSizes(ArrayRef<AffineForOp> band, 68 SmallVectorImpl<unsigned> *tileSizes); 69 70 // Default tile size if nothing is provided. 71 constexpr static unsigned kDefaultTileSize = 4; 72 constexpr static uint64_t kDefaultCacheMemCapacity = 512 * 1024UL; 73 74 // Capacity of the cache to tile for. 75 uint64_t cacheSizeBytes; 76 // If true, tile sizes are set to avoid max/min in bounds if possible. 77 bool avoidMaxMinBounds; 78 }; 79 80 } // end anonymous namespace 81 82 /// Creates a pass to perform loop tiling on all suitable loop nests of a 83 /// Function. 84 std::unique_ptr<OpPassBase<FuncOp>> 85 mlir::createLoopTilingPass(uint64_t cacheSizeBytes) { 86 return std::make_unique<LoopTiling>(cacheSizeBytes); 87 } 88 std::unique_ptr<OpPassBase<FuncOp>> mlir::createLoopTilingPass() { 89 return std::make_unique<LoopTiling>(); 90 } 91 92 // Move the loop body of AffineForOp 'src' from 'src' into the specified 93 // location in destination's body, ignoring the terminator. 94 static inline void moveLoopBody(AffineForOp src, AffineForOp dest, 95 Block::iterator loc) { 96 auto &insts = src.getBody()->getOperations(); 97 dest.getBody()->getOperations().splice(loc, insts, insts.begin(), 98 std::prev(insts.end())); 99 } 100 101 // Move the loop body of AffineForOp 'src' from 'src' to the start of dest's 102 // body. 103 static inline void moveLoopBody(AffineForOp src, AffineForOp dest) { 104 moveLoopBody(src, dest, dest.getBody()->begin()); 105 } 106 107 /// Constructs and sets new loop bounds after tiling for the case of 108 /// hyper-rectangular index sets, where the bounds of one dimension do not 109 /// depend on other dimensions. Bounds of each dimension can thus be treated 110 /// independently, and deriving the new bounds is much simpler and faster 111 /// than for the case of tiling arbitrary polyhedral shapes. 112 static void 113 constructTiledIndexSetHyperRect(MutableArrayRef<AffineForOp> origLoops, 114 MutableArrayRef<AffineForOp> newLoops, 115 ArrayRef<unsigned> tileSizes) { 116 assert(!origLoops.empty()); 117 assert(origLoops.size() == tileSizes.size()); 118 119 OpBuilder b(origLoops[0].getOperation()); 120 unsigned width = origLoops.size(); 121 122 // Bounds for tile space loops. 123 for (unsigned i = 0; i < width; i++) { 124 auto lbOperands = origLoops[i].getLowerBoundOperands(); 125 auto ubOperands = origLoops[i].getUpperBoundOperands(); 126 SmallVector<Value, 4> newLbOperands(lbOperands); 127 SmallVector<Value, 4> newUbOperands(ubOperands); 128 newLoops[i].setLowerBound(newLbOperands, origLoops[i].getLowerBoundMap()); 129 newLoops[i].setUpperBound(newUbOperands, origLoops[i].getUpperBoundMap()); 130 newLoops[i].setStep(tileSizes[i]); 131 } 132 // Bounds for intra-tile loops. 133 for (unsigned i = 0; i < width; i++) { 134 int64_t largestDiv = getLargestDivisorOfTripCount(origLoops[i]); 135 auto mayBeConstantCount = getConstantTripCount(origLoops[i]); 136 // The lower bound is just the tile-space loop. 137 AffineMap lbMap = b.getDimIdentityMap(); 138 newLoops[width + i].setLowerBound( 139 /*operands=*/newLoops[i].getInductionVar(), lbMap); 140 141 // Set the upper bound. 142 if (mayBeConstantCount.hasValue() && 143 mayBeConstantCount.getValue() < tileSizes[i]) { 144 // Trip count is less than tile size; upper bound is the trip count. 145 auto ubMap = b.getConstantAffineMap(mayBeConstantCount.getValue()); 146 newLoops[width + i].setUpperBoundMap(ubMap); 147 } else if (largestDiv % tileSizes[i] != 0) { 148 // Intra-tile loop ii goes from i to min(i + tileSize, ub_i). 149 // Construct the upper bound map; the operands are the original operands 150 // with 'i' (tile-space loop) appended to it. The new upper bound map is 151 // the original one with an additional expression i + tileSize appended. 152 auto ub = origLoops[i].getUpperBound(); 153 SmallVector<Value, 4> ubOperands; 154 ubOperands.reserve(ub.getNumOperands() + 1); 155 auto origUbMap = ub.getMap(); 156 // Add dim operands from original upper bound. 157 for (unsigned j = 0, e = origUbMap.getNumDims(); j < e; ++j) { 158 ubOperands.push_back(ub.getOperand(j)); 159 } 160 // Add dim operand for new loop upper bound. 161 ubOperands.push_back(newLoops[i].getInductionVar()); 162 // Add symbol operands from original upper bound. 163 for (unsigned j = 0, e = origUbMap.getNumSymbols(); j < e; ++j) { 164 ubOperands.push_back(ub.getOperand(origUbMap.getNumDims() + j)); 165 } 166 SmallVector<AffineExpr, 4> boundExprs; 167 boundExprs.reserve(1 + origUbMap.getNumResults()); 168 auto dim = b.getAffineDimExpr(origUbMap.getNumDims()); 169 // The new upper bound map is the original one with an additional 170 // expression i + tileSize appended. 171 boundExprs.push_back(dim + tileSizes[i]); 172 boundExprs.append(origUbMap.getResults().begin(), 173 origUbMap.getResults().end()); 174 auto ubMap = AffineMap::get(origUbMap.getNumDims() + 1, 175 origUbMap.getNumSymbols(), boundExprs); 176 newLoops[width + i].setUpperBound(/*operands=*/ubOperands, ubMap); 177 } else { 178 // No need of the min expression. 179 auto dim = b.getAffineDimExpr(0); 180 auto ubMap = AffineMap::get(1, 0, dim + tileSizes[i]); 181 newLoops[width + i].setUpperBound(newLoops[i].getInductionVar(), ubMap); 182 } 183 } 184 } 185 186 /// Tiles the specified band of perfectly nested loops creating tile-space loops 187 /// and intra-tile loops. A band is a contiguous set of loops. 188 // TODO(bondhugula): handle non hyper-rectangular spaces. 189 LogicalResult mlir::tileCodeGen(MutableArrayRef<AffineForOp> band, 190 ArrayRef<unsigned> tileSizes, 191 SmallVectorImpl<AffineForOp> *tiledNest) { 192 // Check if the supplied for op's are all successively nested. 193 assert(!band.empty() && "no loops in band"); 194 assert(band.size() == tileSizes.size() && "Too few/many tile sizes"); 195 196 for (unsigned i = 1, e = band.size(); i < e; i++) 197 assert(band[i].getParentOp() == band[i - 1] && "not a perfect nest / band"); 198 199 auto origLoops = band; 200 201 AffineForOp rootAffineForOp = origLoops[0]; 202 auto loc = rootAffineForOp.getLoc(); 203 // Note that width is at least one since band isn't empty. 204 unsigned width = band.size(); 205 206 SmallVector<AffineForOp, 6> tiledLoops(2 * width); 207 208 // The outermost among the loops as we add more.. 209 auto *topLoop = rootAffineForOp.getOperation(); 210 AffineForOp innermostPointLoop; 211 212 // Add intra-tile (or point) loops. 213 for (unsigned i = 0; i < width; i++) { 214 OpBuilder b(topLoop); 215 // Loop bounds will be set later. 216 auto pointLoop = b.create<AffineForOp>(loc, 0, 0); 217 pointLoop.getBody()->getOperations().splice( 218 pointLoop.getBody()->begin(), topLoop->getBlock()->getOperations(), 219 topLoop); 220 tiledLoops[2 * width - 1 - i] = pointLoop; 221 topLoop = pointLoop.getOperation(); 222 if (i == 0) 223 innermostPointLoop = pointLoop; 224 } 225 226 // Add tile space loops; 227 for (unsigned i = width; i < 2 * width; i++) { 228 OpBuilder b(topLoop); 229 // Loop bounds will be set later. 230 auto tileSpaceLoop = b.create<AffineForOp>(loc, 0, 0); 231 tileSpaceLoop.getBody()->getOperations().splice( 232 tileSpaceLoop.getBody()->begin(), topLoop->getBlock()->getOperations(), 233 topLoop); 234 tiledLoops[2 * width - i - 1] = tileSpaceLoop; 235 topLoop = tileSpaceLoop.getOperation(); 236 } 237 238 // Move the loop body of the original nest to the new one. 239 moveLoopBody(origLoops[origLoops.size() - 1], innermostPointLoop); 240 241 SmallVector<Value, 8> origLoopIVs; 242 extractForInductionVars(band, &origLoopIVs); 243 SmallVector<Optional<Value>, 6> ids(origLoopIVs.begin(), origLoopIVs.end()); 244 FlatAffineConstraints cst; 245 getIndexSet(band, &cst); 246 247 if (!cst.isHyperRectangular(0, width)) { 248 llvm::dbgs() << "tiled code generation unimplemented for the " 249 "non-hyperrectangular case, op:" 250 << *rootAffineForOp << "\n"; 251 return failure(); 252 } 253 254 constructTiledIndexSetHyperRect(origLoops, tiledLoops, tileSizes); 255 256 // Replace original IVs with intra-tile loop IVs. 257 for (unsigned i = 0; i < width; i++) 258 origLoopIVs[i].replaceAllUsesWith(tiledLoops[i + width].getInductionVar()); 259 260 // Erase the old loop nest. 261 rootAffineForOp.erase(); 262 263 if (tiledNest) 264 *tiledNest = std::move(tiledLoops); 265 266 return success(); 267 } 268 269 // Identify valid and profitable bands of loops to tile. This is currently just 270 // a temporary placeholder to test the mechanics of tiled code generation. 271 // Returns all maximal outermost perfect loop nests to tile. 272 static void getTileableBands(FuncOp f, 273 std::vector<SmallVector<AffineForOp, 6>> *bands) { 274 // Get maximal perfect nest of 'affine.for' insts starting from root 275 // (inclusive). 276 auto getMaximalPerfectLoopNest = [&](AffineForOp root) { 277 SmallVector<AffineForOp, 6> band; 278 getPerfectlyNestedLoops(band, root); 279 bands->push_back(band); 280 }; 281 282 for (auto &block : f) 283 for (auto &op : block) 284 if (auto forOp = dyn_cast<AffineForOp>(op)) 285 getMaximalPerfectLoopNest(forOp); 286 } 287 288 // Reduce each tile size to the largest divisor of the corresponding trip count 289 // (if the trip count is known). 290 static void adjustToDivisorsOfTripCounts(ArrayRef<AffineForOp> band, 291 SmallVectorImpl<unsigned> *tileSizes) { 292 assert(band.size() == tileSizes->size() && "invalid tile size count"); 293 for (unsigned i = 0, e = band.size(); i < e; i++) { 294 unsigned &tSizeAdjusted = (*tileSizes)[i]; 295 auto mayConst = getConstantTripCount(band[i]); 296 if (!mayConst.hasValue()) 297 continue; 298 // Adjust the tile size to largest factor of the trip count less than 299 // tSize. 300 uint64_t constTripCount = mayConst.getValue(); 301 if (constTripCount > 1 && tSizeAdjusted > constTripCount / 2) 302 tSizeAdjusted = constTripCount / 2; 303 while (constTripCount % tSizeAdjusted != 0) 304 tSizeAdjusted--; 305 } 306 } 307 308 // Returns tile sizes to use. Checks CL options; if none are specified, sets it 309 // based on a simple model that looks at the memory footprint and determines 310 // tile sizes assuming identity accesses / 1:1 tile size proportional footprint 311 // along each of the dimensions being tiled. 312 // TODO(mlir-team): evolve this model. Tile size determination is a large area 313 // to play with in general. 314 void LoopTiling::getTileSizes(ArrayRef<AffineForOp> band, 315 SmallVectorImpl<unsigned> *tileSizes) { 316 if (band.empty()) 317 return; 318 319 tileSizes->resize(band.size()); 320 321 // Use clTileSize for all loops if specified. 322 if (clTileSize.getNumOccurrences() > 0) { 323 std::fill(tileSizes->begin(), tileSizes->end(), clTileSize); 324 return; 325 } 326 327 // Use clTileSizes and fill them with default tile size if it's short. 328 if (!clTileSizes.empty()) { 329 std::fill(tileSizes->begin(), tileSizes->end(), 330 LoopTiling::kDefaultTileSize); 331 std::copy(clTileSizes.begin(), 332 clTileSizes.begin() + std::min(clTileSizes.size(), band.size()), 333 tileSizes->begin()); 334 return; 335 } 336 337 // The first loop in the band. 338 auto rootForOp = band[0]; 339 (void)rootForOp; 340 341 // Obtain memory footprint and set tile sizes so that a tile fits in 342 // the cache size. This is an approximation with the assumption that the 343 // footprint increases with the tile size linearly in that dimension (i.e., 344 // assumes one-to-one access function). 345 auto fp = getMemoryFootprintBytes(band[0], 0); 346 if (!fp.hasValue()) { 347 // Fill with default tile sizes if footprint is unknown. 348 std::fill(tileSizes->begin(), tileSizes->end(), 349 LoopTiling::kDefaultTileSize); 350 if (avoidMaxMinBounds) 351 adjustToDivisorsOfTripCounts(band, tileSizes); 352 LLVM_DEBUG( 353 rootForOp.emitWarning("memory footprint unknown: using default tile " 354 "sizes adjusted to trip count divisors")); 355 return; 356 } 357 358 // Check how many times larger the cache size is when compared to footprint. 359 uint64_t excessFactor = llvm::divideCeil(fp.getValue(), cacheSizeBytes); 360 if (excessFactor <= 1) { 361 // No need of any tiling - set tile size to 1. 362 std::fill(tileSizes->begin(), tileSizes->end(), 1); 363 return; 364 } 365 366 // Divide all loops equally in an attempt to reduce footprint. 367 // TODO(bondhugula): this is approximate. Ideally, obtain reuse factor / 368 // profitability along each dimension and weight tile sizes based on that as 369 // one possible approach. Or compute a polynomial in tile sizes and solve for 370 // it. 371 372 // For an n-d tileable band, compute n^th root of the excess. 373 unsigned tSize = 374 static_cast<unsigned>(floorl(std::pow(excessFactor, 1.0 / band.size()))); 375 // We'll keep a running product to determine the last tile size better. 376 unsigned cumulProductOfTileSizes = 1; 377 for (unsigned i = 0, e = band.size(); i < e; i++) { 378 if (i < e - 1) 379 (*tileSizes)[i] = tSize; 380 else 381 // Set last tile size to cover the balance. 382 (*tileSizes)[i] = std::max( 383 1U, static_cast<unsigned>(excessFactor / cumulProductOfTileSizes)); 384 cumulProductOfTileSizes *= (*tileSizes)[i]; 385 } 386 if (avoidMaxMinBounds) 387 adjustToDivisorsOfTripCounts(band, tileSizes); 388 } 389 390 void LoopTiling::runOnFunction() { 391 // Override cache size if provided on command line. 392 if (clCacheSizeKiB.getNumOccurrences() > 0) 393 cacheSizeBytes = clCacheSizeKiB * 1024; 394 395 // Bands of loops to tile. 396 std::vector<SmallVector<AffineForOp, 6>> bands; 397 getTileableBands(getFunction(), &bands); 398 399 // Tile each band. 400 for (auto &band : bands) { 401 // Set up tile sizes; fill missing tile sizes at the end with default tile 402 // size or clTileSize if one was provided. 403 SmallVector<unsigned, 6> tileSizes; 404 getTileSizes(band, &tileSizes); 405 if (llvm::DebugFlag) { 406 auto diag = band[0].emitRemark("using tile sizes ["); 407 for (auto tSize : tileSizes) 408 diag << tSize << ' '; 409 diag << "]\n"; 410 } 411 SmallVector<AffineForOp, 6> tiledNest; 412 if (failed(tileCodeGen(band, tileSizes, &tiledNest))) 413 return signalPassFailure(); 414 415 // Separate full and partial tiles. 416 if (clSeparate) { 417 auto intraTileLoops = 418 MutableArrayRef<AffineForOp>(tiledNest).drop_front(band.size()); 419 separateFullTiles(intraTileLoops); 420 } 421 } 422 } 423 424 constexpr unsigned LoopTiling::kDefaultTileSize; 425 constexpr uint64_t LoopTiling::kDefaultCacheMemCapacity; 426