Correct behavior of custom call handler

EnzymeAD · Aug 12, 2021 · 3cb34f1 · 3cb34f1
1 parent d40597a
commit 3cb34f1
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/enzyme/Enzyme/AdjointGenerator.h b/enzyme/Enzyme/AdjointGenerator.h
@@ -5244,7 +5244,6 @@ class AdjointGenerator
         primalNeededInReverse = is_value_needed_in_reverse<ValueType::Primal>(
             TR, gutils, orig, Mode, Seen, oldUnreachable);
       }
-
       if (subretused && primalNeededInReverse) {
         if (normalReturn != newCall) {
           assert(normalReturn->getType() == newCall->getType());
@@ -5255,8 +5254,14 @@ class AdjointGenerator
         normalReturn = gutils->cacheForReverse(BuilderZ, normalReturn,
                                                getIndex(orig, CacheType::Self));
       } else {
-        if (!orig->mayWriteToMemory() ||
-            Mode == DerivativeMode::ReverseModeGradient)
+        if (normalReturn && normalReturn != newCall) {
+          assert(normalReturn->getType() == newCall->getType());
+          assert(Mode != DerivativeMode::ReverseModeGradient);
+          gutils->replaceAWithB(newCall, normalReturn);
+          BuilderZ.SetInsertPoint(newCall->getNextNode());
+          gutils->erase(newCall);
+        } else if (!orig->mayWriteToMemory() ||
+                   Mode == DerivativeMode::ReverseModeGradient)
           eraseIfUnused(*orig, /*erase*/ true, /*check*/ false);
       }
       return;

diff --git a/enzyme/Enzyme/DifferentialUseAnalysis.h b/enzyme/Enzyme/DifferentialUseAnalysis.h
@@ -173,6 +173,11 @@ static inline bool is_use_directly_needed_in_reverse(
       // we still need even if instruction is inactive
       if (F->getName() == "__kmpc_barrier" || F->getName() == "MPI_Barrier")
         return true;
+
+      // Since adjoint of GC preserve is another preserve in reverse
+      // we still need even if instruction is inactive
+      if (F->getName() == "llvm.julia.gc_preserve_begin")
+        return true;
     }
   }