[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:

Can you put all this new code inside of this `if` check?
```
      if ((hoistStaticAllocs || hoistDynamicAllocs) &&
          isa_and_nonnull<bufferization::AllocationOpInterface>(
              orig.getDefiningOp())) {
```

Then we have all the logic related to "dynamic size" in one place.

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