[Mlir-commits] [mlir] [MLIR][SideEffects] Added 'Init' Memory Effect which defines an Idempotent MemWrite effect and modified LICM pass (PR #153281)

[Mo Bagherbeik]

================
@@ -329,13 +356,99 @@ bool mlir::isMemoryEffectFree(Operation *op) {
 
   // Recurse into the regions and ensure that all nested ops are memory effect
   // free.
-  for (Region &region : op->getRegions())
-    for (Operation &op : region.getOps())
-      if (!isMemoryEffectFree(&op))
+  for (Region &region : op->getRegions()) {
+    for (Operation &op : region.getOps()) {
+      if (!isMemoryEffectFree(&op)) {
         return false;
+      }
+    }
+  }
   return true;
 }
 
+bool mlir::isMemoryInitMovable(Operation *op) {
+  if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
+    // 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;
+    }
+
+    std::unordered_map<std::string, int> resources;
+
+    // ensure op only has Init effects and gather unique
+    // resource names
+    for (const MemoryEffects::EffectInstance &effect : effects) {
+      if (!isa<MemoryEffects::Init>(effect.getEffect())) {
+        return false;
+      }
+
+      std::string name = effect.getResource()->getName().str();
+      resources.try_emplace(name, 0);
+    }
+
+    // op itself is good, need to check rest of its parent region
+    Operation *parent = op->getParentOp();
+
+    for (Region &region : parent->getRegions()) {
+      for (Operation &op_i : region.getOps()) {
+        if (hasMemoryEffectInitConflict(&op_i, resources)) {
+          return false;
+        }
+      }
+    }
+    return true;
+  }
+
+  // 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
+  return false;
+}
+
+bool mlir::hasMemoryEffectInitConflict(
+    Operation *op, std::unordered_map<std::string, int> &resources) {
+  if (auto memInterface = dyn_cast<MemoryEffectOpInterface>(op)) {
----------------
mbagherbeikTT wrote:

I'm not sure about this one since there's a loop that comes after that if statement

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