[Mlir-commits] [mlir] 9cf9ed9 - Multiple fixes to affine loop tiling return status and checks
Uday Bondhugula
llvmlistbot at llvm.org
Sat Jan 8 03:22:26 PST 2022
Author: Uday Bondhugula
Date: 2022-01-08T16:50:44+05:30
New Revision: 9cf9ed94ed3f22c4a9e6c5ddc8faba182e324933
URL: https://github.com/llvm/llvm-project/commit/9cf9ed94ed3f22c4a9e6c5ddc8faba182e324933
DIFF: https://github.com/llvm/llvm-project/commit/9cf9ed94ed3f22c4a9e6c5ddc8faba182e324933.diff
LOG: 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
Added:
Modified:
mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
mlir/lib/Transforms/Utils/LoopUtils.cpp
mlir/test/Dialect/Affine/loop-tiling-validity.mlir
mlir/test/Dialect/Affine/loop-tiling.mlir
mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp
Removed:
mlir/test/Dialect/Affine/loop-tiling-unsupported.mlir
################################################################################
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
index 69d21c66e1a48..52b9fc6fcd03c 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp
+++ b/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"));
+ }
}
}
}
diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index 6328b59d9008e..5746c4a588eb5 100644
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/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();
+ }
+
+ // 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
@@ -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);
diff --git a/mlir/test/Dialect/Affine/loop-tiling-unsupported.mlir b/mlir/test/Dialect/Affine/loop-tiling-unsupported.mlir
deleted file mode 100644
index 03f9866813dac..0000000000000
--- a/mlir/test/Dialect/Affine/loop-tiling-unsupported.mlir
+++ /dev/null
@@ -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 at +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
-}
diff --git a/mlir/test/Dialect/Affine/loop-tiling-validity.mlir b/mlir/test/Dialect/Affine/loop-tiling-validity.mlir
index dd08e8c334fd3..2b67f223f01c2 100644
--- a/mlir/test/Dialect/Affine/loop-tiling-validity.mlir
+++ b/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 at above {{tiled code is illegal due to dependences}}
+ // expected-remark at 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>
diff --git a/mlir/test/Dialect/Affine/loop-tiling.mlir b/mlir/test/Dialect/Affine/loop-tiling.mlir
index 312ac41980e20..8ff91ecff7038 100644
--- a/mlir/test/Dialect/Affine/loop-tiling.mlir
+++ b/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)>
diff --git a/mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp b/mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp
index a0e0082db481b..6038ec3b25f4e 100644
--- a/mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp
+++ b/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);
}
}
More information about the Mlir-commits
mailing list