maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[mlir][linalg] Do not call region builder during vectorization. 

All linalg operations having a region builder shall call it during op creation. Calling it during vectorization is obsolete.

Differential Revision: https://reviews.llvm.org/D99168
This commit is contained in:
Tobias Gysi 2021-03-24 14:22:17 +00:00
parent 92417ebbd1
commit 880822255e
1 changed files with 8 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
View File

@ -288,7 +288,7 @@ static AffineMap getTransferReadMap(LinalgOp linalgOp, unsigned argIndex) {
/// Generic vectorization function that rewrites the body of a `linalgOp` into
/// vector form. Generic vectorization proceeds as follows:
/// 1. The region for the linalg op is created if necessary.
/// 1. Verify the `linalgOp` has one non-empty region.
/// 2. Values defined above the region are mapped to themselves and will be
/// broadcasted on a per-need basis by their consumers.
/// 3. Each region argument is vectorized into a vector.transfer_read (or 0-d
@ -299,36 +299,21 @@ static AffineMap getTransferReadMap(LinalgOp linalgOp, unsigned argIndex) {
LogicalResult vectorizeAsLinalgGeneric(
OpBuilder &builder, LinalgOp linalgOp, SmallVectorImpl<Value> &newResults,
ArrayRef<CustomVectorizationHook> customVectorizationHooks = {}) {
// 1. Certain Linalg ops do not have a region but only a region builder.
// If so, build the region so we can vectorize.
std::unique_ptr<Region> owningRegion;
Region *region;
if (linalgOp->getNumRegions() > 0) {
region = &linalgOp->getRegion(0);
} else {
// RAII avoid remaining in block.
OpBuilder::InsertionGuard g(builder);
owningRegion = std::make_unique<Region>();
region = owningRegion.get();
Block *block = builder.createBlock(region);
auto elementTypes = llvm::to_vector<4>(
llvm::map_range(linalgOp.getShapedOperandTypes(),
[](ShapedType t) { return t.getElementType(); }));
block->addArguments(elementTypes);
linalgOp.getRegionBuilder()(*block, /*captures=*/{});
}
Block *block = &region->front();
// 1. Fail to vectorize if the operation does not have one non-empty region.
if (linalgOp->getNumRegions() != 1 || linalgOp->getRegion(0).empty())
return failure();
auto &block = linalgOp->getRegion(0).front();
BlockAndValueMapping bvm;
// 2. Values defined above the region can only be broadcast for now. Make them
// map to themselves.
llvm::SetVector<Value> valuesSet;
mlir::getUsedValuesDefinedAbove(*region, valuesSet);
mlir::getUsedValuesDefinedAbove(linalgOp->getRegion(0), valuesSet);
bvm.map(valuesSet.getArrayRef(), valuesSet.getArrayRef());
// 3. Turn all BBArgs into vector.transfer_read / load.
SmallVector<AffineMap> indexings;
for (auto bbarg : block->getArguments()) {
for (auto bbarg : block.getArguments()) {
Value vectorArg = linalgOp.getShapedOperand(bbarg.getArgNumber());
AffineMap map;
VectorType vectorType = extractVectorTypeFromShapedValue(vectorArg);
@ -360,7 +345,7 @@ LogicalResult vectorizeAsLinalgGeneric(
hooks.push_back(vectorizeYield);
// 5. Iteratively call `vectorizeOneOp` to each op in the slice.
for (Operation &op : block->getOperations()) {
for (Operation &op : block.getOperations()) {
VectorizationResult result = vectorizeOneOp(builder, &op, bvm, hooks);
if (result.status == VectorizationStatus::Failure) {
LLVM_DEBUG(dbgs() << "\n[" DEBUG_TYPE "]: failed to vectorize: " << op);