[Mlir-commits] [mlir] 438a7d4 - [mlir][sparse] expose optimization flags to mini pipeline (#95158)

[llvmlistbot at llvm.org]

Author: Aart Bik
Date: 2024-06-11T14:20:58-07:00
New Revision: 438a7d4c982e0a38aaa6544a5ba6736d54600733

URL: https://github.com/llvm/llvm-project/commit/438a7d4c982e0a38aaa6544a5ba6736d54600733
DIFF: https://github.com/llvm/llvm-project/commit/438a7d4c982e0a38aaa6544a5ba6736d54600733.diff

LOG: [mlir][sparse] expose optimization flags to mini pipeline (#95158)

Some of the options only fed into the full sparse pipeline. However,
some backends prefer to use the sparse minipipeline. This change exposes
some important optimization flags to the pass as well. This prepares
some SIMDization of PyTorch sparsified code.

Added: 
    mlir/test/Dialect/SparseTensor/minipipeline_vector.mlir

Modified: 
    mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
    mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
index c6554e1c94a4a..196110f55571d 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/Transforms/Passes.td
@@ -462,6 +462,18 @@ def SparsificationAndBufferization : Pass<"sparsification-and-bufferization", "M
     "sparse_tensor::SparseTensorDialect",
     "vector::VectorDialect"
   ];
+  // Important optimization options are made visible to the mini-pipeline
+  // so that clients can set these (when not using the full pipeline).
+  let options = [
+    Option<"vectorLength", "vl", "int32_t", "0",
+           "Set the vector length (use 0 to disable vectorization)">,
+    Option<"enableVLAVectorization", "enable-vla-vectorization", "bool", "false",
+           "Enable vector length agnostic vectorization">,
+    Option<"enableSIMDIndex32", "enable-simd-index32", "bool", "false",
+           "Enable i32 indexing into vectors (for efficient gather/scatter)">,
+    Option<"enableGPULibgen", "enable-gpu-libgen", "bool", "false",
+           "Enable GPU acceleration by means of direct library calls">,
+  ];
 }
 
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
index 3a8972072ac3b..13c750e83d045 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparsificationAndBufferizationPass.cpp
@@ -61,20 +61,34 @@ class SparsificationAndBufferizationPass
     : public impl::SparsificationAndBufferizationBase<
           SparsificationAndBufferizationPass> {
 public:
+  // Private pass options only.
   SparsificationAndBufferizationPass(
       const bufferization::OneShotBufferizationOptions &bufferizationOptions,
       const SparsificationOptions &sparsificationOptions,
       bool createSparseDeallocs, bool enableRuntimeLibrary,
-      bool enableBufferInitialization, unsigned vectorLength,
-      bool enableVLAVectorization, bool enableSIMDIndex32, bool enableGPULibgen)
+      bool enableBufferInitialization)
       : bufferizationOptions(bufferizationOptions),
         sparsificationOptions(sparsificationOptions),
         createSparseDeallocs(createSparseDeallocs),
         enableRuntimeLibrary(enableRuntimeLibrary),
-        enableBufferInitialization(enableBufferInitialization),
-        vectorLength(vectorLength),
-        enableVLAVectorization(enableVLAVectorization),
-        enableSIMDIndex32(enableSIMDIndex32), enableGPULibgen(enableGPULibgen) {
+        enableBufferInitialization(enableBufferInitialization) {}
+  // Private pass options and visible pass options.
+  SparsificationAndBufferizationPass(
+      const bufferization::OneShotBufferizationOptions &bufferizationOptions,
+      const SparsificationOptions &sparsificationOptions,
+      bool createSparseDeallocs, bool enableRuntimeLibrary,
+      bool enableBufferInitialization, unsigned vl, bool vla, bool index32,
+      bool gpu)
+      : bufferizationOptions(bufferizationOptions),
+        sparsificationOptions(sparsificationOptions),
+        createSparseDeallocs(createSparseDeallocs),
+        enableRuntimeLibrary(enableRuntimeLibrary),
+        enableBufferInitialization(enableBufferInitialization) {
+    // Set the visible pass options explicitly.
+    vectorLength = vl;
+    enableVLAVectorization = vla;
+    enableSIMDIndex32 = index32;
+    enableGPULibgen = gpu;
   }
 
   /// Bufferize all dense ops. This assumes that no further analysis is needed
@@ -178,10 +192,6 @@ class SparsificationAndBufferizationPass
   bool createSparseDeallocs;
   bool enableRuntimeLibrary;
   bool enableBufferInitialization;
-  unsigned vectorLength;
-  bool enableVLAVectorization;
-  bool enableSIMDIndex32;
-  bool enableGPULibgen;
 };
 
 } // namespace sparse_tensor
