diff --git a/Compiler/src/typeinfer.jl b/Compiler/src/typeinfer.jl index fa5102e8b7b04..cb1090adbc201 100644 --- a/Compiler/src/typeinfer.jl +++ b/Compiler/src/typeinfer.jl @@ -1723,6 +1723,12 @@ function add_codeinsts_to_jit!(interp::AbstractInterpreter, ci, source_mode::UIn markinspected!(workqueue, callee) continue end + let cached = ccall(:jl_get_ci_equiv, Any, (Any, UInt), callee, get_inference_world(workqueue.interp))::CodeInstance + if cached !== callee + markinspected!(workqueue, callee) + continue + end + end src = ci_get_source(interp, callee) if !isa(src, CodeInfo) newcallee = typeinf_ext(workqueue.interp, callee.def, source_mode) # always SOURCE_MODE_ABI @@ -1740,14 +1746,15 @@ function add_codeinsts_to_jit!(interp::AbstractInterpreter, ci, source_mode::UIn sptypes = sptypes_from_meth_instance(mi) collectinvokes!(workqueue, src, sptypes) if iszero(ccall(:jl_mi_cache_has_ci, Cint, (Any, Any), mi, callee)) - cached = ccall(:jl_get_ci_equiv, Any, (Any, UInt), callee, get_inference_world(workqueue.interp))::CodeInstance - if cached === callee - # make sure callee is gc-rooted and cached, as required by jl_add_codeinsts_to_jit - code_cache(workqueue.interp)[mi] = callee - else - # use an existing CI from the cache, if there is available one that is compatible - callee === ci && (ci = cached) - callee = cached + let cached = ccall(:jl_get_ci_equiv, Any, (Any, UInt), callee, 0x0)::CodeInstance + if cached === callee + # make sure callee is gc-rooted and cached, as required by jl_add_codeinsts_to_jit + code_cache(workqueue.interp)[mi] = callee + else + # use an existing CI from the cache, if there is available one that is compatible + callee === ci && (ci = cached) + callee = cached + end end end push!(codeinsts, callee) @@ -1836,12 +1843,13 @@ function compile!(codeinfos::Vector{Any}, workqueue::CompilationQueue; enqueue_unprepared_invokes) # try to reuse an existing CodeInstance from before to avoid making duplicates in the cache if iszero(ccall(:jl_mi_cache_has_ci, Cint, (Any, Any), mi, callee)) - cached = ccall(:jl_get_ci_equiv, Any, (Any, UInt), callee, world)::CodeInstance - if cached === callee - code_cache(interp)[mi] = callee - else - # Use an existing CI from the cache, if there is available one that is compatible - callee = cached + let cached = ccall(:jl_get_ci_equiv, Any, (Any, UInt), callee, 0x0)::CodeInstance + if cached === callee + code_cache(interp)[mi] = callee + else + # Use an existing CI from the cache, if there is available one that is compatible + callee = cached + end end end push!(codeinfos, callee) @@ -1952,7 +1960,7 @@ function typeinf_ext_toplevel(methods::Vector{Any}, worlds::Vector{UInt}, trim_m ci = ci::CodeInstance mi = get_ci_mi(ci) return !iszero(ccall(:jl_mi_cache_has_ci, Cint, (Any, Any), mi, ci)) || - ccall(:jl_get_ci_equiv, Any, (Any, UInt), ci, ci.min_world)::CodeInstance !== ci + ccall(:jl_get_ci_equiv, Any, (Any, UInt), ci, 0x0)::CodeInstance !== ci end return Core.svec(codeinfos, cis) diff --git a/src/aotcompile.cpp b/src/aotcompile.cpp index 5de83829aaf1a..5d594809fb02d 100644 --- a/src/aotcompile.cpp +++ b/src/aotcompile.cpp @@ -816,22 +816,10 @@ static Function *emit_pkg_plt_thunk(jl_codegen_output_t &out, jl_code_instance_t static jl_compiled_functions_t::iterator get_ci_equiv_compiled(jl_code_instance_t *ci JL_PROPAGATES_ROOT, jl_compiled_functions_t &compiled_functions) JL_NOTSAFEPOINT { - jl_value_t *def = ci->def; - jl_value_t *owner = ci->owner; - jl_value_t *rettype = ci->rettype; - size_t min_world = jl_atomic_load_relaxed(&ci->min_world); - size_t max_world = jl_atomic_load_relaxed(&ci->max_world); for (auto it = compiled_functions.begin(), E = compiled_functions.end(); it != E; ++it) { auto codeinst = it->first; - if (codeinst != ci && - jl_atomic_load_relaxed(&codeinst->inferred) != NULL && - jl_atomic_load_relaxed(&codeinst->min_world) <= min_world && - jl_atomic_load_relaxed(&codeinst->max_world) >= max_world && - jl_egal(codeinst->def, def) && - jl_egal(codeinst->owner, owner) && - jl_egal(codeinst->rettype, rettype)) { + if (codeinst != ci && jl_is_ci_equiv(ci, codeinst, 0)) return it; - } } return compiled_functions.end(); } @@ -847,8 +835,11 @@ static void aot_link_output(jl_codegen_output_t &out) JL_CANSAFEPOINT continue; auto it = out.ci_funcs.find(ci); - if (it == out.ci_funcs.end()) - it = get_ci_equiv_compiled(ci, out.ci_funcs); + if (it == out.ci_funcs.end()) { + auto equiv = get_ci_equiv_compiled(ci, out.ci_funcs); + if (equiv != out.ci_funcs.end()) + it = equiv; + } jl_codeinst_funcs_t funcs; if (it != out.ci_funcs.end()) { funcs = {it->second.invoke_api, it->second.invoke, it->second.specptr}; diff --git a/src/ast.c b/src/ast.c index eb6fc3f21e8fc..5a6867ae33f5d 100644 --- a/src/ast.c +++ b/src/ast.c @@ -1072,7 +1072,7 @@ static jl_value_t *jl_invoke_julia_macro(jl_array_t *args, jl_module_t *inmodule jl_timing_show_macro(mfunc, retry_lno != NULL ? retry_lno : margs[1], inmodule, JL_TIMING_DEFAULT_BLOCK); *ctx = mfunc->def.method->module; - result = jl_invoke(margs[0], &margs[1], nargs - 1, mfunc); + result = jl_invoke_fromdispatch(margs[0], &margs[1], nargs - 1, mfunc); } JL_CATCH { if ((jl_loaderror_type == NULL) || !throw_load_error) { diff --git a/src/gf.c b/src/gf.c index 2085beabc68a1..4dd947d5543ea 100644 --- a/src/gf.c +++ b/src/gf.c @@ -586,8 +586,8 @@ JL_DLLEXPORT jl_code_instance_t *jl_get_method_uninferred( jl_value_t *owner = jl_nothing; // TODO: owner should be arg jl_code_instance_t *codeinst = jl_atomic_load_relaxed(&mi->cache); for (; codeinst; codeinst = jl_atomic_load_relaxed(&codeinst->next)) { - if (jl_atomic_load_relaxed(&codeinst->min_world) == min_world && - jl_atomic_load_relaxed(&codeinst->max_world) == max_world && + if (jl_atomic_load_relaxed(&codeinst->min_world) <= min_world && + jl_atomic_load_relaxed(&codeinst->max_world) >= max_world && jl_egal(codeinst->owner, owner) && jl_egal(codeinst->rettype, rettype)) { if (di == NULL) @@ -599,7 +599,7 @@ JL_DLLEXPORT jl_code_instance_t *jl_get_method_uninferred( if (!(debuginfo && jl_egal((jl_value_t*)debuginfo, (jl_value_t*)di))) continue; } - // TODO: this is implied by the matching worlds, since it is intrinsic, so do we really need to verify it? + // n.b.: this is implied by the matching worlds for min_world, but not max_world == ~0 (which couldn't compute this) jl_svec_t *e = jl_atomic_load_relaxed(&codeinst->edges); if (e && jl_egal((jl_value_t*)e, (jl_value_t*)edges)) return codeinst; @@ -624,24 +624,49 @@ JL_DLLEXPORT int jl_mi_cache_has_ci(jl_method_instance_t *mi, return 0; } -// look for something with an egal ABI and properties that is already in the JIT for a whole edge (target_world=0) or can be added to the JIT with new source just for target_world. -JL_DLLEXPORT jl_code_instance_t *jl_get_ci_equiv(jl_code_instance_t *ci JL_PROPAGATES_ROOT, size_t target_world) JL_NOTSAFEPOINT +// return whether the ci has more restrictions than the other arguments (more edges and narrower worlds) +static int jl_codeinst_edges_sub(jl_code_instance_t *ci, size_t min_world2, size_t max_world2, jl_svec_t *edges2) JL_NOTSAFEPOINT +{ + size_t min_world = jl_atomic_load_relaxed(&ci->min_world); + size_t max_world = jl_atomic_load_relaxed(&ci->max_world); + jl_svec_t *edges = jl_atomic_load_relaxed(&ci->edges); + // n.b.: edges matching is implied sufficiently by the matching worlds for min_world, but not max_world == ~0 (which couldn't compute that) + return min_world >= min_world2 && max_world <= max_world2 && edges; // TODO: && jl_egal((jl_value_t*)edges, (jl_value_t*)edges2); +} + +// return whether the codeinst can be substituted in place of ci for an invoke target in target_world +JL_DLLEXPORT int jl_is_ci_equiv(jl_code_instance_t *ci JL_PROPAGATES_ROOT, jl_code_instance_t *codeinst, size_t target_world) JL_NOTSAFEPOINT { jl_value_t *def = ci->def; - jl_method_instance_t *mi = jl_get_ci_mi(ci); jl_value_t *owner = ci->owner; jl_value_t *rettype = ci->rettype; + if (jl_egal(codeinst->def, def) && + jl_egal(codeinst->owner, owner) && + jl_egal(codeinst->rettype, rettype)) { + if (!target_world || jl_atomic_load_relaxed(&codeinst->invoke) != NULL) { + size_t min_world = jl_atomic_load_relaxed(&ci->min_world); + size_t max_world = jl_atomic_load_relaxed(&ci->max_world); + size_t min_world2 = jl_atomic_load_relaxed(&codeinst->min_world); + size_t max_world2 = jl_atomic_load_relaxed(&codeinst->max_world); + if (target_world || (min_world2 == min_world && max_world2 == max_world)) { + jl_svec_t *edges2 = jl_atomic_load_relaxed(&codeinst->edges); + if (jl_codeinst_edges_sub(ci, min_world2, max_world2, edges2)) { + return 1; + } + } + } + } + return 0; +} + +// look for something with an egal ABI and properties that is already in the JIT for the target_world, or could be added to the JIT instead of ci to satisfy the same invoke edge with the same src. +JL_DLLEXPORT jl_code_instance_t *jl_get_ci_equiv(jl_code_instance_t *ci JL_PROPAGATES_ROOT, size_t target_world) JL_NOTSAFEPOINT +{ + jl_method_instance_t *mi = jl_get_ci_mi(ci); jl_code_instance_t *codeinst = jl_atomic_load_relaxed(&mi->cache); while (codeinst) { - if (codeinst != ci && - jl_atomic_load_relaxed(&codeinst->inferred) != NULL && - jl_atomic_load_relaxed(&codeinst->min_world) <= target_world && - jl_atomic_load_relaxed(&codeinst->max_world) >= target_world && - jl_egal(codeinst->def, def) && - jl_egal(codeinst->owner, owner) && - jl_egal(codeinst->rettype, rettype)) { + if (codeinst != ci && jl_is_ci_equiv(ci, codeinst, target_world)) return codeinst; - } codeinst = jl_atomic_load_relaxed(&codeinst->next); } return ci; @@ -3788,14 +3813,29 @@ static int need_copy_to_mi_cache(jl_method_instance_t *mi, jl_method_instance_t static jl_code_instance_t *copy_to_mi_cache(jl_method_instance_t *mi JL_PROPAGATES_ROOT, jl_code_instance_t *codeinst2) JL_CANSAFEPOINT { + size_t current_world = jl_get_world_counter(); + size_t max_world2 = jl_atomic_load_relaxed(&codeinst2->max_world); + // if codeinst2 is still valid beyond current_world, link codeinst to + // it so that invalidation of codeinst2 also invalidates codeinst + jl_method_t *m = mi->def.method; + jl_svec_t *copy_edge = jl_is_method(m) ? jl_svec2(m->sig, codeinst2) : jl_emptysvec; + JL_GC_PUSH1(©_edge); jl_code_instance_t *codeinst = jl_get_method_uninferred( mi, codeinst2->rettype, jl_atomic_load_relaxed(&codeinst2->min_world), - jl_atomic_load_relaxed(&codeinst2->max_world), // TODO: use min(max_world, current_world) here + max_world2 < current_world ? max_world2 : current_world, jl_atomic_load_relaxed(&codeinst2->debuginfo), - jl_atomic_load_relaxed(&codeinst2->edges)); + copy_edge); + JL_GC_POP(); if (jl_atomic_load_relaxed(&codeinst->invoke) == NULL) { - // TODO: add edges and jl_promote_ci_to_current here + if (max_world2 == ~(size_t)0) { + JL_LOCK(&world_counter_lock); + if (jl_atomic_load_relaxed(&codeinst2->max_world) == ~(size_t)0) { + jl_method_instance_add_backedge(mi, NULL, codeinst); + jl_atomic_store_relaxed(&codeinst->max_world, ~(size_t)0); // jl_promote_ci_to_current + } + JL_UNLOCK(&world_counter_lock); + } jl_gc_write(codeinst, codeinst->rettype_const, jl_value_t, codeinst2->rettype_const); uint8_t specsigflags; jl_callptr_t invoke; @@ -4487,10 +4527,16 @@ JL_DLLEXPORT jl_value_t *jl_invoke(jl_value_t *F, jl_value_t **args, uint32_t na return _jl_invoke(F, args, nargs, mfunc, world, TRIGGER_FOREIGN); } +jl_value_t *jl_invoke_fromdispatch(jl_value_t *F, jl_value_t **args, uint32_t nargs, jl_method_instance_t *mfunc) +{ + size_t world = jl_current_task->world_age; + return _jl_invoke(F, args, nargs, mfunc, world, TRIGGER_DISPATCH); +} + // Used by jl_eval_thunk to invoke top-level thunks. They will be // garbage-collectable as soon as they are invoked, so their ORC symbols must be // unregistered before we enter invoke, which may never return. -JL_DLLEXPORT jl_value_t *jl_invoke_oneshot(jl_value_t *F, jl_value_t **args, uint32_t nargs, jl_method_instance_t *mfunc) +jl_value_t *jl_invoke_oneshot(jl_value_t *F, jl_value_t **args, uint32_t nargs, jl_method_instance_t *mfunc) { size_t world = jl_current_task->world_age; diff --git a/src/jitlayers.cpp b/src/jitlayers.cpp index 1ad2a788d4472..c9c453f9feff5 100644 --- a/src/jitlayers.cpp +++ b/src/jitlayers.cpp @@ -2324,6 +2324,23 @@ bool JuliaOJIT::linkOutput(orc::MaterializationResponsibility &MR, MemoryBufferR RenameDef(Funcs.specptr, S.specptr); } + // Pre-pass: find CI equivalents, and build EquivMap for use in the main + // pass to memoize findCompatibleCI. + DenseMap EquivMap; + for (auto &[Call, T] : Info->call_targets) { + auto [CI, API] = Call; + JL_GC_PROMISE_ROOTED(CI); + if (!Syms.contains(T)) + continue; + if (EquivMap.contains(CI)) + continue; + if (!jl_mi_cache_has_ci(jl_get_ci_mi(CI), CI)) { + jl_code_instance_t *Equiv = findCompatibleCI(CI); + if (Equiv != CI) + EquivMap[CI] = Equiv; + } + } + // Rename referenced CIs in the workqueue. for (auto &[Call, T] : Info->call_targets) { auto [CI, API] = Call; @@ -2339,7 +2356,7 @@ bool JuliaOJIT::linkOutput(orc::MaterializationResponsibility &MR, MemoryBufferR continue; } JL_GC_PROMISE_ROOTED(CI); - auto Dest = linkCallTarget(MR, CI, API); + auto Dest = linkCallTarget(MR, CI, API, EquivMap); if (!Dest) return false; if (auto *DestSym = findLinkGraphSymbolByName(G, Dest); @@ -2398,16 +2415,19 @@ bool JuliaOJIT::linkOutput(orc::MaterializationResponsibility &MR, MemoryBufferR // Must hold LinkerMutex. orc::SymbolStringPtr JuliaOJIT::linkCallTarget(orc::MaterializationResponsibility &MR, - jl_code_instance_t *CI, jl_invoke_api_t API) + jl_code_instance_t *CI, jl_invoke_api_t API, + const DenseMap &EquivMap) { - // This condition should match that in jl_add_codeinst_to_jit!, which will - // add a different, compatible CodeInstance to the JIT but not update the - // invoke statement. - if (!jl_mi_cache_has_ci(jl_get_ci_mi(CI), CI)) - CI = findCompatibleCI(CI); - auto It = CISymbols.find(CI); - if (It != CISymbols.end() && It->second.invoke_api == API) - return It->second.specptr; + { + auto It = EquivMap.find(CI); + if (It != EquivMap.end()) + CI = It->second; + } + { + auto It = CISymbols.find(CI); + if (It != CISymbols.end() && It->second.invoke_api == API) + return It->second.specptr; + } CISymbolPtr *Sym = linkCISymbol(CI); @@ -2438,30 +2458,19 @@ orc::SymbolStringPtr JuliaOJIT::linkCallTarget(orc::MaterializationResponsibilit return Sym->specptr; } -jl_code_instance_t *JuliaOJIT::findCompatibleCI(jl_code_instance_t *CI) +jl_code_instance_t *JuliaOJIT::findCompatibleCI(jl_code_instance_t *ci) { - // add_codeinsts_to_jit! may have added an equivalent CI to the JIT, but + // add_codeinsts_to_jit! may have added an equivalent ci to the JIT, but // the invoke itself won't be updated. - auto MI = jl_get_ci_mi(CI); - jl_value_t *Def = CI->def; - jl_value_t *Owner = CI->owner; - jl_value_t *RetType = CI->rettype; - size_t MinWorld = jl_atomic_load_relaxed(&CI->min_world); - size_t MaxWorld = jl_atomic_load_relaxed(&CI->max_world); - auto IsCompatible = [=](jl_code_instance_t *CI2) JL_NOTSAFEPOINT { - return jl_atomic_load_relaxed(&CI2->min_world) <= MinWorld && - jl_atomic_load_relaxed(&CI2->max_world) >= MaxWorld && - jl_egal(CI2->def, Def) && jl_egal(CI2->owner, Owner) && - jl_egal(CI2->rettype, RetType); - }; - for (auto CI2 = jl_atomic_load_relaxed(&MI->cache); CI2; - CI2 = jl_atomic_load_relaxed(&CI2->next)) { - if (CI2 != CI && IsCompatible(CI2) && - (CISymbols.contains(CI2) || jl_atomic_load_relaxed(&CI2->invoke))) { - return CI2; + auto mi = jl_get_ci_mi(ci); + for (auto ci2 = jl_atomic_load_relaxed(&mi->cache); ci2; + ci2 = jl_atomic_load_relaxed(&ci2->next)) { + if (ci2 != ci && jl_is_ci_equiv(ci, ci2, 0) && + (CISymbols.contains(ci2) || jl_atomic_load_relaxed(&ci2->invoke))) { + return ci2; } } - return CI; + return ci; } CISymbolPtr *JuliaOJIT::linkCISymbol(jl_code_instance_t *CI) diff --git a/src/jitlayers.h b/src/jitlayers.h index d0b34155f881e..c22304ad144da 100644 --- a/src/jitlayers.h +++ b/src/jitlayers.h @@ -890,7 +890,7 @@ class JuliaOJIT { // false after calling MR.failMaterialization(). bool linkOutput(orc::MaterializationResponsibility &MR, MemoryBufferRef ObjBuf, jitlink::LinkGraph &G, - std::unique_ptr Info) JL_NOTSAFEPOINT; + std::unique_ptr Info) JL_CANSAFEPOINT_ENTER_LEAVE; // Return a symbol that should be linked to the call target. The origin of // this symbol depends on the code instance: @@ -904,7 +904,8 @@ class JuliaOJIT { // new module and return a symbol for it. orc::SymbolStringPtr linkCallTarget(orc::MaterializationResponsibility &MR, jl_code_instance_t *CI, - jl_invoke_api_t API) JL_NOTSAFEPOINT; + jl_invoke_api_t API, + const DenseMap &EquivMap) JL_NOTSAFEPOINT; // If the provided CodeInstance is neither compiled nor has an ORC symbol in // CISymbols, look for a compatible CodeInstance in the MethodInstance's diff --git a/src/julia_internal.h b/src/julia_internal.h index 7c2881a2398f3..aae67349e9713 100644 --- a/src/julia_internal.h +++ b/src/julia_internal.h @@ -826,7 +826,8 @@ JL_DLLEXPORT int jl_mi_cache_has_ci(jl_method_instance_t *mi, jl_code_instance_t JL_DLLEXPORT void jl_read_codeinst_invoke(jl_code_instance_t *ci, uint8_t *specsigflags, jl_callptr_t *invoke, void **specptr, int waitcompile) JL_CANSAFEPOINT; JL_DLLEXPORT void jl_add_codeinsts_to_jit(jl_array_t *codeinsts, jl_array_t *srcs) JL_CANSAFEPOINT; -JL_DLLEXPORT jl_value_t *jl_invoke_oneshot(jl_value_t *F, jl_value_t **args, uint32_t nargs, jl_method_instance_t *meth) JL_CANSAFEPOINT; +jl_value_t *jl_invoke_oneshot(jl_value_t *F, jl_value_t **args, uint32_t nargs, jl_method_instance_t *meth) JL_CANSAFEPOINT; +jl_value_t *jl_invoke_fromdispatch(jl_value_t *F, jl_value_t **args, uint32_t nargs, jl_method_instance_t *mfunc) JL_CANSAFEPOINT; JL_DLLEXPORT jl_code_instance_t *jl_new_codeinst_uninit(jl_method_instance_t *mi, jl_value_t *owner) JL_CANSAFEPOINT; JL_DLLEXPORT jl_code_instance_t *jl_new_codeinst( @@ -837,7 +838,7 @@ JL_DLLEXPORT jl_code_instance_t *jl_new_codeinst( uint32_t effects, jl_value_t *analysis_results, jl_debuginfo_t *di, jl_svec_t *edges /* , int absolute_max*/) JL_CANSAFEPOINT; JL_DLLEXPORT jl_code_instance_t *jl_get_ci_equiv(jl_code_instance_t *ci JL_PROPAGATES_ROOT, size_t target_world) JL_NOTSAFEPOINT; - +JL_DLLEXPORT int jl_is_ci_equiv(jl_code_instance_t *ci JL_PROPAGATES_ROOT, jl_code_instance_t *codeinst, size_t target_world) JL_NOTSAFEPOINT; STATIC_INLINE jl_method_instance_t *jl_get_ci_mi(jl_code_instance_t *ci JL_PROPAGATES_ROOT) JL_NOTSAFEPOINT { jl_value_t *def = ci->def;