[mlir][linalg][bufferize][NFC] Allow passing custom BufferizationOptions to pass

[springerm]

Summary

Also remove the pass option for linalg.generic-based memcpys. Users of ModuleBufferization can configure the bufferization as needed.

Depends On D119258

[springerm]

rebase

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h

 } // namespace bufferization
 
 namespace linalg {
 namespace comprehensive_bufferize {
 
 /// Run Module Bufferization on the given module. Performs a simple function
 /// call analysis to determine which function arguments are inplaceable. Then
 /// analyzes and bufferizes FuncOps one-by-one with One-Shot Bufferize.
-LogicalResult runComprehensiveBufferize(
+LogicalResult runModuleBufferize(
     ModuleOp moduleOp,
     std::unique_ptr<bufferization::AnalysisBufferizationOptions> options);
 
+/// Run Module Bufferization on the given module. Performs a simple function
+/// call analysis to determine which function arguments are inplaceable. Then
+/// analyzes and bufferizes FuncOps one-by-one with One-Shot Bufferize.
+///
+/// Note: This bufferization options that are passed to this function overload
+/// must have been prepared with `prepareOptions`.
+LogicalResult runModuleBufferize(
+    ModuleOp moduleOp,
+    const bufferization::AnalysisBufferizationOptions &preparedOptions);
+
+/// Prepare bufferization options. This enqueues all PostAnalysisSteps that are
+/// required for Module Bufferization.
+void prepareOptions(bufferization::AnalysisBufferizationOptions &options);
+
 namespace std_ext {
 
 void registerModuleBufferizationExternalModels(DialectRegistry &registry);
 
 } // namespace std_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_MODULEBUFFERIZATION_H

mlir/include/mlir/Dialect/Linalg/Passes.h

 #ifndef MLIR_DIALECT_LINALG_PASSES_H_
 #define MLIR_DIALECT_LINALG_PASSES_H_
 
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
+namespace bufferization {
+struct AnalysisBufferizationOptions;
+} // namespace bufferization
 
 std::unique_ptr<Pass> createConvertElementwiseToLinalgPass();
 
 std::unique_ptr<Pass> createLinalgFoldUnitExtentDimsPass();
 
 std::unique_ptr<Pass> createLinalgElementwiseOpFusionPass();
 std::unique_ptr<Pass> createFoldReshapeOpsByLinearizationPass();
 

 /// This pass implements a cross-dialect bufferization approach and performs an
 /// analysis to determine which op operands and results may be bufferized in the
 /// same buffers. The analysis is performed on topologically sorted CallOp and
 /// FuncOp within a module. It provides analyses and bufferization across
 /// function boundaries. Within a function boundary, the analysis is performed
 /// on SSA use-def chains starting from function operands that are annotated
 /// with the 'inplaceable' attribute.
 std::unique_ptr<Pass> createLinalgComprehensiveModuleBufferizePass();
-std::unique_ptr<Pass>
-createLinalgComprehensiveModuleBufferizePass(bool useLinalgCopy);
+std::unique_ptr<Pass> createLinalgComprehensiveModuleBufferizePass(
+    std::unique_ptr<bufferization::AnalysisBufferizationOptions> options);
 
 /// Create a pass to convert Linalg operations which work on tensors to use
 /// buffers instead.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgBufferizePass();
 
 /// Create a pass to convert named Linalg operations to Linalg generic
 /// operations.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgGeneralizationPass();

mlir/include/mlir/Dialect/Linalg/Passes.td

Option<"allowReturnMemref", "allow-return-memref", "bool",
             /*default=*/"false",
            "Allows the return of memrefs (for testing purposes only)">,
     Option<"allowUnknownOps", "allow-unknown-ops", "bool",
            /*default=*/"false",
            "Allows unknown (not bufferizable) ops in the input IR.">,
     Option<"useAlloca", "use-alloca", "bool",
            /*default=*/"false",
            "Use stack allocations for memrefs (for testing purposes only)">,
-    Option<"useLinalgCopy", "use-memref.copy", "bool",
-           /*default=*/"false",
-           "Use a copy operation implemented as a Linalg op.">,
     Option<"fullyDynamicLayoutMaps", "fully-dynamic-layout-maps", "bool",
            /*default=*/"true",
            "Generate MemRef types with dynamic offset+strides by default.">,
     Option<"analysisFuzzerSeed", "analysis-fuzzer-seed", "unsigned",
            /*default=*/"0",
            "Analyze ops in random order with a given seed (fuzzer)">,
     Option<"initTensorElimination", "init-tensor-elimination", "bool",
             /*default=*/"false",

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

Value makeComposedPadHighOp(OpBuilder &b, Location loc, RankedTensorType type,
                             Value source, Value pad, bool nofold);
 
 /// Returns a GenericOp that tansposes `inputTensor` into `outputTensor` using
 /// `transposeVector` to permute the `inputTensor` dimensions.
 GenericOp makeTransposeOp(OpBuilder &b, Location loc, Value inputTensor,
                           Value outputTensor,
                           ArrayRef<int64_t> transposeVector);
 
+/// Returns GenericOp that copies an n-D memref. Unlike the current
+/// implementation of memref::CopyOp, this op can further tile, lower to loops
+/// or vectorize.
+GenericOp makeMemRefCopyOp(OpBuilder &b, Location loc, Value from, Value to);
+
 //===----------------------------------------------------------------------===//
 // Fusion / Tiling utilities
 //===----------------------------------------------------------------------===//
 
 /// The type of loops to be generated during tiling.
 enum class LinalgTilingLoopType {
   Loops = 0,
   AffineLoops = 1,

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp

 //===- ModuleBufferization.cpp - Bufferization across Func. Boundaries ----===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Module Bufferization is an extension of Comprehensive Bufferize that
 // bufferizes function boundaries. It provides `BufferizableOpInterface`
 // implementations for FuncOp, CallOp and ReturnOp.
 //
-// Module Bufferization is run via `runComprehensiveBufferize(ModuleOp, ...)`.
-// This function analyzed the given module and determines the order of
-// analysis and bufferization: Functions that are called are processed before
-// their respective callers.
+// Module Bufferization is run via `runModuleBufferize(ModuleOp, ...)`. This
+// function analyzes the given module and determines the order of analysis and
+// bufferization: Functions that are called are processed before their
+// respective callers.
 //
 // After analyzing a FuncOp, additional information about its bbArgs is
 // gathered through PostAnalysisSteps and stored in `ModuleBufferizationState`.
 //
 // * `EquivalentFuncOpBBArgsAnalysis` determines the equivalent bbArg for each
 //   tensor return value (if any).
 // * `FuncOpBbArgReadWriteAnalysis` determines whether or not a tensor bbArg is
 //   read/written.

 annotateOpsWithBufferizationMarkers(FuncOp funcOp,
                                     const BufferizationState &state) {
   auto bufferizableOp = cast<BufferizableOpInterface>(funcOp.getOperation());
   for (BlockArgument bbArg : funcOp.getArguments())
     if (bbArg.getType().isa<TensorType>())
       setInPlaceFuncArgument(bbArg, bufferizableOp.isWritable(bbArg, state));
 }
 
-LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
-    ModuleOp moduleOp, std::unique_ptr<AnalysisBufferizationOptions> options) {
+void mlir::linalg::comprehensive_bufferize::prepareOptions(
+    AnalysisBufferizationOptions &options) {
+  // Collect bbArg/return value information after the analysis.
+  options.postAnalysisSteps.emplace_back(
+      std::make_unique<EquivalentFuncOpBBArgsAnalysis>());
+  options.postAnalysisSteps.emplace_back(
+      std::make_unique<FuncOpBbArgReadWriteAnalysis>());
+}
+
+LogicalResult mlir::linalg::comprehensive_bufferize::runModuleBufferize(
+    ModuleOp moduleOp, const AnalysisBufferizationOptions &preparedOptions) {
   IRRewriter rewriter(moduleOp.getContext());
-  AnalysisBufferizationState state(moduleOp, *options);
+  AnalysisBufferizationState state(moduleOp, preparedOptions);
   ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
   BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, moduleState.orderedFuncOps,
                                       moduleState.callerMap)))
     return failure();
 
-  // Collect bbArg/return value information after the analysis.
-  options->postAnalysisSteps.emplace_back(
-      std::make_unique<EquivalentFuncOpBBArgsAnalysis>());
-  options->postAnalysisSteps.emplace_back(
-      std::make_unique<FuncOpBbArgReadWriteAnalysis>());
-
   // Analyze ops.
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.body().empty())
       continue;
 
     // Now analyzing function.
     moduleState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::InProgress;
 
     // Analyze funcOp.
     if (failed(analyzeOp(funcOp, state)))
       return failure();
 
     // Gather equivalence info for CallOps.
     // TODO: Make this a post-analysis step.
     equivalenceAnalysis(funcOp, aliasInfo, moduleState);
 
     // Mark op as fully analyzed.
     moduleState.analyzedFuncOps[funcOp] = FuncOpAnalysisState::Analyzed;
 
     // Add annotations to function arguments.
-    if (options->testAnalysisOnly)
+    if (preparedOptions.testAnalysisOnly)
       annotateOpsWithBufferizationMarkers(funcOp, state);
   }
 
