[mlir][sparse] Split SparseTensorRewrite into PreSparsificationRewrite and PostSparsificationRewrite.

[bixia]

Rename the two passes.

[wrengr]

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

62

I know this is just mimicking what was there before, but: why don't we pass the other boolean arguments too?

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

86–90

Would be better to rephrase this as PostSparsificationRewritePass(bool enableRuntimeLibrary, bool enableForeach, bool enableConvert) : enableRuntimeLibrary(enableRuntimeLibrary), enableForeach(enableForeach), enableConvert(enableConvert) {}

[aartbik]

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

158

I liked the original layout better, but I guess this is clang formatter?

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

1024–1034

empty line after big header section

[wrengr]

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

158

If it's a big deal, then you can always add // NOLINT just before newlines you want git-clang-format to leave alone. (Though it's best to only do that sparingly, of course)

[bixia]

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

158

Right, that was auto formatted.

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

62

The other two booleans are for testing, disabling rewriting for two ops to simplify FileCheck.
In the real compile pass pipeline, we should never disable these two ops.

[bixia]

Address review comment.

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

86–90

Per offline discussion, this syntax doesn't work.

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

void populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
                                          RewritePatternSet &patterns,
                                          bool enableBufferInitialization);
 
 std::unique_ptr<Pass> createSparseTensorCodegenPass();
 std::unique_ptr<Pass>
 createSparseTensorCodegenPass(bool enableBufferInitialization);
 
 //===----------------------------------------------------------------------===//
-// The SparseTensorRewriting pass.
+// The PreSparsificationRewriting pass.
 //===----------------------------------------------------------------------===//
 
-void populateSparseTensorRewriting(RewritePatternSet &patterns, bool enableRT,
-                                   bool enableForeach, bool enableConvert);
+void populatePreSparsificationRewriting(RewritePatternSet &patterns);
 
-std::unique_ptr<Pass> createSparseTensorRewritePass();
-std::unique_ptr<Pass> createSparseTensorRewritePass(bool enableRT,
-                                                    bool enableForeach = true,
-                                                    bool enableConvert = true);
+std::unique_ptr<Pass> createPreSparsificationRewritePass();
+
+//===----------------------------------------------------------------------===//
+// The PostSparsificationRewriting pass.
+//===----------------------------------------------------------------------===//
+
+void populatePostSparsificationRewriting(RewritePatternSet &patterns,
+                                         bool enableRT, bool enableForeach,
+                                         bool enableConvert);
+
+std::unique_ptr<Pass> createPostSparsificationRewritePass();
+std::unique_ptr<Pass>
+createPostSparsificationRewritePass(bool enableRT, bool enableForeach = true,
 
 //===----------------------------------------------------------------------===//
 // Other rewriting rules and passes.
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<Pass> createDenseBufferizationPass(
     const bufferization::OneShotBufferizationOptions &options);
 

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 SparseTensorRewrite : Pass<"sparse-tensor-rewrite", "ModuleOp"> {
+def PreSparsificationRewrite : Pass<"pre-sparsification-rewrite", "ModuleOp"> {
   let summary = "Applies sparse tensor rewriting rules prior to sparsification";
   let description = [{
     A pass that applies rewriting rules to sparse tensor operations prior
     to running the actual sparsification pass.
   }];
-  let constructor = "mlir::createSparseTensorRewritePass()";
+  let constructor = "mlir::createPreSparsificationRewritePass()";
   let dependentDialects = [
     "arith::ArithDialect",
     "bufferization::BufferizationDialect",
     "linalg::LinalgDialect",
     "memref::MemRefDialect",
     "scf::SCFDialect",
     "sparse_tensor::SparseTensorDialect",
   ];
-  let options = [
-    Option<"enableRuntimeLibrary", "enable-runtime-library", "bool",
-           "true", "Enable runtime library for manipulating sparse tensors">,
-    Option<"enableForeach", "enable-foreach", "bool",
-           "true", "Enable rewriting rules for the foreach operator">,
-    Option<"enableConvert", "enable-convert", "bool",
-           "true", "Enable rewriting rules for the convert operator">,
-  ];
 }
 
 def SparsificationPass : Pass<"sparsification", "ModuleOp"> {
   let summary = "Automatically generate sparse tensor code from sparse tensor types";
   let description = [{
     A pass that implements the core functionality of a **sparse compiler**.
     Each Linalg operation (MLIR's tensor index notation) that operates on
     sparse tensor types is converted into code in which the sparsity is
Option<"parallelization", "parallelization-strategy", "mlir::SparseParallelizationStrategy",
                         "dense-any-loop",
                         "Enable dense parallelization for any loop."),
              clEnumValN(mlir::SparseParallelizationStrategy::kAnyStorageAnyLoop,
                         "any-storage-any-loop",
                         "Enable sparse parallelization for any storage and loop."))}]>
   ];
 }
 
