diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp --- a/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp +++ b/mlir/lib/Dialect/Affine/Transforms/LoopTiling.cpp @@ -178,14 +178,23 @@ diag << "]\n"; } SmallVector 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(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 --- a/mlir/lib/Transforms/Utils/LoopUtils.cpp +++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp @@ -513,17 +513,30 @@ /// Check if the input data is valid and whether tiled code will be legal or /// not. template -void performPreTilingChecks(MutableArrayRef input, - ArrayRef tileSizes) { +LogicalResult performPreTilingChecks(MutableArrayRef input, + ArrayRef 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(); + } + + if (!isPerfectlyNested(input)) { + LLVM_DEBUG(llvm::dbgs() << "input loops not perfectly nested"); + return failure(); + } // Perform tiling legality test. - if (failed(checkTilingLegality(input))) - input[0].emitRemark("tiled code is illegal due to dependences"); + 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 @@ -588,10 +601,16 @@ unsigned width) { FlatAffineValueConstraints cst; SmallVector ops(input.begin(), input.end()); - (void)getIndexSet(ops, &cst); + // 1-d is hyper-rectangular by default. + 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, width)) { - rootAffineForOp.emitError("tiled code generation unimplemented for the " - "non-hyperrectangular case"); + LLVM_DEBUG(llvm::dbgs() + << "Non-hyperrectangular nests not supported for tiling!\n"); return failure(); } return success(); @@ -912,12 +931,19 @@ mlir::tilePerfectlyNested(MutableArrayRef input, ArrayRef tileSizes, SmallVectorImpl *tiledNest) { - performPreTilingChecks(input, tileSizes); + if (input.empty()) + return success(); + + if (failed(performPreTilingChecks(input, tileSizes))) + return failure(); MutableArrayRef 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(); + if (failed(checkIfHyperRectangular(input, rootAffineForOp, width))) + return failure(); + SmallVector tiledLoops(2 * width); // Construct a tiled loop nest without setting their bounds. Bounds are @@ -927,9 +953,6 @@ SmallVector 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,12 +977,16 @@ mlir::tilePerfectlyNestedParametric(MutableArrayRef input, ArrayRef tileSizes, SmallVectorImpl *tiledNest) { - performPreTilingChecks(input, tileSizes); + if (failed(performPreTilingChecks(input, tileSizes))) + return failure(); MutableArrayRef origLoops = input; AffineForOp rootAffineForOp = origLoops[0]; // Note that width is at least one since band isn't empty. unsigned width = input.size(); + if (failed(checkIfHyperRectangular(input, rootAffineForOp, width))) + return failure(); + SmallVector tiledLoops(2 * width); // Construct a tiled loop nest without setting their bounds. Bounds are @@ -969,9 +996,6 @@ SmallVector 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 --- 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@+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 --- a/mlir/test/Dialect/Affine/loop-tiling-validity.mlir +++ b/mlir/test/Dialect/Affine/loop-tiling-validity.mlir @@ -30,12 +30,11 @@ // 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> diff --git a/mlir/test/Dialect/Affine/loop-tiling.mlir b/mlir/test/Dialect/Affine/loop-tiling.mlir --- a/mlir/test/Dialect/Affine/loop-tiling.mlir +++ b/mlir/test/Dialect/Affine/loop-tiling.mlir @@ -240,6 +240,43 @@ 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 --- a/mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp +++ b/mlir/test/lib/Dialect/Affine/TestAffineLoopParametricTiling.cpp @@ -67,7 +67,7 @@ getTileableBands(getFunction(), &bands); // Tile each band. - for (SmallVectorImpl &band : bands) { + for (MutableArrayRef band : bands) { // Capture the tiling parameters from the arguments to the function // enclosing this loop nest. SmallVector tiledNest; @@ -78,9 +78,7 @@ // Get function arguments as tiling parameters. getTilingParameters(band, tilingParameters); - if (failed( - tilePerfectlyNestedParametric(band, tilingParameters, &tiledNest))) - return signalPassFailure(); + (void)tilePerfectlyNestedParametric(band, tilingParameters, &tiledNest); } }