[mlir][bufferize] Decouple promoteBufferResultsToOutParams from One-Shot Bufferize

[springerm]

Summary

Users should explicitly run -buffer-results-to-out-params instead.

The purpose of this change is to remove finalizeBuffers, which made it difficult to extend the bufferization to custom buffer types.

Depends On D126252

[springerm]

update

[springerm]

update

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

struct BufferizationOptions {
   /// * FullyDynamicLayoutMap: All function parameter types and result types
   ///   have a fully dynamic layout map. This option is most efficient because
   ///   any layout map can be casted to a fully dynamic one.
   /// * IdentityLayoutMap: All function parameter types and result types have a
   ///   static identity layout (i.e., no layout map). This option may introduce
   ///   additional buffer allocs and copies because layout maps cannot be casted
   ///   away.
   ///
-  /// If `bufferizeFunctionBoundaries` is not set, this flag has no effect. If
-  /// `promoteBufferResultsToOutParams` is set, `kInferMostPreciseLayoutMap` is
-  /// is an invalid option.
+  /// If `bufferizeFunctionBoundaries` is not set, this flag has no effect.
   ///
   /// Note: Inferred layout maps may not be desireable when interacting with
   /// external functions, because the generated function signatures will be less
   /// predictable.
   LayoutMapOption functionBoundaryTypeConversion =
       LayoutMapOption::InferLayoutMap;
 
   /// This flag controls buffer types on unknown ops (to_memref wrappers) and in
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/Bufferize.h

 
 /// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.
 /// Reuse an existing `BufferizationState`.
 ///
 /// Note: This function overload is useful for extending the bufferization.
 LogicalResult bufferizeOp(Operation *op, BufferizationState &bufferizationState,
                           const OpFilter *opFilter = nullptr);
 
-/// Finalize all buffer allocations: Create alloc/dealloc ops as specified by
-/// the bufferization options.
-LogicalResult finalizeBuffers(Operation *op,
-                              const BufferizationOptions &options);
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZE_H

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

let description = [{
     and supports only simple cases at the moment. In particular:
 
     * Recursive or circular function call graphs are not supported.
     * When a returned tensor can be proven to be equivalent to a tensor function
       argument, the return value disappears. Instead, the buffer of the tensor
       argument is modified in-place.
     * Returning non-equivalent tensors is forbidden by default and must be
       explicitly activated with `allow-return-allocs`.
-    * Non-equivalent returned tensors of fully static size can be promoted to
-      function arguments with `promote-buffer-results-to-out-params`. In that
-      case, buffers for such tensors are allocated at each call site. Instead of
-      returning a buffer, the buffer contents are copied into these allocations.
     * External functions (without bodies) that return a tensor are not
       supported.
     * Function with multiple blocks or multiple ReturnOps are not supported.
     * Layout maps on function signatures can be controlled with a separate
       `function-boundary-type-conversion` option, which is similar to
       `unknown-type-conversion` but supports an additional `infer-layout-map`
       option. `fully-dynamic-layout-map` and `identity-layout-map` ensure that
       function signatures bufferize to easily predictable types, potentially at
Option<"functionBoundaryTypeConversion",
            "Controls layout maps when bufferizing function signatures.">,
     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.">,
     Option<"unknownTypeConversion", "unknown-type-conversion", "std::string",
            /*default=*/"\"fully-dynamic-layout-map\"",
            "Controls layout maps for non-inferrable memref types.">,
   ];
   let constructor = "mlir::bufferization::createOneShotBufferizePass()";
 }
 
