[Mlir-commits] [mlir] [MLIR][Affine] Clean up outer logic of affine loop tiling pass (PR #149750)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Jul 20 17:08:02 PDT 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Uday Bondhugula (bondhugula)
<details>
<summary>Changes</summary>
Clean up outer logic of affine loop tiling pass. A wrongly named temporary method was exposed publicly; fix that. Remove unconditional emission of remarks.
---
Full diff: https://github.com/llvm/llvm-project/pull/149750.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/Affine/LoopUtils.h (-6)
- (modified) mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp (+22-17)
- (modified) mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp (-14)
- (modified) mlir/test/Dialect/Affine/loop-tiling-validity.mlir (+7-5)
- (modified) mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp (+6-2)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Affine/LoopUtils.h b/mlir/include/mlir/Dialect/Affine/LoopUtils.h
index 7fe1f6d48ceeb..9b59af73977d3 100644
--- a/mlir/include/mlir/Dialect/Affine/LoopUtils.h
+++ b/mlir/include/mlir/Dialect/Affine/LoopUtils.h
@@ -98,12 +98,6 @@ void promoteSingleIterationLoops(func::FuncOp f);
LogicalResult affineForOpBodySkew(AffineForOp forOp, ArrayRef<uint64_t> shifts,
bool unrollPrologueEpilogue = false);
-/// Identify valid and profitable bands of loops to tile. This is currently just
-/// a temporary placeholder to test the mechanics of tiled code generation.
-/// Returns all maximal outermost perfect loop nests to tile.
-void getTileableBands(func::FuncOp f,
- std::vector<SmallVector<AffineForOp, 6>> *bands);
-
/// Tiles the specified band of perfectly nested loops creating tile-space loops
/// and intra-tile loops. A band is a contiguous set of loops. This utility
/// doesn't check for the validity of tiling itself, but just performs it.
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
index 8a989ff985c6f..59c630ca11dca 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
@@ -6,7 +6,7 @@
//
//===----------------------------------------------------------------------===//
//
-// This file implements a pass to tile loop nests.
+// This file implements a pass to tile affine loop nests.
//
//===----------------------------------------------------------------------===//
@@ -38,7 +38,7 @@ using namespace mlir::affine;
namespace {
-/// A pass to perform loop tiling on all suitable loop nests of a Function.
+/// A pass to perform loop tiling on all suitable loop nests of a func op.
struct LoopTiling : public affine::impl::AffineLoopTilingBase<LoopTiling> {
LoopTiling() = default;
explicit LoopTiling(uint64_t cacheSizeBytes, bool avoidMaxMinBounds = true)
@@ -59,6 +59,20 @@ struct LoopTiling : public affine::impl::AffineLoopTilingBase<LoopTiling> {
} // namespace
+/// Get bands of loops that are valid to tile from the top-level of `f`.
+static void
+getTopLevelTileableBands(func::FuncOp f,
+ std::vector<SmallVector<AffineForOp, 6>> &bands) {
+ // Get maximal perfect nest of 'affine.for' ops starting from root
+ // (inclusive).
+ for (AffineForOp forOp : f.getOps<AffineForOp>()) {
+ SmallVector<AffineForOp, 6> band;
+ getPerfectlyNestedLoops(band, forOp);
+ if (isTilingValid(band))
+ bands.push_back(band);
+ }
+}
+
/// Creates a pass to perform loop tiling on all suitable loop nests of a
/// Function.
std::unique_ptr<OperationPass<func::FuncOp>>
@@ -122,10 +136,6 @@ void LoopTiling::getTileSizes(ArrayRef<AffineForOp> band,
return;
}
- // The first loop in the band.
- AffineForOp rootForOp = band[0];
- (void)rootForOp;
-
// Obtain memory footprint and set tile sizes so that a tile fits in
// the cache size. This is an approximation with the assumption that the
// footprint increases with the tile size linearly in that dimension (i.e.,
@@ -136,6 +146,9 @@ void LoopTiling::getTileSizes(ArrayRef<AffineForOp> band,
llvm::fill(*tileSizes, LoopTiling::kDefaultTileSize);
if (avoidMaxMinBounds)
adjustToDivisorsOfTripCounts(band, tileSizes);
+ // The first loop in the band.
+ AffineForOp rootForOp = band[0];
+ (void)rootForOp;
LLVM_DEBUG(
rootForOp.emitWarning("memory footprint unknown: using default tile "
"sizes adjusted to trip count divisors"));
@@ -178,23 +191,17 @@ void LoopTiling::getTileSizes(ArrayRef<AffineForOp> band,
void LoopTiling::runOnOperation() {
// Bands of loops to tile.
std::vector<SmallVector<AffineForOp, 6>> bands;
- getTileableBands(getOperation(), &bands);
+ getTopLevelTileableBands(getOperation(), bands);
// Tile each band.
for (auto &band : bands) {
- if (!isTilingValid(band)) {
- band.front().emitRemark("tiling nest is invalid due to dependences");
- continue;
- }
-
// Set up tile sizes; fill missing tile sizes at the end with default tile
// size or tileSize if one was provided.
SmallVector<unsigned, 6> tileSizes;
getTileSizes(band, &tileSizes);
if (llvm::DebugFlag) {
auto diag = band[0].emitRemark("using tile sizes [");
- for (unsigned tSize : tileSizes)
- diag << tSize << ' ';
+ llvm::interleaveComma(tileSizes, llvm::dbgs());
diag << "]\n";
}
SmallVector<AffineForOp, 6> tiledNest;
@@ -213,10 +220,8 @@ void LoopTiling::runOnOperation() {
assert(!intraTileLoops.empty() &&
"guaranteed to succeed on empty bands");
LLVM_DEBUG(intraTileLoops.front()->emitRemark(
- "separation post tiling failed!\n"));
+ "separation post tiling failed!"));
}
}
}
}
-
-constexpr unsigned LoopTiling::kDefaultTileSize;
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
index 0501616ad912c..bc4e2d10f5176 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp
@@ -868,20 +868,6 @@ void mlir::affine::getPerfectlyNestedLoops(
}
}
-/// Identify valid and profitable bands of loops to tile. This is currently just
-/// a temporary placeholder to test the mechanics of tiled code generation.
-/// Returns all maximal outermost perfect loop nests to tile.
-void mlir::affine::getTileableBands(
- func::FuncOp f, std::vector<SmallVector<AffineForOp, 6>> *bands) {
- // Get maximal perfect nest of 'affine.for' insts starting from root
- // (inclusive).
- for (AffineForOp forOp : f.getOps<AffineForOp>()) {
- SmallVector<AffineForOp, 6> band;
- getPerfectlyNestedLoops(band, forOp);
- bands->push_back(band);
- }
-}
-
/// Unrolls this loop completely.
LogicalResult mlir::affine::loopUnrollFull(AffineForOp forOp) {
std::optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);
diff --git a/mlir/test/Dialect/Affine/loop-tiling-validity.mlir b/mlir/test/Dialect/Affine/loop-tiling-validity.mlir
index e2c3832f695cc..cf2e12613ab92 100644
--- a/mlir/test/Dialect/Affine/loop-tiling-validity.mlir
+++ b/mlir/test/Dialect/Affine/loop-tiling-validity.mlir
@@ -7,8 +7,8 @@
// CHECK-DAG: [[$LB:#map[0-9]*]] = affine_map<(d0) -> (d0)>
// CHECK-DAG: [[$UB:#map[0-9]*]] = affine_map<(d0) -> (d0 + 32)>
-// CHECK-LABEL: func @legal_loop()
-func.func @legal_loop() {
+// CHECK-LABEL: func @valid_to_tile()
+func.func @valid_to_tile() {
%0 = memref.alloc() : memref<64xf32>
affine.for %i = 0 to 64 {
@@ -20,8 +20,8 @@ func.func @legal_loop() {
return
}
-// CHECK: affine.for %{{.*}} = 0 to 64 step 32 {
-// CHECK-NEXT: affine.for %{{.*}} = [[$LB]](%{{.*}}) to [[$UB]](%{{.*}}) {
+// CHECK: affine.for %{{.*}} = 0 to 64 step 32 {
+// CHECK-NEXT: affine.for %{{.*}} = [[$LB]](%{{.*}}) to [[$UB]](%{{.*}}) {
// -----
@@ -33,8 +33,10 @@ func.func @legal_loop() {
func.func @illegal_loop_with_diag_dependence() {
%A = memref.alloc() : memref<64x64xf32>
+ // No tiling here.
+ // CHECK: affine.for %{{.*}} = 0 to 64 {
+ // CHECK-NEXT: affine.for %{{.*}} = 0 to 64 {
affine.for %i = 0 to 64 {
- // expected-remark at above {{tiling nest is invalid due to dependences}}
affine.for %j = 0 to 64 {
%0 = affine.load %A[%j, %i] : memref<64x64xf32>
%1 = affine.load %A[%i, %j - 1] : memref<64x64xf32>
diff --git a/mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp b/mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp
index f8e76356c4321..39a8cd953f7ca 100644
--- a/mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp
+++ b/mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp
@@ -74,9 +74,13 @@ getTilingParameters(ArrayRef<AffineForOp> band,
}
void TestAffineLoopParametricTiling::runOnOperation() {
- // Bands of loops to tile.
+ // Get maximal perfect nest of 'affine.for' ops at the top-level.
std::vector<SmallVector<AffineForOp, 6>> bands;
- getTileableBands(getOperation(), &bands);
+ for (AffineForOp forOp : getOperation().getOps<AffineForOp>()) {
+ SmallVector<AffineForOp, 6> band;
+ getPerfectlyNestedLoops(band, forOp);
+ bands.push_back(band);
+ }
// Tile each band.
for (MutableArrayRef<AffineForOp> band : bands) {
``````````
</details>
https://github.com/llvm/llvm-project/pull/149750
More information about the Mlir-commits
mailing list