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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Victor Perez (victor-eds)

<details>
<summary>Changes</summary>

`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.

---
Full diff: https://github.com/llvm/llvm-project/pull/217734.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.cpp (+2-2) 
- (modified) mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize.mlir (+15) 


``````````diff
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferizableOpInterfaceImpl.cpp
index 27e7ddb5d040e..1c802a1c6132d 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.get() == cast<AllocTensorOp>(op).getCopy() &&
            "expected copy operand");
     return true;
   }
 
   bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                                const AnalysisState &state) const {
-    assert(opOperand.getOperandNumber() == op->getNumOperands() - 1 &&
+    assert(opOperand.get() == cast<AllocTensorOp>(op).getCopy() &&
            "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>

``````````

</details>


https://github.com/llvm/llvm-project/pull/217734