[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:40:14 PST 2025
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: None (kimm240)
<details>
<summary>Changes</summary>
Summary
-------
This patch teaches One-Shot Bufferize how to reason about
`bufferization.to_tensor` and `bufferization.to_buffer` without crashing,
and relaxes the previous restriction that rejected `to_tensor` ops without
`restrict`.
The main ideas are:
- Treat `ToTensorOp` / `ToBufferOp` as proper BufferizableOpInterface ops
with well-defined aliasing and read/write semantics.
- Keep tensor-only alias analysis strictly on tensor values.
- Model memref ↔ tensor alias sets for `to_tensor` / `to_buffer` explicitly
in OneShotAnalysis, instead of trying to push memrefs through the generic
tensor alias helpers.
Changes
-------
1) BufferizableOpInterface semantics for ToTensorOp / ToBufferOp
- `ToTensorOp`:
- Implements `getAliasingValues`, `bufferizesToMemoryRead`,
`bufferizesToMemoryWrite` in TableGen.
- The tensor result is modeled as aliasing the memref operand.
- Reads from the memref but does not write.
- `ToBufferOp`:
- Implements `bufferizesToMemoryRead`, `bufferizesToMemoryWrite`,
and `resultBufferizesToMemoryWrite` in a conservative way:
- The op itself does not perform a write; writes are modeled solely
on downstream memref users.
- `getAliasingValues` returns an empty list from the tensor-analysis
point of view; the memref result is not treated as a tensor alias.
2) OneShotAnalysis alias handling for to_tensor / to_buffer
- In `checkPreBufferizationAssumptions`:
- Add explicit alias-set unions for `ToTensorOp` and `ToBufferOp`:
- `to_tensor`: union the tensor result with the memref operand,
and propagate existing aliases of the memref.
- `to_buffer`: union the memref result with the tensor operand,
and propagate existing aliases of the tensor.
- This keeps memref/tensor aliasing localized to OneShotAnalysis where
alias sets are explicitly managed, instead of going through the
tensor-only BufferizableOpInterface helpers.
3) Restrict alias helpers to tensor values
- `AnalysisState::getAliasingOpOperands(Value)`:
- Now early-returns an empty list for non-tensor values.
- This prevents tensor-only alias analysis from trying to follow
aliases starting at memref results (e.g., `ToBufferOp`), which
previously led to assertions in the one-shot pipeline.
Tests
-----
- `mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir`:
- Replace the previous "to_tensor without `restrict` must error" tests
with positive tests that verify `to_tensor` without `restrict`
now bufferizes successfully.
- `mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir`:
- `test_to_tensor_without_restrict_works`:
- Checks that `to_tensor` without `restrict` is lowered to a plain
`memref.load` without extra alloc/copy.
- `forum_to_tensor_to_buffer_example` (renamed/commented generically):
- `tensor<2xf32> -> to_buffer -> external memref call ->
to_tensor -> tensor use`
- Ensures this pattern no longer crashes the one-shot bufferization
pipeline and that the external memref function remains memref-based.
Notes
-----
- This patch focuses on making One-Shot Bufferize robust and correct
for the common `to_tensor` / `to_buffer` patterns (including the
external-call example), while keeping the analysis conservative.
- More aggressive alias reasoning through `to_buffer` could be added
in follow-up work if needed, but is intentionally out of scope here.
Background / Reference
----------------------
- This work was motivated by the crash scenario discussed on the MLIR
discourse thread:
[One-Shot Bufferizer: Bufferization fails in the presence of
`bufferization.to_memref` and `bufferization.to_tensor`][1]
[1]: https://discourse.llvm.org/t/one-shot-bufferizer-bufferization-fails-in-the-presence-of-bufferization-to-memref-and-bufferization-to-tensor/62211
---
Full diff: https://github.com/llvm/llvm-project/pull/170261.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td (+40-1)
- (modified) mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp (+6)
- (modified) mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp (+20-8)
- (modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-invalid.mlir (+12-3)
- (modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize.mlir (+92)
``````````diff
diff --git a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
index a9b2b9f39519d..236b59af23249 100644
--- a/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
+++ b/mlir/include/mlir/Dialect/Bufferization/IR/BufferizationOps.td
@@ -478,6 +478,29 @@ 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.
+ //
+ // 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,
+ 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) {
@@ -549,9 +572,25 @@ def Bufferization_ToBufferOp : Bufferization_Op<"to_buffer", [
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. 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 {
- 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/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);
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 29714e61d336a..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,16 +74,25 @@ 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
}
// -----
+// 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> {
+ %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..4b517a145fa32 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,98 @@ 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>
+}
+
+// -----
+
+// 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>, %[[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
+ %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
+}
+
+// -----
+
+// 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
+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
+}
+
+// -----
+
+// 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.
``````````
</details>
https://github.com/llvm/llvm-project/pull/170261
More information about the Mlir-commits
mailing list