[Mlir-commits] [mlir] [mlir][bufferization] Fix crash when returning an allocation twice (PR #220232)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Sep 2 02:56:38 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/4] 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/4] [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 {
>From 7cca40295d6ad35c344e8e68cc77fcce63651623 Mon Sep 17 00:00:00 2001
From: cubehan3 <cubehan3 at gmail.com>
Date: Wed, 2 Sep 2026 08:55:07 +0000
Subject: [PATCH 3/4] [mlir][bufferization] Simplify duplicate allocation
handling
---
.../Transforms/BufferResultsToOutParams.cpp | 28 ++++++++-----------
1 file changed, 12 insertions(+), 16 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
index 6639e1872a664..aa1f398e7f045 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
@@ -170,24 +170,10 @@ updateReturnOps(func::FuncOp func, ArrayRef<BlockArgument> appendedEntryArgs,
SmallVector<SmallVector<Value>> dynamicSizes;
// Map each returned memref to its primary out parameter.
DenseMap<Value, BlockArgument> primaryArgs;
- // 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 = dynamicSizesCache.find(orig);
- dynIt != dynamicSizesCache.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();
@@ -197,12 +183,22 @@ updateReturnOps(func::FuncOp func, ArrayRef<BlockArgument> appendedEntryArgs,
if ((hoistStaticAllocs || hoistDynamicAllocs) &&
isa_and_nonnull<bufferization::AllocationOpInterface>(
orig.getDefiningOp())) {
- orig.replaceAllUsesWith(arg);
if (hoistDynamicAllocs) {
SmallVector<Value> dynamicSize = getDynamicSize(orig, func);
dynamicSizes.push_back(dynamicSize);
- dynamicSizesCache.try_emplace(orig, std::move(dynamicSize));
}
+
+ auto [it, inserted] = primaryArgs.try_emplace(orig, arg);
+ // Repeated returns of the same allocation reuse the primary out
+ // parameter; copy its contents into the out parameter for this
+ // occurrence.
+ if (!inserted) {
+ if (it->second != arg &&
+ failed(options.memCpyFn(builder, op->getLoc(), it->second, arg)))
+ return WalkResult::interrupt();
+ continue;
+ }
+ orig.replaceAllUsesWith(arg);
toErase.insert(orig.getDefiningOp());
} else {
if (failed(options.memCpyFn(builder, op.getLoc(), orig, arg)))
>From 716e0f835691c63621d84a8a2c9e2ace88fcff68 Mon Sep 17 00:00:00 2001
From: cubehan3 <cubehan3 at gmail.com>
Date: Wed, 2 Sep 2026 09:53:07 +0000
Subject: [PATCH 4/4] [mlir][bufferization] Avoid redundant copies for
duplicate returns
---
.../Transforms/BufferResultsToOutParams.cpp | 5 +--
...ts-to-out-params-hosit-dynamic-allocs.mlir | 23 ++++++++++
...lts-to-out-params-hosit-static-allocs.mlir | 43 +++++++++++++++++++
3 files changed, 68 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
index aa1f398e7f045..f227ec4149f1a 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
@@ -193,14 +193,13 @@ updateReturnOps(func::FuncOp func, ArrayRef<BlockArgument> appendedEntryArgs,
// parameter; copy its contents into the out parameter for this
// occurrence.
if (!inserted) {
- if (it->second != arg &&
- failed(options.memCpyFn(builder, op->getLoc(), it->second, arg)))
+ if (failed(options.memCpyFn(builder, op->getLoc(), it->second, arg)))
return WalkResult::interrupt();
continue;
}
orig.replaceAllUsesWith(arg);
toErase.insert(orig.getDefiningOp());
- } else {
+ } else if (orig != arg) {
if (failed(options.memCpyFn(builder, op.getLoc(), orig, arg)))
return WalkResult::interrupt();
}
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 bbe26f0f91daa..6c08c9cc124d3 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
@@ -105,3 +105,26 @@ func.func @duplicate_dynamic_test(%size : index) {
// 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>) -> ()
+
+// -----
+
+// CHECK-LABEL: func.func private @multiple_dynamic_returns_with_non_alloc(
+// CHECK-SAME: %[[SIZE:.*]]: index, %[[INPUT:.*]]: memref<?xf32>, %[[COND:.*]]: i1, %[[OUT0:.*]]: memref<?xf32>, %[[OUT1:.*]]: memref<?xf32>) {
+// CHECK: cf.cond_br %[[COND]], ^[[THEN:.*]], ^[[ELSE:.*]]
+// CHECK: ^[[THEN]]:
+// CHECK: memref.copy %[[OUT0]], %[[OUT1]] : memref<?xf32> to memref<?xf32>
+// CHECK: return
+// CHECK: ^[[ELSE]]:
+// CHECK: memref.copy %[[INPUT]], %[[OUT0]] : memref<?xf32> to memref<?xf32>
+// CHECK: memref.copy %[[INPUT]], %[[OUT1]] : memref<?xf32> to memref<?xf32>
+// CHECK: return
+func.func private @multiple_dynamic_returns_with_non_alloc(
+ %size: index, %input: memref<?xf32>, %cond: i1)
+ -> (memref<?xf32>, memref<?xf32>) {
+ %a = memref.alloc(%size) : memref<?xf32>
+ cf.cond_br %cond, ^then, ^else
+^then:
+ return %a, %a : memref<?xf32>, memref<?xf32>
+^else:
+ return %input, %input : 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 5a4a7c7eb5d30..f15ac2f4f2921 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
@@ -60,3 +60,46 @@ func.func private @duplicate_return_value()
%a = memref.alloc() : memref<4xf32>
return %a, %a : memref<4xf32>, memref<4xf32>
}
+
+// -----
+
+// CHECK-LABEL: func.func private @multiple_duplicate_returns(
+// CHECK-SAME: %[[COND:.*]]: i1, %[[OUT0:.*]]: memref<4xf32>, %[[OUT1:.*]]: memref<4xf32>) {
+// CHECK: cf.cond_br %[[COND]], ^[[THEN:.*]], ^[[ELSE:.*]]
+// CHECK: ^[[THEN]]:
+// CHECK: memref.copy %[[OUT0]], %[[OUT1]] : memref<4xf32> to memref<4xf32>
+// CHECK: return
+// CHECK: ^[[ELSE]]:
+// CHECK: memref.copy %[[OUT0]], %[[OUT1]] : memref<4xf32> to memref<4xf32>
+// CHECK: return
+func.func private @multiple_duplicate_returns(%cond: i1)
+ -> (memref<4xf32>, memref<4xf32>) {
+ %a = memref.alloc() : memref<4xf32>
+ cf.cond_br %cond, ^then, ^else
+^then:
+ return %a, %a : memref<4xf32>, memref<4xf32>
+^else:
+ return %a, %a : memref<4xf32>, memref<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func.func private @multiple_returns_with_non_alloc(
+// CHECK-SAME: %[[INPUT:.*]]: memref<4xf32>, %{{.*}}: i1, %[[OUT0:.*]]: memref<4xf32>, %[[OUT1:.*]]: memref<4xf32>) {
+// CHECK: cf.cond_br %{{.*}}, ^[[THEN:.*]], ^[[ELSE:.*]]
+// CHECK: ^[[THEN]]:
+// CHECK: memref.copy %[[OUT0]], %[[OUT1]] : memref<4xf32> to memref<4xf32>
+// CHECK: return
+// CHECK: ^[[ELSE]]:
+// CHECK: memref.copy %[[INPUT]], %[[OUT0]] : memref<4xf32> to memref<4xf32>
+// CHECK: memref.copy %[[INPUT]], %[[OUT1]] : memref<4xf32> to memref<4xf32>
+// CHECK: return
+func.func private @multiple_returns_with_non_alloc(%input: memref<4xf32>, %cond: i1)
+ -> (memref<4xf32>, memref<4xf32>) {
+ %a = memref.alloc() : memref<4xf32>
+ cf.cond_br %cond, ^then, ^else
+^then:
+ return %a, %a : memref<4xf32>, memref<4xf32>
+^else:
+ return %input, %input : memref<4xf32>, memref<4xf32>
+}
More information about the Mlir-commits
mailing list