access enzyme_context through gutils

Author: vchuravy

src/Enzyme.jl

@@ -133,6 +133,7 @@ Base.convert(::Type{API.CDerivativeMode}, ::ForwardMode) = API.DEM_ForwardMode
 function guess_activity end
 
 mutable struct EnzymeContext
+    world::UInt64
 end
 
 include("logic.jl")

src/api.jl

@@ -579,6 +579,15 @@ EnzymeCloneFunctionWithoutReturnOrArgs(fn::LLVM.Function, keepret, args) = ccall
 EnzymeGetShadowType(width, T) =
     ccall((:EnzymeGetShadowType, libEnzyme), LLVMTypeRef, (UInt64, LLVMTypeRef), width, T)
 
+function EnzymeGradientUtilsGetExternalContext(gutils)
+    ccall(
+        (:EnzymeGradientUtilsGetExternalContext, libEnzyme),
+        Ptr{Cvoid},
+        (EnzymeGradientUtilsRef,),
+        gutils,
+    )
+end
+
 EnzymeGradientUtilsReplaceAWithB(gutils, a, b) = ccall(
     (:EnzymeGradientUtilsReplaceAWithB, libEnzyme),
     Cvoid,

src/compiler.jl

@@ -223,8 +223,12 @@ include("compiler/utils.jl")
 
 include("compiler/orcv2.jl")
 
-include("gradientutils.jl")
-
+import .Enzyme: GradientUtils, call_samefunc_with_inverted_bundles!,
+                get_width, get_mode, get_runtime_activity,
+                get_strong_zero, get_shadow_type, get_uncacheable,
+                erase_with_placeholder, is_constant_value, is_constant_inst,
+                new_from_original, lookup_value, invert_pointer, debug_from_orig!,
+                add_reverse_block!, set_reverse_block!, enzyme_context
 
 # Julia function to LLVM stem and arity
 const cmplx_known_ops =
@@ -2504,7 +2508,7 @@ function enzyme!(
         convert(API.CDIFFE_TYPE, rt)
     end
 
-    enzyme_context = EnzymeContext()
+    enzyme_context = EnzymeContext(job.world)
     GC.@preserve enzyme_context begin
     LLVM.@dispose logic  = Logic(enzyme_context) begin
 

src/gradientutils.jl

@@ -66,6 +66,7 @@ erase_with_placeholder(
     orig::LLVM.Instruction,
     erase::Bool = true,
 ) = API.EnzymeGradientUtilsEraseWithPlaceholder(gutils, inst, orig, erase)
+
 is_constant_value(gutils::GradientUtils, val::LLVM.Value) =
     API.EnzymeGradientUtilsIsConstantValue(gutils, val) != 0
 
@@ -96,4 +97,16 @@ end
 
 function set_reverse_block!(gutils::GradientUtils, block::LLVM.BasicBlock)
     return LLVM.BasicBlock(API.EnzymeGradientUtilsSetReverseBlock(gutils, block))
-end
+end
+
+function enzyme_context(gutils::GradientUtils)
+    ptr = API.EnzymeGradientUtilsGetExternalContext(gutils)
+    @assert ptr != C_NULL
+    return unsafe_pointer_to_objref(ptr)::EnzymeContext
+end
+
+function enzyme_gutils_context(gutils::API.EnzymeGradientUtilsRef)
+    ptr = API.EnzymeGradientUtilsGetExternalContext(gutils)
+    @assert ptr != C_NULL
+    return unsafe_pointer_to_objref(ptr)::EnzymeContext
+end

src/logic.jl

@@ -18,7 +18,7 @@ function enzyme_context(logic::Logic)
     return logic.ctx::EnzymeContext
 end
 
-function enzyme_context(logic::API.EnzymeLogicRef)
+function enzyme_logic_context(logic::API.EnzymeLogicRef)
     ptr = API.LogicGetExternalContext(logic)
     @assert ptr != C_NULL
     return unsafe_pointer_to_objref(ptr)::EnzymeContext