[MLIR][Affine] Clean up outer logic of affine loop tiling pass

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.

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;

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);

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>

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) {