[Mlir-commits] [mlir] 8d61e15 - [MLIR][Bufferization] Fix out-of-bounds access in setInPlaceOpOperand (#186280)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-13T12:04:49+01:00
New Revision: 8d61e156753331f26b1834cf930499e5194c64ec

URL: https://github.com/llvm/llvm-project/commit/8d61e156753331f26b1834cf930499e5194c64ec
DIFF: https://github.com/llvm/llvm-project/commit/8d61e156753331f26b1834cf930499e5194c64ec.diff

LOG: [MLIR][Bufferization] Fix out-of-bounds access in setInPlaceOpOperand (#186280)

When annotating operations with bufferization markers during analysis,
setInPlaceOpOperand reads the existing __inplace_operands_attr__ and
then sets one entry. If the attribute was provided by the user with
fewer entries than the op has operands (e.g. a return with two tensor
operands but only one entry in the annotation), the function would crash
with an out-of-bounds vector access.

Fix by resizing the vector to the actual operand count before setting
the entry when the existing annotation is too short.

Fixes #128316

Assisted-by: Claude Code

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
index 5dfe3e632b340..b57811868a725 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp
@@ -92,6 +92,11 @@ static void setInPlaceOpOperand(OpOperand &opOperand, bool inPlace) {
   if (auto attr = op->getAttr(kInPlaceOperandsAttrName)) {
     inPlaceVector = SmallVector<StringRef>(llvm::to_vector<4>(
         cast<ArrayAttr>(attr).getAsValueRange<StringAttr>()));
+    // The existing attribute may have fewer entries than the current operand
+    // count (e.g., when user-provided annotations are inconsistent with the
+    // op's actual operand count). Resize to avoid an out-of-bounds access.
+    if (inPlaceVector.size() < op->getNumOperands())
+      inPlaceVector.resize(op->getNumOperands(), "none");
   } else {
     inPlaceVector = SmallVector<StringRef>(op->getNumOperands(), "none");
     for (OpOperand &opOperand : op->getOpOperands())

diff  --git a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
index 454c17aef4d8a..4bbe75e1fe8ac 100644
--- a/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
+++ b/mlir/test/Dialect/Bufferization/Transforms/one-shot-bufferize-analysis.mlir
@@ -186,3 +186,18 @@ func.func @materialize_in_destination(%t: tensor<?xf32>, %sz: index) -> tensor<?
   %r = bufferization.materialize_in_destination %buffer in %buffer : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
   return %r : tensor<?xf32>
 }
+
+// -----
+
+// Regression test for https://github.com/llvm/llvm-project/issues/128316.
+// When an op has a user-provided __inplace_operands_attr__ with fewer entries
+// than the op's operand count, setInPlaceOpOperand must not crash.
+
+// CHECK-LABEL: func @shorter_inplace_attr
+// CHECK: return
+// CHECK-SAME: __inplace_operands_attr__ = ["false", "false"]
+func.func @shorter_inplace_attr(%arg0: tensor<4xi32>) -> (tensor<4xi32>, tensor<4xi32>) {
+  // The __inplace_operands_attr__ has only one entry but the return has two
+  // tensor operands.  The analysis must expand the annotation without crashing.
+  return {__inplace_operands_attr__ = ["false"]} %arg0, %arg0 : tensor<4xi32>, tensor<4xi32>
+}