[mlir][linalg][bufferize][NFC] Refactor BufferizationOption ownership

[springerm]

Summary

Pass unique_ptr<BufferizationOption> to the bufferization. This allows the bufferization to enqueue additional PostAnalysisSteps. When running bufferization a second time, a new BufferizationOptions must be constructed.

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

public:
   const BufferizationOptions &getOptions() const { return options; }
 
   /// Return a reference to the OpBuilder.
   OpBuilder &getBuilder() { return builder; }
 
 private:
   friend LogicalResult
   runComprehensiveBufferize(Operation *op, const BufferizationOptions &options,
-                            BufferizationState &state,
-                            const PostAnalysisStepList &extraSteps);
+                            BufferizationState &state);
 
   friend LogicalResult
   runComprehensiveBufferize(ModuleOp moduleOp,
-                            const BufferizationOptions &options);
+                            std::unique_ptr<BufferizationOptions> options);
 
   /// `aliasInfo` keeps track of aliasing and equivalent values. Only internal
   /// functions and `runComprehensiveBufferize` may access this object.
   BufferizationAliasInfo aliasInfo;
 
   /// Dialect-specific bufferization state.
   DenseMap<StringRef, std::unique_ptr<DialectBufferizationState>> dialectState;
 

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

 
 namespace mlir {
 
 namespace linalg {
 namespace comprehensive_bufferize {
 
 struct BufferizationOptions;
 class BufferizationState;
-struct PostAnalysisStep;
 
 /// Bufferize the given operation. Reuses an existing BufferizationState object.
-LogicalResult runComprehensiveBufferize(
-    Operation *op, const BufferizationOptions &options,
-    BufferizationState &state,
-    const std::vector<std::unique_ptr<PostAnalysisStep>> &extraSteps);
+/// This function overload is for internal usage only.
+LogicalResult runComprehensiveBufferize(Operation *op,
+                                        const BufferizationOptions &options,
+                                        BufferizationState &state);
 
 /// Bufferize the given operation.
-LogicalResult runComprehensiveBufferize(Operation *op,
-                                        const BufferizationOptions &options);
+LogicalResult
+runComprehensiveBufferize(Operation *op,
+                          std::unique_ptr<BufferizationOptions> options);
 
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVE_BUFFERIZE_H

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

 //===- ModuleBufferization.h - 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_MODULE_BUFFERIZATION_H
 #define MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_MODULE_BUFFERIZATION_H
 
+#include <memory>
+
 namespace mlir {
 
 class DialectRegistry;
 struct LogicalResult;
 class ModuleOp;
 
 namespace linalg {
 namespace comprehensive_bufferize {
 
 struct BufferizationOptions;
 
 /// Bufferize the given module. This bufferizations performs a simple function
 /// call analysis to determine which function arguments are inplaceable.
-LogicalResult runComprehensiveBufferize(ModuleOp moduleOp,
-                                        const BufferizationOptions &options);
+LogicalResult
+runComprehensiveBufferize(ModuleOp moduleOp,
+                          std::unique_ptr<BufferizationOptions> options);
 
 namespace std_ext {
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace std_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_MODULE_BUFFERIZATION_H

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

annotateOpsWithBufferizationMarkers(Operation *op,
   op->walk([&](Operation *op) {
     for (OpResult opResult : op->getResults())
       if (opResult.getType().isa<TensorType>())
         setInPlaceOpResult(opResult, aliasInfo.isInPlace(opResult));
   });
 }
 
 LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
-    Operation *op, const BufferizationOptions &options) {
-  BufferizationState state(op, options);
-  PostAnalysisStepList extraSteps;
-  return runComprehensiveBufferize(op, options, state, extraSteps);
+    Operation *op, std::unique_ptr<BufferizationOptions> options) {
+  BufferizationState state(op, *options);
+  return runComprehensiveBufferize(op, *options, state);
 }
 
 LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
     Operation *op, const BufferizationOptions &options,
-    BufferizationState &state, const PostAnalysisStepList &extraSteps) {
+    BufferizationState &state) {
 
   DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.aliasInfo;
 
   if (failed(checkAliasInfoConsistency(op, domInfo, state, aliasInfo)))
     return failure();
 
   // If the analysis fails, just return.
   if (failed(inPlaceAnalysis(op, aliasInfo, state, domInfo,
                              options.analysisFuzzerSeed)))
     return failure();
   equivalenceAnalysis(op, aliasInfo, state);
 
-  auto runPostAnalysisSteps = [&](const PostAnalysisStepList &steps) {
-    for (const std::unique_ptr<PostAnalysisStep> &step : steps) {
-      SmallVector<Operation *> newOps;
-      if (failed(step->run(op, state, aliasInfo, newOps)))
-        return failure();
-      // Analyze ops that were created by the PostAnalysisStep.
-      if (failed(inPlaceAnalysis(newOps, aliasInfo, state, domInfo)))
-        return failure();
-      equivalenceAnalysis(newOps, aliasInfo, state);
-    }
-    return success();
-  };
-
-  if (failed(runPostAnalysisSteps(extraSteps)))
-    return failure();
-  if (failed(runPostAnalysisSteps(options.postAnalysisSteps)))
-    return failure();
+  for (const std::unique_ptr<PostAnalysisStep> &step :
+       options.postAnalysisSteps) {
 
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly) {
     annotateOpsWithBufferizationMarkers(op, aliasInfo);
     return success();
   }
 
   // Bufferize the op and its nested ops.
   if (failed(bufferize(op, state)))
     return failure();
 
   return success();
 }

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

static void annotateOpsWithBufferizationMarkers(FuncOp funcOp,
                                                 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, const BufferizationOptions &options) {
-  BufferizationState state(moduleOp, options);
+    ModuleOp moduleOp, std::unique_ptr<BufferizationOptions> options) {
+  BufferizationState state(moduleOp, *options);
   ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
   BufferizationAliasInfo &aliasInfo = state.aliasInfo;
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, moduleState.orderedFuncOps,
                                       moduleState.callerMap)))
     return failure();
 
   // Interestingly, all function args that are not visible outside of a module
   // can be fully bufferized inplace by guaranteeing the CallOp is bufferized
   // inplace. Therefore, we just bufferize funcOp as if none of its results were
   // inplaceable, detect which operands are cloned internally and decide what to
   // do at call sites.
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.body().empty())
       continue;
 
-    // Register extra post analysis steps. These cannot be stored in `options`
-    // because `options` is immutable.
-    PostAnalysisStepList extraSteps;
-    extraSteps.emplace_back(std::make_unique<EquivalentFuncOpBBArgsAnalysis>());
+    // Collect bbArg/return value information after the analysis.
+    options->postAnalysisSteps.emplace_back(
+        std::make_unique<EquivalentFuncOpBBArgsAnalysis>());
 
     // Gather equivalence info for CallOps.
     equivalenceAnalysis(funcOp, aliasInfo, moduleState);
 
     // Analyze and bufferize funcOp.
-    if (failed(runComprehensiveBufferize(funcOp, options, state, extraSteps)))
+    if (failed(runComprehensiveBufferize(funcOp, *options, state)))
       return failure();
 
     // Add annotations to function arguments.
-    if (options.testAnalysisOnly)
+    if (options->testAnalysisOnly)
       annotateOpsWithBufferizationMarkers(funcOp, state);
   }
 
-  if (options.testAnalysisOnly)
+  if (options->testAnalysisOnly)
     return success();
 
   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, state)))
       return failure();
 
