[mlir][bufferization] Add extra filter mechanism to bufferizeOp

[springerm]

Summary

This is needed so that sparse tensors can be included during the analysis, while being excluded during the bufferization phase.

Depends On D126568

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

 #include "mlir/Transforms/DialectConversion.h"
 
 namespace mlir {
 namespace bufferization {
 
 class AnalysisState;
 struct BufferizationState;
 struct BufferizationOptions;
+class OpFilter;
 
 /// A helper type converter class that automatically populates the relevant
 /// materializations and type conversions for bufferization.
 class BufferizeTypeConverter : public TypeConverter {
 public:
   BufferizeTypeConverter();
 };
 

 //===----------------------------------------------------------------------===//
 // Helper functions for extending Bufferization
 //===----------------------------------------------------------------------===//
 
 /// Bufferize `op` and its nested ops that implement `BufferizableOpInterface`.
 /// Reuse an existing `BufferizationState`.
 ///
 /// Note: This function overload is useful for extending the bufferization.
-LogicalResult bufferizeOp(Operation *op,
-                          BufferizationState &bufferizationState);
+LogicalResult bufferizeOp(Operation *op, BufferizationState &bufferizationState,
+                          const OpFilter *opFilter = nullptr);
 
 /// Finalize all buffer allocations: Create alloc/dealloc ops as specified by
 /// the bufferization options.
 LogicalResult finalizeBuffers(Operation *op,
                               const BufferizationOptions &options);
 } // namespace bufferization
 } // namespace mlir
 
 #endif // MLIR_DIALECT_BUFFERIZATION_TRANSFORMS_BUFFERIZE_H

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

 }
 
 namespace {
 /// A rewriter that keeps track of extra information during bufferization.
 class BufferizationRewriter : public IRRewriter {
 public:
   BufferizationRewriter(MLIRContext *ctx, DenseSet<Operation *> &erasedOps,
                         DenseSet<Operation *> &toMemrefOps,
-                        const BufferizationOptions &options)
+                        const BufferizationOptions &options,
+                        const OpFilter *opFilter)
       : IRRewriter(ctx), erasedOps(erasedOps), toMemrefOps(toMemrefOps),
-        options(options) {}
+        options(options), opFilter(opFilter) {}
 
 protected:
   void notifyOperationRemoved(Operation *op) override {
     IRRewriter::notifyOperationRemoved(op);
     erasedOps.insert(op);
     // Erase if present.
     toMemrefOps.erase(op);
   }
 
   void notifyOperationInserted(Operation *op) override {
     IRRewriter::notifyOperationInserted(op);
 
     // Keep track of to_memref ops.
     if (isa<ToMemrefOp>(op)) {
       toMemrefOps.insert(op);
       return;
     }
 
     // Skip to_tensor ops.
     if (isa<ToTensorOp>(op))
       return;
 
+    // Skip non-tensor ops.
+    if (!hasTensorSemantics(op))
+      return;
+
+    // Skip ops that are not allowed.
+    if (!options.isOpAllowed(op) || (opFilter && !opFilter->isOpAllowed(op)))
+      return;
+
     // Adding new bufferizable ops is not allowed during bufferization. Such ops
     // would not be analyzed and can lead to surprising behavior.
-    assert((!hasTensorSemantics(op) || !options.isOpAllowed(op)) &&
-           "creating new tensor ops is not allowed during bufferization");
+    llvm_unreachable(
   }
 
 private:
   /// A set of all erased ops.
   DenseSet<Operation *> &erasedOps;
 
   /// A set of all to_memref ops.
   DenseSet<Operation *> &toMemrefOps;
 
   /// The bufferization options.
   /// Used for debug modes.
   LLVM_ATTRIBUTE_UNUSED
   const BufferizationOptions &options;
+
+  const OpFilter *opFilter;
 };
 } // namespace
 
-LogicalResult
-bufferization::bufferizeOp(Operation *op,
-                           BufferizationState &bufferizationState) {
+LogicalResult bufferization::bufferizeOp(Operation *op,
+                                         BufferizationState &bufferizationState,
+                                         const OpFilter *opFilter) {
   const auto &options = bufferizationState.getOptions();
   assert(options.unknownTypeConversion !=
              BufferizationOptions::LayoutMapOption::InferLayoutMap &&
          "invalid layout map option");
 
   // Keep track of to_memref ops.
   DenseSet<Operation *> toMemrefOps;
   op->walk([&](ToMemrefOp toMemrefOp) { toMemrefOps.insert(toMemrefOp); });
if (hasTensorSemantics(op))
       worklist.push_back(op);
   });
 
   // Keep track of all erased ops.
   DenseSet<Operation *> erasedOps;
 
   // Bufferize all ops.
   BufferizationRewriter rewriter(op->getContext(), erasedOps, toMemrefOps,
-                                 bufferizationState.getOptions());
+                                 bufferizationState.getOptions(), opFilter);
   for (unsigned i = 0; i < worklist.size(); ++i) {
     Operation *op = worklist[i];
     // Skip ops that were erased.
     if (erasedOps.contains(op))
       continue;
     // Skip ops that are not bufferizable or not allowed.
     auto bufferizableOp = options.dynCastBufferizableOp(op);
     if (!bufferizableOp)
       continue;
+    if (opFilter && !opFilter->isOpAllowed(op))
+      continue;
     // Skip ops that no longer have tensor semantics.
     if (!hasTensorSemantics(op))
       continue;
     // Bufferize the op.
     rewriter.setInsertionPoint(op);
     if (failed(bufferizableOp.bufferize(rewriter, bufferizationState)))
       return op->emitError("failed to bufferize op");
   }
if (erasedOps.contains(op))
       continue;
     // Ops that no longer have tensor semantics (because they were updated
     // in-place) are allowed.
     if (!hasTensorSemantics(op))
       continue;
     // Continue ops that are not allowed.
     if (!options.isOpAllowed(op))
       continue;
+    if (opFilter && !opFilter->isOpAllowed(op))
+      continue;
     // Ops without any uses and no side effects will fold away.
     if (op->getUses().empty() && MemoryEffectOpInterface::hasNoEffect(op))
       continue;
     return op->emitError("op was not bufferized");
   }
 
   return success();
 }