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

[Uday Bondhugula]

================
@@ -58,13 +60,155 @@ 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 resources if the srcEffect
+/// results in a conflict
+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
+/// (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
+    llvm::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)
+          if (writesConflict)
----------------
bondhugula wrote:

&&

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