[mlir][bufferize] Add noAnalysisFuncFilter to OneShotBufferizationOptions struct

[mamrami]

Summary

This change is needed in order to set the flag when running the pass not via the command line.
It also allows simplifying the signature of some functions.

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

 //===- OneShotAnalysis.h - One-Shot (Single Pass) Analysis ------*- C++ -*-===//
 //
 // 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_BUFFERIZATION_TRANSFORMS_ONESHOTANALYSIS_H
 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTANALYSIS_H
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "llvm/ADT/EquivalenceClasses.h"
+#include <string>
 
 namespace mlir {
 namespace bufferization {
 
 struct OneShotBufferizationOptions;
 class BufferizationAliasInfo;
 struct BufferizationStatistics;
 class OneShotAnalysisState;
 
 /// Options for analysis-enabled bufferization.
 struct OneShotBufferizationOptions : public BufferizationOptions {
   enum class AnalysisHeuristic { BottomUp, TopDown };
 
   OneShotBufferizationOptions() = default;
 
   /// Specifies whether returning newly allocated memrefs should be allowed.
   /// Otherwise, a pass failure is triggered.
   bool allowReturnAllocs = false;
 
   /// The heuristic controls the order in which ops are traversed during the
   /// analysis.
   AnalysisHeuristic analysisHeuristic = AnalysisHeuristic::BottomUp;
+
+  /// Specify the functions that should not be analyzed. copyBeforeWrite will be
+  /// set to true when bufferizing them.
+  llvm::ArrayRef<std::string> noAnalysisFuncFilter;
 };
 
 /// The BufferizationAliasInfo class maintains a list of buffer aliases and
 /// equivalence classes to support bufferization.
 class BufferizationAliasInfo {
 public:
   explicit BufferizationAliasInfo(Operation *rootOp);
 

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

 //===- OneShotModuleBufferize.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_BUFFERIZATION_TRANSFORMS_ONESHOTMODULEBUFFERIZE_H
 #define MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTMODULEBUFFERIZE_H
 
-#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 namespace mlir {
 
 struct LogicalResult;
 class ModuleOp;
 
 namespace bufferization {
 struct BufferizationStatistics;
 class OneShotAnalysisState;
 struct OneShotBufferizationOptions;
 
 /// Analyze `moduleOp` and its nested ops. Bufferization decisions are stored in
 /// `state`.
 LogicalResult analyzeModuleOp(ModuleOp moduleOp, OneShotAnalysisState &state,
                               BufferizationStatistics *statistics = nullptr);
 
 /// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.
 ///
 /// Note: This function does not run One-Shot Analysis. No buffer copies are
 /// inserted except two cases:
 /// - `options.copyBeforeWrite` is set, in which case buffers are copied before
 ///   every write.
-/// - `options.copyBeforeWrite` is not set and `analysisFilterFn` returns true
-///   for some FuncOps. These FuncOps were not analyzed. Buffer copies will be
-///   inserted only to these FuncOps.
-LogicalResult
-bufferizeModuleOp(ModuleOp moduleOp, const OneShotBufferizationOptions &options,
-                  BufferizationStatistics *statistics = nullptr,
-                  OpFilter::Entry::FilterFn analysisFilterFn = nullptr);
+/// - `options.copyBeforeWrite` is not set and `options.noAnalysisFuncFilter`
+///   is not empty. The FuncOps it contains were not analyzed. Buffer copies
+///   will be inserted only to these FuncOps.
+LogicalResult bufferizeModuleOp(ModuleOp moduleOp,
+                                const OneShotBufferizationOptions &options,
+                                BufferizationStatistics *statistics = nullptr);
 
 /// Remove bufferization attributes on every FuncOp arguments in the ModuleOp.
 void removeBufferizationAttributesInModule(ModuleOp moduleOp);
 
 /// Run One-Shot 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 runOneShotModuleBufferize(
     ModuleOp moduleOp,
     const bufferization::OneShotBufferizationOptions &options,
-    BufferizationStatistics *statistics = nullptr,
-    OpFilter::Entry::FilterFn analysisFilterFn = nullptr);
+    BufferizationStatistics *statistics = nullptr);
 
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_ONESHOTMODULEBUFFERIZE_H

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

if (!options) {
       opt.createDeallocs = createDeallocs;
       opt.functionBoundaryTypeConversion =
           parseLayoutMapOption(functionBoundaryTypeConversion);
       if (mustInferMemorySpace)
         opt.defaultMemorySpace = std::nullopt;
       opt.printConflicts = printConflicts;
       opt.testAnalysisOnly = testAnalysisOnly;
       opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
+      opt.noAnalysisFuncFilter = noAnalysisFuncFilter;
 
       // Configure type converter.
       LayoutMapOption unknownTypeConversionOption =
           parseLayoutMapOption(unknownTypeConversion);
       opt.unknownTypeConverterFn = [=](Value value, Attribute memorySpace,
                                        const BufferizationOptions &options) {
         auto tensorType = value.getType().cast<TensorType>();
         if (unknownTypeConversionOption == LayoutMapOption::IdentityLayoutMap)
if (!options) {
       opt.opFilter.allowOperation(filterFn);
     } else {
       opt = *options;
     }
 
     BufferizationStatistics statistics;
     ModuleOp moduleOp = getOperation();
     if (opt.bufferizeFunctionBoundaries) {
-      OpFilter::Entry::FilterFn analysisFilterFn = nullptr;
-      // FuncOps whose names are specified in noAnalysisFuncFilter will not be
-      // analyzed. Ops in these FuncOps will not be analyzed as well.
-      if (this->noAnalysisFuncFilter.hasValue())
-        analysisFilterFn = [=](Operation *op) {
-          auto func = dyn_cast<func::FuncOp>(op);
-          if (!func)
-            func = op->getParentOfType<func::FuncOp>();
-          if (func)
-            return llvm::is_contained(noAnalysisFuncFilter, func.getSymName());
-          return false;
-        };
-      if (failed(runOneShotModuleBufferize(moduleOp, opt, &statistics,
-                                           analysisFilterFn))) {
+      if (failed(runOneShotModuleBufferize(moduleOp, opt, &statistics))) {
         signalPassFailure();
         return;
       }
     } else {
-      assert(!this->noAnalysisFuncFilter.hasValue() &&
+      assert(opt.noAnalysisFuncFilter.empty() &&
              "invalid combination of bufferization flags");
       if (failed(runOneShotBufferize(moduleOp, opt, &statistics))) {
         signalPassFailure();
         return;
       }
     }
 
     // Set pass statistics.

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

void mlir::bufferization::removeBufferizationAttributesInModule(
   moduleOp.walk([&](func::FuncOp op) {
     for (BlockArgument bbArg : op.getArguments())
       removeBufferizationAttributes(bbArg);
   });
 }
 
 LogicalResult mlir::bufferization::bufferizeModuleOp(
     ModuleOp moduleOp, const OneShotBufferizationOptions &options,
-    BufferizationStatistics *statistics,
-    OpFilter::Entry::FilterFn analysisFilterFn) {
+    BufferizationStatistics *statistics) {
   assert(options.bufferizeFunctionBoundaries &&
          "expected that function boundary bufferization is activated");
   IRRewriter rewriter(moduleOp.getContext());
 
   // A list of functions in the order in which they are analyzed + bufferized.
   SmallVector<func::FuncOp> orderedFuncOps;
 
   // A mapping of FuncOps to their callers.
   FuncCallerMap callerMap;
 
   if (failed(getFuncOpsOrderedByCalls(moduleOp, orderedFuncOps, callerMap)))
     return failure();
 
   // Bufferize functions.
   for (func::FuncOp funcOp : orderedFuncOps) {
     // Note: It would be good to apply cleanups here but we cannot as aliasInfo
     // would be invalidated.
-    bool copyBeforeWrite = options.copyBeforeWrite ||
-                           (analysisFilterFn && analysisFilterFn(funcOp));
+    bool copyBeforeWrite =
+        options.copyBeforeWrite ||
+        llvm::is_contained(options.noAnalysisFuncFilter, funcOp.getSymName());
     if (failed(bufferizeOp(funcOp, options, copyBeforeWrite,
                            /*opFilter=*/nullptr, statistics)))
       return failure();
     // Change buffer return types to more precise layout maps.
     if (options.functionBoundaryTypeConversion ==
         LayoutMapOption::InferLayoutMap)
       foldMemRefCasts(funcOp);
   }
 
   // Post-pass cleanup of function argument attributes.
   removeBufferizationAttributesInModule(moduleOp);
 
   return success();
 }
 
 LogicalResult mlir::bufferization::runOneShotModuleBufferize(
     ModuleOp moduleOp, const OneShotBufferizationOptions &options,
-    BufferizationStatistics *statistics,
-    OpFilter::Entry::FilterFn analysisFilterFn) {
+    BufferizationStatistics *statistics) {
   assert(options.bufferizeFunctionBoundaries &&
          "expected that function boundary bufferization is activated");
   assert(!(options.copyBeforeWrite && options.testAnalysisOnly) &&
          "invalid combination of bufferization flags");
   if (!options.copyBeforeWrite) {
-    if (!analysisFilterFn) {
+    if (options.noAnalysisFuncFilter.empty()) {
       if (failed(insertTensorCopies(moduleOp, options, statistics)))
         return failure();
     } else {
+      // FuncOps whose names are specified in options.noAnalysisFuncFilter will
+      // not be analyzed. Ops in these FuncOps will not be analyzed as well.
+      OpFilter::Entry::FilterFn analysisFilterFn = [=](Operation *op) {
+        auto func = dyn_cast<func::FuncOp>(op);
+        if (!func)
+          func = op->getParentOfType<func::FuncOp>();
+        if (func)
+          return llvm::is_contained(options.noAnalysisFuncFilter,
+                                    func.getSymName());
+        return false;
+      };
       OneShotBufferizationOptions updatedOptions(options);
       updatedOptions.opFilter.denyOperation(analysisFilterFn);
       if (failed(insertTensorCopies(moduleOp, updatedOptions, statistics)))
         return failure();
     }
   }
   if (options.testAnalysisOnly)
     return success();
-
-  if (failed(
-          bufferizeModuleOp(moduleOp, options, statistics, analysisFilterFn)))
+  if (failed(bufferizeModuleOp(moduleOp, options, statistics)))
     return failure();
   return success();
 }