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

source code of mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp