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

[mlir] [VectorOps] consolidate all vector utilities to one header/cc file 

Reviewers: nicolasvasilache, andydavis1, dcaballe

Reviewed By: andydavis1, dcaballe

Subscribers: dcaballe, merge_guards_bot, mgorny, mehdi_amini, rriddle, jpienaar, burmako, shauheen, antiagainst, arpith-jacob, mgester, lucyrfox, liufengdb, llvm-commits

Tags: #llvm

Differential Revision: https://reviews.llvm.org/D73593
This commit is contained in:
aartbik 2020-01-29 15:21:30 -08:00
parent c5fffa4da3
commit 228ea1a46c
6 changed files with 126 additions and 132 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
mlir/include/mlir/Dialect/VectorOps/VectorUtils.h
View File

@ -1,4 +1,4 @@
//===- Utils.h - VectorOps Utils ----------------------------*- C++ -*-=======//
//===- VectorUtils.h - VectorOps Utilities ------------------*- C++ -*-=======//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
@ -15,6 +15,7 @@
namespace mlir {
// Forward declarations.
class AffineApplyOp;
class AffineForOp;
class AffineMap;
@ -25,6 +26,29 @@ class Operation;
class Value;
class VectorType;
/// Given the shape and sizes of a vector, returns the corresponding
/// strides for each dimension.
SmallVector<int64_t, 4> computeStrides(ArrayRef<int64_t> shape,
ArrayRef<int64_t> sizes);
/// Given the slice strides together with a linear index in the dimension
/// space, returns the vector-space offsets in each dimension for a
/// de-linearized index.
SmallVector<int64_t, 4> delinearize(ArrayRef<int64_t> sliceStrides,
int64_t linearIndex);
/// Given the target sizes of a vector, together with vector-space offsets,
/// returns the element-space offsets for each dimension.
SmallVector<int64_t, 4>
computeElementOffsetsFromVectorSliceOffsets(ArrayRef<int64_t> sizes,
ArrayRef<int64_t> vectorOffsets);
/// Given the shape, sizes, and element-space offsets of a vector, returns
/// the slize sizes for each dimension.
SmallVector<int64_t, 4> computeSliceSizes(ArrayRef<int64_t> shape,
ArrayRef<int64_t> sizes,
ArrayRef<int64_t> elementOffsets);
/// Computes and returns the multi-dimensional ratio of `superShape` to
/// `subShape`. This is calculated by performing a traversal from minor to major
/// dimensions (i.e. in reverse shape order). If integral division is not

1
mlir/lib/Analysis/CMakeLists.txt
View File

@ -13,7 +13,6 @@ add_llvm_library(MLIRAnalysis STATIC
TestMemRefDependenceCheck.cpp
TestParallelismDetection.cpp
Utils.cpp
VectorAnalysis.cpp
Verifier.cpp
ADDITIONAL_HEADER_DIRS

1
mlir/lib/Dialect/VectorOps/CMakeLists.txt
View File

@ -2,6 +2,7 @@ add_llvm_library(MLIRVectorOps
DialectRegistration.cpp
VectorOps.cpp
VectorTransforms.cpp
VectorUtils.cpp
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/VectorOps

35
mlir/lib/Dialect/VectorOps/VectorOps.cpp
View File

@ -14,6 +14,7 @@
#include "mlir/Dialect/VectorOps/VectorOps.h"
#include "mlir/Dialect/StandardOps/Ops.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/VectorOps/VectorUtils.h"
#include "mlir/IR/AffineExpr.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/Builders.h"
@ -524,37 +525,15 @@ isValidExtractOrInsertSlicesType(Operation *op, VectorType vectorType,
if (sizes.size() != rank || strides.size() != rank)
return op->emitError("requires sizes and strides of rank ") << rank;
// Compute the number of slices in each dimension.
// TODO(andydavis) Move this into a slice generation helper function.
auto shape = vectorType.getShape();
SmallVector<int64_t, 4> dimSliceCounts(rank);
for (unsigned i = 0; i < rank; ++i)
dimSliceCounts[i] = ceilDiv(shape[i], sizes[i]);
// Compute the strides between slices in each dimension.
SmallVector<int64_t, 4> sliceStrides(rank);
sliceStrides[rank - 1] = 1;
for (int i = rank - 2; i >= 0; --i)
sliceStrides[i] = sliceStrides[i + 1] * dimSliceCounts[i + 1];
// Generate each slice shape based on 'sizes', 'strides' and 'vectorType',
// and verify that the same matches the corresponding tuple element 'i'.
auto shape = vectorType.getShape();
auto sliceStrides = computeStrides(shape, sizes);
for (int64_t i = 0, e = tupleType.size(); i < e; ++i) {
// De-linearize w.r.t. 'sliceStrides'.
SmallVector<int64_t, 4> vectorOffsets(rank);
int64_t linearIndex = i;
for (unsigned j = 0; j < rank; ++j) {
vectorOffsets[j] = linearIndex / sliceStrides[j];
linearIndex %= sliceStrides[j];
}
// Convert from unrolled vector-space offsets to element-space offsets.
auto offsets = mlir::functional::zipMap(
[](int64_t v1, int64_t v2) { return v1 * v2; }, vectorOffsets, sizes);
// Initialize 'sliceSizes' to target 'sizes'
SmallVector<int64_t, 4> sliceSizes(sizes.begin(), sizes.end());
for (unsigned j = 0; j < rank; ++j) {
// Based on 'offsets' and 'shape' clip some dim sizes for partial tiles.
sliceSizes[j] = std::min(sliceSizes[j], shape[j] - offsets[j]);
}
auto vectorOffsets = delinearize(sliceStrides, i);
auto elementOffsets =
computeElementOffsetsFromVectorSliceOffsets(sizes, vectorOffsets);
auto sliceSizes = computeSliceSizes(shape, sizes, elementOffsets);
// Create slice VectorType type.
auto sliceVectorType =
VectorType::get(sliceSizes, vectorType.getElementType());

107
mlir/lib/Dialect/VectorOps/VectorTransforms.cpp
View File

@ -29,7 +29,6 @@
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/Types.h"
#include "mlir/Support/Functional.h"
#include "mlir/Support/MathExtras.h"
#include "mlir/Support/STLExtras.h"
#include "llvm/Support/CommandLine.h"
@ -78,22 +77,6 @@ static int64_t linearize(ArrayRef<int64_t> offsets, ArrayRef<int64_t> basis) {
return linearIndex;
}
/// Given a shape with sizes greater than 0 along all dimensions, returns the
/// delinearized components of linearIndex along shape.
static SmallVector<int64_t, 8> delinearize(int64_t linearIndex,
ArrayRef<int64_t> basis) {
SmallVector<int64_t, 8> res;
res.reserve(basis.size());
for (unsigned idx = 0, e = basis.size(); idx < e; ++idx) {
assert(basis[idx] > 0);
res.push_back(linearIndex / basis[idx]);
linearIndex %= basis[idx];
}
// Sanity check.
assert(linearIndex == 0 && "linear index remainder must be 0");
return res;
}
// Clones `op` into a new operations that takes `operands` and returns
// `resultTypes`.
static Operation *cloneOpWithOperandsAndTypes(PatternRewriter &builder,
@ -128,9 +111,8 @@ static TupleType generateExtractSlicesOpResultType(VectorType vectorType,
ArrayRef<int64_t> strides,
PatternRewriter &builder) {
assert(llvm::all_of(strides, [](int64_t s) { return s == 1; }));
unsigned rank = vectorType.getRank();
assert(sizes.size() == rank);
assert(strides.size() == rank);
assert(static_cast<int64_t>(sizes.size()) == vectorType.getRank());
assert(static_cast<int64_t>(strides.size()) == vectorType.getRank());
// Compute shape ratio of 'shape' and 'sizes'.
auto shape = vectorType.getShape();
@ -139,21 +121,14 @@ static TupleType generateExtractSlicesOpResultType(VectorType vectorType,
auto sliceDimCounts = *maybeDimSliceCounts;
// Compute strides w.r.t number of slices in each dimension.
auto basis = computeStrides(sliceDimCounts);
auto sliceStrides = computeStrides(sliceDimCounts);
int64_t sliceCount = computeMaxLinearIndex(sliceDimCounts);
SmallVector<Type, 4> vectorTypes(sliceCount);
for (unsigned i = 0; i < sliceCount; ++i) {
// De-linearize w.r.t. 'basis'.
auto vectorOffsets = delinearize(i, basis);
// Convert from unrolled vector-space offsets to element-space offsets.
auto offsets = zipMap([](int64_t v1, int64_t v2) { return v1 * v2; },
vectorOffsets, sizes);
// Initialize 'sliceSizes' to target 'sizes'
SmallVector<int64_t, 4> sliceSizes(sizes.begin(), sizes.end());
for (unsigned j = 0; j < rank; ++j) {
// Based on 'offsets' and 'shape' clip some dim sizes for partial tiles.
sliceSizes[j] = std::min(sliceSizes[j], shape[j] - offsets[j]);
}
auto vectorOffsets = delinearize(sliceStrides, i);
auto elementOffsets =
computeElementOffsetsFromVectorSliceOffsets(sizes, vectorOffsets);
auto sliceSizes = computeSliceSizes(shape, sizes, elementOffsets);
// Create Vector type and add to 'vectorTypes[i]'.
vectorTypes[i] = VectorType::get(sliceSizes, vectorType.getElementType());
}
@ -333,7 +308,7 @@ static Value unrollSingleResultStructuredOp(Operation *op,
}
// Compute number of total unrolled instances.
auto numUnrolledInstances = computeMaxLinearIndex(unrollFactors);
auto basis = computeStrides(unrollFactors);
auto sliceStrides = computeStrides(unrollFactors);
auto &resultValueState = unrolledVectorState[resultIndex];
auto unrolledResultType = VectorType::get(resultValueState.unrolledShape,
@ -346,11 +321,9 @@ static Value unrollSingleResultStructuredOp(Operation *op,
// Unroll 'numUnrolledInstances' of 'op', storing results in 'caches'.
for (unsigned i = 0; i < numUnrolledInstances; ++i) {
// De-linearize w.r.t. 'basis'.
auto vectorOffsets = delinearize(i, basis);
// Convert from unrolled vector-space offsets to element-space offsets.
auto offsets = zipMap([](int64_t v1, int64_t v2) { return v1 * v2; },
vectorOffsets, targetShape);
auto vectorOffsets = delinearize(sliceStrides, i);
auto elementOffsets =
computeElementOffsetsFromVectorSliceOffsets(targetShape, vectorOffsets);
// Get cached slice (or create slice) for each operand at 'offsets'.
SmallVector<Value, 3> operands;
operands.resize(op->getNumOperands());
@ -360,7 +333,7 @@ static Value unrollSingleResultStructuredOp(Operation *op,
continue; // Output
auto operand = op->getOperand(operandIndex);
operands[operandIndex] = getOrCreateUnrolledVectorSlice(
op->getLoc(), unrolledVectorState[i], vectorOffsets, offsets,
op->getLoc(), unrolledVectorState[i], vectorOffsets, elementOffsets,
vectors[i].indexMap, operand, caches[i], builder);
}
// Create op on sliced vector arguments.
@ -498,22 +471,20 @@ generateTransferOpSlices(VectorType vectorType, TupleType tupleType,
auto maybeDimSliceCounts = shapeRatio(vectorType.getShape(), sizes);
assert(maybeDimSliceCounts.hasValue());
auto sliceDimCounts = *maybeDimSliceCounts;
auto basis = computeStrides(sliceDimCounts);
auto sliceStrides = computeStrides(sliceDimCounts);
int64_t numSlices = tupleType.size();
unsigned numSliceIndices = indices.size();
auto *ctx = rewriter.getContext();
for (unsigned i = 0; i < numSlices; ++i) {
// De-linearize w.r.t. 'basis'.
auto vectorOffsets = delinearize(i, basis);
// Convert from unrolled vector-space offsets to element-space offsets.
auto offsets = zipMap([](int64_t v1, int64_t v2) { return v1 * v2; },
vectorOffsets, sizes);
auto vectorOffsets = delinearize(sliceStrides, i);
auto elementOffsets =
computeElementOffsetsFromVectorSliceOffsets(sizes, vectorOffsets);
// Compute 'sliceIndices' by adding 'sliceOffsets[i]' to 'indices[i]'.
SmallVector<Value, 4> sliceIndices(numSliceIndices);
for (auto it : llvm::enumerate(indices)) {
auto expr = getAffineDimExpr(0, ctx) +
getAffineConstantExpr(offsets[it.index()], ctx);
getAffineConstantExpr(elementOffsets[it.index()], ctx);
auto map = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr);
sliceIndices[it.index()] = rewriter.create<AffineApplyOp>(
it.value().getLoc(), map, ArrayRef<Value>(it.value()));
@ -672,13 +643,11 @@ class ExtractSlicesOpLowering
public:
using OpRewritePattern<vector::ExtractSlicesOp>::OpRewritePattern;
// TODO(ajcbik): refactor slice utilities out into VectorUtils.h
PatternMatchResult matchAndRewrite(vector::ExtractSlicesOp op,
PatternRewriter &rewriter) const override {
auto loc = op.getLoc();
VectorType vectorType = op.getSourceVectorType();
int64_t rank = vectorType.getRank();
auto shape = vectorType.getShape();
SmallVector<int64_t, 4> sizes;
@ -686,26 +655,16 @@ public:
SmallVector<int64_t, 4> strides;
op.getStrides(strides); // all-ones at the moment
// Compute the number of slices in each dimension.
SmallVector<int64_t, 4> sliceDimCounts(rank);
for (int64_t r = 0; r < rank; ++r)
sliceDimCounts[r] = ceilDiv(shape[r], sizes[r]);
// For each element in the tuple, generate the proper strided slice.
auto basis = computeStrides(sliceDimCounts);
TupleType tupleType = op.getResultTupleType();
int64_t tupleSize = tupleType.size();
SmallVector<Value, 4> tupleValues(tupleSize);
auto sliceStrides = computeStrides(shape, sizes);
for (int64_t i = 0; i < tupleSize; ++i) {
// De-linearize w.r.t. 'basis'.
auto vectorOffsets = delinearize(i, basis);
// Convert from unrolled vector-space offsets to element-space offsets.
auto elementOffsets = mlir::functional::zipMap(
[](int64_t v1, int64_t v2) { return v1 * v2; }, vectorOffsets, sizes);
// Compute the size of each slice.
SmallVector<int64_t, 4> sliceSizes(rank);
for (int64_t r = 0; r < rank; ++r)
sliceSizes[r] = std::min(sizes[r], shape[r] - elementOffsets[r]);
auto vectorOffsets = delinearize(sliceStrides, i);
auto elementOffsets =
computeElementOffsetsFromVectorSliceOffsets(sizes, vectorOffsets);
auto sliceSizes = computeSliceSizes(shape, sizes, elementOffsets);
// Insert in tuple.
tupleValues[i] = rewriter.create<vector::StridedSliceOp>(
loc, op.vector(), elementOffsets, sliceSizes, strides);
@ -731,13 +690,11 @@ class InsertSlicesOpLowering : public OpRewritePattern<vector::InsertSlicesOp> {
public:
using OpRewritePattern<vector::InsertSlicesOp>::OpRewritePattern;
// TODO(ajcbik): refactor slice utilities out into VectorUtils.h
PatternMatchResult matchAndRewrite(vector::InsertSlicesOp op,
PatternRewriter &rewriter) const override {
auto loc = op.getLoc();
VectorType vectorType = op.getResultVectorType();
int64_t rank = vectorType.getRank();
auto shape = vectorType.getShape();
SmallVector<int64_t, 4> sizes;
@ -745,11 +702,6 @@ public:
SmallVector<int64_t, 4> strides;
op.getStrides(strides); // all-ones at the moment
// Compute the number of slices in each dimension.
SmallVector<int64_t, 4> sliceDimCounts(rank);
for (int64_t r = 0; r < rank; ++r)
sliceDimCounts[r] = ceilDiv(shape[r], sizes[r]);
// Prepare result.
auto elemType = vectorType.getElementType();
Value zero = rewriter.create<ConstantOp>(loc, elemType,
@ -757,20 +709,13 @@ public:
Value result = rewriter.create<SplatOp>(loc, vectorType, zero);
// For each element in the tuple, extract the proper strided slice.
auto basis = computeStrides(sliceDimCounts);
TupleType tupleType = op.getSourceTupleType();
int64_t tupleSize = tupleType.size();
SmallVector<Value, 4> tupleValues(tupleSize);
auto sliceStrides = computeStrides(shape, sizes);
for (int64_t i = 0; i < tupleSize; ++i) {
// De-linearize w.r.t. 'basis'.
auto vectorOffsets = delinearize(i, basis);
// Convert from unrolled vector-space offsets to element-space offsets.
auto elementOffsets = mlir::functional::zipMap(
[](int64_t v1, int64_t v2) { return v1 * v2; }, vectorOffsets, sizes);
// Compute the size of each slice.
SmallVector<int64_t, 4> sliceSizes(rank);
for (int64_t r = 0; r < rank; ++r)
sliceSizes[r] = std::min(sizes[r], shape[r] - elementOffsets[r]);
auto vectorOffsets = delinearize(sliceStrides, i);
auto elementOffsets =
computeElementOffsetsFromVectorSliceOffsets(sizes, vectorOffsets);
// Extract from tuple into the result.
auto index = rewriter.getI64IntegerAttr(i);
auto tupleGet = rewriter.create<vector::TupleGetOp>(

88
mlir/lib/Analysis/VectorAnalysis.cpp → mlir/lib/Dialect/VectorOps/VectorUtils.cpp
View File

@ -1,37 +1,81 @@
//===- VectorAnalysis.cpp - Analysis for Vectorization --------------------===//
//===- VectorUtils.cpp - MLIR Utilities for VectorOps ------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// 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 utility methods for working with the VectorOps dialect.
//
//===----------------------------------------------------------------------===//
#include "mlir/Analysis/AffineAnalysis.h"
#include "mlir/Dialect/VectorOps/VectorUtils.h"
#include "mlir/Analysis/LoopAnalysis.h"
#include "mlir/Dialect/AffineOps/AffineOps.h"
#include "mlir/Dialect/StandardOps/Ops.h"
#include "mlir/Dialect/VectorOps/VectorOps.h"
#include "mlir/Dialect/VectorOps/VectorUtils.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/IntegerSet.h"
#include "mlir/IR/Operation.h"
#include "mlir/Support/Functional.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Support/MathExtras.h"
#include "mlir/Support/STLExtras.h"
#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SetVector.h"
///
/// Implements Analysis functions specific to vectors which support
/// the vectorization and vectorization materialization passes.
///
using namespace mlir;
using llvm::SetVector;
Optional<SmallVector<int64_t, 4>> mlir::shapeRatio(ArrayRef<int64_t> superShape,
ArrayRef<int64_t> subShape) {
namespace mlir {
SmallVector<int64_t, 4> computeStrides(ArrayRef<int64_t> shape,
ArrayRef<int64_t> sizes) {
int64_t rank = shape.size();
// Compute the count for each dimension.
SmallVector<int64_t, 4> sliceDimCounts(rank);
for (int64_t r = 0; r < rank; ++r)
sliceDimCounts[r] = ceilDiv(shape[r], sizes[r]);
// Use that to compute the slice stride for each dimension.
SmallVector<int64_t, 4> sliceStrides(rank);
sliceStrides[rank - 1] = 1;
for (int64_t r = rank - 2; r >= 0; --r)
sliceStrides[r] = sliceStrides[r + 1] * sliceDimCounts[r + 1];
return sliceStrides;
}
SmallVector<int64_t, 4> delinearize(ArrayRef<int64_t> sliceStrides,
int64_t index) {
int64_t rank = sliceStrides.size();
SmallVector<int64_t, 4> vectorOffsets(rank);
for (int64_t r = 0; r < rank; ++r) {
assert(sliceStrides[r] > 0);
vectorOffsets[r] = index / sliceStrides[r];
index %= sliceStrides[r];
}
return vectorOffsets;
}
SmallVector<int64_t, 4>
computeElementOffsetsFromVectorSliceOffsets(ArrayRef<int64_t> sizes,
ArrayRef<int64_t> vectorOffsets) {
return functional::zipMap([](int64_t v1, int64_t v2) { return v1 * v2; },
vectorOffsets, sizes);
}
SmallVector<int64_t, 4> computeSliceSizes(ArrayRef<int64_t> shape,
ArrayRef<int64_t> sizes,
ArrayRef<int64_t> elementOffsets) {
int64_t rank = shape.size();
SmallVector<int64_t, 4> sliceSizes(rank);
for (unsigned r = 0; r < rank; ++r)
sliceSizes[r] = std::min(sizes[r], shape[r] - elementOffsets[r]);
return sliceSizes;
}
Optional<SmallVector<int64_t, 4>> shapeRatio(ArrayRef<int64_t> superShape,
ArrayRef<int64_t> subShape) {
if (superShape.size() < subShape.size()) {
return Optional<SmallVector<int64_t, 4>>();
}
@ -70,8 +114,8 @@ Optional<SmallVector<int64_t, 4>> mlir::shapeRatio(ArrayRef<int64_t> superShape,
return SmallVector<int64_t, 4>{result.rbegin(), result.rend()};
}
Optional<SmallVector<int64_t, 4>> mlir::shapeRatio(VectorType superVectorType,
VectorType subVectorType) {
Optional<SmallVector<int64_t, 4>> shapeRatio(VectorType superVectorType,
VectorType subVectorType) {
assert(superVectorType.getElementType() == subVectorType.getElementType() &&
"vector types must be of the same elemental type");
return shapeRatio(superVectorType.getShape(), subVectorType.getShape());
@ -157,9 +201,9 @@ static SetVector<Operation *> getEnclosingforOps(Operation *op) {
return getParentsOfType<AffineForOp>(op);
}
AffineMap mlir::makePermutationMap(
Operation *op, ArrayRef<Value> indices,
const DenseMap<Operation *, unsigned> &loopToVectorDim) {
AffineMap
makePermutationMap(Operation *op, ArrayRef<Value> indices,
const DenseMap<Operation *, unsigned> &loopToVectorDim) {
DenseMap<Operation *, unsigned> enclosingLoopToVectorDim;
auto enclosingLoops = getEnclosingforOps(op);
for (auto *forInst : enclosingLoops) {
@ -168,11 +212,11 @@ AffineMap mlir::makePermutationMap(
enclosingLoopToVectorDim.insert(*it);
}
}
return ::makePermutationMap(indices, enclosingLoopToVectorDim);
return makePermutationMap(indices, enclosingLoopToVectorDim);
}
bool mlir::matcher::operatesOnSuperVectorsOf(Operation &op,
VectorType subVectorType) {
bool matcher::operatesOnSuperVectorsOf(Operation &op,
VectorType subVectorType) {
// First, extract the vector type and distinguish between:
// a. ops that *must* lower a super-vector (i.e. vector.transfer_read,
// vector.transfer_write); and
@ -230,3 +274,5 @@ bool mlir::matcher::operatesOnSuperVectorsOf(Operation &op,
return true;
}
} // namespace mlir