[mlir][sparse] Introduce proper sparsification passes

[aartbik]

Summary

This revision migrates more code from Linalg into the new permanent home of
SparseTensor. It replaces the test passes with proper compiler passes.

NOTE: the actual removal of the last glue and clutter in Linalg will follow

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 ///   non-local transformation effects. This allows creating custom fused
 ///   transformations where patterns can be ordered and applied at a finer
 ///   granularity than a sequence of traditional compiler passes.
 LogicalResult applyStagedPatterns(
     Operation *op, ArrayRef<FrozenRewritePatternSet> stage1Patterns,
     const FrozenRewritePatternSet &stage2Patterns,
     function_ref<LogicalResult(Operation *)> stage3Lambda = nullptr);
 
-//===----------------------------------------------------------------------===//
-// Support for sparse tensor code generation.
-//
-// The sparse compiler part of MLIR lowers a tensor expression formulated as a
-// Linalg operation into a sequence of loops depending on what dimensions of the
-// tensors are marked dense or sparse. The generated code distinguishes between:
-// (1) for-loops that iterate over a single dense dimension,
-// (2) for-loops that iterate over a single sparse dimension,
-// (3) while-loops that co-iterate over several sparse dimensions.
-// The for-loops may be subsequently optimized for parallel or vector execution.
-//
-// For more details, the Dialect/Linalg/Transforms/Sparsification.cpp file.
-//===----------------------------------------------------------------------===//
-
-/// Defines a parallelization strategy. Any implicit loop in the Linalg
-/// operation that is marked "parallel" (thus not "reduction") 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
-/// (and not a co-iterating while-loop).
-enum class SparseParallelizationStrategy {
-  kNone,
-  kDenseOuterLoop,
-  kAnyStorageOuterLoop,
-  kDenseAnyLoop,
-  kAnyStorageAnyLoop
-  // TODO: support reduction parallelization too?
-};
-
-/// Defines a vectorization strategy. Any implicit inner loop in the Linalg
-/// operation is a candidate (full SIMD for "parallel" loops and horizontal
-/// SIMD for "reduction" loops). A loop is actually vectorized if (1) allowed
-/// by the strategy, and (2) the emitted code is an actual for-loop (and not
-/// a co-iterating while-loop).
-enum class SparseVectorizationStrategy {
-  kNone,
-  kDenseInnerLoop,
-  kAnyStorageInnerLoop
-};
-
-/// Defines a type for "pointer" and "index" storage in the sparse storage
-/// scheme, with a choice between the native platform-dependent index width
-/// or any of 64-/32-/16-/8-bit integers. A narrow width obviously reduces
-/// the memory footprint of the sparse storage scheme, but the width should
-/// suffice to define the total required range (viz. the maximum number of
-/// stored entries per indirection level for the "pointers" and the maximum
-/// value of each tensor index over all dimensions for the "indices").
-enum class SparseIntType { kNative, kI64, kI32, kI16, kI8 };
-
-/// Sparsification options.
-struct SparsificationOptions {
-  SparsificationOptions(SparseParallelizationStrategy p,
-                        SparseVectorizationStrategy v, unsigned vl,
-                        SparseIntType pt, SparseIntType it, bool fo)
-      : parallelizationStrategy(p), vectorizationStrategy(v), vectorLength(vl),
-        ptrType(pt), indType(it), fastOutput(fo) {}
-  SparsificationOptions()
-      : SparsificationOptions(SparseParallelizationStrategy::kNone,
-                              SparseVectorizationStrategy::kNone, 1u,
-                              SparseIntType::kNative, SparseIntType::kNative,
-                              false) {}
-  SparseParallelizationStrategy parallelizationStrategy;
-  SparseVectorizationStrategy vectorizationStrategy;
-  unsigned vectorLength;
-  SparseIntType ptrType;
-  SparseIntType indType;
-  bool fastOutput; // experimental: fast output buffers
-};
-
-/// Sets up sparsification rewriting rules with the given options.
-void populateSparsificationPatterns(
-    RewritePatternSet &patterns,
-    const SparsificationOptions &options = SparsificationOptions());
-
 } // namespace linalg
 } // namespace mlir
 
 #endif // DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H_

mlir/include/mlir/Dialect/SparseTensor/CMakeLists.txt

 add_subdirectory(IR)
+add_subdirectory(Transforms)

mlir/include/mlir/Dialect/SparseTensor/Transforms/CMakeLists.txt

+set(LLVM_TARGET_DEFINITIONS Passes.td)
+mlir_tablegen(Passes.h.inc -gen-pass-decls -name SparseTensor)
+add_public_tablegen_target(MLIRSparseTensorPassIncGen)
+
+add_mlir_doc(Passes SparseTensorPasses ./ -gen-pass-doc)

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

