[mlir][bufferize][NFC] Make PostAnalysisSteps a function

[springerm]

Summary

They used to be classes with a virtual run function. This was inconvenient because post analysis steps are stored in BufferizationOptions. Because of this design choice, BufferizationOptions were not copyable.

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

struct BufferizationOptions {
   using DeallocationFn =
       std::function<LogicalResult(OpBuilder &, Location, Value)>;
   /// Memcpy function: Generate a memcpy between two buffers.
   using MemCpyFn =
       std::function<LogicalResult(OpBuilder &, Location, Value, Value)>;
 
   BufferizationOptions();
 
-  // BufferizationOptions cannot be copied.
-  BufferizationOptions(const BufferizationOptions &other) = delete;
-
   /// Return `true` if the op is allowed to be bufferized.
   bool isOpAllowed(Operation *op) const {
     if (!hasFilter)
       return true;
     return dialectFilter.contains(op->getDialect()->getNamespace()) ||
            operationFilter.contains(op->getName().getStringRef());
   }
 

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

 /// can be used to implement partial bufferization passes.
 LogicalResult bufferizeOp(Operation *op, const BufferizationOptions &options);
 
 /// Populate the pattern set with a pattern that bufferizes ops that implement
 /// `BufferizableOpInterface`.
 void populateBufferizationPattern(const BufferizationState &state,
                                   RewritePatternSet &patterns);
 
-std::unique_ptr<BufferizationOptions> getPartialBufferizationOptions();
+BufferizationOptions getPartialBufferizationOptions();
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZE_H

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

 
 namespace mlir {
 namespace bufferization {
 
 class AnalysisBufferizationState;
 class BufferizationAliasInfo;
 struct AnalysisBufferizationOptions;
 
-/// PostAnalysisSteps can be registered with `BufferizationOptions` and are
+/// PostAnalysisStepFns can be registered with `BufferizationOptions` and are
 /// executed after the analysis, but before bufferization. They can be used to
-/// implement custom dialect-specific optimizations.
-struct PostAnalysisStep {
-  virtual ~PostAnalysisStep() = default;
-
-  /// Run the post analysis step. This function may modify the IR, but must keep
-  /// `aliasInfo` consistent. Newly created operations and operations that
-  /// should be re-analyzed must be added to `newOps`.
-  virtual LogicalResult run(Operation *op, BufferizationState &state,
-                            BufferizationAliasInfo &aliasInfo,
-                            SmallVector<Operation *> &newOps) = 0;
-};
+/// implement custom dialect-specific optimizations. They may modify the IR, but
+/// must keep `aliasInfo` consistent. Newly created operations and operations
+/// that should be re-analyzed must be added to `newOps`.
+using PostAnalysisStepFn = std::function<LogicalResult(
+    Operation *, BufferizationState &, BufferizationAliasInfo &,
+    SmallVector<Operation *> &)>;
 
-using PostAnalysisStepList = std::vector<std::unique_ptr<PostAnalysisStep>>;
+using PostAnalysisStepList = SmallVector<PostAnalysisStepFn>;
 
 /// Options for analysis-enabled bufferization.
 struct AnalysisBufferizationOptions : public BufferizationOptions {
   AnalysisBufferizationOptions() = default;
 
-  // AnalysisBufferizationOptions cannot be copied.
-  AnalysisBufferizationOptions(const AnalysisBufferizationOptions &) = delete;
-
   /// Register a "post analysis" step. Such steps are executed after the
   /// analysis, but before bufferization.
-  template <typename Step, typename... Args>
-  void addPostAnalysisStep(Args... args) {
-    postAnalysisSteps.emplace_back(
-        std::make_unique<Step>(std::forward<Args>(args)...));
+  void addPostAnalysisStep(PostAnalysisStepFn fn) {
+    postAnalysisSteps.push_back(fn);
   }
 
   /// Registered post analysis steps.
   PostAnalysisStepList postAnalysisSteps;
 };
 
