[mlir][bufferization] Fix crash when returning an allocation twice

Project: llvm/llvm-project | Language: C++ | Concepts: memory-management, refactoring, error-handling, testing-strategy | Difficulty: mid
Link: https://github.com/llvm/llvm-project/pull/220232

The problem

When a function returns the same memory allocation more than once (e.g., return %a, %a), the buffer-results-to-out-params transformation would crash. This was because the code erased the allocation after handling its first occurrence, so subsequent uses would try to access a dangling reference, leading to a use-after-free bug.

Why it's non-trivial

A naive fix might simply try to avoid erasing the allocation multiple times, but that would not address the underlying problem: each returned value must be mapped to a unique out-parameter, and if the same allocation is returned multiple times, their buffers must be kept in sync. Additionally, dynamic allocations require careful management of size information, and the fix must work for both static and dynamic shapes.

The approach

The core logic is in the updateReturnOps function, which was updated to:

• Track, for each returned value, its "primary" out-parameter.
• Defer erasure of allocations until all return operands are processed.
• For duplicate returns, emit a memref.copy from the primary to the duplicate out-parameter.
• Cache dynamic sizes so that the caller allocates one buffer per occurrence.

Snippet 1: Track primary out-params and defer erasure

+    SetVector<Operation *> toErase;
+    DenseMap<Value, BlockArgument> primaryArgs;
+    DenseMap<Value, SmallVector<Value>> dynSizes;

This sets up data structures to track which allocation each out-parameter corresponds to, and which allocations need to be erased after all processing is done.

Snippet 2: Handle duplicate returns and copy contents

+      auto [it, inserted] = primaryArgs.try_emplace(orig, arg);
+      if (!inserted) {
+        if (auto dynIt = dynSizes.find(orig); dynIt != dynSizes.end())
+          dynamicSizes.push_back(dynIt->second);
+        if (it->second != arg &&
+            failed(options.memCpyFn(builder, op->getLoc(), it->second, arg)))
+          return WalkResult::interrupt();
+        continue;
+      }

This checks if the allocation has already been seen. If so, it copies the contents from the primary out-parameter to the duplicate, and skips further processing for this occurrence.

Snippet 3: Defer erasure of allocations

+    for (Operation *eraseOp : toErase)
+      eraseOp->erase();

After all return operands are handled, the allocations are erased. This ensures that no use-after-free occurs when the same allocation is returned multiple times.

Snippet 4: Test cases for static and dynamic allocs

The test files (buffer-results-to-out-params-hosit-dynamic-allocs.mlir and buffer-results-to-out-params-hosit-static-allocs.mlir) add cases like:

func.func private @duplicate_dynamic(%size : index)
    -> (memref<?xf32>, memref<?xf32>) {
  %alloc = memref.alloc(%size) : memref<?xf32>
  return %alloc, %alloc : memref<?xf32>, memref<?xf32>
}

This ensures the transformation emits a memref.copy and allocates two buffers in the caller.

From the review thread

No notable review discussion on this PR.

Takeaways

• When transforming code that manipulates memory allocations, always account for the possibility of multiple uses or returns of the same buffer; erasure must be deferred until all uses are handled.
• For each logical value (like an allocation), track a canonical or "primary" representative when mapping to new function arguments or out-parameters, and copy as needed for duplicates.
• When handling both static and dynamic shapes, cache dynamic size information to ensure correct buffer allocation and avoid recomputation or mismatches.