wip

Author: mbagherbeikTT

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

@@ -348,6 +348,14 @@ struct AlwaysSpeculatableImplTrait
 //===----------------------------------------------------------------------===//
 
 namespace MemoryEffects {
+enum Priority {
+  kDefault = 0,
+  kAllocPriority = 1,
+  kFreePriority = 2,
+  kReadPriority = 3,
+  kWritePriority = 4
+};
+
 /// This class represents the base class used for memory effects.
 struct Effect : public SideEffects::Effect {
   using SideEffects::Effect::Effect;
@@ -357,28 +365,48 @@ struct Effect : public SideEffects::Effect {
   using Base = SideEffects::Effect::Base<DerivedEffect, Effect>;
 
   static bool classof(const SideEffects::Effect *effect);
+
+  /// Return the priority associated with this memory effect.
+  Priority getPriority() const { return priority; }
+
+protected:
+  /// Priority value for this effect. Lower numbers indicate higher precedence.
+  Priority priority = kDefault;
 };
 using EffectInstance = SideEffects::EffectInstance<Effect>;
 
+/// Returns vector of effects sorted by effect stage then priority
+/// priority order: allocate -> free -> read -> write
+llvm::SmallVector<MemoryEffects::EffectInstance>
+getMemoryEffectsSorted(Operation *op);
+
 /// The following effect indicates that the operation allocates from some
 /// resource. An 'allocate' effect implies only allocation of the resource, and
 /// not any visible mutation or dereference.
-struct Allocate : public Effect::Base<Allocate> {};
+struct Allocate : public Effect::Base<Allocate> {
+  Allocate() : Effect::Base<Allocate>() { this->priority = kAllocPriority; }
+};
 
 /// The following effect indicates that the operation frees some resource that
 /// has been allocated. An 'allocate' effect implies only de-allocation of the
 /// resource, and not any visible allocation, mutation or dereference.
-struct Free : public Effect::Base<Free> {};
+struct Free : public Effect::Base<Free> {
+  Free() : Effect::Base<Free>() { this->priority = kFreePriority; }
+};
 
 /// The following effect indicates that the operation reads from some resource.
 /// A 'read' effect implies only dereferencing of the resource, and not any
 /// visible mutation.
-struct Read : public Effect::Base<Read> {};
+struct Read : public Effect::Base<Read> {
+  Read() : Effect::Base<Read>() { this->priority = kReadPriority; }
+};
 
 /// The following effect indicates that the operation writes to some resource. A
 /// 'write' effect implies only mutating a resource, and not any visible
 /// dereference or read.
-struct Write : public Effect::Base<Write> {};
+struct Write : public Effect::Base<Write> {
+  Write() : Effect::Base<Write>() { this->priority = kWritePriority; }
+};
 } // namespace MemoryEffects
 
 //===----------------------------------------------------------------------===//
@@ -470,17 +498,6 @@ bool wouldOpBeTriviallyDead(Operation *op);
 ///
 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
@@ -493,35 +510,6 @@ bool isMemoryEffectMovable(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/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h

@@ -9,9 +9,12 @@
 #ifndef MLIR_TRANSFORMS_LOOPINVARIANTCODEMOTIONUTILS_H
 #define MLIR_TRANSFORMS_LOOPINVARIANTCODEMOTIONUTILS_H
 
+#include "mlir/Interfaces/SideEffectInterfaces.h"
 #include "mlir/Support/LLVM.h"
+#include "mlir/Support/TypeID.h"
 
 #include "llvm/ADT/SmallVector.h"
+#include <utility>
 
 namespace mlir {
 
@@ -21,6 +24,17 @@ class Region;
 class RewriterBase;
 class Value;
 
+/// Gathers potential conflicts on all memory resources used within loop
+///
+/// Given a target loop and an op within it (or the loop op itself), 
+/// gathers op's memory effects and flags potential resource conflicts 
+/// in a map and then recurses into the op's regions to gather nested 
+/// resource conflicts 
+///
+/// First call should use loop = someLoop and op = someLoop.getOperation()
+void gatherResourceConflicts(LoopLikeOpInterface loop, Operation *op,
+    DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>> &resourceConflicts);
+
 /// Given a list of regions, perform loop-invariant code motion. An operation is
 /// loop-invariant if it depends only of values defined outside of the loop.
 /// LICM moves these operations out of the loop body so that they are not
@@ -63,7 +77,7 @@ class Value;
 ///
 /// Returns the number of operations moved.
 size_t moveLoopInvariantCode(
-    ArrayRef<Region *> regions,
+    LoopLikeOpInterface loopLike,
     function_ref<bool(Value, Region *)> isDefinedOutsideRegion,
     function_ref<bool(Operation *, Region *)> shouldMoveOutOfRegion,
     function_ref<void(Operation *, Region *)> moveOutOfRegion);

mlir/lib/Interfaces/SideEffectInterfaces.cpp

@@ -11,6 +11,7 @@
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Dominance.h"
 #include "mlir/IR/SymbolTable.h"
+#include <optional>
 #include <utility>
 
 using namespace mlir;
@@ -346,14 +347,29 @@ std::optional<bool> mlir::isZeroTrip(mlir::LoopLikeOpInterface &loop) {
   return false;
 }
 
-bool mlir::isMemoryEffectMovable(Operation *op) {
-  if (isMemoryEffectFree(op))
-    return true;
+llvm::SmallVector<MemoryEffects::EffectInstance>
+mlir::MemoryEffects::getMemoryEffectsSorted(Operation *op) {
+  auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
 
-  if (isMemoryEffectConflictFree(op))
-    return true;
+  llvm::SmallVector<MemoryEffects::EffectInstance> effectsSorted;
+  memInterface.getEffects(effectsSorted);
 
-  return false;
+  auto sortEffects = 
+    [](llvm::SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+    llvm::stable_sort(effects, [](const MemoryEffects::EffectInstance &a,
+                                  const MemoryEffects::EffectInstance &b) {
+      if (a.getStage() < b.getStage())
+        return true;
+      
+      if (a.getStage() == b.getStage())
+        return a.getEffect()->getPriority() < b.getEffect()->getPriority();
+
+      return false; // b before a
+    });
+  };
+  sortEffects(effectsSorted);
+
+  return effectsSorted;
 }
 
 bool mlir::isMemoryEffectFree(Operation *op) {
@@ -383,101 +399,6 @@ bool mlir::isMemoryEffectFree(Operation *op) {
   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);
-
-    bool 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;
-        }
-      }
-    }
-  } 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
-    // op is conflicting or not.
-    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) {