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

[llvmlistbot at llvm.org]

================
@@ -168,7 +168,26 @@ updateReturnOps(func::FuncOp func, ArrayRef<BlockArgument> appendedEntryArgs,
     }
     OpBuilder builder(op);
     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);
----------------
AkunCube wrote:

 Done — moved the duplicate-allocation and dynamic-size handling inside the hoisting check. I also removed the cache since recomputing the dynamic sizes is cheap and keeps the logic simpler.

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