[MLIR][VectorAnalysis] Add a VectorAnalysis and standalone tests

This CL adds some vector support in prevision of the upcoming vector
materialization pass. In particular this CL adds 2 functions to:
1. compute the multiplicity of a subvector shape in a supervector shape;
2. help match operations on strict super-vectors. This is defined for a given
subvector shape as an operation that manipulates a vector type that is an
integral multiple of the subtype, with multiplicity at least 2.

This CL also adds a TestUtil pass where we can dump arbitrary testing of
functions and analysis that operate at a much smaller granularity than a pass
(e.g. an analysis for which it is convenient to write a bit of artificial MLIR
and write some custom test). This is in order to keep using Filecheck for
things that essentially look and feel like C++ unit tests.

PiperOrigin-RevId: 222250910

1 files changed, 116 insertions, 0 deletions

diff --git a/mlir/lib/Transforms/Vectorization/VectorizerTestPass.cpp b/mlir/lib/Transforms/Vectorization/VectorizerTestPass.cpp
new file mode 100644
index 00000000000..dcd0a072ed3
--- /dev/null
+++ b/mlir/lib/Transforms/Vectorization/VectorizerTestPass.cpp
@@ -0,0 +1,116 @@
+//===- VectorizerTestPass.cpp - VectorizerTestPass Pass Impl ----*- C++
+//-*-====================//
+//
+// Copyright 2019 The MLIR Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// =============================================================================
+//
+// This file implements a simple testing pass for vectorization functionality.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Analysis/MLFunctionMatcher.h"
+#include "mlir/Analysis/VectorAnalysis.h"
+#include "mlir/Pass.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/STLExtras.h"
+#include "mlir/Transforms/Passes.h"
+
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
+
+using namespace mlir;
+
+using llvm::outs;
+using llvm::cl::desc;
+using llvm::cl::list;
+using llvm::cl::ZeroOrMore;
+
+static list<int> clTestVectorMultiplicity(
+    "vector-multiplicity", desc("Specify the HW vector size for vectorization"),
+    ZeroOrMore);
+
+#define DEBUG_TYPE "vectorizer-test"
+
+namespace {
+
+struct VectorizerTestPass : public FunctionPass {
+  VectorizerTestPass() : FunctionPass(&VectorizerTestPass::passID) {}
+
+  PassResult runOnMLFunction(MLFunction *f) override;
+  void testVectorMultiplicity(MLFunction *f);
+
+  // Thread-safe RAII contexts local to pass, BumpPtrAllocator freed on exit.
+  MLFunctionMatcherContext MLContext;
+
+  static char passID;
+};
+
+} // end anonymous namespace
+
+char VectorizerTestPass::passID = 0;
+
+void VectorizerTestPass::testVectorMultiplicity(MLFunction *f) {
+  using matcher::Op;
+  SmallVector<int, 8> shape(clTestVectorMultiplicity.begin(),
+                            clTestVectorMultiplicity.end());
+  auto subVectorType = VectorType::get(shape, Type::getF32(f->getContext()));
+  // Only filter statements that operate on a strict super-vector and have one
+  // return. This makes testing easier.
+  auto filter = [subVectorType](const Statement &stmt) {
+    outs() << "\ntest: " << stmt << " ";
+    auto *opStmt = dyn_cast<OperationStmt>(&stmt);
+    if (!opStmt) {
+      return false;
+    }
+    if (!matcher::operatesOnStrictSuperVectors(*opStmt, subVectorType)) {
+      return false;
+    }
+    if (opStmt->getNumResults() != 1) {
+      return false;
+    }
+    return true;
+  };
+  auto pat = Op(filter);
+  auto matches = pat.match(f);
+  for (auto m : matches) {
+    auto *opStmt = cast<OperationStmt>(m.first);
+    // 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 = opStmt->getResult(0)->getType().cast<VectorType>();
+    auto multiplicity = shapeRatio(superVectorType, subVectorType);
+    assert(multiplicity.hasValue() && "Expected multiplicity");
+    outs() << "\nmatched: " << *opStmt << " with multiplicity: ";
+    interleaveComma(MutableArrayRef<unsigned>(*multiplicity), outs());
+  }
+}
+
+PassResult VectorizerTestPass::runOnMLFunction(MLFunction *f) {
+  if (!clTestVectorMultiplicity.empty()) {
+    testVectorMultiplicity(f);
+  }
+
+  return PassResult::Success;
+}
+
+FunctionPass *mlir::createVectorizerTestPass() {
+  return new VectorizerTestPass();
+}
+
+static PassRegistration<VectorizerTestPass>
+    pass("vectorizer-test", "Tests vectorizer standalone functionality.");
+
+#undef DEBUG_TYPE