Don't remove unreachable reverse (#2032)

EnzymeAD · Aug 6, 2024 · 5cf1bd2 · 5cf1bd2
1 parent 9167020
commit 5cf1bd2
Showing 1 changed file with 15 additions and 11 deletions.
diff --git a/enzyme/Enzyme/EnzymeLogic.cpp b/enzyme/Enzyme/EnzymeLogic.cpp
@@ -4319,13 +4319,15 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
     if (guaranteedUnreachable.find(&oBB) != guaranteedUnreachable.end()) {
       auto newBB = cast<BasicBlock>(gutils->getNewFromOriginal(&oBB));
       SmallVector<BasicBlock *, 4> toRemove;
-      if (auto II = dyn_cast<InvokeInst>(oBB.getTerminator())) {
-        toRemove.push_back(
-            cast<BasicBlock>(gutils->getNewFromOriginal(II->getNormalDest())));
-      } else {
-        for (auto next : successors(&oBB)) {
-          auto sucBB = cast<BasicBlock>(gutils->getNewFromOriginal(next));
-          toRemove.push_back(sucBB);
+      if (key.mode != DerivativeMode::ReverseModeCombined) {
+        if (auto II = dyn_cast<InvokeInst>(oBB.getTerminator())) {
+          toRemove.push_back(cast<BasicBlock>(
+              gutils->getNewFromOriginal(II->getNormalDest())));
+        } else {
+          for (auto next : successors(&oBB)) {
+            auto sucBB = cast<BasicBlock>(gutils->getNewFromOriginal(next));
+            toRemove.push_back(sucBB);
+          }
         }
       }
 
@@ -4354,11 +4356,13 @@ Function *EnzymeLogic::CreatePrimalAndGradient(
                             /*check*/ key.mode ==
                                 DerivativeMode::ReverseModeCombined);
       }
-      if (newBB->getTerminator())
-        gutils->erase(newBB->getTerminator());
-      IRBuilder<> builder(newBB);
-      builder.CreateUnreachable();
 
+      if (key.mode != DerivativeMode::ReverseModeCombined) {
+        if (newBB->getTerminator())
+          gutils->erase(newBB->getTerminator());
+        IRBuilder<> builder(newBB);
+        builder.CreateUnreachable();
+      }
       continue;
     }