[Mlir-commits] [mlir] [mlir][affine] Add fold logic when the affine.yield has IV as operand in the AffineForEmptyLoopFolder (PR #164064)
lonely eagle
llvmlistbot at llvm.org
Sat Oct 18 10:43:09 PDT 2025
https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/164064
>From 539615208126bdf31b4e2fcc2ee5bf8ba6da79cc Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sat, 18 Oct 2025 08:53:14 +0000
Subject: [PATCH 1/4] add loop yield iv logic into AffineForEmptyLoopFolder.
---
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 28 ++++++++++++++++++----
mlir/test/Dialect/Affine/canonicalize.mlir | 13 ++++++++++
2 files changed, 37 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index e0a53cd52f143..768ac0fe55823 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2610,6 +2610,19 @@ static std::optional<uint64_t> getTrivialConstantTripCount(AffineForOp forOp) {
return ub - lb <= 0 ? 0 : (ub - lb + step - 1) / step;
}
+/// Calculate the constant value of the loop's induction variable for its last
+/// trip, construct an OpFoldResult using this value and return it.
+static OpFoldResult getConstantInductionVarForLastTrip(AffineForOp forOp) {
+ std::optional<uint64_t> tripCount = getTrivialConstantTripCount(forOp);
+ if (!tripCount.has_value())
+ return {};
+ int64_t lb = forOp.getConstantLowerBound();
+ int64_t step = forOp.getStepAsInt();
+ int64_t lastTripIv = lb + (tripCount.value() - 1) * step;
+ return OpFoldResult(
+ IntegerAttr::get(IndexType::get(forOp.getContext()), lastTripIv));
+}
+
/// Fold the empty loop.
static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
if (!llvm::hasSingleElement(*forOp.getBody()))
@@ -2622,7 +2635,7 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
// results.
return forOp.getInits();
}
- SmallVector<Value, 4> replacements;
+ SmallVector<OpFoldResult, 4> replacements;
auto yieldOp = cast<AffineYieldOp>(forOp.getBody()->getTerminator());
auto iterArgs = forOp.getRegionIterArgs();
bool hasValDefinedOutsideLoop = false;
@@ -2632,8 +2645,15 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
BlockArgument *iterArgIt = llvm::find(iterArgs, val);
// TODO: It should be possible to perform a replacement by computing the
// last value of the IV based on the bounds and the step.
- if (val == forOp.getInductionVar())
- return {};
+ if (val == forOp.getInductionVar()) {
+ OpFoldResult lastTripIv = getConstantInductionVarForLastTrip(forOp);
+ if (lastTripIv) {
+ replacements.push_back(lastTripIv);
+ continue;
+ } else {
+ return {};
+ }
+ }
if (iterArgIt == iterArgs.end()) {
// `val` is defined outside of the loop.
assert(forOp.isDefinedOutsideOfLoop(val) &&
@@ -2656,7 +2676,7 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
// out of order.
if (tripCount.has_value() && tripCount.value() >= 2 && iterArgsNotInOrder)
return {};
- return llvm::to_vector_of<OpFoldResult>(replacements);
+ return replacements;
}
/// Canonicalize the bounds of the given loop.
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 1169cd1c29d74..997f23b4bd669 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -609,6 +609,19 @@ func.func @fold_zero_iter_loops(%in : index) -> index {
// -----
+// CHECK-LABEL: func @fold_empty_loop_iv
+// CHECK-SAME: %[[INIT:.*]]: index
+func.func @fold_empty_loop_iv(%init: index) -> (index, index) {
+ %res:2 = affine.for %i = 0 to 10 step 1 iter_args(%arg0 = %init, %arg1 = %init) -> (index, index) {
+ affine.yield %i, %arg1 : index, index
+ }
+ // CHECK: %[[C9:.*]] = arith.constant 9 : index
+ // CHECK: return %[[C9]], %[[INIT]] : index, index
+ return %res#0, %res#1 : index, index
+}
+
+// -----
+
// CHECK-DAG: #[[$SET:.*]] = affine_set<(d0, d1)[s0] : (d0 >= 0, -d0 + 1022 >= 0, d1 >= 0, -d1 + s0 - 2 >= 0)>
// CHECK-LABEL: func @canonicalize_affine_if
>From 07393f37be8a5cef8b88e70a9ea21a8a0a57b1c9 Mon Sep 17 00:00:00 2001
From: lonely eagle <2020382038 at qq.com>
Date: Sun, 19 Oct 2025 01:26:00 +0800
Subject: [PATCH 2/4] Update mlir/lib/Dialect/Affine/IR/AffineOps.cpp
Co-authored-by: Jakub Kuderski <kubakuderski at gmail.com>
---
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 768ac0fe55823..769b3be11f52d 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2646,8 +2646,7 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
// TODO: It should be possible to perform a replacement by computing the
// last value of the IV based on the bounds and the step.
if (val == forOp.getInductionVar()) {
- OpFoldResult lastTripIv = getConstantInductionVarForLastTrip(forOp);
- if (lastTripIv) {
+ if (OpFoldResult lastTripIv = getConstantInductionVarForLastTrip(forOp)) {
replacements.push_back(lastTripIv);
continue;
} else {
>From d043f6bf20218cea56918edb98638eb0a2529d85 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sat, 18 Oct 2025 17:27:31 +0000
Subject: [PATCH 3/4] update.
---
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 769b3be11f52d..b08d0838b750c 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2649,9 +2649,8 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
if (OpFoldResult lastTripIv = getConstantInductionVarForLastTrip(forOp)) {
replacements.push_back(lastTripIv);
continue;
- } else {
- return {};
}
+ return {};
}
if (iterArgIt == iterArgs.end()) {
// `val` is defined outside of the loop.
>From 3d3b3a922d7b8e84d5c3054134e86750a85287ed Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sat, 18 Oct 2025 17:42:53 +0000
Subject: [PATCH 4/4] Process trip count equals zero.
---
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 2 ++
1 file changed, 2 insertions(+)
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index b08d0838b750c..90d47766ba572 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2616,6 +2616,8 @@ static OpFoldResult getConstantInductionVarForLastTrip(AffineForOp forOp) {
std::optional<uint64_t> tripCount = getTrivialConstantTripCount(forOp);
if (!tripCount.has_value())
return {};
+ if (tripCount.value() == 0)
+ return {};
int64_t lb = forOp.getConstantLowerBound();
int64_t step = forOp.getStepAsInt();
int64_t lastTripIv = lb + (tripCount.value() - 1) * step;
More information about the Mlir-commits
mailing list