[mlir][affine] Add SimplifyTrivialLoops pattern to AffineForOp #165091
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
| @llvm/pr-subscribers-mlir @llvm/pr-subscribers-mlir-affine Author: lonely eagle (linuxlonelyeagle) ChangesThis PR adds the SimplifyTrivialLoops pattern to affine.for, and removes the promote-single-iter functionality from affine-loop-normalize, since scf.for also has the capability to eliminate loops with a trip count of 1. Full diff: https://github.com/llvm/llvm-project/pull/165091.diff 10 Files Affected:
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td index 12a79358d42f1..e52b7d2090d53 100644 --- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td +++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td @@ -330,6 +330,7 @@ def AffineForOp : Affine_Op<"for", Speculation::Speculatability getSpeculatability(); }]; + let hasCanonicalizer = 1; let hasCustomAssemblyFormat = 1; let hasFolder = 1; let hasRegionVerifier = 1; diff --git a/mlir/include/mlir/Dialect/Affine/Passes.h b/mlir/include/mlir/Dialect/Affine/Passes.h index 2f70f24dd3ef2..0d88ecc462afd 100644 --- a/mlir/include/mlir/Dialect/Affine/Passes.h +++ b/mlir/include/mlir/Dialect/Affine/Passes.h @@ -58,8 +58,7 @@ std::unique_ptr<OperationPass<func::FuncOp>> createRaiseMemrefToAffine(); /// Apply normalization transformations to affine loop-like ops. If /// `promoteSingleIter` is true, single iteration loops are promoted (i.e., the /// loop is replaced by its loop body). -std::unique_ptr<OperationPass<func::FuncOp>> -createAffineLoopNormalizePass(bool promoteSingleIter = false); +std::unique_ptr<OperationPass<func::FuncOp>> createAffineLoopNormalizePass(); /// Performs packing (or explicit copying) of accessed memref regions into /// buffers in the specified faster memory space through either pointwise copies diff --git a/mlir/include/mlir/Dialect/Affine/Passes.td b/mlir/include/mlir/Dialect/Affine/Passes.td index 6ad45b828f657..74081772a8441 100644 --- a/mlir/include/mlir/Dialect/Affine/Passes.td +++ b/mlir/include/mlir/Dialect/Affine/Passes.td @@ -383,10 +383,6 @@ def AffineParallelize : Pass<"affine-parallelize", "func::FuncOp"> { def AffineLoopNormalize : Pass<"affine-loop-normalize", "func::FuncOp"> { let summary = "Apply normalization transformations to affine loop-like ops"; let constructor = "mlir::affine::createAffineLoopNormalizePass()"; - let options = [ - Option<"promoteSingleIter", "promote-single-iter", "bool", - /*default=*/"true", "Promote single iteration loops">, - ]; } def LoopCoalescing : Pass<"affine-loop-coalescing", "func::FuncOp"> { diff --git a/mlir/include/mlir/Dialect/Affine/Utils.h b/mlir/include/mlir/Dialect/Affine/Utils.h index ac11f5a7c24c7..f46f40de4bcc7 100644 --- a/mlir/include/mlir/Dialect/Affine/Utils.h +++ b/mlir/include/mlir/Dialect/Affine/Utils.h @@ -171,11 +171,8 @@ void normalizeAffineParallel(AffineParallelOp op); /// lower bound to zero and loop step to one. The upper bound is set to the trip /// count of the loop. Original loops must have a lower bound with only a single /// result. There is no such restriction on upper bounds. Returns success if the -/// loop has been normalized (or is already in the normal form). If -/// `promoteSingleIter` is true, the loop is simply promoted if it has a single -/// iteration. -LogicalResult normalizeAffineFor(AffineForOp op, - bool promoteSingleIter = false); +/// loop has been normalized (or is already in the normal form). +LogicalResult normalizeAffineFor(AffineForOp op); /// Traverse `e` and return an AffineExpr where all occurrences of `dim` have /// been replaced by either: diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp index e0a53cd52f143..1aea482a5bcea 100644 --- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp +++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp @@ -2716,6 +2716,46 @@ LogicalResult AffineForOp::fold(FoldAdaptor adaptor, return success(folded); } +/// Replaces the given op with the contents of the given single-block region, +/// using the operands of the block terminator to replace operation results. +static void replaceOpWithRegion(PatternRewriter &rewriter, Operation *op, + Region ®ion, ValueRange blockArgs = {}) { + assert(region.hasOneBlock() && "expected single-block region"); + Block *block = ®ion.front(); + Operation *terminator = block->getTerminator(); + ValueRange results = terminator->getOperands(); + rewriter.inlineBlockBefore(block, op, blockArgs); + rewriter.replaceOp(op, results); + rewriter.eraseOp(terminator); +} + +struct SimplifyTrivialLoops : public OpRewritePattern<AffineForOp> { + using OpRewritePattern<AffineForOp>::OpRewritePattern; + LogicalResult matchAndRewrite(AffineForOp forOp, + PatternRewriter &rewriter) const override { + std::optional<uint64_t> tripCount = getTrivialConstantTripCount(forOp); + if (!tripCount.has_value() || tripCount != 1) + return failure(); + + SmallVector<Value> blockArgs; + blockArgs.reserve(forOp.getInits().size() + 1); + rewriter.setInsertionPointToStart(forOp.getBody()); + Value lower = + rewriter.create<AffineApplyOp>(forOp.getLoc(), forOp.getLowerBoundMap(), + forOp.getLowerBoundOperands()); + forOp.getInductionVar().replaceAllUsesWith(lower); + blockArgs.push_back(lower); + llvm::append_range(blockArgs, forOp.getInits()); + replaceOpWithRegion(rewriter, forOp, forOp.getRegion(), blockArgs); + return success(); + } +}; + +void AffineForOp::getCanonicalizationPatterns(RewritePatternSet &results, + MLIRContext *context) { + results.add<SimplifyTrivialLoops>(context); +} + OperandRange AffineForOp::getEntrySuccessorOperands(RegionBranchPoint point) { assert((point.isParent() || point == getRegion()) && "invalid region point"); diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp index 5cc38f7051726..4312a733c95c5 100644 --- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp +++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopNormalize.cpp @@ -33,16 +33,14 @@ namespace { /// that are already in a normalized form. struct AffineLoopNormalizePass : public affine::impl::AffineLoopNormalizeBase<AffineLoopNormalizePass> { - explicit AffineLoopNormalizePass(bool promoteSingleIter) { - this->promoteSingleIter = promoteSingleIter; - } + using Base::Base; void runOnOperation() override { getOperation().walk([&](Operation *op) { if (auto affineParallel = dyn_cast<AffineParallelOp>(op)) normalizeAffineParallel(affineParallel); else if (auto affineFor = dyn_cast<AffineForOp>(op)) - (void)normalizeAffineFor(affineFor, promoteSingleIter); + (void)normalizeAffineFor(affineFor); }); } }; @@ -50,6 +48,6 @@ struct AffineLoopNormalizePass } // namespace std::unique_ptr<OperationPass<func::FuncOp>> -mlir::affine::createAffineLoopNormalizePass(bool promoteSingleIter) { - return std::make_unique<AffineLoopNormalizePass>(promoteSingleIter); +mlir::affine::createAffineLoopNormalizePass() { + return std::make_unique<AffineLoopNormalizePass>(); } diff --git a/mlir/lib/Dialect/Affine/Utils/Utils.cpp b/mlir/lib/Dialect/Affine/Utils/Utils.cpp index 845be20d15b69..445617c8a0e52 100644 --- a/mlir/lib/Dialect/Affine/Utils/Utils.cpp +++ b/mlir/lib/Dialect/Affine/Utils/Utils.cpp @@ -556,10 +556,7 @@ void mlir::affine::normalizeAffineParallel(AffineParallelOp op) { op.setUpperBounds(ranges.getOperands(), newUpperMap); } -LogicalResult mlir::affine::normalizeAffineFor(AffineForOp op, - bool promoteSingleIter) { - if (promoteSingleIter && succeeded(promoteIfSingleIteration(op))) - return success(); +LogicalResult mlir::affine::normalizeAffineFor(AffineForOp op) { // Check if the forop is already normalized. if (op.hasConstantLowerBound() && (op.getConstantLowerBound() == 0) && diff --git a/mlir/test/Dialect/Affine/affine-loop-normalize.mlir b/mlir/test/Dialect/Affine/affine-loop-normalize.mlir index 7d90efec0c21b..a25888d50e6b9 100644 --- a/mlir/test/Dialect/Affine/affine-loop-normalize.mlir +++ b/mlir/test/Dialect/Affine/affine-loop-normalize.mlir @@ -1,5 +1,4 @@ // RUN: mlir-opt %s -affine-loop-normalize -split-input-file | FileCheck %s -// RUN: mlir-opt %s -affine-loop-normalize='promote-single-iter=1' -split-input-file | FileCheck %s --check-prefix=PROMOTE-SINGLE-ITER // Normalize steps to 1 and lower bounds to 0. @@ -37,26 +36,6 @@ func.func @relative_bounds(%arg: index) { // ----- -// Check that single iteration loop is removed and its body is promoted to the -// parent block. - -// CHECK-LABEL: func @promote_single_iter_loop -// PROMOTE-SINGLE-ITER-LABEL: func @promote_single_iter_loop -func.func @promote_single_iter_loop(%in: memref<1xf32>, %out: memref<1xf32>) { - affine.for %i = 0 to 1 { - %1 = affine.load %in[%i] : memref<1xf32> - affine.store %1, %out[%i] : memref<1xf32> - } - return -} - -// PROMOTE-SINGLE-ITER-NEXT: arith.constant -// PROMOTE-SINGLE-ITER-NEXT: affine.load -// PROMOTE-SINGLE-ITER-NEXT: affine.store -// PROMOTE-SINGLE-ITER-NEXT: return - -// ----- - // CHECK-DAG: [[$IV0:#map[0-9]*]] = affine_map<(d0) -> (d0 * 2 + 2)> // CHECK-DAG: [[$IV1:#map[0-9]*]] = affine_map<(d0) -> (d0 * 3)> diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir index 1169cd1c29d74..c99d3772f1648 100644 --- a/mlir/test/Dialect/Affine/canonicalize.mlir +++ b/mlir/test/Dialect/Affine/canonicalize.mlir @@ -1323,12 +1323,10 @@ func.func @simplify_bounds_tiled() { } } } - // CHECK: affine.for - // CHECK-NEXT: affine.for + // CHECK: affine.for %{{.*}} 0 to 2 + // CHECK-NEXT: affine.for %{{.*}} = 0 to 32 step 16 // CHECK-NEXT: affine.for %{{.*}} = 0 to 32 step 16 - // CHECK-NEXT: affine.for %{{.*}} = 0 to 32 step 16 - // CHECK-NEXT: affine.for %{{.*}} = 0 to 2 - // CHECK-NEXT: affine.for %{{.*}} = 0 to 16 step 16 + // CHECK-NEXT: affine.for %{{.*}} = 0 to 2 return } @@ -1348,9 +1346,7 @@ func.func @simplify_min_max_multi_expr() { // CHECK: affine.for affine.for %i = 0 to 2 { // CHECK: affine.for - affine.for %j = 0 to 4 { - // The first upper bound expression will not be lower than -9. So, it's redundant. - // CHECK-NEXT: affine.for %{{.*}} = -10 to -9 + affine.for %j = 0 to 4 { affine.for %k = -10 to min affine_map<(d0, d1) -> (4 * d0 - 3 * d1, -9)>(%i, %j) { "test.foo"() : () -> () } @@ -1370,7 +1366,6 @@ func.func @simplify_min_max_multi_expr() { } } - // CHECK: affine.for %{{.*}} = 0 to 1 affine.for %i = 0 to 2 { affine.for %j = max affine_map<(d0) -> (d0 floordiv 2, 0)>(%i) to 1 { "test.foo"() : () -> () @@ -2401,3 +2396,22 @@ func.func @for_empty_body_folder_iv_yield() -> index { } return %10 : index } + +// ----- + +// Check that single iteration loop is removed and its body is promoted to the +// parent block. + +// CHECK-LABEL: func @promote_single_iter_loop +func.func @promote_single_iter_loop(%in: memref<1xf32>, %out: memref<1xf32>) { + affine.for %i = 0 to 1 { + %1 = affine.load %in[%i] : memref<1xf32> + affine.store %1, %out[%i] : memref<1xf32> + } + return +} + +// CHECK-NEXT: arith.constant +// CHECK-NEXT: affine.load +// CHECK-NEXT: affine.store +// CHECK-NEXT: return diff --git a/mlir/test/Dialect/Affine/raise-memref.mlir b/mlir/test/Dialect/Affine/raise-memref.mlir index 8dc24d99db3e2..c9777de2718ae 100644 --- a/mlir/test/Dialect/Affine/raise-memref.mlir +++ b/mlir/test/Dialect/Affine/raise-memref.mlir @@ -51,21 +51,17 @@ func.func @reduce_window_max() { // CHECK: affine.for %[[arg0:.*]] = // CHECK: affine.for %[[arg1:.*]] = // CHECK: affine.for %[[arg2:.*]] = -// CHECK: affine.for %[[arg3:.*]] = -// CHECK: affine.store %[[cst]], %[[v0]][%[[arg0]], %[[arg1]], %[[arg2]], %[[arg3]]] : +// CHECK: affine.store %[[cst]], %[[v0]][0, %[[arg0]], %[[arg1]], %[[arg2]]] : // CHECK: affine.for %[[a0:.*]] = // CHECK: affine.for %[[a1:.*]] = // CHECK: affine.for %[[a2:.*]] = // CHECK: affine.for %[[a3:.*]] = // CHECK: affine.for %[[a4:.*]] = -// CHECK: affine.for %[[a5:.*]] = -// CHECK: affine.for %[[a6:.*]] = -// CHECK: affine.for %[[a7:.*]] = -// CHECK: %[[lhs:.*]] = affine.load %[[v0]][%[[a0]], %[[a1]], %[[a2]], %[[a3]]] : -// CHECK: %[[rhs:.*]] = affine.load %[[v1]][%[[a0]] + %[[a4]], %[[a1]] * 2 + %[[a5]], %[[a2]] * 2 + %[[a6]], %[[a3]] + %[[a7]]] : -// CHECK: %[[res:.*]] = arith.cmpf ogt, %[[lhs]], %[[rhs]] : f32 -// CHECK: %[[sel:.*]] = arith.select %[[res]], %[[lhs]], %[[rhs]] : f32 -// CHECK: affine.store %[[sel]], %[[v0]][%[[a0]], %[[a1]], %[[a2]], %[[a3]]] : +// CHECK: %[[lhs:.*]] = affine.load %[[v0]][0, %[[a0]], %[[a1]], %[[a2]]] : +// CHECK: %[[rhs:.*]] = affine.load %[[v1]][0, %[[a0]] * 2 + %[[a3]], %[[a1]] * 2 + %[[a4]], %[[a2]]] : +// CHECK: %[[res:.*]] = arith.cmpf ogt, %[[lhs]], %[[rhs]] : f32 +// CHECK: %[[sel:.*]] = arith.select %[[res]], %[[lhs]], %[[rhs]] : f32 +// CHECK: affine.store %[[sel]], %[[v0]][0, %[[a0]], %[[a1]], %[[a2]]] : // CHECK-LABEL: func @symbols( func.func @symbols(%N : index) { |
linuxlonelyeagle requested review from bondhugula, ftynse, joker-eph, krzysz00 and kuhar and removed request for kuhar Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit. This suggestion is invalid because no changes were made to the code. Suggestions cannot be applied while the pull request is closed. Suggestions cannot be applied while viewing a subset of changes. Only one suggestion per line can be applied in a batch. Add this suggestion to a batch that can be applied as a single commit. Applying suggestions on deleted lines is not supported. You must change the existing code in this line in order to create a valid suggestion. Outdated suggestions cannot be applied. This suggestion has been applied or marked resolved. Suggestions cannot be applied from pending reviews. Suggestions cannot be applied on multi-line comments. Suggestions cannot be applied while the pull request is queued to merge. Suggestion cannot be applied right now. Please check back later.
This PR adds the SimplifyTrivialLoops pattern to affine.for, and removes the promote-single-iter functionality from affine-loop-normalize, since scf.for also has the capability to eliminate loops with a trip count of 1.