+//===- Passes.h - Sparse tensor pass entry points ---------------*- C++ -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines prototypes of all sparse tensor passes.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES_H_
+#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES_H_
+
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+
+namespace mlir {
+
+/// 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
+/// (and not a co-iterating while-loop).
+enum class SparseParallelizationStrategy {
+  kNone,
+  kDenseOuterLoop,
+  kAnyStorageOuterLoop,
+  kDenseAnyLoop,
+  kAnyStorageAnyLoop
+  // TODO: support reduction parallelization too?
+};
+
+/// Defines a vectorization strategy. Any inner loop is a candidate (full SIMD
+/// for parallel loops and horizontal SIMD for reduction loops). A loop is
+/// actually vectorized if (1) allowed by the strategy, and (2) the emitted
+/// code is an actual for-loop (and not a co-iterating while-loop).
+enum class SparseVectorizationStrategy {
+  kNone,
+  kDenseInnerLoop,
+  kAnyStorageInnerLoop
+};
+
+/// Defines a type for "pointer" and "index" storage in the sparse storage
+/// scheme, with a choice between the native platform-dependent index width
+/// or any of 64-/32-/16-/8-bit integers. A narrow width obviously reduces
+/// the memory footprint of the sparse storage scheme, but the width should
+/// suffice to define the total required range (viz. the maximum number of
+/// stored entries per indirection level for the "pointers" and the maximum
+/// value of each tensor index over all dimensions for the "indices").
+enum class SparseIntType { kNative, kI64, kI32, kI16, kI8 };
+
+/// Sparsification options.
+struct SparsificationOptions {
+  SparsificationOptions(SparseParallelizationStrategy p,
+                        SparseVectorizationStrategy v, unsigned vl,
+                        SparseIntType pt, SparseIntType it, bool fo)
+      : parallelizationStrategy(p), vectorizationStrategy(v), vectorLength(vl),
+        ptrType(pt), indType(it), fastOutput(fo) {}
+  SparsificationOptions()
+      : SparsificationOptions(SparseParallelizationStrategy::kNone,
+                              SparseVectorizationStrategy::kNone, 1u,
+                              SparseIntType::kNative, SparseIntType::kNative,
+                              false) {}
+  SparseParallelizationStrategy parallelizationStrategy;
+  SparseVectorizationStrategy vectorizationStrategy;
+  unsigned vectorLength;
+  SparseIntType ptrType;
+  SparseIntType indType;
+  bool fastOutput; // experimental: fast output buffers
+};
+
+/// Sets up sparsification rewriting rules with the given options.
+void populateSparsificationPatterns(
+    RewritePatternSet &patterns,
+    const SparsificationOptions &options = SparsificationOptions());
+
+/// Sets up sparse tensor conversion rules.
+void populateSparseTensorConversionPatterns(RewritePatternSet &patterns);
+
+std::unique_ptr<Pass> createSparsificationPass();
+std::unique_ptr<Pass> createSparseTensorConversionPass();
+
+//===----------------------------------------------------------------------===//
+// Registration.
+//===----------------------------------------------------------------------===//
+
+/// Generate the code for registering passes.
+#define GEN_PASS_REGISTRATION
+#include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
+
+} // namespace mlir
+
+#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES_H_

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

+//===-- Passes.td - Sparse tensor pass definition file -----*- tablegen -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES
+#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def Sparsification : Pass<"sparsification", "ModuleOp"> {
+  let summary = "Automatically generate sparse tensor code from annotations";
+  let constructor = "mlir::createSparsificationPass()";
+  let dependentDialects = [
+    "LLVM::LLVMDialect",
+    "memref::MemRefDialect",
+    "scf::SCFDialect",
+    "sparse_tensor::SparseTensorDialect",
+    "vector::VectorDialect",
+  ];
+}
+
+def SparseTensorConversion : Pass<"sparse-tensor-conversion", "ModuleOp"> {
+  let summary = "Apply conversion rules to sparse tensors";
+  let constructor = "mlir::createSparseTensorConversionPass()";
+  let dependentDialects = [
+    "LLVM::LLVMDialect",
+    "memref::MemRefDialect",
+    "scf::SCFDialect",
+    "sparse_tensor::SparseTensorDialect",
+    "vector::VectorDialect",
+  ];
+}
+
+#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES
+

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

-//===- Transforms.h - Sparse tensor transformations -------------*- C++ -*-===//
-//
-// 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
-//
-//===----------------------------------------------------------------------===//
-
-#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_TRANSFORMS_H_
-#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_TRANSFORMS_H_
-
-#include "mlir/IR/PatternMatch.h"
-
-namespace mlir {
-namespace sparse_tensor {
-
-/// Sets up sparsification conversion rules with the given options.
-void populateSparsificationConversionPatterns(RewritePatternSet &patterns);
-
-} // namespace sparse_tensor
-} // namespace mlir
-
-#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_TRANSFORMS_H_

mlir/include/mlir/InitAllPasses.h

 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/LLVMIR/Transforms/Passes.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/Quant/Passes.h"
 #include "mlir/Dialect/SCF/Passes.h"
 #include "mlir/Dialect/SPIRV/Transforms/Passes.h"
 #include "mlir/Dialect/Shape/Transforms/Passes.h"
+#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/StandardOps/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tosa/Transforms/Passes.h"
 #include "mlir/Transforms/Passes.h"
 
 #include <cstdlib>
 
 namespace mlir {
inline void registerAllPasses() {
 
   // Dialect passes
   registerAffinePasses();
   registerAsyncPasses();
   registerGPUPasses();
   registerGpuSerializeToCubinPass();
   registerGpuSerializeToHsacoPass();
   registerLinalgPasses();
+  registerSparseTensorPasses();
   LLVM::registerLLVMPasses();
   memref::registerMemRefPasses();
   quant::registerQuantPasses();
   registerSCFPasses();
   registerShapePasses();
   spirv::registerSPIRVPasses();
   registerStandardPasses();
   tensor::registerTensorPasses();
   tosa::registerTosaOptPasses();
 }
 
 } // namespace mlir
 
 #endif // MLIR_INITALLPASSES_H_

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRLinalgTransforms
   Bufferize.cpp
   CodegenStrategy.cpp
   Detensorize.cpp
   DropUnitDims.cpp
   ElementwiseToLinalg.cpp
   Fusion.cpp
   FusionOnTensors.cpp
   Generalization.cpp
   Hoisting.cpp
   Interchange.cpp
   Loops.cpp
   Promotion.cpp
-  Sparsification.cpp
   Tiling.cpp
   Transforms.cpp
   Vectorization.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   DEPENDS
add_mlir_dialect_library(MLIRLinalgTransforms
   MLIRMemRef
   MLIRLinalgAnalysis
   MLIRLinalgEDSC
   MLIRLinalg
   MLIRLinalgUtils
   MLIRSCF
   MLIRSCFTransforms
   MLIRPass
-  MLIRSparseTensor
   MLIRStandard
   MLIRStandardOpsTransforms
   MLIRStandardToLLVM
   MLIRTensor
   MLIRTransforms
   MLIRTransformUtils
   MLIRVector
   MLIRVectorToSCF
 )

mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRSparseTensorTransforms
-  SparseTensorLowering.cpp
+  Sparsification.cpp
+  SparseTensorConversion.cpp
+  SparseTensorPasses.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
 
+  DEPENDS
+  MLIRSparseTensorPassIncGen
+
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRLLVMIR
+  MLIRLinalg
+  MLIRLinalgTransforms
+  MLIRMemRef
   MLIRPass
+  MLIRSCF
   MLIRStandard
   MLIRSparseTensor
   MLIRTransforms
+  MLIRVector
 )

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

 // elaborate IR code that implements the primitives on the selected sparse
 // tensor storage schemes.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
-#include "mlir/Dialect/SparseTensor/Transforms/Transforms.h"
+#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 
 namespace {
 
 /// Returns function reference (first hit also inserts into module).
 static FlatSymbolRefAttr getFunc(Operation *op, StringRef name, Type result,
                                  ValueRange operands) {
   MLIRContext *context = op->getContext();
   auto module = op->getParentOfType<ModuleOp>();
   auto func = module.lookupSymbol<FuncOp>(name);
   if (!func) {
     OpBuilder moduleBuilder(module.getBodyRegion());
     moduleBuilder
         .create<FuncOp>(op->getLoc(), name,
                         FunctionType::get(context, operands.getTypes(), result))
         .setPrivate();
   }
   return SymbolRefAttr::get(context, name);
 }
 
 /// Sparse conversion rule to remove opaque pointer cast.
-class TensorFromPointerConverter
+class SparseTensorFromPointerConverter
     : public OpConversionPattern<sparse_tensor::FromPointerOp> {
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(sparse_tensor::FromPointerOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOp(op, operands[0]);
     return success();
   }
 };
 
 /// Sparse conversion rule for dimension accesses.
-class TensorToDimSizeConverter : public OpConversionPattern<memref::DimOp> {
+class SparseTensorToDimSizeConverter
+    : public OpConversionPattern<memref::DimOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(memref::DimOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     if (!operands[0].getType().isa<LLVM::LLVMPointerType>())
       return failure();
     Type resType = op.getType();
     StringRef name = "sparseDimSize";
     rewriter.replaceOpWithNewOp<CallOp>(
         op, resType, getFunc(op, name, resType, operands), operands);
     return success();
   }
 };
 
 /// Sparse conversion rule for pointer accesses.
-class TensorToPointersConverter
+class SparseTensorToPointersConverter
     : public OpConversionPattern<sparse_tensor::ToPointersOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(sparse_tensor::ToPointersOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     Type eltType = resType.cast<ShapedType>().getElementType();
else
       return failure();
     rewriter.replaceOpWithNewOp<CallOp>(
         op, resType, getFunc(op, name, resType, operands), operands);
     return success();
   }
 };
 
 /// Sparse conversion rule for index accesses.
-class TensorToIndicesConverter
+class SparseTensorToIndicesConverter
     : public OpConversionPattern<sparse_tensor::ToIndicesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(sparse_tensor::ToIndicesOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     Type eltType = resType.cast<ShapedType>().getElementType();
else
       return failure();
     rewriter.replaceOpWithNewOp<CallOp>(
         op, resType, getFunc(op, name, resType, operands), operands);
     return success();
   }
 };
 
 /// Sparse conversion rule for value accesses.
