[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
Sun Dec 21 16:32:07 PST 2025
================
@@ -58,62 +60,407 @@ 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(MemoryConflictMap &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()->getEffectName()
+ << " on resource <" << srcEffect.getResource()->getName() << ">\""
+ << "\n";
+
+ 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"
+ << "\n";
+ 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 Mem"
+ << dstEffect.getValue() << "\n";
+ 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 conflict status updated to = " << conflict << "\n";
+
+ dstIt->second = std::make_pair(conflict, srcEffect);
+}
+
+/// 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);
+ });
+}
+
+/// 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,
+ MemoryConflictMap *resourceConflicts) {
+ LDBG() << "<Checking for memory effect conflict on op> : "
+ << OpWithFlags(op, OpPrintingFlags().skipRegions());
+
+ auto condSpecInterface = dyn_cast<ConditionallySpeculatable>(op);
+
+ // if op implements ConditionallySpeculatable interface, must be speculatable!
+ if (condSpecInterface && !isSpeculatable(op))
+ return true;
+
+ 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;
+
+ // Ops with Recursive Memory Effects are special-cased here.
+ // For now we'll only allow them to be moved if they're effect
+ // free.
+ // A potential solution is to recursively gather all resources on all
+ // contained ops and then run the for-loop further below. Requires discussions
+ // re: obscure corner cases.
+ if (op->hasTrait<OpTrait::HasRecursiveMemoryEffects>()) {
+ return !isMemoryEffectFree(op);
+ }
+
+ // 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;
+
+ // 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 resConIter = resourceConflicts->find(resourceID);
+ assert(resConIter != resourceConflicts->end());
+
+ bool hasConflict = resConIter->second.first;
+ if (hasConflict) {
+ LDBG() << ". . . . "
+ << "Conflict deteceted on resource <"
+ << effect.getResource()->getName() << "> from Memory Effect Mem"
+ << effect.getValue() << "\n";
+ return true;
+ }
+ }
+
+ return false;
+}
- LDBG() << "Moving loop-invariant op: " << *op;
- moveOutOfRegion(op, region);
- ++numMoved;
+static void
+mapLoopResourceUsage(LoopLikeOpInterface loopLike, Operation *op,
+ MemoryConflictMap &resourceConflicts,
+ llvm::SmallSet<Operation *, 8> &opsWithReadBeforeWrite) {
+
+ LDBG() << "<Mapping resource usage on op> : "
+ << OpWithFlags(op, OpPrintingFlags().skipRegions()) << "\n";
+
+ if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
+ LDBG() << ". . . . "
+ << "op has MemoryEffectsOpInterface"
+ << "\n";
+
+ // gather all effects on op
+ SmallVector<MemoryEffects::EffectInstance> effects =
+ MemoryEffects::getMemoryEffectsSorted(op);
+
+ LDBG() << ". . . . "
+ << "# of effects = " << effects.size();
+
+ if (!effects.empty()) {
+ // 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") << "\n";
+
+ bool hasRead = false;
+
+ for (const MemoryEffects::EffectInstance &effect : effects) {
+ bool inConflict = false;
+
+ if (isa<MemoryEffects::Write>(effect.getEffect())) {
+ if (hasRead) {
+ LDBG() << ". . . . "
+ << "read-before-write detected!"
+ << "\n";
+ LDBG() << ". . . . "
+ << ". . . . "
+ << "Inserting op into set for later checks!"
+ << "\n";
+ opsWithReadBeforeWrite.insert(op);
+ }
+
+ inConflict = writesConflict;
+ }
+
+ mergeResource(resourceConflicts, effect, inConflict);
+
+ // All writes to a resource that follow a read on any other resource
+ // need additional logic to check if the read will result in a conflict
+ // on the following write op(s)'s resource(s).
+ // Need to keep track of ops that have read before writes.
+ // If the resource for the read effect has a conflict after all loop
+ // resource usages have been mapped, then the conflict will be
+ // propagated to the resources used by the following writes. LOGIC: if
+ // the read resource is in conflict, that means the value stored is no
+ // longer loop invariant --> the read could be the data source for the
+ // write --> the write is not guaranteed to be loop invariant.
+ if (isa<MemoryEffects::Read>(effect.getEffect())) {
+ TypeID resourceID = effect.getResource()->getResourceID();
+ auto resConIter = resourceConflicts.find(resourceID);
+
+ if (resConIter != resourceConflicts.end()) {
+ hasRead = true;
+ }
+ }
+ }
+ }
+ }
+
+ for (Region ®ion : op->getRegions())
+ for (Operation &opInner : region.getOps())
+ mapLoopResourceUsage(loopLike, &opInner, resourceConflicts,
+ opsWithReadBeforeWrite);
+}
+
+static void propagateSameOpReadBeforeWriteConflicts(
+ LoopLikeOpInterface loopLike, Operation *op,
+ MemoryConflictMap &resourceConflicts,
+ llvm::SmallSet<Operation *, 8> &opsWithReadBeforeWrite) {
+
+ for (auto *opInner : opsWithReadBeforeWrite) {
+ // gather all effects on op
+ SmallVector<MemoryEffects::EffectInstance> effects =
+ MemoryEffects::getMemoryEffectsSorted(opInner);
+
+ bool writesConflict = false;
- // 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);
+ for (const MemoryEffects::EffectInstance &effect : effects) {
+ if (writesConflict && isa<MemoryEffects::Write>(effect.getEffect())) {
+
+ TypeID resourceID = effect.getResource()->getResourceID();
+ auto resConIter = resourceConflicts.find(resourceID);
+
+ // already has conflict, move on
+ if (resConIter->getSecond().first)
+ continue;
+
+ resConIter->getSecond().first = true;
+ resConIter->getSecond().second = effect;
+ }
+
+ if (isa<MemoryEffects::Read>(effect.getEffect())) {
+ TypeID resourceID = effect.getResource()->getResourceID();
+ auto resConIter = resourceConflicts.find(resourceID);
+
+ if (resConIter != resourceConflicts.end() &&
+ resConIter->getSecond().first) {
+ writesConflict = true;
+ }
+ }
}
}
+}
- return numMoved;
+void mlir::mapResourceConflicts(LoopLikeOpInterface loopLike, Operation *op,
+ MemoryConflictMap &resourceConflicts) {
+ llvm::SmallSet<Operation *, 8> opsWithReadBeforeWrite;
+ mapLoopResourceUsage(loopLike, loopLike.getOperation(), resourceConflicts,
+ opsWithReadBeforeWrite);
+ propagateSameOpReadBeforeWriteConflicts(loopLike, loopLike.getOperation(),
+ resourceConflicts,
+ opsWithReadBeforeWrite);
+}
+
+size_t mlir::moveLoopInvariantCode(
+ LoopLikeOpInterface loopLike,
+ function_ref<bool(Value, Region *)> isDefinedOutsideRegion,
+ function_ref<bool(Operation *, Region *)> shouldMoveSpeculatable,
+ function_ref<bool(Operation *, MemoryConflictMap *)> shouldMoveMemoryEffect,
+ function_ref<void(Operation *, Region *)> moveOutOfRegion) {
+ size_t numMovedTotal = 0;
+
+ // Check that the loop isn't dead.
+ // Two separate methods used to check this, depending on what the loopLike op
+ // implements. If neither is available, we can't guarantee loop liveness.
+ auto isMaybeDead = loopLike.isZeroTrip();
+ auto tripCount = loopLike.getStaticTripCount();
+
+ bool confirmedDead = (isMaybeDead.has_value() && isMaybeDead.value()) ||
+ (tripCount.has_value() && tripCount.value() == 0);
+ bool ambiguousLiveness = !isMaybeDead.has_value() && !tripCount.has_value();
+ bool loopIsLive = !confirmedDead && !ambiguousLiveness;
+
+ LDBG() << "Running LICM on loop op. . . ."
+ << "\n";
+ LDBG() << "<LoopLikeOp> : "
+ << OpWithFlags(loopLike.getOperation(),
+ OpPrintingFlags().skipRegions())
+ << "\n";
+ LDBG() << ". . . . "
+ << "confirmedDead = " << confirmedDead << "\n";
+ LDBG() << ". . . . "
+ << "ambiguousLiveness = " << ambiguousLiveness << "\n";
+ LDBG() << ". . . . "
+ << "loopIsLive = " << loopIsLive << "\n";
+
+ int iteration = 0;
+ int numMoved = 0;
+
+ do {
+ // reset value for iteration
+ numMoved = 0;
+
+ MemoryConflictMap resourceConflicts;
+
+ // For loops that are guaranteed to execute at least one iterations:
+ // 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.
+ if (loopIsLive)
+ mapResourceConflicts(loopLike, loopLike.getOperation(),
+ resourceConflicts);
+
+ auto regions = loopLike.getLoopRegions();
+ for (Region *region : regions) {
+ 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;
+
+ bool isHoistable = canBeHoisted(op, definedOutside);
+ bool movableUnderSpeculabilityPath = shouldMoveSpeculatable(op, region);
+ bool movableUnderMemoryEffectsPath =
+ loopIsLive && shouldMoveMemoryEffect(op, &resourceConflicts);
----------------
mbagherbeikTT wrote:
duplicate
https://github.com/llvm/llvm-project/pull/155344
More information about the Mlir-commits
mailing list