Fix forward rewrite (#2201)

EnzymeAD · Dec 22, 2024 · a03c1d4 · a03c1d4
1 parent b387a38
commit a03c1d4
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/enzyme/Enzyme/CallDerivatives.cpp b/enzyme/Enzyme/CallDerivatives.cpp
@@ -3236,6 +3236,13 @@ bool AdjointGenerator::handleKnownCallDerivatives(
           CI->setCallingConv(call.getCallingConv());
           CI->setTailCallKind(call.getTailCallKind());
           CI->setDebugLoc(dbgLoc);
+
+          if (funcName == "julia.gc_alloc_obj" ||
+              funcName == "jl_gc_alloc_typed" ||
+              funcName == "ijl_gc_alloc_typed") {
+            if (EnzymeShadowAllocRewrite)
+              EnzymeShadowAllocRewrite(wrap(CI), gutils);
+          }
           return CI;
         };