Don't re-infer compiled CodeInstance from external images

Compiler/src/precompile.jl

@@ -435,7 +435,7 @@ function compile_and_emit_native(worlds::Vector{UInt},
     # list for codegen, plus the ordered CodeInstances to place in the method
     # caches of the output image.
     result = try
-        typeinf_ext_toplevel(tocompile, worlds, trim_mode)
+        typeinf_ext_toplevel(tocompile, worlds, trim_mode, external_linkage)
     catch exc
         # Handle trimming failures
         isa(exc, Core.TrimFailure) || rethrow()

Compiler/src/typeinfer.jl

@@ -1602,6 +1602,11 @@ function ci_has_invoke(code::CodeInstance)
     return (@atomic :monotonic code.invoke) !== C_NULL
 end
 
+const CI_FLAGS_FROM_IMAGE = 0b0100
+function ci_from_image(code::CodeInstance)
+    return (@atomic :monotonic code.flags) & CI_FLAGS_FROM_IMAGE != 0
+end
+
 function ci_meets_requirement(interp::AbstractInterpreter, code::CodeInstance, source_mode::UInt8)
     source_mode == SOURCE_MODE_NOT_REQUIRED && return true
     source_mode == SOURCE_MODE_ABI && return ci_has_abi(interp, code)
@@ -1891,14 +1896,20 @@ end
 # collect a list of all code that is needed along with CodeInstance to codegen it fully
 function collectinvokes!(workqueue::CompilationQueue, ci::CodeInfo, sptypes::Vector{VarState};
                          invokelatest_queue::Union{CompilationQueue,Nothing} = nothing,
-                         enqueue_unprepared_invokes::Bool = false)
+                         enqueue_unprepared_invokes::Bool = false,
+                         external_linkage::Bool = false)
     src = ci.code
     for i = 1:length(src)
         stmt = src[i]
         isexpr(stmt, :(=)) && (stmt = stmt.args[2])
         if isexpr(stmt, :invoke) || isexpr(stmt, :invoke_modify)
             edge = stmt.args[1]
+            # If this CodeInstance is already compiled in the image, and we can
+            # link to it, we should do that instead of compiling it again.  With
+            # invoke_modify, we need to compile it regardless.
             if edge isa CodeInstance && has_valid_abi_sparams(get_ci_mi(edge)) &&
+                    (isexpr(stmt, :invoke_modify) ||
+                     !(external_linkage && ci_from_image(edge) && ci_has_invoke(edge))) &&
                     (enqueue_unprepared_invokes ||
                      ci_has_invoke(edge) || ci_has_source(workqueue.interp, edge) ||
                      !iszero(ccall(:jl_mi_cache_has_ci, Cint, (Any, Any), get_ci_mi(edge), edge)))
@@ -2075,6 +2086,7 @@ end
 function compile!(codeinfos::Vector{Any}, workqueue::CompilationQueue;
     invokelatest_queue::Union{CompilationQueue,Nothing} = nothing,
     enqueue_unprepared_invokes::Bool = false,
+    external_linkage::Bool,
 )
     interp = workqueue.interp
     world = get_inference_world(interp)
@@ -2134,7 +2146,7 @@ function compile!(codeinfos::Vector{Any}, workqueue::CompilationQueue;
             if src isa CodeInfo
                 sptypes = sptypes_from_meth_instance(mi)
                 collectinvokes!(workqueue, src, sptypes; invokelatest_queue,
-                                enqueue_unprepared_invokes)
+                                enqueue_unprepared_invokes, external_linkage)
                 # 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 = find_equivalent_cached_ci(
@@ -2160,7 +2172,7 @@ const TRIM_NO = 0x0
 const TRIM_SAFE = 0x1
 const TRIM_UNSAFE = 0x2
 const TRIM_UNSAFE_WARN = 0x3
-function typeinf_ext_toplevel(methods::Vector{Any}, worlds::Vector{UInt}, trim_mode::UInt8)
+function typeinf_ext_toplevel(methods::Vector{Any}, worlds::Vector{UInt}, trim_mode::UInt8, external_linkage::Bool)
     # During `--trim`, infer against an isolated cache namespace. The owner is re-stamped
     # back to `nothing` at serialization time (see `src/staticdata.c`).
     cache_owner = trim_mode == TRIM_NO ? nothing : :trim
@@ -2180,14 +2192,14 @@ function typeinf_ext_toplevel(methods::Vector{Any}, worlds::Vector{UInt}, trim_m
         )
 
         append!(workqueue, methods)
