[Mlir-commits] [mlir] [mlir][bufferization] Fix to_buffer being incorrectly hoisted by LICM (PR #188997)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/188997

bufferization.to_buffer was marked as Pure, which caused LICM to hoist it out of loops even when the returned memref is later written to. This is incorrect: when multiple loop iterations share the same backing buffer, writes from one iteration corrupt another's view of the tensor's storage.

Fix by replacing Pure with AlwaysSpeculatable (to preserve speculation behavior) and a custom getEffects that declares a MemoryEffects::Write effect on the result when read_only is not set. With the write effect, LICM correctly treats to_buffer as having observable side effects and refrains from hoisting it.

When read_only is set, getEffects returns no effects, preserving the ability to hoist and CSE such ops.

Fixes #156032

Assisted-by: Claude Code

>From 43ee6b20464054733a0269f8fc347c0bd6eb685f Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Fri, 27 Mar 2026 04:30:12 -0700
Subject: [PATCH] [mlir][bufferization] Fix to_buffer being incorrectly hoisted
 by LICM

bufferization.to_buffer was marked as Pure, which caused LICM to hoist
it out of loops even when the returned memref is later written to. This
is incorrect: when multiple loop iterations share the same backing
buffer, writes from one iteration corrupt another's view of the tensor's
storage.

Fix by replacing Pure with AlwaysSpeculatable (to preserve speculation
behavior) and a custom getEffects that declares a MemoryEffects::Write
effect on the result when read_only is not set. With the write effect,
LICM correctly treats to_buffer as having observable side effects and
refrains from hoisting it.

When read_only is set, getEffects returns no effects, preserving the
ability to hoist and CSE such ops.

Fixes #156032

Assisted-by: Claude Code
---
 .../Bufferization/IR/BufferizationOps.td      |  9 ++-
 .../Bufferization/IR/BufferizationOps.cpp     | 18 ++++++
 .../Dialect/Bufferization/side-effects.mlir   | 14 +++++
 .../loop-invariant-code-motion.mlir           | 58 +++++++++++++++++++
 4 files changed, 98 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index a9b2b9f39519d..9a02809f48661 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -503,7 +503,8 @@ def Bufferization_ToBufferOp : Bufferization_Op<"to_buffer", [
     BufferizableOpInterface,
     SameOperandsAndResultShape,
     SameOperandsAndResultElementType,
-    Pure,
+    AlwaysSpeculatable,
+    DeclareOpInterfaceMethods<MemoryEffectsOpInterface, ["getEffects"]>,
     Bufferization_TensorAndBufferMatch<"tensor", "buffer">
   ]> {
   let summary = "cast a tensor-like type to buffer-like type";
@@ -522,6 +523,12 @@ def Bufferization_ToBufferOp : Bufferization_Op<"to_buffer", [
     The `read_only` attribute can optionally be set, indicating to the
     bufferization that the buffer returned by this op (or an alias created from
     the returned buffer) will not be written to.
+
+    Memory effects: when `read_only` is not set, this op declares a write
+    effect on its result to prevent passes such as LICM from hoisting it out of
+    loops. (Hoisting would cause all loop iterations to share the same backing
+    buffer.) When `read_only` is set, no write effect is declared and the op
+    may be freely hoisted or CSE'd.
   }];
 
   let arguments = (ins Bufferization_TensorLikeTypeInterface:$tensor, UnitAttr:$read_only);
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
index 71ad8bbf91c3b..3b4a0a5c07994 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp
@@ -882,6 +882,24 @@ void ToBufferOp::getCanonicalizationPatterns(RewritePatternSet &results,
               ToBufferToTensorFolding>(context);
 }
 
+void ToBufferOp::getEffects(
+    SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
+        &effects) {
+  // Declare a write effect on the result memref when read_only is not set.
+  // This prevents passes such as LICM from hoisting to_buffer ops out of
+  // loops: if to_buffer were hoisted, all iterations would share the same
+  // backing buffer, so writes from one iteration would corrupt another's view
+  // of the tensor's storage.
+  //
+  // When read_only is set, no writes can occur through the returned buffer so
+  // sharing across loop iterations is safe and no effects are declared, making
+  // the op eligible for hoisting/CSE.
+  if (!getReadOnly())
+    effects.emplace_back(MemoryEffects::Write::get(),
+                         mlir::cast<OpResult>(getBuffer()),
+                         SideEffects::DefaultResource::get());
+}
+
 LogicalResult ToBufferOp::bufferize(RewriterBase &rewriter,
                                     const BufferizationOptions &options,
                                     BufferizationState &state) {
diff --git a/mlir/test/Dialect/Bufferization/side-effects.mlir b/mlir/test/Dialect/Bufferization/side-effects.mlir
index 129fc8b32c270..dda0d72ae1823 100644
--- a/mlir/test/Dialect/Bufferization/side-effects.mlir
+++ b/mlir/test/Dialect/Bufferization/side-effects.mlir
@@ -7,3 +7,17 @@ func.func @test_side_effects(%arg0: memref<2xi32>) -> memref<2xi32> {
   %0 = bufferization.clone %arg0 : memref<2xi32> to memref<2xi32>
   return %0 : memref<2xi32>
 }
+
+// to_buffer without read_only has a write effect on the result (prevents LICM).
+func.func @test_to_buffer_write_effect(%arg0: tensor<2xi32>) -> memref<2xi32> {
+  // expected-remark @below {{found an instance of 'write' on op result 0, on resource '<Default>'}}
+  %0 = bufferization.to_buffer %arg0 : tensor<2xi32> to memref<2xi32>
+  return %0 : memref<2xi32>
+}
+
+// to_buffer with read_only has no side effects.
+func.func @test_to_buffer_readonly_no_effect(%arg0: tensor<2xi32>) -> memref<2xi32> {
+  // expected-remark @below {{operation has no memory effects}}
+  %0 = bufferization.to_buffer %arg0 read_only : tensor<2xi32> to memref<2xi32>
+  return %0 : memref<2xi32>
+}
diff --git a/mlir/test/Transforms/loop-invariant-code-motion.mlir b/mlir/test/Transforms/loop-invariant-code-motion.mlir
index 31a4f64dd7de0..76b7c8c088a56 100644
--- a/mlir/test/Transforms/loop-invariant-code-motion.mlir
+++ b/mlir/test/Transforms/loop-invariant-code-motion.mlir
@@ -1582,3 +1582,61 @@ func.func @do_not_hoist_vector_transfer_ops_memref(
   }
   func.return %final : vector<4x4xf32>
 }
+
+// -----
+
+// bufferization.to_buffer (without read_only) should not be hoisted out of
+// loops, as it may be used to write to the tensor's buffer.
+// https://github.com/llvm/llvm-project/issues/156032
+
+// CHECK-LABEL: func @do_not_hoist_to_buffer
+// CHECK-NOT: bufferization.to_buffer
+// CHECK: scf.for
+// CHECK:   bufferization.to_buffer
+// CHECK:   linalg.fill
+func.func @do_not_hoist_to_buffer(%arg0: tensor<32xi32>, %lb: index, %ub: index, %step: index) {
+  scf.for %i = %lb to %ub step %step {
+    %buf = bufferization.to_buffer %arg0 : tensor<32xi32> to memref<32xi32>
+    %c = arith.index_cast %i : index to i32
+    linalg.fill ins(%c : i32) outs(%buf : memref<32xi32>)
+  }
+  return
+}
+
+// -----
+
+// bufferization.to_buffer with read_only attribute has no write effects and
+// may be hoisted.
+
+// CHECK-LABEL: func @hoist_to_buffer_readonly
+// CHECK: bufferization.to_buffer{{.*}}read_only
+// CHECK: scf.for
+func.func @hoist_to_buffer_readonly(%arg0: tensor<32xi32>, %lb: index, %ub: index, %step: index) {
+  scf.for %i = %lb to %ub step %step {
+    %buf = bufferization.to_buffer %arg0 read_only : tensor<32xi32> to memref<32xi32>
+    func.call @use(%buf) : (memref<32xi32>) -> ()
+  }
+  return
+}
+func.func private @use(memref<32xi32>)
+
+// -----
+
+// to_buffer (without read_only) should not be hoisted out of scf.forall either.
+// This is the original bug pattern from #156032: tensor.empty + to_buffer
+// inside a parallel loop.
+
+// CHECK-LABEL: func @do_not_hoist_to_buffer_forall
+// CHECK-NOT: bufferization.to_buffer
+// CHECK: scf.forall
+// CHECK:   bufferization.to_buffer
+// CHECK:   linalg.fill
+func.func @do_not_hoist_to_buffer_forall(%n: index) {
+  scf.forall (%i) in (%n) {
+    %t = tensor.empty() : tensor<32xi32>
+    %buf = bufferization.to_buffer %t : tensor<32xi32> to memref<32xi32>
+    %c = arith.index_cast %i : index to i32
+    linalg.fill ins(%c : i32) outs(%buf : memref<32xi32>)
+  }
+  return
+}