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