[Mlir-commits] [mlir] [mlir][bufferization] Add buffer-loop-merging pass (PR #213183)
Jianhui Li
llvmlistbot at llvm.org
Thu Jul 30 17:54:24 PDT 2026
https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/213183
Adds a `buffer-loop-merging` pass to the Bufferization dialect that folds a `scf.for` memref iter_arg whose initial and yielded values are distinct, interchangeable `memref.alloca`s onto a single buffer. This makes the iter_arg loop-invariant so that canonicalization drops it and Mem2Reg can promote the underlying slot.
mem2reg cannot do this itself: it processes one allocation slot at a time and has no cross-allocation view, so unifying two distinct allocas must be a pre-pass. The transform reasons about buffer identity and lifetime, so it lives in Bufferization alongside buffer-hoisting / buffer-loop-hoisting rather than in the SCF (loop-structure) passes.
The rewrite only redirects a buffer's uses when the two allocations remain indistinguishable afterwards: `init` may only be used outside the loop strictly before it, and `yielded` may not be written outside the loop. True ping-pong swaps, partial-buffer accesses, heap allocations, and escaping buffers are left untouched. Tests cover the converging scalar/vector cases, all rejection cases, and regressions for the two outside-the-loop miscompiles.
assisted-by-claude
>From b53e069b2c554e9aabb76146a54ad3d1a5eb4f4b Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 31 Jul 2026 00:45:53 +0000
Subject: [PATCH] [mlir][bufferization] Add buffer-loop-merging pass
Adds a `buffer-loop-merging` pass to the Bufferization dialect that folds a
`scf.for` memref iter_arg whose initial and yielded values are distinct,
interchangeable `memref.alloca`s onto a single buffer. This makes the iter_arg
loop-invariant so that canonicalization drops it and Mem2Reg can promote the
underlying slot.
mem2reg cannot do this itself: it processes one allocation slot at a time and
has no cross-allocation view, so unifying two distinct allocas must be a
pre-pass. The transform reasons about buffer identity and lifetime, so it lives
in Bufferization alongside buffer-hoisting / buffer-loop-hoisting rather than in
the SCF (loop-structure) passes.
The rewrite only redirects a buffer's uses when the two allocations remain
indistinguishable afterwards: `init` may only be used outside the loop strictly
before it, and `yielded` may not be written outside the loop. True ping-pong
swaps, partial-buffer accesses, heap allocations, and escaping buffers are left
untouched. Tests cover the converging scalar/vector cases, all rejection cases,
and regressions for the two outside-the-loop miscompiles.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../Bufferization/Transforms/Passes.td | 12 +
.../Transforms/BufferLoopMerging.cpp | 254 ++++++++++++++++++
.../Bufferization/Transforms/CMakeLists.txt | 1 +
.../Transforms/buffer-loop-merging.mlir | 216 +++++++++++++++
4 files changed, 483 insertions(+)
create mode 100644 mlir/lib/Dialect/Bufferization/Transforms/BufferLoopMerging.cpp
create mode 100644 mlir/test/Dialect/Bufferization/Transforms/buffer-loop-merging.mlir
diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 8408315dda607..37add79bed052 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -303,6 +303,18 @@ def BufferLoopHoistingPass : Pass<"buffer-loop-hoisting", "func::FuncOp"> {
}];
}
+def BufferLoopMergingPass : Pass<"buffer-loop-merging", "func::FuncOp"> {
+ let summary = "Merge scf.for memref iter_args that converge onto one buffer";
+ let description = [{
+ Rewrites `scf.for` memref iter_args whose initial and yielded values are
+ distinct `memref.alloca`s used interchangeably across iterations. Folding the
+ yielded buffer into the initial one makes the iter_arg loop-invariant,
+ allowing canonicalization to drop it and Mem2Reg to promote the underlying
+ slot.
+ }];
+ let dependentDialects = ["memref::MemRefDialect", "scf::SCFDialect"];
+}
+
def BufferResultsToOutParamsPass
: Pass<"buffer-results-to-out-params", "ModuleOp"> {
let summary = "Converts memref-typed function results to out-params";
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferLoopMerging.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferLoopMerging.cpp
new file mode 100644
index 0000000000000..0298da2827fd9
--- /dev/null
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferLoopMerging.cpp
@@ -0,0 +1,254 @@
+//===- BufferLoopMerging.cpp - Merge converging scf.for buffers -----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements merging of `scf.for` memref iter_args that converge onto
+// a single buffer after the first iteration.
+//
+// Code that stages a value through scratch memory across loop iterations often
+// threads one buffer in as the initial value and yields a *different* buffer
+// from the body:
+//
+// %yield = memref.alloca()
+// %init = memref.alloca()
+// store %v, %init[]
+// %r = scf.for ... iter_args(%it = %init) -> (memref<f32>) {
+// %x = load %it[] // %init on iteration 0, %yield afterwards
+// store f(%x), %yield[]
+// scf.yield %yield
+// }
+//
+// The iter_arg is not loop-invariant, so the `scf.for` remains a blocking use of
+// both allocations and Mem2Reg abandons them (`scf.for` implements
+// `PromotableRegionOpInterface`, but is not itself a promotable or aliasing op).
+// The buffers are nonetheless interchangeable: `%it` denotes `%init` only on the
+// first iteration and `%yield` on every later one, and no iteration reads a
+// buffer it has not first written except through `%it`. Rewriting all uses of
+// `%yield` to `%init` makes the iter_arg loop-invariant, after which existing
+// canonicalization drops it and Mem2Reg promotes the slot.
+//
+// Merging redirects the body's stores from `%yield` into `%init`, so it is only
+// sound when the buffers cannot be distinguished outside the loop after that
+// redirection: every non-threading use of `%init` must strictly precede the loop
+// (its contents change from the loop onward), and `%yield` must never be written
+// outside the loop (such a write would be lost or reordered).
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/Dominance.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+namespace mlir {
+namespace bufferization {
+#define GEN_PASS_DEF_BUFFERLOOPMERGINGPASS
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h.inc"
+} // namespace bufferization
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::bufferization;
+using namespace mlir::scf;
+
+namespace {
+
+/// How an operation uses a buffer value.
+enum class BufferUseKind {
+ /// The use is the `scf.yield` of the loop being considered, or an init operand
+ /// of the loop itself. Both are the threading of the buffer through the
+ /// iter_arg, which this transform is rewriting.
+ Yield,
+ /// The operation only reads the buffer contents.
+ Read,
+ /// The operation only writes the buffer contents.
+ Write,
+ /// The operation both reads and writes the buffer contents.
+ ReadWrite,
+ /// The operation may expose the buffer's identity (view, cast, call, ...) or
+ /// its effects on the buffer cannot be determined.
+ Opaque
+};
+
+} // namespace
+
+/// Classifies how `use` accesses the buffer it refers to. Only operations whose
+/// memory effects are fully known and limited to reads and writes of the operand
+/// are safe to redirect; anything else may observe the buffer's address rather
+/// than just its contents, which merging does not preserve.
+static BufferUseKind classifyBufferUse(OpOperand &use, ForOp loop) {
+ Operation *user = use.getOwner();
+ if (isa<scf::YieldOp>(user) && user->getParentOp() == loop)
+ return BufferUseKind::Yield;
+ // The loop's own init operand: the caller has already verified that this
+ // threading is the converging pattern being merged.
+ if (user == loop.getOperation())
+ return BufferUseKind::Yield;
+
+ auto effectOp = dyn_cast<MemoryEffectOpInterface>(user);
+ if (!effectOp)
+ return BufferUseKind::Opaque;
+
+ // Operations with regions could hide accesses that the effect list does not
+ // attribute to this operand.
+ if (user->getNumRegions() != 0)
+ return BufferUseKind::Opaque;
+
+ SmallVector<MemoryEffects::EffectInstance> effects;
+ effectOp.getEffects(effects);
+
+ bool reads = false;
+ bool writes = false;
+ for (const MemoryEffects::EffectInstance &effect : effects) {
+ // An effect that is not pinned to a specific value may apply to this
+ // buffer; conservatively treat the operation as opaque.
+ Value effectValue = effect.getValue();
+ if (!effectValue)
+ return BufferUseKind::Opaque;
+ if (effectValue != use.get())
+ continue;
+ if (isa<MemoryEffects::Read>(effect.getEffect())) {
+ reads = true;
+ continue;
+ }
+ if (isa<MemoryEffects::Write>(effect.getEffect())) {
+ writes = true;
+ continue;
+ }
+ // Allocate/Free on the buffer means the operation controls its lifetime.
+ return BufferUseKind::Opaque;
+ }
+
+ // A use that carries no effect on the buffer is a plain capture of the
+ // pointer, e.g. a view or cast that forwards it.
+ if (!reads && !writes)
+ return BufferUseKind::Opaque;
+ if (reads && writes)
+ return BufferUseKind::ReadWrite;
+ return reads ? BufferUseKind::Read : BufferUseKind::Write;
+}
+
+/// Returns true if `value` is produced by an allocation whose lifetime is the
+/// enclosing scope and which cannot be freed explicitly. Only `memref.alloca`
+/// qualifies, so redirecting its uses is not observable outside the scope.
+static bool isMergeableAlloc(Value value) {
+ return value.getDefiningOp<memref::AllocaOp>() != nullptr;
+}
+
+/// Attempts to merge a converging memref iter_arg of `loop` at `argIdx`.
+/// Returns true if the IR was modified.
+static bool tryMergeIterArg(ForOp loop, unsigned argIdx,
+ DominanceInfo &dominance) {
+ Value init = loop.getInitArgs()[argIdx];
+ auto yieldOp = cast<scf::YieldOp>(loop.getBody()->getTerminator());
+ Value yielded = yieldOp.getOperands()[argIdx];
+ Value iterArg = loop.getRegionIterArgs()[argIdx];
+
+ // Only memref-typed slots participate; the merge is about buffer identity.
+ if (!isa<MemRefType>(init.getType()))
+ return false;
+
+ // Already loop-invariant: canonicalization handles this case on its own.
+ if (yielded == init)
+ return false;
+
+ // A yielded value derived from the block argument means the rotation is
+ // genuinely dynamic (a ping-pong swap), which this transform cannot merge.
+ if (yielded == iterArg)
+ return false;
+
+ if (!isMergeableAlloc(init) || !isMergeableAlloc(yielded))
+ return false;
+
+ // Identical types keep the rewrite a pure use replacement.
+ if (init.getType() != yielded.getType())
+ return false;
+
+ // Both allocations must dominate the loop. One defined in the body would be
+ // fresh per iteration and could not be unified with the initial buffer.
+ if (!loop.isDefinedOutsideOfLoop(init) ||
+ !loop.isDefinedOutsideOfLoop(yielded))
+ return false;
+
+ // The two buffers must not be threaded through any other iter_arg of this
+ // loop, where they could be rotated under a different schedule.
+ for (unsigned i = 0, e = loop.getInitArgs().size(); i < e; ++i) {
+ if (i == argIdx)
+ continue;
+ if (loop.getInitArgs()[i] == init || loop.getInitArgs()[i] == yielded ||
+ yieldOp.getOperands()[i] == init || yieldOp.getOperands()[i] == yielded)
+ return false;
+ }
+
+ // Merging collapses two buffers into one, so the two buffers must remain
+ // indistinguishable once the body's stores are redirected from `yielded` into
+ // `init`. The two allocations play asymmetric roles:
+ //
+ // * `init` is read inside the body only through the iter_arg. Any direct
+ // in-loop use would, after merging, observe this iteration's own store.
+ // Outside the loop, redirecting the body's stores into `init` changes its
+ // contents from the loop onward, so every non-threading use must strictly
+ // precede the loop.
+ // * `yielded` is written inside the body to define what the next iteration
+ // reads through the iter_arg; only pure writes are allowed there. Outside
+ // the loop it may be read (a read observes the same bytes as `init` after
+ // merging) but never written -- an outside write would be lost.
+ auto usesAreMergeable = [&](Value buffer, bool isInit) {
+ for (OpOperand &use : buffer.getUses()) {
+ BufferUseKind kind = classifyBufferUse(use, loop);
+ if (kind == BufferUseKind::Opaque)
+ return false;
+ // The threading through the iter_arg is what this transform rewrites.
+ if (kind == BufferUseKind::Yield)
+ continue;
+
+ if (loop->isProperAncestor(use.getOwner())) {
+ // Inside the loop body.
+ if (isInit)
+ return false;
+ if (kind != BufferUseKind::Write)
+ return false;
+ continue;
+ }
+
+ // Outside the loop.
+ if (isInit) {
+ if (!dominance.properlyDominates(use.getOwner(), loop))
+ return false;
+ } else if (kind != BufferUseKind::Read) {
+ return false;
+ }
+ }
+ return true;
+ };
+
+ if (!usesAreMergeable(init, /*isInit=*/true) ||
+ !usesAreMergeable(yielded, /*isInit=*/false))
+ return false;
+
+ // All preconditions hold: fold `yielded` into `init`. This makes the iter_arg
+ // loop-invariant and leaves `yielded` dead.
+ yielded.replaceAllUsesWith(init);
+ return true;
+}
+
+namespace {
+struct BufferLoopMergingPass
+ : public bufferization::impl::BufferLoopMergingPassBase<
+ BufferLoopMergingPass> {
+ void runOnOperation() override {
+ DominanceInfo dominance(getOperation());
+ getOperation()->walk([&](ForOp loop) {
+ for (unsigned i = 0, e = loop.getInitArgs().size(); i < e; ++i)
+ (void)tryMergeIterArg(loop, i, dominance);
+ });
+ }
+};
+} // namespace
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt
index 006fcd1ce0ec7..87b8bbebe637d 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Bufferization/Transforms/CMakeLists.txt
@@ -1,6 +1,7 @@
add_mlir_dialect_library(MLIRBufferizationTransforms
Bufferize.cpp
BufferDeallocationSimplification.cpp
+ BufferLoopMerging.cpp
BufferOptimizations.cpp
BufferResultsToOutParams.cpp
BufferUtils.cpp
diff --git a/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-merging.mlir b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-merging.mlir
new file mode 100644
index 0000000000000..b2f7167c6fd16
--- /dev/null
+++ b/mlir/test/Dialect/Bufferization/Transforms/buffer-loop-merging.mlir
@@ -0,0 +1,216 @@
+// RUN: mlir-opt %s -buffer-loop-merging -split-input-file | FileCheck %s
+
+// A memref iter_arg whose init and yielded values are distinct allocas that are
+// used interchangeably across iterations is merged onto the init buffer, making
+// the iter_arg loop-invariant.
+
+// CHECK-LABEL: func.func @merge_converging_buffers
+// CHECK: memref.alloca() : memref<f32>
+// CHECK: %[[INIT:.*]] = memref.alloca() : memref<f32>
+// CHECK: memref.store %{{.*}}, %[[INIT]][]
+// CHECK: scf.for {{.*}} iter_args(%{{.*}} = %[[INIT]])
+// CHECK: memref.store %{{.*}}, %[[INIT]][]
+// CHECK: scf.yield %[[INIT]]
+func.func @merge_converging_buffers(%init: f32, %lb: index, %ub: index, %st: index) -> f32 {
+ %alloc_yield = memref.alloca() : memref<f32>
+ %alloc_init = memref.alloca() : memref<f32>
+ memref.store %init, %alloc_init[] : memref<f32>
+ %r = scf.for %i = %lb to %ub step %st iter_args(%lv = %alloc_init) -> (memref<f32>) {
+ %v = memref.load %lv[] : memref<f32>
+ %n = arith.addf %v, %v : f32
+ memref.store %n, %alloc_yield[] : memref<f32>
+ scf.yield %alloc_yield : memref<f32>
+ }
+ %o = memref.load %r[] : memref<f32>
+ return %o : f32
+}
+
+// -----
+
+// Same pattern with whole-buffer vector transfers instead of load/store.
+
+// CHECK-LABEL: func.func @merge_converging_buffers_vector
+// CHECK: memref.alloca() : memref<128xf32>
+// CHECK: %[[INIT:.*]] = memref.alloca() : memref<128xf32>
+// CHECK: vector.transfer_write %{{.*}}, %[[INIT]]
+// CHECK: scf.for {{.*}} iter_args(%{{.*}} = %[[INIT]])
+// CHECK: vector.transfer_write %{{.*}}, %[[INIT]]
+// CHECK: scf.yield %[[INIT]]
+func.func @merge_converging_buffers_vector(%pad: f32, %init: vector<128xf32>,
+ %lb: index, %ub: index, %st: index) -> vector<128xf32> {
+ %c0 = arith.constant 0 : index
+ %alloc_yield = memref.alloca() : memref<128xf32>
+ %alloc_init = memref.alloca() : memref<128xf32>
+ vector.transfer_write %init, %alloc_init[%c0] {in_bounds = [true]} : vector<128xf32>, memref<128xf32>
+ %r = scf.for %i = %lb to %ub step %st iter_args(%lv = %alloc_init) -> (memref<128xf32>) {
+ %v = vector.transfer_read %lv[%c0], %pad {in_bounds = [true]} : memref<128xf32>, vector<128xf32>
+ %n = arith.addf %v, %v : vector<128xf32>
+ vector.transfer_write %n, %alloc_yield[%c0] {in_bounds = [true]} : vector<128xf32>, memref<128xf32>
+ scf.yield %alloc_yield : memref<128xf32>
+ }
+ %o = vector.transfer_read %r[%c0], %pad {in_bounds = [true]} : memref<128xf32>, vector<128xf32>
+ return %o : vector<128xf32>
+}
+
+// -----
+
+// A genuine ping-pong swap rotates both buffers every iteration, so which buffer
+// a given iteration reads is not statically known. Must not merge.
+
+// CHECK-LABEL: func.func @no_merge_ping_pong
+// CHECK: scf.yield %{{.*}}, %{{.*}} : memref<f32>, memref<f32>
+func.func @no_merge_ping_pong(%init: f32, %lb: index, %ub: index, %st: index) -> f32 {
+ %a = memref.alloca() : memref<f32>
+ %b = memref.alloca() : memref<f32>
+ memref.store %init, %a[] : memref<f32>
+ %r:2 = scf.for %i = %lb to %ub step %st iter_args(%src = %a, %dst = %b) -> (memref<f32>, memref<f32>) {
+ %v = memref.load %src[] : memref<f32>
+ %n = arith.addf %v, %v : f32
+ memref.store %n, %dst[] : memref<f32>
+ scf.yield %dst, %src : memref<f32>, memref<f32>
+ }
+ %o = memref.load %r#0[] : memref<f32>
+ return %o : f32
+}
+
+// -----
+
+// The yielded buffer is also read inside the body, so merging would make that
+// read observe the current iteration's store. Must not merge.
+
+// CHECK-LABEL: func.func @no_merge_yielded_is_read
+// CHECK: %[[YIELD:.*]] = memref.alloca() : memref<f32>
+// CHECK: memref.load %[[YIELD]][]
+// CHECK: scf.yield %[[YIELD]]
+func.func @no_merge_yielded_is_read(%init: f32, %lb: index, %ub: index, %st: index) -> f32 {
+ %alloc_yield = memref.alloca() : memref<f32>
+ %alloc_init = memref.alloca() : memref<f32>
+ memref.store %init, %alloc_init[] : memref<f32>
+ %r = scf.for %i = %lb to %ub step %st iter_args(%lv = %alloc_init) -> (memref<f32>) {
+ %v = memref.load %lv[] : memref<f32>
+ %prev = memref.load %alloc_yield[] : memref<f32>
+ %n = arith.addf %v, %prev : f32
+ memref.store %n, %alloc_yield[] : memref<f32>
+ scf.yield %alloc_yield : memref<f32>
+ }
+ %o = memref.load %r[] : memref<f32>
+ return %o : f32
+}
+
+// -----
+
+// Heap allocations may be freed or aliased elsewhere; only allocas are merged.
+
+// CHECK-LABEL: func.func @no_merge_heap_alloc
+// CHECK: %[[YIELD:.*]] = memref.alloc() : memref<f32>
+// CHECK: scf.yield %[[YIELD]]
+func.func @no_merge_heap_alloc(%init: f32, %lb: index, %ub: index, %st: index) -> f32 {
+ %alloc_yield = memref.alloc() : memref<f32>
+ %alloc_init = memref.alloc() : memref<f32>
+ memref.store %init, %alloc_init[] : memref<f32>
+ %r = scf.for %i = %lb to %ub step %st iter_args(%lv = %alloc_init) -> (memref<f32>) {
+ %v = memref.load %lv[] : memref<f32>
+ %n = arith.addf %v, %v : f32
+ memref.store %n, %alloc_yield[] : memref<f32>
+ scf.yield %alloc_yield : memref<f32>
+ }
+ %o = memref.load %r[] : memref<f32>
+ return %o : f32
+}
+
+// -----
+
+// The yielded buffer escapes the function, so its identity is observable.
+
+// CHECK-LABEL: func.func @no_merge_escaping_buffer
+// CHECK: %[[YIELD:.*]] = memref.alloca() : memref<f32>
+// CHECK: return %[[YIELD]]
+func.func @no_merge_escaping_buffer(%init: f32, %lb: index, %ub: index, %st: index) -> memref<f32> {
+ %alloc_yield = memref.alloca() : memref<f32>
+ %alloc_init = memref.alloca() : memref<f32>
+ memref.store %init, %alloc_init[] : memref<f32>
+ %r = scf.for %i = %lb to %ub step %st iter_args(%lv = %alloc_init) -> (memref<f32>) {
+ %v = memref.load %lv[] : memref<f32>
+ %n = arith.addf %v, %v : f32
+ memref.store %n, %alloc_yield[] : memref<f32>
+ scf.yield %alloc_yield : memref<f32>
+ }
+ return %alloc_yield : memref<f32>
+}
+
+// -----
+
+// The init buffer is read *after* the loop. Merging redirects the body's stores
+// into the init buffer, so that post-loop read would observe the last
+// iteration's value instead of the original init contents. Must not merge.
+
+// CHECK-LABEL: func.func @no_merge_init_read_after_loop
+// CHECK: %[[YIELD:.*]] = memref.alloca() : memref<f32>
+// CHECK: %[[INIT:.*]] = memref.alloca() : memref<f32>
+// CHECK: scf.yield %[[YIELD]]
+// CHECK: memref.load %[[INIT]][]
+func.func @no_merge_init_read_after_loop(%init: f32, %lb: index, %ub: index, %st: index) -> (f32, f32) {
+ %alloc_yield = memref.alloca() : memref<f32>
+ %alloc_init = memref.alloca() : memref<f32>
+ memref.store %init, %alloc_init[] : memref<f32>
+ %r = scf.for %i = %lb to %ub step %st iter_args(%lv = %alloc_init) -> (memref<f32>) {
+ %v = memref.load %lv[] : memref<f32>
+ %n = arith.addf %v, %v : f32
+ memref.store %n, %alloc_yield[] : memref<f32>
+ scf.yield %alloc_yield : memref<f32>
+ }
+ %x = memref.load %alloc_init[] : memref<f32>
+ %o = memref.load %r[] : memref<f32>
+ return %x, %o : f32, f32
+}
+
+// -----
+
+// The yielded buffer is written *before* the loop. Merging redirects that store
+// into the init buffer too, changing which initial value the loop reads. Must
+// not merge.
+
+// CHECK-LABEL: func.func @no_merge_yield_written_before_loop
+// CHECK: %[[YIELD:.*]] = memref.alloca() : memref<f32>
+// CHECK: %[[INIT:.*]] = memref.alloca() : memref<f32>
+// CHECK: memref.store %{{.*}}, %[[YIELD]][]
+// CHECK: scf.for {{.*}} iter_args(%{{.*}} = %[[INIT]])
+func.func @no_merge_yield_written_before_loop(%i0: f32, %i1: f32, %lb: index, %ub: index, %st: index) -> f32 {
+ %alloc_yield = memref.alloca() : memref<f32>
+ %alloc_init = memref.alloca() : memref<f32>
+ memref.store %i0, %alloc_init[] : memref<f32>
+ memref.store %i1, %alloc_yield[] : memref<f32>
+ %r = scf.for %i = %lb to %ub step %st iter_args(%lv = %alloc_init) -> (memref<f32>) {
+ %v = memref.load %lv[] : memref<f32>
+ %n = arith.addf %v, %v : f32
+ memref.store %n, %alloc_yield[] : memref<f32>
+ scf.yield %alloc_yield : memref<f32>
+ }
+ %o = memref.load %r[] : memref<f32>
+ return %o : f32
+}
+
+// -----
+
+// The yielded buffer may be read after the loop: after merging it aliases the
+// init buffer, which holds the same bytes, so the read is preserved. Merges.
+
+// CHECK-LABEL: func.func @merge_yield_read_after_loop
+// CHECK: memref.alloca() : memref<f32>
+// CHECK: %[[INIT:.*]] = memref.alloca() : memref<f32>
+// CHECK: scf.for {{.*}} iter_args(%{{.*}} = %[[INIT]])
+// CHECK: scf.yield %[[INIT]]
+// CHECK: memref.load %[[INIT]][]
+func.func @merge_yield_read_after_loop(%init: f32, %lb: index, %ub: index, %st: index) -> f32 {
+ %alloc_yield = memref.alloca() : memref<f32>
+ %alloc_init = memref.alloca() : memref<f32>
+ memref.store %init, %alloc_init[] : memref<f32>
+ %r = scf.for %i = %lb to %ub step %st iter_args(%lv = %alloc_init) -> (memref<f32>) {
+ %v = memref.load %lv[] : memref<f32>
+ %n = arith.addf %v, %v : f32
+ memref.store %n, %alloc_yield[] : memref<f32>
+ scf.yield %alloc_yield : memref<f32>
+ }
+ %o = memref.load %alloc_yield[] : memref<f32>
+ return %o : f32
+}
More information about the Mlir-commits
mailing list