1 //===- Passes.h - Sparse tensor pass 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 passes.
10 //
11 // In general, this file takes the approach of keeping "mechanism" (the
12 // actual steps of applying a transformation) completely separate from
13 // "policy" (heuristics for when and where to apply transformations).
14 // The only exception is in `SparseToSparseConversionStrategy`; for which,
15 // see further discussion there.
16 //
17 //===----------------------------------------------------------------------===//
18 
19 #ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES_H_
20 #define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES_H_
21 
22 #include "mlir/IR/PatternMatch.h"
23 #include "mlir/Pass/Pass.h"
24 
25 namespace mlir {
26 namespace bufferization {
27 struct OneShotBufferizationOptions;
28 } // namespace bufferization
29 
30 // Forward.
31 class TypeConverter;
32 
33 //===----------------------------------------------------------------------===//
34 // The Sparsification pass.
35 //===----------------------------------------------------------------------===//
36 
37 /// Defines a parallelization strategy. Any independent loop is a candidate
38 /// for parallelization. The loop is made parallel if (1) allowed by the
39 /// strategy (e.g., AnyStorageOuterLoop considers either a dense or sparse
40 /// outermost loop only), and (2) the generated code is an actual for-loop
41 /// (and not a co-iterating while-loop).
42 enum class SparseParallelizationStrategy {
43   kNone,
44   kDenseOuterLoop,
45   kAnyStorageOuterLoop,
46   kDenseAnyLoop,
47   kAnyStorageAnyLoop
48   // TODO: support reduction parallelization too?
49 };
50 
51 /// Converts command-line parallelization flag to the strategy enum.
52 SparseParallelizationStrategy sparseParallelizationStrategy(int32_t flag);
53 
54 /// Defines a vectorization strategy. Any inner loop is a candidate (full SIMD
55 /// for parallel loops and horizontal SIMD for reduction loops). A loop is
56 /// actually vectorized if (1) allowed by the strategy, and (2) the emitted
57 /// code is an actual for-loop (and not a co-iterating while-loop).
58 enum class SparseVectorizationStrategy {
59   kNone,
60   kDenseInnerLoop,
61   kAnyStorageInnerLoop
62 };
63 
64 /// Converts command-line vectorization flag to the strategy enum.
65 SparseVectorizationStrategy sparseVectorizationStrategy(int32_t flag);
66 
67 /// Options for the Sparsification pass.
68 struct SparsificationOptions {
SparsificationOptionsSparsificationOptions69   SparsificationOptions(SparseParallelizationStrategy p,
70                         SparseVectorizationStrategy v, unsigned vl, bool e,
71                         bool vla)
72       : parallelizationStrategy(p), vectorizationStrategy(v), vectorLength(vl),
73         enableSIMDIndex32(e), enableVLAVectorization(vla) {}
SparsificationOptionsSparsificationOptions74   SparsificationOptions()
75       : SparsificationOptions(SparseParallelizationStrategy::kNone,
76                               SparseVectorizationStrategy::kNone, 1u, false,
77                               false) {}
78   SparseParallelizationStrategy parallelizationStrategy;
79   SparseVectorizationStrategy vectorizationStrategy;
80   unsigned vectorLength;
81   bool enableSIMDIndex32;
82   bool enableVLAVectorization;
83 };
84 
85 /// Sets up sparsification rewriting rules with the given options.
86 void populateSparsificationPatterns(
87     RewritePatternSet &patterns,
88     const SparsificationOptions &options = SparsificationOptions());
89 
90 std::unique_ptr<Pass> createSparsificationPass();
91 std::unique_ptr<Pass>
92 createSparsificationPass(const SparsificationOptions &options);
93 
94 //===----------------------------------------------------------------------===//
95 // The SparseTensorConversion pass.
96 //===----------------------------------------------------------------------===//
97 
98 /// Defines a strategy for implementing sparse-to-sparse conversion.
99 /// `kAuto` leaves it up to the compiler to automatically determine
100 /// the method used.  `kViaCOO` converts the source tensor to COO and
101 /// then converts the COO to the target format.  `kDirect` converts
102 /// directly via the algorithm in <https://arxiv.org/abs/2001.02609>;
103 /// however, beware that there are many formats not supported by this
104 /// conversion method.
105 ///
106 /// The presence of the `kAuto` option violates our usual goal of keeping
107 /// policy completely separated from mechanism.  The reason it exists is
108 /// because (at present) this strategy can only be specified on a per-file
109 /// basis.  To see why this is a problem, note that `kDirect` cannot
110 /// support certain conversions; so if there is no `kAuto` setting,
111 /// then whenever a file contains a single non-`kDirect`-able conversion
112 /// the user would be forced to use `kViaCOO` for all conversions in
113 /// that file!  In the future, instead of using this enum as a `Pass`
114 /// option, we could instead move it to being an attribute on the
115 /// conversion op; at which point `kAuto` would no longer be necessary.
116 enum class SparseToSparseConversionStrategy { kAuto, kViaCOO, kDirect };
117 
118 /// Converts command-line sparse2sparse flag to the strategy enum.
119 SparseToSparseConversionStrategy sparseToSparseConversionStrategy(int32_t flag);
120 
121 /// SparseTensorConversion options.
122 struct SparseTensorConversionOptions {
SparseTensorConversionOptionsSparseTensorConversionOptions123   SparseTensorConversionOptions(SparseToSparseConversionStrategy s2s)
124       : sparseToSparseStrategy(s2s) {}
SparseTensorConversionOptionsSparseTensorConversionOptions125   SparseTensorConversionOptions()
126       : SparseTensorConversionOptions(SparseToSparseConversionStrategy::kAuto) {
127   }
128   SparseToSparseConversionStrategy sparseToSparseStrategy;
129 };
130 
131 /// Sets up sparse tensor conversion rules.
132 void populateSparseTensorConversionPatterns(
133     TypeConverter &typeConverter, RewritePatternSet &patterns,
134     const SparseTensorConversionOptions &options =
135         SparseTensorConversionOptions());
136 
137 std::unique_ptr<Pass> createSparseTensorConversionPass();
138 std::unique_ptr<Pass>
139 createSparseTensorConversionPass(const SparseTensorConversionOptions &options);
140 
141 //===----------------------------------------------------------------------===//
142 // Other rewriting rules and passes.
143 //===----------------------------------------------------------------------===//
144 
145 void populateSparseTensorRewriting(RewritePatternSet &patterns);
146 
147 std::unique_ptr<Pass> createDenseBufferizationPass(
148     const bufferization::OneShotBufferizationOptions &options);
149 
150 //===----------------------------------------------------------------------===//
151 // Registration.
152 //===----------------------------------------------------------------------===//
153 
154 /// Generate the code for registering passes.
155 #define GEN_PASS_REGISTRATION
156 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h.inc"
157 
158 } // namespace mlir
159 
160 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_PASSES_H_
161