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

[llvmlistbot at llvm.org]

================
@@ -178,17 +183,32 @@ 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);
         }
-        orig.getDefiningOp()->erase();
+
+        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 &&
----------------
AkunCube wrote:

Done. I removed the it->second != arg check because duplicate results always use different out parameters. I added orig != arg to the other path because hoisting can rewrite a later return operand to the same out parameter, so we can avoid emitting that self-copy directly instead of relying on canonicalization later.

Example:
```mlir
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>
}
```

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