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

[Matthias Springer]

================
@@ -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);
----------------
matthias-springer wrote:

What happens `orig` is not in the cache? Can that even happen? I'm having a hard time understanding how many elements are pushed to `dynamicSizes` overall, because there are two different places where that vector is populated.

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