[Mlir-commits] [mlir] 098b2bb - [mlir][bufferization] Fix alloc_tensor copy operand assert with size_hint (#217734)

[llvmlistbot at llvm.org]

Author: Victor Perez
Date: 2026-08-24T15:00:47Z
New Revision: 098b2bb3fa2f947bcdf86259cbf4fa0001b30934

URL: https://github.com/llvm/llvm-project/commit/098b2bb3fa2f947bcdf86259cbf4fa0001b30934
DIFF: https://github.com/llvm/llvm-project/commit/098b2bb3fa2f947bcdf86259cbf4fa0001b30934.diff

LOG: [mlir][bufferization] Fix alloc_tensor copy operand assert with size_hint (#217734)

`AllocTensorOpInterface::bufferizesToMemoryRead` and
`bufferizesToMemoryWrite` assert that the `copy` operand is the last
operand. The ODS argument order is `dynamic_sizes`, `copy`, `size_hint`,
`memory_space`.

`copy` and `size_hint` can be used together, so the input is legal and
thus OSB asserts on legal code. Compare the operand against `getCopy()`
instead.

---

Code authored by Claude Code.

Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>

Added: 
    

Modified: 
    mlir/lib/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.cpp
    mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.cpp
index 27e7ddb5d040e..71fcab29b6f69 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -34,14 +34,14 @@ struct AllocTensorOpInterface
 
   bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
-    assert(opOperand.getOperandNumber() == op->getNumOperands() - 1 &&
+    assert(&opOperand == &cast<AllocTensorOp>(op).getCopyMutable()[0] &&
            "expected copy operand");
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
-    assert(opOperand.getOperandNumber() == op->getNumOperands() - 1 &&
+    assert(&opOperand == &cast<AllocTensorOp>(op).getCopyMutable()[0] &&
            "expected copy operand");
     return false;
   }

diff  --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
index 4cb08b68fffa0..63934812bb2cf 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir
@@ -165,6 +165,21 @@ func.func @alloc_tensor_with_copy(%t: tensor<5xf32>) -> tensor<5xf32> {
 
 // -----
 
+// CHECK-LABEL: func @alloc_tensor_with_copy_and_size_hint(
+//  CHECK-SAME:     %[[t:.*]]: tensor<5xf32>)
+func.func @alloc_tensor_with_copy_and_size_hint(%t: tensor<5xf32>) -> tensor<5xf32> {
+  %sz = arith.constant 5 : index
+  // CHECK: %[[m:.*]] = bufferization.to_buffer %[[t]]
+  // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
+  // CHECK: memref.copy %[[m]], %[[alloc]]
+  %0 = bufferization.alloc_tensor() copy(%t) size_hint=%sz : tensor<5xf32>
+  // CHECK: %[[r:.*]] = bufferization.to_tensor %[[alloc]]
+  // CHECK: return %[[r]]
+  return %0 : tensor<5xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @alloc_tensor_with_memory_space()
 func.func @alloc_tensor_with_memory_space() -> tensor<5xf32> {
   // CHECK: %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32, 1>