[Mlir-commits] [mlir] [mlir][affine] Add store sinking transformation for AffineIfOp (PR #213531)
lonely eagle
llvmlistbot at llvm.org
Sun Aug 2 05:04:40 PDT 2026
https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/213531
>From c5de0775b22d1df04452ea707ff52b2f87b3c6cd Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Fri, 31 Jul 2026 07:30:01 +0000
Subject: [PATCH 1/3] impl analyzeIfOp and applySinkPlan in the affine-scalrep
pass.
---
.../Transforms/AffineScalarReplacement.cpp | 242 +++++++++++++++++-
mlir/test/Dialect/Affine/scalrep.mlir | 83 ++++++
2 files changed, 323 insertions(+), 2 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
index 92a4cf3c0ad76..88b149d819a00 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
@@ -11,12 +11,21 @@
// redundant loads.
//===----------------------------------------------------------------------===//
-#include "mlir/Dialect/Affine/Transforms/Passes.h"
-
#include "mlir/Analysis/AliasAnalysis.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Transforms/Passes.h"
#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Block.h"
#include "mlir/IR/Dominance.h"
+#include "mlir/IR/IntegerSet.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/OperationSupport.h"
+#include "mlir/IR/PatternMatch.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/DebugLog.h"
namespace mlir {
namespace affine {
@@ -31,6 +40,223 @@ using namespace mlir;
using namespace mlir::affine;
namespace {
+
+// Stores matched pairs of AffineStoreOps to be sinked outside the AffineIfOp.
+struct IfStorePairToSink {
+ AffineStoreOp thenStore;
+ AffineStoreOp elseStore;
+ AffineStoreOp parentStore;
+};
+
+// Transformation plan for a AffineIfOp.
+struct IfSinkPlan {
+ AffineIfOp ifOp;
+ llvm::SmallVector<IfStorePairToSink, 4> pairs;
+};
+
+// Returns true if `store` is the last store accessing its memory location
+// within its block.
+static bool isLastStoreInBlock(AffineStoreOp store) {
+ MemRefAccess curAccess(store);
+ for (Operation &op : llvm::reverse(*store->getBlock())) {
+ if (auto otherStore = dyn_cast<AffineStoreOp>(&op)) {
+ if (MemRefAccess(otherStore) == curAccess)
+ return otherStore == store;
+ }
+ }
+ return false;
+}
+
+/// Analyzes a single block (either 'then' or 'else') of an AffineIfOp to
+/// identify store operations that can be sinked outside the if statement.
+///
+/// This function identifies store operations across the following memory access
+/// scenarios by tracing users of the target memref:
+/// 1. Parent + Both Branches: The parentStore (located before the ifOp in the
+/// same outer block), the thenStore, and the elseStore all access the exact
+/// same memory location.
+/// 2. Parent + Single Branch: The parentStore (located before the ifOp in the
+/// same outer block) and a store in either the 'then' or 'else' block access
+/// the exact same memory location.
+///
+/// Found store candidates are recorded into 'plan' for transformation, and
+/// matched stores in the 'else' block are added to 'visited' to prevent
+/// duplicate processing.
+static void analyzeIfBlock(AffineIfOp ifOp, Block *block, IfSinkPlan &plan,
+ SmallVectorImpl<AffineStoreOp> &visited) {
+ // Check whether the target block is the 'then' block of the ifOp (and an
+ // 'else' block exists).
+ bool curThenBlock = ifOp.hasElse() && block == ifOp.getThenBlock();
+
+ // Iterate through operations in the block in reverse order.
+ for (Operation &op : llvm::reverse(*block)) {
+ auto store = dyn_cast<AffineStoreOp>(&op);
+
+ // Skip if it is not an AffineStoreOp or not the last store targeting the
+ // memory location in this.
+ if (!store || !isLastStoreInBlock(store))
+ continue;
+
+ // Skip if this store has already been processed and visited.
+ if (llvm::is_contained(visited, store))
+ continue;
+
+ // Verify index operands are valid outside the ifOp scope.
+ MemRefAccess access(store);
+ if (llvm::any_of(store.getIndices(), [&](Value operand) {
+ Operation *defineOp = operand.getDefiningOp();
+ return !defineOp && defineOp->getBlock() == block;
+ }))
+ continue;
+
+ Value memref = store.getMemRef();
+ AffineStoreOp parentStore = nullptr;
+ AffineStoreOp elseStore = nullptr;
+
+ // Trace all users of the memref to locate matching stores.
+ for (Operation *user : memref.getUsers()) {
+ auto userStore = dyn_cast<AffineStoreOp>(user);
+ if (!userStore || userStore == store)
+ continue;
+
+ if (MemRefAccess(userStore) != access)
+ continue;
+
+ // Ensure the `parentStore` is in the same block as ifOp and precedes it,
+ // update parentStore to keep the closest one preceding ifOp.
+ if (userStore->getBlock() == ifOp->getBlock() &&
+ userStore->isBeforeInBlock(ifOp))
+ if (!parentStore || parentStore->isBeforeInBlock(userStore))
+ parentStore = userStore;
+
+ // Identify matching stores in the 'else' block when analyzing the 'then'
+ // block.
+ bool userInElseBlock =
+ ifOp.hasElse() && userStore->getBlock() == ifOp.getElseBlock();
+ if (curThenBlock && userInElseBlock &&
+ (!elseStore || isLastStoreInBlock(userStore)))
+ elseStore = userStore;
+ }
+
+ // Skip if no matching parent store was found.
+ if (!parentStore)
+ continue;
+
+ // Record the candidate store pair into the sink plan based on the current
+ // branch being analyzed.
+ if (curThenBlock) {
+ plan.pairs.push_back({store, elseStore, parentStore});
+ if (elseStore)
+ visited.push_back(elseStore);
+ } else {
+ plan.pairs.push_back({nullptr, store, parentStore});
+ }
+ }
+}
+
+/// Analyzes an AffineIfOp to build a store sinking plan.
+static void analyzeIfOp(AffineIfOp ifOp, IfSinkPlan &plan) {
+ /// maintaining a `visited` tracking vector across both blocks to prevent
+ /// store operations in the `else` block from being processed twice.
+ SmallVector<AffineStoreOp> visited;
+ analyzeIfBlock(ifOp, ifOp.getThenBlock(), plan, visited);
+ if (ifOp.hasElse())
+ analyzeIfBlock(ifOp, ifOp.getElseBlock(), plan, visited);
+}
+
+/// Applies the store sinking plan to rewrite the target AffineIfOp. This
+/// transformation replaces the existing `ifOp` with a new `AffineIfOp` that
+/// yields the values to be stored across branches. It updates the
+/// `AffineYieldOp` terminators in both 'then' and 'else' blocks, erases
+/// internal store operations, and emits a single unified `AffineStoreOp`
+/// immediately following the new `ifOp` for each sinked pair.
+static void applySinkPlan(RewriterBase &rewriter, IfSinkPlan &plan) {
+ if (plan.pairs.empty())
+ return;
+
+ AffineIfOp ifOp = plan.ifOp;
+
+ // Collect new result types for the new AffineIfOp.
+ SmallVector<Type, 4> newTypes(ifOp.getResultTypes());
+ for (IfStorePairToSink &pair : plan.pairs)
+ newTypes.push_back(pair.parentStore.getValue().getType());
+
+ // Create a new AffineIfOp with 'withElse = true' because yielded values must
+ // be passed through both branches (e.g., propagating parentStore value in the
+ // else branch).
+ auto newIf =
+ AffineIfOp::create(rewriter, ifOp->getLoc(), newTypes,
+ ifOp.getIntegerSet(), ifOp->getOperands(), true);
+
+ // Take blocks from oldIf into newIf.
+ newIf.getThenRegion().takeBody(ifOp.getThenRegion());
+ if (ifOp.hasElse()) {
+ newIf.getElseRegion().takeBody(ifOp.getElseRegion());
+ } else {
+ // If we create an else block, we need to explicitly insert a yield.
+ rewriter.setInsertionPointToEnd(newIf.getElseBlock());
+ AffineYieldOp::create(rewriter, newIf->getLoc());
+ }
+
+ // Update yield operands.
+ Block *thenBlock = newIf.getThenBlock();
+ Block *elseBlock = newIf.getElseBlock();
+
+ AffineYieldOp thenYield = cast<AffineYieldOp>(thenBlock->getTerminator());
+ AffineYieldOp elseYield = cast<AffineYieldOp>(elseBlock->getTerminator());
+
+ // Preserve existing yield values from both branches.
+ SmallVector<Value, 4> thenYieldVals(thenYield.getOperands());
+ SmallVector<Value, 4> elseYieldVals(elseYield.getOperands());
+
+ for (auto [thenStore, elseStore, parentStore] : plan.pairs) {
+ // 1. Both branches have stores: yield their respective store values.
+ // 2 & 3. Only one branch has a store: yield that store value in its branch,
+ // and propagate parentStore's value through the other branch.
+ if (thenStore && elseStore) {
+ thenYieldVals.push_back(thenStore.getValue());
+ elseYieldVals.push_back(elseStore.getValue());
+ } else if (thenStore) {
+ thenYieldVals.push_back(thenStore.getValue());
+ elseYieldVals.push_back(parentStore.getValue());
+ } else {
+ thenYieldVals.push_back(parentStore.getValue());
+ elseYieldVals.push_back(elseStore.getValue());
+ }
+ }
+
+ // Update Then block terminator.
+ rewriter.setInsertionPoint(thenYield);
+ AffineYieldOp::create(rewriter, thenYield->getLoc(), thenYieldVals);
+ rewriter.eraseOp(thenYield);
+
+ // Update Else block terminator.
+ rewriter.setInsertionPoint(elseYield);
+ AffineYieldOp::create(rewriter, elseYield->getLoc(), elseYieldVals);
+ rewriter.eraseOp(elseYield);
+
+ // Emit unified AffineStoreOps after newIf using appended results.
+ rewriter.setInsertionPointAfter(newIf);
+ unsigned baseIdx = ifOp.getNumResults();
+ for (auto it : llvm::enumerate(plan.pairs)) {
+ AffineStoreOp store = it.value().parentStore;
+ Value valueToStore = newIf.getResult(baseIdx + it.index());
+ AffineStoreOp::create(rewriter, store->getLoc(), valueToStore,
+ store.getMemRef(), store.getIndices());
+ if (it.value().elseStore)
+ rewriter.eraseOp(it.value().elseStore);
+ if (it.value().thenStore)
+ rewriter.eraseOp(it.value().thenStore);
+ rewriter.eraseOp(store);
+ }
+
+ // Replace original ifOp results and erase stale oldIf.
+ for (auto [oldValue, newValue] : llvm::zip_equal(
+ ifOp.getResults(), newIf.getResults().take_front(baseIdx)))
+ rewriter.replaceAllUsesWith(oldValue, newValue);
+ rewriter.eraseOp(ifOp);
+}
+
struct AffineScalarReplacement
: public affine::impl::AffineScalarReplacementBase<
AffineScalarReplacement> {
@@ -45,6 +271,18 @@ mlir::affine::createAffineScalarReplacementPass() {
}
void AffineScalarReplacement::runOnOperation() {
+ SmallVector<IfSinkPlan> plans;
+ IRRewriter rewriter(getOperation());
+ getOperation()->walk([&](AffineIfOp ifOp) {
+ IfSinkPlan plan{/*ifOp=*/ifOp, /*pairs=*/{}};
+ analyzeIfOp(ifOp, plan);
+ if (!plan.pairs.empty()) {
+ plans.push_back(plan);
+ }
+ });
+ for (IfSinkPlan &plan : plans)
+ applySinkPlan(rewriter, plan);
+
affineScalarReplace(getOperation(), getAnalysis<DominanceInfo>(),
getAnalysis<PostDominanceInfo>(),
getAnalysis<AliasAnalysis>());
diff --git a/mlir/test/Dialect/Affine/scalrep.mlir b/mlir/test/Dialect/Affine/scalrep.mlir
index fb6eef941790a..3413bb3cf61a3 100644
--- a/mlir/test/Dialect/Affine/scalrep.mlir
+++ b/mlir/test/Dialect/Affine/scalrep.mlir
@@ -6,6 +6,7 @@
// CHECK-DAG: [[$MAP3:#map[0-9]*]] = affine_map<(d0, d1) -> (d0 - 1)>
// CHECK-DAG: [[$MAP4:#map[0-9]*]] = affine_map<(d0) -> (d0 + 1)>
// CHECK-DAG: [[$IDENT:#map[0-9]*]] = affine_map<(d0) -> (d0)>
+// CHECK-DAG: [[$SET:#set[0-9]*]] = affine_set<(d0) : (d0 - 1 >= 0)>
// CHECK-LABEL: func @simple_store_load() {
func.func @simple_store_load() {
@@ -1034,3 +1035,85 @@ func.func @vector_store_dead_elim_same_type(%arg0: memref<20x1xi64>) {
affine.vector_store %cst2, %arg0[%c0, %c0] : memref<20x1xi64>, vector<5xi64>
return
}
+
+#set = affine_set<(d0) : (d0 - 1 >= 0)>
+
+// CHECK-LABEL: func @sink_stores_both_branches
+// CHECK-SAME: %[[ARG0:.*]]: memref<10xf32>, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32, %[[ARG3:.*]]: f32, %[[ARG4:.*]]: index)
+func.func @sink_stores_both_branches(%mem: memref<10xf32>, %arg0: f32, %arg1: f32, %arg2: f32, %N: index) {
+ %c0 = arith.constant 0 : index
+ affine.store %arg0, %mem[%c0] : memref<10xf32>
+ affine.if #set(%N) {
+ affine.store %arg1, %mem[%c0] : memref<10xf32>
+ } else {
+ affine.store %arg2, %mem[%c0] : memref<10xf32>
+ }
+ return
+}
+// CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
+// CHECK: %[[IF_0:.*]] = affine.if [[$SET]](%[[ARG4]]) -> f32 {
+// CHECK: affine.yield %[[ARG2]] : f32
+// CHECK: } else {
+// CHECK: affine.yield %[[ARG3]] : f32
+// CHECK: }
+// CHECK: affine.store %[[IF_0]], %[[ARG0]]{{\[}}%[[CONSTANT_0]]] : memref<10xf32>
+
+// CHECK-LABEL: func @sink_stores_then_only
+// CHECK-SAME: %[[ARG0:.*]]: memref<10xf32>, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32, %[[ARG3:.*]]: index)
+func.func @sink_stores_then_only(%mem: memref<10xf32>, %arg0: f32, %arg1: f32, %N: index) {
+ %c0 = arith.constant 0 : index
+ affine.store %arg0, %mem[%c0] : memref<10xf32>
+ affine.if #set(%N) {
+ affine.store %arg1, %mem[%c0] : memref<10xf32>
+ }
+ return
+}
+// CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
+// CHECK: %[[IF_0:.*]] = affine.if [[$SET]](%[[ARG3]]) -> f32 {
+// CHECK: affine.yield %[[ARG1]] : f32
+// CHECK: } else {
+// CHECK: affine.yield %[[ARG2]] : f32
+// CHECK: }
+// CHECK: affine.store %[[IF_0]], %[[ARG0]]{{\[}}%[[CONSTANT_0]]] : memref<10xf32>
+
+// CHECK-LABEL: func @sink_stores_else_only
+// CHECK-SAME: %[[ARG0:.*]]: memref<10xf32>, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32, %[[ARG3:.*]]: index)
+func.func @sink_stores_else_only(%mem: memref<10xf32>, %arg0: f32, %arg2: f32, %N: index) {
+ %c0 = arith.constant 0 : index
+ affine.store %arg0, %mem[%c0] : memref<10xf32>
+ affine.if #set(%N) {
+ } else {
+ affine.store %arg2, %mem[%c0] : memref<10xf32>
+ }
+ return
+}
+// CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
+// CHECK: %[[IF_0:.*]] = affine.if [[$SET]](%[[ARG3]]) -> f32 {
+// CHECK: affine.yield %[[ARG1]] : f32
+// CHECK: } else {
+// CHECK: affine.yield %[[ARG2]] : f32
+// CHECK: }
+// CHECK: affine.store %[[IF_0]], %[[ARG0]]{{\[}}%[[CONSTANT_0]]] : memref<10xf32>
+
+// CHECK-LABEL: func @sink_multiple_pairs
+// CHECK-SAME: %[[ARG0:.*]]: memref<10xf32>, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32, %[[ARG3:.*]]: index)
+func.func @sink_multiple_pairs(%mem: memref<10xf32>, %arg0: f32, %arg1: f32, %N: index) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ affine.store %arg0, %mem[%c0] : memref<10xf32>
+ affine.store %arg0, %mem[%c1] : memref<10xf32>
+ affine.if #set(%N) {
+ affine.store %arg1, %mem[%c0] : memref<10xf32>
+ affine.store %arg1, %mem[%c1] : memref<10xf32>
+ }
+ return
+}
+// CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
+// CHECK: %[[CONSTANT_1:.*]] = arith.constant 1 : index
+// CHECK: %[[IF_0:.*]]:2 = affine.if [[$SET]](%[[ARG3]]) -> (f32, f32) {
+// CHECK: affine.yield %[[ARG1]], %[[ARG1]] : f32, f32
+// CHECK: } else {
+// CHECK: affine.yield %[[ARG2]], %[[ARG2]] : f32, f32
+// CHECK: }
+// CHECK: affine.store %[[VAL_0:.*]]#0, %[[ARG0]]{{\[}}%[[CONSTANT_1]]] : memref<10xf32>
+// CHECK: affine.store %[[VAL_0]]#1, %[[ARG0]]{{\[}}%[[CONSTANT_0]]] : memref<10xf32>
>From 3efaae7dabe6daf6560c7a5e729762d6cfe9183d Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sun, 2 Aug 2026 11:17:47 +0000
Subject: [PATCH 2/3] fix nit.
---
.../Affine/Transforms/AffineScalarReplacement.cpp | 9 ---------
1 file changed, 9 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
index 88b149d819a00..cbdda4e6df236 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
@@ -12,20 +12,11 @@
//===----------------------------------------------------------------------===//
#include "mlir/Analysis/AliasAnalysis.h"
-#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Transforms/Passes.h"
#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/Block.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/IntegerSet.h"
-#include "mlir/IR/Operation.h"
-#include "mlir/IR/OperationSupport.h"
-#include "mlir/IR/PatternMatch.h"
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/DebugLog.h"
namespace mlir {
namespace affine {
>From bc46160aa314ce031b255deafa215512a307cfbb Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sun, 2 Aug 2026 12:04:23 +0000
Subject: [PATCH 3/3] fix test.
---
.../Transforms/AffineScalarReplacement.cpp | 6 ++---
mlir/test/Dialect/Affine/scalrep.mlir | 25 +++++++++++++------
2 files changed, 20 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
index cbdda4e6df236..e47aecfc7e403 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
@@ -161,11 +161,12 @@ static void analyzeIfOp(AffineIfOp ifOp, IfSinkPlan &plan) {
/// `AffineYieldOp` terminators in both 'then' and 'else' blocks, erases
/// internal store operations, and emits a single unified `AffineStoreOp`
/// immediately following the new `ifOp` for each sinked pair.
-static void applySinkPlan(RewriterBase &rewriter, IfSinkPlan &plan) {
+static void applySinkPlan(IfSinkPlan &plan) {
if (plan.pairs.empty())
return;
AffineIfOp ifOp = plan.ifOp;
+ IRRewriter rewriter(ifOp);
// Collect new result types for the new AffineIfOp.
SmallVector<Type, 4> newTypes(ifOp.getResultTypes());
@@ -263,7 +264,6 @@ mlir::affine::createAffineScalarReplacementPass() {
void AffineScalarReplacement::runOnOperation() {
SmallVector<IfSinkPlan> plans;
- IRRewriter rewriter(getOperation());
getOperation()->walk([&](AffineIfOp ifOp) {
IfSinkPlan plan{/*ifOp=*/ifOp, /*pairs=*/{}};
analyzeIfOp(ifOp, plan);
@@ -272,7 +272,7 @@ void AffineScalarReplacement::runOnOperation() {
}
});
for (IfSinkPlan &plan : plans)
- applySinkPlan(rewriter, plan);
+ applySinkPlan(plan);
affineScalarReplace(getOperation(), getAnalysis<DominanceInfo>(),
getAnalysis<PostDominanceInfo>(),
diff --git a/mlir/test/Dialect/Affine/scalrep.mlir b/mlir/test/Dialect/Affine/scalrep.mlir
index 3413bb3cf61a3..1c1c4f9613547 100644
--- a/mlir/test/Dialect/Affine/scalrep.mlir
+++ b/mlir/test/Dialect/Affine/scalrep.mlir
@@ -1040,7 +1040,7 @@ func.func @vector_store_dead_elim_same_type(%arg0: memref<20x1xi64>) {
// CHECK-LABEL: func @sink_stores_both_branches
// CHECK-SAME: %[[ARG0:.*]]: memref<10xf32>, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32, %[[ARG3:.*]]: f32, %[[ARG4:.*]]: index)
-func.func @sink_stores_both_branches(%mem: memref<10xf32>, %arg0: f32, %arg1: f32, %arg2: f32, %N: index) {
+func.func @sink_stores_both_branches(%mem: memref<10xf32>, %arg0: f32, %arg1: f32, %arg2: f32, %N: index) -> f32 {
%c0 = arith.constant 0 : index
affine.store %arg0, %mem[%c0] : memref<10xf32>
affine.if #set(%N) {
@@ -1048,7 +1048,8 @@ func.func @sink_stores_both_branches(%mem: memref<10xf32>, %arg0: f32, %arg1: f3
} else {
affine.store %arg2, %mem[%c0] : memref<10xf32>
}
- return
+ %res = affine.load %mem[%c0] : memref<10xf32>
+ return %res : f32
}
// CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
// CHECK: %[[IF_0:.*]] = affine.if [[$SET]](%[[ARG4]]) -> f32 {
@@ -1057,16 +1058,18 @@ func.func @sink_stores_both_branches(%mem: memref<10xf32>, %arg0: f32, %arg1: f3
// CHECK: affine.yield %[[ARG3]] : f32
// CHECK: }
// CHECK: affine.store %[[IF_0]], %[[ARG0]]{{\[}}%[[CONSTANT_0]]] : memref<10xf32>
+// CHECK: return %[[IF_0]] : f32
// CHECK-LABEL: func @sink_stores_then_only
// CHECK-SAME: %[[ARG0:.*]]: memref<10xf32>, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32, %[[ARG3:.*]]: index)
-func.func @sink_stores_then_only(%mem: memref<10xf32>, %arg0: f32, %arg1: f32, %N: index) {
+func.func @sink_stores_then_only(%mem: memref<10xf32>, %arg0: f32, %arg1: f32, %N: index) -> f32 {
%c0 = arith.constant 0 : index
affine.store %arg0, %mem[%c0] : memref<10xf32>
affine.if #set(%N) {
affine.store %arg1, %mem[%c0] : memref<10xf32>
}
- return
+ %res = affine.load %mem[%c0] : memref<10xf32>
+ return %res : f32
}
// CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
// CHECK: %[[IF_0:.*]] = affine.if [[$SET]](%[[ARG3]]) -> f32 {
@@ -1075,17 +1078,19 @@ func.func @sink_stores_then_only(%mem: memref<10xf32>, %arg0: f32, %arg1: f32, %
// CHECK: affine.yield %[[ARG2]] : f32
// CHECK: }
// CHECK: affine.store %[[IF_0]], %[[ARG0]]{{\[}}%[[CONSTANT_0]]] : memref<10xf32>
+// CHECK: return %[[IF_0]] : f32
// CHECK-LABEL: func @sink_stores_else_only
// CHECK-SAME: %[[ARG0:.*]]: memref<10xf32>, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32, %[[ARG3:.*]]: index)
-func.func @sink_stores_else_only(%mem: memref<10xf32>, %arg0: f32, %arg2: f32, %N: index) {
+func.func @sink_stores_else_only(%mem: memref<10xf32>, %arg0: f32, %arg2: f32, %N: index) -> f32 {
%c0 = arith.constant 0 : index
affine.store %arg0, %mem[%c0] : memref<10xf32>
affine.if #set(%N) {
} else {
affine.store %arg2, %mem[%c0] : memref<10xf32>
}
- return
+ %res = affine.load %mem[%c0] : memref<10xf32>
+ return %res : f32
}
// CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
// CHECK: %[[IF_0:.*]] = affine.if [[$SET]](%[[ARG3]]) -> f32 {
@@ -1094,10 +1099,11 @@ func.func @sink_stores_else_only(%mem: memref<10xf32>, %arg0: f32, %arg2: f32, %
// CHECK: affine.yield %[[ARG2]] : f32
// CHECK: }
// CHECK: affine.store %[[IF_0]], %[[ARG0]]{{\[}}%[[CONSTANT_0]]] : memref<10xf32>
+// CHECK: return %[[IF_0]] : f32
// CHECK-LABEL: func @sink_multiple_pairs
// CHECK-SAME: %[[ARG0:.*]]: memref<10xf32>, %[[ARG1:.*]]: f32, %[[ARG2:.*]]: f32, %[[ARG3:.*]]: index)
-func.func @sink_multiple_pairs(%mem: memref<10xf32>, %arg0: f32, %arg1: f32, %N: index) {
+func.func @sink_multiple_pairs(%mem: memref<10xf32>, %arg0: f32, %arg1: f32, %N: index) -> (f32, f32) {
%c0 = arith.constant 0 : index
%c1 = arith.constant 1 : index
affine.store %arg0, %mem[%c0] : memref<10xf32>
@@ -1106,7 +1112,9 @@ func.func @sink_multiple_pairs(%mem: memref<10xf32>, %arg0: f32, %arg1: f32, %N:
affine.store %arg1, %mem[%c0] : memref<10xf32>
affine.store %arg1, %mem[%c1] : memref<10xf32>
}
- return
+ %res0 = affine.load %mem[%c0] : memref<10xf32>
+ %res1 = affine.load %mem[%c1] : memref<10xf32>
+ return %res0, %res1 : f32, f32
}
// CHECK: %[[CONSTANT_0:.*]] = arith.constant 0 : index
// CHECK: %[[CONSTANT_1:.*]] = arith.constant 1 : index
@@ -1117,3 +1125,4 @@ func.func @sink_multiple_pairs(%mem: memref<10xf32>, %arg0: f32, %arg1: f32, %N:
// CHECK: }
// CHECK: affine.store %[[VAL_0:.*]]#0, %[[ARG0]]{{\[}}%[[CONSTANT_1]]] : memref<10xf32>
// CHECK: affine.store %[[VAL_0]]#1, %[[ARG0]]{{\[}}%[[CONSTANT_0]]] : memref<10xf32>
+// CHECK: return %[[IF_0]]#1, %[[IF_0]]#0 : f32, f32
More information about the Mlir-commits
mailing list