From ba915e0b08d3ba8de9b53728aa783f36f7ccd0f6 Mon Sep 17 00:00:00 2001
From: Billy Moses <wmoses@google.com>
Date: Tue, 7 Jul 2026 12:56:45 -0500
Subject: [PATCH] Fix legalRecompute of ficticious phi

---
 enzyme/Enzyme/GradientUtils.cpp | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
index 6648a57690a..e430ce52264 100644
--- a/enzyme/Enzyme/GradientUtils.cpp
+++ b/enzyme/Enzyme/GradientUtils.cpp
@@ -3974,9 +3974,9 @@ bool GradientUtils::legalRecompute(const Value *val,
 
     auto found = fictiousPHIs.find(const_cast<llvm::PHINode *>(phi));
     if (found != fictiousPHIs.end()) {
-      auto orig = found->second;
-      if (isa<AtomicRMWInst>(orig))
-        return false;
+      if (auto orig = found->second)
+        return legalRecompute(orig, available, BuilderM, reverse,
+                              legalRecomputeCache);
     }
 
     if (phi->getNumIncomingValues() == 0) {
@@ -4251,6 +4251,16 @@ bool GradientUtils::shouldRecompute(const Value *val,
 
   const Instruction *inst = cast<Instruction>(val);
 
+  if (auto phi = dyn_cast<PHINode>(val)) {
+    auto found = fictiousPHIs.find(const_cast<llvm::PHINode *>(phi));
+    if (found != fictiousPHIs.end()) {
+      auto orig = found->second;
+      if (!orig)
+        return false;
+      return shouldRecompute(orig, available, BuilderM);
+    }
+  }
+
   if (TapesToPreventRecomputation.count(inst))
     return false;