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

Menooker · 2024-08-06T02:47:12Z

buffer-results-to-out-params pass will have a referencing nullptr error when hoist-static-allocs option is on, when the return value of a function is a parameter of the function. This PR fixes this issue.

llvmbot · 2024-08-06T02:47:41Z

@llvm/pr-subscribers-mlir-bufferization

@llvm/pr-subscribers-mlir

Author: Menooker (Menooker)

Changes

buffer-results-to-out-params pass will have a nullptr error when hoist-static-allocs option is on, when the return value of a function is a parameter of the function. This PR fixes this issue and let the pass remove the return value in the ReturnOp when

the value type is memref
and the value is a function parameter

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

3 Files Affected:

(modified) mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td (+5)
(modified) mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp (+15-2)
(modified) mlir/test/Transforms/buffer-results-to-out-params-elim.mlir (+12-1)

diff --git a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
index 1cece818dbbbc..d6f13b2153828 100644
--- a/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td
@@ -321,6 +321,11 @@ def BufferResultsToOutParams : Pass<"buffer-results-to-out-params", "ModuleOp">
     This optimization applies on the returned memref which has static shape and
     is allocated by memref.alloc in the function. It will use the memref given
     in function argument to replace the allocated memref.
+
+    If the hoist-static-allocs option is on, and a function returns a memref
+    from the function argument, the pass will avoid the memory-copy from
+    the input function argument to the "out param", and leave the "out param"
+    unused. 
   }];
   let options = [
     Option<"addResultAttribute", "add-result-attr", "bool",
diff --git a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
index b19636adaa69e..16a42ca779381 100644
--- a/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
+++ b/mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp
@@ -102,6 +102,14 @@ updateFuncOp(func::FuncOp func,
   return success();
 }
 
+static bool isFunctionArgument(mlir::Value value) {
+  // Check if the value is a Function argument
+  if (auto blockArg = dyn_cast<mlir::BlockArgument>(value)) {
+    return blockArg.getOwner()->isEntryBlock();
+  }
+  return false;
+}
+
 // Updates all ReturnOps in the scope of the given func::FuncOp by either
 // keeping them as return values or copying the associated buffer contents into
 // the given out-params.
@@ -120,10 +128,15 @@ static LogicalResult updateReturnOps(func::FuncOp func,
     }
     OpBuilder builder(op);
     for (auto [orig, arg] : llvm::zip(copyIntoOutParams, appendedEntryArgs)) {
-      if (hoistStaticAllocs && isa<memref::AllocOp>(orig.getDefiningOp()) &&
-          mlir::cast<MemRefType>(orig.getType()).hasStaticShape()) {
+      bool mayHoistStaticAlloc =
+          hoistStaticAllocs &&
+          mlir::cast<MemRefType>(orig.getType()).hasStaticShape();
+      if (mayHoistStaticAlloc &&
+          isa_and_nonnull<memref::AllocOp>(orig.getDefiningOp())) {
         orig.replaceAllUsesWith(arg);
         orig.getDefiningOp()->erase();
+      } else if (mayHoistStaticAlloc && isFunctionArgument(orig)) {
+        // do nothing but remove the value from the return op.
       } else {
         if (failed(memCpyFn(builder, op.getLoc(), orig, arg)))
           return WalkResult::interrupt();
diff --git a/mlir/test/Transforms/buffer-results-to-out-params-elim.mlir b/mlir/test/Transforms/buffer-results-to-out-params-elim.mlir
index f77dbfaa6cb11..2bd9a9a045531 100644
--- a/mlir/test/Transforms/buffer-results-to-out-params-elim.mlir
+++ b/mlir/test/Transforms/buffer-results-to-out-params-elim.mlir
@@ -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>
+}

cxy-1993

Thanks for the patch. Looks good to me in general, except some logical errors.

mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp

cxy-1993 · 2024-08-06T13:13:09Z

mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp

        orig.replaceAllUsesWith(arg);
        orig.getDefiningOp()->erase();
+      } else if (mayHoistStaticAlloc && isFunctionArgument(orig)) {


e don't need to set the hoistStaticAllocs option to true to do this optimization: the return value doesn't involve any allocations that weren't already present in the function parameters.

It indeed changes the behavior of the program. It is not sure to me what is the expected behavior of this pass without hoistStaticAllocs, if we return the memref from an arg. If we should copy the data from arg to arg?

We dont have to. My point is we dont need hoistStaticAllocs to eliminate copy.

Actually I am not sure if it is ok to avoid the memcpy for this case...

Consider the simple function:

func add1(%0: memref<...>) -> memref<...> { my.add(%0) return %0 }

This pass will change the signature to:

func add1(%0: memref<...>, %out: memref<...>) { .... }

If the user of this pass is unaware of the in-place-return behavior, they may pass different buffers for %0 and %out`. The safest way is to copy %0 to %out. It will always meet the expectation of the users.

If we remove the memcpy from %0 to %out, actually %out is an unused parameter, and we should expect the user aware of this in-place-return behavior. If we turn it on by default, will it be a surprise to them?

You're right, this will confuse users. We should avoid doing this.

cxy-1993 · 2024-08-06T13:13:23Z

mlir/include/mlir/Dialect/Bufferization/Transforms/Passes.td

@@ -321,6 +321,11 @@ def BufferResultsToOutParams : Pass<"buffer-results-to-out-params", "ModuleOp">
    This optimization applies on the returned memref which has static shape and
    is allocated by memref.alloc in the function. It will use the memref given
    in function argument to replace the allocated memref.
+
+    If the hoist-static-allocs option is on, and a function returns a memref


…hoist_static_allocs

Menooker · 2024-08-20T13:35:10Z

Hi @cxy-1993 I have updated this PR for the entry-block issue. Would you please have another look? Thanks!

mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp

mlir/test/Transforms/buffer-results-to-out-params-elim.mlir

Menooker · 2024-09-03T07:06:29Z

Hi @cxy-1993 , thanks for your review and comments. I have removed the confusing optimization in this patch. And I am focusing on the bug fix now. Would you please help to review again? Thanks!

mlir/test/Transforms/buffer-results-to-out-params-elim.mlir

cxy-1993

LGTM. Please fix test style before merge.

fix return arg

41ea098

Menooker requested a review from matthias-springer as a code owner August 6, 2024 02:47

llvmbot added mlir mlir:bufferization Bufferization infrastructure labels Aug 6, 2024

cxy-1993 requested changes Aug 6, 2024

View reviewed changes

Mei, Yijie added 2 commits August 20, 2024 21:29

Merge branch 'main' of https://github.com/llvm/llvm-project into fix_…

e5cf2f2

…hoist_static_allocs

fix comment

0aa9133

cxy-1993 reviewed Aug 22, 2024

View reviewed changes

mlir/lib/Dialect/Bufferization/Transforms/BufferResultsToOutParams.cpp Outdated Show resolved Hide resolved

mlir/test/Transforms/buffer-results-to-out-params-elim.mlir Show resolved Hide resolved

Mei, Yijie added 2 commits September 3, 2024 14:43

Merge branch 'main' into fix_hoist_static_allocs

6845cd4

remove optimizations

30d7d86

cxy-1993 reviewed Sep 3, 2024

View reviewed changes

mlir/test/Transforms/buffer-results-to-out-params-elim.mlir Show resolved Hide resolved

cxy-1993 approved these changes Sep 3, 2024

View reviewed changes

add seperator

a943762

Menooker merged commit 26645ae into llvm:main Sep 4, 2024
8 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

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

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

Uh oh!

Menooker commented Aug 6, 2024 •

edited

Loading

Uh oh!

llvmbot commented Aug 6, 2024 •

edited

Loading

Uh oh!

cxy-1993 left a comment

Uh oh!

Uh oh!

cxy-1993 Aug 6, 2024

Uh oh!

Menooker Aug 6, 2024

Uh oh!

cxy-1993 Aug 8, 2024

Uh oh!

Menooker Aug 9, 2024

Uh oh!

cxy-1993 Aug 11, 2024

Uh oh!

cxy-1993 Aug 6, 2024

Uh oh!

Menooker commented Aug 20, 2024

Uh oh!

Uh oh!

Uh oh!

Menooker commented Sep 3, 2024

Uh oh!

Uh oh!

cxy-1993 left a comment •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

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

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

Uh oh!

Conversation

Menooker commented Aug 6, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Aug 6, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cxy-1993 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

cxy-1993 Aug 6, 2024

Choose a reason for hiding this comment

Uh oh!

Menooker Aug 6, 2024

Choose a reason for hiding this comment

Uh oh!

cxy-1993 Aug 8, 2024

Choose a reason for hiding this comment

Uh oh!

Menooker Aug 9, 2024

Choose a reason for hiding this comment

Uh oh!

cxy-1993 Aug 11, 2024

Choose a reason for hiding this comment

Uh oh!

cxy-1993 Aug 6, 2024

Choose a reason for hiding this comment

Uh oh!

Menooker commented Aug 20, 2024

Uh oh!

Uh oh!

Uh oh!

Menooker commented Sep 3, 2024

Uh oh!

Uh oh!

cxy-1993 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Menooker commented Aug 6, 2024 •

edited

Loading

llvmbot commented Aug 6, 2024 •

edited

Loading

cxy-1993 left a comment •

edited

Loading