1 | //===- VectorizerTestPass.cpp - VectorizerTestPass Pass Impl --------------===// |
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 file implements a simple testing pass for vectorization functionality. |
10 | // |
11 | //===----------------------------------------------------------------------===// |
12 | |
13 | #include "mlir/Analysis/SliceAnalysis.h" |
14 | #include "mlir/Dialect/Affine/Analysis/AffineAnalysis.h" |
15 | #include "mlir/Dialect/Affine/Analysis/NestedMatcher.h" |
16 | #include "mlir/Dialect/Affine/IR/AffineOps.h" |
17 | #include "mlir/Dialect/Affine/LoopUtils.h" |
18 | #include "mlir/Dialect/Affine/Utils.h" |
19 | #include "mlir/Dialect/Func/IR/FuncOps.h" |
20 | #include "mlir/Dialect/Utils/IndexingUtils.h" |
21 | #include "mlir/Dialect/Vector/IR/VectorOps.h" |
22 | #include "mlir/Dialect/Vector/Utils/VectorUtils.h" |
23 | #include "mlir/IR/Builders.h" |
24 | #include "mlir/IR/BuiltinTypes.h" |
25 | #include "mlir/IR/Diagnostics.h" |
26 | #include "mlir/Pass/Pass.h" |
27 | #include "mlir/Transforms/Passes.h" |
28 | |
29 | #include "llvm/ADT/STLExtras.h" |
30 | #include "llvm/Support/CommandLine.h" |
31 | #include "llvm/Support/Debug.h" |
32 | |
33 | #define DEBUG_TYPE "affine-super-vectorizer-test" |
34 | |
35 | using namespace mlir; |
36 | using namespace mlir::affine; |
37 | |
38 | static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options" ); |
39 | |
40 | namespace { |
41 | struct VectorizerTestPass |
42 | : public PassWrapper<VectorizerTestPass, OperationPass<func::FuncOp>> { |
43 | MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VectorizerTestPass) |
44 | |
45 | static constexpr auto kTestAffineMapOpName = "test_affine_map" ; |
46 | static constexpr auto kTestAffineMapAttrName = "affine_map" ; |
47 | void getDependentDialects(DialectRegistry ®istry) const override { |
48 | registry.insert<vector::VectorDialect>(); |
49 | } |
50 | StringRef getArgument() const final { return "affine-super-vectorizer-test" ; } |
51 | StringRef getDescription() const final { |
52 | return "Tests vectorizer standalone functionality." ; |
53 | } |
54 | |
55 | VectorizerTestPass() = default; |
56 | VectorizerTestPass(const VectorizerTestPass &pass) : PassWrapper(pass){}; |
57 | |
58 | ListOption<int> clTestVectorShapeRatio{ |
59 | *this, "vector-shape-ratio" , |
60 | llvm::cl::desc("Specify the HW vector size for vectorization" )}; |
61 | Option<bool> clTestForwardSlicingAnalysis{ |
62 | *this, "forward-slicing" , |
63 | llvm::cl::desc( |
64 | "Enable testing forward static slicing and topological sort " |
65 | "functionalities" )}; |
66 | Option<bool> clTestBackwardSlicingAnalysis{ |
67 | *this, "backward-slicing" , |
68 | llvm::cl::desc("Enable testing backward static slicing and " |
69 | "topological sort functionalities" )}; |
70 | Option<bool> clTestSlicingAnalysis{ |
71 | *this, "slicing" , |
72 | llvm::cl::desc("Enable testing static slicing and topological sort " |
73 | "functionalities" )}; |
74 | Option<bool> clTestComposeMaps{ |
75 | *this, "compose-maps" , |
76 | llvm::cl::desc("Enable testing the composition of AffineMap where each " |
77 | "AffineMap in the composition is specified as the " |
78 | "affine_map attribute " |
79 | "in a constant op." )}; |
80 | Option<bool> clTestVecAffineLoopNest{ |
81 | *this, "vectorize-affine-loop-nest" , |
82 | llvm::cl::desc( |
83 | "Enable testing for the 'vectorizeAffineLoopNest' utility by " |
84 | "vectorizing the outermost loops found" )}; |
85 | |
86 | void runOnOperation() override; |
87 | void testVectorShapeRatio(llvm::raw_ostream &outs); |
88 | void testForwardSlicing(llvm::raw_ostream &outs); |
89 | void testBackwardSlicing(llvm::raw_ostream &outs); |
90 | void testSlicing(llvm::raw_ostream &outs); |
91 | void testComposeMaps(llvm::raw_ostream &outs); |
92 | |
93 | /// Test for 'vectorizeAffineLoopNest' utility. |
94 | void testVecAffineLoopNest(llvm::raw_ostream &outs); |
95 | }; |
96 | |
97 | } // namespace |
98 | |
99 | void VectorizerTestPass::testVectorShapeRatio(llvm::raw_ostream &outs) { |
100 | auto f = getOperation(); |
101 | using affine::matcher::Op; |
102 | SmallVector<int64_t, 8> shape(clTestVectorShapeRatio.begin(), |
103 | clTestVectorShapeRatio.end()); |
104 | auto subVectorType = |
105 | VectorType::get(shape, FloatType::getF32(f.getContext())); |
106 | // Only filter operations that operate on a strict super-vector and have one |
107 | // return. This makes testing easier. |
108 | auto filter = [&](Operation &op) { |
109 | assert(subVectorType.getElementType().isF32() && |
110 | "Only f32 supported for now" ); |
111 | if (!mlir::matcher::operatesOnSuperVectorsOf(op, subVectorType: subVectorType)) { |
112 | return false; |
113 | } |
114 | if (op.getNumResults() != 1) { |
115 | return false; |
116 | } |
117 | return true; |
118 | }; |
119 | auto pat = Op(filter); |
120 | SmallVector<NestedMatch, 8> matches; |
121 | pat.match(op: f, matches: &matches); |
122 | for (auto m : matches) { |
123 | auto *opInst = m.getMatchedOperation(); |
124 | // This is a unit test that only checks and prints shape ratio. |
125 | // As a consequence we write only Ops with a single return type for the |
126 | // purpose of this test. If we need to test more intricate behavior in the |
127 | // future we can always extend. |
128 | auto superVectorType = cast<VectorType>(opInst->getResult(idx: 0).getType()); |
129 | auto ratio = |
130 | computeShapeRatio(superVectorType.getShape(), subVectorType.getShape()); |
131 | if (!ratio) { |
132 | opInst->emitRemark(message: "NOT MATCHED" ); |
133 | } else { |
134 | outs << "\nmatched: " << *opInst << " with shape ratio: " ; |
135 | llvm::interleaveComma(c: MutableArrayRef<int64_t>(*ratio), os&: outs); |
136 | } |
137 | } |
138 | } |
139 | |
140 | static NestedPattern patternTestSlicingOps() { |
141 | using affine::matcher::Op; |
142 | // Match all operations with the kTestSlicingOpName name. |
143 | auto filter = [](Operation &op) { |
144 | // Just use a custom op name for this test, it makes life easier. |
145 | return op.getName().getStringRef() == "slicing-test-op" ; |
146 | }; |
147 | return Op(filter); |
148 | } |
149 | |
150 | void VectorizerTestPass::testBackwardSlicing(llvm::raw_ostream &outs) { |
151 | auto f = getOperation(); |
152 | outs << "\n" << f.getName(); |
153 | |
154 | SmallVector<NestedMatch, 8> matches; |
155 | patternTestSlicingOps().match(op: f, matches: &matches); |
156 | for (auto m : matches) { |
157 | SetVector<Operation *> backwardSlice; |
158 | getBackwardSlice(op: m.getMatchedOperation(), backwardSlice: &backwardSlice); |
159 | outs << "\nmatched: " << *m.getMatchedOperation() |
160 | << " backward static slice: " ; |
161 | for (auto *op : backwardSlice) |
162 | outs << "\n" << *op; |
163 | } |
164 | } |
165 | |
166 | void VectorizerTestPass::testForwardSlicing(llvm::raw_ostream &outs) { |
167 | auto f = getOperation(); |
168 | outs << "\n" << f.getName(); |
169 | |
170 | SmallVector<NestedMatch, 8> matches; |
171 | patternTestSlicingOps().match(op: f, matches: &matches); |
172 | for (auto m : matches) { |
173 | SetVector<Operation *> forwardSlice; |
174 | getForwardSlice(op: m.getMatchedOperation(), forwardSlice: &forwardSlice); |
175 | outs << "\nmatched: " << *m.getMatchedOperation() |
176 | << " forward static slice: " ; |
177 | for (auto *op : forwardSlice) |
178 | outs << "\n" << *op; |
179 | } |
180 | } |
181 | |
182 | void VectorizerTestPass::testSlicing(llvm::raw_ostream &outs) { |
183 | auto f = getOperation(); |
184 | outs << "\n" << f.getName(); |
185 | |
186 | SmallVector<NestedMatch, 8> matches; |
187 | patternTestSlicingOps().match(op: f, matches: &matches); |
188 | for (auto m : matches) { |
189 | SetVector<Operation *> staticSlice = getSlice(op: m.getMatchedOperation()); |
190 | outs << "\nmatched: " << *m.getMatchedOperation() << " static slice: " ; |
191 | for (auto *op : staticSlice) |
192 | outs << "\n" << *op; |
193 | } |
194 | } |
195 | |
196 | static bool customOpWithAffineMapAttribute(Operation &op) { |
197 | return op.getName().getStringRef() == |
198 | VectorizerTestPass::kTestAffineMapOpName; |
199 | } |
200 | |
201 | void VectorizerTestPass::testComposeMaps(llvm::raw_ostream &outs) { |
202 | auto f = getOperation(); |
203 | |
204 | using affine::matcher::Op; |
205 | auto pattern = Op(filter: customOpWithAffineMapAttribute); |
206 | SmallVector<NestedMatch, 8> matches; |
207 | pattern.match(op: f, matches: &matches); |
208 | SmallVector<AffineMap, 4> maps; |
209 | maps.reserve(N: matches.size()); |
210 | for (auto m : llvm::reverse(C&: matches)) { |
211 | auto *opInst = m.getMatchedOperation(); |
212 | auto map = |
213 | cast<AffineMapAttr>(opInst->getDiscardableAttr( |
214 | name: VectorizerTestPass::kTestAffineMapAttrName)) |
215 | .getValue(); |
216 | maps.push_back(Elt: map); |
217 | } |
218 | if (maps.empty()) |
219 | // Nothing to compose |
220 | return; |
221 | AffineMap res; |
222 | for (auto m : maps) { |
223 | res = res ? res.compose(map: m) : m; |
224 | } |
225 | simplifyAffineMap(map: res).print(os&: outs << "\nComposed map: " ); |
226 | } |
227 | |
228 | /// Test for 'vectorizeAffineLoopNest' utility. |
229 | void VectorizerTestPass::testVecAffineLoopNest(llvm::raw_ostream &outs) { |
230 | std::vector<SmallVector<AffineForOp, 2>> loops; |
231 | gatherLoops(getOperation(), loops); |
232 | |
233 | // Expected only one loop nest. |
234 | if (loops.empty() || loops[0].size() != 1) |
235 | return; |
236 | |
237 | // We vectorize the outermost loop found with VF=4. |
238 | AffineForOp outermostLoop = loops[0][0]; |
239 | VectorizationStrategy strategy; |
240 | strategy.vectorSizes.push_back(Elt: 4 /*vectorization factor*/); |
241 | strategy.loopToVectorDim[outermostLoop] = 0; |
242 | |
243 | ReductionLoopMap reductionLoops; |
244 | SmallVector<LoopReduction, 2> reductions; |
245 | if (!isLoopParallel(outermostLoop, &reductions)) { |
246 | outs << "Outermost loop cannot be parallel\n" ; |
247 | return; |
248 | } |
249 | std::vector<SmallVector<AffineForOp, 2>> loopsToVectorize; |
250 | loopsToVectorize.push_back({outermostLoop}); |
251 | (void)vectorizeAffineLoopNest(loops&: loopsToVectorize, strategy); |
252 | } |
253 | |
254 | void VectorizerTestPass::runOnOperation() { |
255 | // Only support single block functions at this point. |
256 | func::FuncOp f = getOperation(); |
257 | if (!llvm::hasSingleElement(f)) |
258 | return; |
259 | |
260 | std::string str; |
261 | llvm::raw_string_ostream outs(str); |
262 | |
263 | { // Tests that expect a NestedPatternContext to be allocated externally. |
264 | NestedPatternContext mlContext; |
265 | |
266 | if (!clTestVectorShapeRatio.empty()) |
267 | testVectorShapeRatio(outs); |
268 | |
269 | if (clTestForwardSlicingAnalysis) |
270 | testForwardSlicing(outs); |
271 | |
272 | if (clTestBackwardSlicingAnalysis) |
273 | testBackwardSlicing(outs); |
274 | |
275 | if (clTestSlicingAnalysis) |
276 | testSlicing(outs); |
277 | |
278 | if (clTestComposeMaps) |
279 | testComposeMaps(outs); |
280 | } |
281 | |
282 | if (clTestVecAffineLoopNest) |
283 | testVecAffineLoopNest(outs); |
284 | |
285 | if (!outs.str().empty()) { |
286 | emitRemark(UnknownLoc::get(&getContext()), outs.str()); |
287 | } |
288 | } |
289 | |
290 | namespace mlir { |
291 | void registerVectorizerTestPass() { PassRegistration<VectorizerTestPass>(); } |
292 | } // namespace mlir |
293 | |