add some more uses of enzyme_context to access the world

Author: vchuravy

src/rules/activityrules.jl

@@ -31,6 +31,7 @@ function julia_activity_rule(f::LLVM.Function)
         return
     end
     world = enzyme_extract_world(f)
+    # TODO: Access to gutils
 
     # TODO fix the attributor inlining such that this can assert always true
     if expectLen != length(parameters(f))

src/rules/customrules.jl

@@ -145,6 +145,7 @@ function enzyme_custom_setup_args(
 
     ofn = LLVM.parent(LLVM.parent(orig))
     world = enzyme_extract_world(ofn)
+    @assert world == enzyme_context(gutils).world
 
     for arg in jlargs
         @assert arg.cc != RemovedParam
@@ -454,6 +455,7 @@ function enzyme_custom_setup_ret(
     mode = get_mode(gutils)
 
     world = enzyme_extract_world(LLVM.parent(LLVM.parent(orig)))
+    @assert world == enzyme_context(gutils).world
 
     needsShadowP = Ref{UInt8}(0)
     needsPrimalP = Ref{UInt8}(0)
@@ -577,6 +579,7 @@ end
     curent_bb = position(B)
     fn = LLVM.parent(curent_bb)
     world = enzyme_extract_world(fn)
+    @assert world == enzyme_context(gutils).world
 
     llvmf = nested_codegen!(mode, mod, fmi, world)
 
@@ -812,6 +815,7 @@ end
 
     fn = LLVM.parent(LLVM.parent(orig))
     world = enzyme_extract_world(fn)
+    @assert world == enzyme_context(gutils).world
 
     C = EnzymeRules.RevConfig{
         Bool(needsPrimal),
@@ -925,6 +929,7 @@ end
 
     fn = LLVM.parent(LLVM.parent(orig))
     world = enzyme_extract_world(fn)
+    @assert world == enzyme_context(gutils).world
     @safe_debug "Trying to apply custom forward rule" TT isKWCall
 
     functy = if isKWCall
@@ -1035,6 +1040,7 @@ function enzyme_custom_common_rev(
     curent_bb = position(B)
     fn = LLVM.parent(curent_bb)
     world = enzyme_extract_world(fn)
+    @assert world == enzyme_context(gutils).world
 
     mode = get_mode(gutils)
 

src/rules/llvmrules.jl

@@ -1932,6 +1932,7 @@ end
 
     fn = LLVM.parent(LLVM.parent(orig))
     world = enzyme_extract_world(fn)
+    @assert world == enzyme_context(gutils).world
     if !guaranteed_nonactive(ET, world)
         emit_error(B, orig, "Enzyme: element type $ET of generic_memory_copyto is potentially active ($reg) and not presently supported")
     end

src/rules/parallelrules.jl

@@ -227,6 +227,7 @@ end
     modifiedBetween = (mode != API.DEM_ForwardMode, false)
 
     world = enzyme_extract_world(LLVM.parent(position(B)))
+    @assert world == enzyme_context(gutils).world
 
     pfuncT = funcT
 
@@ -550,6 +551,7 @@ end
     tt = Tuple{thunkTy,dfuncT,Bool}
     mode = get_mode(gutils)
     world = enzyme_extract_world(LLVM.parent(position(B)))
+    @assert world == enzyme_context(gutils).world
     entry = nested_codegen!(mode, mod, runtime_pfor_fwd, tt, world)
     push!(function_attributes(entry), EnumAttribute("alwaysinline"))
 
@@ -594,6 +596,7 @@ end
     }
     mode = get_mode(gutils)
     world = enzyme_extract_world(LLVM.parent(position(B)))
+    @assert world == enzyme_context(gutils).world
     entry = nested_codegen!(mode, mod, runtime_pfor_augfwd, tt, world)
     push!(function_attributes(entry), EnumAttribute("alwaysinline"))
 
@@ -627,6 +630,7 @@ end
 @register_rev function threadsfor_rev(B, orig, gutils, tape)
     mod = LLVM.parent(LLVM.parent(LLVM.parent(orig)))
     world = enzyme_extract_world(LLVM.parent(position(B)))
+    @assert world == enzyme_context(gutils).world
     if is_constant_value(gutils, orig) && is_constant_inst(gutils, orig)
         return
     end
@@ -675,6 +679,7 @@ end
     mode = get_mode(gutils)
 
     world = enzyme_extract_world(LLVM.parent(position(B)))
+    @assert world == enzyme_context(gutils).world
 
     ops = collect(operands(orig))
 
@@ -731,6 +736,7 @@ end
     ModifiedBetween = (uncacheable[1] != 0,)
 
     world = enzyme_extract_world(LLVM.parent(position(B)))
+    @assert world == enzyme_context(gutils).world
 
     ops = collect(operands(orig))
 

src/rules/typeunstablerules.jl

@@ -442,6 +442,7 @@ function newstruct_common(fwd, run, offset, B, orig, gutils, normalR, shadowR)
     width = get_width(gutils)
 
     world = enzyme_extract_world(LLVM.parent(position(B)))
+    @assert world == enzyme_context(gutils).world
 
     @assert is_constant_value(gutils, origops[offset])
     icvs = [is_constant_value(gutils, v) for v in origops[offset+1:end-1]]
@@ -932,6 +933,7 @@ end
     else
         @assert legal
         world = enzyme_extract_world(LLVM.parent(position(B)))
+        @assert world == enzyme_context(gutils).world
         if !guaranteed_nonactive(TT, world)
             unsafe_store!(tapeR, shadowres.ref)
         end
@@ -1034,6 +1036,7 @@ end
     if legal
         @assert legal
         world = enzyme_extract_world(LLVM.parent(position(B)))
+    @assert world == enzyme_context(gutils).world
         torun = !guaranteed_nonactive(TT, world)
     else
         torun = true