[mlir][bufferize] Support alloc hoisting across function boundaries

[springerm]

Summary

This change integrates the BufferResultsToOutParamsPass into One-Shot Module Bufferization. This improves memory management (deallocation) when buffers are returned from a function.

Note: This currently only works with statically-sized tensors. The generated code is not very efficient yet and there are opportunities for improvment (fewer copies). By default, this new functionality is deactivated.

Depends On D125375

[tpopp]

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

273

I think this is duplicated.

274

What about having if (!hasLeakingAllocs) return success() before the following 2 cases?

[springerm]

address comments

[springerm]

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

273

this is on purpose. both function calls need an options parameter

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

struct BufferizationOptions {
   /// If set to `true`, does not modify the IR apart from adding attributes (for
   /// checking the results of the analysis) and post analysis steps.
   bool testAnalysisOnly = false;
 
   /// If set to `true`, the IR is annotated with details about RaW conflicts.
   /// For debugging only. Should be used together with `testAnalysisOnly`.
   bool printConflicts = false;
 
+  /// If set to `true`, buffers that are returned from functions are replaced
+  /// with buffer "out" parameters. At the call site, new buffers are allocated.
+  bool promoteBufferResultsToOutParams = false;
+
   /// If set to `true`, an `getAliasingOpResult` will return the corresponding
   /// "out"/"dest" OpOperand for every op that has the notion of an "out"/"dest"
   /// operand. I.e., the aliasing OpOperand of the i-th tensor OpResult is
   /// usually the i-th "out" tensor OpOperand. This is in line with
   /// destination-passing style and the default behavior. Op interface
   /// implementations must follow this contract to avoid surprising behavior.
   ///
   /// If set to `false`, BufferizableOpInterface implementations can try to be

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

 #ifndef MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H
 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_PASSES_H
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
+class ModuleOp;
+
 namespace func {
 class FuncOp;
 } // namespace func
 
 namespace bufferization {
 struct OneShotBufferizationOptions;
 
 //===----------------------------------------------------------------------===//

 
 /// Creates a pass that moves allocations upwards out of loops. This avoids
 /// reallocations inside of loops.
 std::unique_ptr<Pass> createBufferLoopHoistingPass();
 
 /// Creates a pass that converts memref function results to out-params.
 std::unique_ptr<Pass> createBufferResultsToOutParamsPass();
 
+/// Replace buffers that are returned from a function with an out parameter.
+/// Also update all call sites.
+LogicalResult promoteBufferResultsToOutParams(ModuleOp module);
+
 /// Creates a pass that finalizes a partial bufferization by removing remaining
 /// bufferization.to_tensor and bufferization.to_memref operations.
 std::unique_ptr<OperationPass<func::FuncOp>> createFinalizingBufferizePass();
 
