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