diff --git a/src/Enzyme.jl b/src/Enzyme.jl index e4f1a92ac5..137ccfc20d 100644 --- a/src/Enzyme.jl +++ b/src/Enzyme.jl @@ -133,13 +133,17 @@ 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( + 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 69c16604a6..ba479ce03e 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, @@ -1359,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 324d46f0c6..a49054615b 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -248,8 +248,13 @@ 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_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, + batch_call_same_with_inverted_arg_if_active! # Julia function to LLVM stem and arity const cmplx_known_ops = @@ -525,12 +530,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) +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) push!(attributes, StringAttribute("enzymejl_world", string(job.world))) @@ -544,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 @@ -584,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 @@ -598,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 @@ -644,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) @@ -705,7 +713,7 @@ end name = meth.name jlmod = meth.module - julia_activity_rule(llvmfn, method_table) + julia_activity_rule(enzyme_context, llvmfn, world, method_table) if has_custom_rule handleCustom( state, @@ -1083,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, @@ -1125,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 @@ -1151,8 +1160,7 @@ function set_module_types!(interp, mod::LLVM.Module, primalf::Union{Nothing, LLV continue end - world = enzyme_extract_world(f) - + world = interp.world jlargs = classify_arguments( mi.specTypes, ftype, @@ -1302,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 @@ -1318,9 +1326,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] @@ -1347,8 +1355,9 @@ 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) + prepare_llvm(enzyme_context, interp, otherMod, job, meta) entry = name(meta.entry) @@ -1371,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) @@ -1380,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()) @@ -1736,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 @@ -1774,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} @@ -1804,8 +1813,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) + world = enzyme_context(gutils).world if !guaranteed_nonactive(Ty, world) B = LLVM.IRBuilder() position!(B, V) @@ -2480,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) @@ -2772,7 +2816,6 @@ function enzyme!( width, returnPrimal, shadow_init, - world, interp, runtimeActivity, ) @@ -2815,7 +2858,6 @@ function enzyme!( width, false, shadow_init, - world, interp, runtimeActivity ) #=returnPrimal=# @@ -2857,7 +2899,6 @@ function enzyme!( width, returnPrimal, shadow_init, - world, interp, runtimeActivity ) @@ -2903,7 +2944,6 @@ function enzyme!( width, returnPrimal, shadow_init, - world, interp, runtimeActivity ) @@ -2987,10 +3027,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 @@ -3290,6 +3330,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) @@ -3490,7 +3531,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′} @@ -4302,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, @@ -4977,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, @@ -5225,7 +5267,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 @@ -5277,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 @@ -5295,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 @@ -5398,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 @@ -5432,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") @@ -5498,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, @@ -5612,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) @@ -5638,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 @@ -5671,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[]) @@ -5708,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 @@ -5751,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 @@ -5798,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 @@ -5826,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 @@ -5862,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 @@ -5880,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 @@ -5900,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) @@ -5920,7 +5963,7 @@ end adjointf, augmented_primalf, TapeType = enzyme!( enzyme_context, job, - interp, + interp, mod, primalf, TT, @@ -6027,7 +6070,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) @@ -6039,6 +6082,7 @@ end llvmf = functions(mod)[llvmf] + push!(function_attributes(llvmf), StringAttribute("implements", fname)) end end 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 cf38342b4b..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,23 +1164,26 @@ 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 throw(IllegalTypeAnalysisException(msg, mi, world, sval, ir, bt)) elseif errtype == API.ET_NoType @assert B != C_NULL @@ -1190,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}) @@ -1287,7 +1309,7 @@ function julia_error( seen = Dict{LLVM.Value,LLVM.Value}() illegal = false created = LLVM.Instruction[] - world = enzyme_extract_world(LLVM.parent(position(IRBuilder(B)))) + world = enzyme_context(gutils).world width = get_width(gutils) function make_batched(@nospecialize(cur::LLVM.Value), B::LLVM.IRBuilder)::LLVM.Value if width == 1 @@ -1357,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 @@ -1412,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) @@ -1671,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=# @@ -1702,23 +1725,23 @@ end end end - mi = nothing - world = nothing + mi = 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_extract_world(f) 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_extract_world(f) end mode = Enzyme.API.DEM_ReverseModeCombined diff --git a/src/gradientutils.jl b/src/gradientutils.jl index 37e883c5bf..c917f241e6 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 @@ -366,3 +366,21 @@ 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 + +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 ad18cc906f..44b112229d 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) @@ -447,7 +447,7 @@ function byref_from_val_if_mixed(B::LLVM.IRBuilder, @nospecialize(val::LLVM.Valu 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 @@ -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/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/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 diff --git a/src/rules/activityrules.jl b/src/rules/activityrules.jl index 9405fe0ff9..7d56d2e113 100644 --- a/src/rules/activityrules.jl +++ b/src/rules/activityrules.jl @@ -1,9 +1,9 @@ -function julia_activity_rule(f::LLVM.Function, 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) @@ -18,7 +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) jlargs = classify_arguments( mi.specTypes, 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 353f75727f..f8b18c081b 100644 --- a/src/rules/customrules.jl +++ b/src/rules/customrules.jl @@ -334,8 +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) + world = enzyme_context(gutils).world jlargs = classify_arguments( mi.specTypes, @@ -861,7 +860,7 @@ function enzyme_custom_setup_ret( width = get_width(gutils) mode = get_mode(gutils) - world = enzyme_extract_world(LLVM.parent(LLVM.parent(orig))) + world = enzyme_context(gutils).world needsShadowP = Ref{UInt8}(0) needsPrimalP = Ref{UInt8}(0) @@ -939,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, @@ -988,14 +987,14 @@ 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 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)) @@ -1010,8 +1009,7 @@ end width = get_width(gutils) - enzyme_ctx = Enzyme.enzyme_context(get_logic(gutils)) - llvmf = nested_codegen!(enzyme_ctx, mode, mod, fmi, world, true) + llvmf = nested_codegen!(enzyme_context(gutils), mode, mod, fmi, true) orig_swiftself = has_swiftself(LLVM.called_operand(orig)) @@ -1019,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)) @@ -1283,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 @@ -1298,8 +1296,7 @@ end needsShadow end - fn = LLVM.parent(LLVM.parent(orig)) - world = enzyme_extract_world(fn) + world = enzyme_context(gutils).world C = EnzymeRules.RevConfig{ Bool(needsPrimal), @@ -1373,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 @@ -1413,8 +1410,7 @@ end end TT = Tuple{tt...} - fn = LLVM.parent(LLVM.parent(orig)) - world = enzyme_extract_world(fn) + world = enzyme_context(gutils).world @safe_debug "Trying to apply custom forward rule" TT isKWCall functy = if isKWCall @@ -1439,9 +1435,8 @@ 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) - mi, RealRt = enzyme_custom_extract_mi(orig) + world = enzyme_context(gutils).world + 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 @@ -1547,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 @@ -1584,7 +1579,8 @@ function enzyme_custom_common_rev( curent_bb = position(B) fn = LLVM.parent(curent_bb) - world = enzyme_extract_world(fn) + enzyme_ctx = enzyme_context(gutils) + world = enzyme_ctx.world mode = get_mode(gutils) @@ -1596,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) @@ -1650,8 +1646,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!(enzyme_ctx, mode, mod, ami, true) @assert llvmf !== nothing rev_RT = nothing final_mi = ami @@ -1694,8 +1689,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!(enzyme_ctx, mode, mod, rmi, true) final_mi = rmi end @@ -1749,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 2b5eef728d..2e0b547b3c 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(enzyme_context(gutils), 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(enzyme_context(gutils), 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(enzyme_context(gutils), LLVM.parent(LLVM.parent(orig)), false) #=error=# + world = enzyme_context(gutils).world if mi !== nothing err = emit_error( @@ -507,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) @@ -800,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) @@ -1213,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 @@ -1670,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)))) @@ -1785,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)) @@ -1795,8 +1792,7 @@ end ET = eltype(dest_ty) - fn = LLVM.parent(LLVM.parent(orig)) - world = enzyme_extract_world(fn) + 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..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) @@ -225,7 +221,7 @@ end e_tt = Tuple{Const{Int}} modifiedBetween = (mode != API.DEM_ForwardMode, false) - world = enzyme_extract_world(LLVM.parent(position(B))) + world = enzyme_context(gutils).world pfuncT = funcT @@ -632,9 +628,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) push!(function_attributes(entry), EnumAttribute("alwaysinline")) pval = functions(mod)[sname] @@ -681,9 +676,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) push!(function_attributes(entry), EnumAttribute("alwaysinline")) pval = functions(mod)[sname] @@ -719,7 +713,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))) if is_constant_value(gutils, orig) && is_constant_inst(gutils, orig) return end @@ -742,8 +735,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] @@ -772,7 +765,7 @@ end width = get_width(gutils) mode = get_mode(gutils) - world = enzyme_extract_world(LLVM.parent(position(B))) + world = enzyme_context(gutils).world vals = LLVM.Value[ unsafe_to_llvm(B, runtime_newtask_fwd), @@ -826,7 +819,7 @@ end uncacheable = get_uncacheable(gutils, orig) ModifiedBetween = (uncacheable[1] != 0,) - world = enzyme_extract_world(LLVM.parent(position(B))) + 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..f003792bc2 100644 --- a/src/rules/typeunstablerules.jl +++ b/src/rules/typeunstablerules.jl @@ -401,7 +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))) + world = enzyme_context(gutils).world @assert is_constant_value(gutils, operands(orig)[offset]) ops = @view arg_operands_view(orig)[offset+1:end] @@ -870,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)))) @@ -881,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 @@ -892,7 +892,7 @@ end unsafe_store!(tapeR, shadowres.ref) else @assert legal - world = enzyme_extract_world(LLVM.parent(position(B))) + world = enzyme_context(gutils).world if !guaranteed_nonactive(TT, world) unsafe_store!(tapeR, shadowres.ref) end @@ -993,7 +993,7 @@ end torun = false if legal @assert legal - world = enzyme_extract_world(LLVM.parent(position(B))) + world = enzyme_context(gutils).world torun = !guaranteed_nonactive(TT, world) else torun = true @@ -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]))))