[Mlir-commits] [mlir] [mlir][affine] Make AffineForEmptyLoopFolder as folder function (PR #163929)
lonely eagle
llvmlistbot at llvm.org
Fri Oct 17 08:18:20 PDT 2025
https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/163929
>From 1fa00379d916b474c2432f3af071d44e230925a5 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Fri, 17 Oct 2025 08:57:58 +0000
Subject: [PATCH 1/3] make AffineForEmptyLoopFolder as folder function.
---
.../mlir/Dialect/Affine/IR/AffineOps.td | 1 -
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 137 ++++++++----------
2 files changed, 62 insertions(+), 76 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
index e52b7d2090d53..12a79358d42f1 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineOps.td
@@ -330,7 +330,6 @@ def AffineForOp : Affine_Op<"for",
Speculation::Speculatability getSpeculatability();
}];
- let hasCanonicalizer = 1;
let hasCustomAssemblyFormat = 1;
let hasFolder = 1;
let hasRegionVerifier = 1;
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 7e5ce26b5f733..c02a7a56ae87b 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2460,6 +2460,67 @@ static LogicalResult foldLoopBounds(AffineForOp forOp) {
return success(folded);
}
+/// Returns constant trip count in trivial cases.
+static std::optional<uint64_t> getTrivialConstantTripCount(AffineForOp forOp) {
+ int64_t step = forOp.getStepAsInt();
+ if (!forOp.hasConstantBounds() || step <= 0)
+ return std::nullopt;
+ int64_t lb = forOp.getConstantLowerBound();
+ int64_t ub = forOp.getConstantUpperBound();
+ return ub - lb <= 0 ? 0 : (ub - lb + step - 1) / step;
+}
+
+/// Fold the empty loop.
+static LogicalResult AffineForEmptyLoopFolder(AffineForOp forOp) {
+ if (!llvm::hasSingleElement(*forOp.getBody()))
+ return failure();
+ if (forOp.getNumResults() == 0)
+ return success();
+ std::optional<uint64_t> tripCount = getTrivialConstantTripCount(forOp);
+ if (tripCount == 0) {
+ // The initial values of the iteration arguments would be the op's
+ // results.
+ forOp.getResults().replaceAllUsesWith(forOp.getInits());
+ return success();
+ }
+ SmallVector<Value, 4> replacements;
+ auto yieldOp = cast<AffineYieldOp>(forOp.getBody()->getTerminator());
+ auto iterArgs = forOp.getRegionIterArgs();
+ bool hasValDefinedOutsideLoop = false;
+ bool iterArgsNotInOrder = false;
+ for (unsigned i = 0, e = yieldOp->getNumOperands(); i < e; ++i) {
+ Value val = yieldOp.getOperand(i);
+ auto *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 failure();
+ if (iterArgIt == iterArgs.end()) {
+ // `val` is defined outside of the loop.
+ assert(forOp.isDefinedOutsideOfLoop(val) &&
+ "must be defined outside of the loop");
+ hasValDefinedOutsideLoop = true;
+ replacements.push_back(val);
+ } else {
+ unsigned pos = std::distance(iterArgs.begin(), iterArgIt);
+ if (pos != i)
+ iterArgsNotInOrder = true;
+ replacements.push_back(forOp.getInits()[pos]);
+ }
+ }
+ // Bail out when the trip count is unknown and the loop returns any value
+ // defined outside of the loop or any iterArg out of order.
+ if (!tripCount.has_value() &&
+ (hasValDefinedOutsideLoop || iterArgsNotInOrder))
+ return failure();
+ // Bail out when the loop iterates more than once and it returns any iterArg
+ // out of order.
+ if (tripCount.has_value() && tripCount.value() >= 2 && iterArgsNotInOrder)
+ return failure();
+ forOp.getResults().replaceAllUsesWith(replacements);
+ return success();
+}
+
/// Canonicalize the bounds of the given loop.
static LogicalResult canonicalizeLoopBounds(AffineForOp forOp) {
SmallVector<Value, 4> lbOperands(forOp.getLowerBoundOperands());
@@ -2491,81 +2552,6 @@ static LogicalResult canonicalizeLoopBounds(AffineForOp forOp) {
return success();
}
-namespace {
-/// Returns constant trip count in trivial cases.
-static std::optional<uint64_t> getTrivialConstantTripCount(AffineForOp forOp) {
- int64_t step = forOp.getStepAsInt();
- if (!forOp.hasConstantBounds() || step <= 0)
- return std::nullopt;
- int64_t lb = forOp.getConstantLowerBound();
- int64_t ub = forOp.getConstantUpperBound();
- return ub - lb <= 0 ? 0 : (ub - lb + step - 1) / step;
-}
-
-/// This is a pattern to fold trivially empty loop bodies.
-/// TODO: This should be moved into the folding hook.
-struct AffineForEmptyLoopFolder : public OpRewritePattern<AffineForOp> {
- using OpRewritePattern<AffineForOp>::OpRewritePattern;
-
- LogicalResult matchAndRewrite(AffineForOp forOp,
- PatternRewriter &rewriter) const override {
- // Check that the body only contains a yield.
- if (!llvm::hasSingleElement(*forOp.getBody()))
- return failure();
- if (forOp.getNumResults() == 0)
- return success();
- std::optional<uint64_t> tripCount = getTrivialConstantTripCount(forOp);
- if (tripCount == 0) {
- // The initial values of the iteration arguments would be the op's
- // results.
- rewriter.replaceOp(forOp, forOp.getInits());
- return success();
- }
- SmallVector<Value, 4> replacements;
- auto yieldOp = cast<AffineYieldOp>(forOp.getBody()->getTerminator());
- auto iterArgs = forOp.getRegionIterArgs();
- bool hasValDefinedOutsideLoop = false;
- bool iterArgsNotInOrder = false;
- for (unsigned i = 0, e = yieldOp->getNumOperands(); i < e; ++i) {
- Value val = yieldOp.getOperand(i);
- auto *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 failure();
- if (iterArgIt == iterArgs.end()) {
- // `val` is defined outside of the loop.
- assert(forOp.isDefinedOutsideOfLoop(val) &&
- "must be defined outside of the loop");
- hasValDefinedOutsideLoop = true;
- replacements.push_back(val);
- } else {
- unsigned pos = std::distance(iterArgs.begin(), iterArgIt);
- if (pos != i)
- iterArgsNotInOrder = true;
- replacements.push_back(forOp.getInits()[pos]);
- }
- }
- // Bail out when the trip count is unknown and the loop returns any value
- // defined outside of the loop or any iterArg out of order.
- if (!tripCount.has_value() &&
- (hasValDefinedOutsideLoop || iterArgsNotInOrder))
- return failure();
- // Bail out when the loop iterates more than once and it returns any iterArg
- // out of order.
- if (tripCount.has_value() && tripCount.value() >= 2 && iterArgsNotInOrder)
- return failure();
- rewriter.replaceOp(forOp, replacements);
- return success();
- }
-};
-} // namespace
-
-void AffineForOp::getCanonicalizationPatterns(RewritePatternSet &results,
- MLIRContext *context) {
- results.add<AffineForEmptyLoopFolder>(context);
-}
-
OperandRange AffineForOp::getEntrySuccessorOperands(RegionBranchPoint point) {
assert((point.isParent() || point == getRegion()) && "invalid region point");
@@ -2615,6 +2601,7 @@ LogicalResult AffineForOp::fold(FoldAdaptor adaptor,
SmallVectorImpl<OpFoldResult> &results) {
bool folded = succeeded(foldLoopBounds(*this));
folded |= succeeded(canonicalizeLoopBounds(*this));
+ folded |= succeeded(AffineForEmptyLoopFolder(*this));
if (hasTrivialZeroTripCount(*this) && getNumResults() != 0) {
// The initial values of the loop-carried variables (iter_args) are the
// results of the op. But this must be avoided for an affine.for op that
>From 2cf25e87b9a1d07a30606179aef866f3107f9a16 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Fri, 17 Oct 2025 15:08:09 +0000
Subject: [PATCH 2/3] don't use replace use api.
---
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 70 +++++++++++++-----------
1 file changed, 37 insertions(+), 33 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index c02a7a56ae87b..df1c65f8d5c47 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2471,17 +2471,16 @@ static std::optional<uint64_t> getTrivialConstantTripCount(AffineForOp forOp) {
}
/// Fold the empty loop.
-static LogicalResult AffineForEmptyLoopFolder(AffineForOp forOp) {
+static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
if (!llvm::hasSingleElement(*forOp.getBody()))
- return failure();
+ return {};
if (forOp.getNumResults() == 0)
- return success();
+ return {};
std::optional<uint64_t> tripCount = getTrivialConstantTripCount(forOp);
if (tripCount == 0) {
// The initial values of the iteration arguments would be the op's
// results.
- forOp.getResults().replaceAllUsesWith(forOp.getInits());
- return success();
+ return forOp.getInits();
}
SmallVector<Value, 4> replacements;
auto yieldOp = cast<AffineYieldOp>(forOp.getBody()->getTerminator());
@@ -2490,11 +2489,11 @@ static LogicalResult AffineForEmptyLoopFolder(AffineForOp forOp) {
bool iterArgsNotInOrder = false;
for (unsigned i = 0, e = yieldOp->getNumOperands(); i < e; ++i) {
Value val = yieldOp.getOperand(i);
- auto *iterArgIt = llvm::find(iterArgs, val);
+ 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 failure();
+ return {};
if (iterArgIt == iterArgs.end()) {
// `val` is defined outside of the loop.
assert(forOp.isDefinedOutsideOfLoop(val) &&
@@ -2512,13 +2511,14 @@ static LogicalResult AffineForEmptyLoopFolder(AffineForOp forOp) {
// defined outside of the loop or any iterArg out of order.
if (!tripCount.has_value() &&
(hasValDefinedOutsideLoop || iterArgsNotInOrder))
- return failure();
+ return {};
// Bail out when the loop iterates more than once and it returns any iterArg
// out of order.
if (tripCount.has_value() && tripCount.value() >= 2 && iterArgsNotInOrder)
- return failure();
- forOp.getResults().replaceAllUsesWith(replacements);
- return success();
+ return {};
+ return llvm::map_to_vector(replacements, [&](Value replacement) {
+ return OpFoldResult(replacement);
+ });
}
/// Canonicalize the bounds of the given loop.
@@ -2552,6 +2552,32 @@ static LogicalResult canonicalizeLoopBounds(AffineForOp forOp) {
return success();
}
+/// Returns true if the affine.for has zero iterations in trivial cases.
+static bool hasTrivialZeroTripCount(AffineForOp op) {
+ return getTrivialConstantTripCount(op) == 0;
+}
+
+LogicalResult AffineForOp::fold(FoldAdaptor adaptor,
+ SmallVectorImpl<OpFoldResult> &results) {
+ bool folded = succeeded(foldLoopBounds(*this));
+ folded |= succeeded(canonicalizeLoopBounds(*this));
+ if (hasTrivialZeroTripCount(*this) && getNumResults() != 0) {
+ // The initial values of the loop-carried variables (iter_args) are the
+ // results of the op. But this must be avoided for an affine.for op that
+ // does not return any results. Since ops that do not return results cannot
+ // be folded away, we would enter an infinite loop of folds on the same
+ // affine.for op.
+ results.assign(getInits().begin(), getInits().end());
+ folded = true;
+ }
+ SmallVector<OpFoldResult> foldResults = AffineForEmptyLoopFolder(*this);
+ if (!foldResults.empty()) {
+ results.assign(foldResults);
+ folded = true;
+ }
+ return success(folded);
+}
+
OperandRange AffineForOp::getEntrySuccessorOperands(RegionBranchPoint point) {
assert((point.isParent() || point == getRegion()) && "invalid region point");
@@ -2592,28 +2618,6 @@ void AffineForOp::getSuccessorRegions(
regions.push_back(RegionSuccessor(getResults()));
}
-/// Returns true if the affine.for has zero iterations in trivial cases.
-static bool hasTrivialZeroTripCount(AffineForOp op) {
- return getTrivialConstantTripCount(op) == 0;
-}
-
-LogicalResult AffineForOp::fold(FoldAdaptor adaptor,
- SmallVectorImpl<OpFoldResult> &results) {
- bool folded = succeeded(foldLoopBounds(*this));
- folded |= succeeded(canonicalizeLoopBounds(*this));
- folded |= succeeded(AffineForEmptyLoopFolder(*this));
- if (hasTrivialZeroTripCount(*this) && getNumResults() != 0) {
- // The initial values of the loop-carried variables (iter_args) are the
- // results of the op. But this must be avoided for an affine.for op that
- // does not return any results. Since ops that do not return results cannot
- // be folded away, we would enter an infinite loop of folds on the same
- // affine.for op.
- results.assign(getInits().begin(), getInits().end());
- folded = true;
- }
- return success(folded);
-}
-
AffineBound AffineForOp::getLowerBound() {
return AffineBound(*this, getLowerBoundOperands(), getLowerBoundMap());
}
>From 54a9d61a00fd888eeb18ae3f64592aac496cffa5 Mon Sep 17 00:00:00 2001
From: lonely eagle <2020382038 at qq.com>
Date: Fri, 17 Oct 2025 23:18:11 +0800
Subject: [PATCH 3/3] Update mlir/lib/Dialect/Affine/IR/AffineOps.cpp
Co-authored-by: Jakub Kuderski <kubakuderski at gmail.com>
---
mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 4 +---
1 file changed, 1 insertion(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index df1c65f8d5c47..570a7dbf153f4 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -2516,9 +2516,7 @@ static SmallVector<OpFoldResult> AffineForEmptyLoopFolder(AffineForOp forOp) {
// out of order.
if (tripCount.has_value() && tripCount.value() >= 2 && iterArgsNotInOrder)
return {};
- return llvm::map_to_vector(replacements, [&](Value replacement) {
- return OpFoldResult(replacement);
- });
+ return llvm::to_vector_of<OpFoldResult>(replacements);
}
/// Canonicalize the bounds of the given loop.
More information about the Mlir-commits
mailing list