Import MLIR into the LLVM tree

1 files changed, 292 insertions, 0 deletions

diff --git a/mlir/test/lib/Transforms/TestVectorizationUtils.cpp b/mlir/test/lib/Transforms/TestVectorizationUtils.cpp
new file mode 100644
index 00000000000..6f4d948e55f
--- /dev/null
+++ b/mlir/test/lib/Transforms/TestVectorizationUtils.cpp
@@ -0,0 +1,292 @@
+//===- VectorizerTestPass.cpp - VectorizerTestPass Pass Impl --------------===//
+//
+// Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a simple testing pass for vectorization functionality.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/AffineAnalysis.h"
+#include "mlir/Analysis/NestedMatcher.h"
+#include "mlir/Analysis/SliceAnalysis.h"
+#include "mlir/Dialect/AffineOps/AffineOps.h"
+#include "mlir/Dialect/VectorOps/Utils.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/StandardTypes.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/Functional.h"
+#include "mlir/Support/STLExtras.h"
+#include "mlir/Transforms/Passes.h"
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "affine-vectorizer-test"
+
+using namespace mlir;
+
+using llvm::SetVector;
+
+using functional::map;
+
+static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options");
+
+static llvm::cl::list<int> clTestVectorShapeRatio(
+    "vector-shape-ratio",
+    llvm::cl::desc("Specify the HW vector size for vectorization"),
+    llvm::cl::ZeroOrMore, llvm::cl::cat(clOptionsCategory));
+static llvm::cl::opt<bool> clTestForwardSlicingAnalysis(
+    "forward-slicing",
+    llvm::cl::desc("Enable testing forward static slicing and topological sort "
+                   "functionalities"),
+    llvm::cl::cat(clOptionsCategory));
+static llvm::cl::opt<bool> clTestBackwardSlicingAnalysis(
+    "backward-slicing",
+    llvm::cl::desc("Enable testing backward static slicing and "
+                   "topological sort functionalities"),
+    llvm::cl::cat(clOptionsCategory));
+static llvm::cl::opt<bool> clTestSlicingAnalysis(
+    "slicing",
+    llvm::cl::desc("Enable testing static slicing and topological sort "
+                   "functionalities"),
+    llvm::cl::cat(clOptionsCategory));
+static llvm::cl::opt<bool> clTestComposeMaps(
+    "compose-maps",
+    llvm::cl::desc(
+        "Enable testing the composition of AffineMap where each "
+        "AffineMap in the composition is specified as the affine_map attribute "
+        "in a constant op."),
+    llvm::cl::cat(clOptionsCategory));
+static llvm::cl::opt<bool> clTestNormalizeMaps(
+    "normalize-maps",
+    llvm::cl::desc(
+        "Enable testing the normalization of AffineAffineApplyOp "
+        "where each AffineAffineApplyOp in the composition is a single output "
+        "operation."),
+    llvm::cl::cat(clOptionsCategory));
+
+namespace {
+struct VectorizerTestPass : public FunctionPass<VectorizerTestPass> {
+  static constexpr auto kTestAffineMapOpName = "test_affine_map";
+  static constexpr auto kTestAffineMapAttrName = "affine_map";
+
+  void runOnFunction() override;
+  void testVectorShapeRatio(llvm::raw_ostream &outs);
+  void testForwardSlicing(llvm::raw_ostream &outs);
+  void testBackwardSlicing(llvm::raw_ostream &outs);
+  void testSlicing(llvm::raw_ostream &outs);
+  void testComposeMaps(llvm::raw_ostream &outs);
+  void testNormalizeMaps();
+};
+
+} // end anonymous namespace
+
+void VectorizerTestPass::testVectorShapeRatio(llvm::raw_ostream &outs) {
+  auto f = getFunction();
+  using matcher::Op;
+  SmallVector<int64_t, 8> shape(clTestVectorShapeRatio.begin(),
+                                clTestVectorShapeRatio.end());
+  auto subVectorType =
+      VectorType::get(shape, FloatType::getF32(f.getContext()));
+  // Only filter operations that operate on a strict super-vector and have one
+  // return. This makes testing easier.
+  auto filter = [&](Operation &op) {
+    assert(subVectorType.getElementType().isF32() &&
+           "Only f32 supported for now");
+    if (!matcher::operatesOnSuperVectorsOf(op, subVectorType)) {
+      return false;
+    }
+    if (op.getNumResults() != 1) {
+      return false;
+    }
+    return true;
+  };
+  auto pat = Op(filter);
+  SmallVector<NestedMatch, 8> matches;
+  pat.match(f, &matches);
+  for (auto m : matches) {
+    auto *opInst = m.getMatchedOperation();
+    // This is a unit test that only checks and prints shape ratio.
+    // As a consequence we write only Ops with a single return type for the
+    // purpose of this test. If we need to test more intricate behavior in the
+    // future we can always extend.
+    auto superVectorType = opInst->getResult(0)->getType().cast<VectorType>();
+    auto ratio = shapeRatio(superVectorType, subVectorType);
+    if (!ratio.hasValue()) {
+      opInst->emitRemark("NOT MATCHED");
+    } else {
+      outs << "\nmatched: " << *opInst << " with shape ratio: ";
+      interleaveComma(MutableArrayRef<int64_t>(*ratio), outs);
+    }
+  }
+}
+
+static NestedPattern patternTestSlicingOps() {
+  using functional::map;
+  using matcher::Op;
+  // Match all operations with the kTestSlicingOpName name.
+  auto filter = [](Operation &op) {
+    // Just use a custom op name for this test, it makes life easier.
+    return op.getName().getStringRef() == "slicing-test-op";
+  };
+  return Op(filter);
+}
+
+void VectorizerTestPass::testBackwardSlicing(llvm::raw_ostream &outs) {
+  auto f = getFunction();
+  outs << "\n" << f.getName();
+
+  SmallVector<NestedMatch, 8> matches;
+  patternTestSlicingOps().match(f, &matches);
+  for (auto m : matches) {
+    SetVector<Operation *> backwardSlice;
+    getBackwardSlice(m.getMatchedOperation(), &backwardSlice);
+    outs << "\nmatched: " << *m.getMatchedOperation()
+         << " backward static slice: ";
+    for (auto *op : backwardSlice)
+      outs << "\n" << *op;
+  }
+}
+
+void VectorizerTestPass::testForwardSlicing(llvm::raw_ostream &outs) {
+  auto f = getFunction();
+  outs << "\n" << f.getName();
+
+  SmallVector<NestedMatch, 8> matches;
+  patternTestSlicingOps().match(f, &matches);
+  for (auto m : matches) {
+    SetVector<Operation *> forwardSlice;
+    getForwardSlice(m.getMatchedOperation(), &forwardSlice);
+    outs << "\nmatched: " << *m.getMatchedOperation()
+         << " forward static slice: ";
+    for (auto *op : forwardSlice)
+      outs << "\n" << *op;
+  }
+}
+
+void VectorizerTestPass::testSlicing(llvm::raw_ostream &outs) {
+  auto f = getFunction();
+  outs << "\n" << f.getName();
+
+  SmallVector<NestedMatch, 8> matches;
+  patternTestSlicingOps().match(f, &matches);
+  for (auto m : matches) {
+    SetVector<Operation *> staticSlice = getSlice(m.getMatchedOperation());
+    outs << "\nmatched: " << *m.getMatchedOperation() << " static slice: ";
+    for (auto *op : staticSlice)
+      outs << "\n" << *op;
+  }
+}
+
+static bool customOpWithAffineMapAttribute(Operation &op) {
+  return op.getName().getStringRef() ==
+         VectorizerTestPass::kTestAffineMapOpName;
+}
+
+void VectorizerTestPass::testComposeMaps(llvm::raw_ostream &outs) {
+  auto f = getFunction();
+
+  using matcher::Op;
+  auto pattern = Op(customOpWithAffineMapAttribute);
+  SmallVector<NestedMatch, 8> matches;
+  pattern.match(f, &matches);
+  SmallVector<AffineMap, 4> maps;
+  maps.reserve(matches.size());
+  for (auto m : llvm::reverse(matches)) {
+    auto *opInst = m.getMatchedOperation();
+    auto map = opInst->getAttr(VectorizerTestPass::kTestAffineMapAttrName)
+                   .cast<AffineMapAttr>()
+                   .getValue();
+    maps.push_back(map);
+  }
+  AffineMap res;
+  for (auto m : maps) {
+    res = res ? res.compose(m) : m;
+  }
+  simplifyAffineMap(res).print(outs << "\nComposed map: ");
+}
+
+static bool affineApplyOp(Operation &op) { return isa<AffineApplyOp>(op); }
+
+static bool singleResultAffineApplyOpWithoutUses(Operation &op) {
+  auto app = dyn_cast<AffineApplyOp>(op);
+  return app && app.use_empty();
+}
+
+void VectorizerTestPass::testNormalizeMaps() {
+  using matcher::Op;
+
+  auto f = getFunction();
+
+  // Save matched AffineApplyOp that all need to be erased in the end.
+  auto pattern = Op(affineApplyOp);
+  SmallVector<NestedMatch, 8> toErase;
+  pattern.match(f, &toErase);
+  {
+    // Compose maps.
+    auto pattern = Op(singleResultAffineApplyOpWithoutUses);
+    SmallVector<NestedMatch, 8> matches;
+    pattern.match(f, &matches);
+    for (auto m : matches) {
+      auto app = cast<AffineApplyOp>(m.getMatchedOperation());
+      OpBuilder b(m.getMatchedOperation());
+      SmallVector<Value, 8> operands(app.getOperands());
+      makeComposedAffineApply(b, app.getLoc(), app.getAffineMap(), operands);
+    }
+  }
+  // We should now be able to erase everything in reverse order in this test.
+  for (auto m : llvm::reverse(toErase)) {
+    m.getMatchedOperation()->erase();
+  }
+}
+
+void VectorizerTestPass::runOnFunction() {
+  // Thread-safe RAII local context, BumpPtrAllocator freed on exit.
+  NestedPatternContext mlContext;
+
+  // Only support single block functions at this point.
+  FuncOp f = getFunction();
+  if (f.getBlocks().size() != 1)
+    return;
+
+  std::string str;
+  llvm::raw_string_ostream outs(str);
+
+  if (!clTestVectorShapeRatio.empty())
+    testVectorShapeRatio(outs);
+
+  if (clTestForwardSlicingAnalysis)
+    testForwardSlicing(outs);
+
+  if (clTestBackwardSlicingAnalysis)
+    testBackwardSlicing(outs);
+
+  if (clTestSlicingAnalysis)
+    testSlicing(outs);
+
+  if (clTestComposeMaps)
+    testComposeMaps(outs);
+
+  if (clTestNormalizeMaps)
+    testNormalizeMaps();
+
+  if (!outs.str().empty()) {
+    emitRemark(UnknownLoc::get(&getContext()), outs.str());
+  }
+}
+
+std::unique_ptr<OpPassBase<FuncOp>> mlir::createVectorizerTestPass() {
+  return std::make_unique<VectorizerTestPass>();
+}
+
+static PassRegistration<VectorizerTestPass>
+    pass("affine-vectorizer-test",
+         "Tests vectorizer standalone functionality.");
+
+#undef DEBUG_TYPE