Google Git Google Git
Sign in
llvm/llvm-project/3104f29fa41ff4bb932bc1bf48ca032f1e9dcafa/./mlir/lib/Dialect/SCF/Transforms/TileUsingInterface.cpp
blob: a8efe6cdb4415e79f79369cf5efbdf9f4ff16407 [file]
//===- Tiling.cpp - Implementation of tiling using 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
//
//===----------------------------------------------------------------------===//
//
// This file implements the tiling using TilingInterface.
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/SCF/Transforms/TileUsingInterface.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Analysis/TopologicalSortUtils.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/SCF/Utils/Utils.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/IR/Dominance.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/DestinationStyleOpInterface.h"
#include "mlir/Interfaces/TilingInterface.h"
#include "mlir/Rewrite/FrozenRewritePatternSet.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Debug.h"
#include <optional>
#define DEBUG_TYPE "tile-using-interface"
using namespace mlir;
scf::SCFTilingOptions &
scf::SCFTilingOptions::setTileSizes(ArrayRef<OpFoldResult> ts) {
assert(!tileSizeComputationFunction && "tile sizes already set");
auto tileSizes = llvm::to_vector(ts);
tileSizeComputationFunction = [tileSizes](OpBuilder &b, Operation *op) {
return tileSizes;
};
return *this;
}
scf::SCFTilingOptions &
scf::SCFTilingOptions::setNumThreads(ArrayRef<OpFoldResult> nt) {
assert(!numThreadsComputationFunction && "num tiles already set");
auto numThreads = llvm::to_vector(nt);
numThreadsComputationFunction = [numThreads](OpBuilder &b, Operation *op) {
return numThreads;
};
return *this;
}
/// Helper method to adjust the interchange vector to match the iteration
/// domain.
static SmallVector<int64_t>
fillInterchangeVector(ArrayRef<int64_t> interchangeVector,
size_t iterationDomainSize) {
SmallVector<int64_t> filledVector = llvm::to_vector(interchangeVector);
if (filledVector.size() < iterationDomainSize) {
auto range = llvm::seq<int64_t>(filledVector.size(), iterationDomainSize);
filledVector.append(range.begin(), range.end());
}
if (filledVector.size() > iterationDomainSize)
filledVector.resize(iterationDomainSize);
return filledVector;
}
//===----------------------------------------------------------------------===//
// tileUsingSCF implementation.
//===----------------------------------------------------------------------===//
/// Verify the tile size options are set in a consistent manner.
static LogicalResult verifyOptions(RewriterBase &rewriter, Location loc,
const scf::SCFTilingOptions &options) {
// Specifying number of threads is only supported on `scf.forall` op.
if (options.numThreadsComputationFunction &&
options.loopType != scf::SCFTilingOptions::LoopType::ForallOp) {
return rewriter.notifyMatchFailure(
loc, "number of threads can only by specified when loop type is "
"set to use `scf.forall`");
}
// If specified, check that the interchange vector is a permutation.
if (!options.interchangeVector.empty()) {
if (!isPermutationVector(options.interchangeVector)) {
return rewriter.notifyMatchFailure(
loc, "invalid interchange vector, not a permutation of the entire "
"iteration space");
}
}
return success();
}
/// Method to instantiate the tile sizes and/or number of threads specified
/// by the user.
static std::tuple<SmallVector<OpFoldResult>, SmallVector<OpFoldResult>>
getUserTileSizesAndNumThreads(RewriterBase &rewriter, TilingInterface op,
ArrayRef<Range> iterationDomain,
const scf::SCFTilingOptions &options) {
OpFoldResult zero = rewriter.getIndexAttr(0);
SmallVector<OpFoldResult> tileSizes, numThreads;
size_t numLoops = iterationDomain.size();
// Check whether the number of tiles to use is specified.
if (options.numThreadsComputationFunction) {
numThreads = options.numThreadsComputationFunction(rewriter, op);
numThreads.resize(numLoops, zero);
// If the number of tiles is also specified, use that.
if (options.tileSizeComputationFunction) {
tileSizes = options.tileSizeComputationFunction(rewriter, op);
tileSizes.resize(numLoops, zero);
return {tileSizes, numThreads};
}
// Compute the tile sizes from the iteration domain and number
// of tiles as follows
// - niters = ceilDiv(ub - lb, step)
// - tileSize = ceilDiv(niters, numThreads)
AffineExpr s0, s1, s2;
bindSymbols(rewriter.getContext(), s0, s1, s2);
// TODO: The step here is assumed to be 1.
AffineExpr numItersExpr = (s1 - s0);
AffineExpr tileSizeExpr = numItersExpr.ceilDiv(s2);
tileSizes.resize(numLoops, zero);
for (auto [index, range, nt] :
llvm::enumerate(iterationDomain, numThreads)) {
if (isZeroInteger(nt))
continue;
tileSizes[index] = affine::makeComposedFoldedAffineApply(
rewriter, op.getLoc(), tileSizeExpr, {range.offset, range.size, nt});
}
tileSizes.resize(numLoops, zero);
return {tileSizes, numThreads};
}
// Enforce the convention that "tiling by zero"
// skips tiling a particular dimension. This convention is significantly
// simpler to handle instead of adjusting affine maps to account for missing
// dimensions.
assert(options.tileSizeComputationFunction &&
"expected tile sizes to be specified");
tileSizes = options.tileSizeComputationFunction(rewriter, op);
tileSizes.resize(numLoops, zero);
return {tileSizes, numThreads};
}
/// Checks if any of the tiled loops are not parallel.
static LogicalResult checkTileSizes(TilingInterface op,
scf::SCFTilingOptions::LoopType loopType,
ReductionTilingStrategy reductionStrategy,
ArrayRef<OpFoldResult> givenTileSizes,
ArrayRef<OpFoldResult> numThreads) {
auto iterators = op.getLoopIteratorTypes();
assert(iterators.size() == givenTileSizes.size() &&
"expected as many tile size values as number of loops");
assert((numThreads.empty() || (numThreads.size() == iterators.size())) &&
"when specified, expected number of threads to use for each loop");
bool isParallelTiling = false;
for (auto [index, iterator, givenTileSize] :
llvm::enumerate(iterators, givenTileSizes)) {
if (!isConstantIntValue(givenTileSize, 0)) {
isParallelTiling |= iterator == utils::IteratorType::parallel;
}
if (loopType == scf::SCFTilingOptions::LoopType::ForallOp &&
reductionStrategy == ReductionTilingStrategy::FullReduction) {
// If num threads is specified, check that it is greater than one only for
// parallel dimensions.
if (!numThreads.empty()) {
if (std::optional<int64_t> constNumThreads =
getConstantIntValue(numThreads[index])) {
if (constNumThreads.value() > 1 &&
iterator != utils::IteratorType::parallel) {
op.emitWarning() << "tiling is not thread safe at axis #" << index;
}
}
continue;
}
if (std::optional<int64_t> constTileSize =
getConstantIntValue(givenTileSize)) {
if (constTileSize.value() > 0 &&
iterator != utils::IteratorType::parallel) {
op.emitWarning() << "tiling is not thread safe at axis #" << index;
}
}
}
}
if (reductionStrategy != ReductionTilingStrategy::FullReduction) {
if (isParallelTiling) {
return op->emitOpError("tiling parallel dimensions is not supported with "
"partial reduction tiling strategies");
}
}
return success();
}
/// Get the reduction dims that are tiled. This accounts for reduction dims
/// that are specified as tiled, but the tile size is 0.
static SetVector<unsigned>
getSanitizedReductionDims(ArrayRef<OpFoldResult> givenTileSizes,
const scf::SCFTilingOptions &options) {
SetVector<unsigned> reductionDims;
for (auto dim : options.reductionDims) {
if (isConstantIntValue(givenTileSizes[dim], 0))
continue;
reductionDims.insert(dim);
}
return reductionDims;
}
/// Check if `stride` evenly divides the trip count `size - offset`.
static bool tileDividesIterationDomain(Range loopRange) {
std::optional<int64_t> offsetAsInt = getConstantIntValue(loopRange.offset);
if (!offsetAsInt)
return false;
std::optional<int64_t> sizeAsInt = getConstantIntValue(loopRange.size);
if (!sizeAsInt)
return false;
std::optional<int64_t> strideAsInt = getConstantIntValue(loopRange.stride);
if (!strideAsInt)
return false;
return ((sizeAsInt.value() - offsetAsInt.value()) % strideAsInt.value() == 0);
}
/// Returns the bounded tile size given the current `offset`, `loopRange` and
/// `tileSize`, i.e., `min(tileSize, range.end() - offset)`.
static OpFoldResult getBoundedTileSize(OpBuilder &b, Location loc,
Range loopRange, OpFoldResult offset,
OpFoldResult givenTileSize) {
std::optional<int64_t> ts = getConstantIntValue(givenTileSize);
if (ts && ts.value() == 1)
return givenTileSize;
if (tileDividesIterationDomain(
Range{loopRange.offset, loopRange.size, givenTileSize}))
return givenTileSize;
// The tile size to use (to avoid out of bounds access) is minimum of
// `tileSize` and `ub - iv`, where `iv` is the induction variable of the tiled
// loop.
AffineExpr s0, s1, d0;
bindDims(b.getContext(), d0);
bindSymbols(b.getContext(), s0, s1);
AffineMap minMap = AffineMap::get(1, 2, {s0 - d0, s1}, b.getContext());
Value size = getValueOrCreateConstantIndexOp(b, loc, loopRange.size);
return affine::makeComposedFoldedAffineMin(
b, loc, minMap, SmallVector<OpFoldResult>{offset, size, givenTileSize});
}
/// Returns true if the maximum tile offset `tileSize * numThreads-1` is less
/// than `iterationSize`.
static bool canOmitTileOffsetInBoundsCheck(OpFoldResult givenTileSize,
OpFoldResult numThreads,
OpFoldResult iterationSize) {
std::optional<int64_t> tileSizeConst = getConstantIntValue(givenTileSize);
std::optional<int64_t> numThreadsConst = getConstantIntValue(numThreads);
std::optional<int64_t> iterSizeConst = getConstantIntValue(iterationSize);
if (!tileSizeConst || !numThreadsConst || !iterSizeConst)
return false;
return *tileSizeConst * (*numThreadsConst - 1) < *iterSizeConst;
}
/// Compute the `OpFoldResult`s that represents the multi-dimensional
/// `offset`s and `size`s of the tile of the iteration space that the
/// innermost loop body of the generated tiled loops corresponds to.
static std::tuple<SmallVector<OpFoldResult>, SmallVector<OpFoldResult>>
getTileOffsetAndSizes(RewriterBase &rewriter, Location loc, ValueRange ivs,
ArrayRef<Range> iterationDomain,
ArrayRef<OpFoldResult> givenTileSizes) {
SmallVector<OpFoldResult> offsets, sizes;
int materializedLoopNum = 0;
for (auto [givenTileSize, loopRange] :
llvm::zip_equal(givenTileSizes, iterationDomain)) {
// Non-tiled cases, set the offset and size to the
// `loopRange.offset/size`.
if (isZeroInteger(givenTileSize)) {
offsets.push_back(loopRange.offset);
sizes.push_back(loopRange.size);
continue;
}
Value iv = ivs[materializedLoopNum++];
OpFoldResult offset = getAsOpFoldResult(iv);
offsets.push_back(offset);
OpFoldResult size =
getBoundedTileSize(rewriter, loc, loopRange, offset, givenTileSize);
sizes.push_back(size);
}
return {offsets, sizes};
}
/// Function to return the bounds of the loops to be generated.
static std::tuple<SmallVector<OpFoldResult>, SmallVector<OpFoldResult>,
SmallVector<OpFoldResult>>
getLoopBounds(RewriterBase &rewriter, Location loc, ArrayRef<Range> loopRanges,
ArrayRef<OpFoldResult> givenTileSizes) {
SmallVector<OpFoldResult> lbs, ubs, steps;
for (auto [loopRange, givenTileSize] :
llvm::zip_equal(loopRanges, givenTileSizes)) {
// No loop if the tile size is 0.
if (isZeroInteger(givenTileSize))
continue;
lbs.push_back(loopRange.offset);
ubs.push_back(loopRange.size);
steps.push_back(givenTileSize);
}
return {lbs, ubs, steps};
}
/// Typedef for function that allows returning additional yielded values during
/// `yieldTiledValuesAndReplace`.
/// - `ivs` induction variable for the loop.
/// - `newBbArgs` basic block arguments corresponding to newly added iter_args.
/// - `tiledValues` the tiled values to return. Must be of same size as
/// `newbbArgs`, each element of this array is inserted into the corresponding
/// element in `newbbArgs`.
/// - `resultOffsets` is of the same size as `tiledValues` and represents
/// the offsets to use when inserting corresponding element from `tiledValues`
/// into the element from `newBbArgs`.
/// - `resultSizes` is of the same size as `tiledValues` and represents
/// the size of the corresponding element from `tiledValues` inserted into
/// the element from `newBbArgs`.
/// In case the method needs to return `failure()` the method is expected
/// to clean up any inserted operations.
using YieldTiledValuesFn = std::function<LogicalResult(
RewriterBase &rewriter, Location loc, ValueRange ivs, ValueRange newBbArgs,
SmallVector<Value> &tiledValues,
SmallVector<SmallVector<OpFoldResult>> &resultOffsets,
SmallVector<SmallVector<OpFoldResult>> &resultSizes)>;
/// Typedef for function that implements the body of a tiled loop.
/// - `ivs` induction variable for the loop.
/// - `tileOffsets` represents offsets for the tiled iteration space.
/// - `tileSizes` represents the sizes for the tiled iteraiton space.
/// - `outerDestinationTensors` tensor that holds the result. Is same size
/// as the destination operands of the original operations.
/// - `tiledResults` results of the tiled computation, corresponds to
/// tiles of the original operation computed by the loop body.
/// Should be same size as the `destinationTensors`
/// - `resultOffsets` is of the same size as `tiledResults` and represents
/// the offset to use when writing the corresponding element from
/// `tiledResults` into `destinationTensors`.
/// - `resultOffsets` is of the same size as `tiledResults` and represents
/// the size to use when writing the corresponding element from
/// `tiledResults` into `destinationTensors`.
/// In case the method needs to return `failure()` the method is expected
/// to clean up any inserted operations.
using GenerateTiledBodyFn = std::function<LogicalResult(
RewriterBase &rewriter, Location Loc, ValueRange ivs,
ArrayRef<OpFoldResult> tileOffsets, ArrayRef<OpFoldResult> tileSizes,
ValueRange outerDestinationTensors, SmallVector<Value> &tiledResults,
SmallVector<SmallVector<OpFoldResult>> &resultOffsets,
SmallVector<SmallVector<OpFoldResult>> &resultSizes)>;
/// Clones the operation and updates the destination if the operation
/// implements the `DestinationStyleOpInterface`.
static Operation *cloneOpAndUpdateDestinationArgs(RewriterBase &rewriter,
Operation *op,
ValueRange newDestArgs) {
Operation *clonedOp = rewriter.clone(*op);
if (newDestArgs.empty())
return clonedOp;
if (auto destinationStyleOp = dyn_cast<DestinationStyleOpInterface>(clonedOp))
destinationStyleOp.getDpsInitsMutable().assign(newDestArgs);
return clonedOp;
}
/// Generate the tile-loop nest using `scf.for` operation.
/// - `loopRanges` specifies the lb, ub and step of the untiled iteration space.
/// - `givenTileSizes` is the tile sizes to use. Zero represent untiled loops.
/// - `outerDestinationTensors` are the init values to use for the outer most
/// loop.
/// - `tiledBodyFn` is called to generated the loop body of the inner
/// most
/// loop.
/// Returns the generated `scf.for` loops on success.
static FailureOr<SmallVector<LoopLikeOpInterface>> generateLoopNestUsingForOp(
RewriterBase &rewriter, Location loc, ArrayRef<Range> loopRanges,
ArrayRef<OpFoldResult> givenTileSizes, ValueRange outerDestinationTensors,
GenerateTiledBodyFn tiledBodyFn) {
assert(!loopRanges.empty() && "unexpected empty loop ranges");
assert(loopRanges.size() == givenTileSizes.size() &&
"expected as many tile sizes as loop ranges");
OpBuilder::InsertionGuard guard(rewriter);
SmallVector<OpFoldResult> lbs, ubs, steps;
std::tie(lbs, ubs, steps) =
getLoopBounds(rewriter, loc, loopRanges, givenTileSizes);
SmallVector<Value> lbVals =
getValueOrCreateConstantIndexOp(rewriter, loc, lbs);
SmallVector<Value> ubVals =
getValueOrCreateConstantIndexOp(rewriter, loc, ubs);
SmallVector<Value> stepVals =
getValueOrCreateConstantIndexOp(rewriter, loc, steps);
SmallVector<Value> ivs;
SmallVector<LoopLikeOpInterface> loops;
ValueRange innerDestinationTensors(outerDestinationTensors);
for (auto [lb, ub, step] : llvm::zip_equal(lbVals, ubVals, stepVals)) {
auto loop =
scf::ForOp::create(rewriter, loc, lb, ub, step, innerDestinationTensors,
[](OpBuilder &bodyBuilder, Location bodyLoc,
Value iv, ValueRange /*iterArgs*/) {});
loops.push_back(loop);
ivs.push_back(loop.getInductionVar());
rewriter.setInsertionPointToEnd(loop.getBody());
innerDestinationTensors = loop.getRegionIterArgs();
}
if (loops.empty())
return success();
// Compute the `offsets` and `sizes` to use for tiling.
SmallVector<OpFoldResult> offsets, sizes;
std::tie(offsets, sizes) =
getTileOffsetAndSizes(rewriter, loc, ivs, loopRanges, givenTileSizes);
SmallVector<Value> tiledResults;
SmallVector<SmallVector<OpFoldResult>> resultOffsets, resultSizes;
if (failed(tiledBodyFn(rewriter, loc, ivs, offsets, sizes,
innerDestinationTensors, tiledResults, resultOffsets,
resultSizes))) {
return rewriter.notifyMatchFailure(
loc, "failed to generate inner tile loop body");
}
if (loops.empty())
return loops;
assert(tiledResults.size() == innerDestinationTensors.size() &&
"Number of results of body should be equal to number of iter args");
// 6. Yield all the results of the tiled operation.
SmallVector<Value> yieldedValues;
for (auto [tiledValue, destinationTensor, resultOffset, resultSize] :
llvm::zip_equal(tiledResults, innerDestinationTensors, resultOffsets,
resultSizes)) {
SmallVector<OpFoldResult> resultStride(resultOffset.size(),
rewriter.getIndexAttr(1));
auto insertSlice = tensor::InsertSliceOp::create(
rewriter, loc, tiledValue, destinationTensor, resultOffset, resultSize,
resultStride);
yieldedValues.push_back(insertSlice);
}
scf::YieldOp::create(rewriter, loc, yieldedValues);
// Add the scf.yield operations for all the outer loops.
for (auto [outerLoop, innerLoop] :
llvm::zip_equal(MutableArrayRef(loops).drop_back(),
MutableArrayRef(loops).drop_front())) {
rewriter.setInsertionPointToEnd(
cast<scf::ForOp>(outerLoop.getOperation()).getBody());
scf::YieldOp::create(rewriter, outerLoop.getLoc(), innerLoop->getResults());
}
return loops;
}
/// Compute the `OpFoldResult`s that represents the multi-dimensional
/// `offset`s and `size`s of the tile of the iteration space that the
/// innermost loop body of the generated tiled loops corresponds to
/// when tiling using `forall` op. This is handle separately due to
/// the special case handling needed for when the tiling is done by
/// specifying number of threads.
static std::tuple<SmallVector<OpFoldResult>, SmallVector<OpFoldResult>>
getTileOffsetAndSizesWithForAllOp(RewriterBase &rewriter, Location loc,
ValueRange ivs,
ArrayRef<Range> iterationDomain,
ArrayRef<OpFoldResult> givenTileSizes,
ArrayRef<OpFoldResult> numThreads) {
if (numThreads.empty()) {
return getTileOffsetAndSizes(rewriter, loc, ivs, iterationDomain,
givenTileSizes);
}
SmallVector<OpFoldResult> offsets, sizes;
int materializedLoopNum = 0;
AffineExpr d0, d1, s0, s1;
AffineExpr offsetExpr, residualTileSizeExpr;
bindDims(rewriter.getContext(), d0, d1);
bindSymbols(rewriter.getContext(), s0, s1);
offsetExpr = d0 + d1 * s0;
residualTileSizeExpr = s1 - (d0 + d1 * s0);
for (auto [index, nt, givenTileSize, loopRange] :
llvm::enumerate(numThreads, givenTileSizes, iterationDomain)) {
// Non-tiled cases, set the offset and size to the
// `loopRange.offset/size`.
if (isZeroInteger(nt)) {
offsets.push_back(loopRange.offset);
sizes.push_back(loopRange.size);
continue;
}
Value iv = ivs[materializedLoopNum++];
OpFoldResult offset = affine::makeComposedFoldedAffineApply(
rewriter, loc, offsetExpr,
ArrayRef<OpFoldResult>{loopRange.offset, iv, givenTileSize});
OpFoldResult residualTileSize = affine::makeComposedFoldedAffineApply(
rewriter, loc, residualTileSizeExpr,
{loopRange.offset, nt, givenTileSize, loopRange.size});
OpFoldResult size = givenTileSize;
if (!isZeroInteger(residualTileSize)) {
OpFoldResult sizeMinusOffsetPerThread =
affine::makeComposedFoldedAffineApply(rewriter, loc, s0 - d0,
{offset, loopRange.size});
size = affine::makeComposedFoldedAffineMin(
rewriter, loc,
AffineMap::getMultiDimIdentityMap(2, rewriter.getContext()),
{sizeMinusOffsetPerThread, givenTileSize});
}
// Consider the case where the original loop was `[0, 100)`.
// If number of threads are `7`, the tile size would be computed as
// `ceilDiv(100, 7) = 15`. For the last thread (thread_id = 6)
// - `offset = 0 + 6 * 15 = 105`
// - `tileSize = min(15, 100 - 105) = -5`
// To avoid negative tile sizes, we need to do a further
// `nonNegativeTileSize = affine.max(0, tileSize)`.
// This `max` can be avoided if
// `offset + tileSize * (numThreads - 1) < (ub - lb)`
if (!canOmitTileOffsetInBoundsCheck(givenTileSize, nt, loopRange.size)) {
AffineMap maxMap =
AffineMap::getMultiDimIdentityMap(2, rewriter.getContext());
size = affine::makeComposedFoldedAffineMax(
rewriter, loc, maxMap, {rewriter.getIndexAttr(0), size});
}
offsets.push_back(offset);
sizes.push_back(size);
}
return {offsets, sizes};
}
/// Generate the tile-loop nest using `scf.forall` operation.
/// - `loopRanges` specifies the lb, ub and step of the untiled iteration space.
/// - `giventileSizes` is the tile sizes to use. Zero represent untiled loops.
/// - `outerDestinationTensors` are the init values to use for the loop.
/// - `mappingVector` is the mapping attributes to use for loop construction.
/// Can be empty.
/// - `tiledBodyFn` is called to generated the loop body of the inner
/// most
/// loop.
/// Returns the generated `scf.forall` loop on success.
static FailureOr<SmallVector<LoopLikeOpInterface>>
generateLoopNestUsingForallOp(RewriterBase &rewriter, Location loc,
ArrayRef<Range> loopRanges,
ArrayRef<OpFoldResult> givenTileSizes,
ArrayRef<OpFoldResult> numThreads,
ArrayRef<Attribute> mappingVector,
ValueRange outerDestinationTensors,
GenerateTiledBodyFn tiledBodyFn) {
assert(!loopRanges.empty() && "unexpected empty loop ranges");
assert(loopRanges.size() == givenTileSizes.size() &&
"expected as many tile sizes as loop ranges");
OpBuilder::InsertionGuard guard(rewriter);
std::optional<ArrayAttr> mappingAttr;
if (!mappingVector.empty())
mappingAttr = rewriter.getArrayAttr(mappingVector);
scf::ForallOp forallOp;
bool useNumThreads = !numThreads.empty();
SmallVector<LoopLikeOpInterface> loops;
if (useNumThreads) {
// Prune the zero numthreads.
SmallVector<OpFoldResult> nonZeroNumThreads;
for (auto nt : numThreads) {
if (isZeroInteger(nt))
continue;
nonZeroNumThreads.push_back(nt);
}
forallOp = scf::ForallOp::create(rewriter, loc, nonZeroNumThreads,
outerDestinationTensors, mappingAttr);
} else {
SmallVector<OpFoldResult> lbs, ubs, steps;
std::tie(lbs, ubs, steps) =
getLoopBounds(rewriter, loc, loopRanges, givenTileSizes);
forallOp = scf::ForallOp::create(rewriter, loc, lbs, ubs, steps,
outerDestinationTensors, mappingAttr);
}
loops.push_back(forallOp);
rewriter.setInsertionPoint(forallOp.getTerminator());
ValueRange innerDestinationTensors = forallOp.getRegionOutArgs();
SmallVector<Value> ivs = forallOp.getInductionVars();
// Compute the `offsets` and `sizes` to use for tiling.
SmallVector<OpFoldResult> offsets, sizes;
std::tie(offsets, sizes) = getTileOffsetAndSizesWithForAllOp(
rewriter, loc, ivs, loopRanges, givenTileSizes, numThreads);
SmallVector<Value> tiledResults;
SmallVector<SmallVector<OpFoldResult>> resultOffsets, resultSizes;
if (failed(tiledBodyFn(rewriter, loc, ivs, offsets, sizes,
innerDestinationTensors, tiledResults, resultOffsets,
resultSizes)))
return rewriter.notifyMatchFailure(loc, "failed to generate loop body");
rewriter.setInsertionPointToEnd(forallOp.getTerminator().getBody());
for (auto [tiledValue, destinationTensor, resultOffset, resultSize] :
llvm::zip_equal(tiledResults, innerDestinationTensors, resultOffsets,
resultSizes)) {
SmallVector<OpFoldResult> resultStride(resultOffset.size(),
rewriter.getIndexAttr(1));
tensor::ParallelInsertSliceOp::create(rewriter, loc, tiledValue,
destinationTensor, resultOffset,
resultSize, resultStride);
}
return loops;
}
/// Generate the tile-loop nest using custom loop operation.
/// - `loopRanges` specifies the lb, ub and step of the untiled iteration space.
/// - `tileSizes` is the tile sizes to use. Zero represent untiled loops.
/// - `destinationTensors` are the init values to use for the outer most loop.
/// - `mappingVector` is the mapping attributes to use for loop construction.
/// Can be empty.
/// - `tiledBodyFn` is called to generated the loop body of the inner
/// most
/// loop.
/// Returns the generated `scf.forall` loop on success.
static FailureOr<SmallVector<LoopLikeOpInterface>>
generateLoopNestUsingCustomOp(
RewriterBase &rewriter, Location loc, ArrayRef<Range> loopRanges,
ArrayRef<OpFoldResult> givenTileSizes, ValueRange outerDestinationTensors,
const scf::SCFTilingOptions::GenerateLoopHeaderFn &generateLoopHeaderFn,
const scf::SCFTilingOptions::GenerateLoopTerminatorFn
&generateLoopTerminatorFn,
GenerateTiledBodyFn tiledBodyFn) {
assert(!loopRanges.empty() && "unexpected empty loop ranges");
assert(loopRanges.size() == givenTileSizes.size() &&
"expected as many tile sizes as loop ranges");
assert(generateLoopHeaderFn && generateLoopTerminatorFn &&
"expected loop header/terminator generation function");
OpBuilder::InsertionGuard guard(rewriter);
FailureOr<scf::SCFTilingOptions::CustomLoopHeaderInfo> loopHeaderInfo =
generateLoopHeaderFn(rewriter, loc, loopRanges, givenTileSizes,
outerDestinationTensors);
if (failed(loopHeaderInfo)) {
return failure();
}
SmallVector<Value> ivs;
SmallVector<Value> tiledResults;
SmallVector<SmallVector<OpFoldResult>> resultOffsets, resultSizes;
if (failed(tiledBodyFn(rewriter, loc, ivs, loopHeaderInfo->tileOffset,
loopHeaderInfo->tileSizes,
loopHeaderInfo->destinationTensors, tiledResults,
resultOffsets, resultSizes))) {
return failure();
}
if (failed(generateLoopTerminatorFn(rewriter, loc, loopHeaderInfo->loops,
tiledResults, resultOffsets, resultSizes,
loopHeaderInfo->destinationTensors))) {
return failure();
}
return loopHeaderInfo->loops;
}
/// Generate the tile-loop nest using the loop construct specifed in `options`.
/// - `options`: Tiling options specified.
/// - `loopRanges` specifies the lb, ub and step of the untiled iteration space.
/// - `tileSizes` is the tile sizes to use. Zero represent untiled loops.
/// - `outerDestinationTensors` are the init values to use for the outer most
/// loop.
/// - `yieldTiledValuesFn` is called to generated the loop body of the inner
/// most
/// loop.
/// Returns the generated loops on success.
static FailureOr<SmallVector<LoopLikeOpInterface>> generateLoopNest(
RewriterBase &rewriter, Location loc, const scf::SCFTilingOptions &options,
ArrayRef<Range> loopRanges, ArrayRef<OpFoldResult> givenTileSizes,
ArrayRef<OpFoldResult> numThreads, ValueRange destinationTensors,
GenerateTiledBodyFn tiledBodyFn) {
// If the tile sizes are all zero, no loops are generated. Just call the
// callback function to handle untiled case.
if (llvm::all_of(givenTileSizes, isZeroInteger)) {
SmallVector<Value> tiledResults;
SmallVector<SmallVector<OpFoldResult>> resultOffsets, resultSizes;
auto tileOffsets =
llvm::map_to_vector(loopRanges, [](Range r) { return r.offset; });
auto tileSizes =
llvm::map_to_vector(loopRanges, [](Range r) { return r.size; });
if (failed(tiledBodyFn(rewriter, loc, ValueRange{}, tileOffsets, tileSizes,
destinationTensors, tiledResults, resultOffsets,
resultSizes))) {
return failure();
}
return SmallVector<LoopLikeOpInterface>{};
}
if (options.loopType == scf::SCFTilingOptions::LoopType::ForOp) {
return generateLoopNestUsingForOp(rewriter, loc, loopRanges, givenTileSizes,
destinationTensors, tiledBodyFn);
}
if (options.loopType == scf::SCFTilingOptions::LoopType::ForallOp) {
return generateLoopNestUsingForallOp(
rewriter, loc, loopRanges, givenTileSizes, numThreads,
options.mappingVector, destinationTensors, tiledBodyFn);
}
if (options.loopType == scf::SCFTilingOptions::LoopType::CustomOp) {
return generateLoopNestUsingCustomOp(
rewriter, loc, loopRanges, givenTileSizes, destinationTensors,
options.generateLoopHeaderFn, options.generateLoopTerminatorFn,
tiledBodyFn);
}
return rewriter.notifyMatchFailure(loc, "unhandled loop type");
}
static FailureOr<SmallVector<Value>> createInitialTensorsForTiling(
RewriterBase &rewriter, TilingInterface op,
ReductionTilingStrategy reductionStrategy, ArrayRef<Range> iterationDomain,
ArrayRef<OpFoldResult> numThreads, ArrayRef<OpFoldResult> givenTileSizes,
const SetVector<unsigned> &reductionDims) {
SmallVector<Value> initTensors;
Location loc = op->getLoc();
if (reductionStrategy == ReductionTilingStrategy::FullReduction) {
if (failed(tensor::getOrCreateDestinations(rewriter, loc, op, initTensors)))
return failure();
return initTensors;
}
auto redOp = dyn_cast<PartialReductionOpInterface>(op.getOperation());
if (!redOp) {
return op->emitOpError(
"PartialReductionOuterReduction tiling strategy is only supported for "
"operations implementing PartialReductionOpInterface");
}
SmallVector<OpFoldResult> sizes(iterationDomain.size());
AffineExpr s0, s1, s2;
bindSymbols(rewriter.getContext(), s0, s1, s2);
AffineExpr sizeExpr = ((s0 - s1).ceilDiv(s2));
AffineExpr divExpr = s0.ceilDiv(s1);
for (auto [index, domain, tileSize] :
llvm::enumerate(iterationDomain, givenTileSizes)) {
if (!numThreads.empty()) {
// Untiled case.
if (isConstantIntValue(numThreads[index], 0)) {
sizes[index] = affine::makeComposedFoldedAffineApply(
rewriter, op.getLoc(), sizeExpr,
{domain.size, domain.offset, domain.stride});
continue;
}
sizes[index] = numThreads[index];
continue;
}
// Non reduction dimensions/non-tiled dimensions.
if (!reductionDims.contains(index) || isConstantIntValue(tileSize, 0)) {
sizes[index] = affine::makeComposedFoldedAffineApply(
rewriter, op.getLoc(), sizeExpr,
{domain.size, domain.offset, domain.stride});
continue;
}
if (reductionStrategy ==
ReductionTilingStrategy::PartialReductionOuterReduction) {
sizes[index] = tileSize;
continue;
}
assert(reductionStrategy ==
ReductionTilingStrategy::PartialReductionOuterParallel);
OpFoldResult normalizedRange = affine::makeComposedFoldedAffineApply(
rewriter, op.getLoc(), sizeExpr,
{domain.size, domain.offset, domain.stride});
sizes[index] = affine::makeComposedFoldedAffineApply(
rewriter, op.getLoc(), divExpr, {normalizedRange, tileSize});
}
return redOp.generateInitialTensorForPartialReduction(rewriter, loc, sizes,
reductionDims);
}
/// For the case of `ReductionTilingStrategy::PartialReductionOuterParallel`
/// the `PartialReductionOpInterface` methods need the index of the parallel
/// split reduction being executed.
static SmallVector<OpFoldResult>
getSplitReductionIvs(RewriterBase &rewriter, Location loc,
ReductionTilingStrategy reductionStrategy, ValueRange ivs,
ArrayRef<OpFoldResult> numThreads,
ArrayRef<OpFoldResult> givenTileSizes,
const SetVector<unsigned> &reductionDims) {
SmallVector<OpFoldResult> splitReductionIvs;
splitReductionIvs.resize(reductionDims.size(), rewriter.getIndexAttr(0));
AffineExpr s0, s1;
bindSymbols(rewriter.getContext(), s0, s1);
AffineExpr divExpr = s0.floorDiv(s1);
int ivIndex = 0;
if (reductionStrategy ==
ReductionTilingStrategy::PartialReductionOuterParallel) {
for (auto [index, reductionDim] : llvm::enumerate(reductionDims)) {
if (!numThreads.empty()) {
splitReductionIvs[index] = ivs[ivIndex++];
continue;
}
splitReductionIvs[index] = affine::makeComposedFoldedAffineApply(
rewriter, loc, divExpr,
ArrayRef<OpFoldResult>{ivs[ivIndex++], givenTileSizes[reductionDim]});
}
}
return splitReductionIvs;
}
static FailureOr<TilingResult>
getTiledImplementation(RewriterBase &rewriter, TilingInterface op,
ReductionTilingStrategy reductionStrategy,
ValueRange regionIterArg, ArrayRef<OpFoldResult> offsets,
ArrayRef<OpFoldResult> sizes, ValueRange ivs,
ArrayRef<OpFoldResult> numThreads,
ArrayRef<OpFoldResult> givenTileSizes,
ArrayRef<InnerTileAlignment> innerTileAlignments,
const SetVector<unsigned> &reductionDims) {
if (reductionStrategy == ReductionTilingStrategy::FullReduction) {
return op.getTiledImplementation(rewriter, offsets, sizes,
innerTileAlignments);
}
auto redOp = dyn_cast<PartialReductionOpInterface>(op.getOperation());
if (!redOp) {
return rewriter.notifyMatchFailure(
op, "PartialReductionOuterReduction tiling strategy is only "
"supported for operations "
"implementing PartialReductionOpInterface");
}
SmallVector<OpFoldResult> splitReductionIvs =
getSplitReductionIvs(rewriter, op.getLoc(), reductionStrategy, ivs,
numThreads, givenTileSizes, reductionDims);
return redOp.tileToPartialReduction(rewriter, op.getLoc(), reductionStrategy,
regionIterArg, offsets, sizes,
reductionDims, splitReductionIvs);
}
static LogicalResult getResultTilePosition(
RewriterBase &rewriter, ReductionTilingStrategy reductionStrategy,
int64_t index, Value tiledResult, TilingInterface op,
ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
ValueRange ivs, ArrayRef<OpFoldResult> numThreads,
ArrayRef<OpFoldResult> givenTileSizes,
const SetVector<unsigned> &reductionDims,
SmallVector<OpFoldResult> &resultOffset,
SmallVector<OpFoldResult> &resultSize) {
if (reductionStrategy == ReductionTilingStrategy::FullReduction) {
return op.getResultTilePosition(rewriter, index, offsets, sizes,
resultOffset, resultSize);
}
auto redOp = dyn_cast<PartialReductionOpInterface>(op.getOperation());
if (!redOp) {
return rewriter.notifyMatchFailure(
op, "PartialReductionOuterReduction tiling strategy is only supported"
"for operations implementing PartialReductionOpInterface");
}
SmallVector<OpFoldResult> splitReductionIvs =
getSplitReductionIvs(rewriter, op.getLoc(), reductionStrategy, ivs,
numThreads, givenTileSizes, reductionDims);
return redOp.getPartialResultTilePosition(
rewriter, index, reductionStrategy, offsets, sizes, reductionDims,
splitReductionIvs, resultOffset, resultSize);
}
static FailureOr<MergeResult>
mergeTilingResults(RewriterBase &rewriter, TilingInterface op,
ReductionTilingStrategy reductionStrategy,
const SetVector<unsigned> &reductionDims,
ValueRange partialResults) {
assert(reductionStrategy != ReductionTilingStrategy::FullReduction &&
"expected merge to be called for only partial reduction cases");
auto redOp = dyn_cast<PartialReductionOpInterface>(op.getOperation());
if (!redOp) {
return rewriter.notifyMatchFailure(
op, "PartialReductionOuterReduction tiling strategy is only "
"supported for operations "
"implementing PartialReductionOpInterface");
}
return redOp.mergeReductions(rewriter, op.getLoc(), partialResults,
reductionDims);
}
/// Append the specified additional `newInitOperands` operands to the
/// loops existing `init` operands (or similar), and replace `loopOp` with
/// the new loop that has the additional init operands. The loop body of
/// this loop is moved over to the new loop. `yieldTiledValuesFn`
/// is called to get the new tiled values returned, and the offset
/// and sizes at which the tiled value is inserted into the
/// new region iter_args that correspond to the newly added init operands.
template <typename LoopType>
static FailureOr<LoopLikeOpInterface>
yieldTiledValuesAndReplaceLoop(LoopType loopOp, RewriterBase &rewriter,
ValueRange newInitOperands,
YieldTiledValuesFn yieldTiledValuesFn) {
return rewriter.notifyMatchFailure(loopOp, "unhandled loop type");
}
/// Implementation of `yieldTiledValuesAndReplaceLoop` for `scf.for`.
template <>
FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop<scf::ForOp>(
scf::ForOp loopOp, RewriterBase &rewriter, ValueRange newInitOperands,
YieldTiledValuesFn yieldTiledValuesFn) {
OpBuilder::InsertionGuard g(rewriter);
Location loc = loopOp.getLoc();
rewriter.setInsertionPoint(loopOp);
auto inits = llvm::to_vector(loopOp.getInitArgs());
inits.append(newInitOperands.begin(), newInitOperands.end());
auto newLoop = scf::ForOp::create(
rewriter, loc, loopOp.getLowerBound(), loopOp.getUpperBound(),
loopOp.getStep(), inits, [](OpBuilder &, Location, Value, ValueRange) {},
loopOp.getUnsignedCmp());
// Move the loop body to the new op.
Block *loopBody = loopOp.getBody();
Block *newLoopBody = newLoop.getBody();
rewriter.mergeBlocks(
loopBody, newLoopBody,
newLoopBody->getArguments().take_front(loopBody->getNumArguments()));
auto yieldOp = cast<scf::YieldOp>(newLoopBody->getTerminator());
rewriter.setInsertionPoint(yieldOp);
SmallVector<Value> tiledValues;
SmallVector<SmallVector<OpFoldResult>> resultOffsets, resultSizes;
ValueRange newRegionIterArgs =
newLoop.getRegionIterArgs().take_back(newInitOperands.size());
if (failed(yieldTiledValuesFn(rewriter, loc, newLoop.getInductionVar(),
newRegionIterArgs, tiledValues, resultOffsets,
resultSizes))) {
rewriter.eraseOp(newLoop);
return rewriter.notifyMatchFailure(loopOp, "failed to get tiled values");
}
SmallVector<Value> newYieldValues = llvm::to_vector(yieldOp.getOperands());
for (auto [tiledValue, regionIterArg, resultOffset, resultSize] :
llvm::zip_equal(tiledValues, newRegionIterArgs, resultOffsets,
resultSizes)) {
SmallVector<OpFoldResult> resultStride(resultOffset.size(),
rewriter.getIndexAttr(1));
Value insert = tensor::InsertSliceOp::create(
rewriter, yieldOp->getLoc(), tiledValue, regionIterArg, resultOffset,
resultSize, resultStride);
newYieldValues.push_back(insert);
}
rewriter.replaceOpWithNewOp<scf::YieldOp>(yieldOp, newYieldValues);
rewriter.replaceOp(loopOp,
newLoop->getResults().take_front(loopOp.getNumResults()));
return cast<LoopLikeOpInterface>(newLoop.getOperation());
}
/// Implementation of `yieldTiledValuesAndReplaceLoop` for `scf.forall`
template <>
FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop<scf::ForallOp>(
scf::ForallOp loopOp, RewriterBase &rewriter, ValueRange newInitOperands,
YieldTiledValuesFn yieldTiledValuesFn) {
OpBuilder::InsertionGuard g(rewriter);
Location loc = loopOp.getLoc();
rewriter.setInsertionPoint(loopOp);
auto inits = llvm::to_vector(loopOp.getOutputs());
inits.append(newInitOperands.begin(), newInitOperands.end());
auto newLoop = scf::ForallOp::create(
rewriter, loc, loopOp.getMixedLowerBound(), loopOp.getMixedUpperBound(),
loopOp.getMixedStep(), inits, loopOp.getMapping(),
[](OpBuilder &, Location, ValueRange) {});
// Move the region of the current block to the newly created op.
Block *loopBody = loopOp.getBody();
Block *newLoopBody = newLoop.getBody();
rewriter.mergeBlocks(
loopBody, newLoopBody,
newLoopBody->getArguments().take_front(loopBody->getNumArguments()));
auto terminator = cast<scf::InParallelOp>(newLoopBody->getTerminator());
rewriter.setInsertionPoint(terminator);
SmallVector<Value> tiledValues;
SmallVector<SmallVector<OpFoldResult>> resultOffsets, resultSizes;
ValueRange regionIterArgs =
newLoop.getRegionIterArgs().take_back(newInitOperands.size());
if (failed(yieldTiledValuesFn(rewriter, loc, newLoop.getInductionVars(),
regionIterArgs, tiledValues, resultOffsets,
resultSizes))) {
rewriter.eraseOp(newLoop);
return rewriter.notifyMatchFailure(loopOp,
"failed to get yielded tiled values");
}
// Update the terminator.
rewriter.setInsertionPointToEnd(terminator.getBody());
for (auto [tiledValue, iterArg, resultOffset, resultSize] : llvm::zip_equal(
tiledValues, regionIterArgs, resultOffsets, resultSizes)) {
SmallVector<OpFoldResult> resultStride(resultOffset.size(),
rewriter.getIndexAttr(1));
tensor::ParallelInsertSliceOp::create(rewriter, terminator.getLoc(),
tiledValue, iterArg, resultOffset,
resultSize, resultStride);
}
rewriter.replaceOp(loopOp,
newLoop->getResults().take_front(loopOp.getNumResults()));
return cast<LoopLikeOpInterface>(newLoop.getOperation());
}
/// Implementation of `yieldTiledValuesAndReplaceLoop` for
/// `LoopLikeOpInterface`, that just dispatches to the implementation for each
/// supported loop type.
static FailureOr<LoopLikeOpInterface> yieldTiledValuesAndReplaceLoop(
LoopLikeOpInterface loopLikeOp, RewriterBase &rewriter,
ValueRange newInitOperands, YieldTiledValuesFn yieldTiledValuesFn) {
return TypeSwitch<Operation *, FailureOr<LoopLikeOpInterface>>(
loopLikeOp.getOperation())
.Case<scf::ForOp, scf::ForallOp>(
[&](auto loopOp) -> FailureOr<LoopLikeOpInterface> {
return yieldTiledValuesAndReplaceLoop(
loopOp, rewriter, newInitOperands, yieldTiledValuesFn);
})
.Default([&](auto loopOp) -> FailureOr<LoopLikeOpInterface> {
return rewriter.notifyMatchFailure(loopOp, "unhandled loop type");
});
}
/// Method to add new init values to a loop nest. Updates `loops` in-place
/// with new loops that use the `newInitValues`. The outer-loops are updated
/// to yield the new result values of the inner loop. For the innermost loop,
/// the call back `getNewYields` is invoked to get the additional values to
/// yield form the innermost loop.
static LogicalResult addInitOperandsToLoopNest(
RewriterBase &rewriter, MutableArrayRef<LoopLikeOpInterface> loops,
ValueRange newInitValues, YieldTiledValuesFn getNewTiledYieldsFn) {
if (loops.empty())
return success();
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPoint(loops.front());
SmallVector<Value> ivs;
for (auto &loop : loops.drop_back()) {
rewriter.setInsertionPoint(loop);
// if loops.size() > 1 we assume that scf.for is used for the loops.
auto forLoop = cast<scf::ForOp>(loop.getOperation());
// Create a new loop with the new init values for this loop.
SmallVector<Value> newInits = llvm::to_vector(forLoop.getInitArgs());
newInits.append(newInitValues.begin(), newInitValues.end());
auto newLoop = scf::ForOp::create(
rewriter, forLoop.getLoc(), forLoop.getLowerBound(),
forLoop.getUpperBound(), forLoop.getStep(), newInits,
[&](OpBuilder &b, Location loc, Value iv, ValueRange iterArgs) {},
forLoop.getUnsignedCmp());
// Merge the body of the new loop with the body of the old loops.
SmallVector<Value> sourceBlockArgs;
sourceBlockArgs.push_back(newLoop.getInductionVar());
auto newRegionIterArgs = newLoop.getRegionIterArgs();
sourceBlockArgs.append(
newRegionIterArgs.begin(),
std::next(newRegionIterArgs.begin(), forLoop.getNumResults()));
rewriter.mergeBlocks(forLoop.getBody(), newLoop.getBody(), sourceBlockArgs);
rewriter.replaceOp(
forLoop, newLoop.getResults().take_front(forLoop.getNumResults()));
loop = newLoop;
ivs.push_back(newLoop.getInductionVar());
newInitValues = newLoop.getRegionIterArgs().take_back(newInitValues.size());
}
// Update the loop body of the innermost loop to get new yield values.
LoopLikeOpInterface innerMostLoop = loops.back();
FailureOr<LoopLikeOpInterface> newInnerMostLoop =
yieldTiledValuesAndReplaceLoop(innerMostLoop, rewriter, newInitValues,
getNewTiledYieldsFn);
if (failed(newInnerMostLoop))
return innerMostLoop.emitOpError("failed to return additional yields");
loops.back() = newInnerMostLoop.value();
// Make all other loops except the innermost loops yield the values returned
// by the inner loop.
for (auto [outerLoop, innerLoop] :
llvm::zip_equal(loops.drop_back(), loops.drop_front())) {
// Again assume that all the outer loops are scf.for operations.
auto outerForLoop = cast<scf::ForOp>(outerLoop.getOperation());
auto outerLoopYield =
cast<scf::YieldOp>(outerForLoop.getBody()->getTerminator());
SmallVector<Value> newYields =
llvm::to_vector(outerLoopYield.getOperands());
ValueRange additionalYields =
innerLoop->getResults().take_back(newInitValues.size());
newYields.append(additionalYields.begin(), additionalYields.end());
rewriter.setInsertionPoint(outerLoopYield);
rewriter.replaceOpWithNewOp<scf::YieldOp>(outerLoopYield, newYields);
}
return success();
}
/// Implementation of tiling transformation of `op` that implements the
/// `TilingInterface` using `scf.for` to iterate over the tiles.
FailureOr<scf::SCFTilingResult>
mlir::scf::tileUsingSCF(RewriterBase &rewriter, TilingInterface op,
const scf::SCFTilingOptions &options) {
if (failed(verifyOptions(rewriter, op.getLoc(), options))) {
return failure();
}
OpBuilder::InsertionGuard guard(rewriter);
rewriter.setInsertionPointAfter(op);
// 1. Get the range of the loops that are represented by the operation.
SmallVector<Range> iterationDomain = op.getIterationDomain(rewriter);
// 2. Materialize the tile sizes and/or number of threads;
SmallVector<OpFoldResult> givenTileSizes, numThreads;
std::tie(givenTileSizes, numThreads) =
getUserTileSizesAndNumThreads(rewriter, op, iterationDomain, options);
// Check if it is safe to tile. This is hold over from previous iterations
// of tile to for-all. Consider dropping it.
if (failed(checkTileSizes(op, options.loopType, options.reductionStrategy,
givenTileSizes, numThreads))) {
return failure();
}
// Get the reduction dims
SetVector<unsigned> reductionDims =
getSanitizedReductionDims(givenTileSizes, options);
// 3. If there is an interchange specified, permute the iteration domain and
// the tile sizes.
SmallVector<int64_t> interchangeVector;
if (!options.interchangeVector.empty()) {
interchangeVector = fillInterchangeVector(options.interchangeVector,
iterationDomain.size());
assert(isPermutationVector(interchangeVector) &&
"expected interchange vector to be a permutation");
applyPermutationToVector(iterationDomain, interchangeVector);
applyPermutationToVector(givenTileSizes, interchangeVector);
if (!numThreads.empty())
applyPermutationToVector(numThreads, interchangeVector);
}
FailureOr<TilingResult> tilingResult;
// 4. Define the lambda function used later to generate the body of the
// innermost tiled loop.
GenerateTiledBodyFn innerYieldTiledValuesFn =
[&](RewriterBase &rewriter, Location loc, ValueRange ivs,
ArrayRef<OpFoldResult> tileOffsets, ArrayRef<OpFoldResult> tileSizes,
ValueRange regionIterArgs, SmallVector<Value> &tiledResults,
SmallVector<SmallVector<OpFoldResult>> &resultOffsets,
SmallVector<SmallVector<OpFoldResult>> &resultSizes)
-> LogicalResult {
// 4b. If interchange was provided, apply inverse of the interchange
// to get back the offsets/sizes in the order to be specified.
SmallVector<OpFoldResult> tileOffsetsVec = llvm::to_vector(tileOffsets);
SmallVector<OpFoldResult> tileSizesVec = llvm::to_vector(tileSizes);
if (!interchangeVector.empty()) {
auto inversePermutation = invertPermutationVector(interchangeVector);
applyPermutationToVector(tileOffsetsVec, inversePermutation);
applyPermutationToVector(tileSizesVec, inversePermutation);
}
// 5. Generate the tiled implementation within the inner most loop.
// 5a. Clone the operation within the loop body.
auto clonedOp = cast<TilingInterface>(
cloneOpAndUpdateDestinationArgs(rewriter, op, regionIterArgs));
// 5b. Early return cloned op if tiling is not happening. We can not
// return the original op because it could lead to `rewriter.replaceOp(op,
// op->getResults())` and users would get crash.
if (llvm::all_of(givenTileSizes, isZeroInteger)) {
tiledResults.append(clonedOp->result_begin(), clonedOp->result_end());
tilingResult =
TilingResult{/*tiledOps=*/{clonedOp}, clonedOp->getResults(),
/*generatedSlices=*/{}};
return success();
}
// 5c. Tile the cloned operation.
SmallVector<InnerTileAlignment> innerTileAlignments =
options.innerTileAlignmentFn
? options.innerTileAlignmentFn(clonedOp, givenTileSizes,
/*slices=*/{})
: SmallVector<InnerTileAlignment>{};
tilingResult = getTiledImplementation(
rewriter, clonedOp, options.reductionStrategy, regionIterArgs,
tileOffsetsVec, tileSizesVec, ivs, numThreads, givenTileSizes,
innerTileAlignments, reductionDims);
if (failed(tilingResult)) {
rewriter.eraseOp(clonedOp);
return op.emitOpError("failed to tile operation");
}
// 5d. Delete the cloned operation.
rewriter.eraseOp(clonedOp);
// 5e. Compute the offsets at which the result values are to be inserted
// back into its destinations.
for (auto [index, tiledValue] :
llvm::enumerate(tilingResult->tiledValues)) {
tiledResults.push_back(tiledValue);
SmallVector<OpFoldResult> resultOffset, resultSize;
if (failed(getResultTilePosition(
rewriter, options.reductionStrategy, index, tiledValue, op,
tileOffsetsVec, tileSizesVec, ivs, numThreads, givenTileSizes,
reductionDims, resultOffset, resultSize))) {
for (auto op : tilingResult->tiledOps) {
rewriter.eraseOp(op);
}
return rewriter.notifyMatchFailure(
op, "failed to get slice of result produced");
}
resultOffsets.emplace_back(std::move(resultOffset));
resultSizes.emplace_back(std::move(resultSize));
}
return success();
};
// 6. Find the destination tensors to use for the operation.
FailureOr<SmallVector<Value>> maybeInits = createInitialTensorsForTiling(
rewriter, op, options.reductionStrategy, iterationDomain, numThreads,
givenTileSizes, reductionDims);
if (failed(maybeInits)) {
return rewriter.notifyMatchFailure(
op, "unable to create initial tensors for tiling");
}
SmallVector<Value> &initTensors = maybeInits.value();
// 7. Generate the tiled loops nest using the callback defined above.
SmallVector<LoopLikeOpInterface> loops;
{
FailureOr<SmallVector<LoopLikeOpInterface>> loopsOr = generateLoopNest(
rewriter, op.getLoc(), options, iterationDomain, givenTileSizes,
numThreads, initTensors, innerYieldTiledValuesFn);
if (failed(loopsOr))
return op.emitOpError("failed to generate tiling loops");
assert(succeeded(tilingResult) &&
"expected tiling result to be computed after loop generation");
std::swap(loops, loopsOr.value());
}
if (loops.empty()) {
// If loops are empty, the tiled op is used as the replacement for the
// untiled op.
return scf::SCFTilingResult{tilingResult->tiledOps,
initTensors,
loops,
tilingResult->tiledValues,
tilingResult->generatedSlices,
{}};
}
auto loopResults = llvm::map_to_vector(loops.front()->getResults(),
[](OpResult r) -> Value { return r; });
// For the full reduction case, there is nothing more to do.
if (options.reductionStrategy == ReductionTilingStrategy::FullReduction) {
return scf::SCFTilingResult{
tilingResult->tiledOps, initTensors, loops, loopResults,
tilingResult->generatedSlices, {}};
}
// The results of the loop needs to be merged.
FailureOr<MergeResult> mergeResult = mergeTilingResults(
rewriter, op, options.reductionStrategy, reductionDims, loopResults);
if (failed(mergeResult)) {
return rewriter.notifyMatchFailure(
op, "Failed to merge partial results from tiling");
}
return scf::SCFTilingResult{tilingResult->tiledOps,
initTensors,
loops,
mergeResult->replacements,
tilingResult->generatedSlices,
mergeResult->mergeOps};
}
FailureOr<scf::SCFTilingResult>
mlir::scf::tileReductionUsingScf(RewriterBase &b,
PartialReductionOpInterface op,
ArrayRef<OpFoldResult> tileSize) {
scf::SCFTilingOptions options;
options.setLoopType(scf::SCFTilingOptions::LoopType::ForOp);
options.setReductionTilingStrategy(
ReductionTilingStrategy::PartialReductionOuterReduction);
options.setTileSizes(tileSize);
SmallVector<unsigned> reductionDims;
for (auto [index, iteratorType] : llvm::enumerate(op.getLoopIteratorTypes()))
if (iteratorType == utils::IteratorType::reduction)
reductionDims.push_back(index);
options.setReductionDims(reductionDims);
return tileUsingSCF(b, op, options);
}
//===----------------------------------------------------------------------===//
// tileConsumerAndFuseProducersUsingSCF implementation.
//===----------------------------------------------------------------------===//
/// Return the untiled producer whose slice is used in a tiled consumer. The
/// method traverses the tile loop nest (`loops`) if needed, and returns the
/// `iter_args` of the outer most that is encountered. Traversing the
/// iter_args indicates that this is a destination operand of the consumer. If
/// there was no loop traversal needed, the second value of the returned tuple
/// is empty.
static std::tuple<OpResult, std::optional<OpOperand *>>
getUntiledProducerFromSliceSource(OpOperand *source,
ArrayRef<LoopLikeOpInterface> loops) {
std::optional<OpOperand *> destinationIterArg;
assert(!loops.empty() && "expected non empty loops container");
auto loopIt = loops.rbegin();
while (loopIt != loops.rend() && isa<BlockArgument>(source->get())) {
auto iterArg = cast<BlockArgument>(source->get());
auto loop = *loopIt;
if (iterArg.getOwner()->getParentOp() != loop)
break;
source = loop.getTiedLoopInit(iterArg);
loopIt++;
}
if (loopIt == loops.rend())
destinationIterArg = source;
auto result = dyn_cast<OpResult>(source->get());
if (result) {
Operation *producer = result.getOwner();
Operation *innermostLoop = loops.back();
// If the producer is already inside the innermost loop (where the slice
// is), it has already been fused. Skip it to avoid infinite loops.
if (innermostLoop->isProperAncestor(producer))
return {OpResult(), std::nullopt};
}
return {result, destinationIterArg};
}
/// Implementation of fusing producer of a single slice by computing the
/// slice of the producer in-place.
std::optional<scf::SCFFuseProducerOfSliceResult>
mlir::scf::tileAndFuseProducerOfSlice(
RewriterBase &rewriter, tensor::ExtractSliceOp candidateSliceOp,
MutableArrayRef<LoopLikeOpInterface> loops,
const InnerTileAlignmentFnTy &fn) {
// 1. Get the producer of the source (potentially walking through
// `iter_args` of nested `scf.for`)
auto [fusableProducer, destinationInitArg] =
getUntiledProducerFromSliceSource(&candidateSliceOp.getSourceMutable(),
loops);
if (!fusableProducer)
return std::nullopt;
unsigned resultNumber = fusableProducer.getResultNumber();
// Resolve the inner-tile alignment hint for the producer in its own iteration
// domain via the control function (consulted only by pack/unpack).
SmallVector<InnerTileAlignment> innerTileAlignments;
if (fn)
if (auto producer = dyn_cast<TilingInterface>(fusableProducer.getOwner()))
innerTileAlignments =
fn(producer, /*tileSizes=*/{}, {candidateSliceOp.getOperation()});
OpBuilder::InsertionGuard g(rewriter);
rewriter.setInsertionPoint(candidateSliceOp);
// 2. Clone the fused producer
// 2a. Compute the destination operands to use for the cloned operation.
SmallVector<Value> origDestinationTensors, clonedOpDestinationTensors;
Operation *fusableProducerOp = fusableProducer.getOwner();
if (isa<DestinationStyleOpInterface>(fusableProducerOp) &&
failed(tensor::getOrCreateDestinations(
rewriter, fusableProducerOp->getLoc(), fusableProducerOp,
origDestinationTensors)))
return std::nullopt;
clonedOpDestinationTensors = origDestinationTensors;
if (destinationInitArg &&
isa<DestinationStyleOpInterface>(fusableProducerOp)) {
// 2b. If the producer is also destination style, then to maintain the
// destination passing style, update the destination of the producer to be
// the source of the slice.
clonedOpDestinationTensors[resultNumber] = candidateSliceOp.getSource();
}
// 2c. Clone the fused producer.
Operation *clonedProducerOp = cloneOpAndUpdateDestinationArgs(
rewriter, fusableProducerOp, clonedOpDestinationTensors);
// 2d. Update the source of the candidateSlice to be the cloned producer.
// Easier to just clone the slice with different source since
// replacements and DCE of cloned ops becomes easier
SmallVector<Value> candidateSliceOpOperands =
llvm::to_vector(candidateSliceOp->getOperands());
candidateSliceOpOperands[0] = clonedProducerOp->getResult(resultNumber);
tensor::ExtractSliceOp clonedCandidateSliceOp =
mlir::clone(rewriter, candidateSliceOp,
candidateSliceOp->getResultTypes(), candidateSliceOpOperands);
// 3. Generate the tiled implementation of the producer of the source
FailureOr<TilingResult> tileAndFuseResult =
tensor::replaceExtractSliceWithTiledProducer(
rewriter, clonedCandidateSliceOp,
clonedProducerOp->getResult(resultNumber), innerTileAlignments);
if (failed(tileAndFuseResult))
return std::nullopt;
// Note: Do not delete the candidateSliceOp, since its passed in from the
// caller.
rewriter.replaceAllUsesWith(candidateSliceOp,
tileAndFuseResult->tiledValues[0]);
rewriter.eraseOp(clonedCandidateSliceOp);
rewriter.eraseOp(clonedProducerOp);
// 3. If the slice is for a destination operand, for example,
//
// ```mlir
// %0 = linalg.init
// %1 = linalg.fill .. outs(%0 : )
// %2 = scf.for .. iter_args(%arg0 = %1) {
// %3 = scf.for .. iter_args(%arg1 = %arg0) {
// %4 = tensor.extract_slice %arg1 [..]
// .. = linalg.matmul .. outs(%4 : )
// }
// }
// ```
//
// the IR is currently
//
// ```
// %0 = linalg.init
// %1 = linalg.fill
// %2 = scf.for .. iter_args(%arg0 = %1 /* incorrect value */ ) {
// %3 = scf.for .. iter_args(%arg1 = %arg0) {
// %4 = tensor.extract_slice %arg1[..]
// %5 = linalg.fill .. outs(%4 : )
// .. = linalg.matmul .. outs(%5 : )
// }
// }
// ```
//
// The untiled `linalg.fill` is still used as the `init_value` since it
// was originally a destination operand of the untiled `linalg.matmul`.
// When fusing an operand that is a destination operand, the iter_arg of
// the outer most loop should be changed to use the destination of the
// fused operation. With this the IR will be.
//
// ```
// %0 = linalg.init
// %1 = scf.for .. iter_args(%arg0 = %0 /* corrected value */ ) {
// %2 = scf.for .. iter_args(%arg1 = %arg0) {
// %3 = tensor.extract_slice %arg1[..]
// %4 = linalg.fill .. outs(%3 : )
// .. = linalg.matmul .. outs(%4 : )
// }
// }
// ```
if (destinationInitArg &&
isa<DestinationStyleOpInterface>(fusableProducerOp) && !loops.empty()) {
loops.front()
->getOpOperands()[destinationInitArg.value()->getOperandNumber()]
.set(origDestinationTensors[resultNumber]);
}
return scf::SCFFuseProducerOfSliceResult{
fusableProducer, tileAndFuseResult->tiledValues[0],
tileAndFuseResult->tiledOps, tileAndFuseResult->generatedSlices};
}
/// Reconstruct the fused producer from within the tiled-and-fused code.
FailureOr<SmallVector<Operation *>> mlir::scf::yieldReplacementForFusedProducer(
RewriterBase &rewriter, tensor::ExtractSliceOp sliceOp,
scf::SCFFuseProducerOfSliceResult fusedProducerInfo,
MutableArrayRef<LoopLikeOpInterface> loops,
ArrayRef<unsigned> yieldResultNumber) {
if (loops.empty())
return success();
Operation *originalOwner = fusedProducerInfo.origProducer.getOwner(),
*tiledOwner = fusedProducerInfo.tiledOps[0];
Location loc = originalOwner->getLoc();
// a. collect all init Value to be appended
SmallVector<unsigned> initNumberList =
yieldResultNumber.empty() ? llvm::to_vector(llvm::seq<unsigned>(
0, originalOwner->getNumResults()))
: llvm::to_vector(yieldResultNumber);
SmallVector<Value> initValueList;
for (const auto &resultNumber : initNumberList) {
FailureOr<Value> initValue = tensor::getOrCreateDestination(
rewriter, loc, originalOwner->getResult(resultNumber));
if (succeeded(initValue)) {
initValueList.push_back(initValue.value());
} else {
return failure();
}
}
SmallVector<Operation *> generatedSlices;
YieldTiledValuesFn newYieldValuesFn =
[&](RewriterBase &innerRewriter, Location loc, ValueRange /*ivs*/,
ValueRange newRegionIterArgs, SmallVector<Value> &tiledResult,
SmallVector<SmallVector<OpFoldResult>> &tiledOffset,
SmallVector<SmallVector<OpFoldResult>> &tiledSizes) -> LogicalResult {
OpBuilder::InsertionGuard g(innerRewriter);
// get sliceOp tile information
SmallVector<OpFoldResult> sliceOffset = sliceOp.getMixedOffsets(),
sliceSizes = sliceOp.getMixedSizes();
// expect all strides of sliceOp being 1
if (!llvm::all_of(sliceOp.getMixedStrides(), isOneInteger))
return failure();
unsigned sliceResultNumber =
fusedProducerInfo.origProducer.getResultNumber();
auto tilableOp = cast<TilingInterface>(originalOwner);
// b. get iterDomain Offset and Sizes based on sliceOp tile
SmallVector<OpFoldResult> iterDomainOffset, iterDomainSizes;
// Set insertion point before any operations that might create new SSA
// values used in offset/size computations. This ensures all values created
// by getIterationDomainTileFromResultTile and getResultTilePosition
// dominate the extract_slice operations created later.
if (auto tiledDestStyleOp =
dyn_cast<DestinationStyleOpInterface>(tiledOwner)) {
rewriter.setInsertionPoint(tiledDestStyleOp);
}
// skip tensor.pack/unpack/pad, which expects single opResult
if (tilableOp->getNumResults() > 1 &&
failed(tilableOp.getIterationDomainTileFromResultTile(
rewriter, sliceResultNumber, sliceOffset, sliceSizes,
iterDomainOffset, iterDomainSizes))) {
// In theory, it is unnecessary to raise an error here. Actually
// although it fails to reconstruct the result tensor, it should not
// broke current fusion anyway. The reason why we must return failure
// currently is that the callback function `newYieldValuesFn` will be
// called after new init operand(s) has already been appended. It will
// take more refactoring to make sure the init operands are added
// consistently in the future. For more details, please refer to:
// https://github.com/llvm/llvm-project/pull/93144#discussion_r1643760814
return failure();
}
// c. calculate offsets and sizes info of all OpResults respectively based
// on iteration Domain Tile
SmallVector<SmallVector<OpFoldResult>> offsetList, sizesList;
for (const auto &resultNumber : initNumberList) {
if (resultNumber == sliceResultNumber) {
offsetList.push_back(sliceOffset);
sizesList.push_back(sliceSizes);
} else {
assert(!iterDomainOffset.empty() && !iterDomainSizes.empty());
// infer result tile according to the iteration domain tile
SmallVector<OpFoldResult> offset, sizes;
if (failed(tilableOp.getResultTilePosition(
rewriter, resultNumber, iterDomainOffset, iterDomainSizes,
offset, sizes))) {
return failure();
}
offsetList.push_back(offset);
sizesList.push_back(sizes);
}
}
// d. create `extract_slice` for `iter_args` for DPS operation if
// necessary
if (auto tiledDestStyleOp =
dyn_cast<DestinationStyleOpInterface>(tiledOwner)) {
for (const auto &&[index, newRegionArg] :
llvm::enumerate(newRegionIterArgs)) {
auto destSlice = tensor::ExtractSliceOp::create(
rewriter, loc, newRegionArg, offsetList[index], sizesList[index],
SmallVector<OpFoldResult>(offsetList[index].size(),
rewriter.getIndexAttr(1)));
generatedSlices.push_back(destSlice);
unsigned resultNumber = initNumberList[index];
rewriter.modifyOpInPlace(tiledDestStyleOp, [&]() {
tiledDestStyleOp.getDpsInitsMutable()[resultNumber].set(destSlice);
});
}
}
// e. prepare tiled offset and sizes for later `insert_slice` creation by
// caller
Block *block = rewriter.getInsertionPoint()->getBlock();
rewriter.setInsertionPoint(block->getTerminator());
for (const auto &&[index, resultNumber] : llvm::enumerate(initNumberList)) {
tiledResult.push_back(tiledOwner->getResult(resultNumber));
tiledOffset.emplace_back(offsetList[index]);
tiledSizes.emplace_back(sizesList[index]);
}
return success();
};
if (failed(addInitOperandsToLoopNest(rewriter, loops, initValueList,
newYieldValuesFn))) {
return failure();
}
return generatedSlices;
}
namespace {
//===----------------------------------------------------------------------===//
// SliceTrackingListener
//===----------------------------------------------------------------------===//
/// This class is a listener for tracking the insertion and removal of
/// `tensor.extract_slice` ops in a worklist. This can be used in a greedy
/// fusion algorithm to apply cleanup patterns in between fusion steps.
class SliceTrackingListener : public RewriterBase::Listener {
public:
explicit SliceTrackingListener(
std::optional<FrozenRewritePatternSet> patterns);
SliceTrackingListener() = default;
/// Adds the given list of operations to the worklist, and if present,
/// applies the list of `patterns` to the newly added operations. This only
/// processes the given operations and any newly inserted ones by the
/// pattern set.
LogicalResult insertAndApplyPatterns(ArrayRef<Operation *> newOps);
/// Add to the new operation worklist if it is an extract_slice.
void notifyOperationInserted(Operation *op,
OpBuilder::InsertPoint previous) override;
/// Shared helper for operation removal from the worklist.
void removeOp(Operation *op);
/// Remove the operation from the worklist.
void notifyOperationErased(Operation *op) override;
/// Remove the operation from the worklist.
void notifyOperationReplaced(Operation *op, ValueRange replacement) override;
/// The worklist for this transformation keeps track of the slices to visit
/// next for fusion.
std::deque<tensor::ExtractSliceOp> worklist;
private:
/// Optional pattern set to apply when adding new operations to the
/// worklist.
std::optional<FrozenRewritePatternSet> patterns = std::nullopt;
};
SliceTrackingListener::SliceTrackingListener(
std::optional<FrozenRewritePatternSet> p) {
patterns = std::move(p);
}
LogicalResult
SliceTrackingListener::insertAndApplyPatterns(ArrayRef<Operation *> ops) {
for (Operation *op : ops) {
if (auto slice = dyn_cast<tensor::ExtractSliceOp>(op))
worklist.push_back(slice);
}
if (!patterns)
return success();
return applyOpPatternsGreedily(
ops, patterns.value(),
GreedyRewriteConfig().setListener(this).setStrictness(
GreedyRewriteStrictness::ExistingAndNewOps));
}
void SliceTrackingListener::notifyOperationInserted(
Operation *op, OpBuilder::InsertPoint previous) {
auto slice = dyn_cast<tensor::ExtractSliceOp>(op);
if (!slice)
return;
worklist.push_back(slice);
}
// Scan the worklist for the given op and remove it if present. The
// expectation is for the worklist to be small and for removal to be
// relatively rare.
void SliceTrackingListener::removeOp(Operation *op) {
if (!isa<tensor::ExtractSliceOp>(op))
return;
auto iter = worklist.begin();
while (iter != worklist.end()) {
if (*iter == op)
break;
iter++;
}
if (iter == worklist.end())
return;
worklist.erase(iter);
}
void SliceTrackingListener::notifyOperationErased(Operation *op) {
removeOp(op);
}
void SliceTrackingListener::notifyOperationReplaced(Operation *op,
ValueRange replacement) {
removeOp(op);
}
//===----------------------------------------------------------------------===//
// ReplacementListener
//===----------------------------------------------------------------------===//
/// Listener that tracks updates replacements for values which can be mutated.
/// This listener runs on top of the existing listener for the rewriter,
/// to make sure external users can still run listeners.
class ReplacementListener : public RewriterBase::ForwardingListener {
public:
ReplacementListener(DenseMap<Value, Value> &replacements,
OpBuilder::Listener *listener)
: ForwardingListener(listener), replacements(replacements) {}
void updateReplacementValues(ValueRange origValues,
ValueRange replaceValues) {
// This can probably be written better, but just iterates over the map
// and the new replacements for now.
for (auto &[key, val] : replacements) {
for (auto [orig, replace] : llvm::zip_equal(origValues, replaceValues)) {
if (val == orig) {
val = replace;
}
}
}
}
void notifyOperationReplaced(Operation *op, Operation *newOp) override {
ForwardingListener::notifyOperationReplaced(op, newOp);
updateReplacementValues(op->getResults(), newOp->getResults());
}
void notifyOperationReplaced(Operation *op, ValueRange values) override {
ForwardingListener::notifyOperationReplaced(op, values);
updateReplacementValues(op->getResults(), values);
}
private:
DenseMap<Value, Value> &replacements;
};
} // namespace
/// Implementation of tile consumer and fuse producer greedily.
FailureOr<scf::SCFTileAndFuseResult>
mlir::scf::tileConsumerAndFuseProducersUsingSCF(
RewriterBase &rewriter, TilingInterface consumer,
const scf::SCFTileAndFuseOptions &options) {
// This transformation is only valid for ops that return values (i.e. not
// valid to use with operations that have memref operands).
if (!consumer->getNumResults()) {
return rewriter.notifyMatchFailure(
consumer, "invalid pattern for op with no results");
}
// 1. First tile the consumer.
SetVector<Operation *> fusedProducers, tiledAndFusedOps;
FailureOr<scf::SCFTilingResult> tilingResult =
tileUsingSCF(rewriter, consumer, options.tilingOptions);
if (failed(tilingResult))
return rewriter.notifyMatchFailure(consumer, "failed to tile consumer");
tiledAndFusedOps.insert_range(tilingResult->tiledOps);
DenseMap<Value, Value> replacements;
for (auto [origVal, replacement] :
llvm::zip_equal(consumer->getResults(), tilingResult->replacements)) {
replacements[origVal] = replacement;
}
// If there are no loops generated, fusion is immaterial.
auto &loops = tilingResult->loops;
if (loops.empty()) {
return scf::SCFTileAndFuseResult{fusedProducers, tiledAndFusedOps, loops,
replacements};
}
// Since the loop gets potentially replaced during fusion, we need to track
// the mutation of replacement values. To do this, we attach a listener to
// update the replacements as they happen.
OpBuilder::Listener *previousListener = rewriter.getListener();
llvm::scope_exit resetListener(
[&]() { rewriter.setListener(previousListener); });
ReplacementListener replaceListener(replacements, previousListener);
rewriter.setListener(&replaceListener);
// 2. Typically, the operands of the tiled operation are slices of the
// operands of the untiled operation. These are expressed in IR using
// `tensor.extract_slice` operations with source being the operands of
// the untiled operation. Create a worklist of these
// `tensor.extract_slice` operations. If the producers of the source of
// the `tensor.extract_slice` can be tiled such that the tiled value is
// generated in-place, that effectively tiles + fuses the operations.
struct WorklistItem {
tensor::ExtractSliceOp candidateSlice;
SCFTileAndFuseOptions::ControlFnResult controlFnResult;
};
SliceTrackingListener sliceTracker =
SliceTrackingListener(options.cleanupPatterns);
if (failed(
sliceTracker.insertAndApplyPatterns(tilingResult->generatedSlices))) {
return rewriter.notifyMatchFailure(consumer, "cleanup patterns failed");
}
OpBuilder::InsertionGuard g(rewriter);
while (!sliceTracker.worklist.empty()) {
auto candidateSlice = sliceTracker.worklist.front();
sliceTracker.worklist.pop_front();
auto [fusableProducer, destinationInitArg] =
getUntiledProducerFromSliceSource(&candidateSlice.getSourceMutable(),
loops);
if (!fusableProducer)
continue;
std::optional<SCFTileAndFuseOptions::ControlFnResult> controlFnResult =
options.fusionControlFn(candidateSlice, fusableProducer,
destinationInitArg.has_value());
if (!controlFnResult)
continue;
WorklistItem worklistItem = {candidateSlice, controlFnResult.value()};
// The operands of the fused producer might themselved be slices of
// values produced by operations that implement the `TilingInterface`.
// Add these operations to the worklist.
std::optional<scf::SCFFuseProducerOfSliceResult> fusedResult =
tileAndFuseProducerOfSlice(rewriter, worklistItem.candidateSlice, loops,
options.tilingOptions.innerTileAlignmentFn);
if (!fusedResult)
continue;
SmallVector<Operation *> worklistCandidates = fusedResult->generatedSlices;
if (worklistItem.controlFnResult.yieldProducerReplacement) {
// Reconstruct and yield all opResult of fusableProducerOp by default.
// The caller can specific which one to yield by designating optional
// argument named `yieldResultNumber` of
// `yieldReplacementForFusedProducer`.
Operation *fusableProducerOp = fusedResult->origProducer.getOwner();
FailureOr<SmallVector<Operation *>> newSlices =
yieldReplacementForFusedProducer(rewriter,
worklistItem.candidateSlice,
fusedResult.value(), loops);
if (failed(newSlices)) {
return rewriter.notifyMatchFailure(
fusableProducerOp, "failed to replacement value for this "
"operation from within the tiled loop");
}
worklistCandidates.append(newSlices.value());
for (auto [index, result] :
llvm::enumerate(fusableProducerOp->getResults())) {
replacements[result] = loops.front()->getResult(
loops.front()->getNumResults() -
fusableProducerOp->getNumResults() + index);
}
}
if (Operation *tiledAndFusedOp =
fusedResult->tiledAndFusedProducer.getDefiningOp()) {
fusedProducers.insert(fusedResult->origProducer.getDefiningOp());
tiledAndFusedOps.insert(tiledAndFusedOp);
}
if (failed(sliceTracker.insertAndApplyPatterns(worklistCandidates))) {
return rewriter.notifyMatchFailure(consumer, "cleanup patterns failed");
}
}
return scf::SCFTileAndFuseResult{fusedProducers, tiledAndFusedOps, loops,
replacements};
}
//===----------------------------------------------------------------------===//
// tileAndFuseConsumerUsingSCF implementation.
//===----------------------------------------------------------------------===//
/// A utility function that checks whether the only use of the result of a
/// tensor.insert_slice op is in a scf.yield op.
static LogicalResult
checkAssumptionForFusingConsumer(tensor::InsertSliceOp candidateSliceOp) {
Value result = candidateSliceOp.getResult();
Value::use_range uses = result.getUses();
if (!llvm::hasSingleElement(uses)) {
LLVM_DEBUG(llvm::dbgs() << "Too many uses of the candidate slice op\n");
return failure();
}
OpOperand &operandUse = (*uses.begin());
Operation *userOp = operandUse.getOwner();
if (!isa<scf::YieldOp>(userOp)) {
LLVM_DEBUG(llvm::dbgs()
<< "Expected scf.yield to be the only user, but got -> "
<< (*userOp));
return failure();
}
if (result.getDefiningOp()->getBlock() != userOp->getBlock()) {
LLVM_DEBUG(llvm::dbgs() << "Expected tensor.insert_slice and scf.yield to "
"be in the same block\n");
return failure();
}
return success();
}
/// An utility to get the first user of the given loopOp. If any of user stay
/// in different block of loopOp, return failure.
static FailureOr<Operation *> getFirstUserOfLoop(Operation *loopOp) {
if (!isa<LoopLikeOpInterface>(loopOp))
return failure();
Operation *firstUserOfLoop = nullptr;
for (Operation *userOp : loopOp->getUsers()) {
// `ParallelInsertSlice` located inside `InParallelOp` has no same parent
// block with any other types of operation. Thus, just redirecting to its
// parent `InParallelOp`. E.g.
//
// ```
// %1 = scf.for {
// ...
// }
// %2 = consumerOp ins(%1, ...)
// scf.forall.in_parallel {
// tensor.parallel_insert_slice %1
// }
// ```
// where `InParallelOp` but not `ParallelInsertSlice` stays in the same
// same block with `consumerOp`.
if (isa<tensor::ParallelInsertSliceOp>(userOp))
userOp = userOp->getParentOfType<scf::InParallelOp>();
if (loopOp->getBlock() != userOp->getBlock())
return failure();
if (!firstUserOfLoop || userOp->isBeforeInBlock(firstUserOfLoop))
firstUserOfLoop = userOp;
}
return firstUserOfLoop;
}
/// This utility currently checks whether the first userOp of loop is NOT
/// before the last defineOp of consumer operand. Because that we need to move
/// the whole loop structure right before the `firstUserOfLoop`. This utility
/// thus helps ensuring that no invalid IR is formed, i.e. no backward slice
/// of consumerOp is dominated by the `firstUserOfLoop`. Saying that:
///
/// ```
/// %0 = scf.for() {
/// ...
/// }
/// ...
/// %1 = firstUserOfLoop(%0)
/// ...
/// %2 = lastDefOfConsumerOperand
/// ...
/// %3 = consumerOp(%2)
/// ```
///
/// If the `firstUserOfLoop` is before `lastDefOfConsumerOperand`, then it
/// would be invalid to move the `loopOp` right before the `firstUserOfLoop`,
/// a.k.a. use-def chain violation:
///
/// ```
/// %0:2 = scf.for() {
/// // use before define error
/// %3 = tiledConsumerOp(%2)
/// }
/// %1 = firstUserOfLoop(%0)
/// ...
/// %2 = lastDefOfConsumerOperand
/// ```
///
/// @param loopOp: loop operation
/// @param consumerOp: consumer operation
/// @param reorderOperations: the flag controls whether to reorder the
/// backward slice w.r.t. the defineOp of `consumerOp` operands.
/// @return: computed backward slice of consumerOp, but excluding those
/// already dominates `firstUserOfLoop`.
static FailureOr<llvm::SetVector<Operation *>>
checkAssumptionForLoop(Operation *loopOp, Operation *consumerOp,
bool reorderOperations) {
FailureOr<Operation *> firstUserOfLoop = getFirstUserOfLoop(loopOp);
if (failed(firstUserOfLoop))
return failure();
BackwardSliceOptions options;
DominanceInfo dominanceInfo;
options.inclusive = true;
options.omitBlockArguments = true;
bool includeLoopOp = false;
options.filter = [&](Operation *op) {
if (op == loopOp) {
includeLoopOp = true;
return false;
}
// Cut off the slice to not include any operation that already dominates
// firstUserOfLoop.
return !dominanceInfo.properlyDominates(op, *firstUserOfLoop);
};
llvm::SetVector<Operation *> slice;
for (auto operand : consumerOp->getOperands()) {
LogicalResult result = getBackwardSlice(operand, &slice, options);
assert(result.succeeded() && "expected a backward slice");
(void)result;
}
if (!slice.empty()) {
// If consumerOp has one producer, which is also the user of loopOp.
// E.g.
// ```
// %0 = %loopOp
// %1 = consumerOp1 ins(%0)
// %2 = consumerOp2 ins(%0, %1)
// ```
// We can not fuse consumerOp2 into loopOp due to UD chain, unless
// consumerOp1 has already been fused into loopOp before.
if (includeLoopOp || !reorderOperations)
return failure();
}
return slice;
}
/// Fetches the OpOperand of the first valid user (and use) of the value `val`
/// which implements `TilingInterface` and `DestinationStyleOpInterface`.
/// Returns failure otherwise.
static FailureOr<OpOperand *> getConsumerFromLoopUses(RewriterBase &rewriter,
Operation *loopOp,
unsigned resultNumber) {
if (!isa<LoopLikeOpInterface>(loopOp))
return failure();
Value val = loopOp->getResult(resultNumber);
Block *loopBlock = loopOp->getBlock();
for (OpOperand &opOperand : val.getUses()) {
Operation *consumerOp = opOperand.getOwner();
// Step 1. Check if the user is tilable.
if (!isa<TilingInterface>(consumerOp) ||
!isa<DestinationStyleOpInterface>(consumerOp)) {
// TODO: We have to init result of consumer before scf.for, use
// DestinationStyleOpInterface to get result shape from init for now.
// Add support for other op such as op has InferTypeOpInterface.
continue;
}
// Step 2. Check if user stay in the same block.
if (loopBlock != consumerOp->getBlock())
continue;
// Step 3. Check if user has succeeding user. Otherwise, it usually
// represents already tiled.
if (consumerOp->use_empty())
continue;
// Step 4. Check assumption for loop with `reorderOperations` enabled.
FailureOr<llvm::SetVector<Operation *>> slice =
checkAssumptionForLoop(loopOp, consumerOp, true);
if (failed(slice))
continue;
// Step 5. If backward sice is not empty, move them before
// firstUserOfLoop.
if (!slice->empty()) {
mlir::topologicalSort(*slice);
FailureOr<Operation *> firstUserOfLoop = getFirstUserOfLoop(loopOp);
assert(succeeded(firstUserOfLoop) && "First user of loop is not found");
for (auto op : *slice) {
rewriter.moveOpBefore(op, *firstUserOfLoop);
}
}
return &opOperand;
}
return failure();
}
/// Fetch the untiled consumer of the outermost scf.for's result which is
/// yielded by a tensor.insert_slice from the innermost scf.for. This function
/// makes the following assumptions :
/// 1. tensor.insert_slice has scf.yield as its only user.
/// 2. scf.for's corresponding result has only one use.
/// 3. The `loops` passed in are perfectly nested `scf.for` operations.
static FailureOr<OpOperand *>
getUntiledConsumerFromSlice(RewriterBase &rewriter,
tensor::InsertSliceOp candidateSliceOp,
MutableArrayRef<LoopLikeOpInterface> loops) {
assert(!loops.empty() && "unexpected loops to be empty");
// 1. Expect slice to be part of the body of the inner most loop.
Operation *containingOp = candidateSliceOp->getParentOp();
if (containingOp != loops.back()) {
return rewriter.notifyMatchFailure(
candidateSliceOp,
"expected slice to be within body of inner-most loop");
}
// 2. Check that the loop is perfectly nested.
if (!isPerfectlyNestedForLoops(loops)) {
return rewriter.notifyMatchFailure(
candidateSliceOp, "expected passed loops to be perfectly nested.");
}
if (failed(checkAssumptionForFusingConsumer(candidateSliceOp)))
return failure();
Value sliceResult = candidateSliceOp.getResult();
// 3. Fetch the corresponding output.
OpOperand &yieldOpOperand = (*sliceResult.getUses().begin());
unsigned resultNumber = yieldOpOperand.getOperandNumber();
scf::ForOp topLevelForOp = cast<scf::ForOp>(loops.front().getOperation());
return getConsumerFromLoopUses(rewriter, topLevelForOp, resultNumber);
}
/// Fetch the first untiled consumer of a scf.forall's result which is yielded
/// by a tensor.parallel_insert_slice.
static FailureOr<OpOperand *>
getUntiledConsumerFromSlice(RewriterBase &rewriter,
tensor::ParallelInsertSliceOp candidateSliceOp,
MutableArrayRef<LoopLikeOpInterface> loops) {
assert(!loops.empty() && "unexpected loops to be empty");
// 1. Check that the surrounding loop is a single scf.forall loop.
if (loops.size() != 1) {
return rewriter.notifyMatchFailure(
candidateSliceOp, "expected single surrounding scf.forall");
}
auto forallOp = dyn_cast<scf::ForallOp>(loops.front().getOperation());
if (!forallOp) {
return rewriter.notifyMatchFailure(
candidateSliceOp, "expected single surrounding scf.forall");
}
// 2. Fetch the corresponding output
Value sliceDest = candidateSliceOp.getDest();
auto iterArg = dyn_cast<BlockArgument>(sliceDest);
if (!iterArg)
return failure();
if (iterArg.getOwner()->getParentOp() != forallOp)
return failure();
unsigned resultNumber =
forallOp.getTiedOpResult(forallOp.getTiedOpOperand(iterArg))
.getResultNumber();
return getConsumerFromLoopUses(rewriter, forallOp, resultNumber);
}
/// A utility to fetch an untiled consumer of
/// tensor.insert_slice/tensor.parallel_insert_slice.
static FailureOr<SmallVector<OpOperand *>> getUntiledConsumerOperandsFromSlices(
RewriterBase &rewriter, ArrayRef<Operation *> sliceOps,
MutableArrayRef<LoopLikeOpInterface> loops) {
assert(!loops.empty() && "unexpected empty loops");
assert(!sliceOps.empty() && "unexpected empty list of candidate slices");
SmallVector<OpOperand *> fusedOperands;
for (auto sliceOp : sliceOps) {
FailureOr<OpOperand *> fusedOperand =
TypeSwitch<Operation *, FailureOr<OpOperand *>>(sliceOp)
.Case<tensor::InsertSliceOp, tensor::ParallelInsertSliceOp>(
[&](auto op) {
return getUntiledConsumerFromSlice(rewriter, op, loops);
})
.Default([&](Operation *op) {
return rewriter.notifyMatchFailure(op, "unhandled slice type");
});
if (failed(fusedOperand)) {
return failure();
}
if (!fusedOperands.empty() &&
fusedOperand.value()->getOwner() != fusedOperands.front()->getOwner()) {
return rewriter.notifyMatchFailure(
fusedOperand.value()->getOwner(),
"all candidate slices must be to the same consumer");
}
fusedOperands.push_back(fusedOperand.value());
}
return fusedOperands;
}
template <typename InsertSliceOpTy>
static tensor::InsertSliceOp cloneAsInsertSlice(RewriterBase &rewriter,
InsertSliceOpTy sliceOp);
template <>
tensor::InsertSliceOp
cloneAsInsertSlice<tensor::InsertSliceOp>(RewriterBase &rewriter,
tensor::InsertSliceOp insertSliceOp) {
return cast<tensor::InsertSliceOp>(
rewriter.clone(*insertSliceOp.getOperation()));
}
template <>
tensor::InsertSliceOp cloneAsInsertSlice<tensor::ParallelInsertSliceOp>(
RewriterBase &rewriter, tensor::ParallelInsertSliceOp insertSliceOp) {
return tensor::InsertSliceOp::create(
rewriter, insertSliceOp->getLoc(), insertSliceOp.getSource(),
insertSliceOp.getDest(), insertSliceOp.getMixedOffsets(),
insertSliceOp.getMixedSizes(), insertSliceOp.getMixedStrides());
}
static SmallVector<tensor::InsertSliceOp>
cloneAsInsertSlices(RewriterBase &rewriter,
ArrayRef<Operation *> candidateSlices) {
assert(!candidateSlices.empty() &&
"unexpected empty list of slices to clone");
SmallVector<tensor::InsertSliceOp> clonedSlices;
for (auto sliceOp : candidateSlices) {
TypeSwitch<Operation *>(sliceOp)
.Case<tensor::InsertSliceOp, tensor::ParallelInsertSliceOp>(
[&](auto op) {
auto clonedOp = cloneAsInsertSlice(rewriter, op);
clonedSlices.push_back(clonedOp);
})
// Assert here assuming this has already been checked.
.DefaultUnreachable(
"unexpected slice type while cloning as insert slice");
}
return clonedSlices;
}
static FailureOr<scf::SCFFuseConsumerOfSliceResult>
tileAndFuseConsumerOfSlicesImpl(RewriterBase &rewriter, Operation *consumerOp,
ArrayRef<OpOperand *> consumerOpOperands,
ArrayRef<Operation *> candidateSlices,
MutableArrayRef<LoopLikeOpInterface> loops,
const mlir::scf::InnerTileAlignmentFnTy &fn) {
assert(!loops.empty() && "expected loops to be not empty");
// Resolve the inner-tile alignment hint for the consumer in its own iteration
// domain via the control function (consulted only by pack/unpack).
SmallVector<InnerTileAlignment> innerTileAlignments;
if (fn)
if (auto consumer = dyn_cast<TilingInterface>(consumerOp))
innerTileAlignments = fn(consumer, /*tileSizes=*/{}, candidateSlices);
// 1. Check assumption for loop with `reorderOperations` disabled.
if (failed(checkAssumptionForLoop(loops.front(), consumerOp, false))) {
return rewriter.notifyMatchFailure(
loops.front(), "the first user of loop should not dominate any define "
"of consumer operand(s)");
}
LoopLikeOpInterface outerMostLoop = loops.front();
LoopLikeOpInterface innerMostLoop = loops.back();
OpBuilder::InsertionGuard g(rewriter);
// 2. Check consumer is not using scf loop's output as init.
auto dstOp = dyn_cast<DestinationStyleOpInterface>(consumerOp);
if (!dstOp)
return rewriter.notifyMatchFailure(consumerOp,
"consumer op is not DPS operation");
if (llvm::any_of(consumerOpOperands, [&](OpOperand *opOperand) {
return dstOp.isDpsInit(opOperand);
})) {
return rewriter.notifyMatchFailure(
consumerOp,
"consumer op taking the result of scf.for as init is not supported");
}
SmallVector<Value> newInits = llvm::to_vector(dstOp.getDpsInits());
// 3. Move the whole loop structure right before firstUserOfLoop, the
// dominance should be already ensured by `checkAssumptionForLoop`.
FailureOr<Operation *> firstUserOfLoop = getFirstUserOfLoop(outerMostLoop);
if (failed(firstUserOfLoop)) {
return rewriter.notifyMatchFailure(
outerMostLoop, "could not find the first user of outer most loop");
}
rewriter.moveOpBefore(outerMostLoop, *firstUserOfLoop);
// 4. Set insertion point before terminator op of the loop and create a new
// tensor.insert_slice. In the scf.for case this is a clone of the
// candidateSliceOp whereas in the scf.forall case this is created from the
// operands of tensor.parallel_insert_slice.
if (auto sliceOp =
dyn_cast<tensor::ParallelInsertSliceOp>(candidateSlices.front())) {
auto newForallOp = cast<scf::ForallOp>(innerMostLoop.getOperation());
rewriter.setInsertionPoint(newForallOp.getTerminator());
} else {
rewriter.setInsertionPoint(candidateSlices.front());
}
// 5.a. Clone all the candidate slices as equivalent insert slice ops.
SmallVector<tensor::InsertSliceOp> clonedInsertSlices =
cloneAsInsertSlices(rewriter, candidateSlices);
// 5.b. Clone consumer op.
auto clonedConsumerOp = cast<TilingInterface>(rewriter.clone(*consumerOp));
SmallVector<unsigned> operandNumbers =
llvm::map_to_vector(consumerOpOperands, [](OpOperand *opOperand) {
return opOperand->getOperandNumber();
});
SmallVector<OpOperand *> clonedOpFusedOperandsList =
llvm::map_to_vector(operandNumbers, [&](unsigned operandNum) {
return &clonedConsumerOp->getOpOperand(operandNum);
});
// 5.c. Replace all uses of the loop result with the result of the cloned
// tensor.insert_slice.
rewriter.modifyOpInPlace(clonedConsumerOp, [&]() {
for (auto [operandToReplace, clonedSliceOp] :
llvm::zip_equal(clonedOpFusedOperandsList, clonedInsertSlices)) {
operandToReplace->set(clonedSliceOp.getResult());
}
});
// 6. Perform tiling of the cloned consumer and replace the operand at
// `operandNumber` with the source of the cloned tensor.insert_slice op.
FailureOr<TilingResult> tileAndFuseResult =
tensor::replaceInsertSlicesWithTiledConsumer(rewriter, clonedInsertSlices,
clonedOpFusedOperandsList,
innerTileAlignments);
if (failed(tileAndFuseResult)) {
return failure();
}
auto tiledConsumerOp = cast<TilingInterface>(tileAndFuseResult->tiledOps[0]);
for (auto [operandNum, clonedSliceOp] :
llvm::zip_equal(operandNumbers, clonedInsertSlices)) {
rewriter.replaceAllUsesWith(tiledConsumerOp->getOperand(operandNum),
clonedSliceOp.getSource());
}
// 7. Reconstruct [nested] loop with new inits.
YieldTiledValuesFn newYieldValuesFn =
[&](RewriterBase &innerRewriter, Location loc, ValueRange /*ivs*/,
ValueRange newRegionIterArgs, SmallVector<Value> &tiledResult,
SmallVector<SmallVector<OpFoldResult>> &tiledOffset,
SmallVector<SmallVector<OpFoldResult>> &tiledSizes) -> LogicalResult {
OpBuilder::InsertionGuard g(innerRewriter);
// 8. Set inner insertPoint right before tiled consumer op.
innerRewriter.setInsertionPoint(tiledConsumerOp);
SmallVector<SmallVector<OpFoldResult>> allOffsets, allSizes;
for (auto candidateSliceOp : clonedInsertSlices) {
SmallVector<OpFoldResult> offsets = candidateSliceOp.getMixedOffsets();
SmallVector<OpFoldResult> sizes = candidateSliceOp.getMixedSizes();
SmallVector<OpFoldResult> strides = candidateSliceOp.getMixedStrides();
// 9. Check all insert stride is 1.
if (!llvm::all_of(strides, isOneInteger)) {
return rewriter.notifyMatchFailure(
candidateSliceOp, "containingOp's result yield with stride");
}
allOffsets.emplace_back(std::move(offsets));
allSizes.emplace_back(std::move(sizes));
}
// 10. Try to get iter domain position from input position. Use
// clonedConsumerOp instead of tiledConsumerOp, because the iteration
// domain may require index computation based on the result size. The
// sizes and offsets should be the same either way, but using
// tiledConsumerOp could lead to some chained unnecessary extra index
// computation.
SmallVector<OpFoldResult> iterDomainOffsets, iterDomainSizes;
if (failed(clonedConsumerOp.getIterationDomainTileFromOperandTiles(
rewriter, operandNumbers, allOffsets, allSizes, iterDomainOffsets,
iterDomainSizes, innerTileAlignments))) {
return rewriter.notifyMatchFailure(
clonedConsumerOp,
"can't get iter domain position from input position");
}
// 11. Try to fetch the offset and size for all results of the cloned
// consumer. This would then be used to form the corresponding
// tensor.insert_slice/parallel_insert_slice later.
unsigned totalNumResultsOfConsumer = tiledConsumerOp->getNumResults();
SmallVector<SmallVector<OpFoldResult>> resultOffsets(
totalNumResultsOfConsumer);
SmallVector<SmallVector<OpFoldResult>> resultSizes(
totalNumResultsOfConsumer);
for (auto [idx, v] : llvm::enumerate(tiledConsumerOp->getResults())) {
if (failed(tiledConsumerOp.getResultTilePosition(
rewriter, idx, iterDomainOffsets, iterDomainSizes,
resultOffsets[idx], resultSizes[idx]))) {
return rewriter.notifyMatchFailure(
tiledConsumerOp,
"can't get result domain position from iter domain position");
}
}
// 12. Create `extract_slice` for `iter_args` for DPS operation if
// necessary.
if (auto tiledDestStyleOp = dyn_cast<DestinationStyleOpInterface>(
tiledConsumerOp.getOperation())) {
rewriter.setInsertionPoint(tiledDestStyleOp);
for (const auto &&[index, newRegionArg] :
llvm::enumerate(newRegionIterArgs)) {
auto destSlice = tensor::ExtractSliceOp::create(
rewriter, loc, newRegionArg, resultOffsets[index],
resultSizes[index],
SmallVector<OpFoldResult>(resultOffsets[index].size(),
rewriter.getIndexAttr(1)));
// Make a copy of index to avoid a capturing structured binding, which
// is a C++20 extension.
auto dstNumber = index;
rewriter.modifyOpInPlace(tiledDestStyleOp, [&]() {
tiledDestStyleOp.getDpsInitsMutable()[dstNumber].set(destSlice);
});
}
}
// 13. Prepare tiled offset and sizes for later `insert_slice` creation by
// caller.
Block *block = rewriter.getInsertionPoint()->getBlock();
rewriter.setInsertionPoint(block->getTerminator());
for (const auto &&[index, result] :
llvm::enumerate(tiledConsumerOp->getResults())) {
tiledResult.push_back(result);
tiledOffset.emplace_back(resultOffsets[index]);
tiledSizes.emplace_back(resultSizes[index]);
}
return success();
};
// 14. Add new inits to [nested] loops.
if (failed(addInitOperandsToLoopNest(rewriter, loops, newInits,
newYieldValuesFn))) {
return rewriter.notifyMatchFailure(tiledConsumerOp,
"unable to add new inits to nest loop");
}
// 15. Replace the result of scf loop and consumer op with new loop's
// results.
for (auto &&[oldResult, newResult] :
llvm::zip(consumerOp->getResults(),
loops.front()->getResults().take_back(newInits.size()))) {
rewriter.replaceAllUsesWith(oldResult, newResult);
}
// 16. Need to erase the old scf loop and the cloned consumer op.
rewriter.eraseOp(clonedConsumerOp);
SmallVector<OpOperand *> tiledAndFusedOpOperands =
llvm::map_to_vector(operandNumbers, [&](unsigned operandNum) {
return &tileAndFuseResult->tiledOps[0]->getOpOperand(operandNum);
});
auto consumerOpOperandsVec = llvm::to_vector(consumerOpOperands);
return scf::SCFFuseConsumerOfSliceResult{
std::move(consumerOpOperandsVec), std::move(tiledAndFusedOpOperands),
std::move(tileAndFuseResult->tiledOps)};
}
/// Implementation of fusing consumer of a single slice by computing the
/// slice of the consumer in-place for scf loop.
FailureOr<scf::SCFFuseConsumerOfSliceResult>
mlir::scf::tileAndFuseConsumerOfSlices(
RewriterBase &rewriter, ArrayRef<Operation *> candidateSlices,
MutableArrayRef<LoopLikeOpInterface> loops,
const InnerTileAlignmentFnTy &fn) {
if (candidateSlices.empty()) {
return rewriter.notifyMatchFailure(
rewriter.getUnknownLoc(),
"no candidate slices provided for consumer fusion");
}
// Return if `loops` is empty, return an error for now. Caller is expected
// to handle this case.
if (loops.empty()) {
return rewriter.notifyMatchFailure(
candidateSlices.front(),
"cannot call tile and fuse consumer with an empty loop nest");
}
if (!(llvm::all_of(candidateSlices, llvm::IsaPred<tensor::InsertSliceOp>) ||
llvm::all_of(candidateSlices,
llvm::IsaPred<tensor::ParallelInsertSliceOp>))) {
return rewriter.notifyMatchFailure(
candidateSlices.front(),
"candidates slices need to be all `tensor.extract_slice`s or "
"`tensor.parallel_insert_slice`s");
}
// Get the consumer of scf.for for the result yielded by
// tensor.insert_slice/parallel_insert_slice.
FailureOr<SmallVector<OpOperand *>> maybeConsumerOpOperands =
getUntiledConsumerOperandsFromSlices(rewriter, candidateSlices, loops);
if (failed(maybeConsumerOpOperands)) {
return rewriter.notifyMatchFailure(candidateSlices.front(),
"could not fetch consumer to fuse");
}
Operation *consumerOp = maybeConsumerOpOperands->front()->getOwner();
return tileAndFuseConsumerOfSlicesImpl(rewriter, consumerOp,
maybeConsumerOpOperands.value(),
candidateSlices, loops, fn);
}
/// For a given `result` of a `forallOp` return the
/// `tensor.parallel_insert_slice` op (or combining op) that is used to
/// construct this result.
static std::optional<Operation *>
getProducingParallelInsertSlice(scf::ForallOp forallOp, OpResult result) {
if (result.getOwner() != forallOp)
return std::nullopt;
BlockArgument bbArg = forallOp.getTiedBlockArgument(result);
SmallVector<Operation *> combiningOps = forallOp.getCombiningOps(bbArg);
// If the number of combining ops is not 1, then this is unexpected. Return
// nullopt.
if (combiningOps.size() != 1)
return std::nullopt;
return combiningOps[0];
}
/// For a given result of the loop nest that is a tiled loop nest, return the
/// insert slice-like op that is used for consumer fusion
static std::optional<Operation *>
getProducingInsertSliceLikeOp(OpResult result,
ArrayRef<LoopLikeOpInterface> loops) {
assert(!loops.empty() && "Expected loops to be not empty");
LoopLikeOpInterface outerMostLoop = loops.front();
if (auto forallOp = dyn_cast<scf::ForallOp>(outerMostLoop.getOperation())) {
assert(loops.size() == 1 &&
"expected only a single loop when tiling using scf.forall");
return getProducingParallelInsertSlice(forallOp, result);
}
// Assume that the loop nest is a nested `scf.for` that is created through
// tiling and retrieve the `tensor.insert_slice` operation used to construct
// the result.
while (loops.size() != 1) {
LoopLikeOpInterface loop = loops.front();
if (result.getOwner() != loop)
return std::nullopt;
auto forOp = dyn_cast<scf::ForOp>(loop.getOperation());
if (!forOp)
return std::nullopt;
auto yieldOp = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
auto innerForResult =
dyn_cast<OpResult>(yieldOp.getOperand(result.getResultNumber()));
if (!innerForResult)
return std::nullopt;
result = innerForResult;
loops = loops.drop_front();
}
LoopLikeOpInterface loop = loops.front();
if (result.getOwner() != loop)
return std::nullopt;
auto forOp = dyn_cast<scf::ForOp>(loop.getOperation());
if (!forOp)
return std::nullopt;
auto yieldOp = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
auto insertSliceOp = yieldOp.getOperand(result.getResultNumber())
.getDefiningOp<tensor::InsertSliceOp>();
if (!insertSliceOp)
return std::nullopt;
return insertSliceOp;
}
FailureOr<scf::SCFFuseConsumerOfSliceResult>
mlir::scf::tileAndFuseConsumer(RewriterBase &rewriter, Operation *consumer,
MutableArrayRef<LoopLikeOpInterface> loops,
const InnerTileAlignmentFnTy &fn) {
if (!isa<TilingInterface>(consumer)) {
return rewriter.notifyMatchFailure(
consumer, "unhandled consumer that does not implement TilingInterface");
}
// Return if `loops` is empty, return an error for now. Caller is expected
// to handle this case.
if (loops.empty()) {
return rewriter.notifyMatchFailure(
consumer, "cannot call tile and fuse consumer with an empty loop nest");
}
LoopLikeOpInterface outermostLoop = loops.front();
// Collect the operands of the consumer that come from the outermost loop of
// the loop nest.
SmallVector<OpOperand *> consumerFusableOperands;
for (OpOperand &opOperand : consumer->getOpOperands()) {
if (opOperand.get().getDefiningOp() == outermostLoop) {
consumerFusableOperands.push_back(&opOperand);
}
}
// Nothing to fuse. Just return an empty set.
if (consumerFusableOperands.empty()) {
return mlir::scf::SCFFuseConsumerOfSliceResult{consumerFusableOperands,
SmallVector<OpOperand *>{},
SmallVector<Operation *>{}};
}
// Collect the relevant tensor.insert_slice/tensor.parallel_insert_slices
// for fusion.
SmallVector<Operation *> candidateSlices;
candidateSlices.reserve(consumerFusableOperands.size());
for (OpOperand *opOperand : consumerFusableOperands) {
std::optional<Operation *> slice =
getProducingInsertSliceLikeOp(cast<OpResult>(opOperand->get()), loops);
if (!slice) {
return rewriter.notifyMatchFailure(
consumer,
"couldnt find producing insert-slice like operation for operand");
}
candidateSlices.push_back(slice.value());
}
return tileAndFuseConsumerOfSlicesImpl(
rewriter, consumer, consumerFusableOperands, candidateSlices, loops, fn);
}
//===----------------------------------------------------------------------===//
// lowerToLoopsUsingSCFForOp implementation.
//===----------------------------------------------------------------------===//
FailureOr<SmallVector<scf::ForOp>>
mlir::scf::lowerToLoopsUsingSCFForOp(RewriterBase &rewriter,
TilingInterface op) {
// TODO: Handle cases where the op has results if needed.
if (op->getNumResults() > 0) {
return rewriter.notifyMatchFailure(
op, "unable to lower to loops operations with return values");
}
SmallVector<Range> domain = op.getIterationDomain(rewriter);
SmallVector<Value> ivs;
SmallVector<scf::ForOp> loops;
Location loc = op.getLoc();
for (auto loopRange : domain) {
Value offsetVal =
getValueOrCreateConstantIndexOp(rewriter, loc, loopRange.offset);
Value sizeVal =
getValueOrCreateConstantIndexOp(rewriter, loc, loopRange.size);
Value strideVal =
getValueOrCreateConstantIndexOp(rewriter, loc, loopRange.stride);
auto loop = scf::ForOp::create(rewriter, op.getLoc(), offsetVal, sizeVal,
strideVal, ValueRange{});
loops.push_back(loop);
ivs.push_back(loop.getInductionVar());
rewriter.setInsertionPoint(loop.getBody()->getTerminator());
}
if (failed(op.generateScalarImplementation(rewriter, op.getLoc(), ivs))) {
return failure();
}
return loops;
}