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