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