From a03c1d4008369c27536ab8450e808e34d4514a5f Mon Sep 17 00:00:00 2001
From: William Moses <gh@wsmoses.com>
Date: Sat, 21 Dec 2024 19:11:58 -0500
Subject: [PATCH] Fix forward rewrite (#2201)

---
 enzyme/Enzyme/CallDerivatives.cpp | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/enzyme/Enzyme/CallDerivatives.cpp b/enzyme/Enzyme/CallDerivatives.cpp
index 22df3dab9a7..db7ab6ac445 100644
--- 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;
         };