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;
LoopTiling__anond42b48350111::LoopTiling36 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>>
createLoopTilingPass(uint64_t cacheSizeBytes)57 mlir::createLoopTilingPass(uint64_t cacheSizeBytes) {
58 return std::make_unique<LoopTiling>(cacheSizeBytes);
59 }
createLoopTilingPass()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).
adjustToDivisorsOfTripCounts(ArrayRef<AffineForOp> band,SmallVectorImpl<unsigned> * tileSizes)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.
getTileSizes(ArrayRef<AffineForOp> band,SmallVectorImpl<unsigned> * tileSizes)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
runOnOperation()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