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> clTestVecAffineLoopNest(
66     "vectorize-affine-loop-nest",
67     llvm::cl::desc(
68         "Enable testing for the 'vectorizeAffineLoopNest' utility by "
69         "vectorizing the outermost loops found"),
70     llvm::cl::cat(clOptionsCategory));
71 
72 namespace {
73 struct VectorizerTestPass
74     : public PassWrapper<VectorizerTestPass, FunctionPass> {
75   static constexpr auto kTestAffineMapOpName = "test_affine_map";
76   static constexpr auto kTestAffineMapAttrName = "affine_map";
77   void getDependentDialects(DialectRegistry &registry) const override {
78     registry.insert<vector::VectorDialect>();
79   }
80 
81   void runOnFunction() override;
82   void testVectorShapeRatio(llvm::raw_ostream &outs);
83   void testForwardSlicing(llvm::raw_ostream &outs);
84   void testBackwardSlicing(llvm::raw_ostream &outs);
85   void testSlicing(llvm::raw_ostream &outs);
86   void testComposeMaps(llvm::raw_ostream &outs);
87 
88   /// Test for 'vectorizeAffineLoopNest' utility.
89   void testVecAffineLoopNest();
90 };
91 
92 } // end anonymous namespace
93 
94 void VectorizerTestPass::testVectorShapeRatio(llvm::raw_ostream &outs) {
95   auto f = getFunction();
96   using matcher::Op;
97   SmallVector<int64_t, 8> shape(clTestVectorShapeRatio.begin(),
98                                 clTestVectorShapeRatio.end());
99   auto subVectorType =
100       VectorType::get(shape, FloatType::getF32(f.getContext()));
101   // Only filter operations that operate on a strict super-vector and have one
102   // return. This makes testing easier.
103   auto filter = [&](Operation &op) {
104     assert(subVectorType.getElementType().isF32() &&
105            "Only f32 supported for now");
106     if (!matcher::operatesOnSuperVectorsOf(op, subVectorType)) {
107       return false;
108     }
109     if (op.getNumResults() != 1) {
110       return false;
111     }
112     return true;
113   };
114   auto pat = Op(filter);
115   SmallVector<NestedMatch, 8> matches;
116   pat.match(f, &matches);
117   for (auto m : matches) {
118     auto *opInst = m.getMatchedOperation();
119     // This is a unit test that only checks and prints shape ratio.
120     // As a consequence we write only Ops with a single return type for the
121     // purpose of this test. If we need to test more intricate behavior in the
122     // future we can always extend.
123     auto superVectorType = opInst->getResult(0).getType().cast<VectorType>();
124     auto ratio = shapeRatio(superVectorType, subVectorType);
125     if (!ratio.hasValue()) {
126       opInst->emitRemark("NOT MATCHED");
127     } else {
128       outs << "\nmatched: " << *opInst << " with shape ratio: ";
129       llvm::interleaveComma(MutableArrayRef<int64_t>(*ratio), outs);
130     }
131   }
132 }
133 
134 static NestedPattern patternTestSlicingOps() {
135   using matcher::Op;
136   // Match all operations with the kTestSlicingOpName name.
137   auto filter = [](Operation &op) {
138     // Just use a custom op name for this test, it makes life easier.
139     return op.getName().getStringRef() == "slicing-test-op";
140   };
141   return Op(filter);
142 }
143 
144 void VectorizerTestPass::testBackwardSlicing(llvm::raw_ostream &outs) {
145   auto f = getFunction();
146   outs << "\n" << f.getName();
147 
148   SmallVector<NestedMatch, 8> matches;
149   patternTestSlicingOps().match(f, &matches);
150   for (auto m : matches) {
151     SetVector<Operation *> backwardSlice;
152     getBackwardSlice(m.getMatchedOperation(), &backwardSlice);
153     outs << "\nmatched: " << *m.getMatchedOperation()
154          << " backward static slice: ";
155     for (auto *op : backwardSlice)
156       outs << "\n" << *op;
157   }
158 }
159 
160 void VectorizerTestPass::testForwardSlicing(llvm::raw_ostream &outs) {
161   auto f = getFunction();
162   outs << "\n" << f.getName();
163 
164   SmallVector<NestedMatch, 8> matches;
165   patternTestSlicingOps().match(f, &matches);
166   for (auto m : matches) {
167     SetVector<Operation *> forwardSlice;
168     getForwardSlice(m.getMatchedOperation(), &forwardSlice);
169     outs << "\nmatched: " << *m.getMatchedOperation()
170          << " forward static slice: ";
171     for (auto *op : forwardSlice)
172       outs << "\n" << *op;
173   }
174 }
175 
176 void VectorizerTestPass::testSlicing(llvm::raw_ostream &outs) {
177   auto f = getFunction();
178   outs << "\n" << f.getName();
179 
180   SmallVector<NestedMatch, 8> matches;
181   patternTestSlicingOps().match(f, &matches);
182   for (auto m : matches) {
183     SetVector<Operation *> staticSlice = getSlice(m.getMatchedOperation());
184     outs << "\nmatched: " << *m.getMatchedOperation() << " static slice: ";
185     for (auto *op : staticSlice)
186       outs << "\n" << *op;
187   }
188 }
189 
190 static bool customOpWithAffineMapAttribute(Operation &op) {
191   return op.getName().getStringRef() ==
192          VectorizerTestPass::kTestAffineMapOpName;
193 }
194 
195 void VectorizerTestPass::testComposeMaps(llvm::raw_ostream &outs) {
196   auto f = getFunction();
197 
198   using matcher::Op;
199   auto pattern = Op(customOpWithAffineMapAttribute);
200   SmallVector<NestedMatch, 8> matches;
201   pattern.match(f, &matches);
202   SmallVector<AffineMap, 4> maps;
203   maps.reserve(matches.size());
204   for (auto m : llvm::reverse(matches)) {
205     auto *opInst = m.getMatchedOperation();
206     auto map = opInst->getAttr(VectorizerTestPass::kTestAffineMapAttrName)
207                    .cast<AffineMapAttr>()
208                    .getValue();
209     maps.push_back(map);
210   }
211   AffineMap res;
212   for (auto m : maps) {
213     res = res ? res.compose(m) : m;
214   }
215   simplifyAffineMap(res).print(outs << "\nComposed map: ");
216 }
217 
218 /// Test for 'vectorizeAffineLoopNest' utility.
219 void VectorizerTestPass::testVecAffineLoopNest() {
220   std::vector<SmallVector<AffineForOp, 2>> loops;
221   gatherLoops(getFunction(), loops);
222 
223   // Expected only one loop nest.
224   if (loops.empty() || loops[0].size() != 1)
225     return;
226 
227   // We vectorize the outermost loop found with VF=4.
228   AffineForOp outermostLoop = loops[0][0];
229   VectorizationStrategy strategy;
230   strategy.vectorSizes.push_back(4 /*vectorization factor*/);
231   strategy.loopToVectorDim[outermostLoop] = 0;
232   std::vector<SmallVector<AffineForOp, 2>> loopsToVectorize;
233   loopsToVectorize.push_back({outermostLoop});
234   (void)vectorizeAffineLoopNest(loopsToVectorize, strategy);
235 }
236 
237 void VectorizerTestPass::runOnFunction() {
238   // Only support single block functions at this point.
239   FuncOp f = getFunction();
240   if (!llvm::hasSingleElement(f))
241     return;
242 
243   std::string str;
244   llvm::raw_string_ostream outs(str);
245 
246   { // Tests that expect a NestedPatternContext to be allocated externally.
247     NestedPatternContext mlContext;
248 
249     if (!clTestVectorShapeRatio.empty())
250       testVectorShapeRatio(outs);
251 
252     if (clTestForwardSlicingAnalysis)
253       testForwardSlicing(outs);
254 
255     if (clTestBackwardSlicingAnalysis)
256       testBackwardSlicing(outs);
257 
258     if (clTestSlicingAnalysis)
259       testSlicing(outs);
260 
261     if (clTestComposeMaps)
262       testComposeMaps(outs);
263   }
264 
265   if (clTestVecAffineLoopNest)
266     testVecAffineLoopNest();
267 
268   if (!outs.str().empty()) {
269     emitRemark(UnknownLoc::get(&getContext()), outs.str());
270   }
271 }
272 
273 namespace mlir {
274 void registerVectorizerTestPass() {
275   PassRegistration<VectorizerTestPass> pass(
276       "affine-super-vectorizer-test",
277       "Tests vectorizer standalone functionality.");
278 }
279 } // namespace mlir
280