MLIR  15.0.0git
LoopTiling.cpp
Go to the documentation of this file.
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"
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).
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.getValue();
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).
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.getValue(), 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;
Include the generated interface declarations.
static std::string diag(llvm::Value &v)
Optional< int64_t > getMemoryFootprintBytes(AffineForOp forOp, int memorySpace=-1)
Gets the memory footprint of all data touched in the specified memory space in bytes; if the memory s...
Definition: Utils.cpp:1331
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value...
Definition: LogicalResult.h:72
LogicalResult tilePerfectlyNested(MutableArrayRef< AffineForOp > input, ArrayRef< unsigned > tileSizes, SmallVectorImpl< AffineForOp > *tiledNest=nullptr)
Tiles the specified band of perfectly nested loops creating tile-space loops and intra-tile loops...
Definition: LoopUtils.cpp:860
std::unique_ptr< OperationPass< func::FuncOp > > createLoopTilingPass(uint64_t cacheSizeBytes)
Creates a pass to perform tiling on loop nests.
Definition: LoopTiling.cpp:57
static void adjustToDivisorsOfTripCounts(ArrayRef< AffineForOp > band, SmallVectorImpl< unsigned > *tileSizes)
Reduces each tile size to the largest divisor of the corresponding trip count (if the trip count is k...
Definition: LoopTiling.cpp:66
LogicalResult separateFullTiles(MutableArrayRef< AffineForOp > nest, SmallVectorImpl< AffineForOp > *fullTileNest=nullptr)
Separates full tiles from partial tiles for a perfect nest nest by generating a conditional guard tha...
Definition: LoopUtils.cpp:2790
void getTileableBands(func::FuncOp f, std::vector< SmallVector< AffineForOp, 6 >> *bands)
Identify valid and profitable bands of loops to tile.
Definition: LoopUtils.cpp:963
Optional< uint64_t > getConstantTripCount(AffineForOp forOp)
Returns the trip count of the loop if it&#39;s a constant, None otherwise.
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)