4521b11397
This revision significantly rewrites hoisting on tensors. Previously, `vector.transfer_read/write` and `tensor.extract/insert_slice` would be clumped together when looking for candidate pairs. This would significantly increase the complexity of the logic and would not apply independently to `tensor.extract/insert_slice`. The new implementation decouples the cases and starts to cast the problem as a generic matching subset extract/insert, which will be future proof when other such operation pairs are introduced. Lastly, the implementation makes the distinction clear between `vector.transfer_read/write` for which we allow bypasses of the disjoint subsets from `tensor.extract/insert_slice` for which we do not yet allow it. This can be extended in the future and unified once we have subset disjunction implemented more generally. The algorithm can be rewritten to be less of a fixed point with interspersed canonicalizations. As a consequence, the test explicitly adds a canonicalization to clean up the IR and verify we end up in the same state. That extra canonicalization exhibited that one of the uses in one of the tests was dead, so we fix the appropriate test. Differential Revision: https://reviews.llvm.org/D144656
554 lines
24 KiB
C++
554 lines
24 KiB
C++
//===- SubsetHoisting.cpp - Linalg hoisting transformations----------------===//
|
|
//
|
|
// 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 functions concerned with hoisting invariant subset
|
|
// operations in the context of Linalg transformations.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Func/IR/FuncOps.h"
|
|
#include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
|
|
#include "mlir/Dialect/SCF/IR/SCF.h"
|
|
#include "mlir/Dialect/SCF/Utils/Utils.h"
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/Dialect/Utils/StaticValueUtils.h"
|
|
#include "mlir/Dialect/Vector/IR/VectorOps.h"
|
|
#include "mlir/IR/BuiltinOps.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
|
|
#include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
|
|
#include "llvm/Support/Debug.h"
|
|
#include "llvm/Support/ErrorHandling.h"
|
|
|
|
#define DEBUG_TYPE "subset-hoisting"
|
|
|
|
#define DBGS() (llvm::dbgs() << '[' << DEBUG_TYPE << "] ")
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::linalg;
|
|
|
|
/// Return true if the location of the subset defined by the op is invariant of
|
|
/// the loop iteration.
|
|
static bool
|
|
isSubsetLocationLoopInvariant(scf::ForOp forOp,
|
|
vector::TransferWriteOp transferWriteOp) {
|
|
for (Value operand : transferWriteOp.getIndices())
|
|
if (!forOp.isDefinedOutsideOfLoop(operand))
|
|
return false;
|
|
return true;
|
|
}
|
|
|
|
/// Return true if the location of the subset defined by the op is invariant of
|
|
/// the loop iteration.
|
|
static bool isSubsetLocationLoopInvariant(scf::ForOp forOp,
|
|
tensor::InsertSliceOp insertSliceOp) {
|
|
for (Value operand : insertSliceOp->getOperands().drop_front(
|
|
tensor::InsertSliceOp::getOffsetSizeAndStrideStartOperandIndex()))
|
|
if (!forOp.isDefinedOutsideOfLoop(operand))
|
|
return false;
|
|
return true;
|
|
}
|
|
|
|
/// Given an `srcTensor` that is a block argument belong to a loop.
|
|
/// Greedily look for the first read that can be hoisted out of the loop (i.e.
|
|
/// that satisfied the conditions):
|
|
/// - The read is of type `tensor.extract_slice`.
|
|
/// - The read is one of the uses of `srcTensor`.
|
|
/// - The read is to the same subset that `tensor.insert_slice` writes.
|
|
// TODO: Unify implementations once the "bypassing behavior" is the same.
|
|
static FailureOr<tensor::ExtractSliceOp>
|
|
findHoistableMatchingExtractSlice(RewriterBase &rewriter,
|
|
tensor::InsertSliceOp insertSliceOp,
|
|
BlockArgument srcTensor) {
|
|
assert(srcTensor.getType().isa<RankedTensorType>() && "not a ranked tensor");
|
|
|
|
auto forOp = cast<scf::ForOp>(srcTensor.getOwner()->getParentOp());
|
|
|
|
LLVM_DEBUG(DBGS() << "--find matching read for: " << insertSliceOp << "\n";
|
|
DBGS() << "--amongst users of: " << srcTensor << "\n");
|
|
|
|
SmallVector<Operation *> users(srcTensor.getUsers());
|
|
if (forOp.isDefinedOutsideOfLoop(insertSliceOp.getDest()))
|
|
llvm::append_range(users, insertSliceOp.getDest().getUsers());
|
|
|
|
for (Operation *user : users) {
|
|
LLVM_DEBUG(DBGS() << "----inspect user: " << *user << "\n");
|
|
auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(user);
|
|
// Skip ops other than extract_slice with an exact matching of their tensor
|
|
// subset.
|
|
if (extractSliceOp) {
|
|
auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; };
|
|
if (extractSliceOp.getResultType() != insertSliceOp.getSourceType() ||
|
|
!extractSliceOp.isSameAs(insertSliceOp, isSame)) {
|
|
LLVM_DEBUG(DBGS() << "------not a matching extract_slice\n";
|
|
DBGS() << *user << " vs " << *insertSliceOp << "\n");
|
|
continue;
|
|
}
|
|
|
|
// Skip insert_slice whose vector is defined within the loop: we need to
|
|
// hoist that definition first otherwise dominance violations trigger.
|
|
if (!extractSliceOp.getSource().isa<BlockArgument>() &&
|
|
!forOp.isDefinedOutsideOfLoop(extractSliceOp.getSource())) {
|
|
LLVM_DEBUG(DBGS() << "------transfer_read vector is loop-dependent\n");
|
|
continue;
|
|
}
|
|
return extractSliceOp;
|
|
}
|
|
|
|
// TODO: Look through disjoint subsets, similar to vector.transfer_write
|
|
// and unify implementations.
|
|
}
|
|
|
|
LLVM_DEBUG(DBGS() << "----no matching extract_slice");
|
|
return failure();
|
|
}
|
|
|
|
/// Given an `srcTensor` that is a block argument belong to a loop.
|
|
/// Greedily look for the first read that can be hoisted out of the loop (i.e.
|
|
/// that satisfied the conditions):
|
|
/// - The read is of type `tensor.transfer_read`.
|
|
/// - The read is one of the uses of `srcTensor`.
|
|
/// - The read is to the same subset that `tensor.transfer_write` writes.
|
|
// TODO: Unify implementations once the "bypassing behavior" is the same.
|
|
static FailureOr<vector::TransferReadOp>
|
|
findHoistableMatchingTransferRead(RewriterBase &rewriter,
|
|
vector::TransferWriteOp transferWriteOp,
|
|
BlockArgument srcTensor) {
|
|
if (!srcTensor.getType().isa<RankedTensorType>())
|
|
return failure();
|
|
|
|
auto forOp = cast<scf::ForOp>(srcTensor.getOwner()->getParentOp());
|
|
|
|
LLVM_DEBUG(DBGS() << "--find matching read for: " << transferWriteOp << "\n";
|
|
DBGS() << "--amongst users of: " << srcTensor << "\n";);
|
|
|
|
// vector.transfer_write is a bit peculiar: we look through dependencies
|
|
// to disjoint tensor subsets. This requires a while loop.
|
|
// TODO: Look through disjoint subsets for tensor.insert_slice and unify
|
|
// implementations.
|
|
SmallVector<Operation *> users(srcTensor.getUsers());
|
|
// TODO: transferWriteOp.getSource is actually the destination tensor!!
|
|
if (forOp.isDefinedOutsideOfLoop(transferWriteOp.getSource()))
|
|
llvm::append_range(users, transferWriteOp.getSource().getUsers());
|
|
while (!users.empty()) {
|
|
Operation *user = users.pop_back_val();
|
|
LLVM_DEBUG(DBGS() << "----inspect user: " << *user << "\n");
|
|
auto read = dyn_cast<vector::TransferReadOp>(user);
|
|
if (read) {
|
|
// Skip ops other than transfer_read with an exact matching subset.
|
|
if (read.getIndices() != transferWriteOp.getIndices() ||
|
|
read.getVectorType() != transferWriteOp.getVectorType()) {
|
|
LLVM_DEBUG(DBGS() << "------not a transfer_read that matches the "
|
|
"transfer_write: "
|
|
<< *user << "\n\t(vs " << *transferWriteOp << ")\n");
|
|
continue;
|
|
}
|
|
|
|
// transfer_read may be of a vector that is defined within the loop: we
|
|
// traverse it by virtue of bypassing disjoint subset operations rooted at
|
|
// a bbArg and yielding a matching yield.
|
|
if (!read.getSource().isa<BlockArgument>() &&
|
|
!forOp.isDefinedOutsideOfLoop(read.getSource())) {
|
|
LLVM_DEBUG(DBGS() << "------transfer_read vector appears loop "
|
|
"dependent but will be tested for disjointness as "
|
|
"part of the bypass analysis\n");
|
|
}
|
|
LLVM_DEBUG(DBGS() << "------found match\n");
|
|
return read;
|
|
}
|
|
|
|
// As an optimization, we look further through dependencies to disjoint
|
|
// tensor subsets. This creates more opportunities to find a matching read.
|
|
if (isa<vector::TransferWriteOp>(user)) {
|
|
// If we find a write with disjoint indices append all its uses.
|
|
// TODO: Generalize areSubsetsDisjoint and allow other bypass than
|
|
// just vector.transfer_write - vector.transfer_write.
|
|
if (vector::isDisjointTransferIndices(
|
|
cast<VectorTransferOpInterface>(user),
|
|
cast<VectorTransferOpInterface>(
|
|
transferWriteOp.getOperation()))) {
|
|
LLVM_DEBUG(DBGS() << "----follow through disjoint write\n");
|
|
users.append(user->getUsers().begin(), user->getUsers().end());
|
|
} else {
|
|
LLVM_DEBUG(DBGS() << "----skip non-disjoint write\n");
|
|
}
|
|
}
|
|
}
|
|
|
|
LLVM_DEBUG(DBGS() << "--no matching transfer_read\n");
|
|
return rewriter.notifyMatchFailure(transferWriteOp,
|
|
"no matching transfer_read");
|
|
}
|
|
|
|
/// Return the `vector.transfer_write` that produces `yieldOperand`, if:
|
|
/// - The write operates on tensors.
|
|
/// - All indices are defined outside of the loop.
|
|
/// Return failure otherwise.
|
|
///
|
|
/// This is sufficient condition to hoist the `vector.transfer_write`; other
|
|
/// operands can always be yielded by the loop where needed.
|
|
// TODO: generalize beyond scf::ForOp.
|
|
// TODO: Unify implementations once the "bypassing behavior" is the same.
|
|
static FailureOr<vector::TransferWriteOp>
|
|
getLoopInvariantTransferWriteDefining(RewriterBase &rewriter, scf::ForOp forOp,
|
|
BlockArgument bbArg,
|
|
OpOperand &yieldOperand) {
|
|
assert(bbArg.getArgNumber() ==
|
|
forOp.getNumInductionVars() + yieldOperand.getOperandNumber() &&
|
|
"bbArg and yieldOperand must match");
|
|
assert(isa<scf::YieldOp>(yieldOperand.getOwner()) && "must be an scf.yield");
|
|
|
|
Value v = yieldOperand.get();
|
|
auto transferWriteOp = v.getDefiningOp<vector::TransferWriteOp>();
|
|
if (!transferWriteOp)
|
|
return rewriter.notifyMatchFailure(v.getLoc(), "not a transfer_write");
|
|
|
|
if (transferWriteOp->getNumResults() == 0) {
|
|
return rewriter.notifyMatchFailure(v.getLoc(),
|
|
"unsupported transfer_write on buffers");
|
|
}
|
|
|
|
// We do not explicitly check that the destination is a BBarg that matches the
|
|
// yield operand as this would prevent us from bypassing other non-conflicting
|
|
// writes.
|
|
|
|
// Indexing must not depend on `forOp`.
|
|
if (!isSubsetLocationLoopInvariant(forOp, transferWriteOp))
|
|
return rewriter.notifyMatchFailure(
|
|
v.getLoc(), "transfer_write indexing is loop-dependent");
|
|
|
|
return transferWriteOp;
|
|
}
|
|
|
|
/// Return the `tensor.insert_slice` that produces `yieldOperand`, if:
|
|
/// 1. Its destination tensor is a block argument of the `forOp`.
|
|
/// 2. The unique use of its result is a yield with operand number matching
|
|
/// the block argument.
|
|
/// 3. All indices are defined outside of the loop.
|
|
/// Return failure otherwise.
|
|
///
|
|
/// This is sufficient condition to hoist the `tensor.insert_slice`; other
|
|
/// operands can always be yielded by the loop where needed.
|
|
/// Note: 1. + 2. ensure that the yield / iter_args cycle results in proper
|
|
/// semantics (i.e. no ping-ping between iter_args across iterations).
|
|
// TODO: generalize beyond scf::ForOp.
|
|
// TODO: Unify implementations once the "bypassing behavior" is the same.
|
|
static FailureOr<tensor::InsertSliceOp>
|
|
getLoopInvariantInsertSliceDefining(RewriterBase &rewriter, scf::ForOp forOp,
|
|
BlockArgument bbArg,
|
|
OpOperand &yieldOperand) {
|
|
assert(bbArg.getArgNumber() ==
|
|
forOp.getNumInductionVars() + yieldOperand.getOperandNumber() &&
|
|
"bbArg and yieldOperand must match");
|
|
assert(isa<scf::YieldOp>(yieldOperand.getOwner()) && "must be an scf.yield");
|
|
|
|
Value v = yieldOperand.get();
|
|
auto insertSliceOp = v.getDefiningOp<tensor::InsertSliceOp>();
|
|
if (!insertSliceOp)
|
|
return rewriter.notifyMatchFailure(v.getLoc(), "not an insert_slice");
|
|
|
|
// Tensor inserted into must be a BBArg at position matching yield operand.
|
|
// TODO: In the future we should not perform this check if we want to bypass
|
|
// other non-conflicting writes.
|
|
if (bbArg != insertSliceOp.getDest())
|
|
return rewriter.notifyMatchFailure(v.getLoc(), "not a matching bbarg");
|
|
|
|
// Indexing inserted into must not depend on `forOp`.
|
|
if (!isSubsetLocationLoopInvariant(forOp, insertSliceOp))
|
|
return rewriter.notifyMatchFailure(
|
|
v.getLoc(), "insert_slice indexing is loop-dependent");
|
|
|
|
return insertSliceOp;
|
|
}
|
|
|
|
/// Check if the chunk of data inserted by the `writeOp` is read by any other
|
|
/// op than the candidateReadOp. This conflicting operation prevents hoisting,
|
|
/// return it or nullptr if none is found.
|
|
// TODO: Generalize subset disjunction analysis/interface.
|
|
// TODO: Support more subset op types.
|
|
static Operation *isTensorChunkAccessedByUnknownOp(Operation *writeOp,
|
|
Operation *candidateReadOp,
|
|
BlockArgument tensorArg) {
|
|
// Make sure none of the other uses read the part of the tensor modified
|
|
// by the transfer_write.
|
|
llvm::SmallVector<Value::use_range, 1> uses;
|
|
uses.push_back(tensorArg.getUses());
|
|
while (!uses.empty()) {
|
|
for (OpOperand &use : uses.pop_back_val()) {
|
|
Operation *user = use.getOwner();
|
|
// Skip the candidate use, only inspect the "other" uses.
|
|
if (user == candidateReadOp || user == writeOp)
|
|
continue;
|
|
|
|
// TODO: Consider all transitive uses through
|
|
// extract_slice/insert_slice. Atm we just bail because a stronger
|
|
// analysis is needed for these cases.
|
|
if (isa<tensor::ExtractSliceOp, tensor::InsertSliceOp>(user))
|
|
return user;
|
|
|
|
// Consider all transitive uses through a vector.transfer_write.
|
|
if (isa<vector::TransferWriteOp>(writeOp)) {
|
|
if (auto writeUser = dyn_cast<vector::TransferWriteOp>(user)) {
|
|
uses.push_back(writeUser->getResult(0).getUses());
|
|
continue;
|
|
}
|
|
}
|
|
|
|
// Consider all nested uses through an scf::ForOp. We may have
|
|
// pass-through tensor arguments left from previous level of
|
|
// hoisting.
|
|
if (auto forUser = dyn_cast<scf::ForOp>(user)) {
|
|
Value arg = forUser.getLoopBody().getArgument(
|
|
use.getOperandNumber() - forUser.getNumControlOperands() +
|
|
/*iv value*/ 1);
|
|
uses.push_back(arg.getUses());
|
|
continue;
|
|
}
|
|
|
|
// Follow the use yield, only if it doesn't escape the original region.
|
|
scf::YieldOp yieldUser = dyn_cast<scf::YieldOp>(user);
|
|
if (yieldUser &&
|
|
writeOp->getParentOp()->isAncestor(yieldUser->getParentOp())) {
|
|
Value ret = yieldUser->getParentOp()->getResult(use.getOperandNumber());
|
|
uses.push_back(ret.getUses());
|
|
continue;
|
|
}
|
|
|
|
// If the write is a vector::TransferWriteOp, it may have been bypassed
|
|
// and we need to check subset disjunction
|
|
if (isa<vector::TransferWriteOp>(writeOp)) {
|
|
auto read = dyn_cast<vector::TransferReadOp>(user);
|
|
if (!read || !vector::isDisjointTransferIndices(
|
|
cast<VectorTransferOpInterface>(read.getOperation()),
|
|
cast<VectorTransferOpInterface>(writeOp))) {
|
|
return user;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
/// Mechanical hoisting of a matching read / write pair.
|
|
/// Return the newly created scf::ForOp with an extra yields.
|
|
// TODO: Unify implementations once the "bypassing behavior" is the same.
|
|
static scf::ForOp hoistTransferReadWrite(
|
|
RewriterBase &rewriter, vector::TransferReadOp transferReadOp,
|
|
vector::TransferWriteOp transferWriteOp, BlockArgument tensorBBArg) {
|
|
scf::ForOp forOp = cast<scf::ForOp>(tensorBBArg.getOwner()->getParentOp());
|
|
LLVM_DEBUG(DBGS() << "--Start hoisting\n";
|
|
DBGS() << "--Hoist read : " << transferReadOp << "\n";
|
|
DBGS() << "--Hoist write: " << transferWriteOp << "\n";
|
|
DBGS() << "--Involving : " << tensorBBArg << "\n");
|
|
|
|
// TODO: don't hardcode /*numIvs=*/1.
|
|
assert(tensorBBArg.getArgNumber() >= /*numIvs=*/1);
|
|
int64_t initArgNumber = tensorBBArg.getArgNumber() - /*numIvs=*/1;
|
|
|
|
// 1. Hoist the read op. Thanks to our previous checks we know this will not
|
|
// trigger dominance violations once BBArgs are updated.
|
|
// TODO: should the rewriter ever want to track this move ?
|
|
transferReadOp->moveBefore(forOp);
|
|
if (!forOp.isDefinedOutsideOfLoop(transferReadOp.getSource())) {
|
|
rewriter.startRootUpdate(transferReadOp);
|
|
transferReadOp.getSourceMutable().assign(
|
|
forOp.getInitArgs()[initArgNumber]);
|
|
rewriter.finalizeRootUpdate(transferReadOp);
|
|
}
|
|
|
|
// 2. Rewrite `loop` with an additional yield. This is the quantity that is
|
|
// computed iteratively but whose storage has become loop-invariant.
|
|
NewYieldValueFn yieldFn = [&](OpBuilder &b, Location loc,
|
|
ArrayRef<BlockArgument> newBBArgs) {
|
|
return SmallVector<Value>{transferWriteOp.getVector()};
|
|
};
|
|
auto newForOp = replaceLoopWithNewYields(
|
|
rewriter, forOp, {transferReadOp.getVector()}, yieldFn);
|
|
rewriter.eraseOp(forOp);
|
|
|
|
// 3. Update the yield. Invariant: initArgNumber is the destination tensor.
|
|
auto yieldOp =
|
|
cast<scf::YieldOp>(newForOp.getRegion().front().getTerminator());
|
|
// TODO: transferWriteOp.getSource is actually the destination tensor!!
|
|
rewriter.startRootUpdate(yieldOp);
|
|
yieldOp->setOperand(initArgNumber, transferWriteOp.getSource());
|
|
rewriter.finalizeRootUpdate(yieldOp);
|
|
|
|
// 4. Hoist write after and make uses of newForOp.getResult(initArgNumber)
|
|
// flow through it.
|
|
// TODO: should the rewriter ever want to track this move ?
|
|
transferWriteOp->moveAfter(newForOp);
|
|
rewriter.startRootUpdate(transferWriteOp);
|
|
transferWriteOp.getVectorMutable().assign(newForOp.getResults().back());
|
|
// TODO: transferWriteOp.getSource is actually the destination tensor!!
|
|
transferWriteOp.getSourceMutable().assign(newForOp.getResult(initArgNumber));
|
|
rewriter.finalizeRootUpdate(transferWriteOp);
|
|
rewriter.replaceAllUsesExcept(newForOp.getResult(initArgNumber),
|
|
transferWriteOp.getResult(), transferWriteOp);
|
|
return newForOp;
|
|
}
|
|
|
|
/// Mechanical hoisting of a matching read / write pair.
|
|
/// Return the newly created scf::ForOp with an extra yields.
|
|
// TODO: Unify implementations once the "bypassing behavior" is the same.
|
|
static scf::ForOp hoistExtractInsertSlice(RewriterBase &rewriter,
|
|
tensor::ExtractSliceOp extractSliceOp,
|
|
tensor::InsertSliceOp insertSliceOp,
|
|
BlockArgument tensorBBArg) {
|
|
scf::ForOp forOp = cast<scf::ForOp>(tensorBBArg.getOwner()->getParentOp());
|
|
LLVM_DEBUG(DBGS() << "--Start hoisting\n";
|
|
DBGS() << "--Hoist read : " << extractSliceOp << "\n";
|
|
DBGS() << "--Hoist write: " << insertSliceOp << "\n";
|
|
DBGS() << "--Involving : " << tensorBBArg << "\n");
|
|
|
|
// TODO: don't hardcode /*numIvs=*/1.
|
|
assert(tensorBBArg.getArgNumber() >= /*numIvs=*/1);
|
|
int64_t initArgNumber = tensorBBArg.getArgNumber() - /*numIvs=*/1;
|
|
|
|
// 1. Hoist the read op. Thanks to our previous checks we know this will not
|
|
// trigger dominance violations once BBArgs are updated.
|
|
// TODO: should the rewriter ever want to track this move ?
|
|
extractSliceOp->moveBefore(forOp);
|
|
if (!forOp.isDefinedOutsideOfLoop(extractSliceOp.getSource())) {
|
|
assert(extractSliceOp.getSource() == tensorBBArg &&
|
|
"extractSlice source not defined above must be the tracked bbArg");
|
|
rewriter.startRootUpdate(extractSliceOp);
|
|
extractSliceOp.getSourceMutable().assign(
|
|
forOp.getInitArgs()[initArgNumber]);
|
|
rewriter.finalizeRootUpdate(extractSliceOp);
|
|
}
|
|
|
|
// 2. Rewrite `loop` with an additional yield. This is the quantity that is
|
|
// computed iteratively but whose storage has become loop-invariant.
|
|
NewYieldValueFn yieldFn = [&](OpBuilder &b, Location loc,
|
|
ArrayRef<BlockArgument> newBBArgs) {
|
|
return SmallVector<Value>{insertSliceOp.getSource()};
|
|
};
|
|
auto newForOp = replaceLoopWithNewYields(rewriter, forOp,
|
|
extractSliceOp.getResult(), yieldFn);
|
|
rewriter.eraseOp(forOp);
|
|
|
|
// 3. Update the yield. Invariant: initArgNumber is the destination tensor.
|
|
auto yieldOp =
|
|
cast<scf::YieldOp>(newForOp.getRegion().front().getTerminator());
|
|
// TODO: should the rewriter ever want to track this ?
|
|
rewriter.startRootUpdate(yieldOp);
|
|
yieldOp->setOperand(initArgNumber, insertSliceOp.getDest());
|
|
rewriter.finalizeRootUpdate(yieldOp);
|
|
|
|
// 4. Hoist write after and make uses of newForOp.getResult(initArgNumber)
|
|
// flow through it.
|
|
// TODO: should the rewriter ever want to track this move ?
|
|
insertSliceOp->moveAfter(newForOp);
|
|
rewriter.startRootUpdate(insertSliceOp);
|
|
insertSliceOp.getSourceMutable().assign(newForOp.getResults().back());
|
|
insertSliceOp.getDestMutable().assign(newForOp.getResult(initArgNumber));
|
|
rewriter.finalizeRootUpdate(insertSliceOp);
|
|
rewriter.replaceAllUsesExcept(newForOp.getResult(initArgNumber),
|
|
insertSliceOp.getResult(), insertSliceOp);
|
|
return newForOp;
|
|
}
|
|
|
|
/// Greedily hoist redundant subset extract/insert operations on tensors
|
|
/// outside `forOp`.
|
|
/// Return the unmodified `forOp` if no hoisting occurred.
|
|
/// Return a new scf::ForOp if hoisting on tensors occurred.
|
|
scf::ForOp
|
|
mlir::linalg::hoistRedundantSubsetExtractInsert(RewriterBase &rewriter,
|
|
scf::ForOp forOp) {
|
|
LLVM_DEBUG(DBGS() << "Enter hoistRedundantSubsetExtractInsert scf.for\n");
|
|
Operation *yield = forOp.getBody()->getTerminator();
|
|
|
|
LLVM_DEBUG(DBGS() << "\n"; DBGS() << "Consider " << forOp << "\n");
|
|
|
|
scf::ForOp newForOp = forOp;
|
|
do {
|
|
forOp = newForOp;
|
|
for (const auto &it : llvm::enumerate(forOp.getRegionIterArgs())) {
|
|
LLVM_DEBUG(DBGS() << "Consider " << it.value() << "\n");
|
|
|
|
// 1. Find a loop invariant subset write yielding `ret` that we can
|
|
// consider for hoisting.
|
|
// TODO: TypeSwitch when we add more cases.
|
|
OpOperand &ret = yield->getOpOperand(it.index());
|
|
FailureOr<vector::TransferWriteOp> transferWriteOp =
|
|
getLoopInvariantTransferWriteDefining(rewriter, forOp, it.value(),
|
|
ret);
|
|
FailureOr<tensor::InsertSliceOp> insertSliceOp =
|
|
getLoopInvariantInsertSliceDefining(rewriter, forOp, it.value(), ret);
|
|
if (failed(transferWriteOp) && failed(insertSliceOp)) {
|
|
LLVM_DEBUG(DBGS() << "no loop invariant write defining iter_args "
|
|
<< it.value() << "\n");
|
|
continue;
|
|
}
|
|
|
|
Operation *writeOp = succeeded(transferWriteOp)
|
|
? transferWriteOp->getOperation()
|
|
: insertSliceOp->getOperation();
|
|
|
|
// 2. Only accept writes with a single use (i.e. the yield).
|
|
if (!writeOp->hasOneUse()) {
|
|
LLVM_DEBUG(DBGS() << "write with more than 1 use " << *writeOp << "\n");
|
|
continue;
|
|
}
|
|
|
|
LLVM_DEBUG(DBGS() << "Write to hoist: " << *writeOp << "\n");
|
|
|
|
// 3. Find a matching read that can also be hoisted.
|
|
Operation *matchingReadOp = nullptr;
|
|
// TODO: TypeSwitch.
|
|
if (succeeded(transferWriteOp)) {
|
|
auto maybeTransferRead = findHoistableMatchingTransferRead(
|
|
rewriter, *transferWriteOp, it.value());
|
|
if (succeeded(maybeTransferRead))
|
|
matchingReadOp = maybeTransferRead->getOperation();
|
|
} else if (succeeded(insertSliceOp)) {
|
|
auto maybeExtractSlice = findHoistableMatchingExtractSlice(
|
|
rewriter, *insertSliceOp, it.value());
|
|
if (succeeded(maybeExtractSlice))
|
|
matchingReadOp = maybeExtractSlice->getOperation();
|
|
} else {
|
|
llvm_unreachable("unexpected case");
|
|
}
|
|
if (!matchingReadOp) {
|
|
LLVM_DEBUG(DBGS() << "No matching read\n");
|
|
continue;
|
|
}
|
|
|
|
// 4. Make sure no other use reads the part of the modified tensor.
|
|
// This is necessary to guard against hazards when non-conflicting subset
|
|
// ops are bypassed.
|
|
Operation *maybeUnknownOp =
|
|
isTensorChunkAccessedByUnknownOp(writeOp, matchingReadOp, it.value());
|
|
if (maybeUnknownOp) {
|
|
LLVM_DEBUG(DBGS() << "Tensor chunk accessed by unknown op, skip: "
|
|
<< *maybeUnknownOp << "\n");
|
|
continue;
|
|
}
|
|
|
|
// 5. Perform the actual mechanical hoisting.
|
|
// TODO: TypeSwitch.
|
|
LLVM_DEBUG(DBGS() << "Read to hoist: " << *matchingReadOp << "\n");
|
|
if (succeeded(transferWriteOp)) {
|
|
newForOp = hoistTransferReadWrite(
|
|
rewriter, cast<vector::TransferReadOp>(matchingReadOp),
|
|
*transferWriteOp, it.value());
|
|
} else if (succeeded(insertSliceOp)) {
|
|
newForOp = hoistExtractInsertSlice(
|
|
rewriter, cast<tensor::ExtractSliceOp>(matchingReadOp),
|
|
*insertSliceOp, it.value());
|
|
} else {
|
|
llvm_unreachable("unexpected case");
|
|
}
|
|
break;
|
|
}
|
|
} while (forOp != newForOp);
|
|
|
|
return newForOp;
|
|
}
|