[Mlir-commits] [mlir] [mlir][bufferization] Fix crash when returning an allocation twice (PR #220232)

[llvmlistbot at llvm.org]

https://github.com/AkunCube updated https://github.com/llvm/llvm-project/pull/220232

>From f40d086693e79bcce376c8489e7fe309c5e7cf76 Mon Sep 17 00:00:00 2001
From: cubehan3 <cubehan3 at gmail.com>
Date: Tue, 1 Sep 2026 12:13:05 +0000
Subject: [PATCH] [mlir][bufferization] Fix crash when returning an allocation
 twice

---
 .../Transforms/BufferResultsToOutParams.cpp   | 20 ++++++++++++-
 ...ts-to-out-params-hosit-dynamic-allocs.mlir | 28 +++++++++++++++++++
 ...lts-to-out-params-hosit-static-allocs.mlir | 11 ++++++++
 3 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
index 434501b030e4a..2b614ee5f5c42 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
@@ -168,7 +168,20 @@ updateReturnOps(func::FuncOp func, ArrayRef<BlockArgument> appendedEntryArgs,
     }
     OpBuilder builder(op);
     SmallVector<SmallVector<Value>> dynamicSizes;
+    SetVector<Operation *> toErase;
+    DenseMap<Value, BlockArgument> primaryArgs;
+    DenseMap<Value, SmallVector<Value>> dynSizes;
     for (auto [orig, arg] : llvm::zip(copyIntoOutParams, appendedEntryArgs)) {
+      auto [it, inserted] = primaryArgs.try_emplace(orig, arg);
+      if (!inserted) {
+        if (auto dynIt = dynSizes.find(orig); dynIt != dynSizes.end())
+          dynamicSizes.push_back(dynIt->second);
+        if (it->second != arg &&
+            failed(options.memCpyFn(builder, op->getLoc(), it->second, arg)))
+          return WalkResult::interrupt();
+        continue;
+      }
+
       bool hoistStaticAllocs =
           options.hoistStaticAllocs &&
           cast<MemRefType>(orig.getType()).hasStaticShape();
@@ -182,13 +195,18 @@ updateReturnOps(func::FuncOp func, ArrayRef<BlockArgument> appendedEntryArgs,
         if (hoistDynamicAllocs) {
           SmallVector<Value> dynamicSize = getDynamicSize(orig, func);
           dynamicSizes.push_back(dynamicSize);
+          dynSizes.try_emplace(orig, std::move(dynamicSize));
         }
-        orig.getDefiningOp()->erase();
+        toErase.insert(orig.getDefiningOp());
       } else {
         if (failed(options.memCpyFn(builder, op.getLoc(), orig, arg)))
           return WalkResult::interrupt();
       }
     }
+
+    for (Operation *eraseOp : toErase)
+      eraseOp->erase();
+
     func::ReturnOp::create(builder, op.getLoc(), keepAsReturnOperands);
     op.erase();
     auto dynamicSizePair =
diff --git a/mlir/test/Transforms/buffer-results-to-out-params-hosit-dynamic-allocs.mlir b/mlir/test/Transforms/buffer-results-to-out-params-hosit-dynamic-allocs.mlir
index f33eb8e26fbce..bbe26f0f91daa 100644
--- a/mlir/test/Transforms/buffer-results-to-out-params-hosit-dynamic-allocs.mlir
+++ b/mlir/test/Transforms/buffer-results-to-out-params-hosit-dynamic-allocs.mlir
@@ -77,3 +77,31 @@ func.func @complex_alloc_test(%size0 : index, %size1: index) {
 //       CHECK:   call @complex_alloc(%[[size0]], %[[size1]], %[[alloc0]], %[[alloc1]], %[[alloc2]]) : (index, index, memref<?x?xf32>, memref<4xf32>, memref<?xf32>) -> ()
 //       CHECK:   "test.sink"(%[[alloc0]], %[[alloc1]], %[[alloc2]]) : (memref<?x?xf32>, memref<4xf32>, memref<?xf32>) -> ()
 //       CHECK: }
+
+// -----
+
+func.func private @duplicate_dynamic(%size : index)
+    -> (memref<?xf32>, memref<?xf32>) {
+  %alloc = memref.alloc(%size) : memref<?xf32>
+  return %alloc, %alloc : memref<?xf32>, memref<?xf32>
+}
+
+func.func @duplicate_dynamic_test(%size : index) {
+  %alloc0, %alloc1 = call @duplicate_dynamic(%size)
+      : (index) -> (memref<?xf32>, memref<?xf32>)
+  "test.sink"(%alloc0, %alloc1) : (memref<?xf32>, memref<?xf32>) -> ()
+}
+
+// CHECK-LABEL: func.func private @duplicate_dynamic(
+//  CHECK-SAME:   %{{.*}}: index,
+//  CHECK-SAME:   %[[OUT0:.*]]: memref<?xf32>,
+//  CHECK-SAME:   %[[OUT1:.*]]: memref<?xf32>) {
+//       CHECK:   memref.copy %[[OUT0]], %[[OUT1]] : memref<?xf32> to memref<?xf32>
+//       CHECK:   return
+
+// CHECK-LABEL: func.func @duplicate_dynamic_test(
+//  CHECK-SAME:   %[[SIZE:.*]]: index) {
+//       CHECK:   %[[ALLOC0:.*]] = memref.alloc(%[[SIZE]]) : memref<?xf32>
+//       CHECK:   %[[ALLOC1:.*]] = memref.alloc(%[[SIZE]]) : memref<?xf32>
+//       CHECK:   call @duplicate_dynamic(%[[SIZE]], %[[ALLOC0]], %[[ALLOC1]]) : (index, memref<?xf32>, memref<?xf32>) -> ()
+//       CHECK:   "test.sink"(%[[ALLOC0]], %[[ALLOC1]]) : (memref<?xf32>, memref<?xf32>) -> ()
diff --git a/mlir/test/Transforms/buffer-results-to-out-params-hosit-static-allocs.mlir b/mlir/test/Transforms/buffer-results-to-out-params-hosit-static-allocs.mlir
index 97fb9f45f57f8..5a4a7c7eb5d30 100644
--- a/mlir/test/Transforms/buffer-results-to-out-params-hosit-static-allocs.mlir
+++ b/mlir/test/Transforms/buffer-results-to-out-params-hosit-static-allocs.mlir
@@ -49,3 +49,14 @@ func.func private @return_arg(%arg0: memref<128x256xf32>, %arg1: memref<128x256x
   "test.source"(%arg0, %arg1)  : (memref<128x256xf32>, memref<128x256xf32>) -> ()
   return %arg0 : memref<128x256xf32>
 }
+
+// CHECK-LABEL: func.func private @duplicate_return_value(
+// CHECK-SAME:    %[[ARG0:.*]]: memref<4xf32>, %[[ARG1:.*]]: memref<4xf32>) {
+// CHECK-NOT:     memref.alloc
+// CHECK:         memref.copy %[[ARG0]], %[[ARG1]] : memref<4xf32> to memref<4xf32>
+// CHECK:         return
+func.func private @duplicate_return_value()
+    -> (memref<4xf32>, memref<4xf32>) {
+  %a = memref.alloc() : memref<4xf32>
+  return %a, %a : memref<4xf32>, memref<4xf32>
+}