89d9913a20
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
117 lines
3.8 KiB
C++
117 lines
3.8 KiB
C++
//===- 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
|