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