[mlir][bufferize][NFC] Add function boundary bufferization flag to BufferizationOptions

[springerm]

Summary

This makes the API easier to use. Also allows us to check for incorrect API usage for easier debugging.

Depends On D124264

[springerm]

update

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

struct BufferizationOptions {
   }
 
   /// Return whether the op should be bufferized or not.
   ///
   /// If the filter does not have an ALLOW rule, ops are bufferized by default,
   /// unless they are explicitly marked as DENY. If the filter has at least one
   /// ALLOW rule, ops are ignored by default and only bufferized if they match
   /// an ALLOW rule and no DENY rule.
-  bool isOpAllowed(Operation *op) const {
-    bool isAllowed = !filterHasAllowRule();
-    for (const OpFilterEntry &entry : opFilter) {
-      bool filterResult = entry.fn(op);
-      switch (entry.type) {
-      case OpFilterEntry::ALLOW:
-        isAllowed |= filterResult;
-        break;
-      case OpFilterEntry::DENY:
-        if (filterResult)
-          // DENY filter matches. This op is no allowed. (Even if other ALLOW
-          // filters may match.)
-          return false;
-      };
-    }
-    return isAllowed;
-  }
+  bool isOpAllowed(Operation *op) const;
 
   /// Allow the given dialects in the filter.
   ///
   /// This function adds one or multiple ALLOW filters.
   template <typename... DialectTs>
   void allowDialectInFilter() {
     // The following expands a call to allowDialectInFilterImpl for each dialect
     // in 'DialectTs'. This magic is necessary due to a limitation in the places
struct BufferizationOptions {
   Optional<DeallocationFn> deallocationFn;
   Optional<MemCpyFn> memCpyFn;
 
   /// Specifies whether not bufferizable ops are allowed in the input. If so,
   /// bufferization.to_memref and bufferization.to_tensor ops are inserted at
   /// the boundaries.
   bool allowUnknownOps = false;
 
+  /// Specifies whether function boundaries (ops in the func dialect) should be
+  /// bufferized or not.
+  bool bufferizeFunctionBoundaries = false;
+
   /// Specifies whether dealloc ops should be generated along with alloc ops. If
   /// not, new memory allocations will leak.
   bool createDeallocs = true;
 
   /// Seed for the analysis fuzzer. If set to `0`, the fuzzer is deactivated.
   /// Should be used only with `testAnalysisOnly = true`.
   unsigned analysisFuzzerSeed = 0;
 
public:
 
   /// Return true if the given tensor (or an aliasing tensor) is yielded from
   /// the containing block. Also include all aliasing tensors in the same block.
   ///
   /// Note: In the absence of an analysis, an implementation may return true for
   /// any given tensor.
   virtual bool isTensorYielded(Value tensor) const = 0;
 
+  /// Return `true` if the given dialect state exists.
+  bool hasDialectState(StringRef name) const {
+    auto it = dialectState.find(name);
+    return it != dialectState.end();
+  }
+
   /// Return dialect-specific bufferization state.
   template <typename StateT>
   Optional<const StateT *> getDialectState(StringRef name) const {
     auto it = dialectState.find(name);
     if (it == dialectState.end())
       return None;
     return static_cast<const StateT *>(it->getSecond().get());
   }
 
   /// Return dialect-specific analysis state or create one if none exists.
   template <typename StateT>
   StateT &getOrCreateDialectState(StringRef name) {
     // Create state if it does not exist yet.
-    if (!dialectState.count(name))
+    if (!hasDialectState(name))
       dialectState[name] = std::make_unique<StateT>();
     return static_cast<StateT &>(*dialectState[name]);
   }
 
   void insertDialectState(StringRef name,
                           std::unique_ptr<DialectAnalysisState> state) {
     assert(!dialectState.count(name) && "dialect state already initialized");
     dialectState[name] = std::move(state);

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

 
 //===----------------------------------------------------------------------===//
 // BufferizationOptions
 //===----------------------------------------------------------------------===//
 
 // Default constructor for BufferizationOptions.
 BufferizationOptions::BufferizationOptions() = default;
 
+bool BufferizationOptions::isOpAllowed(Operation *op) const {
+  // Special case: If function boundary bufferization is deactivated, do not
+  // allow ops that belong to the `func` dialect.
+  bool isFuncBoundaryOp = isa_and_nonnull<func::FuncDialect>(op->getDialect());
+  if (!bufferizeFunctionBoundaries && isFuncBoundaryOp)
+    return false;
+
+  // All other ops: Allow/disallow according to filter.
+  bool isAllowed = !filterHasAllowRule();
+  for (const OpFilterEntry &entry : opFilter) {
+    bool filterResult = entry.fn(op);
+    switch (entry.type) {
+    case OpFilterEntry::ALLOW:
+      isAllowed |= filterResult;
+      break;
+    case OpFilterEntry::DENY:
+      if (filterResult)
+        // DENY filter matches. This op is no allowed. (Even if other ALLOW
+        // filters may match.)
+        return false;
+    };
+  }
+  return isAllowed;
+}
+
 BufferizableOpInterface
 BufferizationOptions::dynCastBufferizableOp(Operation *op) const {
   if (isOpAllowed(op))
     return dyn_cast<BufferizableOpInterface>(op);
   return nullptr;
 }
 
 BufferizableOpInterface

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

if (!options) {
       // pass.
       opt.allowReturnAllocs = allowReturnAllocs;
       opt.allowUnknownOps = allowUnknownOps;
       opt.analysisFuzzerSeed = analysisFuzzerSeed;
       opt.createDeallocs = createDeallocs;
       opt.fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
       opt.printConflicts = printConflicts;
       opt.testAnalysisOnly = testAnalysisOnly;
+      opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
 
       BufferizationOptions::OpFilterEntry::FilterFn filterFn =
           [&](Operation *op) {
-            // Disallow non-func dialect ops. I.e., no ops related to function
-            // calls. (Unless explicitly activated.)
-            bool isFuncBoundaryOp =
-                isa_and_nonnull<func::FuncDialect>(op->getDialect());
-            if (!this->bufferizeFunctionBoundaries && isFuncBoundaryOp)
-              return false;
             // Filter may be specified via options.
             if (this->dialectFilter.hasValue())
               return llvm::find(this->dialectFilter,
                                 op->getDialect()->getNamespace()) !=
                      this->dialectFilter.end();
             // No filter specified: All other ops are allowed.
             return true;
           };
       opt.allowOperationInFilter(filterFn);
     } else {
       opt = *options;
     }
 
     ModuleOp moduleOp = getOperation();
-    if (bufferizeFunctionBoundaries) {
+    if (opt.bufferizeFunctionBoundaries) {
       if (failed(runOneShotModuleBufferize(moduleOp, opt))) {
         signalPassFailure();
         return;
       }
     } else {
       if (failed(runOneShotBufferize(moduleOp, opt))) {
         signalPassFailure();
         return;
bufferization::finalizeBuffers(Operation *op,
   if (failed(createAllocDeallocOps(op, options)))
     return failure();
 
   return success();
 }
 
 LogicalResult bufferization::bufferizeOp(Operation *op,
                                          const AnalysisState &analysisState) {
+  // Catch incorrect API usage.
+  assert((analysisState.hasDialectState(
+              func::FuncDialect::getDialectNamespace()) ||
+          !analysisState.getOptions().bufferizeFunctionBoundaries) &&
+         "must use ModuleBufferize to bufferize function boundaries");
+
   BufferizationState bufferizationState(analysisState);
   if (failed(bufferizeOp(op, bufferizationState)))
     return failure();
   if (failed(finalizeBuffers(op, analysisState.getOptions())))
     return failure();
   return success();
 }
 

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

 
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 
 #include <random>
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"

 
 LogicalResult bufferization::analyzeOp(Operation *op,
                                        OneShotAnalysisState &state) {
   DominanceInfo domInfo(op);
   BufferizationAliasInfo &aliasInfo = state.getAliasInfo();
   const auto &options =
       static_cast<const OneShotBufferizationOptions &>(state.getOptions());
 
+  // Catch incorrect API usage.
+  assert((state.hasDialectState(func::FuncDialect::getDialectNamespace()) ||
+          !options.bufferizeFunctionBoundaries) &&
+         "must use ModuleBufferize to bufferize function boundaries");
+
   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);

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

static void foldMemRefCasts(func::FuncOp funcOp) {
 
   auto newFuncType = FunctionType::get(
       funcOp.getContext(), funcOp.getFunctionType().getInputs(), resultTypes);
   funcOp.setType(newFuncType);
 }
 
 LogicalResult mlir::bufferization::runOneShotModuleBufferize(
     ModuleOp moduleOp, OneShotBufferizationOptions options) {
+  assert(options.bufferizeFunctionBoundaries &&
+         "expected that function boundary bufferization is activated");
   IRRewriter rewriter(moduleOp.getContext());
   OneShotAnalysisState analysisState(moduleOp, options);
   BufferizationState bufferizationState(analysisState);
   FuncAnalysisState &funcState = getFuncAnalysisState(analysisState);
   BufferizationAliasInfo &aliasInfo = analysisState.getAliasInfo();
 
   // A list of functions in the order in which they are analyzed + bufferized.
   SmallVector<func::FuncOp> orderedFuncOps;

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

if (!options) {
     opt.allowReturnAllocs = allowReturnAllocs;
     opt.allowUnknownOps = allowUnknownOps;
     opt.analysisFuzzerSeed = analysisFuzzerSeed;
     opt.createDeallocs = createDeallocs;
     opt.fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
     opt.printConflicts = printConflicts;
     opt.testAnalysisOnly = testAnalysisOnly;
     opt.alwaysAliasingWithDest = alwaysAliasingWithDest;
+    opt.bufferizeFunctionBoundaries = true;
     if (initTensorElimination) {
       opt.addPostAnalysisStep(insertSliceAnchoredInitTensorEliminationStep);
     }
   } else {
     opt = *options;
   }
 
   ModuleOp moduleOp = getOperation();