[Mlir-commits] [mlir] [MLIR][Bufferization] BufferResultsToOutParams: Add an option to eliminate AllocOp and avoid Copy (PR #90011)

[Wenyi Zhao]

================
@@ -118,10 +119,14 @@ static LogicalResult updateReturnOps(func::FuncOp func,
         keepAsReturnOperands.push_back(operand);
     }
     OpBuilder builder(op);
-    for (auto t : llvm::zip(copyIntoOutParams, appendedEntryArgs)) {
-      if (failed(
-              memCpyFn(builder, op.getLoc(), std::get<0>(t), std::get<1>(t))))
-        return WalkResult::interrupt();
+    for (auto [orig, arg] : llvm::zip(copyIntoOutParams, appendedEntryArgs)) {
+      if (avoidBufferResultAllocAndCopy && isa<memref::AllocOp>(orig.getDefiningOp())) {
----------------
wyzero wrote:

Just curious about this: does it support the case where the shape of the MemRef is unknown at compile time?

I mean for example like following:

```
// the shape of %arg0 and %arg1 may be different.
func.func @main(%arg0: memref<?x?xf32>, %arg1 : memref<?x?xf32>) -> memref<?x?xf32> {
  %0 = scf.if %pred {
    scf.yield %arg0 :  memref<?x?xf32>
  } else {
    scf.yield %arg1 :  memref<?x?xf32>
  }
  %2 = memref.alloc(…)
  concat(%0, %0, %2)
  return %2 : memref<?x?xf32>
}
```

->

```
// the shape of %2 is depends on the runtime value, how does the user to alloc the buffer in advance and pass it to the funciton?
func.func @main(%arg0: memref<?x?xf32>, %arg1 : memref<?x?xf32>, %2 : memref<?x?xf32>) {
  %0 = scf.if %pred {
    scf.yield %arg0 :  memref<?x?xf32>
  } else {
    scf.yield %arg1 :  memref<?x?xf32>
  }
  concat(%0, %0, %2)
  return
}
```

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