[mlir][bufferize] Add filterFn option to BufferResultsToOutParams

[cota]

Summary

This allows users to restrict the transformation to a
subset of the functions in a module.

For example, a user might want to apply the transformation to
a module's entry point, but not to the calls in the module
because those calls might refer to external C functions outside
of their control.

[nicolasvasilache]

strong -1 to reintroducing markers, we've clearly seen their limits.

Since the usage is trivial here, we can use the "deny" mechanism with a lambda to make it dynamic (like dynamic legality).

[springerm]

(You can add the lambda to BufferResultsToOutParamsOptions. E.g., it returns true for functions that should be processed. By default it can return true for all functions.)

[cota]

+1 for the lambda.
Is it OK to keep the markers for testing only though? Otherwise I don't know how I'd test this.
Thanks!

[springerm]

Is it OK to keep the markers for testing only though? Otherwise I don't know how I'd test this.

Ideally, that should be a separate test pass then. But IMO it would be OK to go without a test for this, the extra boilerplate is not worth it.

[nicolasvasilache]

In D137264#3902645, @springerm wrote:

(You can add the lambda to BufferResultsToOutParamsOptions. E.g., it returns true for functions that should be processed. By default it can return true for all functions.)

I was thinking to add the lambda as an optional parameter to the deny mechanism.
I wouldn't add an extra one-off with a default for a specific op (if possible)

[springerm]

I see what you meant. This is not the bufferization pass. We don't have a deny mechanism (or any options) yet.

[cota]

Remove the marker; add filterFn lambda.

Please take another look! Thanks.

[cota]

I also removed the tests as suggested by @springerm.

[springerm]

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

135–137

nit: trivial braces can be removed

188–190

nit: trivial braces

[cota]

if (foo)

bar;

[cota]

Resolve nits.

[nicolasvasilache]

In D137264#3903070, @springerm wrote:

I see what you meant. This is not the bufferization pass. We don't have a deny mechanism (or any options) yet.

Ah ok, nm then, this then looks fine.

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

 /// Creates a pass that moves allocations upwards to reduce the number of
 /// required copies that are inserted during the BufferDeallocation pass.
 std::unique_ptr<Pass> createBufferHoistingPass();
 
 /// Creates a pass that moves allocations upwards out of loops. This avoids
 /// reallocations inside of loops.
 std::unique_ptr<Pass> createBufferLoopHoistingPass();
 
+// Options struct for BufferResultsToOutParams pass.
+// Note: defined only here, not in tablegen.
+struct BufferResultsToOutParamsOptions {
+  // Filter function; returns true if the function should be converted.
+  // Defaults to true, i.e. all functions are converted.
+  llvm::function_ref<bool(func::FuncOp *)> filterFn = [](func::FuncOp *func) {
+    return true;
+  };
+};
+
 /// Creates a pass that converts memref function results to out-params.
-std::unique_ptr<Pass> createBufferResultsToOutParamsPass();
+std::unique_ptr<Pass> createBufferResultsToOutParamsPass(
+    const BufferResultsToOutParamsOptions &options = {});
 
 /// Replace buffers that are returned from a function with an out parameter.
 /// Also update all call sites.
-LogicalResult promoteBufferResultsToOutParams(ModuleOp module);
+LogicalResult
+promoteBufferResultsToOutParams(ModuleOp module,
+                                const BufferResultsToOutParamsOptions &options);
 
 /// Creates a pass that drops memref function results that are equivalent to a
 /// function argument.
 std::unique_ptr<Pass> createDropEquivalentBufferResultsPass();
 
 /// Create a pass that rewrites tensor.empty to bufferization.alloc_tensor.
 std::unique_ptr<Pass> createEmptyTensorToAllocTensorPass();
 

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

for (auto t : llvm::zip(copyIntoOutParams, appendedEntryArgs))
                                      std::get<1>(t));
     builder.create<func::ReturnOp>(op.getLoc(), keepAsReturnOperands);
     op.erase();
   });
 }
 
 // Updates all CallOps in the scope of the given ModuleOp by allocating
 // temporary buffers for newly introduced out params.
-static LogicalResult updateCalls(ModuleOp module) {
+static LogicalResult
+updateCalls(ModuleOp module,
+            const bufferization::BufferResultsToOutParamsOptions &options) {
   bool didFail = false;
+  SymbolTable symtab(module);
   module.walk([&](func::CallOp op) {
+    auto callee = symtab.lookup<func::FuncOp>(op.getCallee());
+    if (!callee) {
+      op.emitError() << "cannot find callee '" << op.getCallee() << "' in "
+                     << "symbol table";
+      didFail = true;
+      return;
+    }
+    if (!options.filterFn(&callee))
+      return;
     SmallVector<Value, 6> replaceWithNewCallResults;
     SmallVector<Value, 6> replaceWithOutParams;
     for (OpResult result : op.getResults()) {
       if (result.getType().isa<MemRefType>())
         replaceWithOutParams.push_back(result);
       else
         replaceWithNewCallResults.push_back(result);
     }
     SmallVector<Value, 6> outParams;
module.walk([&](func::CallOp op) {
     for (auto t : llvm::zip(replaceWithNewCallResults, newCall.getResults()))
       std::get<0>(t).replaceAllUsesWith(std::get<1>(t));
     op.erase();
   });
 
   return failure(didFail);
 }
 
-LogicalResult
-mlir::bufferization::promoteBufferResultsToOutParams(ModuleOp module) {
+LogicalResult mlir::bufferization::promoteBufferResultsToOutParams(
+    ModuleOp module,
+    const bufferization::BufferResultsToOutParamsOptions &options) {
   for (auto func : module.getOps<func::FuncOp>()) {
+    if (!options.filterFn(&func))
+      continue;
     SmallVector<BlockArgument, 6> appendedEntryArgs;
     if (failed(updateFuncOp(func, appendedEntryArgs)))
       return failure();
     if (func.isExternal())
       continue;
     updateReturnOps(func, appendedEntryArgs);
   }
-  if (failed(updateCalls(module)))
+  if (failed(updateCalls(module, options)))
     return failure();
   return success();
 }
 
 namespace {
 struct BufferResultsToOutParamsPass
     : bufferization::impl::BufferResultsToOutParamsBase<
           BufferResultsToOutParamsPass> {
+  explicit BufferResultsToOutParamsPass(
+      const bufferization::BufferResultsToOutParamsOptions &options)
+      : options(options) {}
+
   void runOnOperation() override {
-    if (failed(bufferization::promoteBufferResultsToOutParams(getOperation())))
+    if (failed(bufferization::promoteBufferResultsToOutParams(getOperation(),
+                                                              options)))
       return signalPassFailure();
   }
+
+private:
+  bufferization::BufferResultsToOutParamsOptions options;
 };
 } // namespace
 
-std::unique_ptr<Pass>
-mlir::bufferization::createBufferResultsToOutParamsPass() {
-  return std::make_unique<BufferResultsToOutParamsPass>();
+std::unique_ptr<Pass> mlir::bufferization::createBufferResultsToOutParamsPass(
+    const bufferization::BufferResultsToOutParamsOptions &options) {
+  return std::make_unique<BufferResultsToOutParamsPass>(options);
 }