Fix mincut remat bug (rust-lang#893)

EnzymeAD · Oct 12, 2022 · b359e2c · b359e2c
1 parent e3c7867
commit b359e2c
Showing 1 changed file with 10 additions and 5 deletions.
diff --git a/enzyme/Enzyme/DifferentialUseAnalysis.h b/enzyme/Enzyme/DifferentialUseAnalysis.h
@@ -1020,25 +1020,30 @@ static inline void minCut(const DataLayout &DL, LoopInfo &OrigLI,
     for (auto U : V->users()) {
       if (auto I = dyn_cast<Instruction>(U)) {
         for (auto pair : rematerializableAllocations) {
-          if (Intermediates.count(pair.first) && pair.second.stores.count(I))
-            G[Node(V, true)].insert(Node(pair.first, false));
+          if (Intermediates.count(pair.first) && pair.second.stores.count(I)) {
+            if (V != pair.first)
+              G[Node(V, true)].insert(Node(pair.first, false));
+          }
         }
       }
       if (Intermediates.count(U)) {
-        G[Node(V, true)].insert(Node(U, false));
+        if (V != U)
+          G[Node(V, true)].insert(Node(U, false));
       }
     }
   }
   for (auto pair : rematerializableAllocations) {
     if (Intermediates.count(pair.first)) {
       for (LoadInst *L : pair.second.loads) {
         if (Intermediates.count(L)) {
-          G[Node(pair.first, true)].insert(Node(L, false));
+          if (L != pair.first)
+            G[Node(pair.first, true)].insert(Node(L, false));
         }
       }
       for (auto L : pair.second.loadLikeCalls) {
         if (Intermediates.count(L.loadCall)) {
-          G[Node(pair.first, true)].insert(Node(L.loadCall, false));
+          if (L.loadCall != pair.first)
+            G[Node(pair.first, true)].insert(Node(L.loadCall, false));
         }
       }
     }