 /// The BufferizationAliasInfo class maintains a list of buffer aliases and
 /// equivalence classes to support bufferization.

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

 
 namespace mlir {
 class DialectRegistry;
 
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace linalg_ext {
 
-struct InitTensorEliminationStep : public bufferization::PostAnalysisStep {
-  /// A function that matches anchor OpOperands for InitTensorOp elimination.
-  /// If an OpOperand is matched, the function should populate the SmallVector
-  /// with all values that are needed during `RewriteFn` to produce the
-  /// replacement value.
-  using AnchorMatchFn = std::function<bool(OpOperand &, SmallVector<Value> &)>;
 
-  /// A function that rewrites matched anchors.
-  using RewriteFn = std::function<Value(OpBuilder &, Location, OpOperand &)>;
 
-  /// Try to eliminate InitTensorOps inside `op`.
-  ///
-  /// * `rewriteFunc` generates the replacement for the InitTensorOp.
-  /// * Only InitTensorOps that are anchored on a matching OpOperand as per
-  ///   `anchorMatchFunc` are considered. "Anchored" means that there is a path
-  ///   on the reverse SSA use-def chain, starting from the OpOperand and always
-  ///   following the aliasing  OpOperand, that eventually ends at a single
-  ///   InitTensorOp.
-  /// * The result of `rewriteFunc` must usually be analyzed for inplacability.
-  ///   This analysis can be skipped with `skipAnalysis`.
-  LogicalResult
-  eliminateInitTensors(Operation *op, bufferization::BufferizationState &state,
-                       bufferization::BufferizationAliasInfo &aliasInfo,
-                       AnchorMatchFn anchorMatchFunc, RewriteFn rewriteFunc,
-                       SmallVector<Operation *> &newOps);
-};
 
 /// Try to eliminate InitTensorOps inside `op` that are anchored on an
 /// InsertSliceOp, i.e., if it is eventually inserted into another tensor
 /// (and some other conditions are met).
-struct InsertSliceAnchoredInitTensorEliminationStep
-    : public InitTensorEliminationStep {
-  LogicalResult run(Operation *op, bufferization::BufferizationState &state,
-                    bufferization::BufferizationAliasInfo &aliasInfo,
-                    SmallVector<Operation *> &newOps) override;
-};
+LogicalResult insertSliceAnchoredInitTensorEliminationStep(
+    Operation *op, bufferization::BufferizationState &state,
+    SmallVector<Operation *> &newOps);
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace linalg_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_LINALGINTERFACEIMPL_H

mlir/include/mlir/Dialect/SCF/BufferizableOpInterfaceImpl.h

 //===- BufferizableOpInterfaceImpl.h - Impl. of BufferizableOpInterface ---===//
 //
 // 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_SCF_BUFFERIZABLEOPINTERFACEIMPL_H
 #define MLIR_DIALECT_SCF_BUFFERIZABLEOPINTERFACEIMPL_H
 
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 
 namespace mlir {
 class DialectRegistry;
 
+namespace bufferization {
+class BufferizationState;
+class BufferizationAliasInfo;
+} // namespace bufferization
+
 namespace scf {
 /// Assert that yielded values of an scf.for op are aliasing their corresponding
 /// bbArgs. This is required because the i-th OpResult of an scf.for op is
 /// currently assumed to alias with the i-th iter_arg (in the absence of
 /// conflicts).
-struct AssertScfForAliasingProperties : public bufferization::PostAnalysisStep {
-  LogicalResult run(Operation *op, bufferization::BufferizationState &state,
+LogicalResult
+assertScfForAliasingProperties(Operation *op,
+                               bufferization::BufferizationState &state,
                                bufferization::BufferizationAliasInfo &aliasInfo,
-                    SmallVector<Operation *> &newOps) override;
-};
+                               SmallVector<Operation *> &newOps);
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 } // namespace scf
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SCF_BUFFERIZABLEOPINTERFACEIMPL_H

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

struct ArithmeticBufferizePass
     : public ArithmeticBufferizeBase<ArithmeticBufferizePass> {
   ArithmeticBufferizePass(uint64_t alignment = 0, bool constantOpOnly = false)
       : ArithmeticBufferizeBase<ArithmeticBufferizePass>(),
         constantOpOnly(constantOpOnly) {
     this->alignment = alignment;
   }
 
   void runOnOperation() override {
-    std::unique_ptr<BufferizationOptions> options =
-        getPartialBufferizationOptions();
+    BufferizationOptions options = getPartialBufferizationOptions();
     if (constantOpOnly) {
-      options->addToOperationFilter<arith::ConstantOp>();
+      options.addToOperationFilter<arith::ConstantOp>();
     } else {
-      options->addToDialectFilter<arith::ArithmeticDialect>();
+      options.addToDialectFilter<arith::ArithmeticDialect>();
     }
-    options->bufferAlignment = alignment;
+    options.bufferAlignment = alignment;
 
-    if (failed(bufferizeOp(getOperation(), *options)))
+    if (failed(bufferizeOp(getOperation(), options)))
       signalPassFailure();
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<bufferization::BufferizationDialect, memref::MemRefDialect,
                     arith::ArithmeticDialect>();
     arith::registerBufferizableOpInterfaceExternalModels(registry);
   }

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

LogicalResult bufferization::bufferizeOp(Operation *op,
   return bufferizeOp(op, state);
 }
 
 void bufferization::populateBufferizationPattern(
     const BufferizationState &state, RewritePatternSet &patterns) {
   patterns.add<BufferizationPattern>(patterns.getContext(), state);
 }
 
-std::unique_ptr<BufferizationOptions>
-bufferization::getPartialBufferizationOptions() {
-  auto options = std::make_unique<BufferizationOptions>();
-  options->allowReturnMemref = true;
-  options->allowUnknownOps = true;
-  options->createDeallocs = false;
-  options->fullyDynamicLayoutMaps = false;
+BufferizationOptions bufferization::getPartialBufferizationOptions() {
+  BufferizationOptions options;
+  options.allowReturnMemref = true;
+  options.allowUnknownOps = true;
+  options.createDeallocs = false;
+  options.fullyDynamicLayoutMaps = false;
   return options;
 }

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

 /// In the above example, the first scf.yield op satifies destination-passing
 /// style because the yielded value %0 is defined in the parent block. The
 /// second scf.yield op does not satisfy destination-passing style because the
 /// yielded value %t is defined in the same block as the scf.yield op.
 // TODO: The current implementation checks for equivalent values instead of
 // aliasing values, which is stricter than needed. We can currently not check
 // for aliasing values because the analysis is a maybe-alias analysis and we
 // need a must-alias analysis here.
-struct AssertDestinationPassingStyle : public PostAnalysisStep {
-  LogicalResult run(Operation *op, BufferizationState &state,
+static LogicalResult
+assertDestinationPassingStyle(Operation *op, BufferizationState &state,
                               BufferizationAliasInfo &aliasInfo,
-                    SmallVector<Operation *> &newOps) override {
-    LogicalResult status = success();
-    DominanceInfo domInfo(op);
-    op->walk([&](Operation *returnOp) {
-      if (!isRegionReturnLike(returnOp))
-        return WalkResult::advance();
+                              SmallVector<Operation *> &newOps) {
 
-      for (OpOperand &returnValOperand : returnOp->getOpOperands()) {
-        Value returnVal = returnValOperand.get();
-        // Skip non-tensor values.
-        if (!returnVal.getType().isa<TensorType>())
-          continue;
 
-        bool foundEquivValue = false;
-        aliasInfo.applyOnEquivalenceClass(returnVal, [&](Value equivVal) {
-          if (auto bbArg = equivVal.dyn_cast<BlockArgument>()) {
-            Operation *definingOp = bbArg.getOwner()->getParentOp();
-            if (definingOp->isProperAncestor(returnOp))
-              foundEquivValue = true;
-            return;
-          }
 
-          Operation *definingOp = equivVal.getDefiningOp();
-          if (definingOp->getBlock()->findAncestorOpInBlock(
-                  *returnOp->getParentOp()))
-            // Skip ops that happen after `returnOp` and parent ops.
-            if (happensBefore(definingOp, returnOp, domInfo))
-              foundEquivValue = true;
-        });
 
-        if (!foundEquivValue)
-          status =
-              returnOp->emitError()
-              << "operand #" << returnValOperand.getOperandNumber()
-              << " of ReturnLike op does not satisfy destination passing style";
-      }
 
-      return WalkResult::advance();
-    });
 
-    return status;
-  }
-};
 
 LogicalResult bufferization::analyzeOp(Operation *op,
                                        AnalysisBufferizationState &state) {
   DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
   const auto &options =
       static_cast<const AnalysisBufferizationOptions &>(state.getOptions());
 
   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);
 
-  for (const std::unique_ptr<PostAnalysisStep> &step :
-       options.postAnalysisSteps) {
+  for (const PostAnalysisStepFn &fn : options.postAnalysisSteps) {
     SmallVector<Operation *> newOps;
-    if (failed(step->run(op, state, aliasInfo, newOps)))
+    if (failed(fn(op, state, aliasInfo, newOps)))
       return failure();
-    // Analyze ops that were created by the PostAnalysisStep.
+    // Analyze ops that were created by the PostAnalysisStepFn.
     if (failed(inPlaceAnalysis(newOps, aliasInfo, state, domInfo)))
       return failure();
     equivalenceAnalysis(newOps, aliasInfo, state);
   }
 
   if (!options.allowReturnMemref) {
     SmallVector<Operation *> newOps;
-    if (failed(
-            AssertDestinationPassingStyle().run(op, state, aliasInfo, newOps)))
+    if (failed(assertDestinationPassingStyle(op, state, aliasInfo, newOps)))
       return failure();
   }
 
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly)
     annotateOpsWithBufferizationMarkers(op, aliasInfo, state);
 
   return success();

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

 }
 
 /// Try to eliminate InitTensorOps inside `op`. An InitTensorOp is replaced
 /// with the the result of `rewriteFunc` if it is anchored on a matching
 /// OpOperand. "Anchored" means that there is a path on the reverse SSA use-def
 /// chain, starting from the OpOperand and always following the aliasing
 /// OpOperand, that eventually ends at a single InitTensorOp.
 LogicalResult
-mlir::linalg::comprehensive_bufferize::linalg_ext::InitTensorEliminationStep::
-    eliminateInitTensors(Operation *op, BufferizationState &state,
-                         BufferizationAliasInfo &aliasInfo,
-                         AnchorMatchFn anchorMatchFunc, RewriteFn rewriteFunc,
-                         SmallVector<Operation *> &newOps) {
+mlir::linalg::comprehensive_bufferize::linalg_ext::eliminateInitTensors(
+    Operation *op, BufferizationState &state, BufferizationAliasInfo &aliasInfo,
   OpBuilder b(op->getContext());
 
   WalkResult status = op->walk([&](Operation *op) {
     for (OpOperand &operand : op->getOpOperands()) {
       // Skip operands that do not bufferize inplace.
       if (!aliasInfo.isInPlace(operand))
         continue;
       // All values that are needed to create the replacement op.

 /// * On the reverse use-def chain path from the InsertSliceOp to the
 ///   InitTensorOp, all ops were decided to bufferize inplace and the buffer
 ///   relation is "equivalent" (TODO: can be relaxed if needed).
 /// * The reverse use-def chain has exactly one end, which is the InitTensorOp.
 ///
 /// Note that the newly inserted ExtractSliceOp may have to bufferize
 /// out-of-place due to RaW conflicts.
 LogicalResult mlir::linalg::comprehensive_bufferize::linalg_ext::
-    InsertSliceAnchoredInitTensorEliminationStep::run(
+    insertSliceAnchoredInitTensorEliminationStep(
         Operation *op, BufferizationState &state,
         BufferizationAliasInfo &aliasInfo, SmallVector<Operation *> &newOps) {
   return eliminateInitTensors(
       op, state, aliasInfo,
       /*anchorMatchFunc=*/
       [&](OpOperand &operand, SmallVector<Value> &neededValues) {
         auto insertSliceOp =
             dyn_cast<tensor::InsertSliceOp>(operand.getOwner());

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

 // 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.
 //
 // After analyzing a FuncOp, additional information about its bbArgs is
-// gathered through PostAnalysisSteps and stored in `ModuleBufferizationState`.
+// gathered through PostAnalysisStepFns and stored in
+// `ModuleBufferizationState`.
 //
-// * `EquivalentFuncOpBBArgsAnalysis` determines the equivalent bbArg for each
+// * `equivalentFuncOpBBArgsAnalysis` determines the equivalent bbArg for each
 //   tensor return value (if any).
-// * `FuncOpBbArgReadWriteAnalysis` determines whether or not a tensor bbArg is
+// * `funcOpBbArgReadWriteAnalysis` determines whether or not a tensor bbArg is
 //   read/written.
 //
 // Only tensors that are equivalent to some FuncOp bbArg may be returned.
 // Bufferization currently fails if other tensors (in particular tensors that
 // bufferize out-of-place and result in a new buffer allocation) are returned.
 // In the future, such allocations could be hoisted to the caller.
 //
 // Example: `foo` fails bufferization because %0 is not equivalent to any bbArg.

 //   be written to in-place.
 // * If a tensor operand of a CallOp is read after the CallOp, the operand of
 //   the CallOp must bufferize out-of-place.
 //
 // Example: The tensor.insert op bufferizes in-place because it is allowed to
 // modify the buffer of `%t1` directly. The CallOp in `caller` must bufferize
 // out-of-place because `%t0` is modified by the callee but read by the
 // tensor.extract op. The analysis of CallOps decides whether an OpOperand must
-// bufferize out-of-place based on results of `FuncOpBbArgReadWriteAnalysis`.
+// bufferize out-of-place based on results of `funcOpBbArgReadWriteAnalysis`.
 // ```
 // func @callee(%t1 : tensor<?xf32>) -> tensor<?xf32> {
 //   %f = ... : f32
 //   %0 = tensor.insert %f into %t1[...] : tensor<?xf32>
 //   return %0 : tensor<?xf32>
 // }
 //
 // func @caller() -> () {
 //   %t0 = ... : tensor<?xf32>
 //   %1 = call @callee(%t0) : (tensor<?xf32>) -> (tensor<?xf32>)
 //   %2 = tensor.extract %1[...]  : tensor<?xf32>
 // }
 // ```
 //
-// Note: If a function is external, `FuncOpBbArgReadWriteAnalysis` cannot
+// Note: If a function is external, `funcOpBbArgReadWriteAnalysis` cannot
 // analyze the function body. In such a case, the CallOp analysis conservatively
 // assumes that each tensor OpOperand is both read and written.
 //
 // TODO: Add FuncOp attributes so that bbArgs of external FuncOps can be marked
 // as "not reading" and/or "not writing".
 
 #include "mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h"
 
if (auto candidateOp = dyn_cast<ReturnOp>(b.getTerminator())) {
         return nullptr;
       returnOp = candidateOp;
     }
   }
   return returnOp;
 }
 
 namespace {
-/// Store function BlockArguments that are equivalent to a returned value in
-/// ModuleBufferizationState.
-struct EquivalentFuncOpBBArgsAnalysis : public PostAnalysisStep {
-  /// Annotate IR with the results of the analysis. For testing purposes only.
-  static void annotateReturnOp(OpOperand &returnVal, BlockArgument bbArg) {
-    const char *kEquivalentArgsAttr = "__equivalent_func_args__";
-    Operation *op = returnVal.getOwner();
+
+static void annotateEquivalentReturnBbArg(OpOperand &returnVal,
+                                          BlockArgument bbArg) {
 
-    SmallVector<int64_t> equivBbArgs;
-    if (op->hasAttr(kEquivalentArgsAttr)) {
-      auto attr = op->getAttr(kEquivalentArgsAttr).cast<ArrayAttr>();
-      equivBbArgs = llvm::to_vector<4>(llvm::map_range(attr, [](Attribute a) {
-        return a.cast<IntegerAttr>().getValue().getSExtValue();
-      }));
-    } else {
-      equivBbArgs.append(op->getNumOperands(), -1);
-    }
-    equivBbArgs[returnVal.getOperandNumber()] = bbArg.getArgNumber();
 
-    OpBuilder b(op->getContext());
-    op->setAttr(kEquivalentArgsAttr, b.getI64ArrayAttr(equivBbArgs));
-  }
 
-  LogicalResult run(Operation *op, BufferizationState &state,
+/// Store function BlockArguments that are equivalent to a returned value in
+/// ModuleBufferizationState.
+static LogicalResult
+equivalentFuncOpBBArgsAnalysis(Operation *op, BufferizationState &state,
                                BufferizationAliasInfo &aliasInfo,
-                    SmallVector<Operation *> &newOps) override {
-    ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
+                               SmallVector<Operation *> &newOps) {
 
-    // Support only single return-terminated block in the function.
-    auto funcOp = cast<FuncOp>(op);
-    ReturnOp returnOp = getAssumedUniqueReturnOp(funcOp);
-    assert(returnOp && "expected func with single return op");
 
-    for (OpOperand &returnVal : returnOp->getOpOperands())
-      if (returnVal.get().getType().isa<RankedTensorType>())
-        for (BlockArgument bbArg : funcOp.getArguments())
-          if (bbArg.getType().isa<RankedTensorType>())
-            if (aliasInfo.areEquivalentBufferizedValues(returnVal.get(),
-                                                        bbArg)) {
-              moduleState
-                  .equivalentFuncArgs[funcOp][returnVal.getOperandNumber()] =
-                  bbArg.getArgNumber();
-              if (state.getOptions().testAnalysisOnly)
-                annotateReturnOp(returnVal, bbArg);
-            }
+          if (aliasInfo.areEquivalentBufferizedValues(returnVal.get(), bbArg)) {
+              annotateEquivalentReturnBbArg(returnVal, bbArg);
 
-    return success();
-  }
-};
 
 /// Return true if the buffer of the given tensor value is written to. Must not
 /// be called for values inside not yet analyzed functions. (Post-analysis
 /// steps do not have to be run yet, i.e., "in progress" is also OK.)
 static bool isValueWritten(Value value, const BufferizationState &state,
                            const BufferizationAliasInfo &aliasInfo) {
 #ifndef NDEBUG
   assert(value.getType().isa<TensorType>() && "expected TensorType");
aliasInfo.applyOnAliases(value, [&](Value val) {
     for (OpOperand &use : val.getUses())
       if (state.isInPlace(use) && state.bufferizesToMemoryWrite(use))
         isWritten = true;
   });
   return isWritten;
 }
 
 /// Determine which FuncOp bbArgs are read and which are written. If this
-/// PostAnalysisStep is run on a function with unknown ops, it will
+/// PostAnalysisStepFn is run on a function with unknown ops, it will
 /// conservatively assume that such ops bufferize to a read + write.
-struct FuncOpBbArgReadWriteAnalysis : public PostAnalysisStep {
-  LogicalResult run(Operation *op, BufferizationState &state,
+static LogicalResult
+funcOpBbArgReadWriteAnalysis(Operation *op, BufferizationState &state,
                              BufferizationAliasInfo &aliasInfo,
-                    SmallVector<Operation *> &newOps) override {
-    ModuleBufferizationState &moduleState = getModuleBufferizationState(state);
-    auto funcOp = cast<FuncOp>(op);
+                             SmallVector<Operation *> &newOps) {
 
-    // If the function has no body, conservatively assume that all args are
-    // read + written.
-    if (funcOp.getBody().empty()) {
-      for (BlockArgument bbArg : funcOp.getArguments()) {
-        moduleState.readBbArgs.insert(bbArg);
-        moduleState.writtenBbArgs.insert(bbArg);
-      }
 
-      return success();
-    }
 
-    for (BlockArgument bbArg : funcOp.getArguments()) {
-      if (!bbArg.getType().isa<TensorType>())
-        continue;
-      if (state.isValueRead(bbArg))
-        moduleState.readBbArgs.insert(bbArg);
-      if (isValueWritten(bbArg, state, aliasInfo))
-        moduleState.writtenBbArgs.insert(bbArg);
-    }
 
-    return success();
-  }
-};
 } // namespace
 
 static bool isaTensor(Type t) { return t.isa<TensorType>(); }
 
 /// If `value` is a memref::CastOp, return its source. Otherwise, return
 /// `value` directly.
 static Value getNonCastedValue(Value value) {
   while (auto castOp = value.getDefiningOp<memref::CastOp>())
LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
   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>());
+  options->postAnalysisSteps.push_back(equivalentFuncOpBBArgsAnalysis);
+  options->postAnalysisSteps.push_back(funcOpBbArgReadWriteAnalysis);
 
   // Analyze ops.
   for (FuncOp funcOp : moduleState.orderedFuncOps) {
     // No body => no analysis.
     if (funcOp.body().empty())
       continue;
 
     // Now analyzing function.

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

void LinalgComprehensiveModuleBufferize::runOnOperation() {
   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>();
+    options->addPostAnalysisStep(
+        linalg_ext::insertSliceAnchoredInitTensorEliminationStep);
   }
 
   // Only certain scf.for ops are supported by the analysis.
-  options->addPostAnalysisStep<scf::AssertScfForAliasingProperties>();
+  options->addPostAnalysisStep(scf::assertScfForAliasingProperties);
 
   ModuleOp moduleOp = getOperation();
   applyEnablingTransformations(moduleOp);
 
   if (failed(runComprehensiveBufferize(moduleOp, std::move(options)))) {
     signalPassFailure();
     return;
   }

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
     return success();
   }
 };
 
 } // namespace
 } // namespace scf
 } // namespace mlir
 
-LogicalResult mlir::scf::AssertScfForAliasingProperties::run(
+LogicalResult mlir::scf::assertScfForAliasingProperties(
     Operation *op, BufferizationState &state, BufferizationAliasInfo &aliasInfo,
     SmallVector<Operation *> &newOps) {
   LogicalResult status = success();
 
   op->walk([&](scf::ForOp forOp) {
     auto yieldOp =
         cast<scf::YieldOp>(forOp.getLoopBody().front().getTerminator());
     for (OpOperand &operand : yieldOp->getOpOperands()) {

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

 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace bufferization;
 
 namespace {
 struct TensorBufferizePass : public TensorBufferizeBase<TensorBufferizePass> {
   void runOnOperation() override {
-    std::unique_ptr<BufferizationOptions> options =
-        getPartialBufferizationOptions();
-    options->addToDialectFilter<tensor::TensorDialect>();
+    BufferizationOptions options = getPartialBufferizationOptions();
+    options.addToDialectFilter<tensor::TensorDialect>();
 
-    if (failed(bufferizeOp(getOperation(), *options)))
+    if (failed(bufferizeOp(getOperation(), options)))
       signalPassFailure();
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<bufferization::BufferizationDialect, memref::MemRefDialect,
                     tensor::TensorDialect, scf::SCFDialect,
                     arith::ArithmeticDialect>();
     tensor::registerBufferizableOpInterfaceExternalModels(registry);
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createTensorBufferizePass() {
   return std::make_unique<TensorBufferizePass>();
 }

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

Option<bool> createDeallocs{
       llvm::cl::init(true)};
 };
 } // namespace
 
 void TestComprehensiveFunctionBufferize::runOnOperation() {
   auto options = std::make_unique<AnalysisBufferizationOptions>();
 
   if (!allowReturnMemref)
-    options->addPostAnalysisStep<scf::AssertScfForAliasingProperties>();
+    options->addPostAnalysisStep(scf::assertScfForAliasingProperties);
 
   options->allowReturnMemref = allowReturnMemref;
   options->allowUnknownOps = allowUnknownOps;
   options->testAnalysisOnly = testAnalysisOnly;
   options->analysisFuzzerSeed = analysisFuzzerSeed;
   options->fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
   options->createDeallocs = createDeallocs;