-class TensorToValuesConverter
+class SparseTensorToValuesConverter
     : public OpConversionPattern<sparse_tensor::ToValuesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(sparse_tensor::ToValuesOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     Type eltType = resType.cast<ShapedType>().getElementType();
matchAndRewrite(sparse_tensor::ToValuesOp op, ArrayRef<Value> operands,
     return success();
   }
 };
 
 } // namespace
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
-void sparse_tensor::populateSparsificationConversionPatterns(
-    RewritePatternSet &patterns) {
-  patterns.add<TensorFromPointerConverter, TensorToDimSizeConverter,
-               TensorToPointersConverter, TensorToIndicesConverter,
-               TensorToValuesConverter>(patterns.getContext());
+void mlir::populateSparseTensorConversionPatterns(RewritePatternSet &patterns) {
+  patterns.add<SparseTensorFromPointerConverter, SparseTensorToDimSizeConverter,
+               SparseTensorToPointersConverter, SparseTensorToIndicesConverter,
+               SparseTensorToValuesConverter>(patterns.getContext());
 }

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

+//===- SparsificationPass.cpp - Pass for autogen spares tensor code -------===//
+//
+// 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/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// Passes declaration.
+//===----------------------------------------------------------------------===//
+
+#define GEN_PASS_CLASSES
+#include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
+
+//===----------------------------------------------------------------------===//
+// Passes implementation.
+//===----------------------------------------------------------------------===//
+
+struct SparsificationPass : public SparsificationBase<SparsificationPass> {
+
+  SparsificationPass() = default;
+  SparsificationPass(const SparsificationPass &pass) {}
+
+  Option<int32_t> parallelization{
+      *this, "parallelization-strategy",
+      llvm::cl::desc("Set the parallelization strategy"), llvm::cl::init(0)};
+
+  Option<int32_t> vectorization{
+      *this, "vectorization-strategy",
+      llvm::cl::desc("Set the vectorization strategy"), llvm::cl::init(0)};
+
+  Option<int32_t> vectorLength{
+      *this, "vl", llvm::cl::desc("Set the vector length"), llvm::cl::init(1)};
+
+  Option<int32_t> ptrType{*this, "ptr-type",
+                          llvm::cl::desc("Set the pointer type"),
+                          llvm::cl::init(0)};
+
+  Option<int32_t> indType{*this, "ind-type",
+                          llvm::cl::desc("Set the index type"),
+                          llvm::cl::init(0)};
+
+  Option<bool> fastOutput{*this, "fast-output",
+                          llvm::cl::desc("Allows fast output buffers"),
+                          llvm::cl::init(false)};
+
+  /// Returns parallelization strategy given on command line.
+  SparseParallelizationStrategy parallelOption() {
+    switch (parallelization) {
+    default:
+      return SparseParallelizationStrategy::kNone;
+    case 1:
+      return SparseParallelizationStrategy::kDenseOuterLoop;
+    case 2:
+      return SparseParallelizationStrategy::kAnyStorageOuterLoop;
+    case 3:
+      return SparseParallelizationStrategy::kDenseAnyLoop;
+    case 4:
+      return SparseParallelizationStrategy::kAnyStorageAnyLoop;
+    }
+  }
+
+  /// Returns vectorization strategy given on command line.
+  SparseVectorizationStrategy vectorOption() {
+    switch (vectorization) {
+    default:
+      return SparseVectorizationStrategy::kNone;
+    case 1:
+      return SparseVectorizationStrategy::kDenseInnerLoop;
+    case 2:
+      return SparseVectorizationStrategy::kAnyStorageInnerLoop;
+    }
+  }
+
+  /// Returns the requested integer type.
+  SparseIntType typeOption(int32_t option) {
+    switch (option) {
+    default:
+      return SparseIntType::kNative;
+    case 1:
+      return SparseIntType::kI64;
+    case 2:
+      return SparseIntType::kI32;
+    case 3:
+      return SparseIntType::kI16;
+    case 4:
+      return SparseIntType::kI8;
+    }
+  }
+
+  void runOnOperation() override {
+    auto *ctx = &getContext();
+    RewritePatternSet patterns(ctx);
+    // Translate strategy flags to strategy options.
+    SparsificationOptions options(parallelOption(), vectorOption(),
+                                  vectorLength, typeOption(ptrType),
+                                  typeOption(indType), fastOutput);
+    // Apply rewriting.
+    populateSparsificationPatterns(patterns, options);
+    vector::populateVectorToVectorCanonicalizationPatterns(patterns);
+    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+  }
+};
+
+struct SparseTensorConversionPass
+    : public SparseTensorConversionBase<SparseTensorConversionPass> {
+  void runOnOperation() override {
+    auto *ctx = &getContext();
+    RewritePatternSet conversionPatterns(ctx);
+    ConversionTarget target(*ctx);
+    target
+        .addIllegalOp<sparse_tensor::FromPointerOp, sparse_tensor::ToPointersOp,
+                      sparse_tensor::ToIndicesOp, sparse_tensor::ToValuesOp>();
+    target.addLegalOp<CallOp>();
+    populateSparseTensorConversionPatterns(conversionPatterns);
+    if (failed(applyPartialConversion(getOperation(), target,
+                                      std::move(conversionPatterns))))
+      signalPassFailure();
+  }
+};
+
+} // end anonymous namespace
+
+std::unique_ptr<Pass> mlir::createSparsificationPass() {
+  return std::make_unique<SparsificationPass>();
+}
+
+std::unique_ptr<Pass> mlir::createSparseTensorConversionPass() {
+  return std::make_unique<SparseTensorConversionPass>();
+}

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

 // [Kjolstad20] Fredrik Berg Kjolstad. Sparse Tensor Algebra Compilation.
 // PhD thesis, MIT, February, 2020 (tensor-compiler.org).
 //
 // Implementation detail: We use llvm::SmallVector for vectors with
 // variable lengths and std::vector for vectors with fixed lengths.
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
-#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+#include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/Matchers.h"
+#include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
 
 namespace {
 
 enum class Kind { kTensor, kInvariant, kMulF, kMulI, kAddF, kAddI };
 enum class Dim { kSparse, kDense, kUndef };
 
private:
   std::vector<std::vector<Dim>> dims;
   llvm::SmallVector<TensorExp, 32> tensorExps;
   llvm::SmallVector<LatPoint, 16> latPoints;
   llvm::SmallVector<SmallVector<unsigned, 16>, 8> latSets;
 };
 
 // Code generation.
 struct CodeGen {
-  CodeGen(linalg::SparsificationOptions o, unsigned numTensors,
-          unsigned numLoops)
+  CodeGen(mlir::SparsificationOptions o, unsigned numTensors, unsigned numLoops)
       : options(o), loops(numLoops), sizes(numLoops), buffers(numTensors),
         pointers(numTensors, std::vector<Value>(numLoops)),
         indices(numTensors, std::vector<Value>(numLoops)),
         highs(numTensors, std::vector<Value>(numLoops)),
         pidxs(numTensors, std::vector<Value>(numLoops)),
         idxs(numTensors, std::vector<Value>(numLoops)), redExp(-1u), redVal(),
         curVecLength(1), curVecMask() {}
   /// Sparsification options.
-  linalg::SparsificationOptions options;
+  mlir::SparsificationOptions options;
   /// Universal dense indices and upper bounds (by index). The loops array
   /// is updated with the value of the universal dense index in the current
   /// loop. The sizes array is set once with the inferred dimension sizes.
   std::vector<Value> loops;
   std::vector<Value> sizes;
   /// Buffers for storing dense and sparse numerical values (by tensor).
   /// This array is set once during bufferization of all tensors.
   std::vector<Value> buffers;
static unsigned buildLattices(Merger &merger, linalg::GenericOp op,
   case Kind::kAddF:
   case Kind::kAddI:
     return merger.takeDisj(kind, s0, s1);
   }
   llvm_unreachable("unexpected expression kind");
 }
 
 /// Maps sparse integer option to actual integral storage type.
