Multiple fixes to affine loop tiling return status and checks

[bondhugula]

Summary

Fix crash in the presence of yield values. Multiple fixes to affine loop
tiling pre-condition checks and return status. Do not signal pass
failure on a failure to tile since the IR is still valid. Detect index
set computation failure in checkIfHyperrectangular and return failure.
Replace assertions with proper status return. Move checks to an
appropriate place earlier in the utility before mutation happens.

[bondhugula]

More cleanup.

[bondhugula]

Handle trivial case.

[ayzhuang]

Looks great, thanks!

mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp

for (auto &band : bands) {
     getTileSizes(band, &tileSizes);
     if (llvm::DebugFlag) {
       auto diag = band[0].emitRemark("using tile sizes [");
       for (unsigned tSize : tileSizes)
         diag << tSize << ' ';
       diag << "]\n";
     }
     SmallVector<AffineForOp, 6> tiledNest;
-    if (failed(tilePerfectlyNested(band, tileSizes, &tiledNest)))
-      return signalPassFailure();
+    if (failed(tilePerfectlyNested(band, tileSizes, &tiledNest))) {
+      // An empty band always succeeds.
+      assert(!band.empty() && "guaranteed to succeed on empty bands");
+      LLVM_DEBUG(band.front()->emitRemark("loop tiling failed!\n"));
+      continue;
+    }
 
     // Separate full and partial tiles.
     if (separate) {
       auto intraTileLoops =
           MutableArrayRef<AffineForOp>(tiledNest).drop_front(band.size());
-      (void)separateFullTiles(intraTileLoops);
+      if (failed(separateFullTiles(intraTileLoops))) {
+        assert(!intraTileLoops.empty() &&
+               "guaranteed to succeed on empty bands");
+        LLVM_DEBUG(intraTileLoops.front()->emitRemark(
+            "separation post tiling failed!\n"));
+      }
     }
   }
 }
 
 constexpr unsigned LoopTiling::kDefaultTileSize;

mlir/lib/Transforms/Utils/LoopUtils.cpp

 /// lexicographically increasing order on the vector of loop indices. This
 /// function will return failure when any dependence component is negative along
 /// any of `origLoops`.
 LogicalResult
 checkTilingLegality(MutableArrayRef<mlir::AffineForOp> origLoops) {
   return success(checkTilingLegalityImpl(origLoops));
 }
 
-/// Check if the input data is valid and whether tiled code will be legal or
-/// not.
+/// Checks whether a loop nest is hyper-rectangular or not.
+LogicalResult checkIfHyperRectangular(MutableArrayRef<AffineForOp> input) {
+  FlatAffineValueConstraints cst;
+  SmallVector<Operation *, 8> ops(input.begin(), input.end());
+  // 0-d or 1-d is trivially hyper-rectangular.
+  if (input.size() <= 1)
+    return success();
+  if (failed(getIndexSet(ops, &cst))) {
+    LLVM_DEBUG(llvm::dbgs() << "Index set computation failed!\n");
+    return failure();
+  }
+  if (!cst.isHyperRectangular(0, input.size())) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Non-hyperrectangular nests not supported for tiling!\n");
+    return failure();
+  }
+  return success();
+}
+
+/// Check if the input nest is supported for tiling and whether tiling would be
+/// legal or not.
 template <typename t>
