21 #include "llvm/Support/Debug.h"
22 #include "llvm/Support/DebugLog.h"
25 #define DEBUG_TYPE "licm"
43 for (
OpOperand &operand : child->getOpOperands()) {
45 if (op->
isAncestor(operand.get().getParentRegion()->getParentOp()))
47 if (!condition(operand))
52 return !op->
walk(walkFn).wasInterrupted();
58 op, [&](
OpOperand &operand) {
return definedOutside(operand.
get()); });
68 for (
Region *region : regions) {
69 LDBG() <<
"Original loop:\n" << *region->getParentOp();
71 std::queue<Operation *> worklist;
76 auto definedOutside = [&](
Value value) {
77 return isDefinedOutsideRegion(value, region);
80 while (!worklist.empty()) {
87 LDBG() <<
"Checking op: "
89 if (!shouldMoveOutOfRegion(op, region) ||
93 LDBG() <<
"Moving loop-invariant op: " << *op;
94 moveOutOfRegion(op, region);
100 if (user->getParentRegion() == region)
110 loopLike.getLoopRegions(),
112 return loopLike.isDefinedOutsideOfLoop(value);
115 return isMemoryEffectFree(op) && isSpeculatable(op);
122 class MatchingSubsets {
125 void insert(SubsetOpInterface op,
bool collectHoistableOps =
true) {
126 allSubsetOps.push_back(op);
127 if (!collectHoistableOps)
129 if (
auto extractionOp =
130 dyn_cast<SubsetExtractionOpInterface>(op.getOperation()))
131 insertExtractionOp(extractionOp);
132 if (
auto insertionOp =
133 dyn_cast<SubsetInsertionOpInterface>(op.getOperation()))
134 insertInsertionOp(insertionOp);
141 auto getHoistableSubsetOps() {
142 return llvm::make_filter_range(
143 llvm::zip(extractions, insertions), [&](
auto pair) {
144 auto [extractionOp, insertionOp] = pair;
146 if (extractionOp && insertionOp &&
147 extractionOp->getResult(0).getType() !=
148 insertionOp.getSourceOperand().get().getType())
151 return allDisjoint(extractionOp, insertionOp);
160 LogicalResult populateSubsetOpsAtIterArg(LoopLikeOpInterface loopLike,
162 bool collectHoistableOps =
true);
168 static bool isEquivalent(
Value v1,
Value v2) {
return true; }
173 bool allDisjoint(SubsetExtractionOpInterface extractionOp,
174 SubsetInsertionOpInterface insertionOp)
const {
175 for (SubsetOpInterface other : allSubsetOps) {
176 if (other == extractionOp || other == insertionOp)
179 !other.operatesOnDisjointSubset(extractionOp, isEquivalent))
182 !other.operatesOnDisjointSubset(insertionOp, isEquivalent))
191 void insertExtractionOp(SubsetExtractionOpInterface extractionOp) {
195 auto other = cast<SubsetOpInterface>(it.value().getOperation());
196 if (other.operatesOnEquivalentSubset(extractionOp, isEquivalent)) {
197 extractions[it.index()] = extractionOp;
202 extractions.push_back(extractionOp);
203 insertions.push_back({});
209 void insertInsertionOp(SubsetInsertionOpInterface insertionOp) {
213 auto other = cast<SubsetOpInterface>(it.value().getOperation());
214 if (other.operatesOnEquivalentSubset(insertionOp, isEquivalent)) {
215 insertions[it.index()] = insertionOp;
220 extractions.push_back({});
221 insertions.push_back(insertionOp);
242 MatchingSubsets::populateSubsetOpsAtIterArg(LoopLikeOpInterface loopLike,
244 bool collectHoistableOps) {
246 Value value = iterArg;
256 Value nextValue = {};
259 if (
auto nestedLoop = dyn_cast<LoopLikeOpInterface>(use.getOwner())) {
264 auto nestedIterArg = nestedLoop.getTiedLoopRegionIterArg(&use);
270 if (
failed(populateSubsetOpsAtIterArg(nestedLoop, nestedIterArg,
273 nextValue = nestedLoop.getTiedLoopResult(&use);
277 auto subsetOp = dyn_cast<SubsetOpInterface>(use.getOwner());
282 if (
auto insertionOp =
283 dyn_cast<SubsetInsertionOpInterface>(use.getOwner())) {
287 auto dstOp = dyn_cast<DestinationStyleOpInterface>(use.getOwner());
288 if (!dstOp || !dstOp.hasPureTensorSemantics())
293 if (&use != &insertionOp.getDestinationOperand())
299 nextValue = insertionOp.getUpdatedDestination();
313 if (loopLike.getTiedLoopYieldedValue(iterArg) != yieldedOperand)
324 LoopLikeOpInterface loopLike,
327 BlockArgument *it = llvm::find(loopLike.getRegionIterArgs(), iterArg);
328 int64_t iterArgIdx = std::distance(loopLike.getRegionIterArgs().begin(), it);
329 MatchingSubsets subsets;
330 if (
failed(subsets.populateSubsetOpsAtIterArg(loopLike, iterArg)))
334 for (
auto it : subsets.getHoistableSubsetOps()) {
335 auto extractionOp = std::get<0>(it);
336 auto insertionOp = std::get<1>(it);
341 return loopLike.isDefinedOutsideOfLoop(operand.
get()) ||
342 &operand == &extractionOp.getSourceOperand();
348 return loopLike.isDefinedOutsideOfLoop(operand.
get()) ||
349 &operand == &insertionOp.getSourceOperand() ||
350 &operand == &insertionOp.getDestinationOperand();
358 if (extractionOp && insertionOp) {
363 return {insertionOp.getSourceOperand().get()};
365 FailureOr<LoopLikeOpInterface> newLoop =
366 loopLike.replaceWithAdditionalYields(
367 rewriter, extractionOp.getResult(),
368 true, newYieldValuesFn);
374 iterArg = loopLike.getRegionIterArgs()[iterArgIdx];
375 OpResult loopResult = loopLike.getTiedLoopResult(iterArg);
376 OpResult newLoopResult = loopLike.getLoopResults()->back();
380 insertionOp.getDestinationOperand().get());
381 extractionOp.getSourceOperand().set(
382 loopLike.getTiedLoopInit(iterArg)->get());
384 insertionOp.getUpdatedDestination());
385 insertionOp.getSourceOperand().set(newLoopResult);
386 insertionOp.getDestinationOperand().set(loopResult);
395 LoopLikeOpInterface loopLike) {
400 i < static_cast<int64_t>(loopLike.getRegionIterArgs().size()); ++i) {
402 loopLike.getRegionIterArgs()[i]);
static LoopLikeOpInterface hoistSubsetAtIterArg(RewriterBase &rewriter, LoopLikeOpInterface loopLike, BlockArgument iterArg)
Hoist all subset ops that operate on the idx-th region iter_arg of the given loop-like op and index i...
static OpOperand * getSingleTerminatorUse(Value value)
If the given value has a single use by an op that is a terminator, return that use.
static bool canBeHoisted(Operation *op, function_ref< bool(OpOperand &)> condition)
Checks whether the given op can be hoisted by checking that.
This class represents an argument of a Block.
Block * getOwner() const
Returns the block that owns this argument.
Operation * getParentOp()
Returns the closest surrounding operation that contains this block.
IRValueT get() const
Return the current value being used by this operand.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
This class helps build Operations.
This class represents an operand of an operation.
Set of flags used to control the behavior of the various IR print methods (e.g.
This is a value defined by a result of an operation.
This class provides the API for ops that are known to be terminators.
A wrapper class that allows for printing an operation with a set of flags, useful to act as a "stream...
Operation is the basic unit of execution within MLIR.
bool hasTrait()
Returns true if the operation was registered with a particular trait, e.g.
std::enable_if_t< llvm::function_traits< std::decay_t< FnT > >::num_args==1, RetT > walk(FnT &&callback)
Walk the operation by calling the callback for each nested operation (including this one),...
bool isAncestor(Operation *other)
Return true if this operation is an ancestor of the other operation.
user_range getUsers()
Returns a range of all users.
Region * getParentRegion()
Returns the region to which the instruction belongs.
This class contains a list of basic blocks and a link to the parent operation it is attached to.
This class coordinates the application of a rewrite on a set of IR, providing a way for clients to tr...
void replaceAllUsesWith(Value from, Value to)
Find uses of from and replace them with to.
void moveOpBefore(Operation *op, Operation *existingOp)
Unlink this operation from its current block and insert it right before existingOp which may be in th...
void moveOpAfter(Operation *op, Operation *existingOp)
Unlink this operation from its current block and insert it right after existingOp which may be in the...
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
use_range getUses() const
Returns a range of all uses, which is useful for iterating over all uses.
bool hasOneUse() const
Returns true if this value has exactly one use.
static WalkResult advance()
static WalkResult interrupt()
Operation * getOwner() const
Return the owner of this operand.
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Include the generated interface declarations.
LoopLikeOpInterface hoistLoopInvariantSubsets(RewriterBase &rewriter, LoopLikeOpInterface loopLike)
Hoist loop-invariant tensor subsets (subset extraction and subset insertion ops) from loop-like ops.
std::function< SmallVector< Value >(OpBuilder &b, Location loc, ArrayRef< BlockArgument > newBbArgs)> NewYieldValuesFn
A function that returns the additional yielded values during replaceWithAdditionalYields.
size_t moveLoopInvariantCode(ArrayRef< Region * > regions, function_ref< bool(Value, Region *)> isDefinedOutsideRegion, function_ref< bool(Operation *, Region *)> shouldMoveOutOfRegion, function_ref< void(Operation *, Region *)> moveOutOfRegion)
Given a list of regions, perform loop-invariant code motion.