forked from OSchip/llvm-project
730 lines
29 KiB
C++
730 lines
29 KiB
C++
//===- VectorToSCF.cpp - Conversion from Vector to mix of SCF and Std -----===//
|
|
//
|
|
// Part of the LLVM 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 target-dependent lowering of vector transfer operations.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include <type_traits>
|
|
|
|
#include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
|
|
|
|
#include "../PassDetail.h"
|
|
#include "mlir/Dialect/Affine/EDSC/Intrinsics.h"
|
|
#include "mlir/Dialect/MemRef/EDSC/Intrinsics.h"
|
|
#include "mlir/Dialect/SCF/EDSC/Builders.h"
|
|
#include "mlir/Dialect/SCF/EDSC/Intrinsics.h"
|
|
#include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
|
|
#include "mlir/Dialect/Vector/EDSC/Intrinsics.h"
|
|
#include "mlir/Dialect/Vector/VectorOps.h"
|
|
#include "mlir/Dialect/Vector/VectorUtils.h"
|
|
#include "mlir/IR/AffineExpr.h"
|
|
#include "mlir/IR/AffineMap.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/IR/Matchers.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
#include "mlir/Transforms/Passes.h"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::edsc;
|
|
using namespace mlir::edsc::intrinsics;
|
|
using vector::TransferReadOp;
|
|
using vector::TransferWriteOp;
|
|
|
|
// Return a list of Values that correspond to multiple AffineApplyOp, one for
|
|
// each result of `map`. Each `expr` in `map` is canonicalized and folded
|
|
// greedily according to its operands.
|
|
// TODO: factor out in a common location that both linalg and vector can use.
|
|
static SmallVector<Value, 4>
|
|
applyMapToValues(OpBuilder &b, Location loc, AffineMap map, ValueRange values) {
|
|
SmallVector<Value, 4> res;
|
|
res.reserve(map.getNumResults());
|
|
unsigned numDims = map.getNumDims(), numSym = map.getNumSymbols();
|
|
// For each `expr` in `map`, applies the `expr` to the values extracted from
|
|
// ranges. If the resulting application can be folded into a Value, the
|
|
// folding occurs eagerly. Otherwise, an affine.apply operation is emitted.
|
|
for (auto expr : map.getResults()) {
|
|
AffineMap map = AffineMap::get(numDims, numSym, expr);
|
|
SmallVector<Value, 4> operands(values.begin(), values.end());
|
|
fullyComposeAffineMapAndOperands(&map, &operands);
|
|
canonicalizeMapAndOperands(&map, &operands);
|
|
res.push_back(b.createOrFold<AffineApplyOp>(loc, map, operands));
|
|
}
|
|
return res;
|
|
}
|
|
|
|
namespace {
|
|
/// Helper class captures the common information needed to lower N>1-D vector
|
|
/// transfer operations (read and write).
|
|
/// On construction, this class opens an edsc::ScopedContext for simpler IR
|
|
/// manipulation.
|
|
/// In pseudo-IR, for an n-D vector_transfer_read such as:
|
|
///
|
|
/// ```
|
|
/// vector_transfer_read(%m, %offsets, identity_map, %fill) :
|
|
/// memref<(leading_dims) x (major_dims) x (minor_dims) x type>,
|
|
/// vector<(major_dims) x (minor_dims) x type>
|
|
/// ```
|
|
///
|
|
/// where rank(minor_dims) is the lower-level vector rank (e.g. 1 for LLVM or
|
|
/// higher).
|
|
///
|
|
/// This is the entry point to emitting pseudo-IR resembling:
|
|
///
|
|
/// ```
|
|
/// %tmp = alloc(): memref<(major_dims) x vector<minor_dim x type>>
|
|
/// for (%ivs_major, {0}, {vector_shape}, {1}) { // (N-1)-D loop nest
|
|
/// if (any_of(%ivs_major + %offsets, <, major_dims)) {
|
|
/// %v = vector_transfer_read(
|
|
/// {%offsets_leading, %ivs_major + %offsets_major, %offsets_minor},
|
|
/// %ivs_minor):
|
|
/// memref<(leading_dims) x (major_dims) x (minor_dims) x type>,
|
|
/// vector<(minor_dims) x type>;
|
|
/// store(%v, %tmp);
|
|
/// } else {
|
|
/// %v = splat(vector<(minor_dims) x type>, %fill)
|
|
/// store(%v, %tmp, %ivs_major);
|
|
/// }
|
|
/// }
|
|
/// %res = load(%tmp, %0): memref<(major_dims) x vector<minor_dim x type>>):
|
|
// vector<(major_dims) x (minor_dims) x type>
|
|
/// ```
|
|
///
|
|
template <typename ConcreteOp>
|
|
class NDTransferOpHelper {
|
|
public:
|
|
NDTransferOpHelper(PatternRewriter &rewriter, ConcreteOp xferOp,
|
|
const VectorTransferToSCFOptions &options)
|
|
: rewriter(rewriter), options(options), loc(xferOp.getLoc()),
|
|
scope(std::make_unique<ScopedContext>(rewriter, loc)), xferOp(xferOp),
|
|
op(xferOp.getOperation()) {
|
|
vectorType = xferOp.getVectorType();
|
|
// TODO: when we go to k > 1-D vectors adapt minorRank.
|
|
minorRank = 1;
|
|
majorRank = vectorType.getRank() - minorRank;
|
|
leadingRank = xferOp.getLeadingShapedRank();
|
|
majorVectorType =
|
|
VectorType::get(vectorType.getShape().take_front(majorRank),
|
|
vectorType.getElementType());
|
|
minorVectorType =
|
|
VectorType::get(vectorType.getShape().take_back(minorRank),
|
|
vectorType.getElementType());
|
|
/// Memref of minor vector type is used for individual transfers.
|
|
memRefMinorVectorType = MemRefType::get(
|
|
majorVectorType.getShape(), minorVectorType, {},
|
|
xferOp.getShapedType().template cast<MemRefType>().getMemorySpace());
|
|
}
|
|
|
|
LogicalResult doReplace();
|
|
|
|
private:
|
|
/// Creates the loop nest on the "major" dimensions and calls the
|
|
/// `loopBodyBuilder` lambda in the context of the loop nest.
|
|
void
|
|
emitLoops(llvm::function_ref<void(ValueRange, ValueRange, ValueRange,
|
|
ValueRange, const MemRefBoundsCapture &)>
|
|
loopBodyBuilder);
|
|
|
|
/// Common state to lower vector transfer ops.
|
|
PatternRewriter &rewriter;
|
|
const VectorTransferToSCFOptions &options;
|
|
Location loc;
|
|
std::unique_ptr<ScopedContext> scope;
|
|
ConcreteOp xferOp;
|
|
Operation *op;
|
|
// A vector transfer copies data between:
|
|
// - memref<(leading_dims) x (major_dims) x (minor_dims) x type>
|
|
// - vector<(major_dims) x (minor_dims) x type>
|
|
unsigned minorRank; // for now always 1
|
|
unsigned majorRank; // vector rank - minorRank
|
|
unsigned leadingRank; // memref rank - vector rank
|
|
VectorType vectorType; // vector<(major_dims) x (minor_dims) x type>
|
|
VectorType majorVectorType; // vector<(major_dims) x type>
|
|
VectorType minorVectorType; // vector<(minor_dims) x type>
|
|
MemRefType memRefMinorVectorType; // memref<vector<(minor_dims) x type>>
|
|
};
|
|
|
|
template <typename ConcreteOp>
|
|
void NDTransferOpHelper<ConcreteOp>::emitLoops(
|
|
llvm::function_ref<void(ValueRange, ValueRange, ValueRange, ValueRange,
|
|
const MemRefBoundsCapture &)>
|
|
loopBodyBuilder) {
|
|
/// Loop nest operates on the major dimensions
|
|
MemRefBoundsCapture memrefBoundsCapture(xferOp.source());
|
|
|
|
if (options.unroll) {
|
|
auto shape = majorVectorType.getShape();
|
|
auto strides = computeStrides(shape);
|
|
unsigned numUnrolledInstances = computeMaxLinearIndex(shape);
|
|
ValueRange indices(xferOp.indices());
|
|
for (unsigned idx = 0; idx < numUnrolledInstances; ++idx) {
|
|
SmallVector<int64_t, 4> offsets = delinearize(strides, idx);
|
|
SmallVector<Value, 4> offsetValues =
|
|
llvm::to_vector<4>(llvm::map_range(offsets, [](int64_t off) -> Value {
|
|
return std_constant_index(off);
|
|
}));
|
|
loopBodyBuilder(offsetValues, indices.take_front(leadingRank),
|
|
indices.drop_front(leadingRank).take_front(majorRank),
|
|
indices.take_back(minorRank), memrefBoundsCapture);
|
|
}
|
|
} else {
|
|
VectorBoundsCapture vectorBoundsCapture(majorVectorType);
|
|
auto majorLbs = vectorBoundsCapture.getLbs();
|
|
auto majorUbs = vectorBoundsCapture.getUbs();
|
|
auto majorSteps = vectorBoundsCapture.getSteps();
|
|
affineLoopNestBuilder(
|
|
majorLbs, majorUbs, majorSteps, [&](ValueRange majorIvs) {
|
|
ValueRange indices(xferOp.indices());
|
|
loopBodyBuilder(majorIvs, indices.take_front(leadingRank),
|
|
indices.drop_front(leadingRank).take_front(majorRank),
|
|
indices.take_back(minorRank), memrefBoundsCapture);
|
|
});
|
|
}
|
|
}
|
|
|
|
static Optional<int64_t> extractConstantIndex(Value v) {
|
|
if (auto cstOp = v.getDefiningOp<ConstantIndexOp>())
|
|
return cstOp.getValue();
|
|
if (auto affineApplyOp = v.getDefiningOp<AffineApplyOp>())
|
|
if (affineApplyOp.getAffineMap().isSingleConstant())
|
|
return affineApplyOp.getAffineMap().getSingleConstantResult();
|
|
return None;
|
|
}
|
|
|
|
// Missing foldings of scf.if make it necessary to perform poor man's folding
|
|
// eagerly, especially in the case of unrolling. In the future, this should go
|
|
// away once scf.if folds properly.
|
|
static Value onTheFlyFoldSLT(Value v, Value ub) {
|
|
using namespace mlir::edsc::op;
|
|
auto maybeCstV = extractConstantIndex(v);
|
|
auto maybeCstUb = extractConstantIndex(ub);
|
|
if (maybeCstV && maybeCstUb && *maybeCstV < *maybeCstUb)
|
|
return Value();
|
|
return slt(v, ub);
|
|
}
|
|
|
|
/// 1. Compute the indexings `majorIvs + majorOffsets` and save them in
|
|
/// `majorIvsPlusOffsets`.
|
|
/// 2. Return a value of i1 that determines whether the first
|
|
/// `majorIvs.rank()`
|
|
/// dimensions `majorIvs + majorOffsets` are all within `memrefBounds`.
|
|
static Value
|
|
emitInBoundsCondition(PatternRewriter &rewriter,
|
|
VectorTransferOpInterface xferOp, unsigned leadingRank,
|
|
ValueRange majorIvs, ValueRange majorOffsets,
|
|
const MemRefBoundsCapture &memrefBounds,
|
|
SmallVectorImpl<Value> &majorIvsPlusOffsets) {
|
|
Value inBoundsCondition;
|
|
majorIvsPlusOffsets.reserve(majorIvs.size());
|
|
unsigned idx = 0;
|
|
SmallVector<Value, 4> bounds =
|
|
applyMapToValues(rewriter, xferOp.getLoc(), xferOp.permutation_map(),
|
|
memrefBounds.getUbs());
|
|
for (auto it : llvm::zip(majorIvs, majorOffsets, bounds)) {
|
|
Value iv = std::get<0>(it), off = std::get<1>(it), ub = std::get<2>(it);
|
|
using namespace mlir::edsc::op;
|
|
majorIvsPlusOffsets.push_back(iv + off);
|
|
if (!xferOp.isDimInBounds(leadingRank + idx)) {
|
|
Value inBoundsCond = onTheFlyFoldSLT(majorIvsPlusOffsets.back(), ub);
|
|
if (inBoundsCond)
|
|
inBoundsCondition = (inBoundsCondition)
|
|
? (inBoundsCondition && inBoundsCond)
|
|
: inBoundsCond;
|
|
}
|
|
++idx;
|
|
}
|
|
return inBoundsCondition;
|
|
}
|
|
|
|
// TODO: Parallelism and threadlocal considerations.
|
|
static Value setAllocAtFunctionEntry(MemRefType memRefMinorVectorType,
|
|
Operation *op) {
|
|
auto &b = ScopedContext::getBuilderRef();
|
|
OpBuilder::InsertionGuard guard(b);
|
|
Operation *scope =
|
|
op->getParentWithTrait<OpTrait::AutomaticAllocationScope>();
|
|
assert(scope && "Expected op to be inside automatic allocation scope");
|
|
b.setInsertionPointToStart(&scope->getRegion(0).front());
|
|
Value res = memref_alloca(memRefMinorVectorType);
|
|
return res;
|
|
}
|
|
|
|
template <>
|
|
LogicalResult NDTransferOpHelper<TransferReadOp>::doReplace() {
|
|
Value alloc, result;
|
|
if (options.unroll)
|
|
result = std_splat(vectorType, xferOp.padding());
|
|
else
|
|
alloc = setAllocAtFunctionEntry(memRefMinorVectorType, op);
|
|
|
|
emitLoops([&](ValueRange majorIvs, ValueRange leadingOffsets,
|
|
ValueRange majorOffsets, ValueRange minorOffsets,
|
|
const MemRefBoundsCapture &memrefBounds) {
|
|
/// Lambda to load 1-D vector in the current loop ivs + offset context.
|
|
auto load1DVector = [&](ValueRange majorIvsPlusOffsets) -> Value {
|
|
SmallVector<Value, 8> indexing;
|
|
indexing.reserve(leadingRank + majorRank + minorRank);
|
|
indexing.append(leadingOffsets.begin(), leadingOffsets.end());
|
|
indexing.append(majorIvsPlusOffsets.begin(), majorIvsPlusOffsets.end());
|
|
indexing.append(minorOffsets.begin(), minorOffsets.end());
|
|
Value memref = xferOp.source();
|
|
auto map =
|
|
getTransferMinorIdentityMap(xferOp.getShapedType(), minorVectorType);
|
|
ArrayAttr inBounds;
|
|
if (xferOp.isDimInBounds(xferOp.getVectorType().getRank() - 1)) {
|
|
OpBuilder &b = ScopedContext::getBuilderRef();
|
|
inBounds = b.getBoolArrayAttr({true});
|
|
}
|
|
return vector_transfer_read(minorVectorType, memref, indexing,
|
|
AffineMapAttr::get(map), xferOp.padding(),
|
|
inBounds);
|
|
};
|
|
|
|
// 1. Compute the inBoundsCondition in the current loops ivs + offset
|
|
// context.
|
|
SmallVector<Value, 4> majorIvsPlusOffsets;
|
|
Value inBoundsCondition = emitInBoundsCondition(
|
|
rewriter, cast<VectorTransferOpInterface>(xferOp.getOperation()),
|
|
leadingRank, majorIvs, majorOffsets, memrefBounds, majorIvsPlusOffsets);
|
|
|
|
if (inBoundsCondition) {
|
|
// 2. If the condition is not null, we need an IfOp, which may yield
|
|
// if `options.unroll` is true.
|
|
SmallVector<Type, 1> resultType;
|
|
if (options.unroll)
|
|
resultType.push_back(vectorType);
|
|
|
|
// 3. If in-bounds, progressively lower to a 1-D transfer read, otherwise
|
|
// splat a 1-D vector.
|
|
ValueRange ifResults = conditionBuilder(
|
|
resultType, inBoundsCondition,
|
|
[&]() -> scf::ValueVector {
|
|
Value vector = load1DVector(majorIvsPlusOffsets);
|
|
// 3.a. If `options.unroll` is true, insert the 1-D vector in the
|
|
// aggregate. We must yield and merge with the `else` branch.
|
|
if (options.unroll) {
|
|
vector = vector_insert(vector, result, majorIvs);
|
|
return {vector};
|
|
}
|
|
// 3.b. Otherwise, just go through the temporary `alloc`.
|
|
memref_store(vector, alloc, majorIvs);
|
|
return {};
|
|
},
|
|
[&]() -> scf::ValueVector {
|
|
Value vector = std_splat(minorVectorType, xferOp.padding());
|
|
// 3.c. If `options.unroll` is true, insert the 1-D vector in the
|
|
// aggregate. We must yield and merge with the `then` branch.
|
|
if (options.unroll) {
|
|
vector = vector_insert(vector, result, majorIvs);
|
|
return {vector};
|
|
}
|
|
// 3.d. Otherwise, just go through the temporary `alloc`.
|
|
memref_store(vector, alloc, majorIvs);
|
|
return {};
|
|
});
|
|
|
|
if (!resultType.empty())
|
|
result = *ifResults.begin();
|
|
} else {
|
|
// 4. Guaranteed in-bounds, progressively lower to a 1-D transfer read.
|
|
Value loaded1D = load1DVector(majorIvsPlusOffsets);
|
|
// 5.a. If `options.unroll` is true, insert the 1-D vector in the
|
|
// aggregate.
|
|
if (options.unroll)
|
|
result = vector_insert(loaded1D, result, majorIvs);
|
|
// 5.b. Otherwise, just go through the temporary `alloc`.
|
|
else
|
|
memref_store(loaded1D, alloc, majorIvs);
|
|
}
|
|
});
|
|
|
|
assert((!options.unroll ^ (bool)result) &&
|
|
"Expected resulting Value iff unroll");
|
|
if (!result)
|
|
result =
|
|
memref_load(vector_type_cast(MemRefType::get({}, vectorType), alloc));
|
|
rewriter.replaceOp(op, result);
|
|
|
|
return success();
|
|
}
|
|
|
|
template <>
|
|
LogicalResult NDTransferOpHelper<TransferWriteOp>::doReplace() {
|
|
Value alloc;
|
|
if (!options.unroll) {
|
|
alloc = setAllocAtFunctionEntry(memRefMinorVectorType, op);
|
|
memref_store(xferOp.vector(),
|
|
vector_type_cast(MemRefType::get({}, vectorType), alloc));
|
|
}
|
|
|
|
emitLoops([&](ValueRange majorIvs, ValueRange leadingOffsets,
|
|
ValueRange majorOffsets, ValueRange minorOffsets,
|
|
const MemRefBoundsCapture &memrefBounds) {
|
|
// Lower to 1-D vector_transfer_write and let recursion handle it.
|
|
auto emitTransferWrite = [&](ValueRange majorIvsPlusOffsets) {
|
|
SmallVector<Value, 8> indexing;
|
|
indexing.reserve(leadingRank + majorRank + minorRank);
|
|
indexing.append(leadingOffsets.begin(), leadingOffsets.end());
|
|
indexing.append(majorIvsPlusOffsets.begin(), majorIvsPlusOffsets.end());
|
|
indexing.append(minorOffsets.begin(), minorOffsets.end());
|
|
Value result;
|
|
// If `options.unroll` is true, extract the 1-D vector from the
|
|
// aggregate.
|
|
if (options.unroll)
|
|
result = vector_extract(xferOp.vector(), majorIvs);
|
|
else
|
|
result = memref_load(alloc, majorIvs);
|
|
auto map =
|
|
getTransferMinorIdentityMap(xferOp.getShapedType(), minorVectorType);
|
|
ArrayAttr inBounds;
|
|
if (xferOp.isDimInBounds(xferOp.getVectorType().getRank() - 1)) {
|
|
OpBuilder &b = ScopedContext::getBuilderRef();
|
|
inBounds = b.getBoolArrayAttr({true});
|
|
}
|
|
vector_transfer_write(result, xferOp.source(), indexing,
|
|
AffineMapAttr::get(map), inBounds);
|
|
};
|
|
|
|
// 1. Compute the inBoundsCondition in the current loops ivs + offset
|
|
// context.
|
|
SmallVector<Value, 4> majorIvsPlusOffsets;
|
|
Value inBoundsCondition = emitInBoundsCondition(
|
|
rewriter, cast<VectorTransferOpInterface>(xferOp.getOperation()),
|
|
leadingRank, majorIvs, majorOffsets, memrefBounds, majorIvsPlusOffsets);
|
|
|
|
if (inBoundsCondition) {
|
|
// 2.a. If the condition is not null, we need an IfOp, to write
|
|
// conditionally. Progressively lower to a 1-D transfer write.
|
|
conditionBuilder(inBoundsCondition,
|
|
[&] { emitTransferWrite(majorIvsPlusOffsets); });
|
|
} else {
|
|
// 2.b. Guaranteed in-bounds. Progressively lower to a 1-D transfer write.
|
|
emitTransferWrite(majorIvsPlusOffsets);
|
|
}
|
|
});
|
|
|
|
rewriter.eraseOp(op);
|
|
|
|
return success();
|
|
}
|
|
|
|
} // namespace
|
|
|
|
/// Analyzes the `transfer` to find an access dimension along the fastest remote
|
|
/// MemRef dimension. If such a dimension with coalescing properties is found,
|
|
/// `pivs` and `vectorBoundsCapture` are swapped so that the invocation of
|
|
/// LoopNestBuilder captures it in the innermost loop.
|
|
template <typename TransferOpTy>
|
|
static int computeCoalescedIndex(TransferOpTy transfer) {
|
|
// rank of the remote memory access, coalescing behavior occurs on the
|
|
// innermost memory dimension.
|
|
auto remoteRank = transfer.getShapedType().getRank();
|
|
// Iterate over the results expressions of the permutation map to determine
|
|
// the loop order for creating pointwise copies between remote and local
|
|
// memories.
|
|
int coalescedIdx = -1;
|
|
auto exprs = transfer.permutation_map().getResults();
|
|
for (auto en : llvm::enumerate(exprs)) {
|
|
auto dim = en.value().template dyn_cast<AffineDimExpr>();
|
|
if (!dim) {
|
|
continue;
|
|
}
|
|
auto memRefDim = dim.getPosition();
|
|
if (memRefDim == remoteRank - 1) {
|
|
// memRefDim has coalescing properties, it should be swapped in the last
|
|
// position.
|
|
assert(coalescedIdx == -1 && "Unexpected > 1 coalesced indices");
|
|
coalescedIdx = en.index();
|
|
}
|
|
}
|
|
return coalescedIdx;
|
|
}
|
|
|
|
template <typename TransferOpTy>
|
|
VectorTransferRewriter<TransferOpTy>::VectorTransferRewriter(
|
|
VectorTransferToSCFOptions options, MLIRContext *context)
|
|
: RewritePattern(TransferOpTy::getOperationName(), 1, context),
|
|
options(options) {}
|
|
|
|
/// Used for staging the transfer in a local buffer.
|
|
template <typename TransferOpTy>
|
|
MemRefType VectorTransferRewriter<TransferOpTy>::tmpMemRefType(
|
|
TransferOpTy transfer) const {
|
|
auto vectorType = transfer.getVectorType();
|
|
return MemRefType::get(vectorType.getShape().drop_back(),
|
|
VectorType::get(vectorType.getShape().take_back(),
|
|
vectorType.getElementType()),
|
|
{}, 0);
|
|
}
|
|
|
|
static void emitWithBoundsChecks(
|
|
PatternRewriter &rewriter, VectorTransferOpInterface transfer,
|
|
ValueRange ivs, const MemRefBoundsCapture &memRefBoundsCapture,
|
|
function_ref<void(ArrayRef<Value>)> inBoundsFun,
|
|
function_ref<void(ArrayRef<Value>)> outOfBoundsFun = nullptr) {
|
|
// Permute the incoming indices according to the permutation map.
|
|
SmallVector<Value, 4> indices =
|
|
applyMapToValues(rewriter, transfer.getLoc(), transfer.permutation_map(),
|
|
transfer.indices());
|
|
|
|
// Generate a bounds check if necessary.
|
|
SmallVector<Value, 4> majorIvsPlusOffsets;
|
|
Value inBoundsCondition =
|
|
emitInBoundsCondition(rewriter, transfer, 0, ivs, indices,
|
|
memRefBoundsCapture, majorIvsPlusOffsets);
|
|
|
|
// Apply the permutation map to the ivs. The permutation map may not use all
|
|
// the inputs.
|
|
SmallVector<Value, 4> scalarAccessExprs(transfer.indices().size());
|
|
for (unsigned memRefDim = 0; memRefDim < transfer.indices().size();
|
|
++memRefDim) {
|
|
// Linear search on a small number of entries.
|
|
int loopIndex = -1;
|
|
auto exprs = transfer.permutation_map().getResults();
|
|
for (auto en : llvm::enumerate(exprs)) {
|
|
auto expr = en.value();
|
|
auto dim = expr.dyn_cast<AffineDimExpr>();
|
|
// Sanity check.
|
|
assert((dim || expr.cast<AffineConstantExpr>().getValue() == 0) &&
|
|
"Expected dim or 0 in permutationMap");
|
|
if (dim && memRefDim == dim.getPosition()) {
|
|
loopIndex = en.index();
|
|
break;
|
|
}
|
|
}
|
|
|
|
using namespace edsc::op;
|
|
auto i = transfer.indices()[memRefDim];
|
|
scalarAccessExprs[memRefDim] = loopIndex < 0 ? i : i + ivs[loopIndex];
|
|
}
|
|
|
|
if (inBoundsCondition)
|
|
conditionBuilder(
|
|
/* scf.if */ inBoundsCondition, // {
|
|
[&] { inBoundsFun(scalarAccessExprs); },
|
|
// } else {
|
|
outOfBoundsFun ? [&] { outOfBoundsFun(scalarAccessExprs); }
|
|
: function_ref<void()>()
|
|
// }
|
|
);
|
|
else
|
|
inBoundsFun(scalarAccessExprs);
|
|
}
|
|
|
|
namespace mlir {
|
|
|
|
/// Lowers TransferReadOp into a combination of:
|
|
/// 1. local memory allocation;
|
|
/// 2. perfect loop nest over:
|
|
/// a. scalar load from local buffers (viewed as a scalar memref);
|
|
/// a. scalar store to original memref (with padding).
|
|
/// 3. vector_load from local buffer (viewed as a memref<1 x vector>);
|
|
/// 4. local memory deallocation.
|
|
///
|
|
/// Lowers the data transfer part of a TransferReadOp while ensuring no
|
|
/// out-of-bounds accesses are possible. Out-of-bounds behavior is handled by
|
|
/// padding.
|
|
|
|
/// Performs the rewrite.
|
|
template <>
|
|
LogicalResult VectorTransferRewriter<TransferReadOp>::matchAndRewrite(
|
|
Operation *op, PatternRewriter &rewriter) const {
|
|
using namespace mlir::edsc::op;
|
|
|
|
TransferReadOp transfer = cast<TransferReadOp>(op);
|
|
if (transfer.mask())
|
|
return failure();
|
|
auto memRefType = transfer.getShapedType().dyn_cast<MemRefType>();
|
|
if (!memRefType)
|
|
return failure();
|
|
// Fall back to a loop if the fastest varying stride is not 1 or it is
|
|
// permuted.
|
|
int64_t offset;
|
|
SmallVector<int64_t, 4> strides;
|
|
auto successStrides = getStridesAndOffset(memRefType, strides, offset);
|
|
if (succeeded(successStrides) && strides.back() == 1 &&
|
|
transfer.permutation_map().isMinorIdentity()) {
|
|
// If > 1D, emit a bunch of loops around 1-D vector transfers.
|
|
if (transfer.getVectorType().getRank() > 1)
|
|
return NDTransferOpHelper<TransferReadOp>(rewriter, transfer, options)
|
|
.doReplace();
|
|
// If 1-D this is now handled by the target-specific lowering.
|
|
if (transfer.getVectorType().getRank() == 1)
|
|
return failure();
|
|
}
|
|
|
|
// Conservative lowering to scalar load / stores.
|
|
// 1. Setup all the captures.
|
|
ScopedContext scope(rewriter, transfer.getLoc());
|
|
MemRefIndexedValue remote(transfer.source());
|
|
MemRefBoundsCapture memRefBoundsCapture(transfer.source());
|
|
VectorBoundsCapture vectorBoundsCapture(transfer.vector());
|
|
int coalescedIdx = computeCoalescedIndex(transfer);
|
|
// Swap the vectorBoundsCapture which will reorder loop bounds.
|
|
if (coalescedIdx >= 0)
|
|
vectorBoundsCapture.swapRanges(vectorBoundsCapture.rank() - 1,
|
|
coalescedIdx);
|
|
|
|
auto lbs = vectorBoundsCapture.getLbs();
|
|
auto ubs = vectorBoundsCapture.getUbs();
|
|
SmallVector<Value, 8> steps;
|
|
steps.reserve(vectorBoundsCapture.getSteps().size());
|
|
for (auto step : vectorBoundsCapture.getSteps())
|
|
steps.push_back(std_constant_index(step));
|
|
|
|
// 2. Emit alloc-copy-load-dealloc.
|
|
MLIRContext *ctx = op->getContext();
|
|
Value tmp = setAllocAtFunctionEntry(tmpMemRefType(transfer), transfer);
|
|
MemRefIndexedValue local(tmp);
|
|
loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) {
|
|
auto ivsStorage = llvm::to_vector<8>(loopIvs);
|
|
// Swap the ivs which will reorder memory accesses.
|
|
if (coalescedIdx >= 0)
|
|
std::swap(ivsStorage.back(), ivsStorage[coalescedIdx]);
|
|
|
|
ArrayRef<Value> ivs(ivsStorage);
|
|
Value pos = std_index_cast(IntegerType::get(ctx, 32), ivs.back());
|
|
Value inVector = local(ivs.drop_back());
|
|
auto loadValue = [&](ArrayRef<Value> indices) {
|
|
Value vector = vector_insert_element(remote(indices), inVector, pos);
|
|
local(ivs.drop_back()) = vector;
|
|
};
|
|
auto loadPadding = [&](ArrayRef<Value>) {
|
|
Value vector = vector_insert_element(transfer.padding(), inVector, pos);
|
|
local(ivs.drop_back()) = vector;
|
|
};
|
|
emitWithBoundsChecks(
|
|
rewriter, cast<VectorTransferOpInterface>(transfer.getOperation()), ivs,
|
|
memRefBoundsCapture, loadValue, loadPadding);
|
|
});
|
|
Value vectorValue = memref_load(vector_type_cast(tmp));
|
|
|
|
// 3. Propagate.
|
|
rewriter.replaceOp(op, vectorValue);
|
|
return success();
|
|
}
|
|
|
|
/// Lowers TransferWriteOp into a combination of:
|
|
/// 1. local memory allocation;
|
|
/// 2. vector_store to local buffer (viewed as a memref<1 x vector>);
|
|
/// 3. perfect loop nest over:
|
|
/// a. scalar load from local buffers (viewed as a scalar memref);
|
|
/// a. scalar store to original memref (if in bounds).
|
|
/// 4. local memory deallocation.
|
|
///
|
|
/// More specifically, lowers the data transfer part while ensuring no
|
|
/// out-of-bounds accesses are possible.
|
|
template <>
|
|
LogicalResult VectorTransferRewriter<TransferWriteOp>::matchAndRewrite(
|
|
Operation *op, PatternRewriter &rewriter) const {
|
|
using namespace edsc::op;
|
|
|
|
TransferWriteOp transfer = cast<TransferWriteOp>(op);
|
|
if (transfer.mask())
|
|
return failure();
|
|
auto memRefType = transfer.getShapedType().template dyn_cast<MemRefType>();
|
|
if (!memRefType)
|
|
return failure();
|
|
|
|
// Fall back to a loop if the fastest varying stride is not 1 or it is
|
|
// permuted.
|
|
int64_t offset;
|
|
SmallVector<int64_t, 4> strides;
|
|
auto successStrides = getStridesAndOffset(memRefType, strides, offset);
|
|
if (succeeded(successStrides) && strides.back() == 1 &&
|
|
transfer.permutation_map().isMinorIdentity()) {
|
|
// If > 1D, emit a bunch of loops around 1-D vector transfers.
|
|
if (transfer.getVectorType().getRank() > 1)
|
|
return NDTransferOpHelper<TransferWriteOp>(rewriter, transfer, options)
|
|
.doReplace();
|
|
// If 1-D this is now handled by the target-specific lowering.
|
|
if (transfer.getVectorType().getRank() == 1)
|
|
return failure();
|
|
}
|
|
|
|
// 1. Setup all the captures.
|
|
ScopedContext scope(rewriter, transfer.getLoc());
|
|
MemRefIndexedValue remote(transfer.source());
|
|
MemRefBoundsCapture memRefBoundsCapture(transfer.source());
|
|
Value vectorValue(transfer.vector());
|
|
VectorBoundsCapture vectorBoundsCapture(transfer.vector());
|
|
int coalescedIdx = computeCoalescedIndex(transfer);
|
|
// Swap the vectorBoundsCapture which will reorder loop bounds.
|
|
if (coalescedIdx >= 0)
|
|
vectorBoundsCapture.swapRanges(vectorBoundsCapture.rank() - 1,
|
|
coalescedIdx);
|
|
|
|
auto lbs = vectorBoundsCapture.getLbs();
|
|
auto ubs = vectorBoundsCapture.getUbs();
|
|
SmallVector<Value, 8> steps;
|
|
steps.reserve(vectorBoundsCapture.getSteps().size());
|
|
for (auto step : vectorBoundsCapture.getSteps())
|
|
steps.push_back(std_constant_index(step));
|
|
|
|
// 2. Emit alloc-store-copy-dealloc.
|
|
Value tmp = setAllocAtFunctionEntry(tmpMemRefType(transfer), transfer);
|
|
MemRefIndexedValue local(tmp);
|
|
Value vec = vector_type_cast(tmp);
|
|
memref_store(vectorValue, vec);
|
|
loopNestBuilder(lbs, ubs, steps, [&](ValueRange loopIvs) {
|
|
auto ivsStorage = llvm::to_vector<8>(loopIvs);
|
|
// Swap the ivsStorage which will reorder memory accesses.
|
|
if (coalescedIdx >= 0)
|
|
std::swap(ivsStorage.back(), ivsStorage[coalescedIdx]);
|
|
|
|
ArrayRef<Value> ivs(ivsStorage);
|
|
Value pos =
|
|
std_index_cast(IntegerType::get(op->getContext(), 32), ivs.back());
|
|
auto storeValue = [&](ArrayRef<Value> indices) {
|
|
Value scalar = vector_extract_element(local(ivs.drop_back()), pos);
|
|
remote(indices) = scalar;
|
|
};
|
|
emitWithBoundsChecks(
|
|
rewriter, cast<VectorTransferOpInterface>(transfer.getOperation()), ivs,
|
|
memRefBoundsCapture, storeValue);
|
|
});
|
|
|
|
// 3. Erase.
|
|
rewriter.eraseOp(op);
|
|
return success();
|
|
}
|
|
|
|
void populateVectorToSCFConversionPatterns(
|
|
RewritePatternSet &patterns, const VectorTransferToSCFOptions &options) {
|
|
patterns.add<VectorTransferRewriter<vector::TransferReadOp>,
|
|
VectorTransferRewriter<vector::TransferWriteOp>>(
|
|
options, patterns.getContext());
|
|
}
|
|
|
|
} // namespace mlir
|
|
|
|
namespace {
|
|
|
|
struct ConvertVectorToSCFPass
|
|
: public ConvertVectorToSCFBase<ConvertVectorToSCFPass> {
|
|
ConvertVectorToSCFPass() = default;
|
|
ConvertVectorToSCFPass(const VectorTransferToSCFOptions &options) {
|
|
this->fullUnroll = options.unroll;
|
|
}
|
|
|
|
void runOnFunction() override {
|
|
RewritePatternSet patterns(getFunction().getContext());
|
|
populateVectorToSCFConversionPatterns(
|
|
patterns, VectorTransferToSCFOptions().setUnroll(fullUnroll));
|
|
(void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
|
|
}
|
|
};
|
|
|
|
} // namespace
|
|
|
|
std::unique_ptr<Pass>
|
|
mlir::createConvertVectorToSCFPass(const VectorTransferToSCFOptions &options) {
|
|
return std::make_unique<ConvertVectorToSCFPass>(options);
|
|
}
|