mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp

   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 &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
  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 =
 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)) {
 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(f, &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(0).getType());
 129     auto ratio =
 130         computeShapeRatio(superVectorType.getShape(), subVectorType.getShape());
 131     if (!ratio) {
 132       opInst->emitRemark("NOT MATCHED");
 133     } else {
 134       outs << "\nmatched: " << *opInst << " with shape ratio: ";
 135       llvm::interleaveComma(MutableArrayRef<int64_t>(*ratio), outs);
 136     }
 137   }
 138 }
 139
 140 static 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
 150 void VectorizerTestPass::testBackwardSlicing(llvm::raw_ostream &outs) {
 151   auto f = getOperation();
 152   outs << "\n" << f.getName();
 153
 154   SmallVector<NestedMatch, 8> matches;
 155   patternTestSlicingOps().match(f, &matches);
 156   for (auto m : matches) {
 157     SetVector<Operation *> backwardSlice;
 158     getBackwardSlice(m.getMatchedOperation(), &backwardSlice);
 159     outs << "\nmatched: " << *m.getMatchedOperation()
 160          << " backward static slice: ";
 161     for (auto *op : backwardSlice)
 162       outs << "\n" << *op;
 163   }
 164 }
 165
 166 void VectorizerTestPass::testForwardSlicing(llvm::raw_ostream &outs) {
 167   auto f = getOperation();
 168   outs << "\n" << f.getName();
 169
 170   SmallVector<NestedMatch, 8> matches;
 171   patternTestSlicingOps().match(f, &matches);
 172   for (auto m : matches) {
 173     SetVector<Operation *> forwardSlice;
 174     getForwardSlice(m.getMatchedOperation(), &forwardSlice);
 175     outs << "\nmatched: " << *m.getMatchedOperation()
 176          << " forward static slice: ";
 177     for (auto *op : forwardSlice)
 178       outs << "\n" << *op;
 179   }
 180 }
 181
 182 void VectorizerTestPass::testSlicing(llvm::raw_ostream &outs) {
 183   auto f = getOperation();
 184   outs << "\n" << f.getName();
 185
 186   SmallVector<NestedMatch, 8> matches;
 187   patternTestSlicingOps().match(f, &matches);
 188   for (auto m : matches) {
 189     SetVector<Operation *> staticSlice = getSlice(m.getMatchedOperation());
 190     outs << "\nmatched: " << *m.getMatchedOperation() << " static slice: ";
 191     for (auto *op : staticSlice)
 192       outs << "\n" << *op;
 193   }
 194 }
 195
 196 static bool customOpWithAffineMapAttribute(Operation &op) {
 197   return op.getName().getStringRef() ==
 198          VectorizerTestPass::kTestAffineMapOpName;
 199 }
 200
 201 void VectorizerTestPass::testComposeMaps(llvm::raw_ostream &outs) {
 202   auto f = getOperation();
 203
 204   using affine::matcher::Op;
 205   auto pattern = Op(customOpWithAffineMapAttribute);
 206   SmallVector<NestedMatch, 8> matches;
 207   pattern.match(f, &matches);
 208   SmallVector<AffineMap, 4> maps;
 209   maps.reserve(matches.size());
 210   for (auto m : llvm::reverse(matches)) {
 211     auto *opInst = m.getMatchedOperation();
 212     auto map =
 213         cast<AffineMapAttr>(opInst->getDiscardableAttr(
 214                                 VectorizerTestPass::kTestAffineMapAttrName))
 215             .getValue();
 216     maps.push_back(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(m) : m;
 224   }
 225   simplifyAffineMap(res).print(outs << "\nComposed map: ");
 226 }
 227
 228 /// Test for 'vectorizeAffineLoopNest' utility.
 229 void 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(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(loopsToVectorize, strategy);
 252 }
 253
 254 void 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
 290 namespace mlir {
 291 void registerVectorizerTestPass() { PassRegistration<VectorizerTestPass>(); }
 292 } // namespace mlir