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

[Victor Perez]

https://github.com/victor-eds updated https://github.com/llvm/llvm-project/pull/217734

>From 156f82867d6b3a1999264ed5a2e23fd437004e08 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
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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 address against the copy operand
instead. `copy` is optional, so `getCopyMutable()` returns a
`MutableOperandRange`. Element 0 is the copy operand, and the range checks the
index, so an absent copy also trips an assert.

Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
 .../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..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>