 def PromoteBuffersToStack : Pass<"promote-buffers-to-stack", "func::FuncOp"> {

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

if (!options) {
       opt.alwaysAliasingWithDest = alwaysAliasingWithDest;
       opt.analysisFuzzerSeed = analysisFuzzerSeed;
       opt.createDeallocs = createDeallocs;
       opt.functionBoundaryTypeConversion =
           parseLayoutMapOption(functionBoundaryTypeConversion);
       opt.printConflicts = printConflicts;
       opt.testAnalysisOnly = testAnalysisOnly;
       opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
-      opt.promoteBufferResultsToOutParams = promoteBufferResultsToOutParams;
       opt.unknownTypeConversion = parseLayoutMapOption(unknownTypeConversion);
 
       OpFilter::Entry::FilterFn filterFn =
           [&](Operation *op) {
             // Filter may be specified via options.
             if (this->dialectFilter.hasValue())
               return llvm::is_contained(this->dialectFilter,
                                         op->getDialect()->getNamespace());
if (auto funcOp = dyn_cast<FunctionOpInterface>(op)) {
     return hasTensorArg || hasTensorResult;
   }
 
   bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
   bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
   return hasTensorResult || hasTensorOperand;
 }
 
-LogicalResult
-bufferization::finalizeBuffers(Operation *op,
-                               const BufferizationOptions &options) {
-  // 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();
-
-  return success();
-}
-
 LogicalResult bufferization::bufferizeOp(Operation *op,
                                          const AnalysisState &analysisState) {
   // Catch incorrect API usage.
   assert((analysisState.hasDialectState(
               func::FuncDialect::getDialectNamespace()) ||
           !analysisState.getOptions().bufferizeFunctionBoundaries) &&
          "must use ModuleBufferize to bufferize function boundaries");
 
   BufferizationState bufferizationState(analysisState);
   if (failed(bufferizeOp(op, bufferizationState)))
     return failure();
-  if (failed(finalizeBuffers(op, analysisState.getOptions())))
-    return failure();
   return success();
 }
 
 namespace {
 /// A rewriter that keeps track of extra information during bufferization.
 class BufferizationRewriter : public IRRewriter {
 public:
   BufferizationRewriter(MLIRContext *ctx, DenseSet<Operation *> &erasedOps,

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

if (!options.allowReturnAllocs &&
         llvm::any_of(funcOp.getFunctionType().getResults(), [](Type t) {
           return t.isa<MemRefType, UnrankedMemRefType>();
         })) {
       funcOp->emitError("memref return type is unsupported");
       return failure();
     }
   }
 
-  // Finalize all buffers.
-  if (failed(finalizeBuffers(moduleOp, options)))
-    return failure();
-
   // Post-pass cleanup of function argument attributes.
   moduleOp.walk([&](func::FuncOp op) {
     for (BlockArgument bbArg : op.getArguments())
       removeBufferizationAttributes(bbArg);
   });
 
   return success();
 }

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 function-boundary-type-conversion=fully-dynamic-layout-map" -buffer-deallocation -split-input-file | FileCheck %s
-// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs promote-buffer-results-to-out-params function-boundary-type-conversion=identity-layout-map" -buffer-deallocation -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs function-boundary-type-conversion=fully-dynamic-layout-map" -buffer-results-to-out-params -buffer-deallocation -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs function-boundary-type-conversion=identity-layout-map" -buffer-results-to-out-params -buffer-deallocation -split-input-file | FileCheck %s --check-prefix=CHECK-NO-LAYOUT
 // RUN: mlir-opt %s -one-shot-bufferize="bufferize-function-boundaries allow-return-allocs function-boundary-type-conversion=infer-layout-map" -buffer-deallocation -split-input-file | FileCheck %s --check-prefix=CHECK-BASELINE
 
 // Note: function-boundary-type-conversion=infer-layout-map with
 // promote-buffer-results-to-out-params is an unsupported combination.
 
 // 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, #[[$map1]]>) {
 // This alloc is not needed, but it is inserted due to the out-of-place
 // bufferization of the tensor.insert. With a better layering of the out param
 // promotion pass, this alloc could be avoided.
 //       CHECK:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
 //       CHECK:   memref.copy %[[arg0]], %[[alloc]]
 //       CHECK:   memref.store %{{.*}}, %[[alloc]]
-//       CHECK:   memref.copy %[[alloc]], %[[arg1]]
+//       CHECK:   %[[casted:.*]] = memref.cast %[[alloc]]
+//       CHECK:   memref.copy %[[casted]], %[[arg1]]
 //       CHECK:   memref.dealloc %[[alloc]]
 //       CHECK:   return
 //       CHECK: }
 
 // CHECK-NO-LAYOUT-LABEL: func @callee(
 //  CHECK-NO-LAYOUT-SAME:     %[[arg0:.*]]: memref<5xf32>,
 //  CHECK-NO-LAYOUT-SAME:     %[[arg1:.*]]: memref<5xf32>) {
 //       CHECK-NO-LAYOUT:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<5xf32>
func.func @callee(%t: tensor<5xf32>) -> (tensor<5xf32>, 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:   %[[casted:.*]] = memref.cast %[[alloc]] : memref<5xf32> to memref<5xf32, #[[$map1]]>
 // CHECK:   call @callee(%[[arg0]], %[[casted]])
 // CHECK:   %[[l1:.*]] = memref.load %[[arg0]]
-// CHECK:   %[[l2:.*]] = memref.load %[[alloc]]
+// CHECK:   %[[l2:.*]] = memref.load %[[casted]]
 // CHECK:   memref.dealloc %[[alloc]]
 // CHECK:   return %[[l1]], %[[l2]]
 // CHECK: }
 
 // CHECK-NO-LAYOUT-LABEL: func @main(%{{.*}}: memref<5xf32>) -> (f32, f32) {
 //       CHECK-NO-LAYOUT:   %[[alloc:.*]] = memref.alloc() : memref<5xf32>
 //       CHECK-NO-LAYOUT:   call @callee(%{{.*}}, %[[alloc]])
 func.func @main(%t: tensor<5xf32>) -> (f32, f32) {

 
 // CHECK: #[[$map2a:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
 // CHECK: #[[$map2b:.*]] = affine_map<(d0, d1)[s0] -> (d0 * 20 + s0 + d1)>
 // CHECK-LABEL: func @callee(
 //  CHECK-SAME:     %{{.*}}: index,
 //  CHECK-SAME:     %[[r:.*]]: memref<2x5xf32, #[[$map2a]]>) {
 //       CHECK:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<10x20xf32>
 //       CHECK:   %[[subview:.*]] = memref.subview %[[alloc]]{{.*}} : memref<10x20xf32> to memref<2x5xf32, #[[$map2b]]>
-//       CHECK:   memref.copy %[[subview]], %[[r]]
+//       CHECK:   %[[casted:.*]] = memref.cast %[[subview]]
+//       CHECK:   memref.copy %[[casted]], %[[r]]
 //       CHECK:   memref.dealloc %[[alloc]]
 
 // CHECK-NO-LAYOUT-LABEL: func @callee(
 //  CHECK-NO-LAYOUT-SAME:              %{{.*}}: index,
 //  CHECK-NO-LAYOUT-SAME:              %[[r:.*]]: memref<2x5xf32>) {
 //       CHECK-NO-LAYOUT:   %[[alloc:.*]] = memref.alloc() {{.*}} : memref<10x20xf32>
 //       CHECK-NO-LAYOUT:   %[[subview:.*]] = memref.subview %[[alloc]]
 // Note: This alloc is not needed, but it is inserted before the returned buffer
func.func @callee(%idx: index) -> tensor<2x5xf32> {
   %1 = tensor.extract_slice %0[%idx, %idx][2, 5][1, 1] : tensor<10x20xf32> to tensor<2x5xf32>
   return %1 : tensor<2x5xf32>
 }
 
 // CHECK: func @main(
 // CHECK:   %[[alloc:.*]] = memref.alloc() : memref<2x5xf32>
 // CHECK:   %[[casted:.*]] = memref.cast %[[alloc]] : memref<2x5xf32> to memref<2x5xf32, #[[$map2a]]>
 // CHECK:   call @callee(%{{.*}}, %[[casted]])
-// CHECK:   memref.load %[[alloc]]
+// CHECK:   memref.load %[[casted]]
 // CHECK:   memref.dealloc %[[alloc]]
 
 // CHECK-NO-LAYOUT: func @main(
 // CHECK-NO-LAYOUT:   %[[alloc:.*]] = memref.alloc() : memref<2x5xf32>
 // CHECK-NO-LAYOUT:   call @callee(%{{.*}}, %[[alloc]])
 // CHECK-NO-LAYOUT:   memref.load %[[alloc]]
 // CHECK-NO-LAYOUT:   memref.dealloc