-  if (options->testAnalysisOnly)
+  if (preparedOptions.testAnalysisOnly)
     return success();
 
   // Bufferize function bodies.
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.body().empty())
       continue;
 
     if (failed(bufferizeOp(funcOp, state)))
       return failure();
   }
 
   // Bufferize function boundaries.
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.
     if (failed(bufferizeFuncOpBoundary(funcOp, rewriter, state)))
       return failure();
 
-    if (!options->allowReturnMemref &&
+    if (!preparedOptions.allowReturnMemref &&
         llvm::any_of(funcOp.getType().getResults(), [](Type t) {
           return t.isa<MemRefType, UnrankedMemRefType>();
         })) {
       funcOp->emitError("memref return type is unsupported");
       return failure();
     }
   }
 
   // Perform a post-processing pass of layout modification at function boundary
   // according to the kBufferLayoutAttrName.
   layoutPostProcessing(moduleOp);
 
   // Post-pass cleanup of inplaceable and buffer_layout attributes.
   moduleOp.walk([&](FuncOp op) {
     for (BlockArgument bbArg : op.getArguments())
       removeBufferizationFuncArguments(bbArg);
   });
 
   return success();
 }
+
+LogicalResult mlir::linalg::comprehensive_bufferize::runModuleBufferize(
+    ModuleOp moduleOp, std::unique_ptr<AnalysisBufferizationOptions> options) {
+  prepareOptions(*options);
+  return runModuleBufferize(moduleOp, *options);
+}