 /// Create a pass that bufferizes all ops that implement BufferizableOpInterface
 /// with One-Shot Bufferize.
 std::unique_ptr<Pass> createOneShotBufferizePass();
 

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

Option<"fullyDynamicLayoutMaps", "fully-dynamic-layout-maps", "bool",
            "Generate MemRef types with dynamic offset+strides by default.">,
     Option<"testAnalysisOnly", "test-analysis-only", "bool",
             /*default=*/"false",
            "Test only: Only run inplaceability analysis and annotate IR">,
     Option<"printConflicts", "print-conflicts", "bool",
             /*default=*/"false",
            "Test only: Annotate IR with RaW conflicts. Requires "
            "test-analysis-only.">,
+    Option<"promoteBufferResultsToOutParams",
+           "promote-buffer-results-to-out-params", "bool", /*default=*/"false",
+           "Replace returned buffers (that were not dropped) with out params.">,
   ];
   let constructor = "mlir::bufferization::createOneShotBufferizePass()";
 }
 
 def PromoteBuffersToStack : Pass<"promote-buffers-to-stack", "func::FuncOp"> {
   let summary = "Promotes heap-based allocations to automatically managed "
                 "stack-based allocations";
   let description = [{

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

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);
 }
 
-namespace {
-struct BufferResultsToOutParamsPass
-    : BufferResultsToOutParamsBase<BufferResultsToOutParamsPass> {
-  void runOnOperation() override {
-    ModuleOp module = getOperation();
-
-    for (auto func : module.getOps<func::FuncOp>()) {
-      SmallVector<BlockArgument, 6> appendedEntryArgs;
-      updateFuncOp(func, appendedEntryArgs);
-      if (func.isExternal())
-        continue;
-      updateReturnOps(func, appendedEntryArgs);
-    }
-    if (failed(updateCalls(module)))
+LogicalResult
+mlir::bufferization::promoteBufferResultsToOutParams(ModuleOp module) {
+    return failure();
+  return success();
+}
+
+namespace {
+struct BufferResultsToOutParamsPass
+    : BufferResultsToOutParamsBase<BufferResultsToOutParamsPass> {
+  void runOnOperation() override {
+    if (failed(bufferization::promoteBufferResultsToOutParams(getOperation())))
       return signalPassFailure();
   }
 };
 } // namespace
 
 std::unique_ptr<Pass>
 mlir::bufferization::createBufferResultsToOutParamsPass() {
   return std::make_unique<BufferResultsToOutParamsPass>();
 }

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

if (!options) {
       opt.allowUnknownOps = allowUnknownOps;
       opt.alwaysAliasingWithDest = alwaysAliasingWithDest;
       opt.analysisFuzzerSeed = analysisFuzzerSeed;
       opt.createDeallocs = createDeallocs;
       opt.fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
       opt.printConflicts = printConflicts;
       opt.testAnalysisOnly = testAnalysisOnly;
       opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
+      opt.promoteBufferResultsToOutParams = promoteBufferResultsToOutParams;
 
       BufferizationOptions::OpFilterEntry::FilterFn filterFn =
           [&](Operation *op) {
             // Filter may be specified via options.
             if (this->dialectFilter.hasValue())
               return llvm::find(this->dialectFilter,
                                 op->getDialect()->getNamespace()) !=
                      this->dialectFilter.end();

 
 LogicalResult
 bufferization::finalizeBuffers(Operation *op,
                                const BufferizationOptions &options) {
   // Hoist buffers.
   if (failed(hoistBufferAllocations(op, options)))
     return failure();
 
-  // Deallocate buffers that escape block boundaries ("leaking buffers") with
-  // the buffer deallocation pass.
-  bool hasLeakingAlloc = false;
+  // Create allocation ops for "leaking buffers", i.e., buffer allocations that
+  // escape block boundaries. If there are no leaking allocs, `hasLeakingAllocs`
+  // is set to `false`.
+  bool hasLeakingAllocs = false;
   if (failed(createAllocDeallocOps(op, options, /*onlyLeakingAllocs=*/true,
-                                   &hasLeakingAlloc)))
+                                   &hasLeakingAllocs)))
     return failure();
-  if (options.createDeallocs && hasLeakingAlloc &&
-      failed(deallocateBuffers(op)))
+
+  if (hasLeakingAllocs) {
+    // Promote returned buffers to "out" parameters.
+    // TODO: Pass options to support custom dealloc ops.
+    if (options.promoteBufferResultsToOutParams && isa<ModuleOp>(op) &&
+        failed(promoteBufferResultsToOutParams(cast<ModuleOp>(op))))
+      return failure();
+
+    // Create deallocation ops for all "leaking buffers" and all buffer
+    // allocations that were added during the above promotion process.
+    // TODO: Pass options to support custom dealloc ops.
+    if (options.createDeallocs && failed(deallocateBuffers(op)))
       return failure();
+  }
 
-  // Deallocate all remaining buffers at the end of the block.
+  // Deallocate all remaining buffers at the end of their parent blocks.
   if (failed(createAllocDeallocOps(op, options)))
     return failure();
 
   return success();
 }
 
 LogicalResult bufferization::bufferizeOp(Operation *op,
                                          const AnalysisState &analysisState) {

mlir/test/Dialect/Bufferization/Transforms/one-shot-module-bufferize-out-params.mlir

+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs promote-buffer-results-to-out-params" -split-input-file | FileCheck %s
+
+// Note: This bufferization is not very efficient yet, but it works.
+
+// CHECK: #[[$map1:.*]] = affine_map<(d0)[s0, s1] -> (d0 * s1 + s0)>
+// CHECK-LABEL: func @callee(
+//  CHECK-SAME:              %[[arg0:.*]]: memref<5xf32, #[[$map1]]>,
+//  CHECK-SAME:              %[[arg1:.*]]: memref<5xf32>) {
+//       CHECK:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
+//       CHECK:   memref.copy %[[arg0]], %[[alloc]]
+//       CHECK:   memref.store %{{.*}}, %[[alloc]]
+//       CHECK:   memref.copy %[[alloc]], %[[arg1]]
+//       CHECK:   memref.dealloc %[[alloc]]
+//       CHECK:   return
+//       CHECK: }
+func.func @callee(%t: tensor<5xf32>) -> (tensor<5xf32>, tensor<5xf32>) {
+  %c0 = arith.constant 0 : index
+  %cst = arith.constant 8.0 : f32
+  %1 = tensor.insert %cst into %t[%c0] : tensor<5xf32>
+  return %t, %1 : tensor<5xf32>, tensor<5xf32>
+}
+
+// CHECK: func @main(%[[arg0:.*]]: memref<5xf32, #[[$map1]]>) -> (f32, f32) {
+// CHECK:   %[[alloc:.*]] = memref.alloc() : memref<5xf32>
+// CHECK:   call @callee(%[[arg0]], %[[alloc]])
+// CHECK:   %[[l1:.*]] = memref.load %[[arg0]]
+// CHECK:   %[[l2:.*]] = memref.load %[[alloc]]
+// CHECK:   memref.dealloc %[[alloc]]
+// CHECK:   return %[[l1]], %[[l2]]
+// CHECK: }
+func.func @main(%t: tensor<5xf32>) -> (f32, f32) {
+  %c0 = arith.constant 0 : index
+  %0, %1 = func.call @callee(%t)
+      : (tensor<5xf32>) -> (tensor<5xf32>, tensor<5xf32>)
+  %2 = tensor.extract %0[%c0] : tensor<5xf32>
+  %3 = tensor.extract %1[%c0] : tensor<5xf32>
+  return %2, %3 : f32, f32
+}
+