[MLIR][BufferResultsToOutParamsPass] Add Option to Modify Public Function's Signature (llvm#167248)

Since llvm#162441,
`buffer-results-to-out-params` transforms `private` functions only.

But, as mentioned in
llvm#162441 (comment),
this is a breaking change for pipelines handling C code. Our pipeline
@EfficientComputer is also affected by this breaking change.

Therefore, this PR adds an opt-in flag to allow `public` functions to be
transformed by `BufferResultsToOutParamsPass`.

Author: veera-sivarajan

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

@@ -171,6 +171,9 @@ struct BufferResultsToOutParamsOpts {
   /// If true, the pass eliminates the memref.alloc and memcpy if the returned
   /// memref is allocated in the current function and has dynamic shape.
   bool hoistDynamicAllocs = false;
+
+  /// If true, the pass modifies the function signatures of public functions.
+  bool modifyPublicFunctions = false;
 };
 
 /// Replace buffers that are returned from a function with an out parameter.

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

@@ -258,6 +258,9 @@ def BufferResultsToOutParamsPass
               /*default=*/"false", "Hoist static allocations to call sites.">,
        Option<"hoistDynamicAllocs", "hoist-dynamic-allocs", "bool",
               /*default=*/"false", "Hoist dynamic allocations to call sites.">,
+       Option<"modifyPublicFunctions", "modify-public-functions", "bool",
+              /*default=*/"false", "Modify function signatures of public "
+              "functions.">,
   ];
   let dependentDialects = ["memref::MemRefDialect"];
 }

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

@@ -217,7 +217,9 @@ updateCalls(ModuleOp module, const AllocDynamicSizesMap &map,
     }
     if (!options.filterFn(&callee))
       return;
-    if (callee.isExternal() || callee.isPublic())
+    if (callee.isPublic() && !options.modifyPublicFunctions)
+      return;
+    if (callee.isExternal())
       return;
 
     SmallVector<Value, 6> replaceWithNewCallResults;
@@ -295,7 +297,9 @@ LogicalResult mlir::bufferization::promoteBufferResultsToOutParams(
   // function.
   AllocDynamicSizesMap map;
   for (auto func : module.getOps<func::FuncOp>()) {
-    if (func.isExternal() || func.isPublic())
+    if (func.isPublic() && !options.modifyPublicFunctions)
+      continue;
+    if (func.isExternal())
       continue;
     if (!options.filterFn(&func))
       continue;
@@ -326,6 +330,8 @@ struct BufferResultsToOutParamsPass
       options.hoistStaticAllocs = true;
     if (hoistDynamicAllocs)
       options.hoistDynamicAllocs = true;
+    if (modifyPublicFunctions)
+      options.modifyPublicFunctions = true;
 
     if (failed(bufferization::promoteBufferResultsToOutParams(getOperation(),
                                                               options)))

mlir/test/Transforms/buffer-results-to-out-params-modify-public-functions.mlir

@@ -0,0 +1,40 @@
+// RUN: mlir-opt -p 'builtin.module(buffer-results-to-out-params{modify-public-functions})' %s | FileCheck %s
+
+// Test if `public` functions' return values are transformed into out parameters
+// when `buffer-results-to-out-params` is invoked with `modifyPublicFunctions`.
+
+// CHECK-LABEL:   func.func @basic(
+// CHECK-SAME:                     %[[ARG0:.*]]: memref<f32>) {
+// CHECK:           %[[VAL_0:.*]] = "test.source"() : () -> memref<f32>
+// CHECK:           memref.copy %[[VAL_0]], %[[ARG0]] : memref<f32> to memref<f32>
+// CHECK:           return
+// CHECK:         }
+func.func @basic() -> (memref<f32>) {
+  %0 = "test.source"() : () -> (memref<f32>)
+  return %0 : memref<f32>
+}
+
+// CHECK-LABEL:   func.func @presence_of_existing_arguments(
+// CHECK-SAME:      %[[ARG0:.*]]: memref<1xf32>,
+// CHECK-SAME:      %[[ARG1:.*]]: memref<2xf32>) {
+// CHECK:           %[[VAL_0:.*]] = "test.source"() : () -> memref<2xf32>
+// CHECK:           memref.copy %[[VAL_0]], %[[ARG1]] : memref<2xf32> to memref<2xf32>
+// CHECK:           return
+// CHECK:         }
+func.func @presence_of_existing_arguments(%arg0: memref<1xf32>) -> (memref<2xf32>) {
+  %0 = "test.source"() : () -> (memref<2xf32>)
+  return %0 : memref<2xf32>
+}
+
+// CHECK-LABEL:   func.func @multiple_results(
+// CHECK-SAME:      %[[ARG0:.*]]: memref<1xf32>,
+// CHECK-SAME:      %[[ARG1:.*]]: memref<2xf32>) {
+// CHECK:           %[[VAL_0:.*]]:2 = "test.source"() : () -> (memref<1xf32>, memref<2xf32>)
+// CHECK:           memref.copy %[[VAL_0]]#0, %[[ARG0]] : memref<1xf32> to memref<1xf32>
+// CHECK:           memref.copy %[[VAL_0]]#1, %[[ARG1]] : memref<2xf32> to memref<2xf32>
+// CHECK:           return
+// CHECK:         }
+func.func @multiple_results() -> (memref<1xf32>, memref<2xf32>) {
+  %0, %1 = "test.source"() : () -> (memref<1xf32>, memref<2xf32>)
+  return %0, %1 : memref<1xf32>, memref<2xf32>
+}