[mlir][linalg][bufferize] Add FuncOp bufferization pass

[springerm]

Summary

This passes bufferizes FuncOp bodies, but not FuncOp boundaries.

Depends On D114931

[springerm]

update

[springerm]

update

[springerm]

update

[springerm]

update

[nicolasvasilache]

Please turn this into a test pass we do not want to open the door to evolution towards multi-pass comprehensive bufferize.

[springerm]

update

[springerm]

update

[springerm]

rebase

[springerm]

address comments

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

 /// executed after the analysis, but before bufferization. They can be used
 /// implement custom dialect-specific optimizations.
 struct PostAnalysisStep {
   virtual ~PostAnalysisStep() {}
 
   /// Run the post analysis step. This function may modify the IR, but must keep
   /// `aliasInfo` (inside `state`) consistent. Newly created operations and
   /// operations that should be re-analyzed must be stored in `newOps`.
-  virtual LogicalResult run(FuncOp funcOp, BufferizationState &state,
+  virtual LogicalResult run(Operation *op, BufferizationState &state,
                             SmallVector<Operation *> &newOps) = 0;
 };
 
 /// Options for ComprehensiveBufferize.
 struct BufferizationOptions {
   BufferizationOptions();
 
   // BufferizationOptions cannot be copied.

 /// variants of this struct.
 struct DialectBufferizationState {
   virtual ~DialectBufferizationState() = default;
 };
 
 /// BufferizationState keeps track of bufferization state and provides access to
 /// the results of the analysis.
 struct BufferizationState {
-  BufferizationState(ModuleOp moduleOp, const BufferizationOptions &options)
-      : aliasInfo(moduleOp), options(options) {}
+  BufferizationState(Operation *op, const BufferizationOptions &options)
+      : aliasInfo(op), options(options) {}
 
   // BufferizationState should be passed as a reference.
   BufferizationState(const BufferizationState &) = delete;
 
   /// A function that creates an alloc-dealloc pair. This function may perform
   /// additional optimizations such as buffer allocation hoisting. This function
   /// calls `allocationFn` and `deallocationFn` to create (de)allocations.
   Value createAllocDeallocFn(OpBuilder &builder, Location loc,

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

 namespace mlir {
 
 namespace linalg {
 namespace comprehensive_bufferize {
 
 struct BufferizationOptions;
 struct BufferizationState;
 
+/// Bufferize the given op.
+LogicalResult runComprehensiveBufferize(Operation *op,
+                                        const BufferizationOptions &options);
+
 /// Bufferize the given function. Does not bufferize the function boundary.
-// TODO: This function is meant to be called from ModuleBufferize and not can
-// not yet be called standalone.
-LogicalResult runComprehensiveBufferize(FuncOp funcOp,
+/// Reuses an existing BufferizationState object.
+LogicalResult runComprehensiveBufferize(Operation *op,
                                         const BufferizationOptions &options,
                                         BufferizationState &state);
 
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_COMPREHENSIVE_BUFFERIZE_H

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

 namespace linalg {
 namespace comprehensive_bufferize {
 
 class BufferizationAliasInfo;
 
 namespace linalg_ext {
 
 struct InitTensorEliminationStep : public PostAnalysisStep {
-  /// Try to eliminate InitTensorOps inside `funcOp`.
+  /// 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(
-      FuncOp funcOp, BufferizationState &state,
+      Operation *op, BufferizationState &state,
       std::function<bool(OpOperand &)> anchorMatchFunc,
       std::function<Value(OpBuilder &, Location, OpOperand &)> rewriteFunc,
       SmallVector<Operation *> &newOps);
 };
 
-/// Try to eliminate InitTensorOps inside funcOp that are anchored on an
+/// 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(FuncOp funcOp, BufferizationState &state,
+  LogicalResult run(Operation *op, BufferizationState &state,
                     SmallVector<Operation *> &newOps) override;
 };
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace linalg_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_LINALG_INTERFACE_IMPL_H

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

 
 class DialectRegistry;
 
 namespace linalg {
 namespace comprehensive_bufferize {
 namespace tensor_ext {
 
 struct InplaceInsertSliceOpAnalysis : public PostAnalysisStep {
-  LogicalResult run(FuncOp funcOp, BufferizationState &state,
+  LogicalResult run(Operation *op, BufferizationState &state,
                     SmallVector<Operation *> &newOps) override;
 };
 
 void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry);
 
 } // namespace tensor_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 #endif // MLIR_DIALECT_LINALG_COMPREHENSIVEBUFFERIZE_TENSOR_INTERFACE_IMPL_H

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

 
 /// Create a pass to convert Linalg operations to affine.for loops and
 /// affine_load/affine_store accesses.
 /// Placeholder for now, this is NYI.
 std::unique_ptr<OperationPass<FuncOp>> createConvertLinalgToAffineLoopsPass();
 
 /// 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.
+/// same buffers.
+std::unique_ptr<Pass> createLinalgComprehensiveFunctionBufferizePass();
+
+/// This pass is an extension of the comprehensive function bufferization. It
+/// performs a function call analysis and bufferizes an entire module.
 std::unique_ptr<Pass> createLinalgComprehensiveModuleBufferizePass();
 
 /// 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.

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

let description = [{
     Convert ops with the `ElementwiseMappable` trait to linalg parallel loops.
 
     This pass only converts ops that operate on ranked tensors.
   }];
   let constructor = "mlir::createConvertElementwiseToLinalgPass()";
   let dependentDialects = ["linalg::LinalgDialect", "memref::MemRefDialect"];
 }
 
+def LinalgComprehensiveFunctionBufferize :
+    FunctionPass<"linalg-comprehensive-function-bufferize"> {
+  let summary = "Bufferize (tensor into memref) for a function.";
+  let description = [{
+    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 SSA use-def chains starting from
+    function operands that are annotated with the 'inplaceable' attribute.
+  }];
+  let options = [
+    Option<"testAnalysisOnly", "test-analysis-only", "bool",
+            /*default=*/"false",
+           "Only runs inplaceability analysis (for testing purposes only)">,
+    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<"analysisFuzzerSeed", "analysis-fuzzer-seed", "unsigned",
+           /*default=*/"0",
+           "Analyze ops in random order with a given seed (fuzzer)">
+  ];
+  let constructor = "mlir::createLinalgComprehensiveFunctionBufferizePass()";
+}
+
 def LinalgComprehensiveModuleBufferize :
     Pass<"linalg-comprehensive-module-bufferize", "ModuleOp"> {
   let summary = "Bufferize (tensor into memref) for a Module.";
   let description = [{
-    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
+    This pass is an extension of linalg-comprehensive-function-bufferize. It
+    performs an inter-function analysis 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.
+    function boundaries.
   }];
   let options = [
     Option<"testAnalysisOnly", "test-analysis-only", "bool",
             /*default=*/"false",
            "Only runs inplaceability analysis (for testing purposes only)">,
     Option<"allowReturnMemref", "allow-return-memref", "bool",
             /*default=*/"false",
            "Allows the return of memrefs (for testing purposes only)">,

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

 mlir::linalg::comprehensive_bufferize::bufferize(Operation *op,
                                                  BufferizationState &state) {
   OpBuilder b(op->getContext());
 
   // Check if op has tensor results or operands.
   auto isaTensor = [](Type t) { return t.isa<TensorType>(); };
   bool hasTensorResult = any_of(op->getResultTypes(), isaTensor);
   bool hasTensorOperand = any_of(op->getOperandTypes(), isaTensor);
+  bool isFuncOp = isa<FuncOp>(op);
 
   // No tensor results or operands: Simply bufferize all nested ops.
-  if (!hasTensorResult && !hasTensorOperand) {
+  // Always pass through FuncOps.
+  if (!hasTensorResult && !hasTensorOperand && !isFuncOp) {
     for (Region &region : op->getRegions())
       if (failed(bufferize(&region, state)))
         return failure();
     return success();
   }
 
   // Bufferize using `BufferizableOpInterface`. Interface implementations are
   // responsible for bufferizing nested ops.

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

for (OpOperand &opOperand : op->getOpOperands())
                                                        aliasInfo, domInfo)))
               return failure();
 
   return success();
 }
 
 /// Assert that the current bufferization decisions are consistent.
 static LogicalResult
-checkAliasInfoConsistency(FuncOp funcOp, const DominanceInfo &domInfo,
+checkAliasInfoConsistency(Operation *op, const DominanceInfo &domInfo,
                           const BufferizationAliasInfo &aliasInfo) {
   Operation *inconsistentOp = nullptr;
-  WalkResult walkResult = funcOp.walk([&](Operation *op) {
+  WalkResult walkResult = op->walk([&](Operation *op) {
     if (auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op))
       for (OpOperand &opOperand : op->getOpOperands())
         if (opOperand.get().getType().isa<TensorType>()) {
           OpResult opResult = bufferizableOp.getAliasingOpResult(opOperand);
           if (wouldCreateReadAfterWriteInterference(
                   opOperand, opResult, domInfo, aliasInfo,
                   /*checkConsistencyOnly=*/true)) {
             // This error can happen for two reasons. Either the input IR
for (OpResult opResult : op->getResults()) {
           if (bbArg.getType().isa<TensorType>())
             setInPlaceFuncArgument(bbArg,
                                    aliasInfo.bufferizesToWritableMemory(bbArg));
     }
   });
 }
 
 LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
-    FuncOp funcOp, const BufferizationOptions &options,
+    Operation *op, const BufferizationOptions &options) {
+  BufferizationState state(op, options);
+  return runComprehensiveBufferize(op, options, state);
+}
+
+LogicalResult mlir::linalg::comprehensive_bufferize::runComprehensiveBufferize(
+    Operation *op, const BufferizationOptions &options,
     BufferizationState &state) {
 
-  DominanceInfo domInfo(funcOp);
+  DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.aliasInfo;
 
-  if (funcOp.body().empty())
-    return success();
-
-  if (failed(checkAliasInfoConsistency(funcOp, domInfo, aliasInfo)))
+  if (failed(checkAliasInfoConsistency(op, domInfo, aliasInfo)))
     return failure();
 
   // Collect ops so we can build our own reverse traversal.
   SmallVector<Operation *> ops;
-  funcOp.walk([&](Operation *op) {
+  op->walk([&](Operation *op) {
     // No tensors => no buffers.
     if (none_of(op->getOperandTypes(), isaTensor) &&
         none_of(op->getResultTypes(), isaTensor))
       return;
     ops.push_back(op);
   });
 
   // If the analysis fails, just return.
   if (failed(
           inPlaceAnalysis(ops, aliasInfo, domInfo, options.analysisFuzzerSeed)))
     return failure();
 
   for (const std::unique_ptr<PostAnalysisStep> &step :
        options.postAnalysisSteps) {
     SmallVector<Operation *> newOps;
-    if (failed(step->run(funcOp, state, newOps)))
+    if (failed(step->run(op, state, newOps)))
       return failure();
     // Analyze ops that were created by the PostAnalysisStep.
     if (failed(inPlaceAnalysis(newOps, aliasInfo, domInfo)))
       return failure();
   }
 
   // Annotate operations if we only want to report the analysis.
   if (options.testAnalysisOnly) {
-    annotateOpsWithBufferizationMarkers(funcOp, aliasInfo);
+    annotateOpsWithBufferizationMarkers(op, aliasInfo);
     return success();
   }
 
-  // Bufferize all ops in funcOp.
-  OpBuilder b(funcOp.getContext());
-  auto bufferizableOp =
-      dyn_cast<BufferizableOpInterface>(funcOp.getOperation());
-  assert(bufferizableOp && "must use ModuleBufferization");
-  if (failed(bufferizableOp.bufferize(b, state)))
+  // Bufferize op and all nested ops.
+  if (failed(bufferize(op, state)))
     return failure();
 
   // Erase all obsolete ops.
   state.eraseObsoleteOps();
 
   return success();
 }

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

 
 /// Try to eliminate InitTensorOps inside funcOp. 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(
-        FuncOp funcOp, BufferizationState &state,
+        Operation *op, BufferizationState &state,
         std::function<bool(OpOperand &)> anchorMatchFunc,
         std::function<Value(OpBuilder &, Location, OpOperand &)> rewriteFunc,
         SmallVector<Operation *> &newOps) {
-  OpBuilder b(funcOp->getContext());
+  OpBuilder b(op->getContext());
   BufferizationAliasInfo &aliasInfo = state.aliasInfo;
 
-  WalkResult status = funcOp->walk([&](Operation *op) {
+  WalkResult status = op->walk([&](Operation *op) {
     for (OpOperand &operand : op->getOpOperands()) {
       // Is this a matching OpOperand?
       if (!anchorMatchFunc(operand))
         continue;
 
       SetVector<Value> maybeInitTensor =
           findValueInReverseUseDefChain(operand.get(), [&](Value val) {
             // Continue traversal until this function returns true.

 ///   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(
-        FuncOp funcOp, BufferizationState &state,
+        Operation *op, BufferizationState &state,
         SmallVector<Operation *> &newOps) {
   return eliminateInitTensors(
-      funcOp, state,
+      op, state,
       [&](OpOperand &operand) {
         auto insertSliceOp =
             dyn_cast<tensor::InsertSliceOp>(operand.getOwner());
         if (!insertSliceOp)
           return false;
         // Only inplace bufferized InsertSliceOps are eligible.
         if (!state.aliasInfo.isInPlace(insertSliceOp->getOpResult(0)))
           return false;

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

 };
 
 } // namespace tensor_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 LogicalResult mlir::linalg::comprehensive_bufferize::tensor_ext::
-    InplaceInsertSliceOpAnalysis::run(FuncOp funcOp, BufferizationState &state,
+    InplaceInsertSliceOpAnalysis::run(Operation *op, BufferizationState &state,
                                       SmallVector<Operation *> &newOps) {
   auto &tensorState = getTensorBufferizationState(state);
-  funcOp.walk([&](InsertSliceOp insertSliceOp) {
+  op->walk([&](InsertSliceOp insertSliceOp) {
     // A copy of the source buffer is needed if either:
     //   - The producer of `source` is not inplace. This is the case where a
     //     slice is computed out of place into the inplace full tensor.
     //   - The result is not inplace. This is the case where the whole tensor is
     //     cloned and the clone needs to be updated.
     if (isSourceEquivalentToAMatchingInplaceExtractSliceOp(state.aliasInfo,
                                                            insertSliceOp) &&
         state.aliasInfo.isInPlace(insertSliceOp->getResult(0)))

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

 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 using namespace mlir::linalg::comprehensive_bufferize;
 
-namespace {
-struct LinalgComprehensiveModuleBufferize
-    : public LinalgComprehensiveModuleBufferizeBase<
-          LinalgComprehensiveModuleBufferize> {
-  LinalgComprehensiveModuleBufferize() {}
-
-  LinalgComprehensiveModuleBufferize(
-      const LinalgComprehensiveModuleBufferize &p) {}
-
-  void runOnOperation() override;
+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;
+}
 
+namespace {
+/// A helper struct for FunctionBufferize and ModuleBufferize. Both passes are
+/// mostly identical.
+template <typename PassBase>
+struct BufferizePassHelper : public PassBase {
   void getDependentDialects(DialectRegistry &registry) const override {
-    registry
-        .insert<bufferization::BufferizationDialect, linalg::LinalgDialect,
+    registry.insert<bufferization::BufferizationDialect, linalg::LinalgDialect,
                     memref::MemRefDialect, tensor::TensorDialect,
                     vector::VectorDialect, scf::SCFDialect,
-                arith::ArithmeticDialect, StandardOpsDialect, AffineDialect>();
+                    arith::ArithmeticDialect, AffineDialect>();
     affine_ext::registerBufferizableOpInterfaceExternalModels(registry);
     arith_ext::registerBufferizableOpInterfaceExternalModels(registry);
     bufferization_ext::registerBufferizableOpInterfaceExternalModels(registry);
     linalg_ext::registerBufferizableOpInterfaceExternalModels(registry);
     scf_ext::registerBufferizableOpInterfaceExternalModels(registry);
-    std_ext::registerBufferizableOpInterfaceExternalModels(registry);
     tensor_ext::registerBufferizableOpInterfaceExternalModels(registry);
     vector_ext::registerBufferizableOpInterfaceExternalModels(registry);
   }
-};
-} // end namespace
-
-static void applyEnablingTransformations(ModuleOp moduleOp) {
-  RewritePatternSet patterns(moduleOp.getContext());
-  patterns.add<GeneralizePadTensorOpPattern>(moduleOp.getContext());
-  (void)applyPatternsAndFoldGreedily(moduleOp, std::move(patterns));
-}
-
-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;
-  if (useAlloca) {
+  void prepareOptions(BufferizationOptions &options) {
+    if (this->useAlloca) {
       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,
                                          Value to) {
       b.create<linalg::CopyOp>(loc, from, to);
     };
 
-  options.allowReturnMemref = allowReturnMemref;
-  options.allowUnknownOps = allowUnknownOps;
-  options.analysisFuzzerSeed = analysisFuzzerSeed;
-  options.testAnalysisOnly = testAnalysisOnly;
-
     // Enable InitTensorOp elimination.
     options.addPostAnalysisStep<
         linalg_ext::InsertSliceAnchoredInitTensorEliminationStep>();
-  // TODO: Find a way to enable this step automatically when bufferizing tensor
-  // dialect ops.
+    // TODO: Find a way to enable this step automatically when bufferizing
+    // tensor dialect ops.
     options.addPostAnalysisStep<tensor_ext::InplaceInsertSliceOpAnalysis>();
 
+    options.allowReturnMemref = this->allowReturnMemref;
+    options.allowUnknownOps = this->allowUnknownOps;
+    options.analysisFuzzerSeed = this->analysisFuzzerSeed;
+    options.testAnalysisOnly = this->testAnalysisOnly;
+  }
+
+  void runCleanupPipeline(StringRef name, Operation *op) {
+    OpPassManager cleanupPipeline(name);
+    cleanupPipeline.addPass(createCanonicalizerPass());
+    cleanupPipeline.addPass(createCSEPass());
+    cleanupPipeline.addPass(createLoopInvariantCodeMotionPass());
+    (void)this->runPipeline(cleanupPipeline, op);
+  }
+};
+
+struct LinalgComprehensiveModuleBufferize
+    : public BufferizePassHelper<LinalgComprehensiveModuleBufferizeBase<
+          LinalgComprehensiveModuleBufferize>> {
+  LinalgComprehensiveModuleBufferize() {}
+
+  LinalgComprehensiveModuleBufferize(
+      const LinalgComprehensiveModuleBufferize &p) {}
+
+  void runOnOperation() override;
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    BufferizePassHelper<LinalgComprehensiveModuleBufferizeBase<
+        LinalgComprehensiveModuleBufferize>>::getDependentDialects(registry);
+    // Module bufferization requires additional external models for FuncOp etc.
+    registry.insert<StandardOpsDialect>();
+    std_ext::registerBufferizableOpInterfaceExternalModels(registry);
+  }
+};
+
+struct LinalgComprehensiveFunctionBufferize
+    : public BufferizePassHelper<LinalgComprehensiveFunctionBufferizeBase<
+          LinalgComprehensiveFunctionBufferize>> {
+  LinalgComprehensiveFunctionBufferize() {}
+
+  LinalgComprehensiveFunctionBufferize(
+      const LinalgComprehensiveFunctionBufferize &p) {}
+
+  void runOnFunction() final;
+};
+} // end namespace
+
+static void applyEnablingTransformations(Operation *op) {
+  RewritePatternSet patterns(op->getContext());
+  patterns.add<GeneralizePadTensorOpPattern>(op->getContext());
+  (void)applyPatternsAndFoldGreedily(op, std::move(patterns));
+}
+
+void LinalgComprehensiveModuleBufferize::runOnOperation() {
+  BufferizationOptions options;
+  this->prepareOptions(options);
+
   ModuleOp moduleOp = getOperation();
   applyEnablingTransformations(moduleOp);
 
   if (failed(runComprehensiveBufferize(moduleOp, options))) {
     signalPassFailure();
     return;
   }
 
   if (options.testAnalysisOnly)
     return;
 
-  OpPassManager cleanupPipeline("builtin.module");
-  cleanupPipeline.addPass(createCanonicalizerPass());
-  cleanupPipeline.addPass(createCSEPass());
-  cleanupPipeline.addPass(createLoopInvariantCodeMotionPass());
-  (void)runPipeline(cleanupPipeline, moduleOp);
+  this->runCleanupPipeline("builtin.module", moduleOp);
+}
+
+void LinalgComprehensiveFunctionBufferize::runOnFunction() {
+  BufferizationOptions options;
+  this->prepareOptions(options);
+
+  applyEnablingTransformations(getOperation());
+
+  if (failed(runComprehensiveBufferize(getOperation(), options))) {
+    signalPassFailure();
+    return;
+  }
+
+  if (options.testAnalysisOnly)
+    return;
+
+  this->runCleanupPipeline("builtin.func", getOperation());
 }
 
 std::unique_ptr<Pass> mlir::createLinalgComprehensiveModuleBufferizePass() {
   return std::make_unique<LinalgComprehensiveModuleBufferize>();
 }
+
+std::unique_ptr<Pass> mlir::createLinalgComprehensiveFunctionBufferizePass() {
+  return std::make_unique<LinalgComprehensiveFunctionBufferize>();
+}

mlir/test/Dialect/Linalg/comprehensive-function-bufferize.mlir

+// RUN: mlir-opt %s -linalg-comprehensive-function-bufferize="allow-return-memref allow-unknown-ops" -split-input-file | FileCheck %s
+
+// Run fuzzer with different seeds.
+// RUN: mlir-opt %s -linalg-comprehensive-function-bufferize="test-analysis-only analysis-fuzzer-seed=23" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -linalg-comprehensive-function-bufferize="test-analysis-only analysis-fuzzer-seed=59" -split-input-file -o /dev/null
+// RUN: mlir-opt %s -linalg-comprehensive-function-bufferize="test-analysis-only analysis-fuzzer-seed=91" -split-input-file -o /dev/null
+
+// CHECK-LABEL: func @use_tensor_func_arg(
+//  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>
+func @use_tensor_func_arg(%A : tensor<?xf32>) -> (vector<4xf32>) {
+  %c0 = arith.constant 0 : index
+  %f0 = arith.constant 0.0 : f32
+
+  // CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]]
+  // CHECK: %[[res:.*]] = vector.transfer_read %[[A_memref]]
+  %0 = vector.transfer_read %A[%c0], %f0 : tensor<?xf32>, vector<4xf32>
+
+  // CHECK: return %[[res]]
+  return %0 : vector<4xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @return_tensor(
+//  CHECK-SAME:     %[[A:.*]]: tensor<?xf32>
+func @return_tensor(%A : tensor<?xf32>, %v : vector<4xf32>) -> (tensor<?xf32>) {
+  %c0 = arith.constant 0 : index
+
+  // CHECK: %[[A_memref:.*]] = bufferization.to_memref %[[A]]
+  // CHECK: %[[dim:.*]] = tensor.dim %[[A]]
+  // CHECK: %[[alloc:.*]] = memref.alloc(%[[dim]])
+  // CHECK: %[[casted:.*]] = memref.cast %[[alloc]]
+  // CHECK: linalg.copy(%[[A_memref]], %[[alloc]])
+  // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
+  %0 = vector.transfer_write %v, %A[%c0] : vector<4xf32>, tensor<?xf32>
+
+  // CHECK: %[[res_tensor:.*]] = bufferization.to_tensor %[[casted]]
+  // CHECK: return %[[res_tensor]]
+  return %0 : tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @func_without_tensor_args
+func @func_without_tensor_args(%v : vector<10xf32>) -> () {
+  // CHECK: %[[alloc:.*]] = memref.alloc()
+  %0 = linalg.init_tensor[10] : tensor<10xf32>
+
+  %c0 = arith.constant 0 : index
+  // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
+  %1 = vector.transfer_write %v, %0[%c0] : vector<10xf32>, tensor<10xf32>
+
+  %cst = arith.constant 0.0 : f32
+  // CHECK: vector.transfer_read %[[alloc]]
+  %r = vector.transfer_read %1[%c0], %cst : tensor<10xf32>, vector<11xf32>
+
+  vector.print %r : vector<11xf32>
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func private @private_func
+func private @private_func(tensor<?xf32>) -> ()
+
+// CHECK-LABEL: func @empty_func()
+func @empty_func() -> () {
+  return
+}

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

%r = scf.for %iv = %c0 to %c10 step %c1 iter_args(%bb = %t1) -> (tensor<?xf32>) {
     } else {
       %t2 = vector.transfer_write %v, %bb[%idx] : vector<5xf32>, tensor<?xf32>
       scf.yield %t2 : tensor<?xf32>
     }
     scf.yield %r2 : tensor<?xf32>
   }
   return %r : tensor<?xf32>
 }
+
+// -----
+
+// CHECK-LABEL: func @func_without_tensor_args
+func @func_without_tensor_args(%v : vector<10xf32>) -> () {
+  // CHECK: %[[alloc:.*]] = memref.alloc()
+  %0 = linalg.init_tensor[10] : tensor<10xf32>
+
+  %c0 = arith.constant 0 : index
+  // CHECK: vector.transfer_write %{{.*}}, %[[alloc]]
+  %1 = vector.transfer_write %v, %0[%c0] : vector<10xf32>, tensor<10xf32>
+
+  %cst = arith.constant 0.0 : f32
+  // CHECK: vector.transfer_read %[[alloc]]
+  %r = vector.transfer_read %1[%c0], %cst : tensor<10xf32>, vector<11xf32>
+
+  vector.print %r : vector<11xf32>
+  return
+}
+
+// -----
+
+// CHECK-LABEL: func private @private_func
+func private @private_func(tensor<?xf32>) -> ()
+
+// CHECK-LABEL: func @empty_func()
+func @empty_func() -> () {
+  return
+}