@@ -213,16 +223,13 @@ mlir::getBufferizationOptionsForSparsification(bool analysisOnly) {
 
 std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass() {
   SparsificationOptions sparseOptions;
-  return createSparsificationAndBufferizationPass(
+  return std::make_unique<
+      mlir::sparse_tensor::SparsificationAndBufferizationPass>(
       getBufferizationOptionsForSparsification(/*analysisOnly=*/false),
       sparseOptions,
       /*createSparseDeallocs=*/false,
       /*enableRuntimeLibrary=*/false,
-      /*enableBufferInitialization=*/false,
-      /*vectorLength=*/0,
-      /*enableVLAVectorization=*/false,
-      /*enableSIMDIndex32=*/false,
-      /*enableGPULibgen=*/false);
+      /*enableBufferInitialization=*/false);
 }
 
 std::unique_ptr<mlir::Pass> mlir::createSparsificationAndBufferizationPass(

diff  --git a/mlir/test/Dialect/SparseTensor/minipipeline_vector.mlir b/mlir/test/Dialect/SparseTensor/minipipeline_vector.mlir
new file mode 100755
index 0000000000000..2475aa5139da4
--- /dev/null
+++ b/mlir/test/Dialect/SparseTensor/minipipeline_vector.mlir
@@ -0,0 +1,43 @@
+// RUN: mlir-opt %s --sparsification-and-bufferization        | FileCheck %s --check-prefix=CHECK-NOVEC
+// RUN: mlir-opt %s --sparsification-and-bufferization="vl=8" | FileCheck %s --check-prefix=CHECK-VEC
+
+// Test to ensure we can pass optimization flags into
+// the mini sparsification and bufferization pipeline.
+
+#SV = #sparse_tensor.encoding<{ map = (d0) -> (d0 : compressed) }>
+
+#trait_sum_reduction = {
+  indexing_maps = [
+    affine_map<(i) -> (i)>,  // a
+    affine_map<(i) -> ()>    // x (scalar out)
+  ],
+  iterator_types = ["reduction"],
+  doc = "x += SUM_i a(i)"
+}
+
+//
+// CHECK-NOVEC-LABEL: func.func @sum_reduction
+// CHECK-NOVEC:       scf.for
+// CHECK-NOVEC:         arith.addf %{{.*}} %{{.*}} : f32
+// CHECK-NOVEC:       }
+//
+// CHECK-VEC-LABEL: func.func @sum_reduction
+// CHECK-VEC:       vector.insertelement
+// CHECK-VEC:       scf.for
+// CHECK-VEC:         vector.create_mask
+// CHECK-VEC:         vector.maskedload
+// CHECK-VEC:         arith.addf %{{.*}} %{{.*}} : vector<8xf32>
+// CHECK-VEC:       }
+// CHECK-VEC:       vector.reduction <add>
+//
+func.func @sum_reduction(%arga: tensor<?xf32, #SV>,
+                         %argx: tensor<f32>) -> tensor<f32> {
+  %0 = linalg.generic #trait_sum_reduction
+    ins(%arga: tensor<?xf32, #SV>)
+    outs(%argx: tensor<f32>) {
+      ^bb(%a: f32, %x: f32):
+        %0 = arith.addf %x, %a : f32
+        linalg.yield %0 : f32
+  } -> tensor<f32>
+  return %0 : tensor<f32>
+}