[Mlir-commits] [mlir] 30ceb78 - [mlir][sparse] Expose SparseTensor passes as enums instead of opaque numbers for vectorization and parallelization options.

Nick Kreeger llvmlistbot at llvm.org
Sat Sep 3 18:39:54 PDT 2022


Author: Nick Kreeger
Date: 2022-09-04T01:39:35Z
New Revision: 30ceb783e232ff625cb25b9a264ae60519241f30

URL: https://github.com/llvm/llvm-project/commit/30ceb783e232ff625cb25b9a264ae60519241f30
DIFF: https://github.com/llvm/llvm-project/commit/30ceb783e232ff625cb25b9a264ae60519241f30.diff

LOG: [mlir][sparse] Expose SparseTensor passes as enums instead of opaque numbers for vectorization and parallelization options.

The SparseTensor passes currently use opaque numbers for the CLI, despite using an enum internally. This patch exposes the enums instead of numbered items that are matched back to the enum.

Fixes https://github.com/llvm/llvm-project/issues/53389

Differential Revision: https://reviews.llvm.org/D123876

Please also see:
https://reviews.llvm.org/D118379
https://reviews.llvm.org/D117919

Added: 
    

Modified: 
    mlir/benchmark/python/common.py
    mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
    mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
    mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
    mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
    mlir/test/Dialect/SparseTensor/sparse_parallel.mlir
    mlir/test/Dialect/SparseTensor/sparse_vector.mlir
    mlir/test/Dialect/SparseTensor/sparse_vector_chain.mlir
    mlir/test/Dialect/SparseTensor/sparse_vector_index.mlir
    mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index_dense.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scale.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir
    mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir
    mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py
    mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py
    mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py

Removed: 
    


################################################################################
diff  --git a/mlir/benchmark/python/common.py b/mlir/benchmark/python/common.py
index b2269251480bc..3634641074f3e 100644
--- a/mlir/benchmark/python/common.py
+++ b/mlir/benchmark/python/common.py
@@ -13,8 +13,8 @@ def setup_passes(mlir_module):
     """Setup pass pipeline parameters for benchmark functions.
     """
     opt = (
-        "parallelization-strategy=0"
-        " vectorization-strategy=0 vl=1 enable-simd-index32=False"
+        "parallelization-strategy=none"
+        " vectorization-strategy=none vl=1 enable-simd-index32=False"
     )
     pipeline = f"sparse-compiler{{{opt}}}"
     PassManager.parse(pipeline).run(mlir_module)

