Add world to EnzymeContext

[vchuravy]

• access enzyme_context through gutils
• assert check the world we encode in enzyme_extrace_world

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.

diff --git a/src/api.jl b/src/api.jl
index 51a4fd31..b43bb9ad 100644
--- a/src/api.jl
+++ b/src/api.jl
@@ -580,7 +580,7 @@ EnzymeGetShadowType(width, T) =
     ccall((:EnzymeGetShadowType, libEnzyme), LLVMTypeRef, (UInt64, LLVMTypeRef), width, T)
 
 function EnzymeGradientUtilsGetExternalContext(gutils)
-    ccall(
+    return ccall(
         (:EnzymeGradientUtilsGetExternalContext, libEnzyme),
         Ptr{Cvoid},
         (EnzymeGradientUtilsRef,),
diff --git a/src/compiler.jl b/src/compiler.jl
index 58d2ff4d..a570428f 100644
--- a/src/compiler.jl
+++ b/src/compiler.jl
@@ -224,11 +224,11 @@ include("compiler/utils.jl")
 include("compiler/orcv2.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, enzyme_gutils_context
+    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, enzyme_gutils_context
 
 # Julia function to LLVM stem and arity
 const cmplx_known_ops =
@@ -488,7 +488,7 @@ include("typeutils/make_zero.jl")
 
 function nested_codegen!(ctx::EnzymeContext, mode::API.CDerivativeMode, mod::LLVM.Module, @nospecialize(f), @nospecialize(tt::Type))
     funcspec = my_methodinstance(mode == API.DEM_ForwardMode ? Forward : Reverse, typeof(f), tt, ctx.world)
-    nested_codegen!(ctx, mode, mod, funcspec)
+    return nested_codegen!(ctx, mode, mod, funcspec)
 end
 
 function prepare_llvm(interp, mod::LLVM.Module, job, meta)
@@ -1235,10 +1235,10 @@ const DumpPreNestedOpt = Ref(false)
 const DumpPostNestedOpt = Ref(false)
 
 function nested_codegen!(
-    ctx::EnzymeContext,
+        ctx::EnzymeContext,
     mode::API.CDerivativeMode,
     mod::LLVM.Module,
-    funcspec::Core.MethodInstance,
+        funcspec::Core.MethodInstance,
 )
     world = ctx.world
     # TODO: Put a cache here index on `mod` and f->tt
@@ -2401,7 +2401,7 @@ const DumpPostEnzyme = Ref(false)
 const DumpPostWrap = Ref(false)
 
 function enzyme!(
-    enzyme_context::EnzymeContext,
+        enzyme_context::EnzymeContext,
     job::CompilerJob,
     interp,
     mod::LLVM.Module,
@@ -2586,7 +2586,7 @@ function enzyme!(
 
         if wrap
             augmented_primalf = create_abi_wrapper(
-                enzyme_context,
+                        enzyme_context,
                 augmented_primalf,
                 TT,
                 rt,
@@ -2595,7 +2595,7 @@ function enzyme!(
                 augmented,
                 width,
                 returnPrimal,
-                shadow_init,
+                        shadow_init,
                 interp,
                 runtimeActivity,
             )
@@ -2628,7 +2628,7 @@ function enzyme!(
         ) #=atomicAdd=#
         if wrap
             adjointf = create_abi_wrapper(
-                enzyme_context,
+                        enzyme_context,
                 adjointf,
                 TT,
                 rt,
@@ -2637,7 +2637,7 @@ function enzyme!(
                 augmented,
                 width,
                 false,
-                shadow_init,
+                        shadow_init,
                 interp,
                 runtimeActivity
             ) #=returnPrimal=#
@@ -2669,7 +2669,7 @@ function enzyme!(
         augmented_primalf = nothing
         if wrap
             adjointf = create_abi_wrapper(
-                enzyme_context,
+                        enzyme_context,
                 adjointf,
                 TT,
                 rt,
@@ -2678,7 +2678,7 @@ function enzyme!(
                 nothing,
                 width,
                 returnPrimal,
-                shadow_init,
+                        shadow_init,
                 interp,
                 runtimeActivity
             )
@@ -2714,7 +2714,7 @@ function enzyme!(
         if wrap
             pf = adjointf
             adjointf = create_abi_wrapper(
-                enzyme_context,
+                        enzyme_context,
                 adjointf,
                 TT,
                 rt,
@@ -2723,7 +2723,7 @@ function enzyme!(
                 nothing,
                 width,
                 returnPrimal,
-                shadow_init,
+                        shadow_init,
                 interp,
                 runtimeActivity
             )
@@ -2795,7 +2795,7 @@ function set_subprogram!(f::LLVM.Function, sp)
 end
 
 function create_abi_wrapper(
-    ctx::EnzymeContext,
+        ctx::EnzymeContext,
     enzymefn::LLVM.Function,
     @nospecialize(TT::Type),
     @nospecialize(rettype::Type),
@@ -2804,7 +2804,7 @@ function create_abi_wrapper(
     augmented,
     width::Int,
     returnPrimal::Bool,
-    shadow_init::Bool,
+        shadow_init::Bool,
     interp,
     runtime_activity::Bool
 )
@@ -5117,7 +5117,7 @@ end
         adjointf, augmented_primalf, TapeType = enzyme!(
             ctx,
             job,
-	        interp,
+            interp,
             mod,
             primalf,
             TT,

[wsmoses]

this should already be accessible everywhere gutils is available per enzyme_extract_world(fn) or something

[vchuravy]

This is meant as a proof of concept to demonstrate that yes we can get the world we are compiling from everywhere.

Using a string attribute is fine, but is also a blocker for caching the IR for later use, so I wanted to see if we can use this to avoid adding ephemeral data to potentially cacheable inputs.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 52.84%. Comparing base (a0e9679) to head (b623edf).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2643   +/-   ##
=======================================
  Coverage   52.84%   52.84%           
=======================================
  Files          12       12           
  Lines        1194     1194           
=======================================
  Hits          631      631           
  Misses        563      563           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[vchuravy]

@wsmoses would you be okay with a change where CustomErrorHandler always receives gutils? We can go from gutils to TypeAnalyzer (if I saw that right).

[wsmoses]

Depending on the particular error thrown there may not be a gutils created.

But we can for sure audit all the current error types and see what we can do

[wsmoses]

ok new jll is out

[github-actions]

Benchmark Results

	main	b623edf...	main / b623edf...
basics/make_zero/namedtuple	0.0545 ± 0.0021 μs	0.0533 ± 0.003 μs	1.02 ± 0.069
basics/make_zero/struct	0.259 ± 0.0062 μs	0.258 ± 0.0077 μs	1 ± 0.038
basics/overhead	4.34 ± 0.01 ns	4.64 ± 0.01 ns	0.935 ± 0.003
basics/remake_zero!/namedtuple	0.238 ± 0.0086 μs	0.239 ± 0.0081 μs	0.996 ± 0.049
basics/remake_zero!/struct	0.233 ± 0.0092 μs	0.231 ± 0.012 μs	1 ± 0.064
fold_broadcast/multidim_sum_bcast/1D	10.4 ± 1.8 μs	10.3 ± 0.24 μs	1.01 ± 0.18
fold_broadcast/multidim_sum_bcast/2D	12.1 ± 0.24 μs	12.1 ± 0.26 μs	1 ± 0.029
time_to_load	1.24 ± 0.0088 s	1.26 ± 0.027 s	0.988 ± 0.022

Benchmark Plots

A plot of the benchmark results has been uploaded as an artifact at https://github.com/EnzymeAD/Enzyme.jl/actions/runs/18928507383/artifacts/4413159292.