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

[Victor Perez]

https://github.com/victor-eds created https://github.com/llvm/llvm-project/pull/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.

>From 8791e98d3865d85cd24537297a154032a99bf8bf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20P=C3=A9rez=20Carrasco?=
 <victor.pc.upm at gmail.com>
Date: Thu, 20 Aug 2026 01:43:34 -0700
Subject: [PATCH] [mlir][bufferization] Fix alloc_tensor copy operand assert
 with size_hint

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, so size_hint follows copy and
the assert fires whenever size_hint is present.

AllocTensorOp::verify permits that combination and the dialect already
round-trips size_hint, so the input is legal and one-shot bufferize aborts on
an assertions build. Compare the operand against getCopy() instead. The hooks
are only called for tensor operands and copy is the op's only tensor operand,
so the comparison is exact.
---
 .../Transforms/BufferizableOpInterfaceImpl.cpp    |  4 ++--
 .../Transforms/one-shot-bufferize.mlir            | 15 +++++++++++++++
 2 files changed, 17 insertions(+), 2 deletions(-)

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>