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

the cache miss is reachable and correct, so no assert:

1. When can orig be missing from the cache?

The cache is only populated when the first occurrence of orig takes the hoisted-dynamic-alloc path. A duplicate can legitimately miss it in two cases:

 - orig is a hoisted static-shaped alloc (e.g. return %a, %a with %a = memref.alloc() : memref<4xf32> under hoist-static-allocs). The first occurrence is promoted with just replaceAllUsesWith, so no sizes are computed or cached. Pushing nothing for the duplicate is correct: in updateCalls, dynamicSizesIndex only advances for dynamically-shaped out-params, so a static position consumes no entry. (Covered by the new @duplicate_return_value test.)
 - orig is not hoisted at all (e.g. a function argument, or hoisting is disabled). The first occurrence only emits memref.copy orig -> arg and caches nothing.

 In both cases the duplicate needs no size entry — only the memref.copy from the primary out-param, which is emitted right below. So the miss is the expected case, not an invariant violation


2. How many elements end up in dynamicSizes?

 Exactly one entry per return operand that is a hoisted dynamically-shaped allocation — counted per occurrence (i.e. per out-param position), not per unique value. The two push sites are "first occurrence" (compute getDynamicSize once, push, seed the cache) and "duplicate occurrence" (push the cached entry again, same SSA size values).

 It has to be per-occurrence because updateCalls consumes the list positionally: the caller allocates a separate buffer for each result position and advances dynamicSizesIndex only for dynamically-shaped positions, even when two positions return the same buffer in the callee. E.g. for

 ```mlir
 %a = memref.alloc(%size) : memref<?xf32>   // dynamic
 %b = memref.alloc() : memref<4xf32>        // static
 return %a, %b, %a
```

 dynamicSizes ends up as [sizes(%a), sizes(%a)]: positions 0 and 2 each consume one entry, position 1 consumes none. This is the same invariant as before this PR — the PR only adds that duplicate occurrences also push (via the cache).

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