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