From ebd49d1b62ecbe874f3272ece5487c8b07e4a9d7 Mon Sep 17 00:00:00 2001 From: Uday Bondhugula Date: Mon, 21 Jul 2025 05:15:22 +0530 Subject: [PATCH] [MLIR][Affine] Clean up outer logic of affine loop tiling pass Clean up outer logic of affine loop tiling pass. A wrongly named temporary method was exposed publicly; fix that. Remove unconditional emission of remarks. --- mlir/include/mlir/Dialect/Affine/LoopUtils.h | 6 --- .../Dialect/Affine/Transforms/LoopTiling.cpp | 39 +++++++++++-------- mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp | 14 ------- .../Dialect/Affine/loop-tiling-validity.mlir | 12 +++--- .../Affine/TestAffineLoopParametricTiling.cpp | 8 +++- 5 files changed, 35 insertions(+), 44 deletions(-) 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 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> *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() = default; explicit LoopTiling(uint64_t cacheSizeBytes, bool avoidMaxMinBounds = true) @@ -59,6 +59,20 @@ struct LoopTiling : public affine::impl::AffineLoopTilingBase { } // namespace +/// Get bands of loops that are valid to tile from the top-level of `f`. +static void +getTopLevelTileableBands(func::FuncOp f, + std::vector> &bands) { + // Get maximal perfect nest of 'affine.for' ops starting from root + // (inclusive). + for (AffineForOp forOp : f.getOps()) { + SmallVector 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> @@ -122,10 +136,6 @@ void LoopTiling::getTileSizes(ArrayRef 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 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 band, void LoopTiling::runOnOperation() { // Bands of loops to tile. std::vector> 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 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 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> *bands) { - // Get maximal perfect nest of 'affine.for' insts starting from root - // (inclusive). - for (AffineForOp forOp : f.getOps()) { - SmallVector band; - getPerfectlyNestedLoops(band, forOp); - bands->push_back(band); - } -} - /// Unrolls this loop completely. LogicalResult mlir::affine::loopUnrollFull(AffineForOp forOp) { std::optional 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@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 band, } void TestAffineLoopParametricTiling::runOnOperation() { - // Bands of loops to tile. + // Get maximal perfect nest of 'affine.for' ops at the top-level. std::vector> bands; - getTileableBands(getOperation(), &bands); + for (AffineForOp forOp : getOperation().getOps()) { + SmallVector band; + getPerfectlyNestedLoops(band, forOp); + bands.push_back(band); + } // Tile each band. for (MutableArrayRef band : bands) {