removed memInit and refactored LICM

Author: mbagherbeikTT

mlir/include/mlir/Interfaces/SideEffectInterfaces.h

@@ -15,6 +15,7 @@
 #define MLIR_INTERFACES_SIDEEFFECTINTERFACES_H
 
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/Dominance.h"
 
 namespace mlir {
 namespace SideEffects {
@@ -378,12 +379,6 @@ struct Read : public Effect::Base<Read> {};
 /// dereference or read.
 struct Write : public Effect::Base<Write> {};
 
-// The following effect indicates that the operation initializes some
-// memory resource to a known value i.e., an idempotent MemWrite.
-// An 'init' effect implies only mutating a resource in a way that's
-// identical across calls if inputs are the same, and not any visible
-// dereference or read.
-struct Init : public Effect::Base<Init> {};
 } // namespace MemoryEffects
 
 //===----------------------------------------------------------------------===//
@@ -428,13 +423,15 @@ bool isOpTriviallyDead(Operation *op);
 /// Note: Terminators and symbols are never considered to be trivially dead.
 bool wouldOpBeTriviallyDead(Operation *op);
 
-/// Returns true if the given operation is movable under memory effects.
+/// Returns true if the given operation is allowed to be moved under the
+/// memory effects interface.
 ///
-/// An operation is movable if any of the following are true:
-/// (1) isMemoryEffectFree(op) --> true
-/// (2) isMemoryInitMovable(op) --> true
+/// 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
+/// 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.
@@ -449,32 +446,32 @@ bool isMemoryEffectMovable(Operation *op);
 /// conditions are satisfied.
 bool isMemoryEffectFree(Operation *op);
 
-/// Returns true if the given operation has a collision-free 'Init' memory
-/// effect.
+/// Returns true if the given operation has conflict-free write effects
 ///
-/// An operation is movable if:
-/// (1) it has memory effects AND all of its memory effects are of type 'Init'
-/// (2) there are no other ops with memory effects on any ofthose same resources
-/// within the operation's region(s)
+/// An operation is conflict free:
+/// (1) all of its memory effects are of type Write
+/// (2) there are no other ops with Alloc/Free/Write effects on the same
+/// resources within the ops parent region
+/// (3) all ops in the parent region with Read effects on the same resources
+/// are dominated by the operation
 ///
-/// If the operation meets both criteria, then it is movable
-bool isMemoryInitMovable(Operation *op);
+/// If the operation meets all 3 criteria, then it is conflict free
+bool isMemoryEffectConflictFree(Operation *op);
 
-/// Returns true if op and all operations within its nested regions
-/// have >1 Memory Effects on ANY of the input resources.
+/// Returns true if op and/or any operations within its nested regions
+/// have a memory effect conflict with mainOp as defined below:
 ///
-/// The first call to this function is by an op with >=1 MemInit effect on
-/// >=1 unique resources. To check that none of these resources are in conflict
-/// with other Memory Effects, we scan the entire parent region and maintain
-/// a count of Memory Effects that apply to the resources of the original op.
-/// If any resource has more than 1 Memory Effect in that region, the resource
-/// is in conflict and the op can't be moved by LICM.
+/// 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 no resources are in conflict, the op is movable.
-bool hasMemoryEffectInitConflict(
-    Operation *op, DenseMap<TypeID, int> &resourceCounts);
+/// 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/Interfaces/SideEffectInterfaces.td

@@ -87,17 +87,6 @@ def MemWrite : MemWrite<DefaultResource, 0, PartialEffect>;
 class MemWriteAt<int stage, EffectRange range = PartialEffect>
   : MemWrite<DefaultResource, stage, range>;
 
-// The following effect indicates that the operation initializes some
-// memory resource to a known value i.e., an idempotent MemWrite. 
-// An 'init' effect implies only mutating a resource in a way that's
-// identical across calls if inputs are the same, and not any visible 
-// dereference or read.
-class MemInit<Resource resource, int stage = 0,
-               EffectRange range = PartialEffect>
-  : MemoryEffect<"::mlir::MemoryEffects::Init", resource, stage, range>;
-def MemInit : MemInit<DefaultResource, 0, PartialEffect>;
-class MemInitAt<int stage, EffectRange range = PartialEffect>
-  : MemInit<DefaultResource, stage, range>;
 
 //===----------------------------------------------------------------------===//
 // Effect Traits

mlir/lib/Interfaces/SideEffectInterfaces.cpp

@@ -8,6 +8,7 @@
 
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
+#include "mlir/IR/Dominance.h"
 #include "mlir/IR/SymbolTable.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include <unordered_set>
@@ -27,7 +28,7 @@ using namespace mlir;
 //===----------------------------------------------------------------------===//
 
 bool MemoryEffects::Effect::classof(const SideEffects::Effect *effect) {
-  return isa<Allocate, Free, Read, Write, Init>(effect);
+  return isa<Allocate, Free, Read, Write>(effect);
 }
 
 //===----------------------------------------------------------------------===//
@@ -132,7 +133,6 @@ template bool mlir::hasSingleEffect<MemoryEffects::Allocate>(Operation *);
 template bool mlir::hasSingleEffect<MemoryEffects::Free>(Operation *);
 template bool mlir::hasSingleEffect<MemoryEffects::Read>(Operation *);
 template bool mlir::hasSingleEffect<MemoryEffects::Write>(Operation *);
-template bool mlir::hasSingleEffect<MemoryEffects::Init>(Operation *);
 
 template <typename EffectTy>
 bool mlir::hasSingleEffect(Operation *op, Value value) {
@@ -162,8 +162,6 @@ template bool mlir::hasSingleEffect<MemoryEffects::Read>(Operation *,
                                                          Value value);
 template bool mlir::hasSingleEffect<MemoryEffects::Write>(Operation *,
                                                           Value value);
-template bool mlir::hasSingleEffect<MemoryEffects::Init>(Operation *,
-                                                         Value value);
 
 template <typename ValueTy, typename EffectTy>
 bool mlir::hasSingleEffect(Operation *op, ValueTy value) {
@@ -198,18 +196,13 @@ template bool
 mlir::hasSingleEffect<OpOperand *, MemoryEffects::Write>(Operation *,
                                                          OpOperand *);
 template bool
-mlir::hasSingleEffect<OpOperand *, MemoryEffects::Init>(Operation *,
-                                                        OpOperand *);
-template bool
 mlir::hasSingleEffect<OpResult, MemoryEffects::Allocate>(Operation *, OpResult);
 template bool mlir::hasSingleEffect<OpResult, MemoryEffects::Free>(Operation *,
                                                                    OpResult);
 template bool mlir::hasSingleEffect<OpResult, MemoryEffects::Read>(Operation *,
                                                                    OpResult);
 template bool mlir::hasSingleEffect<OpResult, MemoryEffects::Write>(Operation *,
                                                                     OpResult);
-template bool mlir::hasSingleEffect<OpResult, MemoryEffects::Init>(Operation *,
-                                                                   OpResult);
 template bool
 mlir::hasSingleEffect<BlockArgument, MemoryEffects::Allocate>(Operation *,
                                                               BlockArgument);
@@ -222,9 +215,6 @@ mlir::hasSingleEffect<BlockArgument, MemoryEffects::Read>(Operation *,
 template bool
 mlir::hasSingleEffect<BlockArgument, MemoryEffects::Write>(Operation *,
                                                            BlockArgument);
-template bool
-mlir::hasSingleEffect<BlockArgument, MemoryEffects::Init>(Operation *,
-                                                          BlockArgument);
 
 template <typename... EffectTys>
 bool mlir::hasEffect(Operation *op) {
@@ -241,7 +231,6 @@ template bool mlir::hasEffect<MemoryEffects::Allocate>(Operation *);
 template bool mlir::hasEffect<MemoryEffects::Free>(Operation *);
 template bool mlir::hasEffect<MemoryEffects::Read>(Operation *);
 template bool mlir::hasEffect<MemoryEffects::Write>(Operation *);
-template bool mlir::hasEffect<MemoryEffects::Init>(Operation *);
 template bool
 mlir::hasEffect<MemoryEffects::Write, MemoryEffects::Free>(Operation *);
 
@@ -263,7 +252,6 @@ template bool mlir::hasEffect<MemoryEffects::Allocate>(Operation *,
 template bool mlir::hasEffect<MemoryEffects::Free>(Operation *, Value value);
 template bool mlir::hasEffect<MemoryEffects::Read>(Operation *, Value value);
 template bool mlir::hasEffect<MemoryEffects::Write>(Operation *, Value value);
-template bool mlir::hasEffect<MemoryEffects::Init>(Operation *, Value value);
 template bool
 mlir::hasEffect<MemoryEffects::Write, MemoryEffects::Free>(Operation *,
                                                            Value value);
@@ -289,8 +277,6 @@ template bool mlir::hasEffect<OpOperand *, MemoryEffects::Read>(Operation *,
                                                                 OpOperand *);
 template bool mlir::hasEffect<OpOperand *, MemoryEffects::Write>(Operation *,
                                                                  OpOperand *);
-template bool mlir::hasEffect<OpOperand *, MemoryEffects::Init>(Operation *,
-                                                                OpOperand *);
 template bool
 mlir::hasEffect<OpOperand *, MemoryEffects::Write, MemoryEffects::Free>(
     Operation *, OpOperand *);
@@ -303,8 +289,6 @@ template bool mlir::hasEffect<OpResult, MemoryEffects::Read>(Operation *,
                                                              OpResult);
 template bool mlir::hasEffect<OpResult, MemoryEffects::Write>(Operation *,
                                                               OpResult);
-template bool mlir::hasEffect<OpResult, MemoryEffects::Init>(Operation *,
-                                                             OpResult);
 template bool
 mlir::hasEffect<OpResult, MemoryEffects::Write, MemoryEffects::Free>(
     Operation *, OpResult);
@@ -320,8 +304,6 @@ template bool
 mlir::hasEffect<BlockArgument, MemoryEffects::Write>(Operation *,
                                                      BlockArgument);
 template bool
-mlir::hasEffect<BlockArgument, MemoryEffects::Init>(Operation *, BlockArgument);
-template bool
 mlir::hasEffect<BlockArgument, MemoryEffects::Write, MemoryEffects::Free>(
     Operation *, BlockArgument);
 
@@ -334,7 +316,13 @@ bool mlir::wouldOpBeTriviallyDead(Operation *op) {
 }
 
 bool mlir::isMemoryEffectMovable(Operation *op) {
-  return (isMemoryEffectFree(op) || isMemoryInitMovable(op));
+  if (isMemoryEffectFree(op))
+    return true;
+
+  if (isMemoryEffectConflictFree(op))
+    return true;
+
+  return false;
 }
 
 bool mlir::isMemoryEffectFree(Operation *op) {
@@ -358,16 +346,15 @@ bool mlir::isMemoryEffectFree(Operation *op) {
   // free.
   for (Region &region : op->getRegions())
     for (Operation &op : region.getOps())
-      if (!isMemoryEffectFree(&op))
+      if (!isMemoryEffectMovable(&op))
         return false;
 
   return true;
 }
 
-bool mlir::isMemoryInitMovable(Operation *op) {
+bool mlir::isMemoryEffectConflictFree(Operation *op) {
   auto memInterface = dyn_cast<MemoryEffectOpInterface>(op);
   // op does not implement the memory effect op interface
-  // meaning it doesn't have any memory init effects and
   // shouldn't be flagged as movable to be conservative
   if (!memInterface) return false;
 
@@ -378,60 +365,68 @@ bool mlir::isMemoryInitMovable(Operation *op) {
   // op has interface but no effects, be conservative
   if (effects.empty()) return false;
 
-
   DenseMap<TypeID, int> resourceCounts;
 
-  // ensure op only has Init effects and gather unique
+  // ensure op only has Write effects and gather unique
   // resource names
   for (const MemoryEffects::EffectInstance &effect : effects) {
-    if (!isa<MemoryEffects::Init>(effect.getEffect()))
+    if (!isa<MemoryEffects::Write>(effect.getEffect()))
       return false;
 
     resourceCounts.try_emplace(effect.getResource()->getResourceID(), 0);
   }
 
   // op itself is good, need to check rest of its parent region
   Operation *parent = op->getParentOp();
+  mlir::DominanceInfo dom(parent);
 
   for (Region &region : parent->getRegions())
-    for (Operation &op_i : region.getOps())
-      if (hasMemoryEffectInitConflict(&op_i, resourceCounts))
+    for (Operation &opI : region.getOps())
+      if (hasMemoryEffectConflict(op, &opI, dom, resourceCounts))
         return false;
 
   return true;
 }
 
-bool mlir::hasMemoryEffectInitConflict(
-    Operation *op, DenseMap<TypeID, int> &resourceCounts) {
+bool mlir::hasMemoryEffectConflict(
+    Operation *mainOp, Operation *op,
+    mlir::DominanceInfo &dom, DenseMap<TypeID, int> &resourceCounts) {
 
   if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
-    if (!memInterface.hasNoEffect()) {
-      llvm::SmallVector<MemoryEffects::EffectInstance> effects;
-      memInterface.getEffects(effects);
-
-      // ensure op only has Init effects and gather unique
-      // resource names
-      for (const MemoryEffects::EffectInstance &effect : effects) {
-        if (!isa<MemoryEffects::Init>(effect.getEffect()))
-          return true;
+    
+    llvm::SmallVector<MemoryEffects::EffectInstance> effects;
+    memInterface.getEffects(effects);
+
+    auto isDominated = dom.properlyDominates(mainOp, op);
 
-        // only care about resources of the op that called
-        // this recursive function for the first time
-        auto resourceID = effect.getResource()->getResourceID();
+    // 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 (++resourceCounts[resourceID] > 1)
-            return true;
+      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;
+        }
       }
-      return false;
     }
   }
 
   // Recurse into the regions and ensure that nested ops don't
-  // conflict with each others MemInits
+  // conflict with each other's MemWrites
   for (Region &region : op->getRegions())
-    for (Operation &op : region.getOps()) 
-      if (hasMemoryEffectInitConflict(&op, resourceCounts))
+    for (Operation &opI : region.getOps()) 
+      if (hasMemoryEffectConflict(mainOp, &opI, dom, resourceCounts))
         return true;
 
   return false;