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

[Matthias Springer]

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

Is there a case where `it->second == arg`? When `orig` is inserted, we don't go into this branch. And when we return the same memref a second time, it should be mapped to a different `arg`.

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