[mlir][bufferization][NFC] Move OpFilter out of BufferizationOptions

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

 class OpBuilder;
 
 namespace bufferization {
 
 class AnalysisState;
 class BufferizableOpInterface;
 struct DialectAnalysisState;
 
-/// Options for BufferizableOpInterface-based bufferization.
-struct BufferizationOptions {
-  /// Allocator function: Generate a memref allocation with the given type,
-  /// dynamic extents and alignment.
-  using AllocationFn = std::function<FailureOr<Value>(
-      OpBuilder &, Location, MemRefType, ValueRange, unsigned int)>;
-  /// Deallocator function: Deallocate a buffer that was allocated with
-  /// AllocatorFn.
-  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)>;
-  /// Initializer function for analysis state.
-  using AnalysisStateInitFn = std::function<void(AnalysisState &)>;
-  /// Initializer function for dialect-specific analysis state.
-  using DialectStateInitFn =
-      std::function<std::unique_ptr<DialectAnalysisState>()>;
-
+class OpFilter {
+public:
   /// An op filter entry. Filters can be used to specify which ops should be
   /// processed by the bufferization.
-  struct OpFilterEntry {
+  struct Entry {
     /// If the filter function evaluates to `true`, the filter matches.
     using FilterFn = std::function<bool(Operation *)>;
 
     /// Filter type: A filter can either be a DENY filter or an ALLOW filter.
     enum FilterType : int8_t { DENY = 0, ALLOW = 1 };
 
     FilterFn fn;
     FilterType type;
   };
 
-  enum class LayoutMapOption : int8_t {
-    InferLayoutMap = 0,
-    IdentityLayoutMap = 1,
-    FullyDynamicLayoutMap = 2
-  };
-
-  BufferizationOptions();
-
-  /// Return `true` if the filter has at least one ALLOW rule.
-  bool filterHasAllowRule() const {
-    for (const OpFilterEntry &e : opFilter)
-      if (e.type == OpFilterEntry::FilterType::ALLOW)
-        return true;
-    return false;
-  }
-
-  /// Return whether the op should be bufferized or not.
+  /// Return whether the op is allowed or not.
   ///
-  /// If the filter does not have an ALLOW rule, ops are bufferized by default,
+  /// If the filter does not have an ALLOW rule, ops are allowed 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
+  /// ALLOW rule, ops are denied by default and only allowed if they match
   /// an ALLOW rule and no DENY rule.
   bool isOpAllowed(Operation *op) const;
 
-  /// Allow the given dialects in the filter.
+  /// Allow the given dialects.
   ///
-  /// This function adds one or multiple ALLOW filters.
-  template <typename... DialectTs>
-  void allowDialectInFilter() {
-    // The following expands a call to allowDialectInFilterImpl for each dialect
+  /// This function adds one or multiple ALLOW entries.
+  template <typename... DialectTs> void allowDialect() {
+    // The following expands a call to allowDialectImpl for each dialect
     // in 'DialectTs'. This magic is necessary due to a limitation in the places
     // that a parameter pack can be expanded in c++11.
     // FIXME: In c++17 this can be simplified by using 'fold expressions'.
-    (void)std::initializer_list<int>{
-        0, (allowDialectInFilterImpl<DialectTs>(), 0)...};
+    (void)std::initializer_list<int>{0, (allowDialectImpl<DialectTs>(), 0)...};
   }
 
-  /// Deny the given dialects in the filter.
+  /// Deny the given dialects.
   ///
-  /// This function adds one or multiple DENY filters.
-  template <typename... DialectTs> void denyDialectInFilter() {
+  /// This function adds one or multiple DENY entries.
+  template <typename... DialectTs> void denyDialect() {
     // FIXME: In c++17 this can be simplified by using 'fold expressions'.
-    (void)std::initializer_list<int>{
-        0, (denyDialectInFilterImpl<DialectTs>(), 0)...};
+    (void)std::initializer_list<int>{0, (denyDialectImpl<DialectTs>(), 0)...};
   }
 
-  /// Allow the given dialect in the filter.
+  /// Allow the given dialect.
   ///
-  /// This function adds an ALLOW filter.
-  void allowDialectInFilter(StringRef dialectNamespace) {
-    OpFilterEntry::FilterFn filterFn = [=](Operation *op) {
+  /// This function adds an ALLOW entry.
+  void allowDialect(StringRef dialectNamespace) {
+    Entry::FilterFn filterFn = [=](Operation *op) {
       return op->getDialect()->getNamespace() == dialectNamespace;
     };
-    opFilter.push_back(
-        OpFilterEntry{filterFn, OpFilterEntry::FilterType::ALLOW});
+    entries.push_back(Entry{filterFn, Entry::FilterType::ALLOW});
   }
 
-  /// Allow the given ops in the filter.
+  /// Allow the given ops.
   ///
-  /// This function adds one or multiple ALLOW filters.
-  template <typename... OpTys>
-  void allowOperationInFilter() {
+  /// This function adds one or multiple ALLOW entries.
+  template <typename... OpTys> void allowOperation() {
     // FIXME: In c++17 this can be simplified by using 'fold expressions'.
-    (void)std::initializer_list<int>{
-        0, (allowOperationInFilterImpl<OpTys>(), 0)...};
+    (void)std::initializer_list<int>{0, (allowOperationImpl<OpTys>(), 0)...};
   }
 
-  /// Deny the given ops in the filter.
+  /// Deny the given ops.
   ///
-  /// This function adds one or multiple DENY filters.
-  template <typename... OpTys> void denyOperationInFilter() {
+  /// This function adds one or multiple DENY entries.
+  template <typename... OpTys> void denyOperation() {
     // FIXME: In c++17 this can be simplified by using 'fold expressions'.
-    (void)std::initializer_list<int>{
-        0, (denyOperationInFilterImpl<OpTys>(), 0)...};
+    (void)std::initializer_list<int>{0, (denyOperationImpl<OpTys>(), 0)...};
   }
 
-  /// Allow the given op in the filter.
+  /// Allow the given op.
   ///
-  /// This function adds an ALLOW filter.
-  void allowOperationInFilter(StringRef opName) {
-    OpFilterEntry::FilterFn filterFn = [=](Operation *op) {
+  /// This function adds an ALLOW entry.
+  void allowOperation(StringRef opName) {
+    Entry::FilterFn filterFn = [=](Operation *op) {
       return op->getName().getStringRef() == opName;
     };
-    allowOperationInFilter(filterFn);
+    allowOperation(filterFn);
   }
 
-  /// Deny the given op in the filter.
+  /// Deny the given op.
   ///
-  /// This function adds a DENY filter.
-  void denyOperationInFilter(StringRef opName) {
-    OpFilterEntry::FilterFn filterFn = [=](Operation *op) {
+  /// This function adds a DENY entry.
+  void denyOperation(StringRef opName) {
+    Entry::FilterFn filterFn = [=](Operation *op) {
       return op->getName().getStringRef() == opName;
     };
-    denyOperationInFilter(filterFn);
+    denyOperation(filterFn);
   }
 
-  /// Allow ops that are matched by `fn` in the filter.
+  /// Allow ops that are matched by `fn`.
   ///
-  /// This function adds an ALLOW filter.
-  void allowOperationInFilter(OpFilterEntry::FilterFn fn) {
-    opFilter.push_back(OpFilterEntry{fn, OpFilterEntry::FilterType::ALLOW});
+  /// This function adds an ALLOW entry.
+  void allowOperation(Entry::FilterFn fn) {
+    entries.push_back(Entry{fn, Entry::FilterType::ALLOW});
   }
 
-  /// Deny ops that are matched by `fn` in the filter.
+  /// Deny ops that are matched by `fn`.
   ///
-  /// This function adds a DENY filter.
-  void denyOperationInFilter(OpFilterEntry::FilterFn fn) {
-    opFilter.push_back(OpFilterEntry{fn, OpFilterEntry::FilterType::DENY});
+  /// This function adds a DENY entry.
+  void denyOperation(Entry::FilterFn fn) {
+    entries.push_back(Entry{fn, Entry::FilterType::DENY});
+  }
+
+private:
+  /// Return `true` if the filter has at least one ALLOW rule.
+  bool hasAllowRule() const {
+    for (const Entry &e : entries)
+      if (e.type == Entry::FilterType::ALLOW)
+        return true;
+    return false;
+  }
+
+  /// Allow a dialect.
+  template <typename DialectT> void allowDialectImpl() {
+    allowDialect(DialectT::getDialectNamespace());
+  }
+
+  /// Deny a dialect.
+  template <typename DialectT> void denyDialectImpl() {
+    denyDialect(DialectT::getDialectNamespace());
+  }
+
+  /// Allow an op.
+  template <typename OpTy> void allowOperationImpl() {
+    allowOperation(OpTy::getOperationName());
   }
 
+  /// Deny an op.
+  template <typename OpTy> void denyOperationImpl() {
+    denyOperation(OpTy::getOperationName());
+  }
+
+  /// A list of filter entries that determine whether an op should be allowed or
+  /// denied. If the filter has an ALLOW rule, only ops that are allowed and not
+  /// denied are allowed. If the filter does not have an ALLOW rule, only ops
+  /// that are not denied are allowed.
+  SmallVector<Entry> entries;
+};
+
+/// Options for BufferizableOpInterface-based bufferization.
+struct BufferizationOptions {
+  /// Allocator function: Generate a memref allocation with the given type,
+  /// dynamic extents and alignment.
+  using AllocationFn = std::function<FailureOr<Value>(
+      OpBuilder &, Location, MemRefType, ValueRange, unsigned int)>;
+  /// Deallocator function: Deallocate a buffer that was allocated with
+  /// AllocatorFn.
+  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)>;
+  /// Initializer function for analysis state.
+  using AnalysisStateInitFn = std::function<void(AnalysisState &)>;
+  /// Initializer function for dialect-specific analysis state.
+  using DialectStateInitFn =
+      std::function<std::unique_ptr<DialectAnalysisState>()>;
+
+  enum class LayoutMapOption : int8_t {
+    InferLayoutMap = 0,
+    IdentityLayoutMap = 1,
+    FullyDynamicLayoutMap = 2
+  };
+
+  BufferizationOptions();
+
   /// Try to cast the given op to BufferizableOpInterface if the op is allow
   /// listed.
   BufferizableOpInterface dynCastBufferizableOp(Operation *op) const;
 
   /// Try to cast the given value to BufferizableOpInterface if the op is allow
   /// listed.
   BufferizableOpInterface dynCastBufferizableOp(Value value) const;
 
+  /// A filter that specifies which ops should be bufferized and which ops
+  /// should be ignored.
+  OpFilter opFilter;
+
+  /// Return `true` if the given op should be bufferized.
+  bool isOpAllowed(Operation *op) const;
+
   /// Helper functions for allocation, deallocation, memory copying.
   Optional<AllocationFn> allocationFn;
   Optional<DeallocationFn> deallocationFn;
   Optional<MemCpyFn> memCpyFn;
 
   /// Create a memref allocation with the given type and dynamic extents.
   FailureOr<Value> createAlloc(OpBuilder &b, Location loc, MemRefType type,
                                ValueRange dynShape) const;
struct BufferizationOptions {
   /// result of a `linalg.generic` op could bufferize in-place with an "in"
   /// OpOperand if the corresponding "out" operand is not used within the
   /// computation. Whether this pays off or not can be very input IR-specific.
   bool alwaysAliasingWithDest = true;
 
   /// Buffer alignment for new memory allocations.
   unsigned int bufferAlignment = 128;
 
-  /// A list of op filters that determine whether an op should be processed or
-  /// ignored by the bufferization. If the filter has an ALLOW rule, only ops
-  /// that are allowed and not denied are bufferized. If the filter does not
-  /// have an ALLOW rule, only ops that are not denied are bufferized.
-  SmallVector<OpFilterEntry> opFilter;
-
   /// Initializer functions for analysis state. These can be used to
   /// initialize dialect-specific analysis state.
   SmallVector<AnalysisStateInitFn> stateInitializers;
 
   /// Add a analysis state initializer that initializes the specified
   /// dialect-specific analysis state.
   void addDialectStateInitializer(StringRef name, const DialectStateInitFn &fn);
-
-private:
-  /// Allow a dialect.
-  template <typename DialectT>
-  void allowDialectInFilterImpl() {
-    allowDialectInFilter(DialectT::getDialectNamespace());
-  }
-
-  /// Deny a dialect.
-  template <typename DialectT> void denyDialectInFilterImpl() {
-    denyDialectInFilter(DialectT::getDialectNamespace());
-  }
-
-  /// Allow an op.
-  template <typename OpTy>
-  void allowOperationInFilterImpl() {
-    allowOperationInFilter(OpTy::getOperationName());
-  }
-
-  /// Deny an op.
-  template <typename OpTy> void denyOperationInFilterImpl() {
-    denyOperationInFilter(OpTy::getOperationName());
-  }
 };
 
 /// Specify fine-grain relationship between buffers to enable more analysis.
 enum class BufferRelation {
   None,
   // TODO: ResultContainsOperand,
   // TODO: OperandContainsResult,
   Equivalent

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

ArithmeticBufferizePass(uint64_t alignment = 0, bool constantOpOnly = false)
       : ArithmeticBufferizeBase<ArithmeticBufferizePass>(),
         constantOpOnly(constantOpOnly) {
     this->alignment = alignment;
   }
 
   void runOnOperation() override {
     BufferizationOptions options = getPartialBufferizationOptions();
     if (constantOpOnly) {
-      options.allowOperationInFilter<arith::ConstantOp>();
+      options.opFilter.allowOperation<arith::ConstantOp>();
     } else {
-      options.allowDialectInFilter<arith::ArithmeticDialect>();
+      options.opFilter.allowDialect<arith::ArithmeticDialect>();
     }
     options.bufferAlignment = alignment;
 
     if (failed(bufferizeOp(getOperation(), options)))
       signalPassFailure();
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {

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

 
 /// Attribute name used to mark allocs that are created by the bufferization.
 static const char *kBufferAllocationAttr = "bufferization.allocation";
 
 /// Attribute name used to mark allocs that should not be deallocated.
 static const char *kSkipDeallocAttr = "bufferization.skip_dealloc";
 
 //===----------------------------------------------------------------------===//
-// BufferizationOptions
+// OpFilter
 //===----------------------------------------------------------------------===//
 
-// 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;
-
+bool OpFilter::isOpAllowed(Operation *op) const {
   // All other ops: Allow/disallow according to filter.
-  bool isAllowed = !filterHasAllowRule();
-  for (const OpFilterEntry &entry : opFilter) {
+  bool isAllowed = !hasAllowRule();
+  for (const Entry &entry : entries) {
     bool filterResult = entry.fn(op);
     switch (entry.type) {
-    case OpFilterEntry::ALLOW:
+    case Entry::ALLOW:
       isAllowed |= filterResult;
       break;
-    case OpFilterEntry::DENY:
+    case Entry::DENY:
       if (filterResult)
         // DENY filter matches. This op is no allowed. (Even if other ALLOW
         // filters may match.)
         return false;
     };
   }
   return isAllowed;
 }
 
+//===----------------------------------------------------------------------===//
+// 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;
+
+  return opFilter.isOpAllowed(op);
+}
+
 BufferizableOpInterface
 BufferizationOptions::dynCastBufferizableOp(Operation *op) const {
   auto bufferizableOp = dyn_cast<BufferizableOpInterface>(op);
   if (!bufferizableOp)
     return nullptr;
   if (!isOpAllowed(op))
     return nullptr;
   return bufferizableOp;

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

if (!options) {
       opt.functionBoundaryTypeConversion =
           parseLayoutMapOption(functionBoundaryTypeConversion);
       opt.printConflicts = printConflicts;
       opt.testAnalysisOnly = testAnalysisOnly;
       opt.bufferizeFunctionBoundaries = bufferizeFunctionBoundaries;
       opt.promoteBufferResultsToOutParams = promoteBufferResultsToOutParams;
       opt.unknownTypeConversion = parseLayoutMapOption(unknownTypeConversion);
 
-      BufferizationOptions::OpFilterEntry::FilterFn filterFn =
+      OpFilter::Entry::FilterFn filterFn =
           [&](Operation *op) {
             // 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);
+      opt.opFilter.allowOperation(filterFn);
     } else {
       opt = *options;
     }
 
     ModuleOp moduleOp = getOperation();
     if (opt.bufferizeFunctionBoundaries) {
       if (failed(runOneShotModuleBufferize(moduleOp, opt))) {
         signalPassFailure();

 };
 } // namespace
 
 namespace {
 struct BufferizationBufferizePass
     : public BufferizationBufferizeBase<BufferizationBufferizePass> {
   void runOnOperation() override {
     BufferizationOptions options = getPartialBufferizationOptions();
-    options.allowDialectInFilter<BufferizationDialect>();
+    options.opFilter.allowDialect<BufferizationDialect>();
 
     if (failed(bufferizeOp(getOperation(), options)))
       signalPassFailure();
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry
         .insert<bufferization::BufferizationDialect, memref::MemRefDialect>();

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

 using namespace bufferization;
 
 namespace {
 /// Converts Linalg operations that work on tensor-type operands or results to
 /// work on buffers.
 struct LinalgBufferizePass : public LinalgBufferizeBase<LinalgBufferizePass> {
   void runOnOperation() override {
     BufferizationOptions options = getPartialBufferizationOptions();
-    options.allowDialectInFilter<linalg::LinalgDialect>();
+    options.opFilter.allowDialect<linalg::LinalgDialect>();
 
     if (failed(bufferizeOp(getOperation(), options)))
       signalPassFailure();
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<bufferization::BufferizationDialect, memref::MemRefDialect,
                     tensor::TensorDialect, linalg::LinalgDialect>();
     linalg::registerBufferizableOpInterfaceExternalModels(registry);
   }
 };
 } // namespace
 
 std::unique_ptr<OperationPass<func::FuncOp>> mlir::createLinalgBufferizePass() {
   return std::make_unique<LinalgBufferizePass>();
 }

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

 
 using namespace mlir;
 using namespace bufferization;
 
 namespace {
 struct ShapeBufferizePass : public ShapeBufferizeBase<ShapeBufferizePass> {
   void runOnOperation() override {
     BufferizationOptions options = getPartialBufferizationOptions();
-    options.allowDialectInFilter<shape::ShapeDialect>();
+    options.opFilter.allowDialect<shape::ShapeDialect>();
 
     if (failed(bufferizeOp(getOperation(), options)))
       signalPassFailure();
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<bufferization::BufferizationDialect, memref::MemRefDialect,
                     shape::ShapeDialect>();
     shape::registerBufferizableOpInterfaceExternalModels(registry);
   }
 };
 } // namespace
 
 std::unique_ptr<OperationPass<func::FuncOp>> mlir::createShapeBufferizePass() {
   return std::make_unique<ShapeBufferizePass>();
 }

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

 
 using namespace mlir;
 using namespace bufferization;
 
 namespace {
 struct TensorBufferizePass : public TensorBufferizeBase<TensorBufferizePass> {
   void runOnOperation() override {
     BufferizationOptions options = getPartialBufferizationOptions();
-    options.allowDialectInFilter<tensor::TensorDialect>();
+    options.opFilter.allowDialect<tensor::TensorDialect>();
 
     if (failed(bufferizeOp(getOperation(), options)))
       signalPassFailure();
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<bufferization::BufferizationDialect, memref::MemRefDialect,
                     tensor::TensorDialect, scf::SCFDialect,

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

 
 using namespace mlir;
 using namespace bufferization;
 
 namespace {
 struct VectorBufferizePass : public VectorBufferizeBase<VectorBufferizePass> {
   void runOnOperation() override {
     BufferizationOptions options = getPartialBufferizationOptions();
-    options.allowDialectInFilter<vector::VectorDialect>();
+    options.opFilter.allowDialect<vector::VectorDialect>();
 
     if (failed(bufferizeOp(getOperation(), options)))
       signalPassFailure();
   }
 
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<bufferization::BufferizationDialect, memref::MemRefDialect,
                     tensor::TensorDialect, vector::VectorDialect>();
     vector::registerBufferizableOpInterfaceExternalModels(registry);
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::vector::createVectorBufferizePass() {
   return std::make_unique<VectorBufferizePass>();
 }