5550c82189
The MLIR classes Type/Attribute/Operation/Op/Value support cast/dyn_cast/isa/dyn_cast_or_null functionality through llvm's doCast functionality in addition to defining methods with the same name. This change begins the migration of uses of the method to the corresponding function call as has been decided as more consistent. Note that there still exist classes that only define methods directly, such as AffineExpr, and this does not include work currently to support a functional cast/isa call. Caveats include: - This clang-tidy script probably has more problems. - This only touches C++ code, so nothing that is being generated. Context: - https://mlir.llvm.org/deprecation/ at "Use the free function variants for dyn_cast/cast/isa/…" - Original discussion at https://discourse.llvm.org/t/preferred-casting-style-going-forward/68443 Implementation: This first patch was created with the following steps. The intention is to only do automated changes at first, so I waste less time if it's reverted, and so the first mass change is more clear as an example to other teams that will need to follow similar steps. Steps are described per line, as comments are removed by git: 0. Retrieve the change from the following to build clang-tidy with an additional check: https://github.com/llvm/llvm-project/compare/main...tpopp:llvm-project:tidy-cast-check 1. Build clang-tidy 2. Run clang-tidy over your entire codebase while disabling all checks and enabling the one relevant one. Run on all header files also. 3. Delete .inc files that were also modified, so the next build rebuilds them to a pure state. 4. Some changes have been deleted for the following reasons: - Some files had a variable also named cast - Some files had not included a header file that defines the cast functions - Some files are definitions of the classes that have the casting methods, so the code still refers to the method instead of the function without adding a prefix or removing the method declaration at the same time. ``` ninja -C $BUILD_DIR clang-tidy run-clang-tidy -clang-tidy-binary=$BUILD_DIR/bin/clang-tidy -checks='-*,misc-cast-functions'\ -header-filter=mlir/ mlir/* -fix rm -rf $BUILD_DIR/tools/mlir/**/*.inc git restore mlir/lib/IR mlir/lib/Dialect/DLTI/DLTI.cpp\ mlir/lib/Dialect/Complex/IR/ComplexDialect.cpp\ mlir/lib/**/IR/\ mlir/lib/Dialect/SparseTensor/Transforms/SparseVectorization.cpp\ mlir/lib/Dialect/Vector/Transforms/LowerVectorMultiReduction.cpp\ mlir/test/lib/Dialect/Test/TestTypes.cpp\ mlir/test/lib/Dialect/Transform/TestTransformDialectExtension.cpp\ mlir/test/lib/Dialect/Test/TestAttributes.cpp\ mlir/unittests/TableGen/EnumsGenTest.cpp\ mlir/test/python/lib/PythonTestCAPI.cpp\ mlir/include/mlir/IR/ ``` Differential Revision: https://reviews.llvm.org/D150123
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(isa<RankedTensorType>(srcTensor.getType()) && "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 (!isa<BlockArgument>(extractSliceOp.getSource()) &&
|
|
!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 (!isa<RankedTensorType>(srcTensor.getType()))
|
|
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 (!isa<BlockArgument>(read.getSource()) &&
|
|
!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;
|
|
}
|