mlir/lib/Dialect/Linalg/Transforms/ComprehensiveBufferizePass.cpp

 
 namespace {
 struct LinalgComprehensiveModuleBufferize
     : public LinalgComprehensiveModuleBufferizeBase<
           LinalgComprehensiveModuleBufferize> {
   LinalgComprehensiveModuleBufferize() = default;
 
   LinalgComprehensiveModuleBufferize(
-      const LinalgComprehensiveModuleBufferize &p) = default;
-
-  LinalgComprehensiveModuleBufferize(bool linalgCopy) {
-    this->useLinalgCopy = linalgCopy;
+      const LinalgComprehensiveModuleBufferize &p) {
+    llvm_unreachable("pass cannot be copied");
   }
 
+  explicit LinalgComprehensiveModuleBufferize(
+      std::unique_ptr<AnalysisBufferizationOptions> options);
+
   void runOnOperation() override;
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry
         .insert<bufferization::BufferizationDialect, linalg::LinalgDialect,
                 memref::MemRefDialect, tensor::TensorDialect,
                 vector::VectorDialect, scf::SCFDialect,
                 arith::ArithmeticDialect, StandardOpsDialect, AffineDialect>();
     affine_ext::registerBufferizableOpInterfaceExternalModels(registry);
     arith::registerBufferizableOpInterfaceExternalModels(registry);
     linalg_ext::registerBufferizableOpInterfaceExternalModels(registry);
     scf::registerBufferizableOpInterfaceExternalModels(registry);
     std_ext::registerModuleBufferizationExternalModels(registry);
     mlir::registerBufferizableOpInterfaceExternalModels(registry);
     tensor::registerBufferizableOpInterfaceExternalModels(registry);
     vector::registerBufferizableOpInterfaceExternalModels(registry);
   }
+
+private:
+  std::unique_ptr<AnalysisBufferizationOptions> options;
 };
 } // namespace
 
