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

[Mehdi Amini]

================
@@ -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.
----------------
joker-eph wrote:

I don't quite get it? Are we talking about:
```
func.func @nested_loops_both_having_invariant_code() {
  %m = memref.alloc() : memref<10xf32>
  %cf7 = arith.constant 7.0 : f32
  %cf8 = arith.constant 8.0 : f32

  affine.for %arg0 = 0 to 10 {
    %v0 = arith.addf %cf7, %cf8 : f32
    affine.for %arg1 = 0 to 10 {
      %v1 = arith.addf %v0, %cf8 : f32
      affine.store %v0, %m[%arg0] : memref<10xf32>
    }
  }
  return
}
```

There is no dead loop here?

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