Google GitGoogle Git
Sign in
llvm/llvm-project/86a649acfded74042e2cec07df4fc2e70a02a74d/./mlir/lib/Dialect/Linalg/Transforms/TilingInterfaceImpl.cpp
blob: bbb23c00e5e193f692e9578c1aa19741467b8c23 [file]
//===- TilingInterfaceImpl.cpp - Implementation of TilingInterface -------===//
//
// 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
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Linalg/Transforms/TilingInterfaceImpl.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Utils.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/Interfaces/TilingInterface.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "llvm/Support/Debug.h"
#include <optional>
#define DEBUG_TYPE "linalg-tiling-interface-impl"
using namespace mlir;
using namespace mlir::linalg;
//===----------------------------------------------------------------------===//
// Utility methods for implementation of Tiling Interface for Linalg ops
//===----------------------------------------------------------------------===//
/// Return the SSA values that represent the data point accessed using a given
/// `indexingMap` for a given point in the iteration space represented by `ivs`.
static SmallVector<Value> getIndicesForAccess(OpBuilder &b, Location loc,
AffineMap indexingMap,
ValueRange ivs) {
SmallVector<Value> indices;
indices.reserve(indexingMap.getNumResults());
for (auto result : indexingMap.getResults()) {
AffineMap m = AffineMap::get(indexingMap.getNumDims(),
indexingMap.getNumSymbols(), result);
Value v = affine::AffineApplyOp::create(b, loc, m, ivs);
indices.push_back(v);
}
return indices;
}
/// Method to inline the payload of a `linalgOp` given the iteration space
/// point and values for the arguments of the payload.
static LogicalResult inlinePayload(OpBuilder &b, LinalgOp linalgOp,
ValueRange ivs, ValueRange argValues) {
Block *body = linalgOp.getBlock();
IRMapping map;
map.map(body->getArguments(), argValues);
for (auto &op : body->without_terminator()) {
if (auto indexOp = dyn_cast<IndexOp>(&op)) {
map.map(indexOp.getResult(), ivs[indexOp.getDim()]);
continue;
}
b.clone(op, map);
}
Operation *terminator = body->getTerminator();
Location loc = terminator->getLoc();
for (const auto &operand : llvm::enumerate(terminator->getOperands())) {
Value toStore = map.lookupOrDefault(operand.value());
OpOperand *storeInto = linalgOp.getDpsInitOperand(operand.index());
auto indices = getIndicesForAccess(
b, loc, linalgOp.getMatchingIndexingMap(storeInto), ivs);
memref::StoreOp::create(b, loc, toStore,
linalgOp.getDpsInitOperand(operand.index())->get(),
indices);
}
return success();
}
/// Verify that tiling can be applied in presence of semi-affine maps.
static LogicalResult
validateTilingSemiAffineMaps(LinalgOp linalgOp, ArrayRef<OpFoldResult> sizes) {
// Precompute each dimension's constant tile-size upper bound once.
// A failed entry marks a dynamic tile with no static bound.
SmallVector<FailureOr<int64_t>> tileSizeBounds =
llvm::map_to_vector(sizes, [](OpFoldResult size) {
return ValueBoundsConstraintSet::computeConstantBound(
presburger::BoundType::UB, size,
/*stopCondition=*/nullptr, ValueBoundsOptions{/*closedUB=*/true});
});
SmallVector<int64_t> loopRanges = linalgOp.getStaticLoopRanges();
// Dynamic tiles or dynamic loop ranges are conservatively treated as tiled.
SmallVector<bool> tiledDims(loopRanges.size(), false);
for (auto [pos, tileSize] : llvm::enumerate(tileSizeBounds)) {
if (failed(tileSize)) {
tiledDims[pos] = true;
continue;
}
if (*tileSize == 0)
continue;
tiledDims[pos] =
ShapedType::isDynamic(loopRanges[pos]) || *tileSize < loopRanges[pos];
}
for (AffineMap map : linalgOp.getIndexingMapsArray()) {
for (AffineExpr result : map.getResults()) {
WalkResult status = result.walk([&](AffineExpr expr) -> WalkResult {
auto binExpr = dyn_cast<AffineBinaryOpExpr>(expr);
if (!binExpr)
return WalkResult::advance();
AffineExprKind kind = binExpr.getKind();
if (kind != AffineExprKind::Mod && kind != AffineExprKind::FloorDiv &&
kind != AffineExprKind::CeilDiv)
return WalkResult::advance();
// Skip if the semi-affine expression does not involve any tiled
// dimension: an untiled dimension keeps its full extent in every tile,
// so re-applying the map on the slice is exact.
bool involvesTiledDim = expr.walk([&](AffineExpr e) -> WalkResult {
auto dim = dyn_cast<AffineDimExpr>(e);
if (dim && tiledDims[dim.getPosition()])
return WalkResult::interrupt();
return WalkResult::advance();
})
.wasInterrupted();
if (!involvesTiledDim)
return WalkResult::advance();
// Allow only `d OP C` map where `d` is a dimension and `C` is a
// constant. A compound LHS (e.g. `(d0 + d1)`, `(d0 * 2)`, a nested
// semi-affine expression) or a non-constant step is not provably safe,
// so reject it.
auto dimExpr = dyn_cast<AffineDimExpr>(binExpr.getLHS());
auto stepExpr = dyn_cast<AffineConstantExpr>(binExpr.getRHS());
if (!dimExpr || !stepExpr || stepExpr.getValue() <= 0) {
linalgOp.emitOpError()
<< "tiling is not supported for the semi-affine indexing map: "
"only a single iteration dimension divided by a positive "
"constant step can be tiled over a tiled dimension";
return WalkResult::interrupt();
}
// Tiles are spaced by the full tile size, so tile origins are its
// multiples (0, tileSize, 2*tileSize, ...).
// A tile's indices are `origin + d'`, with `origin` the tile's start
// and `0 <= d' < tileSize`. A trailing partial tile is a full tile
// truncated at the same origin, spanning a subset of the same `d'`, so
// full-tile validity implies partial-tile validity and validating the
// upper-bound tile size suffices.
unsigned dimPos = dimExpr.getPosition();
FailureOr<int64_t> tileSize = tileSizeBounds[dimPos];
// Dynamic tile sizes are assumed to be valid.
// Unit tile is always valid.
if (failed(tileSize) || *tileSize == 1)
return WalkResult::advance();
// Tiled op reuses the same map on a slice whose base offset is
// `m(origin) - m(0)`, so it is correct only when
// `m(origin + d') == (m(origin) - m(0)) + m(d')` for every `d'`.
// Slice origins are tile-size multiples, so this reduces to a relation
// between the tile size and the step `C`:
// - `floordiv`/`mod` are locally affine within a step window (floordiv
// is constant, mod is linear), so they compose when the origin is
// step-aligned (`C | tileSize`) or the whole tile fits in one window
// (`tileSize | C`);
// - `ceildiv` jumps at `k * C + 1` instead of `k * C`, so a
// non-step-aligned origin already straddles the jump. It composes
// only from a step-aligned origin, i.e. `C | tileSize`.
int64_t step = stepExpr.getValue();
bool isCeil = kind == AffineExprKind::CeilDiv;
bool safe = *tileSize % step == 0 || (!isCeil && step % *tileSize == 0);
if (!safe) {
linalgOp.emitOpError()
<< "tiling is not supported for the semi-affine indexing map: "
"tile size "
<< *tileSize << " for dimension d" << dimPos
<< (isCeil ? " must be a multiple of the step "
: " must divide or be divisible by the step ")
<< step;
return WalkResult::interrupt();
}
return WalkResult::advance();
});
if (status.wasInterrupted())
return failure();
}
}
return success();
}
//===----------------------------------------------------------------------===//
// External Model for implementing `TilingInterface` for `LinalgOp`s.
//===----------------------------------------------------------------------===//
namespace {
/// Operation-independent implementation shared by the external models for
/// LinalgOps. External models are used for now until `TilingInterface` is
/// on-par with the current Linalg tiling and fusion patterns. It may then be
/// possible to move this into the op definitions, though there are advantages
/// to leaving it as an external model.
struct LinalgOpTilingInterfaceImpl {
/// Return the iteration domain range.
SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
OpBuilder::InsertionGuard g(b);
b.setInsertionPoint(op);
Location loc = op->getLoc();
LinalgOp linalgOp = cast<LinalgOp>(op);
SmallVector<OpFoldResult> allShapesSizes =
linalgOp.createFlatListOfOperandDims(b, loc);
AffineMap map = linalgOp.getShapesToLoopsMap();
return llvm::map_to_vector(map.getResults(), [&](AffineExpr loopExpr) {
OpFoldResult ofr = affine::makeComposedFoldedAffineApply(b, loc, loopExpr,
allShapesSizes);
return Range{b.getIndexAttr(0), ofr, b.getIndexAttr(1)};
});
}
/// Instantiate the tiled implementation of the operation.
FailureOr<TilingResult>
getTiledImplementation(Operation *op, OpBuilder &b,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) const {
// Leave the `sizeBounds` value empty. That is only needed when the `sizes`
// specified could lead to out of bounds accesses.
Location loc = op->getLoc();
LinalgOp linalgOp = cast<LinalgOp>(op);
// In case of a semi-affine expression, generalized tracking of tiles would
// require a per-tile-position shift that cannot be expressed by the
// symbol-free indexing maps.
// Thus, tiling is allowed only when the semi-affine maps can be proven safe
// for the current tiling configuration. Otherwise, tiling can end up
// producing incorrect results.
if (failed(validateTilingSemiAffineMaps(linalgOp, sizes)))
return failure();
SmallVector<Value> valuesToTile = linalgOp->getOperands();
SmallVector<Value> tiledOperands = makeTiledShapes(
b, loc, linalgOp, valuesToTile, offsets, sizes, {}, true);
SmallVector<Operation *> generatedSlices = llvm::map_to_vector(
llvm::make_filter_range(
tiledOperands,
[](Value v) -> bool {
return isa_and_nonnull<tensor::ExtractSliceOp, memref::SubViewOp>(
v.getDefiningOp());
}),
[](Value v) -> Operation * { return v.getDefiningOp(); });
SmallVector<Type> resultTensorTypes =
getTensorOutputTypes(linalgOp, tiledOperands);
Operation *tiledOp = clone(b, linalgOp, resultTensorTypes, tiledOperands);
offsetIndices(b, cast<LinalgOp>(tiledOp), offsets);
return TilingResult{
{tiledOp}, SmallVector<Value>(tiledOp->getResults()), generatedSlices};
}
/// Utility to fetch the offsets and sizes when applied as per the indexing
/// map of the linalg op. This helps in fusing the linalg op as a consumer of
/// a given slice op.
static LogicalResult
getMappedOffsetAndSize(LinalgOp linalgOp, OpBuilder &b,
ArrayRef<AffineMap> indexingMaps,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes,
SmallVectorImpl<OpFoldResult> &mappedOffsetsVec,
SmallVectorImpl<OpFoldResult> &mappedSizesVec) {
DenseMap<unsigned, OpFoldResult> mappedOffsets, mappedSizes;
for (auto [indexingMap, offsets, sizes] :
llvm::zip_equal(indexingMaps, allOffsets, allSizes)) {
for (auto [resultExpr, offset, size] :
llvm::zip_equal(indexingMap.getResults(), offsets, sizes)) {
auto dimExpr = dyn_cast<AffineDimExpr>(resultExpr);
if (!dimExpr)
return failure();
unsigned position = dimExpr.getPosition();
auto it = mappedOffsets.find(position);
if (it != mappedOffsets.end()) {
OpFoldResult seenOffset = it->second;
OpFoldResult seenSize = mappedSizes.lookup(position);
if (seenOffset != offset || seenSize != size) {
LLVM_DEBUG({
llvm::dbgs() << "inconsistent iteration space mapping from "
"offsets/sizes of operands/results";
});
return failure();
}
} else {
mappedOffsets[position] = offset;
mappedSizes[position] = size;
}
}
}
// Aggregate from the given operand offsets and sizes, or default to
// iteration space values.
SmallVector<Range> iterationDomain =
cast<TilingInterface>(linalgOp.getOperation()).getIterationDomain(b);
mappedOffsetsVec.resize(iterationDomain.size());
mappedSizesVec.resize(iterationDomain.size());
for (auto [index, domain] : llvm::enumerate(iterationDomain)) {
auto it = mappedOffsets.find(index);
if (it != mappedOffsets.end()) {
mappedOffsetsVec[index] = it->second;
mappedSizesVec[index] = mappedSizes.lookup(index);
continue;
}
mappedOffsetsVec[index] = domain.offset;
mappedSizesVec[index] = domain.size;
}
return success();
}
/// Method to return the position of the result tile computed by the tiled
/// operation.
LogicalResult getIterationDomainTileFromOperandTiles(
Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes,
SmallVectorImpl<OpFoldResult> &iterDomainOffsets,
SmallVectorImpl<OpFoldResult> &iterDomainSizes) const {
auto linalgOp = cast<LinalgOp>(op);
SmallVector<AffineMap> indexingMaps =
llvm::map_to_vector(operandNumbers, [&](unsigned operandNumber) {
OpOperand &opOperand = linalgOp->getOpOperand(operandNumber);
return linalgOp.getMatchingIndexingMap(&opOperand);
});
if (failed(getMappedOffsetAndSize(linalgOp, b, indexingMaps, allOffsets,
allSizes, iterDomainOffsets,
iterDomainSizes))) {
return failure();
}
return success();
}
/// Return the details of the output tile generated by the tiled
/// implementation.
LogicalResult
getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes,
SmallVector<OpFoldResult> &resultOffsets,
SmallVector<OpFoldResult> &resultSizes) const {
Location loc = op->getLoc();
LinalgOp linalgOp = cast<LinalgOp>(op);
AffineExpr d0;
bindDims(b.getContext(), d0);
SmallVector<OpFoldResult> subShapeSizes =
llvm::map_to_vector(sizes, [&](OpFoldResult ofr) {
return affine::makeComposedFoldedAffineApply(b, loc, d0 - 1, ofr);
});
OpOperand *outOperand = linalgOp.getDpsInitOperand(resultNumber);
SliceParameters sliceParams = computeSliceParameters(
b, loc, outOperand->get(), sizes,
linalgOp.getMatchingIndexingMap(outOperand), offsets,
/*ubs*/ {}, subShapeSizes, true);
resultOffsets = sliceParams.offsets;
resultSizes = sliceParams.sizes;
return success();
}
LogicalResult getIterationDomainTileFromResultTile(
Operation *op, OpBuilder &b, unsigned resultNumber,
ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
SmallVectorImpl<OpFoldResult> &iterDomainOffsets,
SmallVectorImpl<OpFoldResult> &iterDomainSizes) const {
auto linalgOp = cast<LinalgOp>(op);
// Check that the indexing map used for the output is a projected
// permutation. This could be relaxed with a more general approach that can
// map the offsets and sizes from the result to iteration space tiles
// (filling in full extent for dimensions not used to access the result).
AffineMap indexingMap =
linalgOp.getIndexingMapMatchingResult(op->getResult(resultNumber));
if (!indexingMap.isProjectedPermutation()) {
return op->emitOpError(
"unhandled tiled implementation generation when result is not "
"accessed using a permuted projection");
}
SmallVector<OpFoldResult> allOffsets = llvm::to_vector(offsets);
SmallVector<OpFoldResult> allSizes = llvm::to_vector(sizes);
auto status =
getMappedOffsetAndSize(linalgOp, b, indexingMap, {allOffsets},
{allSizes}, iterDomainOffsets, iterDomainSizes);
(void)status;
assert(succeeded(status) && "unexpected error in offset calculation");
return success();
}
FailureOr<TilingResult>
generateResultTileValue(Operation *op, OpBuilder &b, unsigned resultNumber,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) const {
SmallVector<OpFoldResult> mappedOffsets, mappedSizes;
if (failed(getIterationDomainTileFromResultTile(
op, b, resultNumber, offsets, sizes, mappedOffsets, mappedSizes))) {
return failure();
}
auto tilingInterfaceOp = cast<TilingInterface>(op);
FailureOr<TilingResult> tilingResult =
tilingInterfaceOp.getTiledImplementation(b, mappedOffsets, mappedSizes);
if (failed(tilingResult))
return failure();
if (tilingResult->tiledOps.size() != 1)
return op->emitOpError("failed to generate tiled implementation");
return TilingResult{
tilingResult->tiledOps,
SmallVector<Value>{tilingResult->tiledValues[resultNumber]},
tilingResult->generatedSlices};
}
/// Method to generate the tiled implementation of an operation from the tile
/// of the operand.
FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
SmallVector<OpFoldResult> mappedOffsets, mappedSizes;
if (failed(getIterationDomainTileFromOperandTiles(
op, b, operandNumbers, allOffsets, allSizes, mappedOffsets,
mappedSizes))) {
return failure();
}
return getTiledImplementation(op, b, mappedOffsets, mappedSizes);
}
LogicalResult generateScalarImplementation(Operation *op, OpBuilder &builder,
Location loc,
ValueRange ivs) const {
auto linalgOp = cast<LinalgOp>(op);
if (!linalgOp.hasPureBufferSemantics())
return op->emitOpError("expected operation to have buffer semantics");
SmallVector<Value> indexedValues;
indexedValues.reserve(linalgOp->getNumOperands());
Location linalgOpLoc = op->getLoc();
/// Load the data corresponding to the block arguments that
/// represent input operands.
for (OpOperand &operand : linalgOp->getOpOperands()) {
if (!linalgOp.payloadUsesValueFromOperand(&operand)) {
indexedValues.push_back(nullptr);
continue;
}
if (linalgOp.isScalar(&operand)) {
indexedValues.push_back(operand.get());
continue;
}
SmallVector<Value> indices = getIndicesForAccess(
builder, linalgOpLoc, linalgOp.getMatchingIndexingMap(&operand), ivs);
Value load =
memref::LoadOp::create(builder, linalgOpLoc, operand.get(), indices);
indexedValues.push_back(load);
}
/// Inline the op payload and store the result.
return inlinePayload(builder, linalgOp, ivs, indexedValues);
}
bool isOpFusableWithConsumerSlice(Operation *op, unsigned resultNumber,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) const {
// The verifier gives all the necessary requirements for consumer fusion.
return true;
}
bool isOpFusableWithProducerSlices(
Operation *op, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
auto linalgOp = cast<LinalgOp>(op);
SmallVector<AffineMap> indexingMaps =
llvm::map_to_vector(operandNumbers, [&](unsigned operandNumber) {
OpOperand &opOperand = linalgOp->getOpOperand(operandNumber);
return linalgOp.getMatchingIndexingMap(&opOperand);
});
// Check that offsets/sizes are consistent across all operands.
OpBuilder b(op);
SmallVector<OpFoldResult> mappedOffsets, mappedSizes;
return succeeded(getMappedOffsetAndSize(linalgOp, b, indexingMaps,
allOffsets, allSizes, mappedOffsets,
mappedSizes));
}
};
template <typename LinalgOpTy>
struct LinalgOpTilingInterfaceModel
: public TilingInterface::ExternalModel<
LinalgOpTilingInterfaceModel<LinalgOpTy>, LinalgOpTy>,
public LinalgOpTilingInterfaceImpl {
using ExternalModel =
TilingInterface::ExternalModel<LinalgOpTilingInterfaceModel<LinalgOpTy>,
LinalgOpTy>;
using LinalgOpTilingInterfaceImpl::generateScalarImplementation;
using LinalgOpTilingInterfaceImpl::getIterationDomain;
using LinalgOpTilingInterfaceImpl::getIterationDomainTileFromResultTile;
using LinalgOpTilingInterfaceImpl::getResultTilePosition;
using LinalgOpTilingInterfaceImpl::isOpFusableWithConsumerSlice;
using LinalgOpTilingInterfaceImpl::isOpFusableWithProducerSlices;
/// Return the loop iterator type without a dynamic LinalgOp interface lookup.
SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
return cast<LinalgOpTy>(op).getIteratorTypesArray();
}
// Preserve the hint-bearing ExternalModel defaults while routing the
// no-hint overloads through the shared implementation.
using ExternalModel::generateResultTileValue;
FailureOr<TilingResult>
generateResultTileValue(Operation *op, OpBuilder &b, unsigned resultNumber,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) const {
return LinalgOpTilingInterfaceImpl::generateResultTileValue(
op, b, resultNumber, offsets, sizes);
}
using ExternalModel::getIterationDomainTileFromOperandTiles;
LogicalResult getIterationDomainTileFromOperandTiles(
Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes,
SmallVectorImpl<OpFoldResult> &iterDomainOffsets,
SmallVectorImpl<OpFoldResult> &iterDomainSizes) const {
return LinalgOpTilingInterfaceImpl::getIterationDomainTileFromOperandTiles(
op, b, operandNumbers, allOffsets, allSizes, iterDomainOffsets,
iterDomainSizes);
}
using ExternalModel::getTiledImplementation;
FailureOr<TilingResult>
getTiledImplementation(Operation *op, OpBuilder &b,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) const {
return LinalgOpTilingInterfaceImpl::getTiledImplementation(op, b, offsets,
sizes);
}
using ExternalModel::getTiledImplementationFromOperandTiles;
FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
return LinalgOpTilingInterfaceImpl::getTiledImplementationFromOperandTiles(
op, b, operandNumbers, allOffsets, allSizes);
}
};
//===----------------------------------------------------------------------===//
// External Model for implementing `PartialReductionInterface` for `LinalgOp`s.
//===----------------------------------------------------------------------===//
/// In a given set vector, get the position of a particular element.
std::optional<int> getPositionIn(const llvm::SetVector<unsigned> &reductionDims,
unsigned value) {
for (auto [index, reductionDim] : llvm::enumerate(reductionDims)) {
if (reductionDim == value) {
return index;
}
}
return std::nullopt;
}
/// Return an AffineMaps to use for the `outs` operands of the linalg op
/// generated for partial results. The new AffineMap is the AffineMap of the
/// untiled op with reduction dimensions appended at end in order in which they
/// were specified during tiling.
static SmallVector<AffineMap>
getPartialResultAffineMaps(LinalgOp linalgOp,
const SetVector<unsigned> &reductionDims) {
auto partialReductionMaps = llvm::map_to_vector(
linalgOp.getDpsInitsMutable(), [&](OpOperand &opOperand) {
AffineMap map = linalgOp.getMatchingIndexingMap(&opOperand);
for (auto redPos : reductionDims) {
map =
map.insertResult(getAffineDimExpr(redPos, linalgOp.getContext()),
map.getNumResults());
}
return map;
});
return partialReductionMaps;
}
struct InitSliceInfo {
SmallVector<int64_t> resultShape;
SmallVector<OpFoldResult> offsets;
SmallVector<OpFoldResult> sizes;
SmallVector<OpFoldResult> strides;
};
/// Return the result shape, offsets, sizes and strides of the slice of the
/// `initValue` to use as the destination of the partial reduction op generated
/// with outer reduction strategy.
static InitSliceInfo getInitSliceInfoForOuterReduction(
MLIRContext *context, ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes, const SetVector<unsigned> &reductionDims,
ArrayRef<OpFoldResult> splitReductionIvs, AffineMap partialReductionMap,
ArrayRef<OpFoldResult> initOperandShape) {
int64_t initRank = partialReductionMap.getNumResults();
SmallVector<OpFoldResult> initOffsets, initSizes;
Attribute zero = IntegerAttr::get(IndexType::get(context), 0);
Attribute one = IntegerAttr::get(IndexType::get(context), 1);
SmallVector<OpFoldResult> initStrides(initRank, one);
for (auto [resultIdx, dimExpr] :
llvm::enumerate(partialReductionMap.getResults())) {
if (isa<AffineConstantExpr>(dimExpr)) {
// A constant index in the output map accesses a fixed position; keep
// the full output dimension to match the original output operand shape.
initOffsets.push_back(zero);
initSizes.push_back(initOperandShape[resultIdx]);
continue;
}
unsigned dim = cast<AffineDimExpr>(dimExpr).getPosition();
if (reductionDims.contains(dim)) {
initOffsets.push_back(zero);
} else {
initOffsets.push_back(offsets[dim]);
}
initSizes.push_back(sizes[dim]);
}
SmallVector<int64_t> resultShape;
std::tie(resultShape, std::ignore) = decomposeMixedValues(initSizes);
return {resultShape, initOffsets, initSizes, initStrides};
}
/// Return the result shape, offsets, sizes and strides of the slice of the
/// `initValue` to use as destination of the partial reduction op generated with
/// outer parallel strategy.
static InitSliceInfo getInitSliceInfoForOuterParallel(
MLIRContext *context, ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes, const SetVector<unsigned> &reductionDims,
ArrayRef<OpFoldResult> splitReductionIvs, AffineMap partialReductionMap,
ArrayRef<OpFoldResult> initOperandShape) {
int64_t initRank = partialReductionMap.getNumResults();
SmallVector<OpFoldResult> initOffsets, initSizes;
Attribute zero = IntegerAttr::get(IndexType::get(context), 0);
Attribute one = IntegerAttr::get(IndexType::get(context), 1);
SmallVector<OpFoldResult> initStrides(initRank, one);
SmallVector<OpFoldResult> resultShape;
for (auto [resultIdx, dimExpr] :
llvm::enumerate(partialReductionMap.getResults())) {
if (isa<AffineConstantExpr>(dimExpr)) {
// A constant index accesses a fixed position; keep the full output
// dimension to match the original output operand shape.
initOffsets.push_back(zero);
initSizes.push_back(initOperandShape[resultIdx]);
resultShape.push_back(initOperandShape[resultIdx]);
continue;
}
unsigned dim = cast<AffineDimExpr>(dimExpr).getPosition();
if (std::optional<unsigned> dimPos = getPositionIn(reductionDims, dim)) {
initOffsets.push_back(splitReductionIvs[dimPos.value()]);
initSizes.push_back(one);
} else {
initOffsets.push_back(offsets[dim]);
initSizes.push_back(sizes[dim]);
resultShape.push_back(sizes[dim]);
}
}
SmallVector<int64_t> staticShapes;
std::tie(staticShapes, std::ignore) = decomposeMixedValues(resultShape);
return {staticShapes, initOffsets, initSizes, initStrides};
}
/// Return the result shape, offsets, sizes and strides of the slice of the
/// `initValue` to use as destination of the partial reduction op.
static InitSliceInfo getInitSliceInfo(MLIRContext *context,
ReductionTilingStrategy strategy,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes,
const SetVector<unsigned> &reductionDims,
ArrayRef<OpFoldResult> splitReductionIvs,
AffineMap partialReductionMap,
ArrayRef<OpFoldResult> initOperandShape) {
if (strategy == ReductionTilingStrategy::PartialReductionOuterReduction) {
return getInitSliceInfoForOuterReduction(
context, offsets, sizes, reductionDims, splitReductionIvs,
partialReductionMap, initOperandShape);
}
assert(strategy == ReductionTilingStrategy::PartialReductionOuterParallel &&
"unexpected ReductionTilingStrategy");
return getInitSliceInfoForOuterParallel(
context, offsets, sizes, reductionDims, splitReductionIvs,
partialReductionMap, initOperandShape);
}
/// Operation-independent implementation shared by the
/// PartialReductionInterface external models for LinalgOps.
struct LinalgOpPartialReductionInterfaceImpl {
FailureOr<SmallVector<Value>> generateInitialTensorForPartialReduction(
Operation *op, OpBuilder &b, Location loc, ArrayRef<OpFoldResult> sizes,
const SetVector<unsigned> &reductionDims) const {
auto linalgOp = cast<LinalgOp>(op);
OpBuilder::InsertionGuard guard(b);
if (linalgOp.hasPureBufferSemantics())
return op->emitOpError("expected operation to have tensor semantics");
SmallVector<AffineMap> partialResultMaps =
getPartialResultAffineMaps(linalgOp, reductionDims);
SmallVector<Value> inits;
for (auto [initIdx, result, partialMap] :
llvm::enumerate(linalgOp->getResults(), partialResultMaps)) {
SmallVector<Operation *, 4> combinerOps;
if (!matchReduction(linalgOp.getRegionOutputArgs(), initIdx,
combinerOps) ||
combinerOps.size() != 1)
return op->emitOpError("Failed to anaysis the reduction operation.");
Operation *reductionOp = combinerOps[0];
std::optional<TypedAttr> identity = arith::getNeutralElement(reductionOp);
if (!identity.has_value())
return op->emitOpError(
"Failed to get an identity value for the reduction operation.");
// Append the new partial result dimensions.
SmallVector<OpFoldResult> partialResultShape;
Value initValue = linalgOp.getDpsInits()[initIdx];
SmallVector<OpFoldResult> initShape =
tensor::getMixedSizes(b, loc, initValue);
for (auto [resultIdx, dimExpr] :
llvm::enumerate(partialMap.getResults())) {
if (isa<AffineConstantExpr>(dimExpr)) {
// A constant index in the output map accesses a fixed position; use
// the actual output dimension size (not a hardcoded 1).
partialResultShape.push_back(initShape[resultIdx]);
continue;
}
auto dim = cast<AffineDimExpr>(dimExpr);
partialResultShape.push_back(sizes[dim.getPosition()]);
}
Type elType = getElementTypeOrSelf(result.getType());
Value emptyTensor =
tensor::EmptyOp::create(b, loc, partialResultShape, elType);
Value constantOp = arith::ConstantOp::create(b, loc, *identity);
auto identityTensor =
linalg::FillOp::create(b, loc, constantOp, emptyTensor);
inits.push_back(identityTensor.getResult(0));
}
return inits;
}
FailureOr<TilingResult>
tileToPartialReduction(Operation *op, OpBuilder &b, Location loc,
ReductionTilingStrategy tilingStrategy,
ValueRange init, ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes,
const SetVector<unsigned> &reductionDims,
ArrayRef<OpFoldResult> splitReductionIvs) const {
OpBuilder::InsertionGuard guard(b);
auto linalgOp = cast<LinalgOp>(op);
SmallVector<AffineMap> partialReductionMaps =
getPartialResultAffineMaps(linalgOp, reductionDims);
// Step 1. Extend init maps to have reduction dimension dims, since we
// are converting them to parallel dimensions.
SmallVector<AffineMap> newInitMaps;
if (tilingStrategy ==
ReductionTilingStrategy::PartialReductionOuterReduction) {
newInitMaps = llvm::to_vector(partialReductionMaps);
} else {
newInitMaps = llvm::map_to_vector(
linalgOp.getDpsInitsMutable(), [&](OpOperand &opOperand) {
return linalgOp.getMatchingIndexingMap(&opOperand);
});
}
// Step 2a: Extract a slice of the input operands.
SmallVector<Value> tiledInputs = makeTiledShapes(
b, loc, linalgOp, linalgOp.getDpsInputs(), offsets, sizes, {}, true);
SmallVector<Operation *> generatedSlices = llvm::map_to_vector(
llvm::make_filter_range(
tiledInputs, [](Value v) -> bool { return v.getDefiningOp(); }),
[](Value v) -> Operation * { return v.getDefiningOp(); });
// Step 2b: Extract a slice of the init operands.
SmallVector<Value, 1> tiledInits;
for (auto [partialReductionMap, valueToTile, initOperandValue] :
llvm::zip_equal(partialReductionMaps, init, linalgOp.getDpsInits())) {
// Compute the actual shape of the original init operand for handling
// constant expressions in the partial reduction map.
SmallVector<OpFoldResult> initOperandShape =
tensor::getMixedSizes(b, loc, initOperandValue);
InitSliceInfo sliceInfo = getInitSliceInfo(
b.getContext(), tilingStrategy, offsets, sizes, reductionDims,
splitReductionIvs, partialReductionMap, initOperandShape);
auto valueToTileType = cast<RankedTensorType>(valueToTile.getType());
RankedTensorType sliceResultType = RankedTensorType::get(
sliceInfo.resultShape, valueToTileType.getElementType(),
valueToTileType.getEncoding());
auto sliceOp = tensor::ExtractSliceOp::create(
b, loc, sliceResultType, valueToTile, sliceInfo.offsets,
sliceInfo.sizes, sliceInfo.strides);
tiledInits.push_back(sliceOp.getResult());
generatedSlices.push_back(sliceOp);
}
// Update the indexing maps.
SmallVector<AffineMap> newMaps = linalgOp.getIndexingMapsArray();
for (auto [initOperand, newInitMap] :
llvm::zip_equal(linalgOp.getDpsInitsMutable(), newInitMaps)) {
int mapIdx = linalgOp.getIndexingMapIndex(&initOperand);
newMaps[mapIdx] = newInitMap;
}
// Step 3. Change the reduction dim iterator types.
SmallVector<utils::IteratorType> newIteratorTypes =
linalgOp.getIteratorTypesArray();
if (tilingStrategy ==
ReductionTilingStrategy::PartialReductionOuterReduction) {
for (int dim : reductionDims)
newIteratorTypes[dim] = utils::IteratorType::parallel;
}
// Step 4. Create the new generic op.
Operation *partialReductionOp;
auto resultTypes = ValueRange(tiledInits).getTypes();
if (tilingStrategy ==
ReductionTilingStrategy::PartialReductionOuterReduction) {
auto genericOp = GenericOp::create(b, loc, resultTypes, tiledInputs,
tiledInits, newMaps, newIteratorTypes);
IRMapping mapping;
op->getRegion(0).cloneInto(&genericOp.getRegion(),
genericOp.getRegion().begin(), mapping);
offsetIndices(b, genericOp, offsets);
partialReductionOp = genericOp.getOperation();
} else {
SmallVector<Value> operands = std::move(tiledInputs);
llvm::append_range(operands, tiledInits);
partialReductionOp = mlir::clone(b, op, resultTypes, operands);
offsetIndices(b, cast<LinalgOp>(partialReductionOp), offsets);
}
return TilingResult{
{partialReductionOp},
llvm::map_to_vector(partialReductionOp->getResults(),
[](OpResult r) -> Value { return r; }),
generatedSlices};
}
FailureOr<MergeResult>
mergeReductions(Operation *op, OpBuilder &b, Location loc,
ValueRange partialReduce,
const SetVector<unsigned> &reductionDims) const {
auto linalgOp = cast<LinalgOp>(op);
SmallVector<AffineMap> partialReductionMaps =
getPartialResultAffineMaps(linalgOp, reductionDims);
// Permute the reduction dims as permuted by the partial result map.
SmallVector<Operation *> mergeOperations;
SmallVector<Value> replacements;
for (auto [idx, init, partialResult, partialMap] : llvm::enumerate(
linalgOp.getDpsInits(), partialReduce, partialReductionMaps)) {
unsigned initIdx = idx;
// linalg.reduce's iteration space is the tiled result's iteration space
// (and not the tiled operation's iteration space). To account for this,
// permute the reduction dimensions based on the partial result map of the
// tiled result.
SmallVector<int64_t> partialReductionDims;
for (auto [resultNum, dimExpr] :
llvm::enumerate(partialMap.getResults())) {
if (isa<AffineConstantExpr>(dimExpr))
continue; // Constant dims are never reduction dims.
unsigned dim = cast<AffineDimExpr>(dimExpr).getPosition();
if (llvm::is_contained(reductionDims, dim)) {
partialReductionDims.push_back(resultNum);
}
}
auto reduction = linalg::ReduceOp::create(
b, loc, partialResult, init, partialReductionDims,
[&linalgOp, &initIdx](OpBuilder &b, Location loc, ValueRange inputs) {
// Get the combiner op.
SmallVector<Operation *, 4> combinerOps;
matchReduction(linalgOp.getRegionOutputArgs(), initIdx,
combinerOps);
Operation *clonedReductionOp = b.clone(*combinerOps[0]);
// Combine the input at idx and output at numInits + idx.
clonedReductionOp->setOperand(0, inputs[0]);
clonedReductionOp->setOperand(1, inputs[1]);
linalg::YieldOp::create(b, loc, clonedReductionOp->getResult(0));
});
mergeOperations.push_back(reduction);
replacements.push_back(reduction->getResult(0));
}
return MergeResult{mergeOperations, replacements};
}
LogicalResult getPartialResultTilePosition(
Operation *op, OpBuilder &b, unsigned resultNumber,
ReductionTilingStrategy tilingStrategy, ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes, const SetVector<unsigned> &reductionDims,
ArrayRef<OpFoldResult> splitReductionIvs,
SmallVector<OpFoldResult> &resultOffsets,
SmallVector<OpFoldResult> &resultSizes) const {
auto linalgOp = cast<LinalgOp>(op);
SmallVector<AffineMap> partialReductionMaps =
getPartialResultAffineMaps(linalgOp, reductionDims);
// Compute the actual shape of the init operand for handling constant
// expressions in the partial reduction map.
Value initOperandValue = linalgOp.getDpsInits()[resultNumber];
Location loc = op->getLoc();
SmallVector<OpFoldResult> initOperandShape =
tensor::getMixedSizes(b, loc, initOperandValue);
InitSliceInfo sliceInfo =
getInitSliceInfo(b.getContext(), tilingStrategy, offsets, sizes,
reductionDims, splitReductionIvs,
partialReductionMaps[resultNumber], initOperandShape);
std::swap(resultOffsets, sliceInfo.offsets);
std::swap(resultSizes, sliceInfo.sizes);
return success();
}
};
template <typename LinalgOpTy>
struct LinalgOpPartialReductionInterfaceModel
: public PartialReductionOpInterface::ExternalModel<
LinalgOpPartialReductionInterfaceModel<LinalgOpTy>, LinalgOpTy>,
public LinalgOpPartialReductionInterfaceImpl {
using LinalgOpPartialReductionInterfaceImpl::
generateInitialTensorForPartialReduction;
using LinalgOpPartialReductionInterfaceImpl::getPartialResultTilePosition;
using LinalgOpPartialReductionInterfaceImpl::mergeReductions;
using LinalgOpPartialReductionInterfaceImpl::tileToPartialReduction;
};
template <typename OpTy>
static SmallVector<Range> getPackUnPackIterationDomain(OpTy op,
OpBuilder &builder) {
static_assert(llvm::is_one_of<OpTy, PackOp, UnPackOp>::value,
"applies to only pack or unpack operations");
OpBuilder::InsertionGuard g(builder);
int64_t rank = (std::is_same<OpTy, PackOp>::value) ? op.getSourceRank()
: op.getDestRank();
OpFoldResult zero = builder.getIndexAttr(0);
OpFoldResult one = builder.getIndexAttr(1);
ReifiedRankedShapedTypeDims resultShape;
(void)op.reifyResultShapes(builder, resultShape);
SmallVector<Range> loopBounds(rank);
for (auto dim : llvm::seq<int64_t>(0, rank)) {
loopBounds[dim].offset = zero;
loopBounds[dim].stride = one;
loopBounds[dim].size = resultShape[0][dim];
}
return loopBounds;
}
static void applyPermToRange(SmallVector<OpFoldResult> &offsets,
SmallVector<OpFoldResult> &sizes,
ArrayRef<int64_t> permutation) {
if (permutation.empty())
return;
applyPermutationToVector<OpFoldResult>(offsets, permutation);
applyPermutationToVector<OpFoldResult>(sizes, permutation);
}
/// Compute the permutation vector to interchange `elements` such that the
/// elements at positions in `dimsPos` are moved to the positions `[0, ...,
/// dimsPos.size())` in order.
static SmallVector<int64_t>
computeInterchangeFromDimPos(ArrayRef<int64_t> dimsPos, int64_t rank) {
SmallVector<int64_t> interchangeVector;
interchangeVector.reserve(dimsPos.size());
// First map dims and their position. For example, dims_pos = [2, 0] will map
// to:
// [
// [ key: 2, value: 0]
// [ key: 0, value: 1]
// ]
// where key is the idx in dims_pos while value its position in dims_pos.
DenseMap<int64_t, int64_t> dimsAndPosMapping;
for (int64_t dimsIdx = 0, end = dimsPos.size(); dimsIdx < end; dimsIdx++)
dimsAndPosMapping[dimsPos[dimsIdx]] = dimsIdx;
// Scan the position in order and insert the value in the map
// to compute the interchange vector.
for (int64_t dimsIdx = 0; dimsIdx < rank; dimsIdx++) {
if (dimsAndPosMapping.count(dimsIdx))
interchangeVector.push_back(dimsAndPosMapping[dimsIdx]);
}
return interchangeVector;
}
/// Permute the elements of `vec` starting at position `offset` according to
/// `interchangeVector`. The permutation maps position `i` in the permuted range
/// to position `interchangeVector[i]` in the original range. Elements before
/// `offset` are unchanged.
///
/// Example: interchange([a, b, c, d, e], [2, 0, 1], offset=2)
/// returns [a, b, e, c, d] (permutes the suffix [c, d, e])
///
/// Note: This is similar to `applyPermutationToVector` but supports an offset
/// for permuting a suffix of the vector. It is only used for pack/unpack scalar
/// implementation where we need to permute inner tile dimensions which are
/// stored at the end of the index vector.
template <typename T>
static SmallVector<T> interchange(ArrayRef<T> elements,
ArrayRef<int64_t> interchangeVector,
int offset = 0) {
SmallVector<T> vec = llvm::to_vector(elements);
for (auto [idx, val] : llvm::enumerate(interchangeVector))
vec[idx + offset] = elements[val + offset];
return vec;
}
/// Generate the body of the innermost loop of the scalar implementation
/// of `pack` operation.
static void generatePackOpScalarImplementationBody(PackOp packOp,
OpBuilder &builder,
Location loc,
ValueRange ivs) {
// Note: `ivs` are already in the correct order, possibly interchanged based
// on `dims_pos`. However, connecting the loops with the access patterns is
// difficult - What is the relation between the position of the tile loop and
// the point loop? However, if we interchange `ivs` once more to go to the
// canonical blocking format: ABCabc, this connection becomes trivial: Each
// point loop is pointLoopsOffset + inputRank away from the tiled loop.
ArrayRef<int64_t> dimsToInnerBlock = packOp.getInnerDimsPos();
ArrayRef<int64_t> dimsToOuterBlock = packOp.getOuterDimsPerm();
SmallVector<Value> interchangedIvs = ivs;
SmallVector<int64_t> interchangeVector =
computeInterchangeFromDimPos(dimsToInnerBlock, packOp.getSourceRank());
interchangedIvs = interchange<Value>(interchangedIvs, interchangeVector,
/*offset=*/packOp.getSourceRank());
if (!dimsToOuterBlock.empty()) {
interchangeVector =
computeInterchangeFromDimPos(dimsToOuterBlock, packOp.getSourceRank());
interchangedIvs =
interchange<Value>(interchangedIvs, interchangeVector, /*offset=*/0);
}
DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
packOp.getDimAndTileMapping();
SmallVector<OpFoldResult> sourceIndices;
size_t pointLoopsOffset = 0;
int64_t sourceRank = packOp.getSourceRank();
for (auto dim : llvm::seq<int64_t>(0, sourceRank)) {
if (dimAndTileMapping.contains(dim)) {
AffineExpr i, j, tile;
bindDims(builder.getContext(), i, j);
bindSymbols(builder.getContext(), tile);
OpFoldResult sourceIndex = affine::makeComposedFoldedAffineApply(
builder, loc, i * tile + j,
ArrayRef<OpFoldResult>{
interchangedIvs[dim],
interchangedIvs[pointLoopsOffset + packOp.getSourceRank()],
dimAndTileMapping[dim]});
sourceIndices.push_back(sourceIndex);
++pointLoopsOffset;
} else {
sourceIndices.push_back(interchangedIvs[dim]);
}
}
auto createLoad = [&]() -> Value {
return memref::LoadOp::create(
builder, loc, packOp.getSource(),
getValueOrCreateConstantIndexOp(builder, loc, sourceIndices));
};
Value scalar;
if (auto paddingValue = packOp.getPaddingValue()) {
ArithBuilder arithBuilder(builder, loc);
Value isInBounds;
for (auto dim : llvm::seq<int64_t>(0, sourceRank)) {
Value idx =
getValueOrCreateConstantIndexOp(builder, loc, sourceIndices[dim]);
Value cond = arithBuilder.slt(
idx, createOrFoldDimOp(builder, loc, packOp.getSource(), dim));
isInBounds = dim == 0 ? cond : arithBuilder._and(isInBounds, cond);
}
scalar = scf::IfOp::create(
builder, loc, isInBounds, /*thenBuilder=*/
[&](OpBuilder &b, Location l) {
scf::YieldOp::create(b, l, createLoad());
},
/*elseBuilder=*/
[&](OpBuilder &b, Location l) {
scf::YieldOp::create(b, l, paddingValue);
})
.getResult(0);
} else {
scalar = createLoad();
}
memref::StoreOp::create(builder, loc, scalar, packOp.getDest(), ivs);
}
struct PackOpTiling
: public TilingInterface::ExternalModel<PackOpTiling, linalg::PackOp> {
using Base = TilingInterface::ExternalModel<PackOpTiling, linalg::PackOp>;
using Base::getTiledImplementation;
SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
// Note that here we only consider untiled dimensions and outer tiled data
// dimensions, the inner tiled data dimensions are materialized when
// building the body of the operation.
auto packOp = cast<PackOp>(op);
SmallVector<utils::IteratorType> iteratorTypes(
packOp.getSourceRank(), utils::IteratorType::parallel);
return iteratorTypes;
}
SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
return getPackUnPackIterationDomain<PackOp>(cast<PackOp>(op), b);
}
FailureOr<TilingResult>
getTiledImplementation(Operation *op, OpBuilder &b,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) const {
auto packOp = cast<PackOp>(op);
// TODO: Support Memref PackOp. Temporarily return failure.
if (!packOp.hasPureTensorSemantics())
return failure();
Location loc = packOp.getLoc();
// The tiling is applied on interchanged dimensions. We have to undo the
// interchange to map sizes and offsets to the original input.
int64_t inputRank = packOp.getSourceRank();
SmallVector<OpFoldResult> origOffsets(offsets);
SmallVector<OpFoldResult> origSizes(sizes);
applyPermToRange(origOffsets, origSizes,
invertPermutationVector(packOp.getOuterDimsPerm()));
DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
packOp.getDimAndTileMapping();
SmallVector<OpFoldResult> srcDimValues =
tensor::getMixedSizes(b, loc, packOp.getSource());
SmallVector<OpFoldResult> inputIndices, inputSizes;
for (auto dim : llvm::seq<int64_t>(0, inputRank)) {
using AV = affine::AffineValueExpr;
affine::AffineBuilder ab(b, loc);
AffineExpr dim0, dim1, sym;
bindDims(b.getContext(), dim0, dim1);
bindSymbols(b.getContext(), sym);
if (dimAndTileMapping.count(dim)) {
// If the data dimension is tiled, the i-th index is the product of
// offset_i and tile_i, and the i-th size is the product of sizes_i and
// tile_i.
auto avOffset = AV(dim0).bind(origOffsets[dim]);
auto avSize = AV(dim0).bind(origSizes[dim]);
auto avTileSize = AV(sym).bind(dimAndTileMapping[dim]);
inputIndices.push_back(ab.mul(avOffset, avTileSize));
inputSizes.push_back(ab.mul(avSize, avTileSize));
} else {
inputIndices.push_back(origOffsets[dim]);
inputSizes.push_back(origSizes[dim]);
}
// Limit the size of the input operand for incomplete tiles.
if (packOp.getPaddingValue()) {
OpFoldResult dimSize = srcDimValues[dim];
auto avDimSize = AV(dim0).bind(dimSize);
auto avInputIdx = AV(dim1).bind(inputIndices.back());
inputSizes.back() =
ab.min({inputSizes.back(), ab.sub(avDimSize, avInputIdx)});
}
}
auto oneAttr = b.getI64IntegerAttr(1);
SmallVector<OpFoldResult> strides(inputRank, oneAttr);
SmallVector<Value> tiledOperands;
auto sourceSlice = tensor::ExtractSliceOp::create(
b, loc, packOp.getSource(), inputIndices, inputSizes, strides);
tiledOperands.push_back(sourceSlice);
SmallVector<OpFoldResult> outputOffsets, outputSizes;
if (failed(getResultTilePosition(op, b, 0, offsets, sizes, outputOffsets,
outputSizes)))
return {};
strides.append(packOp.getDestRank() - inputRank, oneAttr);
auto outSlice = tensor::ExtractSliceOp::create(
b, loc, packOp.getDest(), outputOffsets, outputSizes, strides);
tiledOperands.push_back(outSlice);
if (auto val = packOp.getPaddingValue())
tiledOperands.push_back(val);
for (auto tile : packOp.getInnerTiles())
tiledOperands.push_back(tile);
PackOp tiledPackOp =
PackOp::create(b, loc, TypeRange{outSlice.getType()}, tiledOperands,
packOp.getProperties(),
packOp->getDiscardableAttrDictionary().getValue());
return TilingResult{
{tiledPackOp},
SmallVector<Value>(tiledPackOp->getResults()),
llvm::to_vector(ArrayRef<Operation *>{sourceSlice, outSlice})};
}
LogicalResult
getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes,
SmallVector<OpFoldResult> &resultOffsets,
SmallVector<OpFoldResult> &resultSizes) const {
// The iteration domain is over outer dimensions of packed layout. In this
// context, the outer dimensions of `resultOffsets` are `offsets`. The
// inner dimensions of `resultOffsets` are zeros because tiling is not
// applied to them.
auto packOp = cast<PackOp>(op);
int64_t inputRank = packOp.getSourceRank();
int64_t outputRank = packOp.getDestRank();
auto zeroAttr = b.getI64IntegerAttr(0);
resultOffsets.assign(offsets.begin(), offsets.end());
resultOffsets.append(outputRank - inputRank, zeroAttr);
ReifiedRankedShapedTypeDims outputShape;
(void)reifyResultShapes(b, packOp, outputShape);
resultSizes.assign(sizes.begin(), sizes.end());
for (auto dataTileDim : llvm::seq<unsigned>(inputRank, outputRank))
resultSizes.push_back(outputShape[0][dataTileDim]);
return success();
}
FailureOr<TilingResult>
generateResultTileValue(Operation *op, OpBuilder &b, unsigned resultNumber,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) const {
return generateResultTileValue(op, b, resultNumber, offsets, sizes,
/*innerTileAlignments=*/{});
}
FailureOr<TilingResult> generateResultTileValue(
Operation *op, OpBuilder &b, unsigned resultNumber,
ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
ArrayRef<InnerTileAlignment> innerTileAlignments) const {
auto packOp = cast<PackOp>(op);
int64_t numTiles = packOp.getInnerDimsPos().size();
// linalg.pack op is fusible (as a producer) only if full inner tiles are
// iterated or inner dims are not tiled. Otherwise, it will generate a
// sequence of non-trivial ops (for partial tiles).
for (auto offset : offsets.take_back(numTiles))
if (!isZeroInteger(offset))
return failure();
// Each requested inner-dim size must cover a full inner tile. A caller may
// instead assert this via an `Equal` alignment hint. The hint is indexed by
// source dim, matching the consumer-fusion path.
ArrayRef<int64_t> innerDimsPos = packOp.getInnerDimsPos();
SmallVector<OpFoldResult> mixedTiles = packOp.getMixedTiles();
ArrayRef<OpFoldResult> innerSizes = sizes.take_back(numTiles);
for (auto [i, pos] : llvm::enumerate(innerDimsPos)) {
InnerTileAlignment alignment =
pos < static_cast<int64_t>(innerTileAlignments.size())
? innerTileAlignments[pos]
: InnerTileAlignment::Unknown;
if (alignment != InnerTileAlignment::Equal &&
!isEqualConstantIntOrValue(mixedTiles[i], innerSizes[i]))
return failure();
}
FailureOr<TilingResult> tilingResult = getTiledImplementation(
op, b, offsets.drop_back(numTiles), sizes.drop_back(numTiles));
if (failed(tilingResult))
return failure();
return tilingResult.value();
}
LogicalResult generateScalarImplementation(Operation *op, OpBuilder &builder,
Location loc,
ValueRange ivs) const {
auto packOp = cast<PackOp>(op);
assert(packOp.hasPureBufferSemantics() &&
"expected operation to have buffer semantics");
OpBuilder::InsertionGuard g(builder);
// The `ivs` already represent the position into the output for the non
// data-tile dimensions.
SmallVector<Value> ivVec(ivs);
// Get output shape - for memrefs, get dimensions from dest directly.
SmallVector<OpFoldResult> outputShape;
Value dest = packOp.getDest();
for (auto dim : llvm::seq<int64_t>(0, packOp.getDestRank()))
outputShape.push_back(createOrFoldDimOp(builder, loc, dest, dim));
// Generate the loops that iterate over the data tile.
Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
Value one = arith::ConstantIndexOp::create(builder, loc, 1);
// All loops except the innermost are simple loops that just iterate
// over the tile dimensions.
for (auto dataTileDim : llvm::seq<unsigned>(packOp.getSourceRank(),
packOp.getDestRank() - 1)) {
Value ub = getValueOrCreateConstantIndexOp(builder, loc,
outputShape[dataTileDim]);
scf::ForOp loop = scf::ForOp::create(builder, loc, zero, ub, one);
builder.setInsertionPointToStart(loop.getBody());
ivVec.push_back(loop.getInductionVar());
}
// The body of the innermost loops does the actual data movement.
scf::ForOp::create(
builder, loc, zero,
getValueOrCreateConstantIndexOp(builder, loc, outputShape.back()), one,
ValueRange{},
[&](OpBuilder &bodyBuilder, Location bodyLoc, Value iv,
ValueRange regionIterArgs) {
ivVec.push_back(iv);
generatePackOpScalarImplementationBody(packOp, bodyBuilder, bodyLoc,
ivVec);
scf::YieldOp::create(bodyBuilder, bodyLoc);
});
return success();
}
LogicalResult getIterationDomainTileFromOperandTiles(
Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes,
SmallVectorImpl<OpFoldResult> &resultOffsets,
SmallVectorImpl<OpFoldResult> &resultSizes) const {
return getIterationDomainTileFromOperandTiles(
op, b, operandNumbers, allOffsets, allSizes, resultOffsets, resultSizes,
/*innerTileAlignments=*/{});
}
/// Method to return the position of iteration domain tile computed by the
/// tiled operation. In current `linalg.pack` context, the `resultOffsets` and
/// `resultSizes` only cover outer dimensions.
LogicalResult getIterationDomainTileFromOperandTiles(
Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes,
SmallVectorImpl<OpFoldResult> &resultOffsets,
SmallVectorImpl<OpFoldResult> &resultSizes,
ArrayRef<InnerTileAlignment> innerTileAlignments) const {
if (operandNumbers.size() != 1 || operandNumbers[0] != 0) {
LLVM_DEBUG(
{ llvm::dbgs() << "unsupported operands for consumer fusion"; });
return failure();
}
ArrayRef<OpFoldResult> offsets(allOffsets[0]);
ArrayRef<OpFoldResult> sizes(allSizes[0]);
auto packOp = cast<PackOp>(op);
Location loc = packOp.getLoc();
SmallVector<OpFoldResult> outerDimOffsets, outerDimSizes;
DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
packOp.getDimAndTileMapping();
SmallVector<int64_t> outerShapeWithoutTranspose(
packOp.getDestType().getShape().take_front(packOp.getSourceRank()));
if (!packOp.getOuterDimsPerm().empty()) {
applyPermutationToVector(
outerShapeWithoutTranspose,
invertPermutationVector(packOp.getOuterDimsPerm()));
}
for (auto dim : llvm::seq<int64_t>(packOp.getSourceRank())) {
if (dimAndTileMapping.count(dim)) {
FailureOr<int64_t> cstTileSize =
ValueBoundsConstraintSet::computeConstantBound(
presburger::BoundType::UB, sizes[dim],
/*stopCondition=*/nullptr,
ValueBoundsOptions{/*closedUB=*/true});
std::optional<int64_t> cstInnerSize =
getConstantIntValue(dimAndTileMapping[dim]);
// A caller-supplied alignment hint (see InnerTileAlignment) asserts
// that this packed dimension is tiled and how its loop tile size
// relates to the pack op inner tile size.
InnerTileAlignment innerTileAlignment =
dim < static_cast<int64_t>(innerTileAlignments.size())
? innerTileAlignments[dim]
: InnerTileAlignment::Unknown;
// If a dimension is not tiled, it is always valid to fuse the pack op,
// even if the op has padding semantics. Because it always generates a
// full slice along the dimension. The tile sizes are for unpacked
// domain, i.e., `srcDimSize`, so `tileSize < srcDimSize` means that the
// dimension is tiled.
// TODO: It could be untiled if the `srcDimSize` is dynamic. It is a
// hard check to determine if a dimension is tiled or not.
// A non-`Unknown` hint also means the caller asserts the dimension is
// tiled: `cstTileSize` is an upper bound, so a scalable/`min`-shaped
// tile (whose bound equals `srcDimSize`) would otherwise be mistaken
// for untiled and bypass the hint below.
int64_t srcDimSize = packOp.getSourceType().getDimSize(dim);
int64_t destDimSize = outerShapeWithoutTranspose[dim];
bool isTiled = innerTileAlignment != InnerTileAlignment::Unknown ||
failed(cstTileSize) ||
ShapedType::isDynamic(srcDimSize) ||
cstTileSize.value() < srcDimSize;
if (!isTiled) {
outerDimOffsets.push_back(offsets[dim]);
if (ShapedType::isStatic(destDimSize)) {
outerDimSizes.push_back(b.getIndexAttr(destDimSize));
} else {
outerDimSizes.push_back(
b.createOrFold<tensor::DimOp>(loc, packOp.getDest(), dim));
}
continue;
}
// Currently fusing `packOp` as consumer only expects perfect tiling
// scenario because even if without padding semantic, the `packOp` may
// also yield incomplete tiles. E.g. tensor<30xf32> -> tensor<5x6xf32>,
// where the `tileSize` from operand of `packOp` is 5, which is not
// exactly divided by `innerTile`(=6) of `packOp`. As the result:
// 1. the first slice is extracted from (0) to (4) and inserted into
// (0,0)~(0,4) at first row.
// 2. the second slice is extracted from (5) to (9) and SHOULD BE
// respectively inserted into two rows with different length, including
// first row: (0,5) and second row (1,0)~(1,3). It is hard to coordinate
// them, thus adding below constraint to bypass them temporarily. In
// another word, we can only support tiling with consumer if the tile
// size for the producer is a multiple of the inner tile size for the
// packed dimensions at this moment.
// The caller may assert how this packed dimension's loop tile size
// relates to the inner tile size via `innerTileAlignments` (see
// InnerTileAlignment). The hint is the source of truth and is honored
// when present. When both sizes are also statically known we assert the
// hint agrees with them (a contradicting hint is a caller bug). When
// the hint is `Unknown`, fall back to requiring a statically-provable
// multiple.
bool assumeInnerTileSizesMatchTiles =
innerTileAlignment == InnerTileAlignment::Equal;
bool staticallyDecidable =
!failed(cstTileSize) && cstInnerSize.has_value();
if (innerTileAlignment == InnerTileAlignment::Unknown) {
if (!staticallyDecidable || *cstTileSize % *cstInnerSize != 0)
return failure();
} else if (staticallyDecidable) {
assert(*cstTileSize % *cstInnerSize == 0 &&
"InnerTileAlignment hint contradicts statically known tile "
"sizes");
assert((innerTileAlignment != InnerTileAlignment::Equal ||
*cstTileSize == *cstInnerSize) &&
"InnerTileAlignment::Equal contradicts statically known tile "
"sizes");
}
using AV = affine::AffineValueExpr;
affine::AffineBuilder ab(b, loc);
AffineExpr dim0, sym;
bindDims(b.getContext(), dim0);
bindSymbols(b.getContext(), sym);
auto avOffset = AV(dim0).bind(offsets[dim]);
auto avSize = AV(dim0).bind(sizes[dim]);
auto avTileSize = AV(sym).bind(dimAndTileMapping[dim]);
outerDimOffsets.push_back(ab.floor(avOffset, avTileSize));
// If the tile size equals the inner tile size, the outer dims are
// always 1.
outerDimSizes.push_back(assumeInnerTileSizesMatchTiles
? b.getIndexAttr(1)
: ab.ceil(avSize, avTileSize));
} else {
outerDimOffsets.push_back(offsets[dim]);
outerDimSizes.push_back(sizes[dim]);
}
}
applyPermToRange(outerDimOffsets, outerDimSizes, packOp.getOuterDimsPerm());
resultOffsets = outerDimOffsets;
resultSizes = outerDimSizes;
return success();
}
FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
return getTiledImplementationFromOperandTiles(op, b, operandNumbers,
allOffsets, allSizes,
/*innerTileAlignments=*/{});
}
/// Method to return the tiled implementation of linalg.pack as a consumer.
FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes,
ArrayRef<InnerTileAlignment> innerTileAlignments) const {
if (operandNumbers.size() != 1 || operandNumbers[0] != 0) {
LLVM_DEBUG({ llvm::dbgs() << "unhandled operands for consumer fusion"; });
return failure();
}
ArrayRef<OpFoldResult> offsets(allOffsets[0]);
ArrayRef<OpFoldResult> sizes(allSizes[0]);
auto packOp = cast<PackOp>(op);
// TODO: Support Memref UnPackOp. Temporarily return failure.
if (!packOp.hasPureTensorSemantics())
return failure();
Location loc = packOp.getLoc();
int64_t inputRank = packOp.getSourceRank();
auto oneAttr = b.getI64IntegerAttr(1);
SmallVector<OpFoldResult> strides(inputRank, oneAttr);
SmallVector<Value> tiledOperands;
auto sourceSlice = tensor::ExtractSliceOp::create(
b, loc, packOp.getSource(), offsets, sizes, strides);
tiledOperands.push_back(sourceSlice);
SmallVector<OpFoldResult> outerDimOffsets, outerDimSizes;
if (failed(getIterationDomainTileFromOperandTiles(
op, b, operandNumbers, allOffsets, allSizes, outerDimOffsets,
outerDimSizes, innerTileAlignments)))
return failure();
SmallVector<OpFoldResult> outputOffsets, outputSizes;
if (failed(getResultTilePosition(op, b, 0, outerDimOffsets, outerDimSizes,
outputOffsets, outputSizes)))
return failure();
strides.append(packOp.getDestRank() - inputRank, oneAttr);
auto outSlice = tensor::ExtractSliceOp::create(
b, loc, packOp.getDest(), outputOffsets, outputSizes, strides);
tiledOperands.push_back(outSlice);
if (auto val = packOp.getPaddingValue())
tiledOperands.push_back(val);
for (auto tile : packOp.getInnerTiles())
tiledOperands.push_back(tile);
PackOp tiledPackOp =
PackOp::create(b, loc, TypeRange{outSlice.getType()}, tiledOperands,
packOp.getProperties(),
packOp->getDiscardableAttrDictionary().getValue());
return TilingResult{
{tiledPackOp},
SmallVector<Value>(tiledPackOp->getResults()),
llvm::to_vector(ArrayRef<Operation *>{sourceSlice, outSlice})};
}
};
struct UnpackTileDimInfo {
bool isAlignedToInnerTileSize;
OpFoldResult sourceOffset;
OpFoldResult sourceSize;
OpFoldResult resultOffset;
OpFoldResult destExpandedSize;
};
/// Returns the needed information for tiling unpack op on `tileDim` with given
/// `tileOffset` and `tileSize`. For more details, see the comment of the
/// `getTiledImplementation`.
static UnpackTileDimInfo
getUnpackTileDimInfo(OpBuilder &b, UnPackOp unpackOp, int64_t tileDim,
OpFoldResult tileOffset, OpFoldResult tileSize,
InnerTileAlignment innerTileAlignment) {
UnpackTileDimInfo info;
Attribute zeroAttr = b.getIndexAttr(0);
Attribute oneAttr = b.getIndexAttr(1);
DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
unpackOp.getDimAndTileMapping();
// The dimension is not one of packed data dimension.
if (!dimAndTileMapping.count(tileDim)) {
info.isAlignedToInnerTileSize = true;
info.sourceOffset = tileOffset;
info.sourceSize = tileSize;
info.resultOffset = zeroAttr;
info.destExpandedSize = tileSize;
return info;
}
Location loc = unpackOp.getLoc();
using AV = affine::AffineValueExpr;
affine::AffineBuilder ab(b, loc);
AffineExpr dim0, dim1, sym0;
bindDims(b.getContext(), dim0, dim1);
bindSymbols(b.getContext(), sym0);
OpFoldResult innerTileSize = dimAndTileMapping[tileDim];
info.isAlignedToInnerTileSize = false;
FailureOr<int64_t> cstSize = ValueBoundsConstraintSet::computeConstantBound(
presburger::BoundType::UB, tileSize,
/*stopCondition=*/nullptr, ValueBoundsOptions{/*closedUB=*/true});
std::optional<int64_t> cstInnerSize = getConstantIntValue(innerTileSize);
// The caller may assert how this dimension's loop tile size relates to the
// op's inner tile size via `innerTileAlignment` (see InnerTileAlignment). The
// hint is the source of truth and is honored when present: `Equal`/`Multiple`
// both mean the tile is aligned to (a multiple of) the inner tile, and
// `Equal` additionally collapses the source slice to a single inner tile.
// When both sizes are also statically known we assert the hint agrees with
// them (a contradicting hint is a caller bug). When `Unknown`, fall back to
// the static upper-bound path below.
bool assumeInnerTileSizesMatchTiles =
innerTileAlignment == InnerTileAlignment::Equal;
bool staticallyDecidable = !failed(cstSize) && cstInnerSize.has_value();
if (innerTileAlignment != InnerTileAlignment::Unknown) {
info.isAlignedToInnerTileSize = true;
if (staticallyDecidable) {
assert(*cstSize % *cstInnerSize == 0 &&
"InnerTileAlignment hint contradicts statically known tile sizes");
assert((innerTileAlignment != InnerTileAlignment::Equal ||
*cstSize == *cstInnerSize) &&
"InnerTileAlignment::Equal contradicts statically known tile "
"sizes");
}
}
if (info.isAlignedToInnerTileSize || (!failed(cstSize) && cstInnerSize)) {
if (!info.isAlignedToInnerTileSize && *cstSize % *cstInnerSize == 0)
info.isAlignedToInnerTileSize = true;
// If the tiling size equals to the inner tiling size, the outer dims are
// always 1.
if (assumeInnerTileSizesMatchTiles ||
(cstInnerSize && !failed(cstSize) && *cstInnerSize == *cstSize)) {
auto lhs = AV(dim0).bind(tileOffset);
auto rhs = AV(dim1).bind(innerTileSize);
info.sourceOffset = ab.floor(lhs, rhs);
info.sourceSize = oneAttr;
info.resultOffset = zeroAttr;
info.destExpandedSize = tileSize;
return info;
}
}
if (info.isAlignedToInnerTileSize) {
info.sourceOffset =
ab.floor(AV(dim0).bind(tileOffset), AV(dim1).bind(innerTileSize));
info.resultOffset = zeroAttr;
info.destExpandedSize = tileSize;
// The ceilDiv is needed here because there could be incomplete tile even
// it is perfect tiling cases. E.g.,
// %0 = unpack tensor<33x2xf32> into tensor<64xf32>
// If the tiling size is 32, there will be 3 tiles. Two of them have
// size=32; one of them have size=2. The size is represented using
// affine_min op; we need ceilDiv.
info.sourceSize =
ab.ceil(AV(dim0).bind(tileSize), AV(dim1).bind(innerTileSize));
return info;
}
affine::DivModValue firstCoord = affine::getDivMod(
b, loc, getValueOrCreateConstantIndexOp(b, loc, tileOffset),
getValueOrCreateConstantIndexOp(b, loc, innerTileSize));
OpFoldResult tileExclusiveBound =
ab.add(AV(dim0).bind(tileOffset), AV(dim1).bind(tileSize));
affine::DivModValue lastCoord = affine::getDivMod(
b, loc,
getValueOrCreateConstantIndexOp(
b, loc,
ab.sub(AV(dim0).bind(tileExclusiveBound), AV(dim1).bind(oneAttr))),
getValueOrCreateConstantIndexOp(b, loc, innerTileSize));
OpFoldResult lengthMinusOne = ab.sub(AV(dim0).bind(lastCoord.quotient),
AV(dim1).bind(firstCoord.quotient));
info.sourceSize =
ab.add(AV(dim0).bind(lengthMinusOne), AV(dim1).bind(oneAttr));
info.sourceOffset = firstCoord.quotient;
info.resultOffset = firstCoord.remainder;
// Do not create an Affine ops for expanded size because the affine op is too
// complicated which would trigger an issue in affine ops simplification.
info.destExpandedSize = b.createOrFold<arith::MulIOp>(
loc, getValueOrCreateConstantIndexOp(b, loc, info.sourceSize),
getValueOrCreateConstantIndexOp(b, loc, innerTileSize));
return info;
}
struct UnPackOpTiling
: public TilingInterface::ExternalModel<UnPackOpTiling, linalg::UnPackOp> {
using Base = TilingInterface::ExternalModel<UnPackOpTiling, linalg::UnPackOp>;
using Base::getIterationDomainTileFromOperandTiles;
SmallVector<utils::IteratorType> getLoopIteratorTypes(Operation *op) const {
auto unpackOp = cast<UnPackOp>(op);
SmallVector<utils::IteratorType> iteratorTypes(
unpackOp.getDestRank(), utils::IteratorType::parallel);
return iteratorTypes;
}
SmallVector<Range> getIterationDomain(Operation *op, OpBuilder &b) const {
return getPackUnPackIterationDomain<UnPackOp>(cast<UnPackOp>(op), b);
}
/// There are two cases in tiling unpack ops. If the tiling size is aligned to
/// the inner tile size, the corresponding tiles of source are all complete.
/// Otherwise, there are in-complete tiles. We will need to expand the slice
/// of source for getting complete tiles. The tiled unpack op unpacks more
/// data from source, so We'll need an extract_slice op to shift and truncate
/// the output.
/// Take Nn_to_N as an example. Say that N=32, n=8, and tiling_size=15. The
/// coordinates of second tile (i.e., result[15..31]) are
/// [(1, 7), (2, 0,), (2, 1) ... (3, 6), (3, 7)]. The first row and the last
/// row are incomplete tiles. To represent the unpack op, we have to complete
/// the rows. I.e., the input coordinates would start with (1, 0); end with
/// (3, 7). In this context, the tiled unpack produces a (3 * n) elements
/// because there are 3 rows in total. Follow by a tensor.extract_slice op, we
/// can get the actual result.
FailureOr<TilingResult>
getTiledImplementation(Operation *op, OpBuilder &b,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) const {
return getTiledImplementation(op, b, offsets, sizes,
/*innerTileAlignments=*/{});
}
FailureOr<TilingResult> getTiledImplementation(
Operation *op, OpBuilder &b, ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes,
ArrayRef<InnerTileAlignment> innerTileAlignments) const {
auto unpackOp = cast<UnPackOp>(op);
// TODO: Support Memref UnPackOp. Temporarily return failure.
if (!unpackOp.hasPureTensorSemantics())
return failure();
int64_t srcRank = unpackOp.getSourceRank();
int64_t destRank = unpackOp.getDestRank();
int64_t numInnerTiles = srcRank - destRank;
Location loc = unpackOp.getLoc();
// The perfect tiling case indicates that the tiling sizes are multiple of
// inner_tile_size. In this context, no extra data is needed when
// representing the tiled unpack op.
bool isPerfectTilingCase = true;
Attribute oneAttr = b.getIndexAttr(1);
SmallVector<OpFoldResult> sliceSrcStrides(destRank, oneAttr);
SmallVector<OpFoldResult> sliceSrcIndices, sliceSrcSizes;
SmallVector<OpFoldResult> destExpandedSizes, resultOffsetsFromDest;
for (auto dim : llvm::seq<int64_t>(0, destRank)) {
UnpackTileDimInfo info = getUnpackTileDimInfo(
b, unpackOp, dim, offsets[dim], sizes[dim],
dim < static_cast<int64_t>(innerTileAlignments.size())
? innerTileAlignments[dim]
: InnerTileAlignment::Unknown);
if (!info.isAlignedToInnerTileSize)
isPerfectTilingCase = false;
sliceSrcIndices.push_back(info.sourceOffset);
sliceSrcSizes.push_back(info.sourceSize);
destExpandedSizes.push_back(info.destExpandedSize);
resultOffsetsFromDest.push_back(info.resultOffset);
}
// The tiling is applied on destination dimensions. We have to apply the
// interchange on source dimensions if outer_dims_perm is set.
applyPermToRange(sliceSrcIndices, sliceSrcSizes,
unpackOp.getOuterDimsPerm());
Attribute zeroAttr = b.getIndexAttr(0);
sliceSrcIndices.append(numInnerTiles, zeroAttr);
sliceSrcSizes.append(unpackOp.getMixedTiles());
sliceSrcStrides.append(numInnerTiles, oneAttr);
SmallVector<Operation *> generatedSlices;
tensor::ExtractSliceOp sliceSource = tensor::ExtractSliceOp::create(
b, loc, unpackOp.getSource(), sliceSrcIndices, sliceSrcSizes,
sliceSrcStrides);
generatedSlices.push_back(sliceSource);
SmallVector<OpFoldResult> destStrides(destRank, oneAttr);
Value sliceDest;
if (isPerfectTilingCase) {
auto destSliceOp = tensor::ExtractSliceOp::create(
b, loc, unpackOp.getDest(), offsets, sizes, destStrides);
sliceDest = destSliceOp;
generatedSlices.push_back(destSliceOp);
} else {
sliceDest = tensor::EmptyOp::create(
b, loc, destExpandedSizes, unpackOp.getDestType().getElementType());
}
SmallVector<Value> tiledOperands = {sliceSource.getResult(), sliceDest};
for (auto tile : unpackOp.getInnerTiles())
tiledOperands.push_back(tile);
UnPackOp tiledUnpackOp =
UnPackOp::create(b, loc, TypeRange{sliceDest.getType()}, tiledOperands,
unpackOp.getProperties(),
unpackOp->getDiscardableAttrDictionary().getValue());
if (isPerfectTilingCase)
return TilingResult{{tiledUnpackOp},
SmallVector<Value>(tiledUnpackOp->getResults()),
generatedSlices};
auto extractSlice = tensor::ExtractSliceOp::create(
b, loc, tiledUnpackOp->getResult(0), resultOffsetsFromDest, sizes,
destStrides);
return TilingResult{
{tiledUnpackOp}, {extractSlice.getResult()}, generatedSlices};
}
LogicalResult
getResultTilePosition(Operation *op, OpBuilder &b, unsigned resultNumber,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes,
SmallVector<OpFoldResult> &resultOffsets,
SmallVector<OpFoldResult> &resultSizes) const {
resultOffsets = llvm::to_vector(offsets);
resultSizes = llvm::to_vector(sizes);
return success();
}
FailureOr<TilingResult>
generateResultTileValue(Operation *op, OpBuilder &b, unsigned resultNumber,
ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes) const {
return generateResultTileValue(op, b, resultNumber, offsets, sizes,
/*innerTileAlignments=*/{});
}
FailureOr<TilingResult> generateResultTileValue(
Operation *op, OpBuilder &b, unsigned resultNumber,
ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
ArrayRef<InnerTileAlignment> innerTileAlignments) const {
FailureOr<TilingResult> tilingResult =
getTiledImplementation(op, b, offsets, sizes, innerTileAlignments);
if (failed(tilingResult))
return failure();
return tilingResult.value();
}
LogicalResult generateScalarImplementation(Operation *op, OpBuilder &builder,
Location loc,
ValueRange ivs) const {
auto unpackOp = cast<UnPackOp>(op);
assert(unpackOp.hasPureBufferSemantics() &&
"expected operation to have buffer semantics");
assert(ivs.size() == unpackOp.getDestRank() &&
"number of ivs must match the rank of the output tensor");
OpBuilder::InsertionGuard g(builder);
DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
unpackOp.getDimAndTileMapping();
// Untiled loops and tile loops induction variables.
SmallVector<Value> inputIvs;
// Point loops induction variables.
SmallVector<Value> inputIvsPointLoops;
inputIvs.reserve(unpackOp.getDestRank());
inputIvsPointLoops.reserve(dimAndTileMapping.size());
for (auto dim : llvm::seq<int64_t>(0, unpackOp.getDestRank())) {
if (dimAndTileMapping.count(dim)) {
affine::DivModValue divMod =
affine::getDivMod(builder, loc, ivs[dim],
getValueOrCreateConstantIndexOp(
builder, loc, dimAndTileMapping[dim]));
inputIvsPointLoops.push_back(divMod.remainder);
inputIvs.push_back(divMod.quotient);
} else {
inputIvs.push_back(ivs[dim]);
}
}
// TODO: (lorenzo) simplify the logic a bit. There is `ivs`,
// `inputIvsPointLoops` and `inputIvs`.
assert(inputIvsPointLoops.size() + inputIvs.size() ==
unpackOp.getSourceRank() &&
"expect same number of induction variables equals to input rank");
// Interchange the point loops induction variables based on `inner_dim_pos`.
ArrayRef<int64_t> innerDims = unpackOp.getInnerDimsPos();
SmallVector<int64_t> interchangeVector =
computeInterchangeFromDimPos(innerDims, unpackOp.getDestRank());
SmallVector<Value> interchangedInputIvsPointLoops = inputIvsPointLoops;
interchangedInputIvsPointLoops = interchange<Value>(
interchangedInputIvsPointLoops, interchangeVector, /*offset=*/0);
// Interchange the tiled loops induction variables based on
// `outer_dims_perm`.
ArrayRef<int64_t> outerDims = unpackOp.getOuterDimsPerm();
if (!outerDims.empty())
inputIvs = interchange<Value>(inputIvs, outerDims, /*offset=*/0);
llvm::append_range(inputIvs, interchangedInputIvsPointLoops);
Value scalar =
memref::LoadOp::create(builder, loc, unpackOp.getSource(), inputIvs);
memref::StoreOp::create(builder, loc, scalar, unpackOp.getDest(), ivs);
return success();
}
/// Method to return the position of iteration domain tile computed by the
/// tiled operation.
LogicalResult getIterationDomainTileFromOperandTiles(
Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes,
SmallVectorImpl<OpFoldResult> &resultOffsets,
SmallVectorImpl<OpFoldResult> &resultSizes) const {
if (operandNumbers.size() != 1) {
LLVM_DEBUG({ llvm::dbgs() << "unable to handle multiple operands"; });
return failure();
}
auto unPackOp = cast<UnPackOp>(op);
unsigned operandNumber = operandNumbers[0];
ArrayRef<OpFoldResult> offsets(allOffsets[0]);
ArrayRef<OpFoldResult> sizes(allSizes[0]);
// If the operand tile is the dest, then no adjustment is needed.
if (operandNumber == unPackOp.getDestMutable().getOperandNumber()) {
resultOffsets = llvm::to_vector(offsets);
resultSizes = llvm::to_vector(sizes);
return success();
}
Location loc = unPackOp.getLoc();
int64_t numTiles = unPackOp.getInnerDimsPos().size();
auto destOffsets = offsets.drop_back(numTiles);
auto destSizes = sizes.drop_back(numTiles);
// The tiling is applied on interchanged dimensions. We have to undo the
// interchange to map sizes and offsets to the original input.
int64_t outputRank = unPackOp.getDestRank();
ReifiedRankedShapedTypeDims reifiedReturnShapes;
if (failed(reifyResultShapes(b, unPackOp, reifiedReturnShapes)))
return failure();
SmallVector<OpFoldResult> outputMixedSizes = reifiedReturnShapes.front();
SmallVector<OpFoldResult> origOffsets(destOffsets);
SmallVector<OpFoldResult> origSizes(destSizes);
applyPermToRange(origOffsets, origSizes,
invertPermutationVector(unPackOp.getOuterDimsPerm()));
DenseMap<int64_t, OpFoldResult> dimAndTileMapping =
unPackOp.getDimAndTileMapping();
for (auto dim : llvm::seq<int64_t>(0, outputRank)) {
using AV = affine::AffineValueExpr;
affine::AffineBuilder ab(b, loc);
AffineExpr dim0, dim1, sym0;
bindDims(b.getContext(), dim0, dim1);
bindSymbols(b.getContext(), sym0);
if (dimAndTileMapping.count(dim)) {
// If the data dimension is tiled, the i-th index is the product of
// offset_i and tile_i, and the i-th size is the product of sizes_i and
// tile_i. The sizes must be clamped to the sizes of the unpack result.
auto avOffset = AV(dim0).bind(origOffsets[dim]);
auto avSize = AV(dim0).bind(origSizes[dim]);
auto avTileSize = AV(sym0).bind(dimAndTileMapping[dim]);
auto avResultSize = AV(dim0).bind(outputMixedSizes[dim]);
resultOffsets.push_back(ab.mul(avOffset, avTileSize));
auto avResultOffset = AV(dim1).bind(resultOffsets.back());
resultSizes.push_back(ab.min({ab.mul(avSize, avTileSize),
ab.sub(avResultSize, avResultOffset)}));
} else {
resultOffsets.push_back(origOffsets[dim]);
resultSizes.push_back(origSizes[dim]);
}
}
return success();
}
FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes) const {
return getTiledImplementationFromOperandTiles(op, b, operandNumbers,
allOffsets, allSizes,
/*innerTileAlignments=*/{});
}
/// Method to return the tiled implementation of linalg.unpack as a consumer.
FailureOr<TilingResult> getTiledImplementationFromOperandTiles(
Operation *op, OpBuilder &b, ArrayRef<unsigned> operandNumbers,
ArrayRef<SmallVector<OpFoldResult>> allOffsets,
ArrayRef<SmallVector<OpFoldResult>> allSizes,
ArrayRef<InnerTileAlignment> innerTileAlignments) const {
if (operandNumbers.size() != 1 || operandNumbers[0] != 0) {
LLVM_DEBUG({ llvm::dbgs() << "unhandled operands for consumer fusion"; });
return failure();
}
auto unPackOp = cast<UnPackOp>(op);
// TODO: Support Memref UnPackOp. Temporarily return failure.
if (!unPackOp.hasPureTensorSemantics())
return failure();
ArrayRef<OpFoldResult> offsets(allOffsets[0]);
ArrayRef<OpFoldResult> sizes(allSizes[0]);
// linalg.unpack op is fusible (as a consumer) only if the inner dims are
// not tiled, i.e. each inner-dim loop tile size equals the inner tile size.
// The caller may assert this per inner dim via InnerTileAlignment::Equal;
// otherwise we require a statically-provable equality.
int64_t numTiles = unPackOp.getInnerDimsPos().size();
ArrayRef<int64_t> innerDimsPos = unPackOp.getInnerDimsPos();
SmallVector<OpFoldResult> mixedTiles = unPackOp.getMixedTiles();
ArrayRef<OpFoldResult> innerSizes = sizes.take_back(numTiles);
for (int64_t i = 0; i < numTiles; ++i) {
// `innerTileAlignments` is indexed by the unpack iteration domain (the
// dest dims); the i-th inner tile lives on dest dim `innerDimsPos[i]`.
int64_t destDim = innerDimsPos[i];
bool hintedEqual =
destDim < static_cast<int64_t>(innerTileAlignments.size()) &&
innerTileAlignments[destDim] == InnerTileAlignment::Equal;
// The hint is the source of truth: honor a caller `Equal` assertion. When
// both sizes are also statically known, assert the hint agrees with them
// (a contradicting hint is a caller bug) rather than silently ignoring
// it. Without an `Equal` hint, require a statically-provable equality
// (the inner dim must not be tiled).
if (hintedEqual) {
assert((!getConstantIntValue(mixedTiles[i]) ||
!getConstantIntValue(innerSizes[i]) ||
isEqualConstantIntOrValue(mixedTiles[i], innerSizes[i])) &&
"InnerTileAlignment::Equal contradicts statically known tile "
"sizes");
continue;
}
if (isEqualConstantIntOrValue(mixedTiles[i], innerSizes[i]))
continue;
return failure();
}
Location loc = unPackOp.getLoc();
// Fetch offset/size for creating the slice of the dest operand of
// unpack op.
SmallVector<OpFoldResult> outputOffsets, outputSizes;
if (failed(getIterationDomainTileFromOperandTiles(
op, b, operandNumbers, allOffsets, allSizes, outputOffsets,
outputSizes)))
return failure();
auto oneAttr = b.getI64IntegerAttr(1);
int64_t outputRank = unPackOp.getDestRank();
SmallVector<OpFoldResult> strides(outputRank, oneAttr);
SmallVector<Value> tiledOperands;
// Create slice of the dest operand.
auto extractDestSlice = tensor::ExtractSliceOp::create(
b, loc, unPackOp.getDest(), outputOffsets, outputSizes, strides);
tiledOperands.push_back(extractDestSlice);
strides.append(unPackOp.getSourceRank() - outputRank, oneAttr);
// Create slice of the source operand.
auto extractSourceSlice = tensor::ExtractSliceOp::create(
b, loc, unPackOp.getSource(), offsets, sizes, strides);
tiledOperands.insert(tiledOperands.begin(), extractSourceSlice);
for (auto tile : unPackOp.getInnerTiles())
tiledOperands.push_back(tile);
// Create tiled unpack op.
UnPackOp tiledUnPackOp =
UnPackOp::create(b, loc, TypeRange{extractDestSlice.getType()},
tiledOperands, unPackOp.getProperties(),
unPackOp->getDiscardableAttrDictionary().getValue());
return TilingResult{{tiledUnPackOp},
SmallVector<Value>(tiledUnPackOp->getResults()),
llvm::to_vector(ArrayRef<Operation *>{
extractSourceSlice, extractDestSlice})};
}
};
} // namespace
template <typename OpType>
static void registerOne(MLIRContext *ctx) {
OpType::template attachInterface<LinalgOpTilingInterfaceModel<OpType>>(*ctx);
OpType::template attachInterface<
LinalgOpPartialReductionInterfaceModel<OpType>>(*ctx);
}
/// Variadic helper function.
template <typename... OpTypes>
static void registerAll(MLIRContext *ctx) {
(registerOne<OpTypes>(ctx), ...);
}
#define GET_OP_LIST
void mlir::linalg::registerTilingInterfaceExternalModels(
DialectRegistry &registry) {
registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
registerOne<linalg::GenericOp>(ctx);
linalg::PackOp::attachInterface<PackOpTiling>(*ctx);
linalg::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);
registerAll<
#include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
>(ctx);
});
}
void mlir::linalg::registerTilingInterfaceExternalModelsForPackUnPackOps(
DialectRegistry &registry) {
registry.addExtension(+[](MLIRContext *ctx, LinalgDialect *dialect) {
linalg::PackOp::attachInterface<PackOpTiling>(*ctx);
linalg::UnPackOp::attachInterface<UnPackOpTiling>(*ctx);
});
}