[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,62 +60,407 @@ static bool canBeHoisted(Operation *op,
       op, [&](OpOperand &operand) { return definedOutside(operand.get()); });
 }
 
-size_t mlir::moveLoopInvariantCode(
-    ArrayRef<Region *> regions,
-    function_ref<bool(Value, Region *)> isDefinedOutsideRegion,
-    function_ref<bool(Operation *, Region *)> shouldMoveOutOfRegion,
-    function_ref<void(Operation *, Region *)> moveOutOfRegion) {
-  size_t numMoved = 0;
+/// 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(MemoryConflictMap &resourceConflicts,
+                          const MemoryEffects::EffectInstance &srcEffect,
+                          bool srcHasConflict) {
+
+  TypeID srcResourceID = srcEffect.getResource()->getResourceID();
+
+  bool srcIsAllocOrFree = isa<MemoryEffects::Allocate>(srcEffect.getEffect()) ||
+                          isa<MemoryEffects::Free>(srcEffect.getEffect());
+
+  LDBG() << "<Merging Effect> : \"" << srcEffect.getEffect()->getEffectName()
+         << " on resource <" << srcEffect.getResource()->getName() << ">\""
+         << "\n";
+
+  bool conflict = srcHasConflict || srcIsAllocOrFree;
+
+  auto [dstIt, inserted] = resourceConflicts.insert(
+        std::make_pair(srcResourceID, std::make_pair(conflict, srcEffect)));
+  if (inserted) {
+    LDBG() << ". . . . "
+           << "Effect inserted to map"
+           << "\n";
+    return;
+  }
 
-  for (Region *region : regions) {
-    LDBG() << "Original loop:\n" << *region->getParentOp();
+  // resource already in use
+  bool dstHasConflict = dstIt->second.first;
+  auto dstEffect = dstIt->second.second;
 
-    std::queue<Operation *> worklist;
-    // Add top-level operations in the loop body to the worklist.
-    for (Operation &op : region->getOps())
-      worklist.push(&op);
+  if (dstHasConflict) {
+    LDBG() << ". . . . "
+           << "Resource has existing conflict from Effect Mem"
+           << dstEffect.getValue() << "\n";
+    return;
+  }
 
-    auto definedOutside = [&](Value value) {
-      return isDefinedOutsideRegion(value, region);
-    };
+  bool srcWrite = isa<MemoryEffects::Write>(srcEffect.getEffect());
+  bool dstRead = isa<MemoryEffects::Read>(dstEffect.getEffect());
+  bool readBeforeWrite = dstRead && srcWrite;
 
-    while (!worklist.empty()) {
-      Operation *op = worklist.front();
-      worklist.pop();
-      // Skip ops that have already been moved. Check if the op can be hoisted.
-      if (op->getParentRegion() != region)
-        continue;
+  conflict = conflict || readBeforeWrite;
 
-      LDBG() << "Checking op: "
-             << OpWithFlags(op, OpPrintingFlags().skipRegions());
-      if (!shouldMoveOutOfRegion(op, region) ||
-          !canBeHoisted(op, definedOutside))
-        continue;
+  LDBG() << ". . . . "
+         << "Resource conflict status updated to = " << conflict << "\n";
+
+  dstIt->second = std::make_pair(conflict, srcEffect);
+}
+
+/// Returns true if any of op's operands is defined inside the loop.
+static bool hasLoopVariantInput(LoopLikeOpInterface loopLike, Operation *op) {
+  return llvm::any_of(op->getOperands(), [&] (Value v) {
+    return !loopLike.isDefinedOutsideOfLoop(v);
+  });
+}
+
+/// 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,
+                                        MemoryConflictMap *resourceConflicts) {
+  LDBG() << "<Checking for memory effect conflict on op> : "
+         << OpWithFlags(op, OpPrintingFlags().skipRegions());
+
+  auto condSpecInterface = dyn_cast<ConditionallySpeculatable>(op);
+
+  // if op implements ConditionallySpeculatable interface, must be speculatable!
+  if (condSpecInterface && !isSpeculatable(op))
+    return true;
----------------
joker-eph wrote:

For example your test operations are not tagged as Speculatable right now as far as I can tell?

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