26 #include "llvm/Support/CommandLine.h"
27 #include "llvm/Support/Debug.h"
31 #define GEN_PASS_DEF_AFFINELOOPTILING
32 #include "mlir/Dialect/Affine/Passes.h.inc"
37 #define DEBUG_TYPE "affine-loop-tile"
42 struct LoopTiling :
public 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);
69 return std::make_unique<LoopTiling>();
76 assert(band.size() == tileSizes->size() &&
"invalid tile size count");
77 for (
unsigned i = 0, e = band.size(); i < e; i++) {
78 unsigned &tSizeAdjusted = (*tileSizes)[i];
84 uint64_t constTripCount = *mayConst;
85 if (constTripCount > 1 && tSizeAdjusted > constTripCount / 2)
86 tSizeAdjusted = constTripCount / 2;
87 while (constTripCount % tSizeAdjusted != 0)
105 tileSizes->assign(band.size(), tileSize);
110 if (!this->tileSizes.empty()) {
111 tileSizes->assign(this->tileSizes.begin(), this->tileSizes.end());
112 tileSizes->resize(band.size(), kDefaultTileSize);
115 tileSizes->resize(band.size());
118 AffineForOp rootForOp = band[0];
128 std::fill(tileSizes->begin(), tileSizes->end(),
129 LoopTiling::kDefaultTileSize);
130 if (avoidMaxMinBounds)
133 rootForOp.emitWarning(
"memory footprint unknown: using default tile "
134 "sizes adjusted to trip count divisors"));
139 uint64_t cacheSizeBytes = cacheSizeInKiB * 1024;
140 uint64_t excessFactor = llvm::divideCeil(*fp, cacheSizeBytes);
141 if (excessFactor <= 1) {
143 std::fill(tileSizes->begin(), tileSizes->end(), 1);
155 static_cast<unsigned>(floorl(std::pow(excessFactor, 1.0 / band.size())));
157 unsigned cumulProductOfTileSizes = 1;
158 for (
unsigned i = 0, e = band.size(); i < e; i++) {
160 (*tileSizes)[i] = tSize;
164 1U,
static_cast<unsigned>(excessFactor / cumulProductOfTileSizes));
165 cumulProductOfTileSizes *= (*tileSizes)[i];
167 if (avoidMaxMinBounds)
171 void LoopTiling::runOnOperation() {
173 std::vector<SmallVector<AffineForOp, 6>> bands;
177 for (
auto &band : bands) {
181 getTileSizes(band, &tileSizes);
182 if (llvm::DebugFlag) {
183 auto diag = band[0].emitRemark(
"using tile sizes [");
184 for (
unsigned tSize : tileSizes)
185 diag << tSize <<
' ';
191 assert(!band.empty() &&
"guaranteed to succeed on empty bands");
192 LLVM_DEBUG(band.front()->emitRemark(
"loop tiling failed!\n"));
198 auto intraTileLoops =
201 assert(!intraTileLoops.empty() &&
202 "guaranteed to succeed on empty bands");
203 LLVM_DEBUG(intraTileLoops.front()->emitRemark(
204 "separation post tiling failed!\n"));
210 constexpr
unsigned LoopTiling::kDefaultTileSize;
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)
Include the generated interface declarations.
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...
void getTileableBands(func::FuncOp f, std::vector< SmallVector< AffineForOp, 6 >> *bands)
Identify valid and profitable bands of loops to tile.
std::optional< uint64_t > getConstantTripCount(AffineForOp forOp)
Returns the trip count of the loop if it's a constant, std::nullopt otherwise.
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.
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value.
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.