add helpers for backedge updates on CI equivalence swaps and copies

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)

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<Value *> funcs;
         if (it != out.ci_funcs.end()) {
             funcs = {it->second.invoke_api, it->second.invoke, it->second.specptr};

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) {

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(&copy_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;
 

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<jl_code_instance_t *, jl_code_instance_t *> 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<jl_code_instance_t *, jl_code_instance_t *> &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)

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<jl_linker_info_t> Info) JL_NOTSAFEPOINT;
+                    std::unique_ptr<jl_linker_info_t> 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<jl_code_instance_t *, jl_code_instance_t *> &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