diff  --git a/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
index 43abadd5a7627..580b967059681 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Pipelines/Passes.h
@@ -30,12 +30,43 @@ namespace sparse_tensor {
 struct SparseCompilerOptions
     : public PassPipelineOptions<SparseCompilerOptions> {
   // These options must be kept in sync with `SparsificationBase`.
-  PassOptions::Option<int32_t> parallelization{
+  // TODO(57514): These options are duplicated in Passes.td.
+  PassOptions::Option<mlir::SparseParallelizationStrategy> parallelization{
       *this, "parallelization-strategy",
-      desc("Set the parallelization strategy"), init(0)};
-  PassOptions::Option<int32_t> vectorization{
-      *this, "vectorization-strategy", desc("Set the vectorization strategy"),
-      init(0)};
+      ::llvm::cl::desc("Set the parallelization strategy"),
+      ::llvm::cl::init(mlir::SparseParallelizationStrategy::kNone),
+      llvm::cl::values(
+          clEnumValN(mlir::SparseParallelizationStrategy::kNone, "none",
+                     "Turn off sparse parallelization."),
+          clEnumValN(mlir::SparseParallelizationStrategy::kDenseOuterLoop,
+                     "dense-outer-loop",
+                     "Enable dense outer loop sparse parallelization."),
+          clEnumValN(mlir::SparseParallelizationStrategy::kAnyStorageOuterLoop,
+                     "any-storage-outer-loop",
+                     "Enable sparse parallelization regardless of storage for "
+                     "the outer loop."),
+          clEnumValN(mlir::SparseParallelizationStrategy::kDenseAnyLoop,
+                     "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."))};
+  PassOptions::Option<mlir::SparseVectorizationStrategy> vectorization{
+      *this, "vectorization-strategy",
+      ::llvm::cl::desc("Set the vectorization strategy"),
+      ::llvm::cl::init(mlir::SparseVectorizationStrategy::kNone),
+      llvm::cl::values(
+          clEnumValN(mlir::SparseVectorizationStrategy::kNone, "none",
+                     "Turn off sparse vectorization."),
+          clEnumValN(mlir::SparseVectorizationStrategy::kDenseInnerLoop,
+                     "dense-inner-loop",
+                     "Enable vectorization for dense inner loops."),
+          clEnumValN(mlir::SparseVectorizationStrategy::kAnyStorageInnerLoop,
+                     "any-storage-inner-loop",
+                     "Enable sparse vectorization for inner loops with any "
+                     "storage."))};
+
   PassOptions::Option<int32_t> vectorLength{
       *this, "vl", desc("Set the vector length"), init(1)};
   PassOptions::Option<bool> enableSIMDIndex32{
@@ -50,10 +81,8 @@ struct SparseCompilerOptions
 
   /// Projects out the options for `createSparsificationPass`.
   SparsificationOptions sparsificationOptions() const {
-    return SparsificationOptions(sparseParallelizationStrategy(parallelization),
-                                 sparseVectorizationStrategy(vectorization),
-                                 vectorLength, enableSIMDIndex32,
-                                 enableVLAVectorization);
+    return SparsificationOptions(parallelization, vectorization, vectorLength,
+                                 enableSIMDIndex32, enableVLAVectorization);
   }
 
   // These options must be kept in sync with `SparseTensorConversionBase`.

diff  --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
index 523f1fb057d47..227b70a381192 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.h
@@ -26,11 +26,6 @@ namespace bufferization {
 struct OneShotBufferizationOptions;
 } // namespace bufferization
 
-#define GEN_PASS_DECL_SPARSIFICATIONPASS
-#define GEN_PASS_DECL_SPARSETENSORCONVERSIONPASS
-#define GEN_PASS_DECL_SPARSETENSORCODEGEN
-#include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
-
 //===----------------------------------------------------------------------===//
 // The Sparsification pass.
 //===----------------------------------------------------------------------===//
@@ -49,9 +44,6 @@ enum class SparseParallelizationStrategy {
   // TODO: support reduction parallelization too?
 };
 
-/// Converts command-line parallelization flag to the strategy enum.
-SparseParallelizationStrategy sparseParallelizationStrategy(int32_t flag);
-
 /// 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
@@ -62,8 +54,10 @@ enum class SparseVectorizationStrategy {
   kAnyStorageInnerLoop
 };
 
-/// Converts command-line vectorization flag to the strategy enum.
-SparseVectorizationStrategy sparseVectorizationStrategy(int32_t flag);
+#define GEN_PASS_DECL_SPARSIFICATIONPASS
+#define GEN_PASS_DECL_SPARSETENSORCONVERSIONPASS
+#define GEN_PASS_DECL_SPARSETENSORCODEGEN
+#include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
 
 /// Options for the Sparsification pass.
 struct SparsificationOptions {

diff  --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index 8f96280efe83a..cd6b77ea50eea 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -62,11 +62,36 @@ def SparsificationPass : Pass<"sparsification", "ModuleOp"> {
     "sparse_tensor::SparseTensorDialect",
     "vector::VectorDialect",
   ];
+  // TODO(57514): These enum options are duplicated in Passes.h.
   let options = [
-    Option<"parallelization", "parallelization-strategy", "int32_t", "0",
-           "Set the parallelization strategy">,
-    Option<"vectorization", "vectorization-strategy", "int32_t", "0",
-           "Set the vectorization strategy">,
+    Option<"parallelization", "parallelization-strategy", "mlir::SparseParallelizationStrategy",
+           "mlir::SparseParallelizationStrategy::kNone",
+           "Set the parallelization strategy", [{llvm::cl::values(
+             clEnumValN(mlir::SparseParallelizationStrategy::kNone, "none",
+                        "Turn off sparse parallelization."),
+             clEnumValN(mlir::SparseParallelizationStrategy::kDenseOuterLoop,
+                        "dense-outer-loop",
+                        "Enable dense outer loop sparse parallelization."),
+             clEnumValN(mlir::SparseParallelizationStrategy::kAnyStorageOuterLoop,
+                        "any-storage-outer-loop",
+                        "Enable sparse parallelization regardless of storage for the outer loop."),
+             clEnumValN(mlir::SparseParallelizationStrategy::kDenseAnyLoop,
+                        "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."))}]>,
+    Option<"vectorization", "vectorization-strategy", "mlir::SparseVectorizationStrategy",
+           "mlir::SparseVectorizationStrategy::kNone",
+           "Set the vectorization strategy", [{llvm::cl::values(
+             clEnumValN(mlir::SparseVectorizationStrategy::kNone, "none",
+                        "Turn off sparse vectorization."),
+             clEnumValN(mlir::SparseVectorizationStrategy::kDenseInnerLoop,
+                        "dense-inner-loop",
+                        "Enable vectorization for dense inner loops."),
+             clEnumValN(mlir::SparseVectorizationStrategy::kAnyStorageInnerLoop,
+                        "any-storage-inner-loop",
+                        "Enable sparse vectorization for inner loops with any storage."))}]>,
     Option<"vectorLength", "vl", "int32_t", "1",
            "Set the vector length">,
     Option<"enableSIMDIndex32", "enable-simd-index32", "bool", "false",

diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
index 9057921e87c68..c79ca9545615e 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorPasses.cpp
@@ -43,8 +43,8 @@ struct SparsificationPass
   SparsificationPass() = default;
   SparsificationPass(const SparsificationPass &pass) = default;
   SparsificationPass(const SparsificationOptions &options) {
-    parallelization = static_cast<int32_t>(options.parallelizationStrategy);
-    vectorization = static_cast<int32_t>(options.vectorizationStrategy);
+    parallelization = options.parallelizationStrategy;
+    vectorization = options.vectorizationStrategy;
     vectorLength = options.vectorLength;
     enableSIMDIndex32 = options.enableSIMDIndex32;
     enableVLAVectorization = options.enableVLAVectorization;
@@ -57,10 +57,8 @@ struct SparsificationPass
     populateSparseTensorRewriting(prePatterns);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(prePatterns));
     // Translate strategy flags to strategy options.
-    SparsificationOptions options(
-        sparseParallelizationStrategy(parallelization),
-        sparseVectorizationStrategy(vectorization), vectorLength,
-        enableSIMDIndex32, enableVLAVectorization);
+    SparsificationOptions options(parallelization, vectorization, vectorLength,
+                                  enableSIMDIndex32, enableVLAVectorization);
     // Apply sparsification and vector cleanup rewriting.
     RewritePatternSet patterns(ctx);
     populateSparsificationPatterns(patterns, options);
@@ -238,33 +236,6 @@ struct SparseTensorStorageExpansionPass
 // Strategy flag methods.
 //===----------------------------------------------------------------------===//
 
-SparseParallelizationStrategy
-mlir::sparseParallelizationStrategy(int32_t flag) {
-  switch (flag) {
-  default:
-    return SparseParallelizationStrategy::kNone;
-  case 1:
-    return SparseParallelizationStrategy::kDenseOuterLoop;
-  case 2:
-    return SparseParallelizationStrategy::kAnyStorageOuterLoop;
-  case 3:
-    return SparseParallelizationStrategy::kDenseAnyLoop;
-  case 4:
-    return SparseParallelizationStrategy::kAnyStorageAnyLoop;
-  }
-}
-
-SparseVectorizationStrategy mlir::sparseVectorizationStrategy(int32_t flag) {
-  switch (flag) {
-  default:
-    return SparseVectorizationStrategy::kNone;
-  case 1:
-    return SparseVectorizationStrategy::kDenseInnerLoop;
-  case 2:
-    return SparseVectorizationStrategy::kAnyStorageInnerLoop;
-  }
-}
-
 SparseToSparseConversionStrategy
 mlir::sparseToSparseConversionStrategy(int32_t flag) {
   switch (flag) {

diff  --git a/mlir/test/Dialect/SparseTensor/sparse_parallel.mlir b/mlir/test/Dialect/SparseTensor/sparse_parallel.mlir
index 9af037c7829a8..5e0268191fbf6 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_parallel.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_parallel.mlir
@@ -1,12 +1,12 @@
-// RUN: mlir-opt %s -sparsification="parallelization-strategy=0" | \
+// RUN: mlir-opt %s -sparsification="parallelization-strategy=none" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR0
-// RUN: mlir-opt %s -sparsification="parallelization-strategy=1" | \
+// RUN: mlir-opt %s -sparsification="parallelization-strategy=dense-outer-loop" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR1
-// RUN: mlir-opt %s -sparsification="parallelization-strategy=2" | \
+// RUN: mlir-opt %s -sparsification="parallelization-strategy=any-storage-outer-loop" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR2
-// RUN: mlir-opt %s -sparsification="parallelization-strategy=3" | \
+// RUN: mlir-opt %s -sparsification="parallelization-strategy=dense-any-loop" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR3
-// RUN: mlir-opt %s -sparsification="parallelization-strategy=4" | \
+// RUN: mlir-opt %s -sparsification="parallelization-strategy=any-storage-any-loop" | \
 // RUN:   FileCheck %s --check-prefix=CHECK-PAR4
 
 #DenseMatrix = #sparse_tensor.encoding<{

diff  --git a/mlir/test/Dialect/SparseTensor/sparse_vector.mlir b/mlir/test/Dialect/SparseTensor/sparse_vector.mlir
index 816a5c75c7c80..a730e7760e3f9 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_vector.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_vector.mlir
@@ -1,12 +1,12 @@
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=0 vl=16" -cse -split-input-file | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=none vl=16" -cse -split-input-file | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC0
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=1 vl=16" -cse -split-input-file | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=dense-inner-loop vl=16" -cse -split-input-file | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC1
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=16" -cse -split-input-file | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=16" -cse -split-input-file | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC2
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=16 enable-simd-index32=true" -cse -split-input-file | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=16 enable-simd-index32=true" -cse -split-input-file | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC3
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=4 enable-vla-vectorization=true" -cse -split-input-file | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=4 enable-vla-vectorization=true" -cse -split-input-file | \
 // RUN:   FileCheck %s --check-prefix=CHECK-VEC4
 
 #DenseVector = #sparse_tensor.encoding<{ dimLevelType = [ "dense" ] }>

diff  --git a/mlir/test/Dialect/SparseTensor/sparse_vector_chain.mlir b/mlir/test/Dialect/SparseTensor/sparse_vector_chain.mlir
index bd557c3d9e1fe..69eaa7b367f44 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_vector_chain.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_vector_chain.mlir
@@ -1,6 +1,6 @@
 // NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
 
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=8" -canonicalize | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=8" -canonicalize | \
 // RUN:   FileCheck %s
 
 #SparseMatrix = #sparse_tensor.encoding<{dimLevelType = ["dense","compressed"]}>

diff  --git a/mlir/test/Dialect/SparseTensor/sparse_vector_index.mlir b/mlir/test/Dialect/SparseTensor/sparse_vector_index.mlir
index dfea43e73f0e9..fd8ceb74f1bb1 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_vector_index.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_vector_index.mlir
@@ -5,7 +5,7 @@
 // about what constitutes a good test! The CHECK should be
 // minimized and named to reflect the test intent.
 
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=8" -canonicalize | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=8" -canonicalize | \
 // RUN:   FileCheck %s
 
 #SparseVector = #sparse_tensor.encoding<{

diff  --git a/mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir b/mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir
index 65f1f7216bed8..276b8a9b8b933 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_vector_peeled.mlir
@@ -1,4 +1,4 @@
-// RUN: mlir-opt %s -sparsification="vectorization-strategy=2 vl=16" -scf-for-loop-peeling -canonicalize | \
+// RUN: mlir-opt %s -sparsification="vectorization-strategy=any-storage-inner-loop vl=16" -scf-for-loop-peeling -canonicalize | \
 // RUN:   FileCheck %s
 
 #SparseVector = #sparse_tensor.encoding<{

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir
index 569d8b122fc98..a59978267781a 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_cast.mlir
@@ -6,7 +6,7 @@
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2" | \
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=2" | \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir
index 52b8400b81acf..bbe4807aba4cc 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_filter_conv2d.mlir
@@ -5,7 +5,7 @@
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2" | \
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=2" | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir
index 45670496d7328..172d3909f2f5a 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_flatten.mlir
@@ -7,7 +7,7 @@
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4" | \
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/test.tns" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index_dense.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index_dense.mlir
index 7fefeadd2598b..983651f4875db 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index_dense.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_index_dense.mlir
@@ -5,7 +5,7 @@
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4" | \
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir
index 17be14d02a03d..4e0b9afc08710 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_matvec.mlir
@@ -8,7 +8,7 @@
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s \
-// RUN:   --sparse-compiler="vectorization-strategy=2 vl=16 enable-simd-index32" | \
+// RUN:   --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=16 enable-simd-index32" | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/wide.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir
index 74302a7d95b5e..d21208d551629 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_mttkrp.mlir
@@ -7,7 +7,7 @@
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4" | \
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/mttkrp_b.tns" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir
index 8bd0133c874de..7079ec019370e 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_simple.mlir
@@ -7,7 +7,7 @@
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4" | \
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir
index 94ab5e4fdfecc..30e04feb21a82 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_quantized_matmul.mlir
@@ -5,7 +5,7 @@
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2" | \
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=2" | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir
index 10a1859ee1453..220fa18dfb45a 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reductions.mlir
@@ -5,7 +5,7 @@
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=8" | \
+// RUN: mlir-opt %s -sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=8" | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir
index 55fe32bd86d5d..fc3dfa1b18a32 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_matmul.mlir
@@ -8,7 +8,7 @@
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
 // RUN: mlir-opt %s \
-// RUN:   --sparse-compiler="vectorization-strategy=2 vl=4 enable-simd-index32" | \
+// RUN:   --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4 enable-simd-index32" | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/test.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir
index 3c7f741c77fcc..8b2f987081021 100755
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sampled_mm_fusion.mlir
@@ -5,7 +5,7 @@
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=8" | \
+// RUN: mlir-opt %s -sparse-compiler="vl=8" | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scale.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scale.mlir
index 56507151dd4a2..7d46b26bb4df8 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scale.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_scale.mlir
@@ -6,7 +6,7 @@
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=4" | \
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=4" | \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir
index 3b483672d8074..33ce4b99657e2 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_spmm.mlir
@@ -7,7 +7,7 @@
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2" | \
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=2" | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/wide.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir
index 1c29f8cb47817..54cd090a802f7 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_sum.mlir
@@ -7,7 +7,7 @@
 //
 // Do the same run, but now with SIMDization as well. This should not change the outcome.
 //
-// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=2 vl=2" | \
+// RUN: mlir-opt %s --sparse-compiler="vectorization-strategy=any-storage-inner-loop vl=2" | \
 // RUN: TENSOR0="%mlir_integration_test_dir/data/test_symmetric.mtx" \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py
index 4ed741ed3727e..a6abb3cc4605b 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py
@@ -140,22 +140,24 @@ def main():
             ir.AffineMap.get_permutation([0, 1]),
             ir.AffineMap.get_permutation([1, 0])
         ]
+        vec_strategy = [
+          'none', 'dense-inner-loop'
+        ]
         for level in levels:
             for ordering in orderings:
                 for pwidth in [32]:
                     for iwidth in [32]:
-                        for par in [0]:
-                            for vec in [0, 1]:
-                                for e in [True]:
-                                    vl = 1 if vec == 0 else 16
-                                    attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
-                                    opt = (f'parallelization-strategy={par} '
-                                           f'vectorization-strategy={vec} '
-                                           f'vl={vl} enable-simd-index32={e}')
-                                    compiler = sparse_compiler.SparseCompiler(
-                                        options=opt, opt_level=0, shared_libs=[support_lib])
-                                    build_compile_and_run_SDDMMM(attr, compiler)
-                                    count = count + 1
+                        for vec in vec_strategy:
+                            for e in [True]:
+                                vl = 1 if vec == 0 else 16
+                                attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
+                                opt = (f'parallelization-strategy=none '
+                                       f'vectorization-strategy={vec} '
+                                       f'vl={vl} enable-simd-index32={e}')
+                                compiler = sparse_compiler.SparseCompiler(
+                                    options=opt, opt_level=0, shared_libs=[support_lib])
+                                build_compile_and_run_SDDMMM(attr, compiler)
+                                count = count + 1
     # CHECK: Passed 16 tests
     print('Passed ', count, 'tests')
 

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py
index 9712620a4a7a9..57ca3d6e1a94c 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py
@@ -120,12 +120,10 @@ def main():
         # a *single* sparse tensor. Note that we deliberate do not exhaustively
         # search the full state space to reduce runtime of the test. It is
         # straightforward to adapt the code below to explore more combinations.
-        par = 0
-        vec = 0
         vl = 1
         e = False
-        opt = (f'parallelization-strategy={par} '
-               f'vectorization-strategy={vec} '
+        opt = (f'parallelization-strategy=none '
+               f'vectorization-strategy=none '
                f'vl={vl} enable-simd-index32={e}')
         levels = [[st.DimLevelType.dense, st.DimLevelType.dense],
                   [st.DimLevelType.dense, st.DimLevelType.compressed],

diff  --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py
index a030643dcc18b..01add6f4a23c3 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py
@@ -183,8 +183,6 @@ def main():
   # CHECK-LABEL: TEST: test_stress
   print("\nTEST: test_stress")
   with ir.Context() as ctx, ir.Location.unknown():
-    par = 0
-    vec = 0
     vl = 1
     e = False
     # Disable direct sparse2sparse conversion, because it doubles the time!
@@ -193,8 +191,8 @@ def main():
     # `s2s=0` on a regular basis, to ensure that it does continue to work.
     s2s = 1
     sparsification_options = (
-        f'parallelization-strategy={par} '
-        f'vectorization-strategy={vec} '
+        f'parallelization-strategy=none '
+        f'vectorization-strategy=none '
         f'vl={vl} '
         f'enable-simd-index32={e} '
         f's2s-strategy={s2s}')


        


More information about the Mlir-commits mailing list