[mlir][linalg][bufferize] Add dialect filter to BufferizationOptions

[springerm]

Summary

This change adds a new option dialectFilter to BufferizationOptions. Only ops from dialects that are allow-listed in the filter are bufferized. Other ops are left unbufferized. Note: This option requires allowUnknownOps = true.

To make use of dialectFilter, BufferizationOptions or BufferizationState must be passed to various helper functions.

The purpose of this change is to provide a better infrastructure for partial bufferization, which will be fully activated in a subsequent change.

Depends On D114671

[springerm]

update

[springerm]

update

[springerm]

update

[nicolasvasilache]

Please move this to test passes, similar to the parent commit.
Once this is more fleshed out they can replace the existing "partial" bufferization passes (and shoule be renamed to partial at that point).

[springerm]

update

[springerm]

update

[springerm]

rebase

[springerm]

address comments

[springerm]

add test case

[rriddle]

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

97 ↗

(On Diff #392644)

I've always found AllowListed to be an awkward mouthful, why not just isOpAllowed/isOpSupported/etc.

98–100 ↗

(On Diff #392644)

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

193 ↗

(On Diff #392644)

Leftover debugging here.

[springerm]

update

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

 /// analysis to determine which op operands and results may be bufferized in the
 /// 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();
 
+/// This is a pass that bufferizes only tensor dialect ops.
+std::unique_ptr<Pass> createLinalgComprehensiveTensorBufferizePass();
+
+/// This is a pass that bufferizes only vector dialect ops.
+std::unique_ptr<Pass> createLinalgComprehensiveVectorBufferizePass();
+
 /// 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.
 std::unique_ptr<OperationPass<FuncOp>> createLinalgGeneralizationPass();
 

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

Option<"useAlloca", "use-alloca", "bool",
            "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::createLinalgComprehensiveModuleBufferizePass()";
 }
 
+def LinalgComprehensiveTensorBufferize :
+    FunctionPass<"linalg-comprehensive-tensor-bufferize"> {
+  let summary = "Bufferize (tensor into memref) tensor dialect ops.";
+  let description = [{
+    This pass bufferizes only tensor dialect ops.
+    TODO: This will be moved to the bufferization dialect in the future.
+  }];
+  let constructor = "mlir::createLinalgComprehensiveTensorBufferizePass()";
+}
+
+def LinalgComprehensiveVectorBufferize :
+    FunctionPass<"linalg-comprehensive-vector-bufferize"> {
+  let summary = "Bufferize (tensor into memref) vector dialect ops.";
+  let description = [{
+    This pass bufferizes only vector dialect ops.
+    TODO: This will be moved to the bufferization dialect in the future.
+  }];
+  let constructor = "mlir::createLinalgComprehensiveVectorBufferizePass()";
+}
+
 def LinalgFoldUnitExtentDims : FunctionPass<"linalg-fold-unit-extent-dims"> {
   let summary = "Remove unit-extent dimension in Linalg ops on tensors";
   let constructor = "mlir::createLinalgFoldUnitExtentDimsPass()";
   let options = [
     Option<"foldOneTripLoopsOnly", "fold-one-trip-loops-only", "bool",
             /*default=*/"false",
            "Only folds the one-trip loops from Linalg ops on tensors "
            "(for testing purposes only)">

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

: public BufferizePassHelper<LinalgComprehensiveFunctionBufferizeBase<
           LinalgComprehensiveFunctionBufferize>> {
   LinalgComprehensiveFunctionBufferize() {}
 
   LinalgComprehensiveFunctionBufferize(
       const LinalgComprehensiveFunctionBufferize &p) {}
 
   void runOnFunction() final;
 };
+
+struct LinalgComprehensiveTensorBufferize
+    : public LinalgComprehensiveTensorBufferizeBase<
+          LinalgComprehensiveTensorBufferize> {
+  LinalgComprehensiveTensorBufferize() {}
+
+  LinalgComprehensiveTensorBufferize(
+      const LinalgComprehensiveTensorBufferize &p) {}
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<bufferization::BufferizationDialect, memref::MemRefDialect,
+                    tensor::TensorDialect>();
+    bufferization_ext::registerBufferizableOpInterfaceExternalModels(registry);
+    tensor_ext::registerBufferizableOpInterfaceExternalModels(registry);
+  }
+
+  void runOnFunction() override;
+};
+
+struct LinalgComprehensiveVectorBufferize
+    : public LinalgComprehensiveVectorBufferizeBase<
+          LinalgComprehensiveVectorBufferize> {
+  LinalgComprehensiveVectorBufferize() {}
+
+  LinalgComprehensiveVectorBufferize(
+      const LinalgComprehensiveVectorBufferize &p) {}
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<bufferization::BufferizationDialect, memref::MemRefDialect,
+                    vector::VectorDialect>();
+    bufferization_ext::registerBufferizableOpInterfaceExternalModels(registry);
+    vector_ext::registerBufferizableOpInterfaceExternalModels(registry);
+  }
+
+  void runOnFunction() override;
+};
 } // end namespace
 
 static void applyEnablingTransformations(Operation *op) {
   RewritePatternSet patterns(op->getContext());
   patterns.add<GeneralizePadTensorOpPattern>(op->getContext());
   (void)applyPatternsAndFoldGreedily(op, std::move(patterns));
 }
 
void LinalgComprehensiveFunctionBufferize::runOnFunction() {
   }
 
   if (options.testAnalysisOnly)
     return;
 
   this->runCleanupPipeline("builtin.func", getOperation());
 }
 
+void LinalgComprehensiveTensorBufferize::runOnFunction() {
+  BufferizationOptions options;
+  options.addPostAnalysisStep<tensor_ext::InplaceInsertSliceOpAnalysis>();
+  options.allowReturnMemref = true;
+  options.allowUnknownOps = true;
+
+  if (failed(runComprehensiveBufferize(getOperation(), options)))
+    signalPassFailure();
+}
+
+void LinalgComprehensiveVectorBufferize::runOnFunction() {
+  BufferizationOptions options;
+  options.allowReturnMemref = true;
+  options.allowUnknownOps = true;
+
+  if (failed(runComprehensiveBufferize(getOperation(), options)))
+    signalPassFailure();
+}
+
 std::unique_ptr<Pass> mlir::createLinalgComprehensiveModuleBufferizePass() {
   return std::make_unique<LinalgComprehensiveModuleBufferize>();
 }
 
 std::unique_ptr<Pass> mlir::createLinalgComprehensiveFunctionBufferizePass() {
   return std::make_unique<LinalgComprehensiveFunctionBufferize>();
 }
+
+std::unique_ptr<Pass> mlir::createLinalgComprehensiveTensorBufferizePass() {
+  return std::make_unique<LinalgComprehensiveTensorBufferize>();
+}
+
+std::unique_ptr<Pass> mlir::createLinalgComprehensiveVectorBufferizePass() {
+  return std::make_unique<LinalgComprehensiveVectorBufferize>();
+}