1 //===- Passes.h - Sparse tensor pipeline entry points -----------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This header file defines prototypes of all sparse tensor pipelines.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #ifndef MLIR_DIALECT_SPARSETENSOR_PIPELINES_PASSES_H_
14 #define MLIR_DIALECT_SPARSETENSOR_PIPELINES_PASSES_H_
15 
16 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
17 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
18 #include "mlir/Pass/PassOptions.h"
19 
20 using namespace mlir::detail;
21 using namespace llvm::cl;
22 
23 namespace mlir {
24 namespace sparse_tensor {
25 
26 /// Options for the "sparse-compiler" pipeline.  So far this only contains
27 /// a subset of the options that can be set for the underlying passes,
28 /// because it must be manually kept in sync with the tablegen files
29 /// for those passes.
30 struct SparseCompilerOptions
31     : public PassPipelineOptions<SparseCompilerOptions> {
32   // These options must be kept in sync with `SparsificationBase`.
33   PassOptions::Option<int32_t> parallelization{
34       *this, "parallelization-strategy",
35       desc("Set the parallelization strategy"), init(0)};
36   PassOptions::Option<int32_t> vectorization{
37       *this, "vectorization-strategy", desc("Set the vectorization strategy"),
38       init(0)};
39   PassOptions::Option<int32_t> vectorLength{
40       *this, "vl", desc("Set the vector length"), init(1)};
41   PassOptions::Option<bool> enableSIMDIndex32{
42       *this, "enable-simd-index32",
43       desc("Enable i32 indexing into vectors (for efficiency)"), init(false)};
44   PassOptions::Option<bool> enableVLAVectorization{
45       *this, "enable-vla-vectorization",
46       desc("Enable vector length agnostic vectorization"), init(false)};
47   PassOptions::Option<bool> testBufferizationAnalysisOnly{
48       *this, "test-bufferization-analysis-only",
49       desc("Run only the inplacability analysis"), init(false)};
50 
51   /// Projects out the options for `createSparsificationPass`.
sparsificationOptionsSparseCompilerOptions52   SparsificationOptions sparsificationOptions() const {
53     return SparsificationOptions(sparseParallelizationStrategy(parallelization),
54                                  sparseVectorizationStrategy(vectorization),
55                                  vectorLength, enableSIMDIndex32,
56                                  enableVLAVectorization);
57   }
58 
59   // These options must be kept in sync with `SparseTensorConversionBase`.
60   PassOptions::Option<int32_t> sparseToSparse{
61       *this, "s2s-strategy",
62       desc("Set the strategy for sparse-to-sparse conversion"), init(0)};
63 
64   /// Projects out the options for `createSparsificationPass`.
sparseTensorConversionOptionsSparseCompilerOptions65   SparseTensorConversionOptions sparseTensorConversionOptions() const {
66     return SparseTensorConversionOptions(
67         sparseToSparseConversionStrategy(sparseToSparse));
68   }
69 
70   // These options must be kept in sync with `ConvertVectorToLLVMBase`.
71   // TODO(wrengr): does `indexOptimizations` differ from `enableSIMDIndex32`?
72   PassOptions::Option<bool> reassociateFPReductions{
73       *this, "reassociate-fp-reductions",
74       desc("Allows llvm to reassociate floating-point reductions for speed"),
75       init(false)};
76   PassOptions::Option<bool> indexOptimizations{
77       *this, "enable-index-optimizations",
78       desc("Allows compiler to assume indices fit in 32-bit if that yields "
79            "faster code"),
80       init(true)};
81   PassOptions::Option<bool> amx{
82       *this, "enable-amx",
83       desc("Enables the use of AMX dialect while lowering the vector dialect."),
84       init(false)};
85   PassOptions::Option<bool> armNeon{*this, "enable-arm-neon",
86                                     desc("Enables the use of ArmNeon dialect "
87                                          "while lowering the vector dialect."),
88                                     init(false)};
89   PassOptions::Option<bool> armSVE{*this, "enable-arm-sve",
90                                    desc("Enables the use of ArmSVE dialect "
91                                         "while lowering the vector dialect."),
92                                    init(false)};
93   PassOptions::Option<bool> x86Vector{
94       *this, "enable-x86vector",
95       desc("Enables the use of X86Vector dialect while lowering the vector "
96            "dialect."),
97       init(false)};
98 
99   /// Projects out the options for `createConvertVectorToLLVMPass`.
lowerVectorToLLVMOptionsSparseCompilerOptions100   LowerVectorToLLVMOptions lowerVectorToLLVMOptions() const {
101     LowerVectorToLLVMOptions opts{};
102     opts.enableReassociateFPReductions(reassociateFPReductions);
103     opts.enableIndexOptimizations(indexOptimizations);
104     opts.enableArmNeon(armNeon);
105     opts.enableArmSVE(armSVE);
106     opts.enableAMX(amx);
107     opts.enableX86Vector(x86Vector);
108     return opts;
109   }
110 };
111 
112 //===----------------------------------------------------------------------===//
113 // Building and Registering.
114 //===----------------------------------------------------------------------===//
115 
116 /// Adds the "sparse-compiler" pipeline to the `OpPassManager`.  This
117 /// is the standard pipeline for taking sparsity-agnostic IR using
118 /// the sparse-tensor type and lowering it to LLVM IR with concrete
119 /// representations and algorithms for sparse tensors.
120 void buildSparseCompiler(OpPassManager &pm,
121                          const SparseCompilerOptions &options);
122 
123 /// Registers all pipelines for the `sparse_tensor` dialect.  At present,
124 /// this includes only "sparse-compiler".
125 void registerSparseTensorPipelines();
126 
127 } // namespace sparse_tensor
128 } // namespace mlir
129 
130 #endif // MLIR_DIALECT_SPARSETENSOR_PIPELINES_PASSES_H_
131