[mlir][sparse] refactor sparsification and bufferization pass into proper TD pass

[aartbik]

Summary

Registering the SparsificationAndBufferization into a proper TD pass
has the advantage that it can be invoked and tested in isolation. This
change also moves some bufferization specific set up from the pipeline
file into the pass file, keeping the logic more locally.

[Peiming]

mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp

22–24

Are these new header file intended (required by the generated file)?

[Peiming]

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

376

super nit, maybe bufferization-and-sparsification? (that should be the order if I remembered correctly}

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

376

ah, true, but the class is SparsificationAndBufferization

shall we do a general rename later (not to change the current diffs too much)

[Peiming]

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

376

Sure, not a blocker ;-)

[aartbik]

mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp

22–24

yeah due to registration that is now defined by TB code in the inc file

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h

 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 //===----------------------------------------------------------------------===//
 // Include the generated pass header (which needs some early definitions).
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
+
 namespace bufferization {
 struct OneShotBufferizationOptions;
 } // namespace bufferization
 
 /// Defines a parallelization strategy. Any independent loop is a candidate
 /// for parallelization. The loop is made parallel if (1) allowed by the
 /// strategy (e.g., AnyStorageOuterLoop considers either a dense or sparse
 /// outermost loop only), and (2) the generated code is an actual for-loop

 };
 
 void populateStorageSpecifierToLLVMPatterns(TypeConverter &converter,
                                             RewritePatternSet &patterns);
 std::unique_ptr<Pass> createStorageSpecifierToLLVMPass();
 
 //===----------------------------------------------------------------------===//
 // The mini-pipeline for sparsification and bufferization.
-//
-// Note that this mini-pipeline is not defined through the tablegen pass
-// mechanism, and, thus, is not individually available through the command-line.
-// It is solely used as part of the full sparse compiler pipeline.
 //===----------------------------------------------------------------------===//
 
+bufferization::OneShotBufferizationOptions
+getBufferizationOptionsForSparsification(bool analysisOnly);
+
+std::unique_ptr<Pass> createSparsificationAndBufferizationPass();
+
 std::unique_ptr<Pass> createSparsificationAndBufferizationPass(
     const bufferization::OneShotBufferizationOptions &bufferizationOptions,
     const SparsificationOptions &sparsificationOptions,
     const SparseTensorConversionOptions &sparseTensorConversionOptions,
     bool createSparseDeallocs, bool enableRuntimeLibrary,
     bool enableBufferInitialization, unsigned vectorLength,
     bool enableVLAVectorization, bool enableSIMDIndex32);
 

mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td

def StorageSpecifierToLLVM : Pass<"sparse-storage-specifier-to-llvm", "ModuleOp"> {
   let constructor = "mlir::createStorageSpecifierToLLVMPass()";
   let dependentDialects = [
     "arith::ArithDialect",
     "LLVM::LLVMDialect",
     "sparse_tensor::SparseTensorDialect",
   ];
 }
 
+def SparsificationAndBufferization : Pass<"sparsification-and-bufferization", "ModuleOp"> {
+  let summary = "Mini-pipeline that combines bufferization and sparsifiation";
+  let description = [{
+     This pass forms a mini-pipeline that combines bufferization and sparsifiation.
+  }];
+  let constructor = "mlir::createSparsificationAndBufferizationPass()";
+  let dependentDialects = [
+    "affine::AffineDialect",
+    "arith::ArithDialect",
+    "bufferization::BufferizationDialect",
+    "gpu::GPUDialect",
+    "LLVM::LLVMDialect",
+    "linalg::LinalgDialect",
+    "memref::MemRefDialect",
+    "scf::SCFDialect",
+    "sparse_tensor::SparseTensorDialect",
+  ];
+}
+
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES

mlir/lib/Dialect/SparseTensor/Pipelines/SparseTensorPipelines.cpp

 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
 
-using namespace mlir;
-using namespace mlir::sparse_tensor;
-
-/// Return configuration options for One-Shot Bufferize.
-static bufferization::OneShotBufferizationOptions
-getBufferizationOptions(bool analysisOnly) {
-  using namespace bufferization;
-  OneShotBufferizationOptions options;
-  options.bufferizeFunctionBoundaries = true;
-  // TODO(springerm): To spot memory leaks more easily, returning dense allocs
-  // should be disallowed.
-  options.allowReturnAllocs = true;
-  options.setFunctionBoundaryTypeConversion(LayoutMapOption::IdentityLayoutMap);
-  options.unknownTypeConverterFn = [](Value value, Attribute memorySpace,
-                                      const BufferizationOptions &options) {
-    return getMemRefTypeWithStaticIdentityLayout(
-        cast<TensorType>(value.getType()), memorySpace);
-  };
-  if (analysisOnly) {
-    options.testAnalysisOnly = true;
-    options.printConflicts = true;
-  }
-  return options;
-}
-
 //===----------------------------------------------------------------------===//
 // Pipeline implementation.
 //===----------------------------------------------------------------------===//
 
 void mlir::sparse_tensor::buildSparseCompiler(
     OpPassManager &pm, const SparseCompilerOptions &options) {
   pm.addNestedPass<func::FuncOp>(createLinalgGeneralizationPass());
   pm.addPass(createSparsificationAndBufferizationPass(
-      getBufferizationOptions(options.testBufferizationAnalysisOnly),
+      getBufferizationOptionsForSparsification(
+          options.testBufferizationAnalysisOnly),
       options.sparsificationOptions(), options.sparseTensorConversionOptions(),
       options.createSparseDeallocs, options.enableRuntimeLibrary,
       options.enableBufferInitialization, options.vectorLength,
       /*enableVLAVectorization=*/options.armSVE,
       /*enableSIMDIndex32=*/options.force32BitVectorIndices));
   if (options.testBufferizationAnalysisOnly)
     return;
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());

mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp

 //===- SparsificationAndBufferizationPass.cpp - Tensor to Memref Lowering -===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
 #include "mlir/Dialect/Bufferization/Transforms/Passes.h"
 #include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
 
 using namespace mlir;
-using namespace mlir::func;
 
 namespace mlir {
+
+#define GEN_PASS_DEF_SPARSIFICATIONANDBUFFERIZATION
+#include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
+
 namespace sparse_tensor {
 
 /// Return `true` if one of the given types is a sparse tensor type.
 static bool containsSparseTensor(TypeRange types) {
   for (Type t : types)
     if (getSparseTensorEncoding(t))
       return true;
   return false;
 }
 
 /// A pass that lowers tensor ops to memref ops, regardless of whether they are
 /// dense or sparse.
 ///
 /// One-Shot Analysis is used to detect RaW conflicts and to insert buffer
 /// copies of the tensor level (`insertTensorCopies`). Afterwards, the lowering
 /// of tensor ops to memref ops follows a different code path depending on
 /// whether the op is sparse or dense:
 ///
 /// * Sparse tensor ops are lowered through Sparsification and follow-up pass
 ///   that lowers sparse_tensor dialect ops.
 /// * Dense tensor ops are lowered through BufferizableOpInterface
 ///   implementations.
 class SparsificationAndBufferizationPass
-    : public PassWrapper<SparsificationAndBufferizationPass,
-                         OperationPass<ModuleOp>> {
+    : public impl::SparsificationAndBufferizationBase<
+          SparsificationAndBufferizationPass> {
 public:
   SparsificationAndBufferizationPass(
       const bufferization::OneShotBufferizationOptions &bufferizationOptions,
       const SparsificationOptions &sparsificationOptions,
       const SparseTensorConversionOptions &sparseTensorConversionOptions,
       bool createSparseDeallocs, bool enableRuntimeLibrary,
       bool enableBufferInitialization, unsigned vectorLength,
       bool enableVLAVectorization, bool enableSIMDIndex32)
if (failed(bufferization::bufferizeOp(getOperation(), bufferizationOptions,
                                           /*copyBeforeWrite=*/false,
                                           &denseOpFilter)))
       return failure();
 
     bufferization::removeBufferizationAttributesInModule(getOperation());
     return success();
   }
 
-  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
-    registry.insert<bufferization::BufferizationDialect>();
-    registry.insert<gpu::GPUDialect>();
-    registry.insert<LLVM::LLVMDialect>();
-  }
-
   void runOnOperation() override {
     {
       // Run enabling transformations.
       OpPassManager pm("builtin.module");
       pm.addPass(createPreSparsificationRewritePass());
       pm.addNestedPass<func::FuncOp>(
           bufferization::createEmptyTensorToAllocTensorPass());
       if (failed(runPipeline(pm, getOperation())))
private:
   unsigned vectorLength;
   bool enableVLAVectorization;
   bool enableSIMDIndex32;
 };
 
 } // namespace sparse_tensor
 } // namespace mlir
 
-std::unique_ptr<Pass> mlir::createSparsificationAndBufferizationPass(
+mlir::bufferization::OneShotBufferizationOptions
+mlir::getBufferizationOptionsForSparsification(bool analysisOnly) {
+  using namespace mlir::bufferization;
+  OneShotBufferizationOptions options;
+  options.bufferizeFunctionBoundaries = true;
+  // TODO(springerm): To spot memory leaks more easily, returning dense allocs
+  // should be disallowed.
+  options.allowReturnAllocs = true;
+  options.setFunctionBoundaryTypeConversion(LayoutMapOption::IdentityLayoutMap);
+  options.unknownTypeConverterFn = [](Value value, Attribute memorySpace,
+                                      const BufferizationOptions &options) {
+    return getMemRefTypeWithStaticIdentityLayout(
+        cast<TensorType>(value.getType()), memorySpace);
+  };
+  if (analysisOnly) {
+    options.testAnalysisOnly = true;
+    options.printConflicts = true;
+  }
+  return options;
+}
+
+std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass() {
+  SparsificationOptions sparseOptions;
+  SparseTensorConversionOptions convOptions;
+  return createSparsificationAndBufferizationPass(
+      getBufferizationOptionsForSparsification(/*analysisOnly=*/false),
+      sparseOptions, convOptions,
+      /*createSparseDeallocs=*/false,
+      /*enableRuntimeLibrary=*/false,
+      /*enableBufferInitialization=*/false,
+      /*vectorLength=*/0,
+      /*enableVLAVectorization=*/false,
+      /*enableSIMDIndex32=*/false);
+}
+
+std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass(
     const bufferization::OneShotBufferizationOptions &bufferizationOptions,
     const SparsificationOptions &sparsificationOptions,
     const SparseTensorConversionOptions &sparseTensorConversionOptions,
     bool createSparseDeallocs, bool enableRuntimeLibrary,
     bool enableBufferInitialization, unsigned vectorLength,
     bool enableVLAVectorization, bool enableSIMDIndex32) {
   return std::make_unique<
       mlir::sparse_tensor::SparsificationAndBufferizationPass>(
       bufferizationOptions, sparsificationOptions,
       sparseTensorConversionOptions, createSparseDeallocs, enableRuntimeLibrary,
       enableBufferInitialization, vectorLength, enableVLAVectorization,
       enableSIMDIndex32);
 }