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

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

@@ -14,6 +14,7 @@
 #ifndef MLIR_INTERFACES_SIDEEFFECTINTERFACES_H
 #define MLIR_INTERFACES_SIDEEFFECTINTERFACES_H
 
+#include "mlir/IR/Dominance.h"
 #include "mlir/IR/OpDefinition.h"
 
 namespace mlir {
@@ -346,6 +347,20 @@ struct AlwaysSpeculatableImplTrait
 //===----------------------------------------------------------------------===//
 
 namespace MemoryEffects {
+/// Defines the priority of the different memory effects.
+///
+/// Sorting/ordering memory effects of an operation is done based on
+/// their defined stage and priority, in that order. If stage values for two
+/// effect instances are equal, they are then sorted by priority. Lower priority
+/// values indicate higher precedence.
+enum Priority {
+  DefaultPriority = 0,
+  AllocPriority = 1,
+  FreePriority = 2,
+  ReadPriority = 3,
+  WritePriority = 4
+};
+
 /// This class represents the base class used for memory effects.
 struct Effect : public SideEffects::Effect {
   using SideEffects::Effect::Effect;
@@ -355,28 +370,50 @@ 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 = Priority::DefaultPriority;
 };
 using EffectInstance = SideEffects::EffectInstance<Effect>;
 
+/// Returns vector of the op's memory effects sorted in increasing stage order
+/// and in increasing priority order within each stage.
+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 = Priority::AllocPriority;
+  }
+};
 
 /// 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 = Priority::FreePriority; }
+};
 
 /// 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 = Priority::ReadPriority; }
+};
 
 /// 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 = Priority::WritePriority; }
+};
 } // namespace MemoryEffects
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Transforms/LoopInvariantCodeMotionUtils.h

@@ -9,9 +9,11 @@
 #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 +23,40 @@ 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.
+///
+/// Typical usage:
+/// \code
+///   LoopLikeOpInterface myLoop = ...;
+///   DenseMap<TypeID, std::pair<bool, MemoryEffects::EffectInstance>>
+///   myConflicts; gatherResourceConflicts(myLoop, myLoop.getOperation(),
+///   resourceConflicts);
+/// \endcode
+///
+/// \param loop The loop to gather resource conflicts for.
+/// \param op The operation to gather resource conflicts for,
+/// typically the loop op itself via loop.getOperation().
+/// \param resourceConflicts Map to store potential resource conflicts.
+/// Key is the resource ID that effects are applied to. Value is a pair of
+/// a boolean, indicating if the resource has a conflict, and the last effect
+/// that was applied to the resource (if no conflicts exist) or the effect
+/// that caused the conflict (if conflicts exist). This was done so we
+/// check the effect that causes the conflict for debugging purposes.
+/// First call should use loop = someLoop and op = someLoop.getOperation()
+///
+/// resourceConflicts is modified by the function and will be non-empty
+/// as long as there are memory effects within the loop, even if there are
+/// no conflicts.
+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 +99,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

@@ -8,7 +8,9 @@
 
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
+#include "mlir/IR/Dominance.h"
 #include "mlir/IR/SymbolTable.h"
+#include <optional>
 #include <utility>
 
 using namespace mlir;
@@ -317,14 +319,45 @@ bool mlir::wouldOpBeTriviallyDead(Operation *op) {
   return wouldOpBeTriviallyDeadImpl(op);
 }
 
+llvm::SmallVector<MemoryEffects::EffectInstance>
+mlir::MemoryEffects::getMemoryEffectsSorted(Operation *op) {
+  llvm::SmallVector<MemoryEffects::EffectInstance> effectsSorted;
+
+  auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
+
+  if (!memInterface)
+    return effectsSorted; // return empty vec
+
+  memInterface.getEffects(effectsSorted);
+
+  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) {
   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,6 +371,7 @@ bool mlir::isMemoryEffectFree(Operation *op) {
     for (Operation &op : region.getOps())
       if (!isMemoryEffectFree(&op))
         return false;
+
   return true;
 }