Use shouldrecompute of primal in shadow

EnzymeAD · Feb 12, 2025 · 089916d · 089916d
1 parent d6d5243
commit 089916d
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/enzyme/Enzyme/GradientUtils.cpp b/enzyme/Enzyme/GradientUtils.cpp
@@ -3833,6 +3833,12 @@ bool GradientUtils::legalRecompute(const Value *val,
             reverse); // TODO ADD && !TR.intType(getOriginal(dli),
                       // /*mustfind*/false).isPossibleFloat();
       }
+      if (auto ci = dyn_cast<CallInst>(uiv)) {
+        auto called = ci->getCalledFunction();
+        if (ci->hasFnAttr("enzyme_shouldrecompute") ||
+            (called && called->hasFnAttribute("enzyme_shouldrecompute")))
+          return true;
+      }
       if (phi->getNumIncomingValues() == 0) {
         return false;
       }