-        compile!(codeinfos, workqueue; invokelatest_queue,
+        compile!(codeinfos, workqueue; invokelatest_queue, external_linkage,
                  enqueue_unprepared_invokes = trim_mode != TRIM_NO)
     end
 
     if invokelatest_queue !== nothing
         # This queue is intentionally aliased, to handle e.g. a `finalizer` calling `Core.finalizer`
         # (it will enqueue into itself and immediately drain)
-        compile!(codeinfos, invokelatest_queue; invokelatest_queue,
+        compile!(codeinfos, invokelatest_queue; invokelatest_queue, external_linkage,
                  enqueue_unprepared_invokes = trim_mode != TRIM_NO)
     end
 

Compiler/test/inference.jl

@@ -7653,7 +7653,7 @@ function tt57873(a::Vector{String}, pref)
     end
     return ret
 end
-let code = Compiler.typeinf_ext_toplevel(Any[Core.svec(Any,Tuple{typeof(tt57873),Vector{String},Tuple{String}})], [Base.get_world_counter()], Base.Compiler.TRIM_NO)[1]
+let code = Compiler.typeinf_ext_toplevel(Any[Core.svec(Any,Tuple{typeof(tt57873),Vector{String},Tuple{String}})], [Base.get_world_counter()], Base.Compiler.TRIM_NO, false)[1]
     @test !isempty(code)
     ## If we were to run trim here, we should fail with:
     #    Verifier error #1: unresolved invoke from statement tt57873(::Vector{String}, ::Tuple{String, String})::Vector{String}

Compiler/test/verifytrim.jl

@@ -35,6 +35,7 @@ let infos = typeinf_ext_toplevel(
     Any[Base.method_instance(scoped_trim_read, ())],
     [Base.get_world_counter()],
     TRIM_UNSAFE,
+    false,
 )[1]
     errors, _ = get_verify_typeinf_trim(infos)
     @test scoped_trim_read() == 2
@@ -43,7 +44,7 @@ end
 
 finalizer(@nospecialize(f), @nospecialize(o)) = Core.finalizer(f, o)
 
-let infos = typeinf_ext_toplevel(Any[Core.svec(Nothing, Tuple{typeof(finalizer), typeof(identity), Any})], [Base.get_world_counter()], TRIM_UNSAFE)[1]
+let infos = typeinf_ext_toplevel(Any[Core.svec(Nothing, Tuple{typeof(finalizer), typeof(identity), Any})], [Base.get_world_counter()], TRIM_UNSAFE, false)[1]
     errors, parents = get_verify_typeinf_trim(infos)
     @test !isempty(errors) # unresolvable finalizer
 
@@ -63,14 +64,14 @@ end
 
 # test that basic `cfunction` generation is allowed, when the dispatch target can be resolved
 make_cfunction() = @cfunction(+, Float64, (Int64,Int64))
-let infos = typeinf_ext_toplevel(Any[Core.svec(Ptr{Cvoid}, Tuple{typeof(make_cfunction)})], [Base.get_world_counter()], TRIM_UNSAFE)[1]
+let infos = typeinf_ext_toplevel(Any[Core.svec(Ptr{Cvoid}, Tuple{typeof(make_cfunction)})], [Base.get_world_counter()], TRIM_UNSAFE, false)[1]
     errors, parents = get_verify_typeinf_trim(infos)
     @test isempty(errors)
 end
 
 # use TRIM_UNSAFE to bypass verifier inside typeinf_ext_toplevel
 make_cfunction_bad(@nospecialize(f::Any)) = @cfunction($f, Float64, (Int64,Int64))::Base.CFunction
-let infos = typeinf_ext_toplevel(Any[Core.svec(Base.CFunction, Tuple{typeof(make_cfunction_bad), Any})], [Base.get_world_counter()], TRIM_UNSAFE)[1]
+let infos = typeinf_ext_toplevel(Any[Core.svec(Base.CFunction, Tuple{typeof(make_cfunction_bad), Any})], [Base.get_world_counter()], TRIM_UNSAFE, false)[1]
     errors, parents = get_verify_typeinf_trim(infos)
     @test !isempty(errors) # missing cfunction
 
@@ -95,7 +96,7 @@ let infos = typeinf_ext_toplevel(Any[Core.svec(Base.CFunction, Tuple{typeof(make
     @test repr == "unresolved ccallable for Tuple{$(typeof(make_cfunction_bad)), Any} => Base.CFunction\n\n"
 end
 
-let infos = typeinf_ext_toplevel(Any[Core.svec(Base.SecretBuffer, Tuple{Type{Base.SecretBuffer}})], [Base.get_world_counter()], TRIM_UNSAFE)[1]
+let infos = typeinf_ext_toplevel(Any[Core.svec(Base.SecretBuffer, Tuple{Type{Base.SecretBuffer}})], [Base.get_world_counter()], TRIM_UNSAFE, false)[1]
     @test length(infos) > 4
     errors, parents = get_verify_typeinf_trim(infos)
     @test isempty(errors)
@@ -113,7 +114,7 @@ let infos = typeinf_ext_toplevel(Any[Core.svec(Base.SecretBuffer, Tuple{Type{Bas
     @test repr == "unresolved ccallable for Tuple{Type{Base.SecretBuffer}} => Base.SecretBuffer\n\n"
 end
 
-let infos = typeinf_ext_toplevel(Any[Core.svec(Float64, Tuple{typeof(+), Int32, Int64})], [Base.get_world_counter()], TRIM_UNSAFE)[1]
+let infos = typeinf_ext_toplevel(Any[Core.svec(Float64, Tuple{typeof(+), Int32, Int64})], [Base.get_world_counter()], TRIM_UNSAFE, false)[1]
     errors, parents = get_verify_typeinf_trim(infos)
     (warn, desc) = only(errors)
     @test !warn
@@ -125,7 +126,7 @@ let infos = typeinf_ext_toplevel(Any[Core.svec(Float64, Tuple{typeof(+), Int32,
     @test repr == "ccallable declared return type does not match inference for Tuple{typeof(+), Int32, Int64} => Int64\n\n"
 end
 
-let infos = typeinf_ext_toplevel(Any[Core.svec(Int64, Tuple{typeof(ifelse), Bool, Int64, UInt64})], [Base.get_world_counter()], TRIM_UNSAFE)[1]
+let infos = typeinf_ext_toplevel(Any[Core.svec(Int64, Tuple{typeof(ifelse), Bool, Int64, UInt64})], [Base.get_world_counter()], TRIM_UNSAFE, false)[1]
     errors, parents = get_verify_typeinf_trim(infos)
     (warn, desc) = only(errors)
     @test warn  # this is a warning since Union{Int64, UInt64} <: Int64 is false but not an error
@@ -135,17 +136,17 @@ let infos = typeinf_ext_toplevel(Any[Core.svec(Int64, Tuple{typeof(ifelse), Bool
     @test repr == "ccallable declared return type does not match inference for Tuple{typeof(ifelse), Bool, Int64, UInt64} => Union{Int64, UInt64}\n\n"
 end
 
-let infos = typeinf_ext_toplevel(Any[Core.svec(Union{Int64,UInt64}, Tuple{typeof(ifelse), Bool, Int64, UInt64})], [Base.get_world_counter()], TRIM_SAFE)[1]
+let infos = typeinf_ext_toplevel(Any[Core.svec(Union{Int64,UInt64}, Tuple{typeof(ifelse), Bool, Int64, UInt64})], [Base.get_world_counter()], TRIM_SAFE, false)[1]
     errors, parents = get_verify_typeinf_trim(infos)
     @test isempty(errors)
-    infos = typeinf_ext_toplevel(Any[Core.svec(Real, Tuple{typeof(ifelse), Bool, Int64, UInt64})], [Base.get_world_counter()], TRIM_SAFE)[1]
+    infos = typeinf_ext_toplevel(Any[Core.svec(Real, Tuple{typeof(ifelse), Bool, Int64, UInt64})], [Base.get_world_counter()], TRIM_SAFE, false)[1]
     errors, parents = get_verify_typeinf_trim(infos)
     @test isempty(errors)
 end
 
 
 mi = Base.method_instance(sum, (Vector{Union{Int64,Float64, Float32,UInt32}},))
-let infos = typeinf_ext_toplevel(Any[mi], [Base.get_world_counter()], TRIM_UNSAFE)[1]
+let infos = typeinf_ext_toplevel(Any[mi], [Base.get_world_counter()], TRIM_UNSAFE, false)[1]
     errors, parents = get_verify_typeinf_trim(infos)
     @test !isempty(errors)
 end

src/aotcompile.cpp

@@ -827,6 +827,33 @@ static jl_compiled_functions_t::iterator get_ci_equiv_compiled(jl_code_instance_
     return compiled_functions.end();
 }
 
+// Check the global cache for an equivalent CodeInstance with a world age range
+// containing the world age range of the given CodeInstance.
+static jl_code_instance_t *jl_get_ci_equiv_range(jl_code_instance_t *ci JL_PROPAGATES_ROOT) JL_NOTSAFEPOINT
+{
+    size_t target_min = jl_atomic_load_relaxed(&ci->min_world);
+    size_t target_max = jl_atomic_load_relaxed(&ci->max_world);
+    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;
+    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_min &&
+            jl_atomic_load_relaxed(&codeinst->max_world) >= target_max &&
+            jl_egal(codeinst->def, def) &&
+            jl_egal(codeinst->owner, owner) &&
+            jl_egal(codeinst->rettype, rettype)) {
+            return codeinst;
+        }
+        codeinst = jl_atomic_load_relaxed(&codeinst->next);
+    }
+    return ci;
+}
+
+
 // Static version of JuliaOJIT::linkOutput
 static void aot_link_output(jl_codegen_output_t &out) JL_CANSAFEPOINT
 {
@@ -838,11 +865,16 @@ static void aot_link_output(jl_codegen_output_t &out) JL_CANSAFEPOINT
             continue;
 
         auto it = out.ci_funcs.find(ci);
+        // Prefer a equivalent CodeInstance that we are compiling.
         if (it == out.ci_funcs.end()) {
             auto equiv = get_ci_equiv_compiled(ci, out.ci_funcs);
             if (equiv != out.ci_funcs.end())
                 it = equiv;
         }
+        // If that fails, look for an equivalent CodeInstance that we can link to.
+        if (it == out.ci_funcs.end() && out.external_linkage &&
+            !(jl_atomic_load_relaxed(&ci->flags) & JL_CI_FLAGS_FROM_IMAGE))
+            ci = jl_get_ci_equiv_range(ci);
         jl_codeinst_funcs_t<Value *> funcs;
         if (it != out.ci_funcs.end()) {
             funcs = {it->second.invoke_api, it->second.invoke, it->second.specptr};