[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
----------------
mbagherbeikTT wrote:

Kind of, yes. HasRecursiveMemoryEffects still works as it previously did through the isMemoryEffectFree() path but I wanted to discuss how to add support through this part of the path before implementing it. 

The simplest approach for an op with this trait is to recurse down the op's regions and if all nested ops have the MemEffInterface and are conflict free, and the op is speculatable, then the op is movable. The part I'm fuzzy on is whether all nested ops should be required to be speculatable too or if that doesn't matter

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