[Mlir-commits] [mlir] [mlir][affine] Add fold logic when the affine.yield has IV as operand in the AffineForEmptyLoopFolder (PR #165639)
lonely eagle
llvmlistbot at llvm.org
Thu Oct 30 10:56:56 PDT 2025
https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/165639
>From a4e8857d318999f5b196571a36d73cc2c5d17fb4 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/7] 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 0fa939c9ed1e63798c6230ca3a6a809f4620037e 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/7] 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 e4ebc2b4be89aa7d783d86693098b769ad554001 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/7] 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 b0ce9fff62dff7e32a57f42f0f6b2bc2c1c3f74f 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/7] 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;
>From c816d094b00d9a6e727e68cfb834994eda161c87 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Tue, 28 Oct 2025 17:06:05 +0000
Subject: [PATCH 5/7] rebase main and update code.
---
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 19 +++++++++----------
1 file changed, 9 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 90d47766ba572..fd5b37448b421 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2611,18 +2611,18 @@ static std::optional<uint64_t> getTrivialConstantTripCount(AffineForOp forOp) {
}
/// 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) {
+/// trip.
+static std::optional<int64_t>
+getConstantInductionVarForLastTrip(AffineForOp forOp) {
std::optional<uint64_t> tripCount = getTrivialConstantTripCount(forOp);
if (!tripCount.has_value())
- return {};
+ return std::nullopt;
if (tripCount.value() == 0)
- return {};
+ return std::nullopt;
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));
+ return lastTripIv;
}
/// Fold the empty loop.
@@ -2645,11 +2645,10 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
for (unsigned i = 0, e = yieldOp->getNumOperands(); i < e; ++i) {
Value val = yieldOp.getOperand(i);
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()) {
- if (OpFoldResult lastTripIv = getConstantInductionVarForLastTrip(forOp)) {
- replacements.push_back(lastTripIv);
+ if (auto lastTripIv = getConstantInductionVarForLastTrip(forOp)) {
+ replacements.push_back(IntegerAttr::get(
+ IndexType::get(forOp.getContext()), lastTripIv.value()));
continue;
}
return {};
>From 86c0c0b45401f452b19b92fd32da13b69f0e2a52 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Thu, 30 Oct 2025 01:48:50 +0000
Subject: [PATCH 6/7] make replacements from SmallVector<OpFoldResult,4> to
SmallVector<OpFoldResult> type.
---
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index fd5b37448b421..feb95d91a4b80 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2637,7 +2637,7 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
// results.
return forOp.getInits();
}
- SmallVector<OpFoldResult, 4> replacements;
+ SmallVector<OpFoldResult> replacements;
auto yieldOp = cast<AffineYieldOp>(forOp.getBody()->getTerminator());
auto iterArgs = forOp.getRegionIterArgs();
bool hasValDefinedOutsideLoop = false;
>From 22b96ba5aca5ced54d3c6c624ec72edf50f6b54d Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Thu, 30 Oct 2025 17:56:38 +0000
Subject: [PATCH 7/7] use llvm::to_vector return replacements.
---
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 4 ++--
1 file changed, 2 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index feb95d91a4b80..be9ca0657a98f 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2637,7 +2637,7 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
// results.
return forOp.getInits();
}
- SmallVector<OpFoldResult> replacements;
+ SmallVector<OpFoldResult, 4> replacements;
auto yieldOp = cast<AffineYieldOp>(forOp.getBody()->getTerminator());
auto iterArgs = forOp.getRegionIterArgs();
bool hasValDefinedOutsideLoop = false;
@@ -2675,7 +2675,7 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
// out of order.
if (tripCount.has_value() && tripCount.value() >= 2 && iterArgsNotInOrder)
return {};
- return replacements;
+ return llvm::to_vector(replacements);
}
/// Canonicalize the bounds of the given loop.
More information about the Mlir-commits
mailing list