Multiple fixes to affine loop tiling return status and checks

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.

Differential Revision: https://reviews.llvm.org/D116738

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

@@ -178,14 +178,23 @@ void LoopTiling::runOnFunction() {
       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"));
+      }
     }
   }
 }

mlir/lib/Transforms/Utils/LoopUtils.cpp

@@ -510,20 +510,55 @@ 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,
-                            ArrayRef<t> tileSizes) {
-  // Check if the supplied for op's are all successively nested.
-  assert(!input.empty() && "no loops in input band");
+LogicalResult performPreTilingChecks(MutableArrayRef<AffineForOp> input,
+                                     ArrayRef<t> tileSizes) {
   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();
+  }
 
-  // Perform tiling legality test.
-  if (failed(checkTilingLegality(input)))
-    input[0].emitRemark("tiled code is illegal due to dependences");
+  // 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();
+  }
+
+  return success();
 }
 
 /// Move the loop body of AffineForOp 'src' from 'src' into the specified
@@ -582,21 +617,6 @@ void constructTiledLoopNest(MutableArrayRef<AffineForOp> origLoops,
   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,
@@ -912,11 +932,16 @@ 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);
 
@@ -927,9 +952,6 @@ mlir::tilePerfectlyNested(MutableArrayRef<AffineForOp> input,
   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);
 
@@ -954,11 +976,14 @@ 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);
 
@@ -969,9 +994,6 @@ mlir::tilePerfectlyNestedParametric(MutableArrayRef<AffineForOp> input,
   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);

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

@@ -1,15 +0,0 @@
-// 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

@@ -30,12 +30,11 @@ func @legal_loop() {
 // 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>

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

@@ -240,6 +240,43 @@ func @separate_full_tile_2d(%M : index, %N : index) {
   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)>

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

@@ -67,7 +67,7 @@ void TestAffineLoopParametricTiling::runOnFunction() {
   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;
@@ -78,9 +78,7 @@ void TestAffineLoopParametricTiling::runOnFunction() {
     // Get function arguments as tiling parameters.
     getTilingParameters(band, tilingParameters);
 
-    if (failed(
-            tilePerfectlyNestedParametric(band, tilingParameters, &tiledNest)))
-      return signalPassFailure();
+    (void)tilePerfectlyNestedParametric(band, tilingParameters, &tiledNest);
   }
 }