[Mlir-commits] [mlir] [mlir][affine] Avoid fusion when the sibling loop has result usage (PR #216564)
Kunal Dubey
llvmlistbot at llvm.org
Sat Aug 29 08:55:49 PDT 2026
https://github.com/xakep8 updated https://github.com/llvm/llvm-project/pull/216564
>From 3954639785ac09aa176cc87586143896750df7bb Mon Sep 17 00:00:00 2001
From: Kunal Dubey <xakep8 at protonmail.com>
Date: Sun, 16 Aug 2026 17:47:57 +0530
Subject: [PATCH 1/3] [mlir][affine] Avoid fusion when the sibling loop has
result usage
Sibling loop fusion copies sibling loop to the destination loop and
removes the sibling without checking if the sibiling has loop results
being used at other places. Added check for this to reject fusion if the
sibling has usage.
Added test for the same.
---
.../Dialect/Affine/Transforms/LoopFusion.cpp | 6 ++++++
.../Dialect/Affine/loop-fusion-sibling.mlir | 17 +++++++++++++++++
2 files changed, 23 insertions(+)
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 1ec5fbfef50c3..8adbb97563e6f 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -1369,6 +1369,12 @@ struct GreedyFusion {
// TODO: Remove restrict to single load op restriction.
if (sibNode->getLoadOpCount(memref) != 1)
return false;
+
+ // Sibling fusion removes the sibling node after cloning it into the
+ // destination. Do not fuse siblings whose results still have uses.
+ if (!sibNode->op->use_empty())
+ return false;
+
// Skip if there exists a path of dependent edges between
// 'sibNode' and 'dstNode'.
if (mdg->hasDependencePath(sibNode->id, dstNode->id) ||
diff --git a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
index 937c855b86b50..6c3b1db56d0e9 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
@@ -21,3 +21,20 @@ func.func @disjoint_stores(%0: memref<8xf32>) {
// CHECK-NOT: affine.for
return
}
+
+// CHECK-LABEL: func @sibling_with_used_loop_result
+func.func @sibling_with_used_loop_result(%m: memref<4xi32>, %n: memref<4xi32>,
+ %init: i32) -> i32 {
+ // CHECK: %[[RESULT:.*]] = affine.for {{.*}} iter_args
+ %a = affine.for %i = 0 to 4 iter_args(%x = %init) -> (i32) {
+ %v = affine.load %m[%i] : memref<4xi32>
+ %t = arith.addi %x, %v : i32
+ affine.yield %t : i32
+ }
+ affine.for %i = 0 to 4 {
+ %v = affine.load %m[%i] : memref<4xi32>
+ affine.store %v, %n[%i] : memref<4xi32>
+ }
+ // CHECK: return %[[RESULT]]
+ return %a : i32
+}
>From 7ecb9d4c4b0e281fb0894b19eaeb2a0331bc3e37 Mon Sep 17 00:00:00 2001
From: Kunal Dubey <xakep8 at protonmail.com>
Date: Sat, 29 Aug 2026 16:10:39 +0530
Subject: [PATCH 2/3] [mlir][affine] Preserve used sibling loop results during
fusion
Allowing sibling fusion when the source affine.for has used results, by
carrying the cloned results through the destination loop before
replacing uses of the original sibling loop.
Kept the exisiting cleanup path for result-less sibling loops, and
skipping fusion when moving to the destination loop breaks dominance of
result users by moving definition after usage.
---
.../mlir/Dialect/Affine/LoopFusionUtils.h | 9 +-
.../Dialect/Affine/Transforms/LoopFusion.cpp | 171 +++++++++++++++++-
.../Dialect/Affine/Utils/LoopFusionUtils.cpp | 28 +--
.../Dialect/Affine/loop-fusion-sibling.mlir | 103 ++++++++++-
4 files changed, 290 insertions(+), 21 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h b/mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h
index 0ef39fd7d1463..9b81ed443e3b0 100644
--- a/mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h
+++ b/mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h
@@ -22,6 +22,7 @@
namespace mlir {
class Operation;
+class IRMapping;
namespace affine {
class AffineForOp;
@@ -113,11 +114,15 @@ canFuseLoops(AffineForOp srcForOp, AffineForOp dstForOp, unsigned dstLoopDepth,
/// Fuses 'srcForOp' into 'dstForOp' with destination loop block insertion
/// point and source slice loop bounds specified in 'srcSlice'.
-/// `isInnermostSiblingInsertionFusion` enables cleanup of `srcForOp that is a
+/// `isInnermostSiblingInsertionFusion` enables cleanup of `srcForOp` that is a
/// single-iteration reduction loop being sibling-fused into a 'dstForOp'.
+/// If provided, `mapper` records the mapping from the source loop IR to the
+/// cloned loop IR. Single-iteration loop promotion is skipped in that case to
+/// keep those mappings usable by the caller.
void fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
const ComputationSliceState &srcSlice,
- bool isInnermostSiblingInsertionFusion = false);
+ bool isInnermostSiblingInsertionFusion = false,
+ IRMapping *mapper = nullptr);
/// LoopNestStats aggregates various per-loop statistics (eg. loop trip count
/// and operation count) for a loop nest up until (and including) the innermost
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 8adbb97563e6f..300e2cd71469e 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -22,6 +22,8 @@
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Builders.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/PatternMatch.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/CommandLine.h"
@@ -124,6 +126,118 @@ static bool canRemoveSrcNodeAfterFusion(
return true;
}
+static void replaceInits(RewriterBase &rewriter, AffineForOp forOp,
+ ArrayRef<unsigned> initIndices, ValueRange newInits) {
+ assert(initIndices.size() == newInits.size() &&
+ "expected one replacement value per init operand index");
+
+ SmallVector<Value> inits(forOp.getInits());
+ assert(llvm::all_of(initIndices,
+ [&](unsigned index) { return index < inits.size(); }) &&
+ "expected replacement indices to fit in affine.for init operands");
+
+ for (auto [initIndex, newInit] : llvm::zip(initIndices, newInits))
+ inits[initIndex] = newInit;
+ rewriter.modifyOpInPlace(forOp,
+ [&]() { forOp.getInitsMutable().assign(inits); });
+}
+
+static AffineForOp
+getAffineForOpWithResults(ValueRange values,
+ SmallVectorImpl<unsigned> &indices) {
+ assert(!values.empty() && "expected values to forward");
+ auto firstResult = dyn_cast<OpResult>(values.front());
+ assert(firstResult && "expected forwarded values to be operation results");
+
+ auto forOp = dyn_cast<AffineForOp>(firstResult.getOwner());
+ assert(forOp && "expected forwarded values to be affine.for results");
+ assert(forOp.getNumResults() == forOp.getInits().size() &&
+ "expected one init operand per affine.for result");
+
+ indices.clear();
+ for (Value value : values) {
+ auto result = dyn_cast<OpResult>(value);
+ assert(result && result.getOwner() == forOp &&
+ "expected forwarded values to be results of the same affine.for");
+ indices.push_back(result.getResultNumber());
+ }
+ return forOp;
+}
+
+static Operation *getAncestorInBlock(Operation *op, Block *block) {
+ while (op && op->getBlock() != block)
+ op = op->getParentOp();
+ return op;
+}
+
+// Check that all result users remain dominated after sibling fusion moves
+// `definingOp` before `insertPoint`.
+static bool resultUsersRemainDominated(Operation *opWithUsers,
+ Operation *definingOp,
+ Operation *insertPoint) {
+ Block *block = insertPoint->getBlock();
+ bool definingOpWillMove = definingOp != insertPoint;
+ for (Value result : opWithUsers->getResults()) {
+ if (result.use_empty())
+ continue;
+ for (OpOperand &use : result.getUses()) {
+ Operation *userAncestor = getAncestorInBlock(use.getOwner(), block);
+ if (!userAncestor)
+ return false;
+ if (definingOpWillMove && userAncestor->isBeforeInBlock(insertPoint))
+ return false;
+ if (!definingOpWillMove && !definingOp->isBeforeInBlock(userAncestor))
+ return false;
+ }
+ }
+ return true;
+}
+
+// Carry values from the cloned sibling loop through the destination loop nest.
+// Each newly added iter_arg is used as the corresponding init of the loop that
+// currently defines the forwarded values.
+static SmallVector<Value>
+forwardValuesThroughAffineForNest(ValueRange values, ValueRange inits,
+ AffineForOp &outerForOp) {
+ IRRewriter rewriter(outerForOp->getContext());
+ SmallVector<Value> forwardedValues(values.begin(), values.end());
+ assert(!forwardedValues.empty() && "expected values to forward");
+
+ while (true) {
+ Operation *parentOp =
+ forwardedValues.front().getParentBlock()->getParentOp();
+ AffineForOp parentForOp = dyn_cast<AffineForOp>(parentOp);
+ if (!parentForOp)
+ parentForOp = parentOp->getParentOfType<AffineForOp>();
+
+ assert(parentForOp && "expected value to be nested in an affine.for");
+
+ FailureOr<LoopLikeOpInterface> newLoopOrFailure =
+ parentForOp.replaceWithAdditionalYields(
+ rewriter, inits, /*replaceInitOperandUsesInLoop=*/false,
+ [&](OpBuilder &, Location, ValueRange newIterArgs) {
+ SmallVector<unsigned> initIndices;
+ AffineForOp producerForOp =
+ getAffineForOpWithResults(forwardedValues, initIndices);
+ replaceInits(rewriter, producerForOp, initIndices, newIterArgs);
+ return forwardedValues;
+ });
+
+ assert(succeeded(newLoopOrFailure) && "failed to forward loop results");
+
+ AffineForOp newForOp = cast<AffineForOp>(newLoopOrFailure->getOperation());
+ auto newResults = newForOp->getResults().take_back(forwardedValues.size());
+ forwardedValues.assign(newResults.begin(), newResults.end());
+
+ if (parentForOp == outerForOp) {
+ outerForOp = newForOp;
+ break;
+ }
+ }
+
+ return forwardedValues;
+}
+
/// Returns in 'srcIdCandidates' the producer fusion candidates for consumer
/// 'dstId'. Candidates are sorted by node id order. This order corresponds to
/// the program order when the 'mdg' is created. However, program order is not
@@ -1331,15 +1445,65 @@ struct GreedyFusion {
// destination loop. Based on this, the fused loop may be optimized
// further inside `fuseLoops`.
bool isInnermostInsertion = (bestDstLoopDepth == dstLoopDepthTest);
+
+ bool enableSiblingInsertionCleanup =
+ isInnermostInsertion && sibAffineForOp->use_empty();
+
// Fuse computation slice of 'sibLoopNest' into 'dstLoopNest'.
+ bool hasUsedSiblingResults = !sibAffineForOp->use_empty();
+ if (!dstAffineForOp->use_empty() &&
+ !resultUsersRemainDominated(dstAffineForOp, dstAffineForOp,
+ insertPointInst)) {
+ LDBG() << "Destination results have users before the fused loop "
+ "insertion point.";
+ continue;
+ }
+ if (hasUsedSiblingResults &&
+ !resultUsersRemainDominated(sibAffineForOp, dstAffineForOp,
+ insertPointInst)) {
+ LDBG() << "Sibling results have users before the fused loop insertion "
+ "point.";
+ continue;
+ }
+
+ IRMapping mapper;
affine::fuseLoops(sibAffineForOp, dstAffineForOp, bestSlice,
- isInnermostInsertion);
+ enableSiblingInsertionCleanup,
+ hasUsedSiblingResults ? &mapper : nullptr);
auto dstForInst = cast<AffineForOp>(dstNode->op);
// Update operation position of fused loop nest (if needed).
if (insertPointInst != dstForInst)
dstForInst->moveBefore(insertPointInst);
+ // The cloned sibling loop now lives inside the destination loop. Carry
+ // any used results out before replacing the original sibling loop.
+ if (hasUsedSiblingResults) {
+ SmallVector<Value> usedClonedSibResults;
+ SmallVector<Value> usedSibInits;
+ for (auto [oldResult, init] : llvm::zip(sibAffineForOp->getResults(),
+ sibAffineForOp.getInits())) {
+ if (oldResult.use_empty())
+ continue;
+ usedClonedSibResults.push_back(mapper.lookup(oldResult));
+ usedSibInits.push_back(init);
+ }
+
+ SmallVector<Value> replacementResults =
+ forwardValuesThroughAffineForNest(usedClonedSibResults,
+ usedSibInits, dstAffineForOp);
+
+ unsigned nextReplacement = 0;
+ for (Value oldResult : sibAffineForOp->getResults()) {
+ if (oldResult.use_empty())
+ continue;
+ Value newResult = replacementResults[nextReplacement++];
+ oldResult.replaceAllUsesWith(newResult);
+ }
+
+ dstNode->op = dstAffineForOp;
+ }
+
LDBG() << "Fused sibling nest " << sibId << " into destination nest "
<< dstNode->id << " at depth " << bestDstLoopDepth << ":";
LDBG() << dstAffineForOp;
@@ -1370,11 +1534,6 @@ struct GreedyFusion {
if (sibNode->getLoadOpCount(memref) != 1)
return false;
- // Sibling fusion removes the sibling node after cloning it into the
- // destination. Do not fuse siblings whose results still have uses.
- if (!sibNode->op->use_empty())
- return false;
-
// Skip if there exists a path of dependent edges between
// 'sibNode' and 'dstNode'.
if (mdg->hasDependencePath(sibNode->id, dstNode->id) ||
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index 68296ea3368a1..3ceb7f509452c 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -424,10 +424,12 @@ static LogicalResult promoteSingleIterReductionLoop(AffineForOp forOp,
/// and source slice loop bounds specified in 'srcSlice'.
void mlir::affine::fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
const ComputationSliceState &srcSlice,
- bool isInnermostSiblingInsertion) {
+ bool isInnermostSiblingInsertion,
+ IRMapping *resultMapper) {
// Clone 'srcForOp' into 'dstForOp' at 'srcSlice->insertPoint'.
OpBuilder b(srcSlice.insertPoint->getBlock(), srcSlice.insertPoint);
- IRMapping mapper;
+ IRMapping localMapper;
+ IRMapping &mapper = resultMapper ? *resultMapper : localMapper;
b.clone(*srcForOp, mapper);
// Update 'sliceLoopNest' upper and lower bounds from computed 'srcSlice'.
@@ -455,16 +457,18 @@ void mlir::affine::fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
return (buildSliceTripCountMap(srcSlice, &sliceTripCountMap) &&
(getSliceIterationCount(sliceTripCountMap) == 1));
};
- // Fix up and if possible, eliminate single iteration loops.
- for (AffineForOp forOp : sliceLoops) {
- if (isLoopParallelAndContainsReduction(forOp) &&
- isInnermostSiblingInsertion && srcIsUnitSlice())
- // Patch reduction loop - only ones that are sibling-fused with the
- // destination loop - into the parent loop.
- (void)promoteSingleIterReductionLoop(forOp, true);
- else
- // Promote any single iteration slice loops.
- (void)promoteIfSingleIteration(forOp);
+ if (!resultMapper) {
+ // Fix up and if possible, eliminate single iteration loops.
+ for (AffineForOp forOp : sliceLoops) {
+ if (isLoopParallelAndContainsReduction(forOp) &&
+ isInnermostSiblingInsertion && srcIsUnitSlice())
+ // Patch reduction loop - only ones that are sibling-fused with the
+ // destination loop - into the parent loop.
+ (void)promoteSingleIterReductionLoop(forOp, true);
+ else
+ // Promote any single iteration slice loops.
+ (void)promoteIfSingleIteration(forOp);
+ }
}
}
diff --git a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
index 6c3b1db56d0e9..9335df646c39d 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
@@ -25,7 +25,13 @@ func.func @disjoint_stores(%0: memref<8xf32>) {
// CHECK-LABEL: func @sibling_with_used_loop_result
func.func @sibling_with_used_loop_result(%m: memref<4xi32>, %n: memref<4xi32>,
%init: i32) -> i32 {
- // CHECK: %[[RESULT:.*]] = affine.for {{.*}} iter_args
+ // CHECK: %[[RESULT:.*]] = affine.for %[[I:.*]] = 0 to 4 iter_args(%[[CARRIED:.*]] = %{{.*}}) -> (i32) {
+ // CHECK: %[[INNER:.*]] = affine.for %[[J:.*]] = {{.*}}(%[[I]]) to {{.*}}(%[[I]]) iter_args(%[[INNER_ARG:.*]] = %[[CARRIED]]) -> (i32) {
+ // CHECK: %[[LOAD:.*]] = affine.load %{{.*}}[%[[J]]] : memref<4xi32>
+ // CHECK: %[[ADD:.*]] = arith.addi %[[INNER_ARG]], %[[LOAD]] : i32
+ // CHECK: affine.yield %[[ADD]] : i32
+ // CHECK: }
+ // CHECK: affine.yield %[[INNER]] : i32
%a = affine.for %i = 0 to 4 iter_args(%x = %init) -> (i32) {
%v = affine.load %m[%i] : memref<4xi32>
%t = arith.addi %x, %v : i32
@@ -38,3 +44,98 @@ func.func @sibling_with_used_loop_result(%m: memref<4xi32>, %n: memref<4xi32>,
// CHECK: return %[[RESULT]]
return %a : i32
}
+
+// CHECK-LABEL: func @sibling_with_multiple_used_loop_results
+func.func @sibling_with_multiple_used_loop_results(%m: memref<4xi32>,
+ %n: memref<4xi32>,
+ %init0: i32,
+ %init1: i32) -> (i32, i32) {
+ // CHECK: %[[RESULT:.*]]:2 = affine.for %[[I:.*]] = 0 to 4 iter_args(%[[CARRIED0:.*]] = %{{.*}}, %[[CARRIED1:.*]] = %{{.*}}) -> (i32, i32) {
+ // CHECK: %[[INNER:.*]]:2 = affine.for %[[J:.*]] = {{.*}}(%[[I]]) to {{.*}}(%[[I]]) iter_args(%[[INNER_ARG0:.*]] = %[[CARRIED0]], %[[INNER_ARG1:.*]] = %[[CARRIED1]]) -> (i32, i32) {
+ // CHECK: %[[LOAD:.*]] = affine.load %{{.*}}[%[[J]]] : memref<4xi32>
+ // CHECK: %[[ADD0:.*]] = arith.addi %[[INNER_ARG0]], %[[LOAD]] : i32
+ // CHECK: %[[ADD1:.*]] = arith.addi %[[INNER_ARG1]], %[[LOAD]] : i32
+ // CHECK: affine.yield %[[ADD0]], %[[ADD1]] : i32, i32
+ // CHECK: }
+ // CHECK: affine.yield %[[INNER]]#0, %[[INNER]]#1 : i32, i32
+ %a:2 = affine.for %i = 0 to 4 iter_args(%x = %init0, %y = %init1) -> (i32, i32) {
+ %v = affine.load %m[%i] : memref<4xi32>
+ %t0 = arith.addi %x, %v : i32
+ %t1 = arith.addi %y, %v : i32
+ affine.yield %t0, %t1 : i32, i32
+ }
+ affine.for %i = 0 to 4 {
+ %v = affine.load %m[%i] : memref<4xi32>
+ affine.store %v, %n[%i] : memref<4xi32>
+ }
+ // CHECK: return %[[RESULT]]#0, %[[RESULT]]#1
+ return %a#0, %a#1 : i32, i32
+}
+
+// CHECK-LABEL: func @sibling_with_second_of_multiple_loop_results_used
+func.func @sibling_with_second_of_multiple_loop_results_used(%m: memref<4xi32>,
+ %n: memref<4xi32>,
+ %init0: i32,
+ %init1: i32) -> i32 {
+ // CHECK: %[[RESULT:.*]] = affine.for %[[I:.*]] = 0 to 4 iter_args(%[[CARRIED:.*]] = %{{.*}}) -> (i32) {
+ // CHECK: %[[INNER:.*]]:2 = affine.for %[[J:.*]] = {{.*}}(%[[I]]) to {{.*}}(%[[I]]) iter_args(%{{.*}} = %{{.*}}, %[[INNER_ARG:.*]] = %[[CARRIED]]) -> (i32, i32) {
+ // CHECK: %[[LOAD:.*]] = affine.load %{{.*}}[%[[J]]] : memref<4xi32>
+ // CHECK: %[[ADD:.*]] = arith.addi %[[INNER_ARG]], %[[LOAD]] : i32
+ // CHECK: affine.yield %{{.*}}, %[[ADD]] : i32, i32
+ // CHECK: }
+ // CHECK: affine.yield %[[INNER]]#1 : i32
+ %a:2 = affine.for %i = 0 to 4 iter_args(%x = %init0, %y = %init1) -> (i32, i32) {
+ %v = affine.load %m[%i] : memref<4xi32>
+ %t0 = arith.addi %x, %v : i32
+ %t1 = arith.addi %y, %v : i32
+ affine.yield %t0, %t1 : i32, i32
+ }
+ affine.for %i = 0 to 4 {
+ %v = affine.load %m[%i] : memref<4xi32>
+ affine.store %v, %n[%i] : memref<4xi32>
+ }
+ // CHECK: return %[[RESULT]]
+ return %a#1 : i32
+}
+
+// CHECK-LABEL: func @sibling_result_used_before_destination
+func.func @sibling_result_used_before_destination(%m: memref<4xi32>,
+ %n: memref<4xi32>,
+ %init: i32) -> i32 {
+ // CHECK: %[[A:.*]] = affine.for
+ %a = affine.for %i = 0 to 4 iter_args(%x = %init) -> (i32) {
+ %v = affine.load %m[%i] : memref<4xi32>
+ %t = arith.addi %x, %v : i32
+ affine.yield %t : i32
+ }
+ // CHECK: %[[B:.*]] = arith.addi %[[A]],
+ %b = arith.addi %a, %init : i32
+ // CHECK: affine.for
+ affine.for %i = 0 to 4 {
+ %v = affine.load %m[%i] : memref<4xi32>
+ affine.store %v, %n[%i] : memref<4xi32>
+ }
+ // CHECK: return %[[B]]
+ return %b : i32
+}
+
+// CHECK-LABEL: func @destination_result_used_before_sibling
+func.func @destination_result_used_before_sibling(%m: memref<4xi32>,
+ %n: memref<4xi32>,
+ %init: i32) -> i32 {
+ // CHECK: %[[A:.*]] = affine.for
+ %a = affine.for %i = 0 to 4 iter_args(%x = %init) -> (i32) {
+ %v = affine.load %m[%i] : memref<4xi32>
+ affine.store %v, %n[%i] : memref<4xi32>
+ affine.yield %v : i32
+ }
+ // CHECK: %[[B:.*]] = arith.addi %[[A]],
+ %b = arith.addi %a, %init : i32
+ // CHECK: affine.for
+ affine.for %i = 0 to 4 {
+ %v = affine.load %m[%i] : memref<4xi32>
+ affine.store %v, %n[%i] : memref<4xi32>
+ }
+ // CHECK: return %[[B]]
+ return %b : i32
+}
>From acbbd6c06a8206da3e0504804bb5225034e8654b Mon Sep 17 00:00:00 2001
From: Kunal Dubey <xakep8 at protonmail.com>
Date: Sat, 29 Aug 2026 21:24:08 +0530
Subject: [PATCH 3/3] [MLIR][Affine] Preserve used sibling loop results during
fusion
Updated mapper handling so that single-iteration promotion keeps cloned result mappings usable, and add sibling fusion tests for used loop results.
---
.../mlir/Dialect/Affine/LoopFusionUtils.h | 3 +-
.../Dialect/Affine/Transforms/LoopFusion.cpp | 129 +++++++++++++-----
.../Dialect/Affine/Utils/LoopFusionUtils.cpp | 89 ++++++++++--
.../Dialect/Affine/loop-fusion-sibling.mlir | 33 ++---
4 files changed, 189 insertions(+), 65 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h b/mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h
index 9b81ed443e3b0..45a8a9791d2f6 100644
--- a/mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h
+++ b/mlir/include/mlir/Dialect/Affine/LoopFusionUtils.h
@@ -117,8 +117,7 @@ canFuseLoops(AffineForOp srcForOp, AffineForOp dstForOp, unsigned dstLoopDepth,
/// `isInnermostSiblingInsertionFusion` enables cleanup of `srcForOp` that is a
/// single-iteration reduction loop being sibling-fused into a 'dstForOp'.
/// If provided, `mapper` records the mapping from the source loop IR to the
-/// cloned loop IR. Single-iteration loop promotion is skipped in that case to
-/// keep those mappings usable by the caller.
+/// cloned loop IR.
void fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
const ComputationSliceState &srcSlice,
bool isInnermostSiblingInsertionFusion = false,
diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 300e2cd71469e..b9675dc01fd47 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -12,6 +12,7 @@
#include "mlir/Dialect/Affine/Transforms/Passes.h"
+#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/Utils.h"
@@ -62,7 +63,7 @@ struct LoopFusion : public affine::impl::AffineLoopFusionBase<LoopFusion> {
this->affineFusionMode = affineFusionMode;
}
- void runOnBlock(Block *block);
+ LogicalResult runOnBlock(Block *block);
void runOnOperation() override;
};
@@ -147,23 +148,59 @@ getAffineForOpWithResults(ValueRange values,
SmallVectorImpl<unsigned> &indices) {
assert(!values.empty() && "expected values to forward");
auto firstResult = dyn_cast<OpResult>(values.front());
- assert(firstResult && "expected forwarded values to be operation results");
+ if (!firstResult)
+ return nullptr;
auto forOp = dyn_cast<AffineForOp>(firstResult.getOwner());
- assert(forOp && "expected forwarded values to be affine.for results");
+ if (!forOp)
+ return nullptr;
assert(forOp.getNumResults() == forOp.getInits().size() &&
"expected one init operand per affine.for result");
indices.clear();
for (Value value : values) {
auto result = dyn_cast<OpResult>(value);
- assert(result && result.getOwner() == forOp &&
- "expected forwarded values to be results of the same affine.for");
+ if (!result || result.getOwner() != forOp) {
+ indices.clear();
+ return nullptr;
+ }
indices.push_back(result.getResultNumber());
}
return forOp;
}
+static LogicalResult getBackwardSliceForValues(ValueRange values,
+ AffineForOp parentForOp,
+ SetVector<Operation *> &slice) {
+ BackwardSliceOptions options;
+ options.inclusive = true;
+ options.omitBlockArguments = true;
+ options.filter = [&](Operation *op) {
+ return parentForOp->isProperAncestor(op);
+ };
+
+ for (Value value : values)
+ if (Operation *definingOp = value.getDefiningOp()) {
+ if (failed(getBackwardSlice(definingOp, &slice, options)))
+ return failure();
+ }
+ return success();
+}
+
+static void replaceInitUsesInPromotedSlice(RewriterBase &rewriter,
+ const SetVector<Operation *> &slice,
+ ValueRange oldValues,
+ ValueRange newValues) {
+ assert(oldValues.size() == newValues.size() &&
+ "expected one replacement value per old value");
+
+ for (auto [oldValue, newValue] : llvm::zip(oldValues, newValues)) {
+ rewriter.replaceUsesWithIf(oldValue, newValue, [&](OpOperand &use) {
+ return slice.contains(use.getOwner());
+ });
+ }
+}
+
static Operation *getAncestorInBlock(Operation *op, Block *block) {
while (op && op->getBlock() != block)
op = op->getParentOp();
@@ -196,7 +233,7 @@ static bool resultUsersRemainDominated(Operation *opWithUsers,
// Carry values from the cloned sibling loop through the destination loop nest.
// Each newly added iter_arg is used as the corresponding init of the loop that
// currently defines the forwarded values.
-static SmallVector<Value>
+static FailureOr<SmallVector<Value>>
forwardValuesThroughAffineForNest(ValueRange values, ValueRange inits,
AffineForOp &outerForOp) {
IRRewriter rewriter(outerForOp->getContext());
@@ -212,20 +249,34 @@ forwardValuesThroughAffineForNest(ValueRange values, ValueRange inits,
assert(parentForOp && "expected value to be nested in an affine.for");
+ SmallVector<unsigned> initIndices;
+ AffineForOp producerForOp =
+ getAffineForOpWithResults(forwardedValues, initIndices);
+ SetVector<Operation *> promotedSlice;
+ if (!producerForOp && failed(getBackwardSliceForValues(
+ forwardedValues, parentForOp, promotedSlice)))
+ return failure();
+
+ SmallVector<Value> newIterArgsForReplacement;
FailureOr<LoopLikeOpInterface> newLoopOrFailure =
parentForOp.replaceWithAdditionalYields(
rewriter, inits, /*replaceInitOperandUsesInLoop=*/false,
[&](OpBuilder &, Location, ValueRange newIterArgs) {
- SmallVector<unsigned> initIndices;
- AffineForOp producerForOp =
- getAffineForOpWithResults(forwardedValues, initIndices);
- replaceInits(rewriter, producerForOp, initIndices, newIterArgs);
+ newIterArgsForReplacement.assign(newIterArgs.begin(),
+ newIterArgs.end());
+ if (producerForOp)
+ replaceInits(rewriter, producerForOp, initIndices, newIterArgs);
return forwardedValues;
});
- assert(succeeded(newLoopOrFailure) && "failed to forward loop results");
+ if (failed(newLoopOrFailure))
+ return failure();
AffineForOp newForOp = cast<AffineForOp>(newLoopOrFailure->getOperation());
+ if (!producerForOp)
+ replaceInitUsesInPromotedSlice(rewriter, promotedSlice, inits,
+ newIterArgsForReplacement);
+
auto newResults = newForOp->getResults().take_back(forwardedValues.size());
forwardedValues.assign(newResults.begin(), newResults.end());
@@ -923,16 +974,19 @@ struct GreedyFusion {
}
}
/// Run only sibling fusion on the `mdg`.
- void runSiblingFusionOnly() {
- fuseSiblingNodes();
+ LogicalResult runSiblingFusionOnly() {
+ if (failed(fuseSiblingNodes()))
+ return failure();
eraseUnusedMemRefAllocations();
+ return success();
}
/// Run only producer/consumer fusion on the `mdg`.
- void runProducerConsumerFusionOnly() {
+ LogicalResult runProducerConsumerFusionOnly() {
fuseProducerConsumerNodes(
/*maxSrcUserCount=*/std::numeric_limits<unsigned>::max());
eraseUnusedMemRefAllocations();
+ return success();
}
// Run the GreedyFusion pass.
@@ -940,13 +994,15 @@ struct GreedyFusion {
// unique consumer.
// *) Second pass fuses sibling nodes which share no dependence edges.
// *) Third pass fuses any remaining producer nodes into their users.
- void runGreedyFusion() {
+ LogicalResult runGreedyFusion() {
// TODO: Run this repeatedly until a fixed-point is reached.
fuseProducerConsumerNodes(/*maxSrcUserCount=*/1);
- fuseSiblingNodes();
+ if (failed(fuseSiblingNodes()))
+ return failure();
fuseProducerConsumerNodes(
/*maxSrcUserCount=*/std::numeric_limits<unsigned>::max());
eraseUnusedMemRefAllocations();
+ return success();
}
/// Returns true if a private memref can be created for `memref` given
@@ -1290,7 +1346,7 @@ struct GreedyFusion {
// Visits each node in the graph, and for each node, attempts to fuse it with
// its sibling nodes (nodes which share a parent, but no dependence edges).
- void fuseSiblingNodes() {
+ LogicalResult fuseSiblingNodes() {
LDBG() << "--- Sibling Fusion ---";
init();
while (!worklist.empty()) {
@@ -1306,12 +1362,14 @@ struct GreedyFusion {
if (!isa<AffineForOp>(dstNode->op))
continue;
// Attempt to fuse 'dstNode' with its sibling nodes in the graph.
- fuseWithSiblingNodes(dstNode);
+ if (failed(fuseWithSiblingNodes(dstNode)))
+ return failure();
}
+ return success();
}
// Attempt to fuse 'dstNode' with sibling nodes in the graph.
- void fuseWithSiblingNodes(Node *dstNode) {
+ LogicalResult fuseWithSiblingNodes(Node *dstNode) {
DenseSet<unsigned> visitedSibNodeIds;
std::pair<unsigned, Value> idAndMemref;
auto dstAffineForOp = cast<AffineForOp>(dstNode->op);
@@ -1401,7 +1459,7 @@ struct GreedyFusion {
if (!fraction || fraction > 0) {
LDBG() << "Can't perform maximal fusion with a cyclic dependence "
<< "and non-zero additional compute.";
- return;
+ return success();
}
} else {
// Set redundant computation tolerance to zero regardless of what the
@@ -1489,15 +1547,20 @@ struct GreedyFusion {
usedSibInits.push_back(init);
}
- SmallVector<Value> replacementResults =
+ FailureOr<SmallVector<Value>> replacementResults =
forwardValuesThroughAffineForNest(usedClonedSibResults,
usedSibInits, dstAffineForOp);
+ if (failed(replacementResults)) {
+ LDBG() << "Failed to forward sibling loop results through the "
+ "destination loop nest.";
+ return failure();
+ }
unsigned nextReplacement = 0;
for (Value oldResult : sibAffineForOp->getResults()) {
if (oldResult.use_empty())
continue;
- Value newResult = replacementResults[nextReplacement++];
+ Value newResult = (*replacementResults)[nextReplacement++];
oldResult.replaceAllUsesWith(newResult);
}
@@ -1517,6 +1580,7 @@ struct GreedyFusion {
mdg->removeNode(sibNode->id);
op->erase();
}
+ return success();
}
// Searches block argument uses and the graph from 'dstNode' looking for a
@@ -1682,11 +1746,11 @@ struct GreedyFusion {
} // namespace
/// Run fusion on `block`.
-void LoopFusion::runOnBlock(Block *block) {
+LogicalResult LoopFusion::runOnBlock(Block *block) {
MemRefDependenceGraph g(*block);
if (!g.init()) {
LDBG() << "MDG init failed";
- return;
+ return success();
}
std::optional<unsigned> fastMemorySpaceOpt;
@@ -1697,25 +1761,28 @@ void LoopFusion::runOnBlock(Block *block) {
maximalFusion, computeToleranceThreshold);
if (affineFusionMode == FusionMode::ProducerConsumer)
- fusion.runProducerConsumerFusionOnly();
+ return fusion.runProducerConsumerFusionOnly();
else if (affineFusionMode == FusionMode::Sibling)
- fusion.runSiblingFusionOnly();
- else
- fusion.runGreedyFusion();
+ return fusion.runSiblingFusionOnly();
+ return fusion.runGreedyFusion();
}
void LoopFusion::runOnOperation() {
// Call fusion on every op that has at least two affine.for nests (in post
// order).
- getOperation()->walk([&](Operation *op) {
+ WalkResult result = getOperation()->walk([&](Operation *op) {
for (Region ®ion : op->getRegions()) {
for (Block &block : region.getBlocks()) {
auto affineFors = block.getOps<AffineForOp>();
- if (!affineFors.empty() && !llvm::hasSingleElement(affineFors))
- runOnBlock(&block);
+ if (!affineFors.empty() && !llvm::hasSingleElement(affineFors) &&
+ failed(runOnBlock(&block)))
+ return WalkResult::interrupt();
}
}
+ return WalkResult::advance();
});
+ if (result.wasInterrupted())
+ signalPassFailure();
}
std::unique_ptr<Pass> mlir::affine::createLoopFusionPass(
diff --git a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
index 3ceb7f509452c..1d4add9372ac6 100644
--- a/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
+++ b/mlir/lib/Dialect/Affine/Utils/LoopFusionUtils.cpp
@@ -21,6 +21,7 @@
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/PatternMatch.h"
+#include "llvm/ADT/DenseMap.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/DebugLog.h"
#include "llvm/Support/raw_ostream.h"
@@ -352,10 +353,40 @@ FusionResult mlir::affine::canFuseLoops(AffineForOp srcForOp,
return FusionResult::Success;
}
+static void remapPromotedResults(AffineForOp forOp, ValueRange replacements,
+ IRMapping &mapper) {
+ llvm::SmallDenseMap<Value, Value, 4> replacementMap;
+ for (auto [result, replacement] : llvm::zip(forOp.getResults(), replacements))
+ replacementMap[result] = replacement;
+
+ SmallVector<std::pair<Value, Value>> remaps;
+ for (auto [from, to] : mapper.getValueMap()) {
+ auto replacement = replacementMap.find(to);
+ if (replacement != replacementMap.end())
+ remaps.push_back({from, replacement->second});
+ }
+ for (auto [from, to] : remaps)
+ mapper.map(from, to);
+}
+
+static FailureOr<SmallVector<Value>>
+getSingleIterationYieldedValues(AffineForOp forOp) {
+ std::optional<APInt> tripCount = forOp.getStaticTripCount();
+ if (!tripCount || *tripCount != 1)
+ return failure();
+
+ // TODO: extend this for arbitrary affine bounds.
+ if (forOp.getLowerBoundMap().getNumResults() != 1)
+ return failure();
+
+ return SmallVector<Value>(forOp.getBody()->getTerminator()->getOperands());
+}
+
/// Patch the loop body of a forOp that is a single iteration reduction loop
/// into its containing block.
static LogicalResult promoteSingleIterReductionLoop(AffineForOp forOp,
- bool siblingFusionUser) {
+ bool siblingFusionUser,
+ IRMapping *mapper) {
// Check if the reduction loop is a single iteration loop.
std::optional<APInt> tripCount = forOp.getStaticTripCount();
if (!tripCount || *tripCount != 1)
@@ -376,6 +407,11 @@ static LogicalResult promoteSingleIterReductionLoop(AffineForOp forOp,
[&](OpBuilder &b, Location loc, ArrayRef<BlockArgument> newBbArgs) {
return newOperands;
}));
+ if (mapper)
+ remapPromotedResults(
+ forOp,
+ newLoop.getResults().slice(parentOpNumResults, forOp.getNumResults()),
+ *mapper);
// For sibling-fusion users, collect operations that use the results of the
// `forOp` outside the new parent loop that has absorbed all its iter args
@@ -420,6 +456,35 @@ static LogicalResult promoteSingleIterReductionLoop(AffineForOp forOp,
return success();
}
+static LogicalResult
+promoteIfSingleIterationAndUpdateMapper(AffineForOp forOp, IRMapping *mapper) {
+ if (!mapper)
+ return promoteIfSingleIteration(forOp);
+
+ FailureOr<SmallVector<Value>> replacements =
+ getSingleIterationYieldedValues(forOp);
+ if (failed(replacements))
+ return failure();
+
+ llvm::SmallDenseMap<Value, Value, 4> replacementMap;
+ for (auto [result, replacement] :
+ llvm::zip(forOp.getResults(), *replacements))
+ replacementMap[result] = replacement;
+
+ for (auto &mapEntry : mapper->getValueMap()) {
+ Value to = mapEntry.second;
+ auto replacement = replacementMap.find(to);
+ if (replacement == replacementMap.end())
+ continue;
+ Operation *definingOp = replacement->second.getDefiningOp();
+ if (!definingOp || !forOp->isProperAncestor(definingOp))
+ return failure();
+ }
+
+ remapPromotedResults(forOp, *replacements, *mapper);
+ return promoteIfSingleIteration(forOp);
+}
+
/// Fuses 'srcForOp' into 'dstForOp' with destination loop block insertion point
/// and source slice loop bounds specified in 'srcSlice'.
void mlir::affine::fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
@@ -457,18 +522,16 @@ void mlir::affine::fuseLoops(AffineForOp srcForOp, AffineForOp dstForOp,
return (buildSliceTripCountMap(srcSlice, &sliceTripCountMap) &&
(getSliceIterationCount(sliceTripCountMap) == 1));
};
- if (!resultMapper) {
- // Fix up and if possible, eliminate single iteration loops.
- for (AffineForOp forOp : sliceLoops) {
- if (isLoopParallelAndContainsReduction(forOp) &&
- isInnermostSiblingInsertion && srcIsUnitSlice())
- // Patch reduction loop - only ones that are sibling-fused with the
- // destination loop - into the parent loop.
- (void)promoteSingleIterReductionLoop(forOp, true);
- else
- // Promote any single iteration slice loops.
- (void)promoteIfSingleIteration(forOp);
- }
+ // Fix up and if possible, eliminate single iteration loops.
+ for (AffineForOp forOp : sliceLoops) {
+ if (isLoopParallelAndContainsReduction(forOp) &&
+ isInnermostSiblingInsertion && srcIsUnitSlice())
+ // Patch reduction loop - only ones that are sibling-fused with the
+ // destination loop - into the parent loop.
+ (void)promoteSingleIterReductionLoop(forOp, true, resultMapper);
+ else
+ // Promote any single iteration slice loops.
+ (void)promoteIfSingleIterationAndUpdateMapper(forOp, resultMapper);
}
}
diff --git a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
index 9335df646c39d..fedf04e1d7fce 100644
--- a/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
+++ b/mlir/test/Dialect/Affine/loop-fusion-sibling.mlir
@@ -26,12 +26,10 @@ func.func @disjoint_stores(%0: memref<8xf32>) {
func.func @sibling_with_used_loop_result(%m: memref<4xi32>, %n: memref<4xi32>,
%init: i32) -> i32 {
// CHECK: %[[RESULT:.*]] = affine.for %[[I:.*]] = 0 to 4 iter_args(%[[CARRIED:.*]] = %{{.*}}) -> (i32) {
- // CHECK: %[[INNER:.*]] = affine.for %[[J:.*]] = {{.*}}(%[[I]]) to {{.*}}(%[[I]]) iter_args(%[[INNER_ARG:.*]] = %[[CARRIED]]) -> (i32) {
- // CHECK: %[[LOAD:.*]] = affine.load %{{.*}}[%[[J]]] : memref<4xi32>
- // CHECK: %[[ADD:.*]] = arith.addi %[[INNER_ARG]], %[[LOAD]] : i32
- // CHECK: affine.yield %[[ADD]] : i32
- // CHECK: }
- // CHECK: affine.yield %[[INNER]] : i32
+ // CHECK: %[[LOAD:.*]] = affine.load %{{.*}}[%[[I]]] : memref<4xi32>
+ // CHECK: %[[ADD:.*]] = arith.addi %[[CARRIED]], %[[LOAD]] : i32
+ // CHECK: affine.store
+ // CHECK: affine.yield %[[ADD]] : i32
%a = affine.for %i = 0 to 4 iter_args(%x = %init) -> (i32) {
%v = affine.load %m[%i] : memref<4xi32>
%t = arith.addi %x, %v : i32
@@ -51,13 +49,11 @@ func.func @sibling_with_multiple_used_loop_results(%m: memref<4xi32>,
%init0: i32,
%init1: i32) -> (i32, i32) {
// CHECK: %[[RESULT:.*]]:2 = affine.for %[[I:.*]] = 0 to 4 iter_args(%[[CARRIED0:.*]] = %{{.*}}, %[[CARRIED1:.*]] = %{{.*}}) -> (i32, i32) {
- // CHECK: %[[INNER:.*]]:2 = affine.for %[[J:.*]] = {{.*}}(%[[I]]) to {{.*}}(%[[I]]) iter_args(%[[INNER_ARG0:.*]] = %[[CARRIED0]], %[[INNER_ARG1:.*]] = %[[CARRIED1]]) -> (i32, i32) {
- // CHECK: %[[LOAD:.*]] = affine.load %{{.*}}[%[[J]]] : memref<4xi32>
- // CHECK: %[[ADD0:.*]] = arith.addi %[[INNER_ARG0]], %[[LOAD]] : i32
- // CHECK: %[[ADD1:.*]] = arith.addi %[[INNER_ARG1]], %[[LOAD]] : i32
- // CHECK: affine.yield %[[ADD0]], %[[ADD1]] : i32, i32
- // CHECK: }
- // CHECK: affine.yield %[[INNER]]#0, %[[INNER]]#1 : i32, i32
+ // CHECK: %[[LOAD:.*]] = affine.load %{{.*}}[%[[I]]] : memref<4xi32>
+ // CHECK: %[[ADD0:.*]] = arith.addi %[[CARRIED0]], %[[LOAD]] : i32
+ // CHECK: %[[ADD1:.*]] = arith.addi %[[CARRIED1]], %[[LOAD]] : i32
+ // CHECK: affine.store
+ // CHECK: affine.yield %[[ADD0]], %[[ADD1]] : i32, i32
%a:2 = affine.for %i = 0 to 4 iter_args(%x = %init0, %y = %init1) -> (i32, i32) {
%v = affine.load %m[%i] : memref<4xi32>
%t0 = arith.addi %x, %v : i32
@@ -78,12 +74,11 @@ func.func @sibling_with_second_of_multiple_loop_results_used(%m: memref<4xi32>,
%init0: i32,
%init1: i32) -> i32 {
// CHECK: %[[RESULT:.*]] = affine.for %[[I:.*]] = 0 to 4 iter_args(%[[CARRIED:.*]] = %{{.*}}) -> (i32) {
- // CHECK: %[[INNER:.*]]:2 = affine.for %[[J:.*]] = {{.*}}(%[[I]]) to {{.*}}(%[[I]]) iter_args(%{{.*}} = %{{.*}}, %[[INNER_ARG:.*]] = %[[CARRIED]]) -> (i32, i32) {
- // CHECK: %[[LOAD:.*]] = affine.load %{{.*}}[%[[J]]] : memref<4xi32>
- // CHECK: %[[ADD:.*]] = arith.addi %[[INNER_ARG]], %[[LOAD]] : i32
- // CHECK: affine.yield %{{.*}}, %[[ADD]] : i32, i32
- // CHECK: }
- // CHECK: affine.yield %[[INNER]]#1 : i32
+ // CHECK: %[[LOAD:.*]] = affine.load %{{.*}}[%[[I]]] : memref<4xi32>
+ // CHECK: arith.addi %{{.*}}, %[[LOAD]] : i32
+ // CHECK: %[[ADD:.*]] = arith.addi %[[CARRIED]], %[[LOAD]] : i32
+ // CHECK: affine.store
+ // CHECK: affine.yield %[[ADD]] : i32
%a:2 = affine.for %i = 0 to 4 iter_args(%x = %init0, %y = %init1) -> (i32, i32) {
%v = affine.load %m[%i] : memref<4xi32>
%t0 = arith.addi %x, %v : i32
More information about the Mlir-commits
mailing list