-    if (!options.allowReturnMemref &&
+    if (!options->allowReturnMemref &&
         llvm::any_of(funcOp.getType().getResults(), [](Type t) {
           return t.isa<MemRefType, UnrankedMemRefType>();
         })) {
       funcOp->emitError("memref return type is unsupported");
       return failure();
     }
   }
 

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

static Optional<Value> allocationFnUsingAlloca(OpBuilder &b, Location loc,
                                                MemRefType type,
                                                ArrayRef<Value> dynShape) {
   Value allocated = b.create<memref::AllocaOp>(
       loc, type, dynShape, b.getI64IntegerAttr(kBufferAlignments));
   return allocated;
 }
 
 void LinalgComprehensiveModuleBufferize::runOnOperation() {
-  BufferizationOptions options;
+  auto options = std::make_unique<BufferizationOptions>();
   if (useAlloca) {
-    options.allocationFns->allocationFn = allocationFnUsingAlloca;
-    options.allocationFns->deallocationFn = [](OpBuilder &b, Location loc,
+    options->allocationFns->allocationFn = allocationFnUsingAlloca;
+    options->allocationFns->deallocationFn = [](OpBuilder &b, Location loc,
                                                 Value v) {};
   }
   // TODO: Change to memref::CopyOp (default memCpyFn).
-  options.allocationFns->memCpyFn = [](OpBuilder &b, Location loc, Value from,
+  options->allocationFns->memCpyFn = [](OpBuilder &b, Location loc, Value from,
                                         Value to) {
     b.create<linalg::CopyOp>(loc, from, to);
   };
 
-  options.allowReturnMemref = allowReturnMemref;
-  options.allowUnknownOps = allowUnknownOps;
-  options.analysisFuzzerSeed = analysisFuzzerSeed;
-  options.testAnalysisOnly = testAnalysisOnly;
-  options.printConflicts = printConflicts;
+  options->allowReturnMemref = allowReturnMemref;
+  options->allowUnknownOps = allowUnknownOps;
+  options->analysisFuzzerSeed = analysisFuzzerSeed;
+  options->testAnalysisOnly = testAnalysisOnly;
+  options->printConflicts = printConflicts;
 
   // Enable InitTensorOp elimination.
-  options.addPostAnalysisStep<
+  options->addPostAnalysisStep<
       linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
   // TODO: Find a way to enable this step automatically when bufferizing tensor
   // dialect ops.
-  options.addPostAnalysisStep<tensor_ext::InplaceInsertSliceOpAnalysis>();
+  options->addPostAnalysisStep<tensor_ext::InplaceInsertSliceOpAnalysis>();
   if (!allowReturnMemref)
-    options.addPostAnalysisStep<scf_ext::AssertDestinationPassingStyle>();
+    options->addPostAnalysisStep<scf_ext::AssertDestinationPassingStyle>();
 
   ModuleOp moduleOp = getOperation();
   applyEnablingTransformations(moduleOp);
 
-  if (failed(runComprehensiveBufferize(moduleOp, options))) {
+  if (failed(runComprehensiveBufferize(moduleOp, std::move(options)))) {
     signalPassFailure();
     return;
   }
 
-  if (options.testAnalysisOnly)
+  if (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>();
 }

mlir/test/lib/Dialect/Linalg/TestComprehensiveBufferize.cpp

struct TestComprehensiveFunctionBufferize
   ListOption<std::string> dialectFilter{
       *this, "dialect-filter",
       llvm::cl::desc("Bufferize only ops from the specified dialects"),
       llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
 };
 } // namespace
 
 void TestComprehensiveFunctionBufferize::runOnFunction() {
-  BufferizationOptions options;
+  auto options = std::make_unique<BufferizationOptions>();
 
   // Enable InitTensorOp elimination.
-  options.addPostAnalysisStep<
+  options->addPostAnalysisStep<
       linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
   // TODO: Find a way to enable this step automatically when bufferizing
   // tensor dialect ops.
-  options.addPostAnalysisStep<tensor_ext::InplaceInsertSliceOpAnalysis>();
+  options->addPostAnalysisStep<tensor_ext::InplaceInsertSliceOpAnalysis>();
   if (!allowReturnMemref)
-    options.addPostAnalysisStep<scf_ext::AssertDestinationPassingStyle>();
+    options->addPostAnalysisStep<scf_ext::AssertDestinationPassingStyle>();
 
-  options.allowReturnMemref = allowReturnMemref;
-  options.allowUnknownOps = allowUnknownOps;
-  options.testAnalysisOnly = testAnalysisOnly;
-  options.analysisFuzzerSeed = analysisFuzzerSeed;
+  options->allowReturnMemref = allowReturnMemref;
+  options->allowUnknownOps = allowUnknownOps;
+  options->testAnalysisOnly = testAnalysisOnly;
+  options->analysisFuzzerSeed = analysisFuzzerSeed;
 
   if (dialectFilter.hasValue()) {
-    options.dialectFilter.emplace();
+    options->dialectFilter.emplace();
     for (const std::string &dialectNamespace : dialectFilter)
-      options.dialectFilter->insert(dialectNamespace);
+      options->dialectFilter->insert(dialectNamespace);
   }
 
   Operation *op = getFunction().getOperation();
-  if (failed(runComprehensiveBufferize(op, options)))
+  if (failed(runComprehensiveBufferize(op, std::move(options))))
     return;
 
   OpPassManager cleanupPipeline("builtin.func");
   cleanupPipeline.addPass(createCanonicalizerPass());
   cleanupPipeline.addPass(createCSEPass());
   cleanupPipeline.addPass(createLoopInvariantCodeMotionPass());
   (void)this->runPipeline(cleanupPipeline, op);
 }
 
 namespace mlir {
 namespace test {
 void registerTestComprehensiveFunctionBufferize() {
   PassRegistration<TestComprehensiveFunctionBufferize>();
 }
 } // namespace test
 } // namespace mlir