[MLIR][SideEffects][MemoryEffects] Modified LICM to be more aggressive when checking movability of ops with MemWrite effects

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

@@ -14,7 +14,9 @@
 #ifndef MLIR_INTERFACES_SIDEEFFECTINTERFACES_H
 #define MLIR_INTERFACES_SIDEEFFECTINTERFACES_H
 
+#include "mlir/IR/Dominance.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
 
 namespace mlir {
 namespace SideEffects {
@@ -459,6 +461,26 @@ bool isOpTriviallyDead(Operation *op);
 /// Note: Terminators and symbols are never considered to be trivially dead.
 bool wouldOpBeTriviallyDead(Operation *op);
 
+/// Returns TRUE if the loop is dead/zero-trip,
+/// FALSE if loop has constant bounds/steps and has at least 1 iteration
+/// on every dimension, returns nullopt otherwise
+///
+/// Can only infer if loop is dead if it has constant loop bounds/steps.
+/// Otherwise we assume that it's dead to be conservative.
+///
+std::optional<bool> isZeroTrip(mlir::LoopLikeOpInterface &loop);
+
+/// Returns true if the given operation is allowed to be moved under the
+/// memory effects interface.
+///
+/// An operation is movable if either case is true:
+/// (a) free of memory effects as defined in isMemoryEffectFree()
+/// (b) if the operation does have memory effects, it must be conflict-free
+/// as defined in isMemoryEffectConflictFree()
+///
+/// If the operation meets either criteria, then it is movable under memory effects
+bool isMemoryEffectMovable(Operation *op);
+
 /// Returns true if the given operation is free of memory effects.
 ///
 /// An operation is free of memory effects if its implementation of
@@ -471,6 +493,35 @@ bool wouldOpBeTriviallyDead(Operation *op);
 /// conditions are satisfied.
 bool isMemoryEffectFree(Operation *op);
 
+/// Returns true if the given operation has conflict-free write effects
+///
+/// An operation is conflict free:
+/// (1) Parent is a loop with the LoopLikeOpInterface
+/// (2) Parent loop is not a zero trip loop and has constant bounds/steps
+/// (3) all of the op's memory effects are of type Write
+/// (4) there are no other ops with Alloc/Free/Write effects on the same
+/// resources within the op's parent loop region
+/// (5) all ops in the parent loop region with Read effects on the same
+/// resources are dominated by the operation
+///
+/// If the operation meets all criteria, then it is conflict free
+bool isMemoryEffectConflictFree(Operation *op);
+
+/// Returns true if op and/or any operations within its nested regions
+/// have a memory effect conflict with mainOp as defined below:
+///
+/// op has a memory effect conflict with mainOp if op and/or any of 
+/// the operations in its nested regions meet any of these criteria:
+/// (a) they have any Alloc/Free/Write effects on the resources used by mainOp 
+/// (b) they dominate mainOp and have any read effect on the resources used by mainOp
+///
+/// Function mutates resources map
+///
+/// If none of the critera above are met, mainOp and op are conflict free
+bool hasMemoryEffectConflict(
+    Operation *mainOp, Operation *op, 
+    mlir::DominanceInfo &dom, DenseMap<TypeID, int> &resourceCounts);
+
 /// Returns the side effects of an operation. If the operation has
 /// RecursiveMemoryEffects, include all side effects of child operations.
 ///

mlir/lib/Interfaces/CMakeLists.txt

@@ -85,7 +85,21 @@ add_mlir_interface_library(MemorySlotInterfaces)
 add_mlir_interface_library(ParallelCombiningOpInterface)
 add_mlir_interface_library(RuntimeVerifiableOpInterface)
 add_mlir_interface_library(ShapedOpInterfaces)
-add_mlir_interface_library(SideEffectInterfaces)
+
+add_mlir_library(MLIRSideEffectInterfaces
+  SideEffectInterfaces.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Interfaces
+
+  DEPENDS
+  MLIRSideEffectInterfacesIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRIR
+  MLIRDialectUtils
+  MLIRLoopLikeInterface
+)
 
 add_mlir_library(MLIRSubsetOpInterface
   SubsetOpInterface.cpp

mlir/lib/Interfaces/SideEffectInterfaces.cpp

@@ -8,6 +8,8 @@
 
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/IR/Dominance.h"
 #include "mlir/IR/SymbolTable.h"
 #include <utility>
 
@@ -317,14 +319,53 @@ bool mlir::wouldOpBeTriviallyDead(Operation *op) {
   return wouldOpBeTriviallyDeadImpl(op);
 }
 
+std::optional<bool> mlir::isZeroTrip(mlir::LoopLikeOpInterface &loop) {
+  auto lbs = loop.getLoopLowerBounds();
+  auto ubs = loop.getLoopUpperBounds();
+  auto steps = loop.getLoopSteps();
+
+  if (!lbs || !ubs || !steps)
+    return std::nullopt;
+
+  if (lbs->size() != ubs->size() || ubs->size() != steps->size())
+    return std::nullopt;
+
+  for (size_t i = 0; i < steps->size(); ++i) {
+    auto lb = getConstantIntValue((*lbs)[i]);
+    auto ub = getConstantIntValue((*ubs)[i]);
+    auto st = getConstantIntValue((*steps)[i]);
+
+    if (!lb || !ub || !st)
+      return std::nullopt; // non-constant -> unknown
+
+    if (*st >= 0 && *lb >= *ub)
+      return true;
+    if (*st < 0 && *lb <= *ub)
+      return true;
+  }
+  return false;
+}
+
+bool mlir::isMemoryEffectMovable(Operation *op) {
+  if (isMemoryEffectFree(op))
+    return true;
+
+  if (isMemoryEffectConflictFree(op))
+    return true;
+
+  return false;
+}
+
 bool mlir::isMemoryEffectFree(Operation *op) {
   if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
     if (!memInterface.hasNoEffect())
       return false;
+
     // If the op does not have recursive side effects, then it is memory effect
     // free.
     if (!op->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
       return true;
+
   } else if (!op->hasTrait<OpTrait::HasRecursiveMemoryEffects>()) {
     // Otherwise, if the op does not implement the memory effect interface and
     // it does not have recursive side effects, then it cannot be known that the
@@ -338,9 +379,101 @@ bool mlir::isMemoryEffectFree(Operation *op) {
     for (Operation &op : region.getOps())
       if (!isMemoryEffectFree(&op))
         return false;
+
   return true;
 }
 
+bool mlir::isMemoryEffectConflictFree(Operation *op) {
+  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 false;
+
+  // check parent loop to make sure it's not dead
+  Operation *parent = op->getParentOp();
+  if (!parent)
+    return false;
+
+  auto loopInterface = dyn_cast<LoopLikeOpInterface>(parent);
+  if (!loopInterface)
+    return false;
+
+  auto isDead = isZeroTrip(loopInterface);
+  if (!isDead.has_value() || isDead.value())
+    return false;
+
+  // 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 false;
+
+  DenseMap<TypeID, int> resourceCounts;
+
+  // ensure op only has Write effects and gather unique
+  // resource names
+  for (const MemoryEffects::EffectInstance &effect : effects) {
+    if (!isa<MemoryEffects::Write>(effect.getEffect()))
+      return false;
+
+    resourceCounts.try_emplace(effect.getResource()->getResourceID(), 0);
+  }
+
+  mlir::DominanceInfo dom(parent);
+
+  for (Region &region : parent->getRegions())
+    for (Operation &opI : region.getOps())
+      if (hasMemoryEffectConflict(op, &opI, dom, resourceCounts))
+        return false;
+
+  return true;
+}
+
+bool mlir::hasMemoryEffectConflict(
+    Operation *mainOp, Operation *op,
+    mlir::DominanceInfo &dom, DenseMap<TypeID, int> &resourceCounts) {
+
+  if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
+
+    llvm::SmallVector<MemoryEffects::EffectInstance> effects;
+    memInterface.getEffects(effects);
+
+    auto isDominated = dom.properlyDominates(mainOp, op);
+
+    // ensure op only has Write or dominated Read effects
+    // check used resources
+    for (const MemoryEffects::EffectInstance &effect : effects) {
+      auto resourceID = effect.getResource()->getResourceID();
+
+      if (resourceCounts.contains(resourceID)) {
+        if (isa<MemoryEffects::Read>(effect.getEffect())) {
+          if (isDominated) {
+            continue; // skip dominated reads
+          }
+        }
+        else if (!isa<MemoryEffects::Write>(effect.getEffect())) {
+          return true; // count alloc/free in same region as conflict, be conservative
+        }
+
+        // update write counts, should always be <=1 per resource in region
+        if (++resourceCounts[resourceID] > 1) {
+          return true;
+        }
+      }
+    }
+  }
+
+  // Recurse into the regions and ensure that nested ops don't
+  // conflict with each other's MemWrites
+  for (Region &region : op->getRegions())
+    for (Operation &opI : region.getOps()) 
+      if (hasMemoryEffectConflict(mainOp, &opI, dom, resourceCounts))
+        return true;
+
+  return false;
+}
+
 // the returned vector may contain duplicate effects
 std::optional<llvm::SmallVector<MemoryEffects::EffectInstance>>
 mlir::getEffectsRecursively(Operation *rootOp) {

mlir/lib/Transforms/Utils/LoopInvariantCodeMotionUtils.cpp

@@ -112,7 +112,7 @@ size_t mlir::moveLoopInvariantCode(LoopLikeOpInterface loopLike) {
         return loopLike.isDefinedOutsideOfLoop(value);
       },
       [&](Operation *op, Region *) {
-        return isMemoryEffectFree(op) && isSpeculatable(op);
+        return isSpeculatable(op) && isMemoryEffectMovable(op);
       },
       [&](Operation *op, Region *) { loopLike.moveOutOfLoop(op); });
 }