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