23 #include "llvm/Support/CommandLine.h"
24 #include "llvm/Support/Debug.h"
29 #define GEN_PASS_DEF_AFFINELOOPTILING
30 #include "mlir/Dialect/Affine/Passes.h.inc"
37 #define DEBUG_TYPE "affine-loop-tile"
42 struct LoopTiling :
public affine::impl::AffineLoopTilingBase<LoopTiling> {
43 LoopTiling() =
default;
44 explicit LoopTiling(uint64_t cacheSizeBytes,
bool avoidMaxMinBounds =
true)
45 : avoidMaxMinBounds(avoidMaxMinBounds) {
46 this->cacheSizeInKiB = cacheSizeBytes / 1024;
49 void runOnOperation()
override;
54 constexpr
static unsigned kDefaultTileSize = 4;
57 bool avoidMaxMinBounds =
true;
64 std::unique_ptr<OperationPass<func::FuncOp>>
66 return std::make_unique<LoopTiling>(cacheSizeBytes);
68 std::unique_ptr<OperationPass<func::FuncOp>>
70 return std::make_unique<LoopTiling>();
77 assert(band.size() == tileSizes->size() &&
"invalid tile size count");
78 for (
unsigned i = 0, e = band.size(); i < e; i++) {
79 unsigned &tSizeAdjusted = (*tileSizes)[i];
85 uint64_t constTripCount = *mayConst;
86 if (constTripCount > 1 && tSizeAdjusted > constTripCount / 2)
87 tSizeAdjusted = constTripCount / 2;
88 while (constTripCount % tSizeAdjusted != 0)
106 tileSizes->assign(band.size(), tileSize);
111 if (!this->tileSizes.empty()) {
112 tileSizes->assign(this->tileSizes.begin(), this->tileSizes.end());
113 tileSizes->resize(band.size(), kDefaultTileSize);
116 tileSizes->resize(band.size());
120 if (cacheSizeInKiB == 0) {
121 llvm::fill(*tileSizes, 1);
126 AffineForOp rootForOp = band[0];
136 llvm::fill(*tileSizes, LoopTiling::kDefaultTileSize);
137 if (avoidMaxMinBounds)
140 rootForOp.emitWarning(
"memory footprint unknown: using default tile "
141 "sizes adjusted to trip count divisors"));
146 uint64_t cacheSizeBytes = cacheSizeInKiB * 1024;
148 if (excessFactor <= 1) {
150 llvm::fill(*tileSizes, 1);
162 static_cast<unsigned>(floorl(std::pow(excessFactor, 1.0 / band.size())));
164 unsigned cumulProductOfTileSizes = 1;
165 for (
unsigned i = 0, e = band.size(); i < e; i++) {
167 (*tileSizes)[i] = tSize;
171 1U,
static_cast<unsigned>(excessFactor / cumulProductOfTileSizes));
172 cumulProductOfTileSizes *= (*tileSizes)[i];
174 if (avoidMaxMinBounds)
178 void LoopTiling::runOnOperation() {
180 std::vector<SmallVector<AffineForOp, 6>> bands;
184 for (
auto &band : bands) {
186 band.front().emitRemark(
"tiling nest is invalid due to dependences");
194 if (llvm::DebugFlag) {
195 auto diag = band[0].emitRemark(
"using tile sizes [");
196 for (
unsigned tSize : tileSizes)
197 diag << tSize <<
' ';
203 assert(!band.empty() &&
"guaranteed to succeed on empty bands");
204 LLVM_DEBUG(band.front()->emitRemark(
"loop tiling failed!\n"));
210 auto intraTileLoops =
213 assert(!intraTileLoops.empty() &&
214 "guaranteed to succeed on empty bands");
215 LLVM_DEBUG(intraTileLoops.front()->emitRemark(
216 "separation post tiling failed!\n"));
222 constexpr
unsigned LoopTiling::kDefaultTileSize;
static SmallVector< Value > getTileSizes(Location loc, amx::TileType tType, RewriterBase &rewriter)
Maps the 2-dim vector shape to the two 16-bit tile sizes.
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...
static std::string diag(const llvm::Value &value)
static Value max(ImplicitLocOpBuilder &builder, Value value, Value bound)
std::optional< uint64_t > getConstantTripCount(AffineForOp forOp)
Returns the trip count of the loop if it's a constant, std::nullopt otherwise.
bool isTilingValid(ArrayRef< AffineForOp > loops)
Checks whether hyper-rectangular loop tiling of the nest represented by loops is valid.
std::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...
std::unique_ptr< OperationPass< func::FuncOp > > createLoopTilingPass(uint64_t cacheSizeBytes)
Creates a pass to perform tiling on loop nests.
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.
void getTileableBands(func::FuncOp f, std::vector< SmallVector< AffineForOp, 6 >> *bands)
Identify valid and profitable bands of loops to tile.
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...
llvm::TypeSize divideCeil(llvm::TypeSize numerator, uint64_t denominator)
Divides the known min value of the numerator by the denominator and rounds the result up to the next ...
Include the generated interface declarations.