From 1fa00379d916b474c2432f3af071d44e230925a5 Mon Sep 17 00:00:00 2001 From: linuxlonelyeagle <2020382038@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 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 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 replacements; + auto yieldOp = cast(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 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 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 { - using OpRewritePattern::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 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 replacements; - auto yieldOp = cast(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(context); -} - OperandRange AffineForOp::getEntrySuccessorOperands(RegionBranchPoint point) { assert((point.isParent() || point == getRegion()) && "invalid region point"); @@ -2615,6 +2601,7 @@ LogicalResult AffineForOp::fold(FoldAdaptor adaptor, SmallVectorImpl &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@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 getTrivialConstantTripCount(AffineForOp forOp) { } /// Fold the empty loop. -static LogicalResult AffineForEmptyLoopFolder(AffineForOp forOp) { +static SmallVector AffineForEmptyLoopFolder(AffineForOp forOp) { if (!llvm::hasSingleElement(*forOp.getBody())) - return failure(); + return {}; if (forOp.getNumResults() == 0) - return success(); + return {}; std::optional 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 replacements; auto yieldOp = cast(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 &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 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 &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@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 --- 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 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(replacements); } /// Canonicalize the bounds of the given loop.