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