[Mlir-commits] [mlir] [mlir][affine] Add store sinking transformation for AffineIfOp (PR #213531)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sun Aug 2 04:23:48 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-affine
@llvm/pr-subscribers-mlir
Author: lonely eagle (linuxlonelyeagle)
<details>
<summary>Changes</summary>
This PR introduces store sinking for AffineIfOp in the Affine dialect. It identifies AffineStoreOps in then and/or else blocks that target the same memory location as a preceding parentStore in the same outer block, and sinks them past the ifOp. By yielding branch store values via AffineYieldOp, it unifies memory writes into a single post-ifOp AffineStoreOp, eliminating redundant store operations.
---
Full diff: https://github.com/llvm/llvm-project/pull/213531.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp (+231-2)
- (modified) mlir/test/Dialect/Affine/scalrep.mlir (+83)
``````````diff
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
index 92a4cf3c0ad76..cbdda4e6df236 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineScalarReplacement.cpp
@@ -11,12 +11,12 @@
// redundant loads.
//===----------------------------------------------------------------------===//
-#include "mlir/Dialect/Affine/Transforms/Passes.h"
-
#include "mlir/Analysis/AliasAnalysis.h"
+#include "mlir/Dialect/Affine/Transforms/Passes.h"
#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Dominance.h"
+#include "mlir/IR/IntegerSet.h"
namespace mlir {
namespace affine {
@@ -31,6 +31,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 +262,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>
``````````
</details>
https://github.com/llvm/llvm-project/pull/213531
More information about the Mlir-commits
mailing list