36d8e7056e
Enable ops with only read side effects in scf.for to be hoisted with a scf.if guard that checks against the trip count This patch takes a step towards a less conservative LICM in MLIR as discussed in the following discourse thread: [Speculative LICM?](https://discourse.llvm.org/t/speculative-licm/80977) This patch in particular does the following: 1. Relaxes the original constraint for hoisting that only hoists ops without any side effects. This patch also allows the ops with only read side effects to be hoisted into an scf.if guard only if every op in the loop or its nested regions is side-effect free or has only read side effects. This scf.if guard wraps the original scf.for and checks for **trip_count > 0**. 2. To support this, two new interface methods are added to **LoopLikeInterface**: _wrapInTripCountCheck_ and _unwrapTripCountCheck_. Implementation starts with wrapping the scf.for loop into scf.if guard using _wrapInTripCountCheck_ and if there is no op hoisted into the this guard after we are done processing the worklist, it unwraps the guard by calling _unwrapTripCountCheck_.
424 lines
17 KiB
C++
424 lines
17 KiB
C++
//===- LoopInvariantCodeMotionUtils.cpp - LICM Utils ------------*- C++ -*-===//
|
|
//
|
|
// 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 contains the implementation of the core LICM algorithm.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Transforms/LoopInvariantCodeMotionUtils.h"
|
|
|
|
#include "mlir/IR/Operation.h"
|
|
#include "mlir/IR/PatternMatch.h"
|
|
#include "mlir/Interfaces/LoopLikeInterface.h"
|
|
#include "mlir/Interfaces/SideEffectInterfaces.h"
|
|
#include "mlir/Interfaces/SubsetOpInterface.h"
|
|
#include "llvm/Support/Debug.h"
|
|
#include <queue>
|
|
|
|
#define DEBUG_TYPE "licm"
|
|
|
|
using namespace mlir;
|
|
|
|
/// Checks whether the given op can be hoisted by checking that
|
|
/// - the op and none of its contained operations depend on values inside of the
|
|
/// loop (by means of calling definedOutside).
|
|
/// - the op has no side-effects.
|
|
static bool canBeHoisted(Operation *op,
|
|
function_ref<bool(OpOperand &)> condition) {
|
|
// Do not move terminators.
|
|
if (op->hasTrait<OpTrait::IsTerminator>())
|
|
return false;
|
|
|
|
// Walk the nested operations and check that all used values are either
|
|
// defined outside of the loop or in a nested region, but not at the level of
|
|
// the loop body.
|
|
auto walkFn = [&](Operation *child) {
|
|
for (OpOperand &operand : child->getOpOperands()) {
|
|
// Ignore values defined in a nested region.
|
|
if (op->isAncestor(operand.get().getParentRegion()->getParentOp()))
|
|
continue;
|
|
if (!condition(operand))
|
|
return WalkResult::interrupt();
|
|
}
|
|
return WalkResult::advance();
|
|
};
|
|
return !op->walk(walkFn).wasInterrupted();
|
|
}
|
|
|
|
static bool canBeHoisted(Operation *op,
|
|
function_ref<bool(Value)> definedOutside) {
|
|
return canBeHoisted(
|
|
op, [&](OpOperand &operand) { return definedOutside(operand.get()); });
|
|
}
|
|
|
|
size_t mlir::moveLoopInvariantCode(
|
|
ArrayRef<Region *> regions,
|
|
function_ref<bool(Value, Region *)> isDefinedOutsideRegion,
|
|
function_ref<bool(Operation *, Region *)> shouldMoveOutOfRegion,
|
|
function_ref<void(Operation *)> moveOutOfRegionWithoutGuard,
|
|
function_ref<void(Operation *)> moveOutOfRegionWithGuard) {
|
|
size_t numMoved = 0;
|
|
|
|
for (Region *region : regions) {
|
|
LLVM_DEBUG(llvm::dbgs() << "Original loop:\n"
|
|
<< *region->getParentOp() << "\n");
|
|
|
|
bool anyOpHoistedWithGuard = false;
|
|
bool loopSideEffectFreeOrHasOnlyReadSideEffect =
|
|
isMemoryEffectFreeOrOnlyRead(region->getParentOp());
|
|
|
|
std::queue<Operation *> worklist;
|
|
// Add top-level operations in the loop body to the worklist.
|
|
for (Operation &op : region->getOps())
|
|
worklist.push(&op);
|
|
|
|
auto definedOutside = [&](Value value) {
|
|
return isDefinedOutsideRegion(value, region);
|
|
};
|
|
|
|
while (!worklist.empty()) {
|
|
Operation *op = worklist.front();
|
|
worklist.pop();
|
|
// Skip ops that have already been moved. Check if the op can be hoisted.
|
|
if (op->getParentRegion() != region)
|
|
continue;
|
|
|
|
LLVM_DEBUG(llvm::dbgs() << "Checking op: " << *op << "\n");
|
|
|
|
if (!shouldMoveOutOfRegion(op, region) ||
|
|
!canBeHoisted(op, definedOutside))
|
|
continue;
|
|
// Can only hoist pure ops (side-effect free) when there is an op with
|
|
// write and/or unknown side effects in the loop.
|
|
if (!loopSideEffectFreeOrHasOnlyReadSideEffect && !isMemoryEffectFree(op))
|
|
continue;
|
|
|
|
bool moveWithoutGuard = !anyOpHoistedWithGuard && isMemoryEffectFree(op);
|
|
if (moveWithoutGuard) {
|
|
LLVM_DEBUG(llvm::dbgs() << "Moving loop-invariant op: " << *op
|
|
<< " without guard\n");
|
|
moveOutOfRegionWithoutGuard(op);
|
|
} else {
|
|
LLVM_DEBUG(llvm::dbgs()
|
|
<< "Moving loop-invariant op: " << *op << " with guard\n");
|
|
moveOutOfRegionWithGuard(op);
|
|
anyOpHoistedWithGuard = true;
|
|
}
|
|
++numMoved;
|
|
|
|
// Since the op has been moved, we need to check its users within the
|
|
// top-level of the loop body.
|
|
for (Operation *user : op->getUsers())
|
|
if (user->getParentRegion() == region)
|
|
worklist.push(user);
|
|
}
|
|
}
|
|
|
|
return numMoved;
|
|
}
|
|
|
|
size_t mlir::moveLoopInvariantCode(LoopLikeOpInterface loopLike) {
|
|
return moveLoopInvariantCode(
|
|
loopLike.getLoopRegions(),
|
|
[&](Value value, Region *region) {
|
|
return !region->isAncestor(value.getParentRegion());
|
|
},
|
|
[&](Operation *op, Region *) {
|
|
return isSpeculatable(op) && isMemoryEffectFreeOrOnlyRead(op);
|
|
},
|
|
[&](Operation *op) { loopLike.moveOutOfLoop(op); },
|
|
[&](Operation *op) { loopLike.moveOutOfLoopWithGuard(op); });
|
|
}
|
|
|
|
namespace {
|
|
/// Helper data structure that keeps track of equivalent/disjoint subset ops.
|
|
class MatchingSubsets {
|
|
public:
|
|
/// Insert a subset op.
|
|
void insert(SubsetOpInterface op, bool collectHoistableOps = true) {
|
|
allSubsetOps.push_back(op);
|
|
if (!collectHoistableOps)
|
|
return;
|
|
if (auto extractionOp =
|
|
dyn_cast<SubsetExtractionOpInterface>(op.getOperation()))
|
|
insertExtractionOp(extractionOp);
|
|
if (auto insertionOp =
|
|
dyn_cast<SubsetInsertionOpInterface>(op.getOperation()))
|
|
insertInsertionOp(insertionOp);
|
|
}
|
|
|
|
/// Return a range of matching extraction-insertion subset ops. If there is no
|
|
/// matching extraction/insertion op, the respective value is empty. Ops are
|
|
/// skipped if there are other subset ops that are not guaranteed to operate
|
|
/// on disjoint subsets.
|
|
auto getHoistableSubsetOps() {
|
|
return llvm::make_filter_range(
|
|
llvm::zip(extractions, insertions), [&](auto pair) {
|
|
auto [extractionOp, insertionOp] = pair;
|
|
// Hoist only if the extracted and inserted values have the same type.
|
|
if (extractionOp && insertionOp &&
|
|
extractionOp->getResult(0).getType() !=
|
|
insertionOp.getSourceOperand().get().getType())
|
|
return false;
|
|
// Hoist only if there are no conflicting subset ops.
|
|
return allDisjoint(extractionOp, insertionOp);
|
|
});
|
|
}
|
|
|
|
/// Populate subset ops starting from the given region iter_arg. Return
|
|
/// "failure" if non-subset ops are found along the path to the loop yielding
|
|
/// op or if there is no single path to the tied yielded operand. If
|
|
/// `collectHoistableOps` is set to "false", subset ops are gathered
|
|
/// throughout the traversal, but not enumerated by `getHoistableSubsetOps`.
|
|
LogicalResult populateSubsetOpsAtIterArg(LoopLikeOpInterface loopLike,
|
|
BlockArgument iterArg,
|
|
bool collectHoistableOps = true);
|
|
|
|
private:
|
|
/// Helper function for equivalence of tensor values. Since only insertion
|
|
/// subset ops (that are also destination style ops) are followed when
|
|
/// traversing the SSA use-def chain, all tensor values are equivalent.
|
|
static bool isEquivalent(Value v1, Value v2) { return true; }
|
|
|
|
/// Return "true" if the subsets of the given extraction and insertion ops
|
|
/// are operating disjoint from the subsets that all other known subset ops
|
|
/// are operating on.
|
|
bool allDisjoint(SubsetExtractionOpInterface extractionOp,
|
|
SubsetInsertionOpInterface insertionOp) const {
|
|
for (SubsetOpInterface other : allSubsetOps) {
|
|
if (other == extractionOp || other == insertionOp)
|
|
continue;
|
|
if (extractionOp &&
|
|
!other.operatesOnDisjointSubset(extractionOp, isEquivalent))
|
|
return false;
|
|
if (insertionOp &&
|
|
!other.operatesOnDisjointSubset(insertionOp, isEquivalent))
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
/// Insert a subset extraction op. If the subset is equivalent to an existing
|
|
/// subset insertion op, pair them up. (If there is already a paired up subset
|
|
/// extraction op, overwrite the subset extraction op.)
|
|
void insertExtractionOp(SubsetExtractionOpInterface extractionOp) {
|
|
for (auto it : llvm::enumerate(insertions)) {
|
|
if (!it.value())
|
|
continue;
|
|
auto other = cast<SubsetOpInterface>(it.value().getOperation());
|
|
if (other.operatesOnEquivalentSubset(extractionOp, isEquivalent)) {
|
|
extractions[it.index()] = extractionOp;
|
|
return;
|
|
}
|
|
}
|
|
// There is no known equivalent insertion op. Create a new entry.
|
|
extractions.push_back(extractionOp);
|
|
insertions.push_back({});
|
|
}
|
|
|
|
/// Insert a subset insertion op. If the subset is equivalent to an existing
|
|
/// subset extraction op, pair them up. (If there is already a paired up
|
|
/// subset insertion op, overwrite the subset insertion op.)
|
|
void insertInsertionOp(SubsetInsertionOpInterface insertionOp) {
|
|
for (auto it : llvm::enumerate(extractions)) {
|
|
if (!it.value())
|
|
continue;
|
|
auto other = cast<SubsetOpInterface>(it.value().getOperation());
|
|
if (other.operatesOnEquivalentSubset(insertionOp, isEquivalent)) {
|
|
insertions[it.index()] = insertionOp;
|
|
return;
|
|
}
|
|
}
|
|
// There is no known equivalent extraction op. Create a new entry.
|
|
extractions.push_back({});
|
|
insertions.push_back(insertionOp);
|
|
}
|
|
|
|
SmallVector<SubsetExtractionOpInterface> extractions;
|
|
SmallVector<SubsetInsertionOpInterface> insertions;
|
|
SmallVector<SubsetOpInterface> allSubsetOps;
|
|
};
|
|
} // namespace
|
|
|
|
/// If the given value has a single use by an op that is a terminator, return
|
|
/// that use. Otherwise, return nullptr.
|
|
static OpOperand *getSingleTerminatorUse(Value value) {
|
|
if (!value.hasOneUse())
|
|
return nullptr;
|
|
OpOperand &use = *value.getUses().begin();
|
|
if (use.getOwner()->hasTrait<OpTrait::IsTerminator>())
|
|
return &use;
|
|
return nullptr;
|
|
}
|
|
|
|
LogicalResult
|
|
MatchingSubsets::populateSubsetOpsAtIterArg(LoopLikeOpInterface loopLike,
|
|
BlockArgument iterArg,
|
|
bool collectHoistableOps) {
|
|
assert(iterArg.getOwner()->getParentOp() == loopLike && "invalid iter_arg");
|
|
Value value = iterArg;
|
|
|
|
// Traverse use-def chain. Subset ops can be hoisted only if all ops along the
|
|
// use-def chain starting from the region iter_arg are subset extraction or
|
|
// subset insertion ops. The chain must terminate at the corresponding yield
|
|
// operand (e.g., no swapping of iter_args).
|
|
OpOperand *yieldedOperand = nullptr;
|
|
// Iterate until the single use of the current SSA value is a terminator,
|
|
// which is expected to be the yielding operation of the loop.
|
|
while (!(yieldedOperand = getSingleTerminatorUse(value))) {
|
|
Value nextValue = {};
|
|
|
|
for (OpOperand &use : value.getUses()) {
|
|
if (auto nestedLoop = dyn_cast<LoopLikeOpInterface>(use.getOwner())) {
|
|
// Subset ops in nested loops are collected to check if there are only
|
|
// disjoint subset ops, but such subset ops are not subject to hoisting.
|
|
// To hoist subset ops from nested loops, the hoisting transformation
|
|
// should be run on the nested loop.
|
|
auto nestedIterArg = nestedLoop.getTiedLoopRegionIterArg(&use);
|
|
if (!nestedIterArg)
|
|
return failure();
|
|
// Note: `populateSubsetOpsAtIterArg` fails if there is no single SSA
|
|
// use-def chain starting at `nestedIterArg` and terminating in the
|
|
// tied, yielding operand.
|
|
if (failed(populateSubsetOpsAtIterArg(nestedLoop, nestedIterArg,
|
|
/*collectHoistableOps=*/false)))
|
|
return failure();
|
|
nextValue = nestedLoop.getTiedLoopResult(&use);
|
|
continue;
|
|
}
|
|
|
|
auto subsetOp = dyn_cast<SubsetOpInterface>(use.getOwner());
|
|
if (!subsetOp)
|
|
return failure();
|
|
insert(subsetOp);
|
|
|
|
if (auto insertionOp =
|
|
dyn_cast<SubsetInsertionOpInterface>(use.getOwner())) {
|
|
// Current implementation expects that the insertionOp implement
|
|
// the destinationStyleOpInterface as well. Abort if that tha is not
|
|
// the case
|
|
if (!isa<DestinationStyleOpInterface>(use.getOwner())) {
|
|
return failure();
|
|
}
|
|
|
|
// The value must be used as a destination. (In case of a source, the
|
|
// entire tensor would be read, which would prevent any hoisting.)
|
|
if (&use != &insertionOp.getDestinationOperand())
|
|
return failure();
|
|
// There must be a single use-def chain from the region iter_arg to the
|
|
// terminator. I.e., only one insertion op. Branches are not supported.
|
|
if (nextValue)
|
|
return failure();
|
|
nextValue = insertionOp.getUpdatedDestination();
|
|
}
|
|
}
|
|
|
|
// Nothing can be hoisted if the chain does not continue with loop yielding
|
|
// op or a subset insertion op.
|
|
if (!nextValue)
|
|
return failure();
|
|
value = nextValue;
|
|
}
|
|
|
|
// Hoist only if the SSA use-def chain ends in the yielding terminator of the
|
|
// loop and the yielded value is the `idx`-th operand. (I.e., there is no
|
|
// swapping yield.)
|
|
if (loopLike.getTiedLoopYieldedValue(iterArg) != yieldedOperand)
|
|
return failure();
|
|
|
|
return success();
|
|
}
|
|
|
|
/// Hoist all subset ops that operate on the idx-th region iter_arg of the given
|
|
/// loop-like op and index into loop-invariant subset locations. Return the
|
|
/// newly created loop op (that has extra iter_args) or the original loop op if
|
|
/// nothing was hoisted.
|
|
static LoopLikeOpInterface hoistSubsetAtIterArg(RewriterBase &rewriter,
|
|
LoopLikeOpInterface loopLike,
|
|
BlockArgument iterArg) {
|
|
assert(iterArg.getOwner()->getParentOp() == loopLike && "invalid iter_arg");
|
|
auto it = llvm::find(loopLike.getRegionIterArgs(), iterArg);
|
|
int64_t iterArgIdx = std::distance(loopLike.getRegionIterArgs().begin(), it);
|
|
MatchingSubsets subsets;
|
|
if (failed(subsets.populateSubsetOpsAtIterArg(loopLike, iterArg)))
|
|
return loopLike;
|
|
|
|
// Hoist all matching extraction-insertion pairs one-by-one.
|
|
for (auto it : subsets.getHoistableSubsetOps()) {
|
|
auto extractionOp = std::get<0>(it);
|
|
auto insertionOp = std::get<1>(it);
|
|
|
|
// Ops cannot be hoisted if they depend on loop-variant values.
|
|
if (extractionOp) {
|
|
if (!canBeHoisted(extractionOp, [&](OpOperand &operand) {
|
|
return loopLike.isDefinedOutsideOfLoop(operand.get()) ||
|
|
&operand == &extractionOp.getSourceOperand();
|
|
}))
|
|
extractionOp = {};
|
|
}
|
|
if (insertionOp) {
|
|
if (!canBeHoisted(insertionOp, [&](OpOperand &operand) {
|
|
return loopLike.isDefinedOutsideOfLoop(operand.get()) ||
|
|
&operand == &insertionOp.getSourceOperand() ||
|
|
&operand == &insertionOp.getDestinationOperand();
|
|
}))
|
|
insertionOp = {};
|
|
}
|
|
|
|
// Only hoist extraction-insertion pairs for now. Standalone extractions/
|
|
// insertions that are loop-invariant could be hoisted, but there may be
|
|
// easier ways to canonicalize the IR.
|
|
if (extractionOp && insertionOp) {
|
|
// Create a new loop with an additional iter_arg.
|
|
NewYieldValuesFn newYieldValuesFn =
|
|
[&](OpBuilder &b, Location loc,
|
|
ArrayRef<BlockArgument> innerNewBBArgs) -> SmallVector<Value> {
|
|
return {insertionOp.getSourceOperand().get()};
|
|
};
|
|
FailureOr<LoopLikeOpInterface> newLoop =
|
|
loopLike.replaceWithAdditionalYields(
|
|
rewriter, extractionOp.getResult(),
|
|
/*replaceInitOperandUsesInLoop=*/true, newYieldValuesFn);
|
|
if (failed(newLoop))
|
|
return loopLike;
|
|
loopLike = *newLoop;
|
|
|
|
// Hoist the extraction/insertion ops.
|
|
iterArg = loopLike.getRegionIterArgs()[iterArgIdx];
|
|
OpResult loopResult = loopLike.getTiedLoopResult(iterArg);
|
|
OpResult newLoopResult = loopLike.getLoopResults()->back();
|
|
rewriter.moveOpBefore(extractionOp, loopLike);
|
|
rewriter.moveOpAfter(insertionOp, loopLike);
|
|
rewriter.replaceAllUsesWith(insertionOp.getUpdatedDestination(),
|
|
insertionOp.getDestinationOperand().get());
|
|
extractionOp.getSourceOperand().set(
|
|
loopLike.getTiedLoopInit(iterArg)->get());
|
|
rewriter.replaceAllUsesWith(loopResult,
|
|
insertionOp.getUpdatedDestination());
|
|
insertionOp.getSourceOperand().set(newLoopResult);
|
|
insertionOp.getDestinationOperand().set(loopResult);
|
|
}
|
|
}
|
|
|
|
return loopLike;
|
|
}
|
|
|
|
LoopLikeOpInterface
|
|
mlir::hoistLoopInvariantSubsets(RewriterBase &rewriter,
|
|
LoopLikeOpInterface loopLike) {
|
|
// Note: As subset ops are getting hoisted, the number of region iter_args
|
|
// increases. This can enable further hoisting opportunities on the new
|
|
// iter_args.
|
|
for (int64_t i = 0;
|
|
i < static_cast<int64_t>(loopLike.getRegionIterArgs().size()); ++i) {
|
|
loopLike = hoistSubsetAtIterArg(rewriter, loopLike,
|
|
loopLike.getRegionIterArgs()[i]);
|
|
}
|
|
return loopLike;
|
|
}
|