[Mlir-commits] [mlir] [MLIR][SideEffects][MemoryEffects] Modified LICM to be more aggressive when checking movability of ops with MemWrite effects (PR #155344)
Mo Bagherbeik
llvmlistbot at llvm.org
Wed Oct 8 14:13:35 PDT 2025
================
@@ -58,62 +60,256 @@ static bool canBeHoisted(Operation *op,
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 *, Region *)> moveOutOfRegion) {
- size_t numMoved = 0;
+/// Merges srcEffect's Memory Effect on its resource into the
+/// resourceConflicts map, flagging the resource if the srcEffect
+/// results in a conflict.
+///
+/// \param resourceConflicts The map to store resources' conflicts status.
+/// \param srcEffect The effect to merge into the resourceConflicts map.
+/// \param srcHasConflict Whether the srcEffect results in a conflict based
+/// on higher level analysis.
+///
+/// resourceConflicts is modified by the function and will be non-empty
+static void
+mergeResource(DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
+ &resourceConflicts,
+ const MemoryEffects::EffectInstance &srcEffect,
+ bool srcHasConflict) {
+
+ TypeID srcResourceID = srcEffect.getResource()->getResourceID();
+
+ bool srcIsAllocOrFree = isa<MemoryEffects::Allocate>(srcEffect.getEffect()) ||
+ isa<MemoryEffects::Free>(srcEffect.getEffect());
+
+ LDBG() << " Merging Effect \"" << srcEffect.getEffect() << "<" << srcEffect.getResource()->getName() << ">\"";
+
+ bool conflict = srcHasConflict || srcIsAllocOrFree;
+
+ auto [dstIt, inserted] = resourceConflicts.insert(
+ std::make_pair(srcResourceID, std::make_pair(conflict, srcEffect)));
+ if (inserted) {
+ LDBG() << " Effect inserted to map";
+ return;
+ }
- for (Region *region : regions) {
- LDBG() << "Original loop:\n" << *region->getParentOp();
+ // resource already in use
+ bool dstHasConflict = dstIt->second.first;
+ auto dstEffect = dstIt->second.second;
- std::queue<Operation *> worklist;
- // Add top-level operations in the loop body to the worklist.
- for (Operation &op : region->getOps())
- worklist.push(&op);
+ if (dstHasConflict) {
+ LDBG() << " Resource has existing conflict from Effect \"" << dstEffect.getEffect() << "<" << dstEffect.getResource()->getName() << ">\"";
+ return;
+ }
- auto definedOutside = [&](Value value) {
- return isDefinedOutsideRegion(value, region);
- };
+ bool srcWrite = isa<MemoryEffects::Write>(srcEffect.getEffect());
+ bool dstRead = isa<MemoryEffects::Read>(dstEffect.getEffect());
+ bool readBeforeWrite = dstRead && srcWrite;
- 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;
+ conflict = conflict || readBeforeWrite;
- LDBG() << "Checking op: "
- << OpWithFlags(op, OpPrintingFlags().skipRegions());
- if (!shouldMoveOutOfRegion(op, region) ||
- !canBeHoisted(op, definedOutside))
- continue;
+ LDBG() << " Resource " << dstEffect.getResource()->getName() << " updated to conflict status = " << conflict;
+ dstIt->second = std::make_pair(conflict, srcEffect);
+}
- LDBG() << "Moving loop-invariant op: " << *op;
- moveOutOfRegion(op, region);
- ++numMoved;
+/// Returns true if any of op's operands is defined inside the loop.
+static bool hasLoopVariantInput(LoopLikeOpInterface loopLike, Operation *op) {
+ return llvm::any_of(op->getOperands(), [&] (Value v) {
+ return !loopLike.isDefinedOutsideOfLoop(v);
+ });
+}
- // 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);
+/// Returns true if:
+/// (a) any of the resources used by op's Memory Effects have been
+/// flagged as having a conflict within the resourceConflicts map OR
+/// (b) op doesn't have a MemoryEffectOpInterface or has one but
+/// without any specific effects.
+static bool mayHaveMemoryEffectConflict(
+ Operation *op,
+ DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
+ &resourceConflicts) {
+ LDBG() << "Checking for memory effect conflict on op: "
+ << OpWithFlags(op, OpPrintingFlags().skipRegions());
+
+ auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
+
+ // op does not implement the memory effect op interface
+ // shouldn't be flagged as movable to be conservative
+ if (!memInterface)
+ return true;
+
+ // gather all effects on op
+ llvm::SmallVector<MemoryEffects::EffectInstance> effects;
+ memInterface.getEffects(effects);
+
+ // op has interface but no effects, be conservative
+ if (effects.empty())
+ return true;
+
+ // RFC moving ops with HasRecursiveMemoryEffects that have nested ops
+ // with effects is not supported
+
+ // op has no conflicts IFF all resources are flagged as having no conflicts
+ for (const MemoryEffects::EffectInstance &effect : effects) {
+ auto resourceID = effect.getResource()->getResourceID();
+
+ auto resConIt = resourceConflicts.find(resourceID);
+ assert(resConIt != resourceConflicts.end());
+
+ bool hasConflict = resConIt->second.first;
+ if (hasConflict) {
+ LDBG() << " has conflict on resource " << effect.getResource()->getName()
+ << " from Memory Effect Mem" << effect.getEffect();
+ return true;
}
}
- return numMoved;
+ return false;
+}
+
+void mlir::gatherResourceConflicts(
+ LoopLikeOpInterface loopLike, Operation *op,
+ DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
+ &resourceConflicts) {
+
+ LDBG() << "Gather conflicts: on loop "
+ << OpWithFlags(loopLike.getOperation(), OpPrintingFlags().skipRegions())
+ << ", visiting op: " << OpWithFlags(op, OpPrintingFlags().skipRegions());
+
+ if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
+ // gather all effects on op
+ SmallVector<MemoryEffects::EffectInstance> effects =
+ MemoryEffects::getMemoryEffectsSorted(op);
+
+ if (!effects.empty()) {
+ LDBG() << " # of effects = " << effects.size();
+
+ // Any input to the op could be the input data source
+ // for write effects in the same op. E.g.,
+ // scf.for ... {
+ // %0 = foo.bar(...) : ...
+ // foo.baz(%0) // foo.baz has a MemWrite<SomeResource, 0> effect
+ // }
+ // The input %0 could be the data source for the write effect in
+ // foo.baz. Since %0 is loop-variant, this may cause a conflict on
+ // SomeResource as the MemWrite contents may change between loop iterations.
+ // A more complex analysis would be needed to determine
+ // if this is a true conflict or not.
+ bool writesConflict = hasLoopVariantInput(loopLike, op);
+ LDBG() << " has loop-variant input? "
+ << (writesConflict ? "true" : "false");
+
+ for (const MemoryEffects::EffectInstance &effect : effects) {
+ bool conflict = writesConflict && isa<MemoryEffects::Write>(effect.getEffect());
+
+ mergeResource(resourceConflicts, effect, conflict);
+
+ // All writes to a resource that follow a read on any other resource
+ // have to be considered as causing a conflict, similar
+ // to how we handle loop-variant inputs above; guaranteeing that the read
+ // is invariant and won't affect the write requires more robust logic.
+ if (isa<MemoryEffects::Read>(effect.getEffect())) {
+ LDBG() << " has MemRead; subsequent Writes will be treated as conflicting";
+ writesConflict = true;
+ }
+ }
+ }
+ }
+
+ for (Region ®ion : op->getRegions())
+ for (Operation &opInner : region.getOps())
+ gatherResourceConflicts(loopLike, &opInner, resourceConflicts);
+}
+
+size_t mlir::moveLoopInvariantCode(
+ LoopLikeOpInterface loopLike,
+ function_ref<bool(Value, Region *)> isDefinedOutsideRegion,
+ function_ref<bool(Operation *, Region *)> shouldMoveOutOfRegion,
+ function_ref<void(Operation *, Region *)> moveOutOfRegion) {
+ size_t numMovedTotal = 0;
+
+ // TODO: see if this can be spec'd in a meaningful way to add back later.
+ //
+ // check that the loop isn't dead
+ auto isDead = loopLike.isZeroTrip();
+
+ LDBG() << "Loop " << OpWithFlags(loopLike.getOperation(), OpPrintingFlags().skipRegions())
+ << " has constant bounds and steps? isZeroTrip()? " << (isDead.has_value() ? (isDead.value() ? "YES, YES" : "YES, NO") : "NO, NULL");
+
+ if (isDead.has_value() && isDead.value())
+ return numMovedTotal;
+
+ int numMoved = 0;
+
+ do {
+ // reset value for iteration
+ numMoved = 0;
+
+ // Go through loop body and map out resource usages.
+ // op->regions are essentially merged sequentially.
+ // E.g., an if's "then" and "else" regions are treated like one
+ // continuous region --> need to add fork checking.
+ //
+ // loop "do" and "then" regions also merged.
+ DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
+ resourceConflicts;
+ gatherResourceConflicts(loopLike, loopLike.getOperation(), resourceConflicts);
+
+ auto regions = loopLike.getLoopRegions();
+ for (Region *region : regions) {
+ LDBG() << "Original loop:\n" << *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;
+
+ LDBG() << "Checking op: "
+ << OpWithFlags(op, OpPrintingFlags().skipRegions());
+
+ bool noMemoryConflicts =
+ isMemoryEffectFree(op) ||
+ !mayHaveMemoryEffectConflict(op, resourceConflicts);
+
+ if (!noMemoryConflicts || !shouldMoveOutOfRegion(op, region) ||
+ !canBeHoisted(op, definedOutside))
+ continue;
+
+ LDBG() << "Moving loop-invariant op: " << *op;
+ moveOutOfRegion(op, region);
+ ++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);
+ }
+ }
+
+ numMovedTotal += numMoved;
----------------
mbagherbeikTT wrote:
great suggestion: added
https://github.com/llvm/llvm-project/pull/155344
More information about the Mlir-commits
mailing list