From eaf04196fdff67cd8e45a0cedab171680ca7464e Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Thu, 2 Oct 2025 17:27:52 +0200 Subject: [PATCH 01/12] access enzyme_context through gutils --- src/Enzyme.jl | 4 +++- src/api.jl | 9 +++++++++ src/compiler.jl | 10 +++++++--- src/gradientutils.jl | 12 ++++++++++++ src/logic.jl | 2 +- 5 files changed, 32 insertions(+), 5 deletions(-) diff --git a/src/Enzyme.jl b/src/Enzyme.jl index e4f1a92ac5..6a5e16e663 100644 --- a/src/Enzyme.jl +++ b/src/Enzyme.jl @@ -133,10 +133,12 @@ Base.convert(::Type{API.CDerivativeMode}, ::ForwardMode) = API.DEM_ForwardMode function guess_activity end mutable struct EnzymeContext + world::UInt64 modules_to_link::Vector{LLVM.Module} edges::Vector{Any} nested_cache::Dict{Core.MethodInstance, String} - EnzymeContext() = new( + EnzymeContext(world::UInt64) = new( + world, LLVM.Module[], Any[], Dict{Core.MethodInstance, String}() diff --git a/src/api.jl b/src/api.jl index 69c16604a6..ee2cfadd70 100644 --- a/src/api.jl +++ b/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, diff --git a/src/compiler.jl b/src/compiler.jl index 324d46f0c6..5d7d9635ee 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -248,8 +248,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 = @@ -5225,7 +5229,7 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT params = config.params - enzyme_context = EnzymeContext() + enzyme_context = EnzymeContext(job.world) expectedTapeType = params.expectedTapeType mode = params.mode diff --git a/src/gradientutils.jl b/src/gradientutils.jl index 37e883c5bf..630f95f346 100644 --- a/src/gradientutils.jl +++ b/src/gradientutils.jl @@ -366,3 +366,15 @@ function batch_call_same_with_inverted_arg_if_active!( return shadow 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 diff --git a/src/logic.jl b/src/logic.jl index 58a45a9cca..24750e50e5 100644 --- a/src/logic.jl +++ b/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 From 8c04b9628393a55860b30371edd792c1b8f70393 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Thu, 2 Oct 2025 17:32:08 +0200 Subject: [PATCH 02/12] assert check the world we encode in enzyme_extrace_world --- src/compiler.jl | 2 ++ src/errors.jl | 5 ++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/compiler.jl b/src/compiler.jl index 5d7d9635ee..a5ffb32ac7 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -1156,6 +1156,7 @@ function set_module_types!(interp, mod::LLVM.Module, primalf::Union{Nothing, LLV end world = enzyme_extract_world(f) + @assert world == interp.world jlargs = classify_arguments( mi.specTypes, @@ -1810,6 +1811,7 @@ function shadow_alloc_rewrite(V::LLVM.API.LLVMValueRef, gutils::API.EnzymeGradie mode == API.DEM_ReverseModeCombined fn = LLVM.parent(LLVM.parent(V)) world = enzyme_extract_world(fn) + @assert world == enzyme_gutils_context(gutils).world if !guaranteed_nonactive(Ty, world) B = LLVM.IRBuilder() position!(B, V) diff --git a/src/errors.jl b/src/errors.jl index cf38342b4b..20b7d33233 100644 --- a/src/errors.jl +++ b/src/errors.jl @@ -1164,6 +1164,7 @@ function julia_error( ) #=error=# world = enzyme_extract_world(f) end + @assert world == enzyme_gutils_context(gutils).world throw(IllegalTypeAnalysisException(msg, mi, world, sval, ir, bt)) elseif errtype == API.ET_NoType @assert B != C_NULL @@ -1288,6 +1289,7 @@ function julia_error( illegal = false created = LLVM.Instruction[] world = enzyme_extract_world(LLVM.parent(position(IRBuilder(B)))) + @assert world == enzyme_context(gutils).world width = get_width(gutils) function make_batched(@nospecialize(cur::LLVM.Value), B::LLVM.IRBuilder)::LLVM.Value if width == 1 @@ -1702,7 +1704,7 @@ end end end - mi = nothing + mi = nothing world = nothing if isa(val, LLVM.Instruction) @@ -1720,6 +1722,7 @@ end ) #=error=# world = enzyme_extract_world(f) end + @assert world == enzyme_gutils_context(gutils).world mode = Enzyme.API.DEM_ReverseModeCombined if mi !== nothing From e2faadb5a358710653d247cb1c66280bd132b364 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Thu, 2 Oct 2025 20:08:24 +0200 Subject: [PATCH 03/12] fixup! access enzyme_context through gutils --- src/compiler.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/compiler.jl b/src/compiler.jl index a5ffb32ac7..ad96b686e7 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -253,7 +253,7 @@ import .Enzyme: GradientUtils, call_samefunc_with_inverted_bundles!, 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 + add_reverse_block!, set_reverse_block!, enzyme_context, enzyme_gutils_context # Julia function to LLVM stem and arity const cmplx_known_ops = From 3bfd4d7be012858904c5454e2f0ee61851449160 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Fri, 3 Oct 2025 18:08:04 +0200 Subject: [PATCH 04/12] fixup! fixup! access enzyme_context through gutils --- src/compiler.jl | 2 +- src/errors.jl | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/compiler.jl b/src/compiler.jl index ad96b686e7..4089329304 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -1811,7 +1811,7 @@ function shadow_alloc_rewrite(V::LLVM.API.LLVMValueRef, gutils::API.EnzymeGradie mode == API.DEM_ReverseModeCombined fn = LLVM.parent(LLVM.parent(V)) world = enzyme_extract_world(fn) - @assert world == enzyme_gutils_context(gutils).world + @assert world == enzyme_context(gutils).world if !guaranteed_nonactive(Ty, world) B = LLVM.IRBuilder() position!(B, V) diff --git a/src/errors.jl b/src/errors.jl index 20b7d33233..9ba622daf9 100644 --- a/src/errors.jl +++ b/src/errors.jl @@ -1164,7 +1164,8 @@ function julia_error( ) #=error=# world = enzyme_extract_world(f) end - @assert world == enzyme_gutils_context(gutils).world + # TODO: get world from TypeAnalyzer + # @assert world == enzyme_gutils_context(gutils).world throw(IllegalTypeAnalysisException(msg, mi, world, sval, ir, bt)) elseif errtype == API.ET_NoType @assert B != C_NULL @@ -1722,7 +1723,9 @@ end ) #=error=# world = enzyme_extract_world(f) end - @assert world == enzyme_gutils_context(gutils).world + # what is data? + # Can we get world here? + # @assert world == enzyme_context(gutils).world mode = Enzyme.API.DEM_ReverseModeCombined if mi !== nothing From 7ccccd568b2eb1f9d2829791634badd4446f2f8f Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Fri, 3 Oct 2025 19:38:13 +0200 Subject: [PATCH 05/12] add some more uses of enzyme_context to access the world --- src/rules/activityrules.jl | 1 + src/rules/customrules.jl | 11 ++++++++--- src/rules/llvmrules.jl | 1 + src/rules/parallelrules.jl | 14 ++++++++------ src/rules/typeunstablerules.jl | 3 +++ 5 files changed, 21 insertions(+), 9 deletions(-) diff --git a/src/rules/activityrules.jl b/src/rules/activityrules.jl index 9405fe0ff9..fc5205a490 100644 --- a/src/rules/activityrules.jl +++ b/src/rules/activityrules.jl @@ -19,6 +19,7 @@ function julia_activity_rule(f::LLVM.Function, method_table) return end world = enzyme_extract_world(f) + # TODO: Access to gutils jlargs = classify_arguments( mi.specTypes, diff --git a/src/rules/customrules.jl b/src/rules/customrules.jl index 353f75727f..32dda2d3eb 100644 --- a/src/rules/customrules.jl +++ b/src/rules/customrules.jl @@ -336,6 +336,7 @@ function enzyme_custom_setup_args( ofn = LLVM.parent(LLVM.parent(orig)) world = enzyme_extract_world(ofn) + @assert world == enzyme_context(gutils).world jlargs = classify_arguments( mi.specTypes, @@ -862,6 +863,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) @@ -988,7 +990,7 @@ end curent_bb = position(B) fn = LLVM.parent(curent_bb) - world = enzyme_extract_world(fn) + world = enzyme_context(gutils).world # TODO: don't inject the code multiple times for multiple calls @@ -1010,8 +1012,8 @@ end width = get_width(gutils) - enzyme_ctx = Enzyme.enzyme_context(get_logic(gutils)) - llvmf = nested_codegen!(enzyme_ctx, mode, mod, fmi, world, true) + ctx = enzyme_context(gutils) + llvmf = nested_codegen!(ctx, mode, mod, fmi, ctx.world, true) orig_swiftself = has_swiftself(LLVM.called_operand(orig)) @@ -1300,6 +1302,7 @@ end fn = LLVM.parent(LLVM.parent(orig)) world = enzyme_extract_world(fn) + @assert world == enzyme_context(gutils).world C = EnzymeRules.RevConfig{ Bool(needsPrimal), @@ -1415,6 +1418,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 @@ -1585,6 +1589,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) diff --git a/src/rules/llvmrules.jl b/src/rules/llvmrules.jl index 2b5eef728d..aee6e1ddfd 100644 --- a/src/rules/llvmrules.jl +++ b/src/rules/llvmrules.jl @@ -1797,6 +1797,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 diff --git a/src/rules/parallelrules.jl b/src/rules/parallelrules.jl index 5b67f30ea0..6351c54a45 100644 --- a/src/rules/parallelrules.jl +++ b/src/rules/parallelrules.jl @@ -226,6 +226,7 @@ end modifiedBetween = (mode != API.DEM_ForwardMode, false) world = enzyme_extract_world(LLVM.parent(position(B))) + @assert world == enzyme_context(gutils).world pfuncT = funcT @@ -632,9 +633,8 @@ end tt = Tuple{thunkTy,dfuncT,Bool} mode = get_mode(gutils) - world = enzyme_extract_world(LLVM.parent(position(B))) - enzyme_ctx = Enzyme.enzyme_context(get_logic(gutils)) - entry = nested_codegen!(enzyme_ctx, mode, mod, runtime_pfor_fwd, tt, world) + ctx = enzyme_context(gutils) + entry = nested_codegen!(ctx, mode, mod, runtime_pfor_fwd, tt, ctx.world) push!(function_attributes(entry), EnumAttribute("alwaysinline")) pval = functions(mod)[sname] @@ -681,9 +681,8 @@ end Bool, } mode = get_mode(gutils) - world = enzyme_extract_world(LLVM.parent(position(B))) - enzyme_ctx = Enzyme.enzyme_context(get_logic(gutils)) - entry = nested_codegen!(enzyme_ctx, mode, mod, runtime_pfor_augfwd, tt, world) + ctx = enzyme_context(gutils) + entry = nested_codegen!(ctx, mode, mod, runtime_pfor_augfwd, tt, ctx.world) push!(function_attributes(entry), EnumAttribute("alwaysinline")) pval = functions(mod)[sname] @@ -720,6 +719,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 @@ -773,6 +773,7 @@ end mode = get_mode(gutils) world = enzyme_extract_world(LLVM.parent(position(B))) + @assert world == enzyme_context(gutils).world vals = LLVM.Value[ unsafe_to_llvm(B, runtime_newtask_fwd), @@ -827,6 +828,7 @@ end ModifiedBetween = (uncacheable[1] != 0,) world = enzyme_extract_world(LLVM.parent(position(B))) + @assert world == enzyme_context(gutils).world vals = LLVM.Value[ unsafe_to_llvm(B, runtime_newtask_augfwd), diff --git a/src/rules/typeunstablerules.jl b/src/rules/typeunstablerules.jl index 71a0bacd60..b2b9fb0b01 100644 --- a/src/rules/typeunstablerules.jl +++ b/src/rules/typeunstablerules.jl @@ -402,6 +402,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, operands(orig)[offset]) ops = @view arg_operands_view(orig)[offset+1:end] @@ -893,6 +894,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 @@ -994,6 +996,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 From f9d2b2a63527a59fae19ca6a9638b5225b5a4968 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Fri, 3 Oct 2025 19:41:46 +0200 Subject: [PATCH 06/12] add world as an argument to julia_activity_rule --- src/compiler.jl | 2 +- src/rules/activityrules.jl | 4 +--- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/src/compiler.jl b/src/compiler.jl index 4089329304..ccbd07fbab 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -709,7 +709,7 @@ end name = meth.name jlmod = meth.module - julia_activity_rule(llvmfn, method_table) + julia_activity_rule(llvmfn, world, method_table) if has_custom_rule handleCustom( state, diff --git a/src/rules/activityrules.jl b/src/rules/activityrules.jl index fc5205a490..aa14728b77 100644 --- a/src/rules/activityrules.jl +++ b/src/rules/activityrules.jl @@ -1,5 +1,5 @@ -function julia_activity_rule(f::LLVM.Function, method_table) +function julia_activity_rule(f::LLVM.Function, world, method_table) if startswith(LLVM.name(f), "japi3") || startswith(LLVM.name(f), "japi1") return end @@ -18,8 +18,6 @@ function julia_activity_rule(f::LLVM.Function, method_table) if mi.specTypes.parameters[end] === Vararg{Any} return end - world = enzyme_extract_world(f) - # TODO: Access to gutils jlargs = classify_arguments( mi.specTypes, From 8cd70adb1b6f0a6e5f785288d1e350de80b5d6b1 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Fri, 3 Oct 2025 19:43:49 +0200 Subject: [PATCH 07/12] add todo to julia_error --- src/errors.jl | 1 + 1 file changed, 1 insertion(+) diff --git a/src/errors.jl b/src/errors.jl index 9ba622daf9..bf5e4bae44 100644 --- a/src/errors.jl +++ b/src/errors.jl @@ -1723,6 +1723,7 @@ end ) #=error=# world = enzyme_extract_world(f) end + # TODO(vchuravy) # what is data? # Can we get world here? # @assert world == enzyme_context(gutils).world From ea757100b8ff17c02a5aea66793303cecf16d38f Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Tue, 7 Oct 2025 21:23:19 +0200 Subject: [PATCH 08/12] pass context to enzyme! --- src/compiler.jl | 24 ++++++++++++------------ src/rules/customrules.jl | 13 +++++-------- src/rules/parallelrules.jl | 10 ++++------ src/rules/typeunstablerules.jl | 2 +- 4 files changed, 22 insertions(+), 27 deletions(-) diff --git a/src/compiler.jl b/src/compiler.jl index ccbd07fbab..76cbbcfc05 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -529,12 +529,13 @@ include("llvm/transforms.jl") include("llvm/passes.jl") include("typeutils/make_zero.jl") -function nested_codegen!(enzyme_context::EnzymeContext, mode::API.CDerivativeMode, mod::LLVM.Module, @nospecialize(f), @nospecialize(tt::Type), world::UInt) - funcspec = my_methodinstance(mode == API.DEM_ForwardMode ? Forward : Reverse, typeof(f), tt, world) - nested_codegen!(enzyme_context, mode, mod, funcspec, world) +function nested_codegen!(enzyme_context::EnzymeContext, mode::API.CDerivativeMode, mod::LLVM.Module, @nospecialize(f), @nospecialize(tt::Type)) + funcspec = my_methodinstance(mode == API.DEM_ForwardMode ? Forward : Reverse, typeof(f), tt, enzyme_context.world) + nested_codegen!(enzyme_context, mode, mod, funcspec) end function prepare_llvm(interp, mod::LLVM.Module, job, meta) + # TODO: remove enzymejl_world for f in functions(mod) attributes = function_attributes(f) push!(attributes, StringAttribute("enzymejl_world", string(job.world))) @@ -1323,9 +1324,9 @@ function nested_codegen!( mode::API.CDerivativeMode, mod::LLVM.Module, funcspec::Core.MethodInstance, - world::UInt, alwaysinline::Bool=false, ) + world = enzyme_context.world cache_key = funcspec if haskey(enzyme_context.nested_cache, cache_key) fname = enzyme_context.nested_cache[cache_key] @@ -1352,6 +1353,7 @@ function nested_codegen!( GPUCompiler.prepare_job!(job) otherMod, meta = GPUCompiler.emit_llvm(job) + # TODO: interp should be cached since it contains internal caches interp = GPUCompiler.get_interpreter(job) prepare_llvm(interp, otherMod, job, meta) @@ -2778,7 +2780,6 @@ function enzyme!( width, returnPrimal, shadow_init, - world, interp, runtimeActivity, ) @@ -2821,7 +2822,6 @@ function enzyme!( width, false, shadow_init, - world, interp, runtimeActivity ) #=returnPrimal=# @@ -2863,7 +2863,6 @@ function enzyme!( width, returnPrimal, shadow_init, - world, interp, runtimeActivity ) @@ -2909,7 +2908,6 @@ function enzyme!( width, returnPrimal, shadow_init, - world, interp, runtimeActivity ) @@ -2993,10 +2991,10 @@ function create_abi_wrapper( width::Int, returnPrimal::Bool, shadow_init::Bool, - world::UInt, interp, runtime_activity::Bool ) + world = enzyme_context.world is_adjoint = Mode == API.DEM_ReverseModeGradient || Mode == API.DEM_ReverseModeCombined is_split = Mode == API.DEM_ReverseModeGradient || Mode == API.DEM_ReverseModePrimal needs_tape = Mode == API.DEM_ReverseModeGradient @@ -3296,6 +3294,7 @@ function create_abi_wrapper( realparms = LLVM.Value[] i = 1 + # TODO(vchuravy): remove for attr in collect(function_attributes(enzymefn)) if kind(attr) == "enzymejl_world" push!(function_attributes(llvm_f), attr) @@ -3496,7 +3495,7 @@ function create_abi_wrapper( end Func = get_func(T) funcspec = my_methodinstance(Mode == API.DEM_ForwardMode ? Forward : Reverse, Func, Tuple{}, world) - llvmf = nested_codegen!(enzyme_context, Mode, mod, funcspec, world) + llvmf = nested_codegen!(enzyme_context, Mode, mod, funcspec) push!(function_attributes(llvmf), EnumAttribute("alwaysinline", 0)) Func_RT = return_type(interp, funcspec) @assert Func_RT == NTuple{width,T′} @@ -5926,7 +5925,7 @@ end adjointf, augmented_primalf, TapeType = enzyme!( enzyme_context, job, - interp, + interp, mod, primalf, TT, @@ -6033,7 +6032,7 @@ end fname = String(name) * pf if haskey(functions(mod), fname) funcspec = my_methodinstance(Mode == API.DEM_ForwardMode ? Forward : Reverse, fnty, Tuple{JT}, job.world) - llvmf = nested_codegen!(enzyme_context, mode, mod, funcspec, job.world) + llvmf = nested_codegen!(enzyme_context, mode, mod, funcspec) llvmf = LLVM.name(llvmf) @@ -6045,6 +6044,7 @@ end llvmf = functions(mod)[llvmf] + push!(function_attributes(llvmf), StringAttribute("implements", fname)) end end diff --git a/src/rules/customrules.jl b/src/rules/customrules.jl index 32dda2d3eb..6000cd55e3 100644 --- a/src/rules/customrules.jl +++ b/src/rules/customrules.jl @@ -1012,8 +1012,7 @@ end width = get_width(gutils) - ctx = enzyme_context(gutils) - llvmf = nested_codegen!(ctx, mode, mod, fmi, ctx.world, true) + llvmf = nested_codegen!(enzyme_context(gutils), mode, mod, fmi, true) orig_swiftself = has_swiftself(LLVM.called_operand(orig)) @@ -1588,8 +1587,8 @@ 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 + ctx = enzyme_context(gutils) + world = ctx.world mode = get_mode(gutils) @@ -1655,8 +1654,7 @@ function enzyme_custom_common_rev( final_mi = nothing if forward - enzyme_ctx = Enzyme.enzyme_context(get_logic(gutils)) - llvmf = nested_codegen!(enzyme_ctx, mode, mod, ami, world, true) + llvmf = nested_codegen!(ctx, mode, mod, ami, true) @assert llvmf !== nothing rev_RT = nothing final_mi = ami @@ -1699,8 +1697,7 @@ function enzyme_custom_common_rev( rmi = rmi::Core.MethodInstance rev_RT = rev_RT::Type - enzyme_ctx = Enzyme.enzyme_context(get_logic(gutils)) - llvmf = nested_codegen!(enzyme_ctx, mode, mod, rmi, world, true) + llvmf = nested_codegen!(ctx, mode, mod, rmi, true) final_mi = rmi end diff --git a/src/rules/parallelrules.jl b/src/rules/parallelrules.jl index 6351c54a45..7fc66d0c85 100644 --- a/src/rules/parallelrules.jl +++ b/src/rules/parallelrules.jl @@ -634,7 +634,7 @@ end tt = Tuple{thunkTy,dfuncT,Bool} mode = get_mode(gutils) ctx = enzyme_context(gutils) - entry = nested_codegen!(ctx, mode, mod, runtime_pfor_fwd, tt, ctx.world) + entry = nested_codegen!(ctx, mode, mod, runtime_pfor_fwd, tt) push!(function_attributes(entry), EnumAttribute("alwaysinline")) pval = functions(mod)[sname] @@ -682,7 +682,7 @@ end } mode = get_mode(gutils) ctx = enzyme_context(gutils) - entry = nested_codegen!(ctx, mode, mod, runtime_pfor_augfwd, tt, ctx.world) + entry = nested_codegen!(ctx, mode, mod, runtime_pfor_augfwd, tt) push!(function_attributes(entry), EnumAttribute("alwaysinline")) pval = functions(mod)[sname] @@ -718,8 +718,6 @@ 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 @@ -742,8 +740,8 @@ end Bool, } mode = get_mode(gutils) - enzyme_ctx = Enzyme.enzyme_context(get_logic(gutils)) - entry = nested_codegen!(enzyme_ctx, mode, mod, runtime_pfor_rev, tt, world) + ctx = enzyme_context(gutils) + entry = nested_codegen!(ctx, mode, mod, runtime_pfor_rev, tt) push!(function_attributes(entry), EnumAttribute("alwaysinline")) pval = functions(mod)[sname] diff --git a/src/rules/typeunstablerules.jl b/src/rules/typeunstablerules.jl index b2b9fb0b01..1ba435d5b6 100644 --- a/src/rules/typeunstablerules.jl +++ b/src/rules/typeunstablerules.jl @@ -996,7 +996,7 @@ end if legal @assert legal world = enzyme_extract_world(LLVM.parent(position(B))) - @assert world == enzyme_context(gutils).world + @assert world == enzyme_context(gutils).world torun = !guaranteed_nonactive(TT, world) else torun = true From 0958a0af5e80d0241c2174c4aadc73d732fceff1 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Tue, 12 May 2026 14:59:07 +0200 Subject: [PATCH 09/12] fix rebase --- src/rules/customrules.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/rules/customrules.jl b/src/rules/customrules.jl index 6000cd55e3..15e23d7cba 100644 --- a/src/rules/customrules.jl +++ b/src/rules/customrules.jl @@ -1587,8 +1587,8 @@ function enzyme_custom_common_rev( curent_bb = position(B) fn = LLVM.parent(curent_bb) - ctx = enzyme_context(gutils) - world = ctx.world + enzyme_ctx = enzyme_context(gutils) + world = enzyme_ctx.world mode = get_mode(gutils) @@ -1654,7 +1654,7 @@ function enzyme_custom_common_rev( final_mi = nothing if forward - llvmf = nested_codegen!(ctx, mode, mod, ami, true) + llvmf = nested_codegen!(enzyme_ctx, mode, mod, ami, true) @assert llvmf !== nothing rev_RT = nothing final_mi = ami @@ -1697,7 +1697,7 @@ function enzyme_custom_common_rev( rmi = rmi::Core.MethodInstance rev_RT = rev_RT::Type - llvmf = nested_codegen!(ctx, mode, mod, rmi, true) + llvmf = nested_codegen!(enzyme_ctx, mode, mod, rmi, true) final_mi = rmi end From 8cf073aed94f2dea451cb37923813af32a1f5a49 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Tue, 12 May 2026 15:24:55 +0200 Subject: [PATCH 10/12] fixup PR --- src/compiler.jl | 11 ++++------- src/errors.jl | 3 +-- src/gradientutils.jl | 18 +++++++++--------- src/jlrt.jl | 6 +++--- src/rules/customrules.jl | 15 ++++----------- src/rules/llvmrules.jl | 19 +++++++------------ src/rules/parallelrules.jl | 9 +++------ src/rules/typeunstablerules.jl | 13 +++++-------- 8 files changed, 36 insertions(+), 58 deletions(-) diff --git a/src/compiler.jl b/src/compiler.jl index 76cbbcfc05..3aab08e6a6 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -253,7 +253,8 @@ import .Enzyme: GradientUtils, call_samefunc_with_inverted_bundles!, 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 + add_reverse_block!, set_reverse_block!, enzyme_context, enzyme_gutils_context, + batch_call_same_with_inverted_arg_if_active! # Julia function to LLVM stem and arity const cmplx_known_ops = @@ -1156,9 +1157,7 @@ function set_module_types!(interp, mod::LLVM.Module, primalf::Union{Nothing, LLV continue end - world = enzyme_extract_world(f) - @assert world == interp.world - + world = interp.world jlargs = classify_arguments( mi.specTypes, ftype, @@ -1811,9 +1810,7 @@ function shadow_alloc_rewrite(V::LLVM.API.LLVMValueRef, gutils::API.EnzymeGradie if mode == API.DEM_ReverseModePrimal || mode == API.DEM_ReverseModeGradient || mode == API.DEM_ReverseModeCombined - fn = LLVM.parent(LLVM.parent(V)) - world = enzyme_extract_world(fn) - @assert world == enzyme_context(gutils).world + world = enzyme_context(gutils).world if !guaranteed_nonactive(Ty, world) B = LLVM.IRBuilder() position!(B, V) diff --git a/src/errors.jl b/src/errors.jl index bf5e4bae44..fad06cf017 100644 --- a/src/errors.jl +++ b/src/errors.jl @@ -1289,8 +1289,7 @@ function julia_error( seen = Dict{LLVM.Value,LLVM.Value}() illegal = false created = LLVM.Instruction[] - world = enzyme_extract_world(LLVM.parent(position(IRBuilder(B)))) - @assert world == enzyme_context(gutils).world + world = enzyme_context(gutils).world width = get_width(gutils) function make_batched(@nospecialize(cur::LLVM.Value), B::LLVM.IRBuilder)::LLVM.Value if width == 1 diff --git a/src/gradientutils.jl b/src/gradientutils.jl index 630f95f346..8fecf55d2c 100644 --- a/src/gradientutils.jl +++ b/src/gradientutils.jl @@ -154,7 +154,7 @@ function get_or_insert_conditional_execute!(fn::LLVM.Function; force_run=false, ppr = preprocess(builder, args) end res = call!(builder, FT0, fn, rparms) - callconv!(res, callconv(fn)) + LLVM.callconv!(res, LLVM.callconv(fn)) end cmp = icmp!(builder, LLVM.API.LLVMIntNE, parms[1 + extra_rt], parms[1 + cmpidx + extra_rt]) @@ -167,7 +167,7 @@ function get_or_insert_conditional_execute!(fn::LLVM.Function; force_run=false, ppr = preprocess(builder, rparms) end res = call!(builder, FT0, fn, rparms) - callconv!(res, callconv(fn)) + LLVM.callconv!(res, LLVM.callconv(fn)) end if postprocess !== nothing postprocess(builder, res, rparms, ppr) @@ -249,7 +249,7 @@ function call_same_with_inverted_arg_if_active!( valTys, lookup ) - callconv!(res, callconv(orig)) + LLVM.callconv!(res, LLVM.callconv(orig)) debug_from_orig!(gutils, res, orig) if postprocess_const === nothing @@ -277,7 +277,7 @@ function call_same_with_inverted_arg_if_active!( args = collect(LLVM.Value, args) insert!(args, 1, new_from_original(gutils, origops[cmpidx])) newval = nothing - if value_type(orig) != LLVM.VoidType() && postprocess_const === nothing && need_result + if LLVM.value_type(orig) != LLVM.VoidType() && postprocess_const === nothing && need_result newval = new_from_original(gutils, orig) insert!(args, 1, newval) end @@ -298,7 +298,7 @@ function call_same_with_inverted_arg_if_active!( false, ), ) #=lookup=# - callconv!(res, callconv(orig)) + LLVM.callconv!(res, LLVM.callconv(orig)) debug_from_orig!(gutils, res, orig) if movebefore && newval !== nothing @@ -330,9 +330,9 @@ function batch_call_same_with_inverted_arg_if_active!( width = get_width(gutils) - void_rt = value_type(orig) ==LLVM.VoidType() + void_rt = LLVM.value_type(orig) ==LLVM.VoidType() shadow = if !void_rt && need_result - ST = LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(orig))) + ST = LLVM.LLVMType(API.EnzymeGetShadowType(width, LLVM.value_type(orig))) LLVM.UndefValue(ST)::LLVM.Value end @@ -343,7 +343,7 @@ function batch_call_same_with_inverted_arg_if_active!( args2 = collect(LLVM.Value, args) for i in 1:length(valTys) if valTys[i] == API.VT_Shadow - args2[i] = extract_value!(B, args2[i], idx - 1) + args2[i] = LLVM.extract_value!(B, args2[i], idx - 1) end end end @@ -354,7 +354,7 @@ function batch_call_same_with_inverted_arg_if_active!( if width == 1 shadow = res else - shadow = insert_value!(B, shadow, res, idx - 1) + shadow = LLVM.insert_value!(B, shadow, res, idx - 1) if idx == 1 norm = new_from_original(gutils, orig) if norm == res diff --git a/src/jlrt.jl b/src/jlrt.jl index ad18cc906f..d1b7221c85 100644 --- a/src/jlrt.jl +++ b/src/jlrt.jl @@ -378,7 +378,7 @@ function load_if_mixed(oval::OT, val::VT) where {OT, VT} end function val_from_byref_if_mixed(B::LLVM.IRBuilder, gutils::GradientUtils, @nospecialize(oval::LLVM.Value), @nospecialize(val::LLVM.Value))::LLVM.Value - world = enzyme_extract_world(LLVM.parent(position(B))) + world = enzyme_context(gutils).world legal, TT, _ = abs_typeof(oval) if !legal legal, TT, _ = abs_typeof(oval, true) @@ -424,8 +424,8 @@ end end end -function byref_from_val_if_mixed(B::LLVM.IRBuilder, @nospecialize(val::LLVM.Value))::LLVM.Value - world = enzyme_extract_world(LLVM.parent(position(B))) +function byref_from_val_if_mixed(B::LLVM.IRBuilder, gutils::GradientUtils, @nospecialize(val::LLVM.Value))::LLVM.Value + world = enzyme_context(gutils).world legal, TT, _ = abs_typeof(val) if !legal legal, TT, _ = abs_typeof(val, true) diff --git a/src/rules/customrules.jl b/src/rules/customrules.jl index 15e23d7cba..1caa30b263 100644 --- a/src/rules/customrules.jl +++ b/src/rules/customrules.jl @@ -334,9 +334,7 @@ function enzyme_custom_setup_args( alloctx = LLVM.IRBuilder() position!(alloctx, LLVM.BasicBlock(API.EnzymeGradientUtilsAllocationBlock(gutils))) - ofn = LLVM.parent(LLVM.parent(orig)) - world = enzyme_extract_world(ofn) - @assert world == enzyme_context(gutils).world + world = enzyme_context(gutils).world jlargs = classify_arguments( mi.specTypes, @@ -862,8 +860,7 @@ function enzyme_custom_setup_ret( width = get_width(gutils) mode = get_mode(gutils) - world = enzyme_extract_world(LLVM.parent(LLVM.parent(orig))) - @assert world == enzyme_context(gutils).world + world = enzyme_context(gutils).world needsShadowP = Ref{UInt8}(0) needsPrimalP = Ref{UInt8}(0) @@ -1299,9 +1296,7 @@ end needsShadow end - fn = LLVM.parent(LLVM.parent(orig)) - world = enzyme_extract_world(fn) - @assert world == enzyme_context(gutils).world + world = enzyme_context(gutils).world C = EnzymeRules.RevConfig{ Bool(needsPrimal), @@ -1415,9 +1410,7 @@ end end TT = Tuple{tt...} - fn = LLVM.parent(LLVM.parent(orig)) - world = enzyme_extract_world(fn) - @assert world == enzyme_context(gutils).world + world = enzyme_context(gutils).world @safe_debug "Trying to apply custom forward rule" TT isKWCall functy = if isKWCall diff --git a/src/rules/llvmrules.jl b/src/rules/llvmrules.jl index aee6e1ddfd..ec9e5d0050 100644 --- a/src/rules/llvmrules.jl +++ b/src/rules/llvmrules.jl @@ -167,9 +167,8 @@ include("parallelrules.jl") end end - pf = LLVM.parent(LLVM.parent(orig))::LLVM.Function - mi, _ = enzyme_custom_extract_mi(pf, false) #=error=# - world = enzyme_extract_world(pf) + mi, _ = enzyme_custom_extract_mi(LLVM.parent(LLVM.parent(orig)), false) #=error=# + world = enzyme_context(gutils).world if mi !== nothing err = emit_error( @@ -262,9 +261,8 @@ end end end - pf = LLVM.parent(LLVM.parent(orig))::LLVM.Function - mi, _ = enzyme_custom_extract_mi(pf, false) #=error=# - world = enzyme_extract_world(pf) + mi, _ = enzyme_custom_extract_mi(LLVM.parent(LLVM.parent(orig)), false) #=error=# + world = enzyme_context(gutils).world if mi !== nothing err = emit_error( @@ -364,9 +362,8 @@ end end end - pf = LLVM.parent(LLVM.parent(orig))::LLVM.Function - mi, _ = enzyme_custom_extract_mi(pf, false) #=error=# - world = enzyme_extract_world(pf) + mi, _ = enzyme_custom_extract_mi(LLVM.parent(LLVM.parent(orig)), false) #=error=# + world = enzyme_context(gutils).world if mi !== nothing err = emit_error( @@ -1795,9 +1792,7 @@ end ET = eltype(dest_ty) - fn = LLVM.parent(LLVM.parent(orig)) - world = enzyme_extract_world(fn) - @assert world == enzyme_context(gutils).world + 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 diff --git a/src/rules/parallelrules.jl b/src/rules/parallelrules.jl index 7fc66d0c85..0c172eb8d9 100644 --- a/src/rules/parallelrules.jl +++ b/src/rules/parallelrules.jl @@ -225,8 +225,7 @@ end e_tt = Tuple{Const{Int}} modifiedBetween = (mode != API.DEM_ForwardMode, false) - world = enzyme_extract_world(LLVM.parent(position(B))) - @assert world == enzyme_context(gutils).world + world = enzyme_context(gutils).world pfuncT = funcT @@ -770,8 +769,7 @@ end width = get_width(gutils) mode = get_mode(gutils) - world = enzyme_extract_world(LLVM.parent(position(B))) - @assert world == enzyme_context(gutils).world + world = enzyme_context(gutils).world vals = LLVM.Value[ unsafe_to_llvm(B, runtime_newtask_fwd), @@ -825,8 +823,7 @@ end uncacheable = get_uncacheable(gutils, orig) ModifiedBetween = (uncacheable[1] != 0,) - world = enzyme_extract_world(LLVM.parent(position(B))) - @assert world == enzyme_context(gutils).world + world = enzyme_context(gutils).world vals = LLVM.Value[ unsafe_to_llvm(B, runtime_newtask_augfwd), diff --git a/src/rules/typeunstablerules.jl b/src/rules/typeunstablerules.jl index 1ba435d5b6..cd65786c57 100644 --- a/src/rules/typeunstablerules.jl +++ b/src/rules/typeunstablerules.jl @@ -401,8 +401,7 @@ set_fn_max_args(runtime_tuple_rev) 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 + world = enzyme_context(gutils).world @assert is_constant_value(gutils, operands(orig)[offset]) ops = @view arg_operands_view(orig)[offset+1:end] @@ -871,7 +870,7 @@ end vals = [new_from_original(gutils, operands(orig)[1]), val_from_byref_if_mixed(B, gutils, operands(orig)[2], shadowsin)] shadowres = LLVM.call!(B, called_type(orig), LLVM.called_operand(orig), vals) callconv!(shadowres, callconv(orig)) - shadowres = byref_from_val_if_mixed(B, shadowres) + shadowres = byref_from_val_if_mixed(B, gutils, shadowres) else shadowres = UndefValue(LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(orig)))) @@ -882,7 +881,7 @@ end ] tmp = LLVM.call!(B, called_type(orig), LLVM.called_operand(orig), vals) callconv!(tmp, callconv(orig)) - tmp = byref_from_val_if_mixed(B, tmp) + tmp = byref_from_val_if_mixed(B, gutils, tmp) shadowres = insert_value!(B, shadowres, tmp, idx - 1) end end @@ -893,8 +892,7 @@ end unsafe_store!(tapeR, shadowres.ref) else @assert legal - world = enzyme_extract_world(LLVM.parent(position(B))) - @assert world == enzyme_context(gutils).world + world = enzyme_context(gutils).world if !guaranteed_nonactive(TT, world) unsafe_store!(tapeR, shadowres.ref) end @@ -995,8 +993,7 @@ end torun = false if legal @assert legal - world = enzyme_extract_world(LLVM.parent(position(B))) - @assert world == enzyme_context(gutils).world + world = enzyme_context(gutils).world torun = !guaranteed_nonactive(TT, world) else torun = true From 068d1310e018d7b20bbf401b79250fc8e83ee668 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Tue, 12 May 2026 15:32:47 +0200 Subject: [PATCH 11/12] fixup PR --- src/compiler.jl | 2 +- src/errors.jl | 6 +-- src/jlrt.jl | 82 ---------------------------------------- src/rules/customrules.jl | 3 +- 4 files changed, 5 insertions(+), 88 deletions(-) diff --git a/src/compiler.jl b/src/compiler.jl index 3aab08e6a6..74e5549f7e 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -250,7 +250,7 @@ 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, + get_strong_zero, get_atomic_add, 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, diff --git a/src/errors.jl b/src/errors.jl index fad06cf017..0205fbca49 100644 --- a/src/errors.jl +++ b/src/errors.jl @@ -1713,19 +1713,19 @@ end f, false, ) #=error=# - world = enzyme_extract_world(f) + world = enzyme_context(gutils).world elseif isa(val, LLVM.Argument) f = parent_scope(val)::LLVM.Function mi, rt = enzyme_custom_extract_mi( f, false, ) #=error=# - world = enzyme_extract_world(f) + world = enzyme_context(gutils).world end # TODO(vchuravy) # what is data? # Can we get world here? - # @assert world == enzyme_context(gutils).world + @assert world == enzyme_context(gutils).world mode = Enzyme.API.DEM_ReverseModeCombined if mi !== nothing diff --git a/src/jlrt.jl b/src/jlrt.jl index d1b7221c85..f9945fefe7 100644 --- a/src/jlrt.jl +++ b/src/jlrt.jl @@ -573,88 +573,6 @@ function emit_jltypeof!(B::LLVM.IRBuilder, @nospecialize(arg::LLVM.Value))::LLVM call!(B, FT, fn, [arg]) end -function emit_methodinstance!(B::LLVM.IRBuilder, @nospecialize(func), args::Vector{LLVM.Value})::LLVM.Value - curent_bb = position(B) - fn = LLVM.parent(curent_bb) - mod = LLVM.parent(fn) - - world = enzyme_extract_world(fn) - - sizeT = convert(LLVMType, Csize_t) - psizeT = LLVM.PointerType(sizeT) - - primalvaltys = LLVM.Value[unsafe_to_llvm(B, Core.Typeof(func))] - for a in args - push!(primalvaltys, emit_jltypeof!(B, a)) - end - - meth = only(methods(func)) - tag = emit_apply_type!(B, Tuple, primalvaltys) - - # TT = meth.sig - # while TT isa UnionAll - # TT = TT.body - # end - # parms = TT.parameters - # - # tosv = primalvaltys - # if length(parms) > 0 && typeof(parms[end]) == Core.TypeofVararg - # tosv = LLVM.Value[tosv[1:length(parms)-1]..., emit_apply_type!(B, Tuple, tosv[length(parms):end])] - # end - # sv = emit_svec!(B, tosv[2:end]) - # - - meth = unsafe_to_llvm(B, meth) - - T_jlvalue = LLVM.StructType(LLVMType[]) - T_prjlvalue = LLVM.PointerType(T_jlvalue, Tracked) - worlds, FT = get_function!( - mod, - "jl_gf_invoke_lookup_worlds", - LLVM.FunctionType(T_prjlvalue, [T_prjlvalue, T_prjlvalue, sizeT, psizeT, psizeT]), - ) - EB = LLVM.IRBuilder() - position!(EB, first(LLVM.instructions(LLVM.entry(fn)))) - minworld = alloca!(EB, sizeT) - maxworld = alloca!(EB, sizeT) - store!(B, LLVM.ConstantInt(sizeT, 0), minworld) - store!(B, LLVM.ConstantInt(sizeT, -1), maxworld) - methodmatch = call!( - B, - FT, - worlds, - LLVM.Value[ - tag, - unsafe_to_llvm(B, nothing), - LLVM.ConstantInt(sizeT, world), - minworld, - maxworld, - ], - ) - # emit_jl!(B, methodmatch) - # emit_jl!(B, emit_jltypeof!(B, methodmatch)) - offset = 1 - AT = LLVM.ArrayType(T_prjlvalue, offset + 1) - methodmatch = addrspacecast!(B, methodmatch, LLVM.PointerType(T_jlvalue, Derived)) - methodmatch = bitcast!(B, methodmatch, LLVM.PointerType(AT, Derived)) - gep = LLVM.inbounds_gep!( - B, - AT, - methodmatch, - LLVM.Value[LLVM.ConstantInt(0), LLVM.ConstantInt(offset)], - ) - sv = LLVM.load!(B, T_prjlvalue, gep) - - fn, FT = get_function!( - mod, - "jl_specializations_get_linfo", - LLVM.FunctionType(T_prjlvalue, [T_prjlvalue, T_prjlvalue, T_prjlvalue]), - ) - - mi = call!(B, FT, fn, [meth, tag, sv]) - - return mi -end function emit_writebarrier!(B::LLVM.IRBuilder, T::Vector{LLVM.Value}) curent_bb = position(B) diff --git a/src/rules/customrules.jl b/src/rules/customrules.jl index 1caa30b263..69823a729d 100644 --- a/src/rules/customrules.jl +++ b/src/rules/customrules.jl @@ -1435,8 +1435,7 @@ end end @inline function has_easy_rule_from_call(orig::LLVM.CallInst, gutils::GradientUtils)::Bool - fn = LLVM.parent(LLVM.parent(orig)) - world = enzyme_extract_world(fn) + world = enzyme_context(gutils).world mi, RealRt = enzyme_custom_extract_mi(orig) specTypes = Interpreter.simplify_kw(mi.specTypes) return EnzymeRules.has_easy_rule_from_sig(specTypes; world) From f442732db70ac4e64bc883b5cb312319c03de690 Mon Sep 17 00:00:00 2001 From: Valentin Churavy Date: Tue, 9 Jun 2026 16:47:59 +0200 Subject: [PATCH 12/12] Thread EnzymeContext for world/MI lookup instead of LLVM attributes Carry world and an MI cache on EnzymeContext, reachable from GradientUtils/TypeAnalyzer via new external-context accessors, and use them in place of per-function `enzymejl_world` attribute extraction (`enzyme_extract_world`) and attribute-scraping `enzyme_custom_extract_mi`. Drops the now-unused `emit_methodinstance!`. Also avoids a crash in the runtime-activity error path: `world` is now initialized from the context rather than left `nothing` (and a tautological assert removed), so reporting no longer asserts when the erroring value is neither an Instruction nor an Argument. Co-Authored-By: Claude Opus 4.8 --- src/Enzyme.jl | 4 +- src/absint.jl | 65 +++++++++++----------- src/api.jl | 9 ++++ src/compiler.jl | 99 ++++++++++++++++++++++++---------- src/compiler/validation.jl | 40 +++++++------- src/errors.jl | 55 ++++++++++++------- src/gradientutils.jl | 6 +++ src/jlrt.jl | 8 +-- src/llvm/transforms.jl | 4 +- src/rules/activityrules.jl | 4 +- src/rules/allocrules.jl | 2 +- src/rules/customrules.jl | 18 +++---- src/rules/jitrules.jl | 12 ++--- src/rules/llvmrules.jl | 30 +++++------ src/rules/parallelrules.jl | 6 +-- src/rules/typeunstablerules.jl | 8 +-- 16 files changed, 223 insertions(+), 147 deletions(-) diff --git a/src/Enzyme.jl b/src/Enzyme.jl index 6a5e16e663..137ccfc20d 100644 --- a/src/Enzyme.jl +++ b/src/Enzyme.jl @@ -137,11 +137,13 @@ mutable struct EnzymeContext modules_to_link::Vector{LLVM.Module} edges::Vector{Any} nested_cache::Dict{Core.MethodInstance, String} + mi_cache::Dict{String, Tuple{Core.MethodInstance, Any}} EnzymeContext(world::UInt64) = new( world, LLVM.Module[], Any[], - Dict{Core.MethodInstance, String}() + Dict{Core.MethodInstance, String}(), + Dict{String, Tuple{Core.MethodInstance, Any}}() ) end diff --git a/src/absint.jl b/src/absint.jl index 003e3cabbe..100467170f 100644 --- a/src/absint.jl +++ b/src/absint.jl @@ -13,7 +13,7 @@ function unbind(@nospecialize(val)) end end -function absint(@nospecialize(arg::LLVM.Value), partial::Bool = false, istracked::Bool=false, typetag::Bool=false)::Tuple{Bool, Any} +function absint(@nospecialize(arg::LLVM.Value), partial::Bool = false, istracked::Bool=false, typetag::Bool=false, enzyme_context::Union{EnzymeContext, Nothing} = nothing)::Tuple{Bool, Any} if (value_type(arg) == LLVM.PointerType(LLVM.StructType(LLVMType[]), Tracked)) || (value_type(arg) == LLVM.PointerType(LLVM.StructType(LLVMType[]), Derived)) || istracked ce, _ = get_base_and_offset(arg; offsetAllowed = false, inttoptr = true) if isa(ce, GlobalVariable) @@ -66,11 +66,11 @@ function absint(@nospecialize(arg::LLVM.Value), partial::Bool = false, istracked if isa(arg, ConstantExpr) if opcode(arg) == LLVM.API.LLVMAddrSpaceCast || opcode(arg) == LLVM.API.LLVMBitCast - return absint(operands(arg)[1], partial, false, typetag) + return absint(operands(arg)[1], partial, false, typetag, enzyme_context) end end if isa(arg, LLVM.BitCastInst) || isa(arg, LLVM.AddrSpaceCastInst) || isa(arg, LLVM.IntToPtrInst) - return absint(operands(arg)[1], partial, false, typetag) + return absint(operands(arg)[1], partial, false, typetag, enzyme_context) end if isa(arg, LLVM.CallInst) fn = LLVM.called_operand(arg) @@ -102,13 +102,13 @@ function absint(@nospecialize(arg::LLVM.Value), partial::Bool = false, istracked end end if nm == "julia.pointer_from_objref" - return absint(operands(arg)[1], partial) + return absint(operands(arg)[1], partial, false, false, enzyme_context) end if nm == "julia.gc_loaded" - return absint(operands(arg)[2], partial) + return absint(operands(arg)[2], partial, false, false, enzyme_context) end if nm == "jl_typeof" || nm == "ijl_typeof" - vals = abs_typeof(operands(arg)[1], partial) + vals = abs_typeof(operands(arg)[1], partial, Set{LLVM.PHIInst}(), enzyme_context) return (vals[1], vals[2]) end if LLVM.callconv(arg) == 37 || nm == "julia.call" @@ -121,10 +121,10 @@ function absint(@nospecialize(arg::LLVM.Value), partial::Bool = false, istracked if nm == "jl_f_apply_type" || nm == "ijl_f_apply_type" index += 1 found = Any[] - legal, Ty = absint(operands(arg)[index], partial) + legal, Ty = absint(operands(arg)[index], partial, false, false, enzyme_context) unionalls = TypeVar[] for sarg in @view arg_operands_view(arg)[index+1:end] - slegal, foundv = absint(sarg, partial) + slegal, foundv = absint(sarg, partial, false, false, enzyme_context) if slegal push!(found, foundv) elseif partial @@ -150,7 +150,7 @@ function absint(@nospecialize(arg::LLVM.Value), partial::Bool = false, istracked found = Any[] legal = true for sarg in @view arg_operands_view(arg)[index:end] - slegal, foundv = absint(sarg, partial) + slegal, foundv = absint(sarg, partial, false, false, enzyme_context) if slegal push!(found, foundv) else @@ -361,7 +361,8 @@ const TypesNotToDisect = Set{Type}([BigFloat]) function abs_typeof( @nospecialize(arg::LLVM.Value), - partial::Bool = false, seenphis = Set{LLVM.PHIInst}() + partial::Bool = false, seenphis = Set{LLVM.PHIInst}(), + enzyme_context::Union{EnzymeContext, Nothing} = nothing )::Union{Tuple{Bool, Type, GPUCompiler.ArgumentCC}, Tuple{Bool, Nothing, Nothing}} if (value_type(arg) == LLVM.PointerType(LLVM.StructType(LLVMType[]), Tracked)) || (value_type(arg) == LLVM.PointerType(LLVM.StructType(LLVMType[]), Derived)) ce, _ = get_base_and_offset(arg; offsetAllowed = false, inttoptr = true) @@ -406,11 +407,11 @@ function abs_typeof( if isa(arg, ConstantExpr) if opcode(arg) == LLVM.API.LLVMAddrSpaceCast || opcode(arg) == LLVM.API.LLVMBitCast - return abs_typeof(operands(arg)[1], partial, seenphis) + return abs_typeof(operands(arg)[1], partial, seenphis, enzyme_context) end end if isa(arg, LLVM.BitCastInst) || isa(arg, LLVM.AddrSpaceCastInst) || isa(arg, LLVM.IntToPtrInst) - return abs_typeof(operands(arg)[1], partial, seenphis) + return abs_typeof(operands(arg)[1], partial, seenphis, enzyme_context) end if isa(arg, LLVM.AllocaInst) || isa(arg, LLVM.CallInst) @@ -433,11 +434,11 @@ function abs_typeof( end if nm == "julia.pointer_from_objref" - return abs_typeof(operands(arg)[1], partial, seenphis) + return abs_typeof(operands(arg)[1], partial, seenphis, enzyme_context) end if nm == "julia.gc_loaded" - legal, res, byref = abs_typeof(operands(arg)[2], partial, seenphis) + legal, res, byref = abs_typeof(operands(arg)[2], partial, seenphis, enzyme_context) return legal, res, byref end @@ -466,14 +467,14 @@ function abs_typeof( if nm == "julia.gc_alloc_obj" || nm == "jl_gc_alloc_typed" || nm == "ijl_gc_alloc_typed" - vals = absint(operands(arg)[3], partial, false, #=typetag=#true) + vals = absint(operands(arg)[3], partial, false, true, enzyme_context) @assert !(vals[2] isa Core.Binding) return (vals[1], vals[2], vals[1] ? GPUCompiler.BITS_REF : nothing) end # Type tag is arg 3 if nm == "jl_alloc_genericmemory_unchecked" || nm == "ijl_alloc_genericmemory_unchecked" - vals = absint(operands(arg)[3], partial, true, #=typetag=#true) + vals = absint(operands(arg)[3], partial, true, true, enzyme_context) @assert !(vals[2] isa Core.Binding) return (vals[1], vals[2], vals[1] ? GPUCompiler.MUT_REF : nothing) end @@ -488,13 +489,13 @@ function abs_typeof( nm == "ijl_new_array" || nm == "jl_alloc_genericmemory" || nm == "ijl_alloc_genericmemory" - vals = absint(operands(arg)[1], partial, false, #=typetag=#true) + vals = absint(operands(arg)[1], partial, false, true, enzyme_context) @assert !(vals[2] isa Core.Binding) return (vals[1], vals[2], vals[1] ? GPUCompiler.MUT_REF : nothing) end if nm == "jl_new_structt" || nm == "ijl_new_structt" - vals = absint(operands(arg)[1], partial, false, #=typetag=#true) + vals = absint(operands(arg)[1], partial, false, true, enzyme_context) @assert !(vals[2] isa Core.Binding) return (vals[1], vals[2], vals[1] ? GPUCompiler.MUT_REF : nothing) end @@ -513,7 +514,7 @@ function abs_typeof( if nm == "jl_new_structv" || nm == "ijl_new_structv" @assert index == 2 - vals = absint(operands(arg)[index], partial, false, #=typetag=#true) + vals = absint(operands(arg)[index], partial, false, true, enzyme_context) @assert !(vals[2] isa Core.Binding) return (vals[1], vals[2], vals[1] ? GPUCompiler.MUT_REF : nothing) end @@ -524,7 +525,7 @@ function abs_typeof( unionalls = TypeVar[] legal = true for sarg in @view arg_operands_view(arg)[index:end] - slegal, foundv, _ = abs_typeof(sarg, partial, seenphis) + slegal, foundv, _ = abs_typeof(sarg, partial, seenphis, enzyme_context) if slegal push!(found, foundv) elseif partial @@ -547,11 +548,11 @@ function abs_typeof( if nm == "jl_f__apply_iterate" || nm == "ijl_f__apply_iterate" index += 1 - legal, iterfn = absint(operands(arg)[index]) + legal, iterfn = absint(operands(arg)[index], false, false, false, enzyme_context) iterfn = unbind(iterfn) index += 1 if legal && iterfn == Base.iterate - legal0, combfn = absint(operands(arg)[index]) + legal0, combfn = absint(operands(arg)[index], false, false, false, enzyme_context) combfn = unbind(combfn) index += 1 if legal0 && combfn == Core.apply_type && partial @@ -559,7 +560,7 @@ function abs_typeof( end resvals = Type[] while index != length(operands(arg)) - legal, pval, _ = abs_typeof(operands(arg)[index], partial, seenphis) + legal, pval, _ = abs_typeof(operands(arg)[index], partial, seenphis, enzyme_context) if !legal break end @@ -585,7 +586,7 @@ function abs_typeof( end if nm == "jl_array_copy" || nm == "ijl_array_copy" - legal, RT, _ = abs_typeof(operands(arg)[1], partial, seenphis) + legal, RT, _ = abs_typeof(operands(arg)[1], partial, seenphis, enzyme_context) if legal if !(RT <: Array) return (false, nothing, nothing) @@ -597,7 +598,7 @@ function abs_typeof( @static if VERSION < v"1.11-" else if nm == "jl_genericmemory_copy_slice" || nm == "ijl_genericmemory_copy_slice" - legal, RT, _ = abs_typeof(operands(arg)[1], partial, seenphis) + legal, RT, _ = abs_typeof(operands(arg)[1], partial, seenphis, enzyme_context) if legal @assert RT <: Memory return (legal, RT, GPUCompiler.MUT_REF) @@ -606,7 +607,11 @@ function abs_typeof( end end - _, RT = enzyme_custom_extract_mi(arg, false) + mi, RT = if enzyme_context !== nothing + enzyme_custom_extract_mi(enzyme_context, arg, false) + else + enzyme_custom_extract_mi(arg, false) + end if RT !== nothing llrt, sret, returnRoots = get_return_info(RT) if sret !== nothing @@ -632,7 +637,7 @@ function abs_typeof( end end larg, offset = get_base_and_offset(operands(arg)[1]) - legal, typ, byref = abs_typeof(larg, false, seenphis) + legal, typ, byref = abs_typeof(larg, false, seenphis, enzyme_context) dl = LLVM.datalayout(LLVM.parent(LLVM.parent(LLVM.parent(arg)))) @@ -784,7 +789,7 @@ function abs_typeof( indptrs = LLVM.API.LLVMGetIndices(arg) numind = LLVM.API.LLVMGetNumIndices(arg) offset = Cuint[unsafe_load(indptrs, i) for i in 1:numind] - found, typ, byref = abs_typeof(larg, partial, seenphis) + found, typ, byref = abs_typeof(larg, partial, seenphis, enzyme_context) if !found return (false, nothing, nothing) end @@ -870,7 +875,7 @@ function abs_typeof( seenphis2 = copy(seenphis) push!(seenphis2, arg) for op in ops - tmp = abs_typeof(op, partial, seenphis2) + tmp = abs_typeof(op, partial, seenphis2, enzyme_context) if resvals == nothing resvals = tmp else @@ -903,7 +908,7 @@ function abs_typeof( end end - legal, val = absint(arg, partial) + legal, val = absint(arg, partial, false, false, enzyme_context) if legal val = unbind(val) return (true, Core.Typeof(val), GPUCompiler.BITS_REF) diff --git a/src/api.jl b/src/api.jl index ee2cfadd70..ba479ce03e 100644 --- a/src/api.jl +++ b/src/api.jl @@ -1368,6 +1368,15 @@ function EnzymeTypeAnalyzerToString(typeanalyzer) ) end +function EnzymeTypeAnalyzerGetExternalContext(typeanalyzer) + ccall( + (:EnzymeTypeAnalyzerGetExternalContext, libEnzyme), + Ptr{Cvoid}, + (EnzymeTypeAnalyzerRef,), + typeanalyzer, + ) +end + function EnzymeGradientUtilsInvertedPointersToString(gutils) ccall( (:EnzymeGradientUtilsInvertedPointersToString, libEnzyme), diff --git a/src/compiler.jl b/src/compiler.jl index 74e5549f7e..a49054615b 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -535,7 +535,7 @@ function nested_codegen!(enzyme_context::EnzymeContext, mode::API.CDerivativeMod nested_codegen!(enzyme_context, mode, mod, funcspec) end -function prepare_llvm(interp, mod::LLVM.Module, job, meta) +function prepare_llvm(enzyme_context::EnzymeContext, interp, mod::LLVM.Module, job, meta) # TODO: remove enzymejl_world for f in functions(mod) attributes = function_attributes(f) @@ -550,6 +550,8 @@ function prepare_llvm(interp, mod::LLVM.Module, job, meta) RT = return_type(interp, mi) + enzyme_context.mi_cache[LLVM.name(llvmfn)] = (mi, RT) + _, _, returnRoots0 = get_return_info(RT) returnRoots = returnRoots0 !== nothing @@ -590,7 +592,7 @@ function prepare_llvm(interp, mod::LLVM.Module, job, meta) end end - fixup_1p12_sret!(llvmfn) + fixup_1p12_sret!(enzyme_context, llvmfn) end # We explicitly save the type of alloca's before they get lowered @@ -604,7 +606,7 @@ function prepare_llvm(interp, mod::LLVM.Module, job, meta) continue end if LLVM.name(fn) == "julia.gc_alloc_obj" - legal, RT, _ = abs_typeof(inst) + legal, RT, _ = abs_typeof(inst, false, Set{LLVM.PHIInst}(), enzyme_context) if legal metadata(inst)["enzymejl_gc_alloc_rt"] = MDNode(LLVM.Metadata[MDString(string(convert(UInt, unsafe_to_pointer(RT))))]) end @@ -650,7 +652,7 @@ function handleCustom(state::HandlerState, custom, k_name::String, llvmfn::LLVM. nothing end -function handle_compiled(state::HandlerState, edges::Vector, run_enzyme::Bool, mode::API.CDerivativeMode, world::UInt, method_table, custom::Dict{String, LLVM.API.LLVMLinkage}, mod::LLVM.Module, mi::Core.MethodInstance, k_name::String, @nospecialize(rettype::Type))::Nothing +function handle_compiled(enzyme_context::EnzymeContext, state::HandlerState, edges::Vector, run_enzyme::Bool, mode::API.CDerivativeMode, world::UInt, method_table, custom::Dict{String, LLVM.API.LLVMLinkage}, mod::LLVM.Module, mi::Core.MethodInstance, k_name::String, @nospecialize(rettype::Type))::Nothing has_custom_rule = false specTypes = Interpreter.simplify_kw(mi.specTypes) @@ -711,7 +713,7 @@ end name = meth.name jlmod = meth.module - julia_activity_rule(llvmfn, world, method_table) + julia_activity_rule(enzyme_context, llvmfn, world, method_table) if has_custom_rule handleCustom( state, @@ -1089,6 +1091,7 @@ end if sret cur = llvmfn == state.primalf llvmfn, _, state.boxedArgs, state.loweredArgs, state.removedRoots = lower_convention( + enzyme_context, mi.specTypes, mod, llvmfn, @@ -1131,13 +1134,13 @@ end return end -function set_module_types!(interp, mod::LLVM.Module, primalf::Union{Nothing, LLVM.Function}, job, edges, run_enzyme, mode::API.CDerivativeMode)::Tuple{Dict{String,LLVM.API.LLVMLinkage}, HandlerState} +function set_module_types!(enzyme_context::EnzymeContext, interp, mod::LLVM.Module, primalf::Union{Nothing, LLVM.Function}, job, edges, run_enzyme, mode::API.CDerivativeMode)::Tuple{Dict{String,LLVM.API.LLVMLinkage}, HandlerState} for f in functions(mod) if startswith(LLVM.name(f), "japi3") || startswith(LLVM.name(f), "japi1") continue end - mi, RT = enzyme_custom_extract_mi(f, false) + mi, RT = enzyme_custom_extract_mi(enzyme_context, f, false) if mi === nothing continue end @@ -1307,7 +1310,7 @@ function set_module_types!(interp, mod::LLVM.Module, primalf::Union{Nothing, LLV end end if mi !== nothing && RT !== nothing - handle_compiled(state, edges, run_enzyme, mode, world, method_table, custom, mod, mi, fname, RT) + handle_compiled(enzyme_context, state, edges, run_enzyme, mode, world, method_table, custom, mod, mi, fname, RT) end end @@ -1354,7 +1357,7 @@ function nested_codegen!( # TODO: interp should be cached since it contains internal caches interp = GPUCompiler.get_interpreter(job) - prepare_llvm(interp, otherMod, job, meta) + prepare_llvm(enzyme_context, interp, otherMod, job, meta) entry = name(meta.entry) @@ -1377,7 +1380,7 @@ function nested_codegen!( API.EnzymeDumpModuleRef(otherMod.ref) end - check_ir(interp, job, otherMod) + check_ir(enzyme_context, interp, job, otherMod) if DumpPreNestedOpt[] API.EnzymeDumpModuleRef(otherMod.ref) @@ -1386,7 +1389,7 @@ function nested_codegen!( # Skipped inline of blas run_enzyme = false - set_module_types!(interp, otherMod, nothing, job, edges, run_enzyme, mode) + set_module_types!(enzyme_context, interp, otherMod, nothing, job, edges, run_enzyme, mode) # Apply first stage of optimization's so that this module is at the same stage as `mod` optimize!(otherMod, JIT.get_tm()) @@ -1742,7 +1745,7 @@ function shadow_alloc_rewrite(V::LLVM.API.LLVMValueRef, gutils::API.EnzymeGradie V = LLVM.CallInst(V) gutils = GradientUtils(gutils) mode = get_mode(gutils) - has, Ty, byref = abs_typeof(V) + has, Ty, byref = abs_typeof(V, false, Set{LLVM.PHIInst}(), enzyme_context(gutils)) partial = false count = nothing if !has @@ -1780,10 +1783,10 @@ function shadow_alloc_rewrite(V::LLVM.API.LLVMValueRef, gutils::API.EnzymeGradie if nm == "jl_f_apply_type" || nm == "ijl_f_apply_type" index += 1 found = Any[] - legal, Ty = absint(operands(arg)[index], partial) + legal, Ty = absint(operands(arg)[index], partial, false, false, enzyme_context(gutils)) Ty = unbind(Ty) if legal && Ty == NTuple - legal, Ty = absint(operands(arg)[index+2]) + legal, Ty = absint(operands(arg)[index+2], false, false, false, enzyme_context(gutils)) Ty = unbind(Ty) if legal # count should represent {the total size in bytes, the aligned size of each element} @@ -2485,6 +2488,42 @@ function enzyme_extract_world(fn::LLVM.Function)::UInt throw(AssertionError("Enzyme: could not find world in $(string(fn))")) end +function enzyme_custom_extract_mi(enzyme_context::EnzymeContext, orig::LLVM.CallInst, error::Bool = true) + operand = LLVM.called_operand(orig) + if isa(operand, LLVM.Function) + return enzyme_custom_extract_mi(enzyme_context, operand::LLVM.Function, error) + elseif error + GPUCompiler.@safe_error "Enzyme: Custom handler, could not find fn", orig + end + return nothing, nothing +end + +function enzyme_custom_extract_mi(enzyme_context::EnzymeContext, orig::LLVM.Function, error::Bool = true) + name = LLVM.name(orig) + if haskey(enzyme_context.mi_cache, name) + return enzyme_context.mi_cache[name] + end + + mi = nothing + RT = nothing + for fattr in collect(function_attributes(orig)) + if isa(fattr, LLVM.StringAttribute) + if kind(fattr) == "enzymejl_mi" + ptr = reinterpret(Ptr{Cvoid}, parse(UInt, LLVM.value(fattr))) + mi = Base.unsafe_pointer_to_objref(ptr) + end + if kind(fattr) == "enzymejl_rt" + ptr = reinterpret(Ptr{Cvoid}, parse(UInt, LLVM.value(fattr))) + RT = Base.unsafe_pointer_to_objref(ptr) + end + end + end + if error && mi === nothing + GPUCompiler.@safe_error "Enzyme: Custom handler, could not find mi", orig + end + return mi, RT +end + function enzyme_custom_extract_mi(orig::LLVM.CallInst, error::Bool = true) operand = LLVM.called_operand(orig) if isa(operand, LLVM.Function) @@ -4304,6 +4343,7 @@ end # Modified from GPUCompiler/src/irgen.jl:365 lower_byval function lower_convention( + enzyme_context::EnzymeContext, @nospecialize(functy::Type), mod::LLVM.Module, entry_f::LLVM.Function, @@ -4979,7 +5019,7 @@ function lower_convention( fixup_metadata!(entry_f) - mi, rt = enzyme_custom_extract_mi(entry_f) + mi, rt = enzyme_custom_extract_mi(enzyme_context, entry_f) attributes = function_attributes(wrapper_f) push!( attributes, @@ -5279,7 +5319,7 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT edges = enzyme_context.edges primal_interp = GPUCompiler.get_interpreter(primal_job) - prepare_llvm(primal_interp, mod, primal_job, meta) + prepare_llvm(enzyme_context, primal_interp, mod, primal_job, meta) for f in functions(mod) permit_inlining!(f) end @@ -5297,7 +5337,7 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT API.EnzymeDumpModuleRef(mod.ref) end interp = GPUCompiler.get_interpreter(job) - check_ir(interp, job, mod) + check_ir(enzyme_context, interp, job, mod) if DumpPostCheck[] API.EnzymeDumpModuleRef(mod.ref) end @@ -5400,7 +5440,7 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT run!(GlobalOptPass(), mod) end - custom, state = set_module_types!(interp, mod, primalf, job, edges, params.run_enzyme, mode) + custom, state = set_module_types!(enzyme_context, interp, mod, primalf, job, edges, params.run_enzyme, mode) primalf = state.primalf must_wrap = state.must_wrap @@ -5434,7 +5474,7 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT push!(return_attributes(wrapper_f), attr) end - mi, rt = enzyme_custom_extract_mi(primalf) + mi, rt = enzyme_custom_extract_mi(enzyme_context, primalf) let builder = IRBuilder() entry = BasicBlock(wrapper_f, "entry") @@ -5500,6 +5540,7 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT if state.lowerConvention primalf, returnRoots, boxedArgs, loweredArgs, removedRoots, actualRetType = lower_convention( + enzyme_context, source_sig, mod, primalf, @@ -5614,7 +5655,7 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT if length(calluse) > 0 RTs = Union{Nothing, Type}[] for cu in calluse - _, RT = enzyme_custom_extract_mi(cu, false) + _, RT = enzyme_custom_extract_mi(enzyme_context, cu, false) push!(RTs, RT) end @assert all(RTs[1] == RT for RT in RTs) @@ -5640,7 +5681,7 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT if !API.HasFromStack(inst) && ((isa(inst, LLVM.CallInst) && (!isa(fn, LLVM.Function) || isempty(blocks(fn))) ) || isa(inst, LLVM.LoadInst) || isa(inst, LLVM.AllocaInst) || isa(inst, LLVM.ExtractValueInst)) - legal, source_typ, byref = abs_typeof(inst) + legal, source_typ, byref = abs_typeof(inst, false, Set{LLVM.PHIInst}(), enzyme_context) codegen_typ = value_type(inst) if legal if codegen_typ isa LLVM.PointerType || codegen_typ isa LLVM.IntegerType @@ -5673,7 +5714,7 @@ function GPUCompiler.compile_unhooked(output::Symbol, job::CompilerJob{<:EnzymeT @static if VERSION < v"1.11-" else - legal2, obj = absint(inst) + legal2, obj = absint(inst, false, false, false, enzyme_context) obj = unbind(obj) if legal2 && is_memory_instance(obj) metadata(inst)["nonnull"] = MDNode(LLVM.Metadata[]) @@ -5710,7 +5751,7 @@ end intr == LLVM.Intrinsic("llvm.memmove").id || intr == LLVM.Intrinsic("llvm.memset").id base, offset = get_base_and_offset(operands(inst)[1]) - legal, jTy, byref = abs_typeof(base) + legal, jTy, byref = abs_typeof(base, false, Set{LLVM.PHIInst}(), enzyme_context) sz = if intr == LLVM.Intrinsic("llvm.memcpy").id || intr == LLVM.Intrinsic("llvm.memmove").id @@ -5753,7 +5794,7 @@ end continue end - legal, jTy, byref = abs_typeof(inst, true) + legal, jTy, byref = abs_typeof(inst, true, Set{LLVM.PHIInst}(), enzyme_context) if !legal continue end @@ -5800,7 +5841,7 @@ end end if !mayWriteToMemory(user) - slegal, foundv, byref = abs_typeof(user) + slegal, foundv, byref = abs_typeof(user, false, Set{LLVM.PHIInst}(), enzyme_context) if slegal reg2 = active_reg(foundv, job.world) if reg2 == ActiveState || reg2 == AnyState @@ -5828,7 +5869,7 @@ end end # we are storing into the variable if operands(user)[2] == cur - slegal, foundv, byref = abs_typeof(operands(user)[1]) + slegal, foundv, byref = abs_typeof(operands(user)[1], false, Set{LLVM.PHIInst}(), enzyme_context) if slegal reg2 = active_reg(foundv, job.world) if reg2 == AnyState @@ -5864,7 +5905,7 @@ end continue end if is_readonly(called) - slegal, foundv, byref = abs_typeof(user) + slegal, foundv, byref = abs_typeof(user, false, Set{LLVM.PHIInst}(), enzyme_context) if slegal reg2 = active_reg(foundv, job.world) if reg2 == ActiveState || reg2 == AnyState @@ -5882,7 +5923,7 @@ end push!(todo, parm) end end - slegal, foundv, byref = abs_typeof(user) + slegal, foundv, byref = abs_typeof(user, false, Set{LLVM.PHIInst}(), enzyme_context) if slegal reg2 = active_reg(foundv, job.world) if reg2 == ActiveState || reg2 == AnyState @@ -5902,7 +5943,7 @@ end string(user) * ", using " * string(cur) - slegal, foundv = absint(cur) + slegal, foundv = absint(cur, false, false, false, enzyme_context) if slegal foundv = unbind(foundv) resstr *= "of type " * string(foundv) diff --git a/src/compiler/validation.jl b/src/compiler/validation.jl index 7d0d52fc0f..d2d572ac51 100644 --- a/src/compiler/validation.jl +++ b/src/compiler/validation.jl @@ -181,15 +181,15 @@ function restore_lookups(mod::LLVM.Module)::Nothing return end -function check_ir(interp, @nospecialize(job::CompilerJob), mod::LLVM.Module) - errors = check_ir!(interp, job, IRError[], mod) +function check_ir(enzyme_context::EnzymeContext, interp, @nospecialize(job::CompilerJob), mod::LLVM.Module) + errors = check_ir!(enzyme_context, interp, job, IRError[], mod) unique!(errors) return if !isempty(errors) throw(InvalidIRError(job, errors)) end end -function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRError}, mod::LLVM.Module) +function check_ir!(enzyme_context::EnzymeContext, interp, @nospecialize(job::CompilerJob), errors::Vector{IRError}, mod::LLVM.Module) imported = Set(String[]) if haskey(functions(mod), "malloc") f = functions(mod)["malloc"] @@ -213,7 +213,7 @@ function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRErr if in(f, del) continue end - check_ir!(interp, job, errors, imported, f, del, mod) + check_ir!(enzyme_context, interp, job, errors, imported, f, del, mod) end for d in del LLVM.API.LLVMDeleteFunction(d) @@ -224,7 +224,7 @@ function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRErr if in(f, del) continue end - check_ir!(interp, job, errors, imported, f, del, mod) + check_ir!(enzyme_context, interp, job, errors, imported, f, del, mod) end for d in del LLVM.API.LLVMDeleteFunction(d) @@ -233,7 +233,7 @@ function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRErr return errors end -function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRError}, imported::Set{String}, f::LLVM.Function, deletedfns::Vector{LLVM.Function}, mod::LLVM.Module) +function check_ir!(enzyme_context::EnzymeContext, interp, @nospecialize(job::CompilerJob), errors::Vector{IRError}, imported::Set{String}, f::LLVM.Function, deletedfns::Vector{LLVM.Function}, mod::LLVM.Module) calls = LLVM.CallInst[] isInline = API.EnzymeGetCLBool(cglobal((:EnzymeInline, API.libEnzyme))) != 0 mod = LLVM.parent(f) @@ -537,7 +537,7 @@ function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRErr while length(calls) > 0 inst = pop!(calls) - check_ir!(interp, job, errors, imported, inst, calls, mod) + check_ir!(enzyme_context, interp, job, errors, imported, inst, calls, mod) end return errors @@ -686,7 +686,7 @@ end import GPUCompiler: DYNAMIC_CALL, DELAYED_BINDING, RUNTIME_FUNCTION, UNKNOWN_FUNCTION, POINTER_FUNCTION import GPUCompiler: backtrace, isintrinsic -function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRError}, imported::Set{String}, inst::LLVM.CallInst, calls::Vector{LLVM.CallInst}, mod::LLVM.Module) +function check_ir!(enzyme_context::EnzymeContext, interp, @nospecialize(job::CompilerJob), errors::Vector{IRError}, imported::Set{String}, inst::LLVM.CallInst, calls::Vector{LLVM.CallInst}, mod::LLVM.Module) world = job.world method_table = Core.Compiler.method_table(interp) bt = backtrace(inst) @@ -857,7 +857,7 @@ function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRErr fname = ops[2] if isa(flib, LLVM.ConstantExpr) || isa(flib, LLVM.GlobalVariable) - legal, flib2 = absint(flib) + legal, flib2 = absint(flib, false, false, false, enzyme_context) if legal flib = unbind(flib2) end @@ -1009,12 +1009,12 @@ function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRErr # Add 1 to account for function being first arg iteroff = 2 - legal, iterlib = absint(operands(inst)[iteroff + 1]) + legal, iterlib = absint(operands(inst)[iteroff + 1], false, false, false, enzyme_context) iterlib = unbind(iterlib) if legal && iterlib == Base.iterate - legal, GT, byref = abs_typeof(operands(inst)[4 + 1], true) + legal, GT, byref = abs_typeof(operands(inst)[4 + 1], true, Set{LLVM.PHIInst}(), enzyme_context) funcoff = 3 - legal2, funclib, byref2 = abs_typeof(operands(inst)[funcoff + 1]) + legal2, funclib, byref2 = abs_typeof(operands(inst)[funcoff + 1], false, Set{LLVM.PHIInst}(), enzyme_context) if legal && (GT <: Vector || GT <: Tuple) if legal2 tys = Union{Type, Core.TypeofVararg}[funclib, Vararg{Any}] @@ -1089,17 +1089,17 @@ function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRErr if isa(dest, LLVM.Function) && in(LLVM.name(dest), keys(generic_method_offsets)) offset, start = generic_method_offsets[LLVM.name(dest)] # Add 1 to account for function being first arg - legal, flibty, byref = abs_typeof(operands(inst)[offset + 1]) + legal, flibty, byref = abs_typeof(operands(inst)[offset + 1], false, Set{LLVM.PHIInst}(), enzyme_context) if legal tys = Union{Type, Core.TypeofVararg}[flibty] for op in @view arg_operands_view(inst)[(start + 1):end] - legal, typ, byref2 = abs_typeof(op, true) + legal, typ, byref2 = abs_typeof(op, true, Set{LLVM.PHIInst}(), enzyme_context) if !legal typ = Any end push!(tys, typ) end - legal, flib = absint(operands(inst)[offset + 1]) + legal, flib = absint(operands(inst)[offset + 1], false, false, false, enzyme_context) flib = unbind(flib) if legal && isa(flib, Core.MethodInstance) if !Base.isvarargtype(flib.specTypes.parameters[end]) @@ -1244,17 +1244,17 @@ function check_ir!(interp, @nospecialize(job::CompilerJob), errors::Vector{IRErr if isa(dest, LLVM.Function) && in(LLVM.name(dest), keys(generic_method_offsets)) offset, start = generic_method_offsets[LLVM.name(dest)] - legal, flibty, byref = abs_typeof(operands(inst)[offset]) + legal, flibty, byref = abs_typeof(operands(inst)[offset], false, Set{LLVM.PHIInst}(), enzyme_context) if legal tys = Union{Type, Core.TypeofVararg}[flibty] for op in @view arg_operands_view(inst)[start:end] - legal, typ, byref2 = abs_typeof(op, true) + legal, typ, byref2 = abs_typeof(op, true, Set{LLVM.PHIInst}(), enzyme_context) if !legal typ = Any end push!(tys, typ) end - legal, flib = absint(operands(inst)[offset + 1]) + legal, flib = absint(operands(inst)[offset + 1], false, false, false, enzyme_context) flib = unbind(flib) if legal && isa(flib, Core.MethodInstance) if !Base.isvarargtype(flib.specTypes.parameters[end]) @@ -1383,7 +1383,7 @@ function rewrite_union_returns_as_ref(enzymefn::LLVM.Function, off::Int64, world end if nm == "julia.gc_alloc_obj" - legal, Ty, byref = abs_typeof(cur) + legal, Ty, byref = abs_typeof(cur, false, Set{LLVM.PHIInst}(), enzyme_context) @assert legal if !guaranteed_nonactive(Ty, world) NTy = Base.RefValue{Ty} @@ -1460,7 +1460,7 @@ function rewrite_union_returns_as_ref(enzymefn::LLVM.Function, off::Int64, world if length(off) == 0 && value_type(cur) == LLVM.PointerType(LLVM.StructType(LLVMType[]), Tracked) - legal, typ, byref = abs_typeof(cur) + legal, typ, byref = abs_typeof(cur, false, Set{LLVM.PHIInst}(), enzyme_context) if legal if guaranteed_nonactive(typ, world) continue diff --git a/src/errors.jl b/src/errors.jl index 0205fbca49..baeb5dbea5 100644 --- a/src/errors.jl +++ b/src/errors.jl @@ -1079,20 +1079,37 @@ function julia_error( mi = nothing world = nothing + gutils_ref = API.EnzymeGradientUtilsRef(data) + ctx = if gutils_ref != C_NULL + enzyme_gutils_context(gutils_ref) + else + nothing + end + if isa(val, LLVM.Instruction) f = LLVM.parent(LLVM.parent(val))::LLVM.Function mi, rt = enzyme_custom_extract_mi( + ctx, f, false, ) #=error=# - world = enzyme_extract_world(f) + world = if ctx !== nothing + ctx.world + else + enzyme_extract_world(f) + end elseif isa(val, LLVM.Argument) f = parent_scope(val)::LLVM.Function mi, rt = enzyme_custom_extract_mi( + ctx, f, false, ) #=error=# - world = enzyme_extract_world(f) + world = if ctx !== nothing + ctx.world + else + enzyme_extract_world(f) + end end if mi !== nothing emit_error(B, nothing, (msg2, mi, world), EnzymeNoDerivativeError{Core.MethodInstance, UInt}, data2) @@ -1113,7 +1130,7 @@ function julia_error( print(io, "Current scope: \n") print(io, ir) end - legal, obj = absint(val) + legal, obj = absint(val, false, false, false, enzyme_context(gutils)) if legal obj0 = obj obj = unbind(obj) @@ -1147,22 +1164,23 @@ function julia_error( API.EnzymeStringFree(ip) mi = nothing - world = nothing + ctx = enzyme_context(data) + world = ctx.world if isa(val, LLVM.Instruction) f = LLVM.parent(LLVM.parent(val))::LLVM.Function mi, rt = enzyme_custom_extract_mi( + ctx, f, false, ) #=error=# - world = enzyme_extract_world(f) elseif isa(val, LLVM.Argument) f = parent_scope(val)::LLVM.Function mi, rt = enzyme_custom_extract_mi( + ctx, f, false, ) #=error=# - world = enzyme_extract_world(f) end # TODO: get world from TypeAnalyzer # @assert world == enzyme_gutils_context(gutils).world @@ -1192,29 +1210,31 @@ function julia_error( println(io) end pscope = parent_scope(val)::LLVM.Function - mi, rt = enzyme_custom_extract_mi(pscope, false) #=error=# + ctx = enzyme_context(API.EnzymeTypeAnalyzerRef(data)) + mi, rt = enzyme_custom_extract_mi(ctx, pscope, false) #=error=# if mi !== nothing println(io, "within ", mi) end end mi = nothing - world = nothing + ctx = enzyme_context(data) + world = ctx.world if isa(val, LLVM.Instruction) f = LLVM.parent(LLVM.parent(val))::LLVM.Function mi, rt = enzyme_custom_extract_mi( + ctx, f, false, ) #=error=# - world = enzyme_extract_world(f) elseif isa(val, LLVM.Argument) f = parent_scope(val)::LLVM.Function mi, rt = enzyme_custom_extract_mi( + ctx, f, false, ) #=error=# - world = enzyme_extract_world(f) end if mi !== nothing emit_error(B, nothing, (msg2, mi, world), EnzymeNoTypeError{Core.MethodInstance, UInt}) @@ -1359,7 +1379,7 @@ function julia_error( return make_batched(ncur, prevbb) end - legal2, obj = absint(cur) + legal2, obj = absint(cur, false, false, false, enzyme_context(gutils)) obj0 = obj # Only do so for the immediate operand/etc to a phi, since otherwise we will make multiple if legal2 @@ -1414,7 +1434,7 @@ else if isa(cur, LLVM.LoadInst) larg, off = get_base_and_offset(operands(cur)[1]) if isa(larg, LLVM.LoadInst) - legal2, obj = absint(larg) + legal2, obj = absint(larg, false, false, false, enzyme_context(gutils)) obj = unbind(obj) if legal2 && is_memory_instance(obj) return make_batched(ncur, prevbb) @@ -1673,6 +1693,7 @@ end end if LLVM.API.LLVMIsAReturnInst(val) != C_NULL mi, rt = enzyme_custom_extract_mi( + enzyme_context(gutils), LLVM.parent(LLVM.parent(val))::LLVM.Function, false, ) #=error=# @@ -1705,27 +1726,23 @@ end end mi = nothing - world = nothing + world = enzyme_context(gutils).world if isa(val, LLVM.Instruction) f = LLVM.parent(LLVM.parent(val))::LLVM.Function mi, rt = enzyme_custom_extract_mi( + enzyme_context(gutils), f, false, ) #=error=# - world = enzyme_context(gutils).world elseif isa(val, LLVM.Argument) f = parent_scope(val)::LLVM.Function mi, rt = enzyme_custom_extract_mi( + enzyme_context(gutils), f, false, ) #=error=# - world = enzyme_context(gutils).world end - # TODO(vchuravy) - # what is data? - # Can we get world here? - @assert world == enzyme_context(gutils).world mode = Enzyme.API.DEM_ReverseModeCombined if mi !== nothing diff --git a/src/gradientutils.jl b/src/gradientutils.jl index 8fecf55d2c..c917f241e6 100644 --- a/src/gradientutils.jl +++ b/src/gradientutils.jl @@ -378,3 +378,9 @@ function enzyme_gutils_context(gutils::API.EnzymeGradientUtilsRef) @assert ptr != C_NULL return unsafe_pointer_to_objref(ptr)::EnzymeContext end + +function enzyme_context(analyzer::API.EnzymeTypeAnalyzerRef) + ptr = API.EnzymeTypeAnalyzerGetExternalContext(analyzer) + @assert ptr != C_NULL + return unsafe_pointer_to_objref(ptr)::EnzymeContext +end diff --git a/src/jlrt.jl b/src/jlrt.jl index f9945fefe7..44b112229d 100644 --- a/src/jlrt.jl +++ b/src/jlrt.jl @@ -447,7 +447,7 @@ function byref_from_val_if_mixed(B::LLVM.IRBuilder, gutils::GradientUtils, @nosp end end -function emit_apply_type!(B::LLVM.IRBuilder, @nospecialize(Ty::Type), args::Vector{LLVM.Value})::LLVM.Value +function emit_apply_type!(B::LLVM.IRBuilder, @nospecialize(Ty::Type), args::Vector{LLVM.Value}, enzyme_context::Union{EnzymeContext, Nothing} = nothing)::LLVM.Value curent_bb = position(B) fn = LLVM.parent(curent_bb) mod = LLVM.parent(fn) @@ -455,7 +455,7 @@ function emit_apply_type!(B::LLVM.IRBuilder, @nospecialize(Ty::Type), args::Vect legal = true found = Any[] for arg in args - slegal, foundv = absint(arg) + slegal, foundv = absint(arg, false, false, false, enzyme_context) if slegal push!(found, unbind(foundv)) else @@ -502,7 +502,7 @@ function emit_apply_type!(B::LLVM.IRBuilder, @nospecialize(Ty::Type), args::Vect return tag end -function emit_tuple!(B::LLVM.IRBuilder, args::Vector{LLVM.Value})::LLVM.Value +function emit_tuple!(B::LLVM.IRBuilder, args::Vector{LLVM.Value}, enzyme_context::Union{EnzymeContext, Nothing} = nothing)::LLVM.Value curent_bb = position(B) fn = LLVM.parent(curent_bb) mod = LLVM.parent(fn) @@ -510,7 +510,7 @@ function emit_tuple!(B::LLVM.IRBuilder, args::Vector{LLVM.Value})::LLVM.Value legal = true found = Any[] for arg in args - slegal, foundv = absint(arg) + slegal, foundv = absint(arg, false, false, false, enzyme_context) if slegal push!(found, unbind(foundv)) else diff --git a/src/llvm/transforms.jl b/src/llvm/transforms.jl index 808e70b867..811b149890 100644 --- a/src/llvm/transforms.jl +++ b/src/llvm/transforms.jl @@ -253,11 +253,11 @@ function rewrite_ccalls!(mod::LLVM.Module) end end -function fixup_1p12_sret!(f::LLVM.Function) +function fixup_1p12_sret!(enzyme_context::EnzymeContext, f::LLVM.Function) if VERSION < v"1.12" return end - mi, RT = enzyme_custom_extract_mi(f, false) + mi, RT = enzyme_custom_extract_mi(enzyme_context, f, false) if mi === nothing return end diff --git a/src/rules/activityrules.jl b/src/rules/activityrules.jl index aa14728b77..7d56d2e113 100644 --- a/src/rules/activityrules.jl +++ b/src/rules/activityrules.jl @@ -1,9 +1,9 @@ -function julia_activity_rule(f::LLVM.Function, world, method_table) +function julia_activity_rule(enzyme_context::EnzymeContext, f::LLVM.Function, world, method_table) if startswith(LLVM.name(f), "japi3") || startswith(LLVM.name(f), "japi1") return end - mi, RT = enzyme_custom_extract_mi(f) + mi, RT = enzyme_custom_extract_mi(enzyme_context, f) llRT, sret, returnRoots = get_return_info(RT) retRemoved, parmsRemoved = removed_ret_parms(f) diff --git a/src/rules/allocrules.jl b/src/rules/allocrules.jl index c2220f99ef..74a8fd7f5f 100644 --- a/src/rules/allocrules.jl +++ b/src/rules/allocrules.jl @@ -11,7 +11,7 @@ function array_shadow_handler( ctx = LLVM.context(LLVM.Value(OrigCI)) gutils = GradientUtils(gutils) - legal, typ, byref = abs_typeof(inst) + legal, typ, byref = abs_typeof(inst, false, Set{LLVM.PHIInst}(), enzyme_context(gutils)) if !legal throw( AssertionError( diff --git a/src/rules/customrules.jl b/src/rules/customrules.jl index 69823a729d..f8b18c081b 100644 --- a/src/rules/customrules.jl +++ b/src/rules/customrules.jl @@ -938,7 +938,7 @@ function enzyme_custom_setup_ret( if active_reg(RealRt, world) == MixedState && B !== nothing bt = GPUCompiler.backtrace(orig) msg2 = sprint(Base.Fix2(Base.show_backtrace, bt)) - mi, _ = enzyme_custom_extract_mi(orig) + mi, _ = enzyme_custom_extract_mi(enzyme_context(gutils), orig) emit_error( B, orig, @@ -994,7 +994,7 @@ end fmi, (args, TT, fwd_RT, kwtup, RT, needsPrimal, RealRt, origNeedsPrimal, activity, C) = fwd_mi(orig, gutils, B) if kwtup !== nothing && kwtup <: Duplicated - mi, _ = enzyme_custom_extract_mi(orig) + mi, _ = enzyme_custom_extract_mi(enzyme_context(gutils), orig) bt = GPUCompiler.backtrace(orig) msg2 = sprint(Base.Fix2(Base.show_backtrace, bt)) @@ -1017,7 +1017,7 @@ end if swiftself pushfirst!(reinsert_gcmarker!(fn, B)) end - _, sret, returnRoots0 = get_return_info(enzyme_custom_extract_mi(llvmf)[2]) + _, sret, returnRoots0 = get_return_info(enzyme_custom_extract_mi(enzyme_context(gutils), llvmf)[2]) returnRoots = returnRoots0 if sret !== nothing sret_lty = convert(LLVMType, eltype(sret)) @@ -1281,7 +1281,7 @@ end width = get_width(gutils) # 1) extract out the MI from attributes - mi, RealRt = enzyme_custom_extract_mi(orig) + mi, RealRt = enzyme_custom_extract_mi(enzyme_context(gutils), orig) isKWCall = isKWCallSignature(mi.specTypes) # 2) Create activity, and annotate function spec @@ -1370,7 +1370,7 @@ end @nospecialize(B::Union{Nothing, LLVM.IRBuilder}) = nothing, ) # 1) extract out the MI from attributes - mi, RealRt = enzyme_custom_extract_mi(orig) + mi, RealRt = enzyme_custom_extract_mi(enzyme_context(gutils), orig) kwfunc = nothing @@ -1436,7 +1436,7 @@ end @inline function has_easy_rule_from_call(orig::LLVM.CallInst, gutils::GradientUtils)::Bool world = enzyme_context(gutils).world - mi, RealRt = enzyme_custom_extract_mi(orig) + mi, RealRt = enzyme_custom_extract_mi(enzyme_context(gutils), orig) specTypes = Interpreter.simplify_kw(mi.specTypes) return EnzymeRules.has_easy_rule_from_sig(specTypes; world) end @@ -1542,7 +1542,7 @@ function enzyme_custom_common_rev( # TODO: don't inject the code multiple times for multiple calls # 1) extract out the MI from attributes - mi, RealRt = enzyme_custom_extract_mi(orig) + mi, RealRt = enzyme_custom_extract_mi(enzyme_context(gutils), orig) isKWCall = isKWCallSignature(mi.specTypes) # 2) Create activity, and annotate function spec @@ -1592,7 +1592,7 @@ function enzyme_custom_common_rev( ) aug_RT = return_type(interp, ami) if kwtup !== nothing && kwtup <: Duplicated - mi, _ = enzyme_custom_extract_mi(orig) + mi, _ = enzyme_custom_extract_mi(enzyme_context(gutils), orig) bt = GPUCompiler.backtrace(orig) msg2 = sprint(Base.Fix2(Base.show_backtrace, bt)) emit_error(B, orig, (msg2, mi, world), NonConstantKeywordArgException) @@ -1743,7 +1743,7 @@ function enzyme_custom_common_rev( orig_swiftself = has_swiftself(LLVM.called_operand(orig)) swiftself = has_swiftself(llvmf) - miRT = enzyme_custom_extract_mi(llvmf)[2] + miRT = enzyme_custom_extract_mi(enzyme_context(gutils), llvmf)[2] _, sret, returnRoots0 = get_return_info(miRT) returnRoots = returnRoots0 sret_union = is_sret_union(miRT) diff --git a/src/rules/jitrules.jl b/src/rules/jitrules.jl index e9904da18d..6d562e9dba 100644 --- a/src/rules/jitrules.jl +++ b/src/rules/jitrules.jl @@ -1960,8 +1960,8 @@ function generic_setup( if runtime_activity pushfirst!(vals, unsafe_to_llvm(B, Val(get_runtime_activity(gutils)))) end - etup0 = emit_tuple!(B, ActivityList) - etup = emit_apply_type!(B, Base.Val, LLVM.Value[etup0]) + etup0 = emit_tuple!(B, ActivityList, enzyme_context(gutils)) + etup = emit_apply_type!(B, Base.Val, LLVM.Value[etup0], enzyme_context(gutils)) if isa(etup, LLVM.Instruction) @assert length(collect(LLVM.uses(etup0))) == 1 end @@ -2405,9 +2405,9 @@ function common_apply_iterate_fwd(offset, B, orig, gutils, normalR, shadowR) return true end - v, isiter = absint(operands(orig)[offset+1]) + v, isiter = absint(operands(orig)[offset+1], false, false, false, enzyme_context(gutils)) isiter = unbind(isiter) - v2, istup = absint(operands(orig)[offset+2]) + v2, istup = absint(operands(orig)[offset+2], false, false, false, enzyme_context(gutils)) istup = unbind(istup) width = get_width(gutils) @@ -2558,8 +2558,8 @@ function common_apply_iterate_augfwd(offset, B, orig, gutils, normalR, shadowR, return true end - v, isiter = absint(operands(orig)[offset+1]) - v2, istup = absint(operands(orig)[offset+2]) + v, isiter = absint(operands(orig)[offset+1], false, false, false, enzyme_context(gutils)) + v2, istup = absint(operands(orig)[offset+2], false, false, false, enzyme_context(gutils)) isiter = unbind(isiter) istup = unbind(istup) diff --git a/src/rules/llvmrules.jl b/src/rules/llvmrules.jl index ec9e5d0050..2e0b547b3c 100644 --- a/src/rules/llvmrules.jl +++ b/src/rules/llvmrules.jl @@ -167,7 +167,7 @@ include("parallelrules.jl") end end - mi, _ = enzyme_custom_extract_mi(LLVM.parent(LLVM.parent(orig)), false) #=error=# + mi, _ = enzyme_custom_extract_mi(enzyme_context(gutils), LLVM.parent(LLVM.parent(orig)), false) #=error=# world = enzyme_context(gutils).world if mi !== nothing @@ -261,7 +261,7 @@ end end end - mi, _ = enzyme_custom_extract_mi(LLVM.parent(LLVM.parent(orig)), false) #=error=# + mi, _ = enzyme_custom_extract_mi(enzyme_context(gutils), LLVM.parent(LLVM.parent(orig)), false) #=error=# world = enzyme_context(gutils).world if mi !== nothing @@ -362,7 +362,7 @@ end end end - mi, _ = enzyme_custom_extract_mi(LLVM.parent(LLVM.parent(orig)), false) #=error=# + mi, _ = enzyme_custom_extract_mi(enzyme_context(gutils), LLVM.parent(LLVM.parent(orig)), false) #=error=# world = enzyme_context(gutils).world if mi !== nothing @@ -504,7 +504,7 @@ end shadowres = UndefValue(LLVM.LLVMType(API.EnzymeGetShadowType(width, value_type(orig)))) - found, arty, byref = abs_typeof(origops[1]) + found, arty, byref = abs_typeof(origops[1], false, Set{LLVM.PHIInst}(), enzyme_context(gutils)) needs_runtime_zero = !found && !(eltype(arty) <: Base.IEEEFloat) @@ -797,7 +797,7 @@ end len = new_from_original(gutils, origops[3]) - found, arty, byref = abs_typeof(origops[1]) + found, arty, byref = abs_typeof(origops[1], false, Set{LLVM.PHIInst}(), enzyme_context(gutils)) needs_runtime_zero = !found && !(eltype(arty) <: Base.IEEEFloat) @@ -1210,21 +1210,21 @@ end " " * string(orig) * " result: " * - string(absint(orig)) * + string(absint(orig, false, false, false, enzyme_context(gutils))) * " " * - string(abs_typeof(orig, true)) * + string(abs_typeof(orig, true, Set{LLVM.PHIInst}(), enzyme_context(gutils))) * " dict: " * - string(absint(origh)) * + string(absint(origh, false, false, false, enzyme_context(gutils))) * " " * - string(abs_typeof(origh, true)) * + string(abs_typeof(origh, true, Set{LLVM.PHIInst}(), enzyme_context(gutils))) * " key " * - string(absint(origkey)) * + string(absint(origkey, false, false, false, enzyme_context(gutils))) * " " * - string(abs_typeof(origkey, true)) * + string(abs_typeof(origkey, true, Set{LLVM.PHIInst}(), enzyme_context(gutils))) * " dflt " * - string(absint(origdflt)) * + string(absint(origdflt, false, false, false, enzyme_context(gutils))) * " " * - string(abs_typeof(origdflt, true)), + string(abs_typeof(origdflt, true, Set{LLVM.PHIInst}(), enzyme_context(gutils))), ) end @@ -1667,7 +1667,7 @@ end args = LLVM.Value[anti, offset] - found, arty, byref = abs_typeof(origops[1]) + found, arty, byref = abs_typeof(origops[1], false, Set{LLVM.PHIInst}(), enzyme_context(gutils)) anti = shadowin elSize = if found LLVM.ConstantInt(Csize_t(actual_size(eltype(arty)))) @@ -1782,7 +1782,7 @@ end end width = get_width(gutils) - legal, dest_ty, _ = abs_typeof(first(operands(orig))) + legal, dest_ty, _ = abs_typeof(first(operands(orig)), false, Set{LLVM.PHIInst}(), enzyme_context(gutils)) if !legal emit_error(B, orig, "Enzyme: could not deduce element type of value within generic_memory_copyto of " * string(first(operands(orig))) * " within " * string(orig)) diff --git a/src/rules/parallelrules.jl b/src/rules/parallelrules.jl index 0c172eb8d9..39dd9def87 100644 --- a/src/rules/parallelrules.jl +++ b/src/rules/parallelrules.jl @@ -190,7 +190,7 @@ end mod = LLVM.parent(LLVM.parent(LLVM.parent(orig))) llvmfn = LLVM.called_operand(orig) - mi = nothing + mi, _ = enzyme_custom_extract_mi(enzyme_context(gutils), llvmfn) fwdmodenm = nothing augfwdnm = nothing adjointnm = nothing @@ -198,10 +198,6 @@ end attributes = function_attributes(llvmfn) for fattr in collect(attributes) if isa(fattr, LLVM.StringAttribute) - if kind(fattr) == "enzymejl_mi" - ptr = reinterpret(Ptr{Cvoid}, parse(UInt, LLVM.value(fattr))) - mi = Base.unsafe_pointer_to_objref(ptr) - end if kind(fattr) == "enzymejl_tapetype" ptr = reinterpret(Ptr{Cvoid}, parse(UInt, LLVM.value(fattr))) TapeType = Base.unsafe_pointer_to_objref(ptr) diff --git a/src/rules/typeunstablerules.jl b/src/rules/typeunstablerules.jl index cd65786c57..f003792bc2 100644 --- a/src/rules/typeunstablerules.jl +++ b/src/rules/typeunstablerules.jl @@ -1450,7 +1450,7 @@ function common_jl_getfield_augfwd(offset, B, orig, gutils, normalR, shadowR, ta push!(vals, inps[1]) sym = new_from_original(gutils, ops[3]) - sym = emit_apply_type!(B, Base.Val, LLVM.Value[sym]) + sym = emit_apply_type!(B, Base.Val, LLVM.Value[sym], enzyme_context(gutils)) push!(vals, sym) push!(vals, unsafe_to_llvm(B, Val(is_constant_value(gutils, ops[2])))) @@ -1547,7 +1547,7 @@ function common_jl_getfield_rev(offset, B, orig, gutils, tape) sym = new_from_original(gutils, ops[3]) sym = lookup_value(gutils, sym, B) - sym = emit_apply_type!(B, Base.Val, LLVM.Value[sym]) + sym = emit_apply_type!(B, Base.Val, LLVM.Value[sym], enzyme_context(gutils)) push!(vals, sym) push!(vals, unsafe_to_llvm(B, Val(is_constant_value(gutils, ops[2])))) @@ -1645,7 +1645,7 @@ end sym = new_from_original(gutils, operands(orig)[2]) sym = (sizeof(Int) == sizeof(Int64) ? emit_box_int64! : emit_box_int32!)(B, sym) - sym = emit_apply_type!(B, Base.Val, LLVM.Value[sym]) + sym = emit_apply_type!(B, Base.Val, LLVM.Value[sym], enzyme_context(gutils)) push!(vals, sym) # TODO properly handle runtime activity here @@ -1744,7 +1744,7 @@ end sym = new_from_original(gutils, operands(orig)[2]) sym = lookup_value(gutils, sym, B) sym = (sizeof(Int) == sizeof(Int64) ? emit_box_int64! : emit_box_int32!)(B, sym) - sym = emit_apply_type!(B, Base.Val, LLVM.Value[sym]) + sym = emit_apply_type!(B, Base.Val, LLVM.Value[sym], enzyme_context(gutils)) push!(vals, sym) push!(vals, unsafe_to_llvm(B, Val(is_constant_value(gutils, operands(orig)[1]))))