+LinalgComprehensiveModuleBufferize::LinalgComprehensiveModuleBufferize(
+    std::unique_ptr<AnalysisBufferizationOptions> options)
+    : options(std::move(options)) {
+  // Only certain scf.for ops are supported by the analysis.
+  this->options->addPostAnalysisStep<scf::AssertScfForAliasingProperties>();
+  prepareOptions(*this->options);
+}
+
 static void applyEnablingTransformations(ModuleOp moduleOp) {
   RewritePatternSet patterns(moduleOp.getContext());
   patterns.add<GeneralizePadOpPattern>(moduleOp.getContext());
   (void)applyPatternsAndFoldGreedily(moduleOp, std::move(patterns));
 }
 
 static FailureOr<Value> allocationFnUsingAlloca(OpBuilder &b, Location loc,
                                                 MemRefType type,
                                                 ValueRange dynShape,
                                                 unsigned int bufferAlignment) {
   Value allocated = b.create<memref::AllocaOp>(
       loc, type, dynShape, b.getI64IntegerAttr(bufferAlignment));
   return allocated;
 }
 
-/// Create a linalg::GenericOp version of an n-D copy that can further tile,
-/// lower to loops or vectorize, unlike the current implementation of
-/// memref::CopyOp.
-/// Do not depend on memref::CopyOp that is getting deprecated.
-static LogicalResult createLinalgCopyOp(OpBuilder &b, Location loc, Value from,
-                                        Value to) {
-  auto memrefTypeFrom = from.getType().cast<MemRefType>();
-  auto memrefTypeTo = to.getType().cast<MemRefType>();
-  if (!memrefTypeFrom || !memrefTypeTo ||
-      memrefTypeFrom.getRank() != memrefTypeTo.getRank())
-    return failure();
-  AffineMap id =
-      AffineMap::getMultiDimIdentityMap(memrefTypeTo.getRank(), b.getContext());
-  SmallVector<StringRef> iteratorTypes(memrefTypeTo.getRank(),
-                                       getParallelIteratorTypeName());
-  b.create<linalg::GenericOp>(loc,
-                              /*inputs=*/from,
-                              /*outputs=*/to,
-                              /*indexingMaps=*/llvm::makeArrayRef({id, id}),
-                              /*iteratorTypes=*/iteratorTypes,
-                              [](OpBuilder &b, Location loc, ValueRange args) {
-                                b.create<linalg::YieldOp>(loc, args.front());
-                              });
-  return success();
-}
-
 void LinalgComprehensiveModuleBufferize::runOnOperation() {
-  auto options = std::make_unique<AnalysisBufferizationOptions>();
+  if (!options) {
+    // Make new bufferization options if none were provided when creating the
+    // pass.
+    options = std::make_unique<AnalysisBufferizationOptions>();
     if (useAlloca) {
       options->allocationFn = allocationFnUsingAlloca;
       options->deallocationFn = [](OpBuilder &b, Location loc, Value v) {
         return success();
       };
     }
-  // TODO: atm memref::CopyOp can be 200x slower than linalg::GenericOp.
-  // Once this perf bug is fixed more systematically, we can revisit.
-  if (useLinalgCopy)
-    options->memCpyFn = createLinalgCopyOp;
-
     options->allowReturnMemref = allowReturnMemref;
     options->allowUnknownOps = allowUnknownOps;
     options->analysisFuzzerSeed = analysisFuzzerSeed;
     options->createDeallocs = createDeallocs;
     options->fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
     options->printConflicts = printConflicts;
     options->testAnalysisOnly = testAnalysisOnly;
-
-  // Enable InitTensorOp elimination.
     if (initTensorElimination) {
       options->addPostAnalysisStep<
           linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
     }
-
     // Only certain scf.for ops are supported by the analysis.
     options->addPostAnalysisStep<scf::AssertScfForAliasingProperties>();
+    prepareOptions(*options);
+  }
 
   ModuleOp moduleOp = getOperation();
   applyEnablingTransformations(moduleOp);
 
-  if (failed(runComprehensiveBufferize(moduleOp, std::move(options)))) {
+  if (failed(runModuleBufferize(moduleOp, *options))) {
     signalPassFailure();
     return;
   }
 
-  if (testAnalysisOnly)
+  if (options->testAnalysisOnly)
     return;
 
   OpPassManager cleanupPipeline("builtin.module");
   cleanupPipeline.addPass(createCanonicalizerPass());
   cleanupPipeline.addPass(createCSEPass());
   cleanupPipeline.addPass(createLoopInvariantCodeMotionPass());
   (void)runPipeline(cleanupPipeline, moduleOp);
 }
 
 std::unique_ptr<Pass> mlir::createLinalgComprehensiveModuleBufferizePass() {
   return std::make_unique<LinalgComprehensiveModuleBufferize>();
 }
 
-std::unique_ptr<Pass>
-mlir::createLinalgComprehensiveModuleBufferizePass(bool useLinalgCopy) {
-  return std::make_unique<LinalgComprehensiveModuleBufferize>(useLinalgCopy);
+std::unique_ptr<Pass> mlir::createLinalgComprehensiveModuleBufferizePass(
+    std::unique_ptr<AnalysisBufferizationOptions> options) {
+  return std::make_unique<LinalgComprehensiveModuleBufferize>(
+      std::move(options));
 }

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

GenericOp makeTransposeOp(OpBuilder &b, Location loc, Value inputTensor,
 
   // Create the body of the transpose operation.
   OpBuilder::InsertionGuard g(b);
   b.setInsertionPointToEnd(&body.front());
   b.create<YieldOp>(loc, transposeOp.getRegion().front().getArgument(0));
   return transposeOp;
 }
 
+GenericOp makeMemRefCopyOp(OpBuilder &b, Location loc, Value from, Value to) {
+  auto memrefTypeTo = to.getType().cast<MemRefType>();
+#ifndef NDEBUG
+  auto memrefTypeFrom = from.getType().cast<MemRefType>();
+  assert(memrefTypeFrom.getRank() == memrefTypeTo.getRank() &&
+         "`from` and `to` memref must have the same rank");
+#endif // NDEBUG
+
+  AffineMap id =
+      AffineMap::getMultiDimIdentityMap(memrefTypeTo.getRank(), b.getContext());
+  SmallVector<StringRef> iteratorTypes(memrefTypeTo.getRank(),
+                                       getParallelIteratorTypeName());
+  return b.create<linalg::GenericOp>(
+      loc,
+      /*inputs=*/from,
+      /*outputs=*/to,
+      /*indexingMaps=*/llvm::makeArrayRef({id, id}),
+      /*iteratorTypes=*/iteratorTypes,
+      [](OpBuilder &b, Location loc, ValueRange args) {
+        b.create<linalg::YieldOp>(loc, args.front());
+      });
+}
+
 /// Specialization to build an scf "for" nest.
 template <>
 void GenerateLoopNest<scf::ForOp>::doit(
     OpBuilder &b, Location loc, ArrayRef<Range> loopRanges, LinalgOp linalgOp,
     ArrayRef<Attribute> iteratorTypes,
     function_ref<scf::ValueVector(OpBuilder &, Location, ValueRange,
                                   ValueRange)>
         bodyBuilderFn,