+def PostSparsificationRewrite : Pass<"post-sparsification-rewrite", "ModuleOp"> {
+  let summary = "Applies sparse tensor rewriting rules after sparsification";
+  let description = [{
+    A pass that applies rewriting rules to sparse tensor operations after
+    running the actual sparsification pass.
+  }];
+  let constructor = "mlir::createPostSparsificationRewritePass()";
+  let dependentDialects = [
+    "arith::ArithDialect",
+    "bufferization::BufferizationDialect",
+    "linalg::LinalgDialect",
+    "memref::MemRefDialect",
+    "scf::SCFDialect",
+    "sparse_tensor::SparseTensorDialect",
+  ];
+  let options = [
+    Option<"enableRuntimeLibrary", "enable-runtime-library", "bool",
+           "true", "Enable runtime library for manipulating sparse tensors">,
+    Option<"enableForeach", "enable-foreach", "bool",
+           "true", "Enable rewriting rules for the foreach operator">,
+    Option<"enableConvert", "enable-convert", "bool",
+           "true", "Enable rewriting rules for the convert operator">,
+  ];
+}
+
 def SparseTensorConversionPass : Pass<"sparse-tensor-conversion", "ModuleOp"> {
   let summary = "Convert sparse tensors and primitives to library calls";
   let description = [{
     A pass that converts sparse tensor primitives into calls into a runtime
     support library. Sparse tensor types are converted into opaque pointers
     to the underlying sparse storage schemes.
 
     The use of opaque pointers together with runtime support library keeps

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

 void mlir::sparse_tensor::buildSparseCompiler(
     OpPassManager &pm, const SparseCompilerOptions &options) {
   pm.addNestedPass<func::FuncOp>(createLinalgGeneralizationPass());
   pm.addPass(
       bufferization::createTensorCopyInsertionPass(getBufferizationOptions(
           /*analysisOnly=*/options.testBufferizationAnalysisOnly)));
   if (options.testBufferizationAnalysisOnly)
     return;
-  pm.addPass(createSparseTensorRewritePass(options.enableRuntimeLibrary));
+  pm.addPass(createPreSparsificationRewritePass());
   pm.addPass(createSparsificationPass(options.sparsificationOptions()));
+  pm.addPass(createPostSparsificationRewritePass(options.enableRuntimeLibrary));
   if (options.enableRuntimeLibrary) {
     pm.addPass(createSparseTensorConversionPass(
         options.sparseTensorConversionOptions()));
   } else {
     pm.addPass(
         createSparseTensorCodegenPass(options.enableBufferInitialization));
     pm.addPass(
         createSparseBufferRewritePass(options.enableBufferInitialization));

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

 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/Transforms/Transforms.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 namespace mlir {
-#define GEN_PASS_DEF_SPARSETENSORREWRITE
+#define GEN_PASS_DEF_PRESPARSIFICATIONREWRITE
 #define GEN_PASS_DEF_SPARSIFICATIONPASS
+#define GEN_PASS_DEF_POSTSPARSIFICATIONREWRITE
 #define GEN_PASS_DEF_SPARSETENSORCONVERSIONPASS
 #define GEN_PASS_DEF_SPARSETENSORCODEGEN
 #define GEN_PASS_DEF_SPARSEBUFFERREWRITE
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
 } // namespace mlir
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 namespace {
 
 //===----------------------------------------------------------------------===//
 // Passes implementation.
 //===----------------------------------------------------------------------===//
 
-struct SparseTensorRewritePass
-    : public impl::SparseTensorRewriteBase<SparseTensorRewritePass> {
+struct PreSparsificationRewritePass
+    : public impl::PreSparsificationRewriteBase<PreSparsificationRewritePass> {
 
-  SparseTensorRewritePass() = default;
-  SparseTensorRewritePass(const SparseTensorRewritePass &pass) = default;
-  SparseTensorRewritePass(bool enableRT, bool foreach, bool convert) {
-    enableRuntimeLibrary = enableRT;
-    enableForeach = foreach;
-    enableConvert = convert;
-  }
+  PreSparsificationRewritePass() = default;
+  PreSparsificationRewritePass(const PreSparsificationRewritePass &pass) =
+      default;
 
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
-    populateSparseTensorRewriting(patterns, enableRuntimeLibrary, enableForeach,
-                                  enableConvert);
+    populatePreSparsificationRewriting(patterns);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 };
 
 struct SparsificationPass
     : public impl::SparsificationPassBase<SparsificationPass> {
 
   SparsificationPass() = default;
void runOnOperation() override {
     RewritePatternSet patterns(ctx);
     populateSparsificationPatterns(patterns, options);
     vector::populateVectorToVectorCanonicalizationPatterns(patterns);
     scf::ForOp::getCanonicalizationPatterns(patterns, ctx);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
   }
 };
 
+struct PostSparsificationRewritePass
+    : public impl::PostSparsificationRewriteBase<
+          PostSparsificationRewritePass> {
+
+  PostSparsificationRewritePass() = default;
+  PostSparsificationRewritePass(const PostSparsificationRewritePass &pass) =
+      default;
+  PostSparsificationRewritePass(bool enableRT, bool foreach, bool convert) {
+    enableRuntimeLibrary = enableRT;
+    enableForeach = foreach;
+    enableConvert = convert;
+  }
+
+  void runOnOperation() override {
+    auto *ctx = &getContext();
+    RewritePatternSet patterns(ctx);
+    populatePostSparsificationRewriting(patterns, enableRuntimeLibrary,
+                                        enableForeach, enableConvert);
+    (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
+  }
+};
+
 struct SparseTensorConversionPass
     : public impl::SparseTensorConversionPassBase<SparseTensorConversionPass> {
 
   SparseTensorConversionPass() = default;
   SparseTensorConversionPass(const SparseTensorConversionPass &pass) = default;
   SparseTensorConversionPass(const SparseTensorConversionOptions &options) {
     sparseToSparse = static_cast<int32_t>(options.sparseToSparseStrategy);
   }
case 2:
     return SparseToSparseConversionStrategy::kDirect;
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Pass creation methods.
 //===----------------------------------------------------------------------===//
 
-std::unique_ptr<Pass> mlir::createSparseTensorRewritePass() {
-  return std::make_unique<SparseTensorRewritePass>();
-}
-
-std::unique_ptr<Pass> mlir::createSparseTensorRewritePass(bool enableRT,
-                                                          bool enableForeach,
-                                                          bool enableConvert) {
-  return std::make_unique<SparseTensorRewritePass>(enableRT, enableForeach,
-                                                   enableConvert);
+std::unique_ptr<Pass> mlir::createPreSparsificationRewritePass() {
+  return std::make_unique<PreSparsificationRewritePass>();
 }
 
 std::unique_ptr<Pass> mlir::createSparsificationPass() {
   return std::make_unique<SparsificationPass>();
 }
 
 std::unique_ptr<Pass>
 mlir::createSparsificationPass(const SparsificationOptions &options) {
   return std::make_unique<SparsificationPass>(options);
 }
 
+std::unique_ptr<Pass> mlir::createPostSparsificationRewritePass() {
+  return std::make_unique<PostSparsificationRewritePass>();
+}
+
+std::unique_ptr<Pass>
+mlir::createPostSparsificationRewritePass(bool enableRT, bool enableForeach,
+                                          bool enableConvert) {
+  return std::make_unique<PostSparsificationRewritePass>(
+      enableRT, enableForeach, enableConvert);
+}
+
 std::unique_ptr<Pass> mlir::createSparseTensorConversionPass() {
   return std::make_unique<SparseTensorConversionPass>();
 }
 
 std::unique_ptr<Pass> mlir::createSparseTensorConversionPass(
     const SparseTensorConversionOptions &options) {
   return std::make_unique<SparseTensorConversionPass>(options);
 }

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

struct OutRewriter : public OpRewritePattern<OutOp> {
   }
 };
 
 } // namespace
 
 //===---------------------------------------------------------------------===//
 // Methods that add patterns described in this file to a pattern list.
 //===---------------------------------------------------------------------===//
-void mlir::populateSparseTensorRewriting(RewritePatternSet &patterns,
-                                         bool enableRT, bool enableForeach,
+
+void mlir::populatePreSparsificationRewriting(RewritePatternSet &patterns) {
+  patterns.add<FoldInvariantYield, FuseSparseMultiplyOverAdd>(
+      patterns.getContext());
+}
+
+void mlir::populatePostSparsificationRewriting(RewritePatternSet &patterns,
+                                               bool enableRT,
+                                               bool enableForeach,
                                                bool enableConvert) {
-  patterns.add<FoldInvariantYield, FuseSparseMultiplyOverAdd,
-               ReshapeRewriter<tensor::ExpandShapeOp>,
+  patterns.add<ReshapeRewriter<tensor::ExpandShapeOp>,
                ReshapeRewriter<tensor::CollapseShapeOp>>(patterns.getContext());
   if (enableForeach)
     patterns.add<ForeachRewriter>(patterns.getContext());
   // TODO: If RT not enabled, rewrite concatenate ops, etc here.
   if (!enableRT) {
     patterns.add<ConcatenateRewriter, NewRewriter, OutRewriter,
                  Sparse2SparseReshapeRewriter<tensor::ExpandShapeOp>,
                  Sparse2SparseReshapeRewriter<tensor::CollapseShapeOp>>(
         patterns.getContext());
     if (enableConvert)
       patterns.add<ConvertRewriter>(patterns.getContext());
   }
 }

mlir/test/Dialect/SparseTensor/convert_dense2sparse.mlir

 // RUN: mlir-opt %s --sparse-tensor-conversion --canonicalize --cse | FileCheck %s
-// RUN: mlir-opt %s --sparse-tensor-rewrite="enable-runtime-library=false enable-foreach=false" \
+// RUN: mlir-opt %s --post-sparsification-rewrite="enable-runtime-library=false enable-foreach=false" \
 // RUN: --canonicalize --cse | FileCheck %s --check-prefix=CHECK-RWT
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"]
 }>
 
 #CSR = #sparse_tensor.encoding<{
   dimLevelType = ["dense", "compressed"]

mlir/test/Dialect/SparseTensor/convert_sparse2dense.mlir

 // RUN: mlir-opt %s --sparse-tensor-conversion --canonicalize --cse | FileCheck %s
 
-// RUN: mlir-opt %s --sparse-tensor-rewrite="enable-runtime-library=false enable-foreach=false" \
+// RUN: mlir-opt %s --post-sparsification-rewrite="enable-runtime-library=false enable-foreach=false" \
 // RUN: --canonicalize --cse | FileCheck %s --check-prefix=CHECK-RWT
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"]
 }>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = ["dense", "compressed"]

mlir/test/Dialect/SparseTensor/convert_sparse2sparse.mlir

 // First use with `kViaCOO` for sparse2sparse conversion (the old way).
 // RUN: mlir-opt %s --sparse-tensor-conversion="s2s-strategy=1" \
 // RUN:    --canonicalize --cse | FileCheck %s -check-prefix=CHECK-COO
 //
 // Now again with `kAuto` (the new default).
 // RUN: mlir-opt %s --sparse-tensor-conversion="s2s-strategy=0" \
 // RUN:    --canonicalize --cse | FileCheck %s -check-prefixes=CHECK-AUTO,CHECK
 
-// RUN: mlir-opt %s --sparse-tensor-rewrite="enable-runtime-library=false enable-foreach=false" \
+// RUN: mlir-opt %s --post-sparsification-rewrite="enable-runtime-library=false enable-foreach=false" \
 // RUN: --canonicalize --cse | FileCheck %s --check-prefix=CHECK-RWT
 
 #SparseVector64 = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"],
   pointerBitWidth = 64,
   indexBitWidth = 64
 }>
 

mlir/test/Dialect/SparseTensor/rewriting.mlir

-// RUN: mlir-opt %s -sparse-tensor-rewrite | FileCheck %s
+// RUN: mlir-opt %s -post-sparsification-rewrite | FileCheck %s
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"]
 }>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = ["compressed", "compressed"]
 }>

mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir

-// RUN: mlir-opt %s -sparse-tensor-rewrite="enable-runtime-library=false enable-convert=false" |\
+// RUN: mlir-opt %s -post-sparsification-rewrite="enable-runtime-library=false enable-convert=false" |\
 // RUN: FileCheck %s
 
 #CSR = #sparse_tensor.encoding<{
   dimLevelType = ["dense", "compressed"]
 }>
 
 // CHECK-LABEL:   func.func @sparse_new(
 // CHECK-SAME:    %[[A:.*]]: !llvm.ptr<i8>) -> tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> {

mlir/test/Dialect/SparseTensor/sparse_concat_codegen.mlir

-// RUN: mlir-opt %s --sparse-tensor-rewrite="enable-runtime-library=false enable-convert=false" \
+// RUN: mlir-opt %s --post-sparsification-rewrite="enable-runtime-library=false enable-convert=false" \
 // RUN: --sparsification | FileCheck %s
 
 #DCSR = #sparse_tensor.encoding<{dimLevelType = ["compressed", "compressed"]}>
 
 // CHECK-LABEL: @concat_sparse_sparse(
 //  CHECK-SAME:  %[[TMP_arg0:.*]]: tensor<2x4xf64, #sparse_tensor
 //  CHECK-SAME:  %[[TMP_arg1:.*]]: tensor<3x4xf64, #sparse_tensor
 //  CHECK-SAME:  %[[TMP_arg2:.*]]: tensor<4x4xf64, #sparse_tensor

mlir/test/Dialect/SparseTensor/sparse_fill_zero.mlir

-// RUN: mlir-opt %s --linalg-generalize-named-ops --sparse-tensor-rewrite --sparsification --sparse-tensor-conversion --canonicalize --cse | FileCheck %s
+// RUN: mlir-opt %s --linalg-generalize-named-ops --pre-sparsification-rewrite --sparsification --sparse-tensor-conversion --canonicalize --cse | FileCheck %s
 
 #DCSR = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
 
 // CHECK-LABEL:   func.func @fill_zero_after_alloc(
 // CHECK-SAME:      %[[Arg0:.*]]: !llvm.ptr<i8>,
 // CHECK-SAME:      %[[Arg1:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8> {
 // CHECK-DAG:       %[[F0:.*]] = arith.constant 0.000000e+00 : f64
 // CHECK-DAG:       %[[C1:.*]] = arith.constant 1 : i32

mlir/test/Dialect/SparseTensor/sparse_reshape.mlir

 // RUN: mlir-opt %s | mlir-opt | FileCheck %s --check-prefix=CHECK-ROUND
 // RUN: mlir-opt %s --sparse-tensor-conversion --cse --canonicalize | FileCheck %s --check-prefix=CHECK-CONV
-// RUN: mlir-opt %s --sparse-tensor-rewrite="enable-runtime-library=false enable-convert=false" \
+// RUN: mlir-opt %s --post-sparsification-rewrite="enable-runtime-library=false enable-convert=false" \
 // RUN: --cse --canonicalize  | FileCheck %s --check-prefix=CHECK-RWT
 
 #SparseVector = #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>
 #SparseMatrix = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
 
 //
 // roundtrip:
 //

mlir/test/Dialect/SparseTensor/sparse_sddmm.mlir

-// RUN: mlir-opt %s  --tensor-copy-insertion --sparse-tensor-rewrite --sparsification --cse | FileCheck %s
+// RUN: mlir-opt %s  --tensor-copy-insertion --pre-sparsification-rewrite --sparsification --cse | FileCheck %s
 
 #SM = #sparse_tensor.encoding<{ dimLevelType = [ "compressed", "compressed" ] }>
 
 #trait_matmul = {
   indexing_maps = [
     affine_map<(d0, d1, d2) -> (d1, d0)>,
     affine_map<(d0, d1, d2) -> (d0, d2)>,
     affine_map<(d0, d1, d2) -> (d1, d2)>