[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)
+            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;
 
+  // check that the loop isn't dead
+  // auto isDead = loopLike.isZeroTrip();
+  // if (!isDead.has_value() || isDead.value())
+  //   return numMoved;
+
+  // 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;
+  mlir::gatherResourceConflicts(loopLike, loopLike.getOperation(), resourceConflicts);
----------------
bondhugula wrote:

clang-format.

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