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