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