diff --git a/mlir/include/mlir/Dialect/VectorOps/VectorUtils.h b/mlir/include/mlir/Dialect/VectorOps/VectorUtils.h --- a/mlir/include/mlir/Dialect/VectorOps/VectorUtils.h +++ b/mlir/include/mlir/Dialect/VectorOps/VectorUtils.h @@ -1,4 +1,4 @@ -//===- Utils.h - VectorOps Utils ----------------------------*- C++ -*-=======// +//===- Utils.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,28 @@ class Value; class VectorType; +/// Given the shape and sizes of a vector, returns the slice strides +/// for each dimension. +SmallVector makeSliceStrides(ArrayRef shape, + ArrayRef 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 delinearize(ArrayRef sliceStrides, + int64_t linearIndex); + +/// Given the sizes of a vector, together with vector-space offsets, returns +/// the element-space offsets for each dimension. +SmallVector toElementSpace(ArrayRef sizes, + ArrayRef vectorOffsets); + +/// Given the shape, sizes, and element-space offsets of a vector, returns +/// the slize sizes for each dimension. +SmallVector makeSliceSizes(ArrayRef shape, + ArrayRef sizes, + ArrayRef 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 diff --git a/mlir/lib/Analysis/CMakeLists.txt b/mlir/lib/Analysis/CMakeLists.txt --- a/mlir/lib/Analysis/CMakeLists.txt +++ b/mlir/lib/Analysis/CMakeLists.txt @@ -13,7 +13,6 @@ TestMemRefDependenceCheck.cpp TestParallelismDetection.cpp Utils.cpp - VectorAnalysis.cpp Verifier.cpp ADDITIONAL_HEADER_DIRS diff --git a/mlir/lib/Dialect/VectorOps/CMakeLists.txt b/mlir/lib/Dialect/VectorOps/CMakeLists.txt --- a/mlir/lib/Dialect/VectorOps/CMakeLists.txt +++ b/mlir/lib/Dialect/VectorOps/CMakeLists.txt @@ -2,6 +2,7 @@ DialectRegistration.cpp VectorOps.cpp VectorTransforms.cpp + VectorUtils.cpp ADDITIONAL_HEADER_DIRS ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/VectorOps diff --git a/mlir/lib/Dialect/VectorOps/VectorOps.cpp b/mlir/lib/Dialect/VectorOps/VectorOps.cpp --- a/mlir/lib/Dialect/VectorOps/VectorOps.cpp +++ b/mlir/lib/Dialect/VectorOps/VectorOps.cpp @@ -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,14 @@ 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 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 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 = makeSliceStrides(shape, sizes); for (int64_t i = 0, e = tupleType.size(); i < e; ++i) { - // De-linearize w.r.t. 'sliceStrides'. - SmallVector 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 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 = toElementSpace(sizes, vectorOffsets); + auto sliceSizes = makeSliceSizes(shape, sizes, elementOffsets); // Create slice VectorType type. auto sliceVectorType = VectorType::get(sliceSizes, vectorType.getElementType()); diff --git a/mlir/lib/Dialect/VectorOps/VectorTransforms.cpp b/mlir/lib/Dialect/VectorOps/VectorTransforms.cpp --- a/mlir/lib/Dialect/VectorOps/VectorTransforms.cpp +++ b/mlir/lib/Dialect/VectorOps/VectorTransforms.cpp @@ -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 @@ return linearIndex; } -/// Given a shape with sizes greater than 0 along all dimensions, returns the -/// delinearized components of linearIndex along shape. -static SmallVector delinearize(int64_t linearIndex, - ArrayRef basis) { - SmallVector 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 @@ ArrayRef 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(sizes.size()) == vectorType.getRank()); + assert(static_cast(strides.size()) == vectorType.getRank()); // Compute shape ratio of 'shape' and 'sizes'. auto shape = vectorType.getShape(); @@ -139,21 +121,13 @@ 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 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 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 = toElementSpace(sizes, vectorOffsets); + auto sliceSizes = makeSliceSizes(shape, sizes, elementOffsets); // Create Vector type and add to 'vectorTypes[i]'. vectorTypes[i] = VectorType::get(sliceSizes, vectorType.getElementType()); } @@ -333,7 +307,7 @@ } // 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 +320,8 @@ // 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 = toElementSpace(targetShape, vectorOffsets); // Get cached slice (or create slice) for each operand at 'offsets'. SmallVector operands; operands.resize(op->getNumOperands()); @@ -360,7 +331,7 @@ 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 +469,19 @@ 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 = toElementSpace(sizes, vectorOffsets); // Compute 'sliceIndices' by adding 'sliceOffsets[i]' to 'indices[i]'. SmallVector 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( it.value().getLoc(), map, ArrayRef(it.value())); @@ -672,13 +640,11 @@ public: using OpRewritePattern::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 sizes; @@ -686,26 +652,15 @@ SmallVector strides; op.getStrides(strides); // all-ones at the moment - // Compute the number of slices in each dimension. - SmallVector 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 tupleValues(tupleSize); + auto sliceStrides = makeSliceStrides(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 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 = toElementSpace(sizes, vectorOffsets); + auto sliceSizes = makeSliceSizes(shape, sizes, elementOffsets); // Insert in tuple. tupleValues[i] = rewriter.create( loc, op.vector(), elementOffsets, sliceSizes, strides); @@ -731,13 +686,11 @@ public: using OpRewritePattern::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 sizes; @@ -745,11 +698,6 @@ SmallVector strides; op.getStrides(strides); // all-ones at the moment - // Compute the number of slices in each dimension. - SmallVector 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(loc, elemType, @@ -757,20 +705,12 @@ Value result = rewriter.create(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 tupleValues(tupleSize); + auto sliceStrides = makeSliceStrides(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 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 = toElementSpace(sizes, vectorOffsets); // Extract from tuple into the result. auto index = rewriter.getI64IntegerAttr(i); auto tupleGet = rewriter.create( diff --git a/mlir/lib/Analysis/VectorAnalysis.cpp b/mlir/lib/Dialect/VectorOps/VectorUtils.cpp rename from mlir/lib/Analysis/VectorAnalysis.cpp rename to mlir/lib/Dialect/VectorOps/VectorUtils.cpp --- a/mlir/lib/Analysis/VectorAnalysis.cpp +++ b/mlir/lib/Dialect/VectorOps/VectorUtils.cpp @@ -1,37 +1,80 @@ -//===- 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 llvm::SetVector; -using namespace mlir; +namespace mlir { -using llvm::SetVector; +SmallVector makeSliceStrides(ArrayRef shape, + ArrayRef sizes) { + int64_t rank = shape.size(); + // Compute the count for each dimension. + SmallVector 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 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; +} -Optional> mlir::shapeRatio(ArrayRef superShape, - ArrayRef subShape) { +SmallVector delinearize(ArrayRef sliceStrides, + int64_t index) { + int64_t rank = sliceStrides.size(); + SmallVector 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 toElementSpace(ArrayRef sizes, + ArrayRef vectorOffsets) { + return functional::zipMap([](int64_t v1, int64_t v2) { return v1 * v2; }, + vectorOffsets, sizes); +} + +SmallVector makeSliceSizes(ArrayRef shape, + ArrayRef sizes, + ArrayRef elementOffsets) { + int64_t rank = shape.size(); + SmallVector sliceSizes(rank); + for (unsigned r = 0; r < rank; ++r) + sliceSizes[r] = std::min(sizes[r], shape[r] - elementOffsets[r]); + return sliceSizes; +} + +Optional> shapeRatio(ArrayRef superShape, + ArrayRef subShape) { if (superShape.size() < subShape.size()) { return Optional>(); } @@ -70,8 +113,8 @@ return SmallVector{result.rbegin(), result.rend()}; } -Optional> mlir::shapeRatio(VectorType superVectorType, - VectorType subVectorType) { +Optional> 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 +200,9 @@ return getParentsOfType(op); } -AffineMap mlir::makePermutationMap( - Operation *op, ArrayRef indices, - const DenseMap &loopToVectorDim) { +AffineMap +makePermutationMap(Operation *op, ArrayRef indices, + const DenseMap &loopToVectorDim) { DenseMap enclosingLoopToVectorDim; auto enclosingLoops = getEnclosingforOps(op); for (auto *forInst : enclosingLoops) { @@ -168,11 +211,11 @@ 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 +273,5 @@ return true; } + +} // namespace mlir