[Mlir-commits] [mlir] [mlir][Bufferization]Support to_tensor / to_buffer in One-Shot Bufferize analysis (PR #170261)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Dec 2 18:39:41 PST 2025
https://github.com/kimm240 updated https://github.com/llvm/llvm-project/pull/170261
>From 7ebea8565857fdeaccd3e2816221df12c944a9a6 Mon Sep 17 00:00:00 2001
From: hyun gyu kim <kimm240 at telepix.net>
Date: Tue, 2 Dec 2025 15:15:46 +0900
Subject: [PATCH 1/7] Add tests for to_tensor/to_buffer op handling issues
- Add test case for to_tensor without restrict attribute that returns a tensor
This test verifies that to_tensor ops without restrict cause an error
as expected by One-Shot Analysis (OneShotAnalysis.cpp:1208-1217)
- Add test case for to_buffer copy creation
This test verifies that to_buffer ops create a copy because they cannot
be analyzed (one-shot-module-bufferize.mlir:696-702)
These tests reproduce the to_tensor/to_memref op handling issues from PR1.
---
.../one-shot-module-bufferize-invalid.mlir | 10 ++++++++++
.../Transforms/one-shot-module-bufferize.mlir | 17 +++++++++++++++++
2 files changed, 27 insertions(+)
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir
index 29714e61d336a..06aa4589a6be5 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir
@@ -84,6 +84,16 @@ func.func @to_tensor_op_unsupported(%m: memref<?xf32>, %idx: index) -> (f32) {
// -----
+// Test case: to_tensor without restrict attribute that returns a tensor
+func.func @test_to_tensor_without_restrict(%m: memref<?xf32>) -> tensor<?xf32> {
+ // to_tensor without restrict should cause an error
+ // expected-error @+1 {{to_tensor ops without `restrict` are not supported by One-Shot Analysis}}
+ %t = bufferization.to_tensor %m : memref<?xf32> to tensor<?xf32>
+ return %t : tensor<?xf32>
+}
+
+// -----
+
func.func @yield_alloc_dominance_test_2(%cst : f32, %idx : index,
%idx2 : index) -> f32 {
%1 = bufferization.alloc_tensor(%idx) : tensor<?xf32>
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
index 8db1ebb87a1e5..6293b93b4dc35 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
@@ -710,6 +710,23 @@ func.func @to_buffer_op_unsupported(
// -----
+// Test case: to_buffer creates a copy because it cannot be analyzed
+// CHECK-LABEL: func @test_to_buffer_copy(
+// CHECK-SAME: %[[arg0:.*]]: memref<?xf32,
+func.func @test_to_buffer_copy(
+ %t: tensor<?xf32> {bufferization.writable = true}, %idx: index) -> vector<5xf32> {
+ // to_buffer cannot be analyzed, so a copy is created
+ // CHECK: %[[alloc:.*]] = memref.alloc
+ // CHECK: memref.copy %[[arg0]], %[[alloc]]
+ %m = bufferization.to_buffer %t : tensor<?xf32> to memref<?xf32>
+ %c0 = arith.constant 0 : index
+ %cst = arith.constant 0.0 : f32
+ %r = vector.transfer_read %t[%c0], %cst : tensor<?xf32>, vector<5xf32>
+ return %r : vector<5xf32>
+}
+
+// -----
+
// Note: The cf.br canonicalizes away, so there's nothing to check here. There
// is a detailed test in ControlFlow/bufferize.mlir.
>From a6dab52884549404d73c9bac2c8fe343e439a6c8 Mon Sep 17 00:00:00 2001
From: hyun gyu kim <kimm240 at telepix.net>
Date: Tue, 2 Dec 2025 16:00:39 +0900
Subject: [PATCH 2/7] Enable to_tensor/to_buffer alias analysis in one-shot
bufferize
- Implement BufferizableOpInterface behavior for ToTensorOp/ToBufferOp in TableGen
(aliasing, read/write semantics)
- Set up explicit alias unions for to_tensor/to_buffer in OneShotAnalysis
- Remove restriction that rejected to_tensor without restrict and update tests
to cover the new behavior
---
.../Bufferization/IR/BufferizationOps.td | 21 +++++++++++++-
.../Transforms/OneShotAnalysis.cpp | 28 +++++++++++++------
.../one-shot-module-bufferize-invalid.mlir | 9 +++---
.../Transforms/one-shot-module-bufferize.mlir | 18 ++++++++++++
4 files changed, 62 insertions(+), 14 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index a9b2b9f39519d..7ab2354ca4923 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -478,6 +478,24 @@ def Bufferization_ToTensorOp : Bufferization_Op<"to_tensor", [
bool isWritable(Value value, const AnalysisState &state);
+ AliasingValueList getAliasingValues(OpOperand &opOperand,
+ const AnalysisState &state) const {
+ // The output tensor aliases with the input memref.
+ return {{const_cast<ToTensorOp *>(this)->getResult(), BufferRelation::Equivalent}};
+ }
+
+ bool bufferizesToMemoryRead(OpOperand &opOperand,
+ const AnalysisState &state) const {
+ // to_tensor reads from the memref.
+ return true;
+ }
+
+ bool bufferizesToMemoryWrite(OpOperand &opOperand,
+ const AnalysisState &state) const {
+ // to_tensor does not write to the memref.
+ return false;
+ }
+
FailureOr<BufferLikeType> getBufferType(
Value value, const BufferizationOptions &options,
const BufferizationState &state, SmallVector<Value> &invocationStack) {
@@ -551,7 +569,8 @@ def Bufferization_ToBufferOp : Bufferization_Op<"to_buffer", [
AliasingValueList getAliasingValues(
OpOperand &opOperand, const AnalysisState &state) const {
- return {};
+ // The output memref aliases with the input tensor.
+ return {{const_cast<ToBufferOp *>(this)->getResult(), BufferRelation::Equivalent}};
}
LogicalResult bufferize(RewriterBase &rewriter,
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
index 5dfe3e632b340..2aa1a87d0b027 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
@@ -1205,15 +1205,27 @@ checkPreBufferizationAssumptions(Operation *op, const DominanceInfo &domInfo,
if (!options.isOpAllowed(op.getOperation()))
return WalkResult::advance();
- // Input IR may not contain any ToTensorOps without the "restrict"
- // attribute. Such tensors may alias any other tensor, which is currently
- // not handled in the analysis.
+ // Set up alias relationships for to_tensor and to_buffer ops.
if (auto toTensorOp = dyn_cast<ToTensorOp>(op.getOperation())) {
- if (!toTensorOp.getRestrict() && !toTensorOp->getUses().empty()) {
- op->emitOpError("to_tensor ops without `restrict` are not supported by "
- "One-Shot Analysis");
- return WalkResult::interrupt();
- }
+ Value inputMemRef = toTensorOp.getBuffer();
+ Value outputTensor = toTensorOp.getResult();
+ // Union alias sets: the output tensor aliases with the input memref.
+ state.unionAliasSets(outputTensor, inputMemRef);
+ // Handle recursive aliasing: if input memref already aliases with other
+ // tensors, propagate those aliases to the output tensor.
+ state.applyOnAliases(inputMemRef, [&](Value alias) {
+ state.unionAliasSets(outputTensor, alias);
+ });
+ } else if (auto toBufferOp = dyn_cast<ToBufferOp>(op.getOperation())) {
+ Value inputTensor = toBufferOp.getTensor();
+ Value outputMemRef = toBufferOp.getResult();
+ // Union alias sets: the output memref aliases with the input tensor.
+ state.unionAliasSets(outputMemRef, inputTensor);
+ // Handle recursive aliasing: if input tensor already aliases with other
+ // memrefs, propagate those aliases to the output memref.
+ state.applyOnAliases(inputTensor, [&](Value alias) {
+ state.unionAliasSets(outputMemRef, alias);
+ });
}
for (OpOperand &opOperand : op->getOpOperands()) {
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir
index 06aa4589a6be5..b9b41fc6f75c8 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir
@@ -74,10 +74,10 @@ func.func @scf_while_non_equiv_yield(%arg0: tensor<5xi1>,
// -----
-func.func @to_tensor_op_unsupported(%m: memref<?xf32>, %idx: index) -> (f32) {
- // expected-error @+1 {{to_tensor ops without `restrict` are not supported by One-Shot Analysis}}
+// Note: to_tensor without restrict is now supported with proper alias analysis.
+// This test verifies that to_tensor ops work correctly without restrict.
+func.func @to_tensor_without_restrict_works(%m: memref<?xf32>, %idx: index) -> (f32) {
%0 = bufferization.to_tensor %m : memref<?xf32> to tensor<?xf32>
-
%1 = tensor.extract %0[%idx] : tensor<?xf32>
return %1 : f32
}
@@ -85,9 +85,8 @@ func.func @to_tensor_op_unsupported(%m: memref<?xf32>, %idx: index) -> (f32) {
// -----
// Test case: to_tensor without restrict attribute that returns a tensor
+// This should now work correctly with alias analysis.
func.func @test_to_tensor_without_restrict(%m: memref<?xf32>) -> tensor<?xf32> {
- // to_tensor without restrict should cause an error
- // expected-error @+1 {{to_tensor ops without `restrict` are not supported by One-Shot Analysis}}
%t = bufferization.to_tensor %m : memref<?xf32> to tensor<?xf32>
return %t : tensor<?xf32>
}
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
index 6293b93b4dc35..53b49bb27826b 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
@@ -727,6 +727,24 @@ func.func @test_to_buffer_copy(
// -----
+// Test case: to_tensor without restrict works correctly
+// Verify that to_tensor without restrict can be used without errors
+// CHECK-LABEL: func @test_to_tensor_without_restrict_works(
+// CHECK-SAME: %[[arg0:.*]]: memref<?xf32,
+func.func @test_to_tensor_without_restrict_works(
+ %m: memref<?xf32>, %idx: index) -> f32 {
+ // to_tensor without restrict should work with alias analysis
+ %t = bufferization.to_tensor %m : memref<?xf32> to tensor<?xf32>
+ // Read from tensor - should read from the original memref due to alias
+ %result = tensor.extract %t[%idx] : tensor<?xf32>
+ // CHECK: memref.load %[[arg0]]
+ // CHECK-NOT: memref.alloc
+ // CHECK-NOT: memref.copy
+ return %result : f32
+}
+
+// -----
+
// Note: The cf.br canonicalizes away, so there's nothing to check here. There
// is a detailed test in ControlFlow/bufferize.mlir.
>From 54cdc46ac858fccbd44e7b70106d4f20c9104570 Mon Sep 17 00:00:00 2001
From: hyun gyu kim <kimm240 at telepix.net>
Date: Tue, 2 Dec 2025 16:30:51 +0900
Subject: [PATCH 3/7] Cover forum to_tensor/to_buffer example in tests
Add a one-shot bufferization test that mirrors the forum IR example
with to_buffer (to_memref), an external memref function call, and a
subsequent to_tensor use. This ensures that our alias handling and
crash behavior can be validated against the real-world scenario that
motivated this change.
Also document the use of const_cast in ToTensorOp/ToBufferOp
BufferizableOpInterface implementations so that the rationale is clear
in the generated C++ code.
---
.../Bufferization/IR/BufferizationOps.td | 14 +++++++--
.../Transforms/one-shot-module-bufferize.mlir | 30 +++++++++++++++++++
2 files changed, 42 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index 7ab2354ca4923..e1f8c267a9d19 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -481,7 +481,12 @@ def Bufferization_ToTensorOp : Bufferization_Op<"to_tensor", [
AliasingValueList getAliasingValues(OpOperand &opOperand,
const AnalysisState &state) const {
// The output tensor aliases with the input memref.
- return {{const_cast<ToTensorOp *>(this)->getResult(), BufferRelation::Equivalent}};
+ //
+ // Note: getResult() is a non-const helper. Use const_cast here to avoid
+ // duplicating accessors while still satisfying the const interface
+ // contract of BufferizableOpInterface.
+ return {{const_cast<ToTensorOp *>(this)->getResult(),
+ BufferRelation::Equivalent}};
}
bool bufferizesToMemoryRead(OpOperand &opOperand,
@@ -570,7 +575,12 @@ def Bufferization_ToBufferOp : Bufferization_Op<"to_buffer", [
AliasingValueList getAliasingValues(
OpOperand &opOperand, const AnalysisState &state) const {
// The output memref aliases with the input tensor.
- return {{const_cast<ToBufferOp *>(this)->getResult(), BufferRelation::Equivalent}};
+ //
+ // Note: getResult() is a non-const helper. Use const_cast here to avoid
+ // duplicating accessors while still satisfying the const interface
+ // contract of BufferizableOpInterface.
+ return {{const_cast<ToBufferOp *>(this)->getResult(),
+ BufferRelation::Equivalent}};
}
LogicalResult bufferize(RewriterBase &rewriter,
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
index 53b49bb27826b..95f7662ab78f1 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
@@ -745,6 +745,36 @@ func.func @test_to_tensor_without_restrict_works(
// -----
+// Forum example: to_buffer (to_memref) + external memref call + to_tensor
+// This mirrors the IR discussed in the MLIR forum thread:
+// "Bufferization fails in the presence of `bufferization.to_memref`
+// and `bufferization.to_tensor`".
+//
+// CHECK-LABEL: func @forum_to_tensor_to_buffer_example(
+// CHECK-SAME: %[[ARG0:.*]]: memref<2xf32
+func.func @forum_to_tensor_to_buffer_example(%arg0: tensor<2xf32>) {
+ %m = bufferization.to_buffer %arg0 : tensor<2xf32> to memref<2xf32>
+ func.call @some_func_operating_on_memref(%m)
+ : (memref<2xf32>) -> ()
+ %t = bufferization.to_tensor %m : memref<2xf32> to tensor<2xf32>
+ %c0 = arith.constant 0 : index
+ %v = tensor.extract %t[%c0] : tensor<2xf32>
+ vector.print %v : f32
+ return
+}
+
+// The external memref function should remain a memref-based function after
+// bufferization and there should be no remaining to_tensor/to_buffer ops.
+// CHECK: func.func private @some_func_operating_on_memref(
+// CHECK-SAME: memref<2xf32>
+// CHECK-NOT: bufferization.to_tensor
+// CHECK-NOT: bufferization.to_buffer
+func.func private @some_func_operating_on_memref(%m: memref<2xf32>) -> () {
+ return
+}
+
+// -----
+
// Note: The cf.br canonicalizes away, so there's nothing to check here. There
// is a detailed test in ControlFlow/bufferize.mlir.
>From d5431cbc8f5c6ea464e7309ac56fc1d02e0f5154 Mon Sep 17 00:00:00 2001
From: hyun gyu kim <kimm240 at telepix.net>
Date: Tue, 2 Dec 2025 17:15:07 +0900
Subject: [PATCH 4/7] Restrict alias analysis helpers to tensor values for
ToBufferOp
Teach AnalysisState::getAliasingOpOperands to early-return for
non-tensor values so that tensor-only alias analysis does not try to
follow aliases through memref results such as ToBufferOp. This keeps
non-tensor aliasing (e.g., between ToBufferOp results and their tensor
operands) in OneShotAnalysis where alias sets are explicitly unioned.
Also update ToBufferOp's BufferizableOpInterface implementation so that
neither its operand nor its result are treated as implicit memory
writes; writes are modeled solely on downstream memref users.
---
.../Bufferization/IR/BufferizationOps.td | 27 +++++++++++++------
.../IR/BufferizableOpInterface.cpp | 6 +++++
2 files changed, 25 insertions(+), 8 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index e1f8c267a9d19..f146097f556b0 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -569,18 +569,29 @@ def Bufferization_ToBufferOp : Bufferization_Op<"to_buffer", [
bool bufferizesToMemoryWrite(OpOperand &opOperand,
const AnalysisState &state) {
- return !getReadOnly();
+ // The returned memref may be written through later memref ops, but
+ // ToBufferOp itself does not perform a write. The write semantics are
+ // modeled on the users of the memref result, not on this op.
+ return false;
+ }
+
+ bool resultBufferizesToMemoryWrite(OpResult opResult,
+ const AnalysisState &state) {
+ // The ToBufferOp result is a memref view of an existing tensor buffer.
+ // The op itself does not perform a memory write; any writes are modeled
+ // on the users of the memref. Override the default implementation to
+ // avoid treating the result as an implicit write in tensor-only analysis.
+ return false;
}
AliasingValueList getAliasingValues(
OpOperand &opOperand, const AnalysisState &state) const {
- // The output memref aliases with the input tensor.
- //
- // Note: getResult() is a non-const helper. Use const_cast here to avoid
- // duplicating accessors while still satisfying the const interface
- // contract of BufferizableOpInterface.
- return {{const_cast<ToBufferOp *>(this)->getResult(),
- BufferRelation::Equivalent}};
+ // ToBufferOp does not produce a tensor result, so from the perspective
+ // of tensor-based BufferizableOpInterface analysis there is no tensor
+ // value that aliases with its tensor operand. The aliasing relationship
+ // between the tensor operand and the memref result is modeled separately
+ // in OneShotAnalysis via explicit alias set unions.
+ return {};
}
LogicalResult bufferize(RewriterBase &rewriter,
diff --git a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
index 9b11270e7bbe2..2a833c5d7971d 100644
--- a/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
+++ b/mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp
@@ -434,6 +434,12 @@ static void setInsertionPointAfter(OpBuilder &b, Value value) {
/// Determine which OpOperand* will alias with `value` if the op is bufferized
/// in place. Return all tensor OpOperand* if the op is not bufferizable.
AliasingOpOperandList AnalysisState::getAliasingOpOperands(Value value) const {
+ // This helper is intended for tensor values. Non-tensor alias relationships
+ // (e.g., between ToBufferOp results and their tensor operands) are modeled
+ // separately in OneShotAnalysis via explicit alias set unions.
+ if (!llvm::isa<TensorType>(value.getType()))
+ return {};
+
if (Operation *op = getOwnerOfValue(value))
if (auto bufferizableOp = getOptions().dynCastBufferizableOp(op))
return bufferizableOp.getAliasingOpOperands(value, *this);
>From a856885762bb96d7bb6a0a0252b9d6aaddd45fdd Mon Sep 17 00:00:00 2001
From: hyun gyu kim <kimm240 at telepix.net>
Date: Tue, 2 Dec 2025 17:17:53 +0900
Subject: [PATCH 5/7] Clarify to_buffer/to_tensor external-call test
description
Update the comment around the to_buffer + external memref call +
to_tensor test to describe the scenario in generic terms instead of
referencing the original forum thread. The test now simply documents
that this pattern must not crash the one-shot bufferization pipeline
and keeps the IR example self-contained.
---
.../Transforms/one-shot-module-bufferize.mlir | 7 +++----
1 file changed, 3 insertions(+), 4 deletions(-)
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
index 95f7662ab78f1..894acf94acbe0 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
@@ -745,10 +745,9 @@ func.func @test_to_tensor_without_restrict_works(
// -----
-// Forum example: to_buffer (to_memref) + external memref call + to_tensor
-// This mirrors the IR discussed in the MLIR forum thread:
-// "Bufferization fails in the presence of `bufferization.to_memref`
-// and `bufferization.to_tensor`".
+// Example: to_buffer (to_memref) + external memref call + to_tensor
+// Ensure that to_buffer/to_tensor around an external memref call do not crash
+// the one-shot bufferization pipeline.
//
// CHECK-LABEL: func @forum_to_tensor_to_buffer_example(
// CHECK-SAME: %[[ARG0:.*]]: memref<2xf32
>From 54a9d029994a07384857ad6edc5fd40098949375 Mon Sep 17 00:00:00 2001
From: hyun gyu kim <kimm240 at telepix.net>
Date: Wed, 3 Dec 2025 11:32:02 +0900
Subject: [PATCH 6/7] Fix ToBufferOp to correctly handle memory writes for copy
insertion
This fixes the copy omission issue identified by the reviewer where
memref.copy operations were missing after bufferization.to_buffer,
leading to incorrect values being loaded and violating tensor SSA
immutability.
Changes:
- ToBufferOp::bufferizesToMemoryWrite now returns !getReadOnly() instead
of false, ensuring that when read_only is not set, the operation is
treated as a memory write, which triggers proper copy insertion.
- Added ToBufferOp::resultBufferizesToMemoryWrite override to handle
memref results correctly, avoiding assertion failures when the default
implementation tries to call getAliasingOpOperands on memref types.
- ToBufferOp::getAliasingValues now returns the input tensor to enable
RaW conflict detection when the result memref is written to and the
input tensor is read from.
- Fixed test CHECK pattern to match actual output format.
This ensures that when a new tensor (e.g., from linalg.fill) is
converted to a memref via to_buffer, the memref does not share memory
with the original tensor buffer, maintaining tensor immutability.
---
.../Bufferization/IR/BufferizationOps.td | 25 +++++++++----------
.../Transforms/one-shot-module-bufferize.mlir | 2 +-
2 files changed, 13 insertions(+), 14 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index f146097f556b0..236b59af23249 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -569,29 +569,28 @@ def Bufferization_ToBufferOp : Bufferization_Op<"to_buffer", [
bool bufferizesToMemoryWrite(OpOperand &opOperand,
const AnalysisState &state) {
- // The returned memref may be written through later memref ops, but
- // ToBufferOp itself does not perform a write. The write semantics are
- // modeled on the users of the memref result, not on this op.
- return false;
+ return !getReadOnly();
}
bool resultBufferizesToMemoryWrite(OpResult opResult,
const AnalysisState &state) {
// The ToBufferOp result is a memref view of an existing tensor buffer.
// The op itself does not perform a memory write; any writes are modeled
- // on the users of the memref. Override the default implementation to
- // avoid treating the result as an implicit write in tensor-only analysis.
- return false;
+ // on the users of the memref. However, we need to indicate that the result
+ // may be written to if read_only is not set, so that proper copy insertion
+ // happens. Since the result is a memref (not a tensor), we cannot use
+ // getAliasingOpOperands which expects tensor types.
+ return !getReadOnly();
}
AliasingValueList getAliasingValues(
OpOperand &opOperand, const AnalysisState &state) const {
- // ToBufferOp does not produce a tensor result, so from the perspective
- // of tensor-based BufferizableOpInterface analysis there is no tensor
- // value that aliases with its tensor operand. The aliasing relationship
- // between the tensor operand and the memref result is modeled separately
- // in OneShotAnalysis via explicit alias set unions.
- return {};
+ // The result memref aliases with the input tensor. However, getAliasingValues
+ // only returns tensor Values, so we cannot return the memref result here.
+ // The alias relationship is established in OneShotAnalysis via unionAliasSets.
+ // Returning the input tensor itself ensures that RaW conflicts are detected
+ // when the result memref is written to and the input tensor is read from.
+ return {{opOperand.get(), BufferRelation::Equivalent}};
}
LogicalResult bufferize(RewriterBase &rewriter,
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
index 894acf94acbe0..7de3de488767d 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
@@ -730,7 +730,7 @@ func.func @test_to_buffer_copy(
// Test case: to_tensor without restrict works correctly
// Verify that to_tensor without restrict can be used without errors
// CHECK-LABEL: func @test_to_tensor_without_restrict_works(
-// CHECK-SAME: %[[arg0:.*]]: memref<?xf32,
+// CHECK-SAME: %[[arg0:.*]]: memref<?xf32>, %[[arg1:.*]]: index
func.func @test_to_tensor_without_restrict_works(
%m: memref<?xf32>, %idx: index) -> f32 {
// to_tensor without restrict should work with alias analysis
>From 1996ef58dc4570560fa2cebb1ae005a8a877e899 Mon Sep 17 00:00:00 2001
From: hyun gyu kim <kimm240 at telepix.net>
Date: Wed, 3 Dec 2025 11:38:13 +0900
Subject: [PATCH 7/7] Add test case for linalg.fill + to_buffer copy insertion
This adds a test case based on the reviewer's specific IR example that
demonstrates the fix for copy omission. The test verifies that when a
new tensor (from linalg.fill) is converted to a memref via to_buffer,
and then the memref is written to while the original tensor is read
from, a copy is correctly inserted to maintain tensor immutability.
The test case:
- Creates a tensor with tensor.empty()
- Fills it with linalg.fill
- Converts to memref with to_buffer
- Writes to the memref
- Reads from the original tensor
This ensures memref.copy is inserted between the fill and the store,
maintaining the immutability of the original tensor.
---
.../Transforms/one-shot-module-bufferize.mlir | 28 +++++++++++++++++++
1 file changed, 28 insertions(+)
diff --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
index 7de3de488767d..4b517a145fa32 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir
@@ -774,6 +774,34 @@ func.func private @some_func_operating_on_memref(%m: memref<2xf32>) -> () {
// -----
+// Test case: to_buffer after linalg.fill requires copy insertion
+// This test verifies that when a new tensor (from linalg.fill) is converted
+// to a memref via to_buffer, and then the memref is written to while the
+// original tensor is read from, a copy is inserted to maintain tensor
+// immutability.
+// CHECK-LABEL: func @test_fill_to_buffer_requires_copy(
+func.func @test_fill_to_buffer_requires_copy() {
+ %0 = tensor.empty() : tensor<10xf32>
+ %cst = arith.constant 5.0 : f32
+ %cst2 = arith.constant 6.0 : f32
+ %c0 = arith.constant 0 : index
+ %2 = linalg.fill ins(%cst : f32) outs(%0 : tensor<10xf32>) -> tensor<10xf32>
+ %3 = bufferization.to_buffer %2 : tensor<10xf32> to memref<10xf32>
+ memref.store %cst2, %3[%c0] : memref<10xf32>
+ %4 = tensor.extract %2 [%c0] : tensor<10xf32>
+ vector.print %4 : f32
+ return
+}
+// CHECK: %[[alloc:.*]] = memref.alloc
+// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%[[alloc]]
+// CHECK: %[[alloc_copy:.*]] = memref.alloc
+// CHECK: memref.copy %[[alloc]], %[[alloc_copy]]
+// CHECK: memref.store %{{.*}}, %[[alloc_copy]]
+// CHECK: memref.load %[[alloc]]
+// CHECK: vector.print
+
+// -----
+
// Note: The cf.br canonicalizes away, so there's nothing to check here. There
// is a detailed test in ControlFlow/bufferize.mlir.
More information about the Mlir-commits
mailing list