-static Type genIntType(PatternRewriter &rewriter, linalg::SparseIntType tp) {
+static Type genIntType(PatternRewriter &rewriter, SparseIntType tp) {
   switch (tp) {
-  case linalg::SparseIntType::kNative:
+  case SparseIntType::kNative:
     return rewriter.getIndexType();
-  case linalg::SparseIntType::kI64:
+  case SparseIntType::kI64:
     return rewriter.getIntegerType(64);
-  case linalg::SparseIntType::kI32:
+  case SparseIntType::kI32:
     return rewriter.getIntegerType(32);
-  case linalg::SparseIntType::kI16:
+  case SparseIntType::kI16:
     return rewriter.getIntegerType(16);
-  case linalg::SparseIntType::kI8:
+  case SparseIntType::kI8:
     return rewriter.getIntegerType(8);
   }
   llvm_unreachable("unexpected SparseIntType");
 }
 
 /// Generates buffer for the output tensor.
 static Value genOutputBuffer(CodeGen &codegen, PatternRewriter &rewriter,
                              linalg::GenericOp op, MemRefType denseTp,
static bool genInit(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
   return needsUniv;
 }
 
 /// Returns vectorization strategy. Any implicit inner loop in the Linalg
 /// operation is a candidate. Whether it is actually converted to SIMD code
 /// depends on the requested strategy.
 static bool isVectorFor(CodeGen &codegen, bool isInner, bool isSparse) {
   switch (codegen.options.vectorizationStrategy) {
-  case linalg::SparseVectorizationStrategy::kNone:
+  case SparseVectorizationStrategy::kNone:
     return false;
-  case linalg::SparseVectorizationStrategy::kDenseInnerLoop:
+  case SparseVectorizationStrategy::kDenseInnerLoop:
     return isInner && !isSparse;
-  case linalg::SparseVectorizationStrategy::kAnyStorageInnerLoop:
+  case SparseVectorizationStrategy::kAnyStorageInnerLoop:
     return isInner;
   }
   llvm_unreachable("unexpected vectorization strategy");
 }
 
 /// Returns parallelization strategy. Any implicit loop in the Linalg operation
 /// that is marked "parallel" is a candidate. Whether it is actually converted
 /// to a parallel operation depends on the requested strategy.
 static bool isParallelFor(CodeGen &codegen, bool isOuter, bool isReduction,
                           bool isSparse, bool isVector) {
   switch (codegen.options.parallelizationStrategy) {
-  case linalg::SparseParallelizationStrategy::kNone:
+  case SparseParallelizationStrategy::kNone:
     return false;
-  case linalg::SparseParallelizationStrategy::kDenseOuterLoop:
+  case SparseParallelizationStrategy::kDenseOuterLoop:
     return isOuter && !isSparse && !isReduction && !isVector;
-  case linalg::SparseParallelizationStrategy::kAnyStorageOuterLoop:
+  case SparseParallelizationStrategy::kAnyStorageOuterLoop:
     return isOuter && !isReduction && !isVector;
-  case linalg::SparseParallelizationStrategy::kDenseAnyLoop:
+  case SparseParallelizationStrategy::kDenseAnyLoop:
     return !isSparse && !isReduction && !isVector;
-  case linalg::SparseParallelizationStrategy::kAnyStorageAnyLoop:
+  case SparseParallelizationStrategy::kAnyStorageAnyLoop:
     return !isReduction && !isVector;
   }
   llvm_unreachable("unexpected parallelization strategy");
 }
 
 /// Checks unit strides for dense tensors. The iteration graph may have ignored
 /// dense access patterns in order to avoid cycles (sparse access patterns are
 /// always placed innermost), but that means dense access has become strided.
static void genStmt(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
   codegen.loops[idx] = Value();
 }
 
 namespace {
 
 /// Sparse rewriting rule for generic Lingalg operation.
 struct GenericOpSparsifier : public OpRewritePattern<linalg::GenericOp> {
 public:
-  GenericOpSparsifier(MLIRContext *context, linalg::SparsificationOptions o)
+  GenericOpSparsifier(MLIRContext *context, SparsificationOptions o)
       : OpRewritePattern<linalg::GenericOp>(context), options(o) {}
 
   LogicalResult matchAndRewrite(linalg::GenericOp op,
                                 PatternRewriter &rewriter) const override {
     // Detects sparse annotations and translate the per-dimension sparsity
     // information for all tensors to loop indices in the kernel.
     if (!op.hasSparseSemantics())
       return failure();
LogicalResult matchAndRewrite(linalg::GenericOp op,
     Value result = rewriter.create<memref::TensorLoadOp>(
         op.getLoc(), codegen.buffers.back());
     rewriter.replaceOp(op, result);
     return success();
   }
 
 private:
   /// Options to control sparse code generation.
-  linalg::SparsificationOptions options;
+  SparsificationOptions options;
 };
 
 } // namespace
 
 /// Populates the given patterns list with rewriting rules required for
 /// the sparsification of linear algebra operations.
-void linalg::populateSparsificationPatterns(
+void mlir::populateSparsificationPatterns(
     RewritePatternSet &patterns, const SparsificationOptions &options) {
   patterns.add<GenericOpSparsifier>(patterns.getContext(), options);
 }

mlir/test/Dialect/SparseTensor/conversion.mlir

-// RUN: mlir-opt --test-sparsification="lower" %s | FileCheck %s
+// RUN: mlir-opt --sparse-tensor-conversion %s | FileCheck %s
 
 !SparseTensor = type !llvm.ptr<i8>
 
 // CHECK-LABEL: func @sparse_pointers(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: %[[C:.*]] = constant 1 : index
 //       CHECK: %[[T:.*]] = call @sparsePointers64(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
 //       CHECK: return %[[T]] : memref<?xindex>

mlir/test/Dialect/SparseTensor/sparse_1d.mlir

 // NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
-// RUN: mlir-opt %s -test-sparsification | FileCheck %s
+// RUN: mlir-opt %s -sparsification | FileCheck %s
 
 #trait_d = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> (i)>   // x (out)
   ],
   sparse = [
     [ "D" ],  // a

mlir/test/Dialect/SparseTensor/sparse_2d.mlir

 // NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
-// RUN: mlir-opt %s -test-sparsification | FileCheck %s
+// RUN: mlir-opt %s -sparsification | FileCheck %s
 
 #trait_dd = {
   indexing_maps = [
     affine_map<(i,j) -> (i,j)>,  // A
     affine_map<(i,j) -> (i,j)>,  // B
     affine_map<(i,j) -> (i,j)>   // X (out)
   ],
   sparse = [

mlir/test/Dialect/SparseTensor/sparse_3d.mlir

 // NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
-// RUN: mlir-opt %s -test-sparsification | FileCheck %s
+// RUN: mlir-opt %s -sparsification | FileCheck %s
 
 #trait_ddd = {
   indexing_maps = [
     affine_map<(i,j,k) -> (i,j,k)>,  // A
     affine_map<(i,j,k) -> (i,j,k)>,  // B
     affine_map<(i,j,k) -> (i,j,k)>   // X (out)
   ],
   sparse = [

mlir/test/Dialect/SparseTensor/sparse_lower.mlir

-// RUN: mlir-opt %s -test-sparsification | \
+// RUN: mlir-opt %s -sparsification | \
 // RUN:   FileCheck %s --check-prefix=CHECK-HIR
 //
-// RUN: mlir-opt %s -test-sparsification="lower" --convert-linalg-to-loops | \
+// RUN: mlir-opt %s -sparsification \
+// RUN:   --sparse-tensor-conversion --convert-linalg-to-loops | \
 // RUN:   FileCheck %s --check-prefix=CHECK-MIR
 //
-// RUN: mlir-opt %s -test-sparsification="lower" --convert-linalg-to-loops \
+// RUN: mlir-opt %s -sparsification \
+// RUN:   --sparse-tensor-conversion --convert-linalg-to-loops \
 // RUN:   --func-bufferize --tensor-constant-bufferize \
 // RUN:   --tensor-bufferize --finalizing-bufferize  | \
 // RUN:   FileCheck %s --check-prefix=CHECK-LIR
 //
-// RUN: mlir-opt %s -test-sparsification="lower fast-output" --convert-linalg-to-loops \
+// RUN: mlir-opt %s -sparsification="fast-output" \
+// RUN:   --sparse-tensor-conversion --convert-linalg-to-loops \
 // RUN:   --func-bufferize --tensor-constant-bufferize \
 // RUN:   --tensor-bufferize --finalizing-bufferize  | \
 // RUN:   FileCheck %s --check-prefix=CHECK-FAST
 
 #trait_matvec = {
   indexing_maps = [
     affine_map<(i,j) -> (i,j)>,  // A
     affine_map<(i,j) -> (j)>,    // b

mlir/test/Dialect/SparseTensor/sparse_nd.mlir

 // NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
-// RUN: mlir-opt %s -test-sparsification | FileCheck %s
+// RUN: mlir-opt %s -sparsification | FileCheck %s
 
 // Example with cyclic iteration graph with sparse and dense constraints,
 // but an acyclic iteration graph using sparse constraints only.
 #trait_mul = {
   indexing_maps = [
     affine_map<(i,j,k,l,m,n,o,p) -> (i,j,k,l,m,n,o,p)>,  // A
     affine_map<(i,j,k,l,m,n,o,p) -> (p,o,n,m,l,k,j,i)>,  // B
     affine_map<(i,j,k,l,m,n,o,p) -> (i,j,k,l,m,n,o,p)>   // X

mlir/test/Dialect/SparseTensor/sparse_parallel.mlir

-// RUN: mlir-opt %s -test-sparsification="parallelization-strategy=0" | \
+// RUN: mlir-opt %s -sparsification="parallelization-strategy=0" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR0
-// RUN: mlir-opt %s -test-sparsification="parallelization-strategy=1" | \
+// RUN: mlir-opt %s -sparsification="parallelization-strategy=1" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR1
-// RUN: mlir-opt %s -test-sparsification="parallelization-strategy=2" | \
+// RUN: mlir-opt %s -sparsification="parallelization-strategy=2" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR2
-// RUN: mlir-opt %s -test-sparsification="parallelization-strategy=3" | \
+// RUN: mlir-opt %s -sparsification="parallelization-strategy=3" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR3
-// RUN: mlir-opt %s -test-sparsification="parallelization-strategy=4" | \
+// RUN: mlir-opt %s -sparsification="parallelization-strategy=4" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR4
 
 #trait_dd = {
   indexing_maps = [
     affine_map<(i,j) -> (i,j)>,  // A
     affine_map<(i,j) -> (i,j)>   // X (out)
   ],
   sparse = [

mlir/test/Dialect/SparseTensor/sparse_storage.mlir

-// RUN: mlir-opt %s -test-sparsification="ptr-type=1 ind-type=1" | \
+// RUN: mlir-opt %s -sparsification="ptr-type=1 ind-type=1" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-TYPE0
-// RUN: mlir-opt %s -test-sparsification="ptr-type=1 ind-type=2" | \
+// RUN: mlir-opt %s -sparsification="ptr-type=1 ind-type=2" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-TYPE1
-// RUN: mlir-opt %s -test-sparsification="ptr-type=2 ind-type=1" | \
+// RUN: mlir-opt %s -sparsification="ptr-type=2 ind-type=1" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-TYPE2
-// RUN: mlir-opt %s -test-sparsification="ptr-type=2 ind-type=2" | \
+// RUN: mlir-opt %s -sparsification="ptr-type=2 ind-type=2" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-TYPE3
-// RUN: mlir-opt %s -test-sparsification="ptr-type=3 ind-type=3" | \
+// RUN: mlir-opt %s -sparsification="ptr-type=3 ind-type=3" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-TYPE4
-// RUN: mlir-opt %s -test-sparsification="ptr-type=4 ind-type=4" | \
+// RUN: mlir-opt %s -sparsification="ptr-type=4 ind-type=4" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-TYPE5
 
 #trait_mul_1d = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> (i)>,  // b
     affine_map<(i) -> (i)>   // x (out)
   ],

mlir/test/Dialect/SparseTensor/sparse_vector.mlir

-// RUN: mlir-opt %s -test-sparsification="vectorization-strategy=0 ptr-type=2 ind-type=2 vl=16" | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=0 ptr-type=2 ind-type=2 vl=16" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC0
-// RUN: mlir-opt %s -test-sparsification="vectorization-strategy=1 ptr-type=2 ind-type=2 vl=16" | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=1 ptr-type=2 ind-type=2 vl=16" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC1
-// RUN: mlir-opt %s -test-sparsification="vectorization-strategy=2 ptr-type=2 ind-type=2 vl=16" | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 ptr-type=2 ind-type=2 vl=16" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC2
-// RUN: mlir-opt %s -test-sparsification="vectorization-strategy=2 ptr-type=0 ind-type=0 vl=16" | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 ptr-type=0 ind-type=0 vl=16" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC3
 
 #trait_scale_d = {
   indexing_maps = [
     affine_map<(i) -> (i)>,  // a
     affine_map<(i) -> (i)>   // x (out)
   ],
   sparse = [

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir

 // RUN: mlir-opt %s \
-// RUN:   --test-sparsification="lower ptr-type=4 ind-type=4" \
+// RUN:   --sparsification="ptr-type=4 ind-type=4" --sparse-tensor-conversion \
 // RUN:   --convert-linalg-to-loops --convert-vector-to-scf --convert-scf-to-std \
 // RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
 // RUN:   --std-bufferize --finalizing-bufferize  \
 // RUN:   --convert-vector-to-llvm --convert-std-to-llvm | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/wide.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 //
 // RUN: mlir-opt %s \
-// RUN:   --test-sparsification="lower vectorization-strategy=2 ptr-type=4 ind-type=4 vl=16" \
+// RUN:   --sparsification="vectorization-strategy=2 ptr-type=4 ind-type=4 vl=16" --sparse-tensor-conversion \
 // RUN:   --convert-linalg-to-loops --convert-vector-to-scf --convert-scf-to-std \
 // RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
 // RUN:   --std-bufferize --finalizing-bufferize  \
 // RUN:   --convert-vector-to-llvm --convert-std-to-llvm | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/wide.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir

 // RUN: mlir-opt %s \
-// RUN:   --test-sparsification="lower ptr-type=2 ind-type=2 fast-output" \
+// RUN:   --sparsification="ptr-type=2 ind-type=2 fast-output" --sparse-tensor-conversion \
 // RUN:   --convert-linalg-to-loops --convert-vector-to-scf --convert-scf-to-std \
 // RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
 // RUN:   --std-bufferize --finalizing-bufferize  \
 // RUN:   --convert-vector-to-llvm --convert-std-to-llvm | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir

 // RUN: mlir-opt %s \
-// RUN:   --test-sparsification="lower" \
+// RUN:   --sparsification --sparse-tensor-conversion \
 // RUN:   --convert-linalg-to-loops --convert-vector-to-scf --convert-scf-to-std \
 // RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
 // RUN:   --std-bufferize --finalizing-bufferize  \
 // RUN:   --convert-vector-to-llvm --convert-std-to-llvm | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \

mlir/test/lib/Transforms/CMakeLists.txt

add_mlir_library(MLIRTestTransforms
   TestLoopUnrolling.cpp
   TestNumberOfExecutions.cpp
   TestOpaqueLoc.cpp
   TestPolynomialApproximation.cpp
   TestMemRefBoundCheck.cpp
   TestMemRefDependenceCheck.cpp
   TestMemRefStrideCalculation.cpp
   TestSCFUtils.cpp
-  TestSparsification.cpp
   TestVectorTransforms.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Transforms
 
   DEPENDS

mlir/test/lib/Transforms/TestSparsification.cpp

-//===- TestSparsification.cpp - Test sparsification of tensors ------------===//
-//
-// 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/LLVMIR/LLVMDialect.h"
-#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
-#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
-#include "mlir/Dialect/SparseTensor/Transforms/Transforms.h"
-#include "mlir/Dialect/Vector/VectorOps.h"
-#include "mlir/Pass/Pass.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-
-using namespace mlir;
-
-namespace {
-
-struct TestSparsification
-    : public PassWrapper<TestSparsification, OperationPass<ModuleOp>> {
-
-  TestSparsification() = default;
-  TestSparsification(const TestSparsification &pass) {}
-
-  Option<int32_t> parallelization{
-      *this, "parallelization-strategy",
-      llvm::cl::desc("Set the parallelization strategy"), llvm::cl::init(0)};
-
-  Option<int32_t> vectorization{
-      *this, "vectorization-strategy",
-      llvm::cl::desc("Set the vectorization strategy"), llvm::cl::init(0)};
-
-  Option<int32_t> vectorLength{
-      *this, "vl", llvm::cl::desc("Set the vector length"), llvm::cl::init(1)};
-
-  Option<int32_t> ptrType{*this, "ptr-type",
-                          llvm::cl::desc("Set the pointer type"),
-                          llvm::cl::init(0)};
-
-  Option<int32_t> indType{*this, "ind-type",
-                          llvm::cl::desc("Set the index type"),
-                          llvm::cl::init(0)};
-
-  Option<bool> fastOutput{*this, "fast-output",
-                          llvm::cl::desc("Allows fast output buffers"),
-                          llvm::cl::init(false)};
-
-  Option<bool> lower{*this, "lower", llvm::cl::desc("Lower sparse primitives"),
-                     llvm::cl::init(false)};
-
-  /// Registers all dialects required by testing.
-  void getDependentDialects(DialectRegistry &registry) const override {
-    registry.insert<memref::MemRefDialect, scf::SCFDialect,
-                    sparse_tensor::SparseTensorDialect, vector::VectorDialect,
-                    LLVM::LLVMDialect>();
-  }
-
-  /// Returns parallelization strategy given on command line.
-  linalg::SparseParallelizationStrategy parallelOption() {
-    switch (parallelization) {
-    default:
-      return linalg::SparseParallelizationStrategy::kNone;
-    case 1:
-      return linalg::SparseParallelizationStrategy::kDenseOuterLoop;
-    case 2:
-      return linalg::SparseParallelizationStrategy::kAnyStorageOuterLoop;
-    case 3:
-      return linalg::SparseParallelizationStrategy::kDenseAnyLoop;
-    case 4:
-      return linalg::SparseParallelizationStrategy::kAnyStorageAnyLoop;
-    }
-  }
-
-  /// Returns vectorization strategy given on command line.
-  linalg::SparseVectorizationStrategy vectorOption() {
-    switch (vectorization) {
-    default:
-      return linalg::SparseVectorizationStrategy::kNone;
-    case 1:
-      return linalg::SparseVectorizationStrategy::kDenseInnerLoop;
-    case 2:
-      return linalg::SparseVectorizationStrategy::kAnyStorageInnerLoop;
-    }
-  }
-
-  /// Returns the requested integer type.
-  linalg::SparseIntType typeOption(int32_t option) {
-    switch (option) {
-    default:
-      return linalg::SparseIntType::kNative;
-    case 1:
-      return linalg::SparseIntType::kI64;
-    case 2:
-      return linalg::SparseIntType::kI32;
-    case 3:
-      return linalg::SparseIntType::kI16;
-    case 4:
-      return linalg::SparseIntType::kI8;
-    }
-  }
-
-  /// Runs the test on a function.
-  void runOnOperation() override {
-    auto *ctx = &getContext();
-    RewritePatternSet patterns(ctx);
-    // Translate strategy flags to strategy options.
-    linalg::SparsificationOptions options(parallelOption(), vectorOption(),
-                                          vectorLength, typeOption(ptrType),
-                                          typeOption(indType), fastOutput);
-    // Apply rewriting.
-    linalg::populateSparsificationPatterns(patterns, options);
-    vector::populateVectorToVectorCanonicalizationPatterns(patterns);
-    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
-    // Lower sparse primitives to calls into runtime support library.
-    if (lower) {
-      RewritePatternSet conversionPatterns(ctx);
-      ConversionTarget target(*ctx);
-      target.addIllegalOp<
-          sparse_tensor::FromPointerOp, sparse_tensor::ToPointersOp,
-          sparse_tensor::ToIndicesOp, sparse_tensor::ToValuesOp>();
-      target.addLegalOp<CallOp>();
-      sparse_tensor::populateSparsificationConversionPatterns(
-          conversionPatterns);
-      if (failed(applyPartialConversion(getOperation(), target,
-                                        std::move(conversionPatterns))))
-        signalPassFailure();
-    }
-  }
-};
-
-} // end anonymous namespace
-
-namespace mlir {
-namespace test {
-
-void registerTestSparsification() {
-  PassRegistration<TestSparsification> sparsificationPass(
-      "test-sparsification", "Test automatic generation of sparse tensor code");
-}
-
-} // namespace test
-} // namespace mlir

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestMemRefStrideCalculation();
 void registerTestNumberOfBlockExecutionsPass();
 void registerTestNumberOfOperationExecutionsPass();
 void registerTestOpaqueLoc();
 void registerTestPDLByteCodePass();
 void registerTestPreparationPassWithAllowedMemrefResults();
 void registerTestRecursiveTypesPass();
 void registerTestSCFUtilsPass();
-void registerTestSparsification();
 void registerTestVectorConversions();
 } // namespace test
 } // namespace mlir
 
 #ifdef MLIR_INCLUDE_TESTS
 void registerTestPasses() {
   registerConvertToTargetEnvPass();
   registerPassManagerTestPass();
#endif
   test::registerTestMemRefDependenceCheck();
   test::registerTestMemRefStrideCalculation();
   test::registerTestNumberOfBlockExecutionsPass();
   test::registerTestNumberOfOperationExecutionsPass();
   test::registerTestOpaqueLoc();
   test::registerTestPDLByteCodePass();
   test::registerTestRecursiveTypesPass();
   test::registerTestSCFUtilsPass();
-  test::registerTestSparsification();
   test::registerTestVectorConversions();
 }
 #endif
 
 int main(int argc, char **argv) {
   registerAllPasses();
 #ifdef MLIR_INCLUDE_TESTS
   registerTestPasses();