Batched no jlcall (#2226)

* Batched no jlcall * fix
EnzymeAD · Dec 24, 2024 · 1ee964a · 1ee964a
1 parent a2db43b
commit 1ee964a
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 2 deletions.
diff --git a/src/llvm/attributes.jl b/src/llvm/attributes.jl
@@ -1,4 +1,5 @@
 const nofreefns = Set{String}((
+    "jl_egal__unboxed", "ijl_egal__unboxed",
     "jl_restore_excstack",
     "ijl_restore_excstack",
     "ClientGetDevice",
@@ -154,6 +155,7 @@ const nofreefns = Set{String}((
 ))
 
 const inactivefns = Set{String}((
+    "jl_egal__unboxed", "ijl_egal__unboxed",
     "ClientGetDevice",
     "BufferOnCPU",
     "pcre2_match_data_create_from_pattern_8",

diff --git a/src/rules/llvmrules.jl b/src/rules/llvmrules.jl
@@ -177,7 +177,17 @@ include("parallelrules.jl")
     normal =
         (unsafe_load(normalR) != C_NULL) ? LLVM.Instruction(unsafe_load(normalR)) : nothing
     if shadowR != C_NULL && normal !== nothing
-        unsafe_store!(shadowR, normal.ref)
+        t_shadow = normal
+        width = get_width(gutils)
+        if width != 1 
+            t_shadow = UndefValue(
+                LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(normal))),
+            )
+            for idx = 1:width
+                t_shadow = insert_value!(B, t_shadow, normal, idx - 1)
+            end
+        end
+        unsafe_store!(shadowR, t_shadow.ref)
     end
     # Delete the primal code
     if normal !== nothing
@@ -246,7 +256,17 @@ end
     normal =
         (unsafe_load(normalR) != C_NULL) ? LLVM.Instruction(unsafe_load(normalR)) : nothing
     if shadowR != C_NULL && normal !== nothing
-        unsafe_store!(shadowR, normal.ref)
+        t_shadow = normal
+        width = get_width(gutils)
+        if width != 1 
+            t_shadow = UndefValue(
+                LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(normal))),
+            )
+            for idx = 1:width
+                t_shadow = insert_value!(B, t_shadow, normal, idx - 1)
+            end
+        end
+        unsafe_store!(shadowR, t_shadow.ref)
     end
     # Delete the primal code
     if normal !== nothing