Distribute immediately after wrap (#207)

* Fix barriers betting removed when they are required for interchange

* Fix distributing after wrap

* Make dsitrib after for respect mincut setting

* Ignore CacheLoadOp's memory effects when checking recomputability

* Use FullyRecomputable function in while wrap and interchange

* Make CacheLoad a no side effect op

* Reuse code for recomputability check

* clang-format

* Remove NoSideEffects from CacheLoad, handle it explicitly

* Fix collectEffect behaviour for CanonicalizeFor

* Fix distributing around the wrong barrier after a wrap

* Forgot to set singleExecution=true for recomputable check before ifs

* Update tests

* Comments

* Do not replicate reads if possible in distribute around barriers

* Fix some resultless operations such as yield not getting recomputed

* Fix op recalculation when distributing around barrier

* Improved check for recomputability in distribute

* Update tests

* clang format

Author: ivanradanov

include/polygeist/Ops.h

@@ -24,6 +24,11 @@
 #include "mlir/IR/PatternMatch.h"
 #include "llvm/Support/CommandLine.h"
 
+bool collectEffects(
+    mlir::Operation *op,
+    llvm::SmallVectorImpl<mlir::MemoryEffects::EffectInstance> &effects,
+    bool ignoreBarriers);
+
 bool getEffectsBefore(
     mlir::Operation *op,
     llvm::SmallVectorImpl<mlir::MemoryEffects::EffectInstance> &effects,

lib/polygeist/Ops.cpp

@@ -42,12 +42,18 @@ LogicalResult verify(BarrierOp) { return success(); }
 /// Collect the memory effects of the given op in 'effects'. Returns 'true' it
 /// could extract the effect information from the op, otherwise returns 'false'
 /// and conservatively populates the list with all possible effects.
-static bool
-collectEffects(Operation *op,
-               SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+bool collectEffects(Operation *op,
+                    SmallVectorImpl<MemoryEffects::EffectInstance> &effects,
+                    bool ignoreBarriers) {
   // Skip over barriers to avoid infinite recursion (those barriers would ask
   // this barrier again).
-  if (isa<BarrierOp>(op))
+  if (ignoreBarriers && isa<BarrierOp>(op))
+    return true;
+
+  // Ignore CacheLoads as they are already guaranteed to not have side effects
+  // in the context of a parallel op, these only exist while we are in the
+  // CPUifyPass
+  if (isa<CacheLoad>(op))
     return true;
 
   // Collect effect instances the operation. Note that the implementation of
@@ -64,7 +70,7 @@ collectEffects(Operation *op,
     for (auto &region : op->getRegions()) {
       for (auto &block : region) {
         for (auto &innerOp : block)
-          if (!collectEffects(&innerOp, effects))
+          if (!collectEffects(&innerOp, effects, ignoreBarriers))
             return false;
       }
     }
@@ -94,7 +100,7 @@ bool getEffectsBefore(Operation *op,
         else
           continue;
       }
-      if (!collectEffects(it, effects))
+      if (!collectEffects(it, effects, /* ignoreBarriers */ true))
         return false;
     }
 
@@ -114,7 +120,7 @@ bool getEffectsBefore(Operation *op,
     op->getParentOp()->walk([&](Operation *in) {
       if (conservative)
         return WalkResult::interrupt();
-      if (!collectEffects(in, effects)) {
+      if (!collectEffects(in, effects, /* ignoreBarriers */ true)) {
         conservative = true;
         return WalkResult::interrupt();
       }
@@ -134,7 +140,7 @@ bool getEffectsAfter(Operation *op,
           return true;
         continue;
       }
-      if (!collectEffects(it, effects))
+      if (!collectEffects(it, effects, /* ignoreBarriers */ true))
         return false;
     }
 
@@ -154,7 +160,7 @@ bool getEffectsAfter(Operation *op,
     op->getParentOp()->walk([&](Operation *in) {
       if (conservative)
         return WalkResult::interrupt();
-      if (!collectEffects(in, effects)) {
+      if (!collectEffects(in, effects, /* ignoreBarriers */ true)) {
         conservative = true;
         return WalkResult::interrupt();
       }

lib/polygeist/Passes/CanonicalizeFor.cpp

@@ -1521,39 +1521,6 @@ struct MoveWhileDown3 : public OpRewritePattern<WhileOp> {
   }
 };
 
-static bool
-collectEffects(Operation *op,
-               SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
-  // Collect effect instances the operation. Note that the implementation of
-  // getEffects erases all effect instances that have the type other than the
-  // template parameter so we collect them first in a local buffer and then
-  // copy.
-  if (auto iface = dyn_cast<MemoryEffectOpInterface>(op)) {
-    SmallVector<MemoryEffects::EffectInstance> localEffects;
-    iface.getEffects(localEffects);
-    llvm::append_range(effects, localEffects);
-    return true;
-  }
-  if (op->hasTrait<OpTrait::HasRecursiveSideEffects>()) {
-    for (auto &region : op->getRegions()) {
-      for (auto &block : region) {
-        for (auto &innerOp : block)
-          if (!collectEffects(&innerOp, effects))
-            return false;
-      }
-    }
-    return true;
-  }
-
-  // We need to be conservative here in case the op doesn't have the interface
-  // and assume it can have any possible effect.
-  effects.emplace_back(MemoryEffects::Effect::get<MemoryEffects::Read>());
-  effects.emplace_back(MemoryEffects::Effect::get<MemoryEffects::Write>());
-  effects.emplace_back(MemoryEffects::Effect::get<MemoryEffects::Allocate>());
-  effects.emplace_back(MemoryEffects::Effect::get<MemoryEffects::Free>());
-  return false;
-}
-
 // Rewritten from LoopInvariantCodeMotion.cpp
 struct WhileLICM : public OpRewritePattern<WhileOp> {
   using OpRewritePattern<WhileOp>::OpRewritePattern;
@@ -1573,10 +1540,10 @@ struct WhileLICM : public OpRewritePattern<WhileOp> {
         if (isReadOnly(op) && !isSpeculatable) {
 
           SmallVector<MemoryEffects::EffectInstance> whileEffects;
-          collectEffects(whileOp, whileEffects);
+          collectEffects(whileOp, whileEffects, /*ignoreBarriers*/ false);
 
           SmallVector<MemoryEffects::EffectInstance> opEffects;
-          collectEffects(op, opEffects);
+          collectEffects(op, opEffects, /*ignoreBarriers*/ false);
 
           bool conflict = false;
           for (auto before : opEffects)