[Mlir-commits] [mlir] [mlir][affine] Add SimplifyTrivialLoops pattern to AffineForOp (PR #165091)
lonely eagle
llvmlistbot at llvm.org
Sat Oct 25 04:03:22 PDT 2025
https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/165091
>From 80a7348b2e3912d4e3c9b61cb832ccdc0c80ecac Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sat, 25 Oct 2025 10:29:48 +0000
Subject: [PATCH 1/2] add SimplifyTrivialLoops pattern to AffineForOp.
---
.../mlir/Dialect/Affine/IR/AffineOps.td | 1 +
mlir/include/mlir/Dialect/Affine/Passes.h | 3 +-
mlir/include/mlir/Dialect/Affine/Passes.td | 4 --
mlir/include/mlir/Dialect/Affine/Utils.h | 7 +---
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 40 +++++++++++++++++++
.../Affine/Transforms/AffineLoopNormalize.cpp | 10 ++---
mlir/lib/Dialect/Affine/Utils/Utils.cpp | 5 +--
.../Dialect/Affine/affine-loop-normalize.mlir | 21 ----------
mlir/test/Dialect/Affine/canonicalize.mlir | 32 ++++++++++-----
mlir/test/Dialect/Affine/raise-memref.mlir | 16 +++-----
10 files changed, 78 insertions(+), 61 deletions(-)
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) {
>From 0cc8b06e847604182ca31c2cc4ae7e42426f57f3 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sat, 25 Oct 2025 11:03:07 +0000
Subject: [PATCH 2/2] fix compiler error.
---
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 6 +++---
mlir/test/Dialect/Affine/canonicalize.mlir | 7 +++----
2 files changed, 6 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 1aea482a5bcea..db1e6084093b0 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2740,9 +2740,9 @@ struct SimplifyTrivialLoops : public OpRewritePattern<AffineForOp> {
SmallVector<Value> blockArgs;
blockArgs.reserve(forOp.getInits().size() + 1);
rewriter.setInsertionPointToStart(forOp.getBody());
- Value lower =
- rewriter.create<AffineApplyOp>(forOp.getLoc(), forOp.getLowerBoundMap(),
- forOp.getLowerBoundOperands());
+ Value lower = AffineApplyOp::create(
+ rewriter, forOp.getLoc(), forOp.getLowerBoundMap(),
+ ValueRange(forOp.getLowerBoundOperands()));
forOp.getInductionVar().replaceAllUsesWith(lower);
blockArgs.push_back(lower);
llvm::append_range(blockArgs, forOp.getInits());
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index c99d3772f1648..2f0fcc804bd59 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -2403,6 +2403,7 @@ func.func @for_empty_body_folder_iv_yield() -> index {
// parent block.
// CHECK-LABEL: func @promote_single_iter_loop
+// CHECK-SAME: %[[IN:.*]]: memref<1xf32>, %[[OUT:.*]]: memref<1xf32>
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>
@@ -2411,7 +2412,5 @@ func.func @promote_single_iter_loop(%in: memref<1xf32>, %out: memref<1xf32>) {
return
}
-// CHECK-NEXT: arith.constant
-// CHECK-NEXT: affine.load
-// CHECK-NEXT: affine.store
-// CHECK-NEXT: return
+// CHECK-NEXT: %[[DATA:.*]] = affine.load %[[IN]][0]
+// CHECK-NEXT: affine.store %[[DATA]], %[[OUT]][0]
More information about the Mlir-commits
mailing list