[Mlir-commits] [mlir] [mlir][memref] Fix hoist-static-allocs option of buffer-results-to-out-params when function parameters are returned (PR #102093)

[llvmlistbot at llvm.org]

================
@@ -34,4 +34,15 @@ func.func @basic_dynamic(%d: index) -> (memref<?xf32>) {
   %b = memref.alloc(%d) : memref<?xf32>
   "test.source"(%b)  : (memref<?xf32>) -> ()
   return %b : memref<?xf32>
-}
\ No newline at end of file
+}
+
+// CHECK-LABEL:   func @return_arg(
+// CHECK-SAME:        %[[ARG0:.*]]: memref<128x256xf32>, %[[ARG1:.*]]: memref<128x256xf32>, %[[ARG2:.*]]: memref<128x256xf32>) {
+// CHECK:           "test.source"(%[[ARG0]], %[[ARG1]])
+// CHECK-NOT:       memref.copy
+// CHECK:           return
+// CHECK:         }
+func.func @return_arg(%arg0: memref<128x256xf32>, %arg1: memref<128x256xf32>) -> memref<128x256xf32> {
+  "test.source"(%arg0, %arg1)  : (memref<128x256xf32>, memref<128x256xf32>) -> ()
+  return %arg0 : memref<128x256xf32>
+}
----------------
Menooker wrote:

>  Originally, users expected the value in the returned memref to be identical to %arg0. After the copy eliminated, it's different.

Yes, but maybe we can expect the users to know that it is a in-place returning function when they specify this `hoist-static-alloc` option?

Maybe we should rename hoist-static-alloc to something like avoid-copy-if-you-can. In current main branch, it will avoid the memcpy and remove the temp-buffer allocation if a memref is returned. We are adding another optimization (and bug fix) in this PR, to remove the arg-to-arg memref copy. I think the behaviors are similar.

Another way is to add another option in this pass, like `elim-arg-copy` to exactly control this optimization.

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