[Mlir-commits] [mlir] [MLIR][SideEffects][MemoryEffects] Modified LICM to be more aggressive when checking movability of ops with MemWrite effects (PR #155344)

[Mo Bagherbeik]

================
@@ -58,13 +60,170 @@ static bool canBeHoisted(Operation *op,
       op, [&](OpOperand &operand) { return definedOutside(operand.get()); });
 }
 
+/// 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());
+
+  bool conflict = srcHasConflict || srcIsAllocOrFree;
+
+  auto dstIt = resourceConflicts.find(srcResourceID);
+
+  // if it doesn't already exist, create entry for resource in map
+  if (dstIt == resourceConflicts.end()) {
+    resourceConflicts.insert(
+        std::make_pair(srcResourceID, std::make_pair(conflict, srcEffect)));
+    return;
+  }
+
+  // resource already in use
+  bool dstHasConflict = dstIt->second.first;
+  auto dstEffect = dstIt->second.second;
+
+  if (dstHasConflict)
+    return;
+
+  bool srcWrite = isa<MemoryEffects::Write>(srcEffect.getEffect());
+  bool dstRead = isa<MemoryEffects::Read>(dstEffect.getEffect());
+  bool readBeforeWrite = dstRead && srcWrite;
+
+  conflict = conflict || readBeforeWrite;
+
+  dstIt->second = std::make_pair(conflict, srcEffect);
+}
+
+/// Returns true if any of op's OpOperands are defined outside of loopLike
+static bool hasLoopVariantInput(LoopLikeOpInterface loopLike, Operation *op) {
+  for (OpOperand &input : op->getOpOperands())
+    if (!loopLike.isDefinedOutsideOfLoop(input.get()))
+      return true;
+
+  return false;
+}
+
+/// 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) {
+
+  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);
+    if (resConIt == resourceConflicts.end())
+      return true; // RFC realistically shouldn't reach here but just in case?
+
+    bool hasConflict = resConIt->second.first;
+    if (hasConflict)
+      return true;
+  }
+
+  return false;
+}
+
+void mlir::gatherResourceConflicts(
+    LoopLikeOpInterface loopLike, Operation *op,
+    DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
+        &resourceConflicts) {
+
+  if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
+    // gather all effects on op
+    SmallVector<MemoryEffects::EffectInstance> effects =
+        MemoryEffects::getMemoryEffectsSorted(op);
+
+    if (!effects.empty()) {
+      // any variant input to the op could be the data source
+      // for writes, set all writes to conflict
+      bool writesConflict = hasLoopVariantInput(loopLike, op);
+
+      for (const MemoryEffects::EffectInstance &effect : effects) {
+        bool conflict = false;
+        bool isWrite = isa<MemoryEffects::Write>(effect.getEffect());
+
+        // all writes to a resource that follow a read on any other resource
+        // have to be considered a conflict as guaranteeing that the read
+        // is invariant and won't affect the write requires more robust logic
+        if (isa<MemoryEffects::Read>(effect.getEffect()))
+          writesConflict = true;
+
+        if (isWrite && writesConflict)
+          conflict = true;
+
+        mergeResource(resourceConflicts, effect, conflict);
+      }
+    }
+  }
+
+  for (Region &region : op->getRegions())
+    for (Operation &opInner : region.getOps())
+      gatherResourceConflicts(loopLike, &opInner, resourceConflicts);
+}
+
 size_t mlir::moveLoopInvariantCode(
-    ArrayRef<Region *> regions,
+    LoopLikeOpInterface loopLike,
     function_ref<bool(Value, Region *)> isDefinedOutsideRegion,
     function_ref<bool(Operation *, Region *)> shouldMoveOutOfRegion,
     function_ref<void(Operation *, Region *)> moveOutOfRegion) {
   size_t numMoved = 0;
 
+  // TODO: see if this can be spec'd in a meaningful way to add back later.
----------------
mbagherbeikTT wrote:

You're correct, the loops aren't dead. But, when I query their bounds via the methods below, lbs and ubs aren't populated with the loop bounds and isZeroTrip() returns null since it can't guarantee anything. Those methods work with scf.for but not with affine.for. We'd have to use other methods to get the bounds, which could blow-up if every LoopLike needs a different path. I think the best approach for that is to either add a standard way of getting the bounds from ops that have the LoopLikeOpInterface or add an isZeroTrip() method for all of them. I'm not sure if that falls under the scope of this PR. 

` 
auto lbs = loop.getLoopLowerBounds();
auto ubs = loop.getLoopUpperBounds();
auto steps = loop.getLoopSteps();
`

https://github.com/llvm/llvm-project/pull/155344