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

>  So in this case, the op isn't arbitrarily speculatable but is movable under the right conditions.

"always Speculatable" does not mean we don't need to check on memory effects, it just means it won't trigger undefined behavior, so:

> Maybe something similar to "Pure", but instead flags the op as having only memory effects and no other side effects.

is somehow the definition of speculatable actually.


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