[Mlir-commits] [mlir] [mlir][bufferization] Fix crash when returning an allocation twice (PR #220232)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Sep 1 18:59:15 PDT 2026
https://github.com/AkunCube updated https://github.com/llvm/llvm-project/pull/220232
>From 8f3292d21941085362fbf485d0a14c39f18310c3 Mon Sep 17 00:00:00 2001
From: cubehan3 <cubehan3 at gmail.com>
Date: Tue, 1 Sep 2026 12:13:05 +0000
Subject: [PATCH 1/2] 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>
+}
>From b4b9f0769e350c2843569a51d5ca5618c6c71595 Mon Sep 17 00:00:00 2001
From: cubehan3 <cubehan3 at gmail.com>
Date: Wed, 2 Sep 2026 01:55:04 +0000
Subject: [PATCH 2/2] [mlir] Clarify duplicate return allocation handling
---
.../Transforms/BufferResultsToOutParams.cpp | 14 ++++++++++----
1 file changed, 10 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
index 2b614ee5f5c42..6639e1872a664 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
@@ -168,13 +168,19 @@ updateReturnOps(func::FuncOp func, ArrayRef<BlockArgument> appendedEntryArgs,
}
OpBuilder builder(op);
SmallVector<SmallVector<Value>> dynamicSizes;
- SetVector<Operation *> toErase;
+ // Map each returned memref to its primary out parameter.
DenseMap<Value, BlockArgument> primaryArgs;
- DenseMap<Value, SmallVector<Value>> dynSizes;
+ // Cache dynamic sizes by result so duplicate dynamic results can reuse
+ // their sizes when adding entries to `dynamicSizes`.
+ DenseMap<Value, SmallVector<Value>> dynamicSizesCache;
+ // Delay erasure until all return operands have been inspected because an
+ // allocation may be returned multiple times.
+ SetVector<Operation *> toErase;
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())
+ if (auto dynIt = dynamicSizesCache.find(orig);
+ dynIt != dynamicSizesCache.end())
dynamicSizes.push_back(dynIt->second);
if (it->second != arg &&
failed(options.memCpyFn(builder, op->getLoc(), it->second, arg)))
@@ -195,7 +201,7 @@ 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));
+ dynamicSizesCache.try_emplace(orig, std::move(dynamicSize));
}
toErase.insert(orig.getDefiningOp());
} else {
More information about the Mlir-commits
mailing list