-void performPreTilingChecks(MutableArrayRef<AffineForOp> input,
+LogicalResult performPreTilingChecks(MutableArrayRef<AffineForOp> input,
                                      ArrayRef<t> tileSizes) {
-  // Check if the supplied for op's are all successively nested.
-  assert(!input.empty() && "no loops in input band");
   assert(input.size() == tileSizes.size() && "Too few/many tile sizes");
 
-  assert(isPerfectlyNested(input) && "input loops not perfectly nested");
+  if (llvm::any_of(input,
+                   [](AffineForOp op) { return op.getNumResults() > 0; })) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Cannot tile nest where a loop has yield values\n");
+    return failure();
+  }
+
+  // Check if the supplied `for` ops are all successively nested.
+  if (!isPerfectlyNested(input)) {
+    LLVM_DEBUG(llvm::dbgs() << "input loops not perfectly nested");
+    return failure();
+  }
+
+  if (failed(checkIfHyperRectangular(input)))
+    return failure();
+
+  // Check if tiling is legal.
+  if (failed(checkTilingLegality(input))) {
+    input[0].emitRemark("tiling code is illegal due to dependences");
+    return failure();
+  }
 
-  // Perform tiling legality test.
-  if (failed(checkTilingLegality(input)))
-    input[0].emitRemark("tiled code is illegal due to dependences");
+  return success();
 }
 
 /// Move the loop body of AffineForOp 'src' from 'src' into the specified
 /// location in destination's body, ignoring the terminator.
 static void moveLoopBodyImpl(AffineForOp src, AffineForOp dest,
                              Block::iterator loc) {
   auto &ops = src.getBody()->getOperations();
   dest.getBody()->getOperations().splice(loc, ops, ops.begin(),
for (unsigned i = width; i < 2 * width; i++) {
     tiledLoops[2 * width - i - 1] = tileSpaceLoop;
     topLoop = tileSpaceLoop.getOperation();
   }
 
   // Move the loop body of the original nest to the new one.
   moveLoopBody(origLoops.back(), innermostPointLoop);
 }
 
-/// Checks whether a loop nest is hyper-rectangular or not.
-LogicalResult checkIfHyperRectangular(MutableArrayRef<AffineForOp> input,
-                                      AffineForOp rootAffineForOp,
-                                      unsigned width) {
-  FlatAffineValueConstraints cst;
-  SmallVector<Operation *, 8> ops(input.begin(), input.end());
-  (void)getIndexSet(ops, &cst);
-  if (!cst.isHyperRectangular(0, width)) {
-    rootAffineForOp.emitError("tiled code generation unimplemented for the "
-                              "non-hyperrectangular case");
-    return failure();
-  }
-  return success();
-}
-
 /// Set lower and upper bounds of intra-tile loops for parametric tiling.
 //  TODO: Handle non-constant lower bounds.
 static void setIntraTileBoundsParametric(OpBuilder &b, AffineForOp origLoop,
                                          AffineForOp newInterTileLoop,
                                          AffineForOp newIntraTileLoop,
                                          Value tileSize) {
   // The lower bound for the intra-tile loop is represented by an affine map
   // as (%i, %t0)->((%i - %origlb) * %t0 + %origlb). Similarly, the upper bound

 
 /// Tiles the specified band of perfectly nested loops creating tile-space loops
 /// and intra-tile loops. A band is a contiguous set of loops.
 //  TODO: handle non hyper-rectangular spaces.
 LogicalResult
 mlir::tilePerfectlyNested(MutableArrayRef<AffineForOp> input,
                           ArrayRef<unsigned> tileSizes,
                           SmallVectorImpl<AffineForOp> *tiledNest) {
-  performPreTilingChecks(input, tileSizes);
+  if (input.empty())
+    return success();
+
+  if (failed(performPreTilingChecks(input, tileSizes)))
+    return failure();
 
   MutableArrayRef<AffineForOp> origLoops = input;
   AffineForOp rootAffineForOp = origLoops[0];
-  // Note that width is at least one since band isn't empty.
+
+  // Note that width is at least one since the band isn't empty.
   unsigned width = input.size();
   SmallVector<AffineForOp, 6> tiledLoops(2 * width);
 
   // Construct a tiled loop nest without setting their bounds. Bounds are
   // set later.
   constructTiledLoopNest(origLoops, rootAffineForOp, width, tiledLoops);
 
   SmallVector<Value, 8> origLoopIVs;
   extractForInductionVars(input, &origLoopIVs);
 
-  if (failed(checkIfHyperRectangular(input, rootAffineForOp, width)))
-    return failure();
-
   // Set loop bounds for the tiled loop nest.
   constructTiledIndexSetHyperRect(origLoops, tiledLoops, tileSizes);
 
   // Replace original IVs with intra-tile loop IVs.
   for (unsigned i = 0; i < width; i++)
     origLoopIVs[i].replaceAllUsesWith(tiledLoops[i + width].getInductionVar());
 
   // Erase the old loop nest.
   rootAffineForOp.erase();
 
   if (tiledNest)
     *tiledNest = std::move(tiledLoops);
 
   return success();
 }
 
 /// Tiles the specified band of perfectly nested loops creating tile-space
 /// loops and intra-tile loops, using SSA values as tiling parameters. A band
 /// is a contiguous set of loops.
 //  TODO: handle non hyper-rectangular spaces.
 LogicalResult
 mlir::tilePerfectlyNestedParametric(MutableArrayRef<AffineForOp> input,
                                     ArrayRef<Value> tileSizes,
                                     SmallVectorImpl<AffineForOp> *tiledNest) {
-  performPreTilingChecks(input, tileSizes);
+  if (input.empty())
+    return success();
+
+  if (failed(performPreTilingChecks(input, tileSizes)))
+    return failure();
 
   MutableArrayRef<AffineForOp> origLoops = input;
   AffineForOp rootAffineForOp = origLoops[0];
-  // Note that width is at least one since band isn't empty.
   unsigned width = input.size();
   SmallVector<AffineForOp, 6> tiledLoops(2 * width);
 
   // Construct a tiled loop nest without setting their bounds. Bounds are
   // set later.
   constructTiledLoopNest(origLoops, rootAffineForOp, width, tiledLoops);
 
   SmallVector<Value, 8> origLoopIVs;
   extractForInductionVars(input, &origLoopIVs);
 
-  if (failed(checkIfHyperRectangular(input, rootAffineForOp, width)))
-    return failure();
-
   // Set loop bounds for the tiled loop nest.
   constructParametricallyTiledIndexSetHyperRect(origLoops, tiledLoops,
                                                 tileSizes);
 
   // Replace original IVs with intra-tile loop IVs.
   for (unsigned i = 0; i < width; i++)
     origLoopIVs[i].replaceAllUsesWith(tiledLoops[i + width].getInductionVar());
 

mlir/test/Dialect/Affine/loop-tiling-unsupported.mlir

-// RUN: mlir-opt %s -affine-loop-tile="tile-size=32" -split-input-file -verify-diagnostics
-
-// -----
-
-#ub = affine_map<(d0)[s0] -> (d0, s0)>
-func @non_hyperrect_loop() {
-  %N = arith.constant 128 : index
-  // expected-error@+1 {{tiled code generation unimplemented for the non-hyperrectangular case}}
-  affine.for %i = 0 to %N {
-    affine.for %j = 0 to min #ub(%i)[%N] {
-      affine.yield
-    }
-  }
-  return
-}

mlir/test/Dialect/Affine/loop-tiling-validity.mlir

 
 // -----
 
 // There are dependences along the diagonal of the 2d iteration space,
 // specifically, they are of direction (+, -).
 // The default tiling method (hyper-rect) will violate tiling legality.
 // We expect a remark that points that issue out to be emitted.
 
-// CHECK-LABEL: func @illegal_loop_with_diag_dependence
 func @illegal_loop_with_diag_dependence() {
   %A = memref.alloc() : memref<64x64xf32>
 
   affine.for %i = 0 to 64 {
-    // expected-remark@above {{tiled code is illegal due to dependences}}
+    // expected-remark@above {{tiling code is illegal 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>
       %2 = arith.addf %0, %1 : f32
       affine.store %2, %A[%i, %j] : memref<64x64xf32>
     }
   }
 
   return
 }

mlir/test/Dialect/Affine/loop-tiling.mlir

func @separate_full_tile_2d(%M : index, %N : index) {
   affine.for %i = 0 to %M {
     affine.for %j = 0 to %N {
       "test.foo"() : () -> ()
     }
   }
   return
 }
 
+// -----
+
+#ub = affine_map<(d0)[s0] -> (d0, s0)>
+// CHECK-LABEL: func @non_hyperrectangular_loop
+func @non_hyperrectangular_loop() {
+  %N = arith.constant 128 : index
+  affine.for %i = 0 to %N {
+    affine.for %j = 0 to min #ub(%i)[%N] {
+      "test.foo"() : () -> ()
+    }
+ }
+  // No tiling is performed here.
+  // CHECK:      arith.constant
+  // CHECK-NEXT: affine.for
+  // CHECK-NEXT:   affine.for
+  // CHECK-NEXT:     test.foo
+  return
+}
+
+// -----
+
+// No tiling supported on loops with yield values.
+
+// CHECK-LABEL: func @yield_values
+func @yield_values(%init : index) {
+  %r = affine.for %i = 0 to 10 iter_args(%s = %init) -> index {
+    "test.foo"() : () -> ()
+    affine.yield %s : index
+  }
+  // No tiling here.
+  // CHECK-NEXT: affine.for {{.*}} {
+  // CHECK-NEXT:   test.foo
+  return
+}
+
+// -----
+
 // SEPARATE-DAG: #[[$SEP_COND:.*]] = affine_set<(d0, d1)[s0, s1] : (-d0 + s0 - 32 >= 0, -d1 + s1 - 32 >= 0)>
 // SEPARATE-DAG: #[[$LB:.*]] = affine_map<(d0) -> (d0)>
 // SEPARATE-DAG: #[[$FULL_TILE_UB:.*]] = affine_map<(d0) -> (d0 + 32)>
 // SEPARATE-DAG: #[[$PART_TILE_UB:.*]] = affine_map<(d0)[s0] -> (d0 + 32, s0)>
 
 // SEPARATE-LABEL: func @separate_full_tile_2d(
 // SEPARATE: %[[M:.*]]: index, %[[N:.*]]: index
 

mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp

 }
 
 void TestAffineLoopParametricTiling::runOnFunction() {
   // Bands of loops to tile.
   std::vector<SmallVector<AffineForOp, 6>> bands;
   getTileableBands(getFunction(), &bands);
 
   // Tile each band.
-  for (SmallVectorImpl<AffineForOp> &band : bands) {
+  for (MutableArrayRef<AffineForOp> band : bands) {
     // Capture the tiling parameters from the arguments to the function
     // enclosing this loop nest.
     SmallVector<AffineForOp, 6> tiledNest;
     SmallVector<Value, 6> tilingParameters;
     // Check if tiling parameters are present.
     checkIfTilingParametersExist(band);
 
     // Get function arguments as tiling parameters.
     getTilingParameters(band, tilingParameters);
 
-    if (failed(
-            tilePerfectlyNestedParametric(band, tilingParameters, &tiledNest)))
-      return signalPassFailure();
+    (void)tilePerfectlyNestedParametric(band, tilingParameters, &tiledNest);
   }
 }
 
 namespace mlir {
 namespace test {
 void registerTestAffineLoopParametricTilingPass() {
   PassRegistration<TestAffineLoopParametricTiling>();
 }
 } // namespace test
 } // namespace mlir