precompile: fix a number of mistakes in handling precompiled code instance caches

Get the real list of external code instances worthwhile to re-cache
directly from `native_functions` instead of calling
`jl_compute_new_ext_cis` a second time if we have `native_functions`. If
not saving native code though, we can still call it to get a list of
code instances that might be worth including even without native code.

Author: vtjnash

Compiler/src/bindinginvalidations.jl

@@ -188,7 +188,7 @@ function scan_new_method!(method::Method, image_backedges_only::Bool)
     @atomic method.did_scan_source |= 0x1
 end
 
-function scan_new_methods!(extext_methods::Vector{Any}, internal_methods::Vector{Any}, image_backedges_only::Bool)
+function scan_new_methods!(internal_methods::Vector{Any}, image_backedges_only::Bool)
     if image_backedges_only && generating_output(true)
         # Replacing image bindings is forbidden during incremental precompilation - skip backedge insertion
         return
@@ -198,7 +198,4 @@ function scan_new_methods!(extext_methods::Vector{Any}, internal_methods::Vector
            scan_new_method!(method, image_backedges_only)
         end
     end
-    for tme::Core.TypeMapEntry in extext_methods
-        scan_new_method!(tme.func::Method, image_backedges_only)
-    end
 end

Compiler/src/precompile.jl

@@ -426,9 +426,6 @@ function compile_and_emit_native(worlds::Vector{UInt},
         invokelatest(exit, 1)
     end
 
-    # Step 5: Always set newly_inferred global for serialization use
-    #ccall(:jl_set_newly_inferred, Cvoid, (Any,), filter(ci -> ci isa CodeInstance, codeinfos))
-
     return codeinfos
 
 end

Compiler/src/reinfer.jl

@@ -2,7 +2,7 @@
 
 using ..Compiler.Base
 using ..Compiler: _findsup, store_backedges, JLOptions, get_world_counter,
-    _methods_by_ftype, get_methodtable, get_ci_mi, should_instrument,
+    _methods_by_ftype, get_methodtable, get_ci_mi, ci_has_abi, should_instrument,
     morespecific, RefValue, get_require_world, Vector, IdDict
 using .Core: CodeInstance, MethodInstance
 
@@ -69,37 +69,46 @@ end
 # Restore backedges to external targets
 # `edges` = [caller1, ...], the list of worklist-owned code instances internally
 # `ext_ci_list` = [caller1, ...], the list of worklist-owned code instances externally
-function insert_backedges(edges::Vector{Any}, ext_ci_list::Union{Nothing,Vector{Any}}, extext_methods::Vector{Any}, internal_methods::Vector{Any})
+function insert_backedges(ext_ci_list::Union{Nothing,Vector{Any}}, internal_methods::Vector{Any})
     # determine which CodeInstance objects are still valid in our image
     # to enable any applicable new codes
     backedges_only = unsafe_load(cglobal(:jl_first_image_replacement_world, UInt)) == typemax(UInt)
-    scan_new_methods!(extext_methods, internal_methods, backedges_only)
+    scan_new_methods!(internal_methods, backedges_only)
     workspace = VerifyMethodWorkspace()
-    _insert_backedges(edges, workspace)
+    scan_new_code!(internal_methods, workspace)
     if ext_ci_list !== nothing
-        _insert_backedges(ext_ci_list, workspace, #=external=#true)
+        insert_ext_cache(ext_ci_list)
     end
+    nothing
 end
 
-function _insert_backedges(edges::Vector{Any}, workspace::VerifyMethodWorkspace, external::Bool=false)
-    for i = 1:length(edges)
-        codeinst = edges[i]::CodeInstance
+function scan_new_code!(internal_methods::Vector{Any}, workspace::VerifyMethodWorkspace)
+    for i = 1:length(internal_methods)
+        codeinst = internal_methods[i]
+        codeinst isa CodeInstance || continue
+        # codeinst.owner === nothing || continue
         validation_world = get_world_counter()
         verify_method_graph(codeinst, validation_world, workspace)
         # After validation, under the world_counter_lock, set max_world to typemax(UInt) for all dependencies
         # (recursively). From that point onward the ordinary backedge mechanism is responsible for maintaining
         # validity.
         @ccall jl_promote_ci_to_current(codeinst::Any, validation_world::UInt)::Cvoid
+    end
+end
+
+function insert_ext_cache(edges::Vector{Any})
+    for i = 1:length(edges)
+        codeinst = edges[i]::CodeInstance
         minvalid = codeinst.min_world
         maxvalid = codeinst.max_world
         # Finally, if this CI is still valid in some world age and belongs to an external method(specialization),
-        # poke it that mi's cache
-        if maxvalid ≥ minvalid && external
+        # poke it in that mi's cache, unless there is already one there that has an invoke pointer
+        if maxvalid ≥ minvalid
             caller = get_ci_mi(codeinst)
             @assert isdefined(codeinst, :inferred) # See #53586, #53109
             inferred = @ccall jl_rettype_inferred(
                 codeinst.owner::Any, caller::Any, minvalid::UInt, maxvalid::UInt)::Any
-            if inferred !== nothing
+            if inferred !== nothing && (ci_has_abi(inferred::CodeInstance) || !ci_has_abi(codeinst))
                 # We already got a code instance for this world age range from
                 # somewhere else - we don't need this one.
             else
@@ -658,7 +667,7 @@ function verify_invokesig(@nospecialize(invokesig), expected::Method, world::UIn
 end
 
 # Wrapper to call insert_backedges in typeinf_world for external calls
-function insert_backedges_typeinf(edges::Vector{Any}, ext_ci_list::Union{Nothing,Vector{Any}}, extext_methods::Vector{Any}, internal_methods::Vector{Any})
-    args = Any[insert_backedges, edges, ext_ci_list, extext_methods, internal_methods]
+function insert_backedges_typeinf(ext_ci_list::Union{Nothing,Vector{Any}}, internal_methods::Vector{Any})
+    args = Any[insert_backedges, ext_ci_list, internal_methods]
     return ccall(:jl_call_in_typeinf_world, Any, (Ptr{Any}, Cint), args, length(args))
 end

Compiler/src/typeinfer.jl

@@ -1288,6 +1288,8 @@ prepared interp to be able to provide source code for it.
 """
 const SOURCE_MODE_GET_SOURCE = 0xf
 
+ci_has_abi(code::CodeInstance) = (@atomic :acquire code.invoke) !== C_NULL
+
 """
     ci_has_abi(interp::AbstractInterpreter, code::CodeInstance)
 
@@ -1296,7 +1298,7 @@ interp gave it to the runtime system (either because it already has an ->invoke
 ptr, or because interp has source that could be compiled).
 """
 function ci_has_abi(interp::AbstractInterpreter, code::CodeInstance)
-    (@atomic :acquire code.invoke) !== C_NULL && return true
+    ci_has_abi(code) && return true
     return ci_has_source(interp, code)
 end
 

base/loading.jl

@@ -1283,12 +1283,10 @@ function _include_from_serialized(pkg::PkgId, path::String, ocachepath::Union{No
         end
 
         sv = sv::SimpleVector
-        edges = sv[3]::Vector{Any}
-        ext_edges = sv[4]::Union{Nothing,Vector{Any}}
-        extext_methods = sv[5]::Vector{Any}
-        internal_methods = sv[6]::Vector{Any}
+        ext_edges = sv[3]::Union{Nothing,Vector{Any}}
+        internal_methods = sv[4]::Vector{Any}
         Compiler.@zone "CC: INSERT_BACKEDGES" begin
-            ReinferUtils.insert_backedges_typeinf(edges, ext_edges, extext_methods, internal_methods)
+            ReinferUtils.insert_backedges_typeinf(ext_edges, internal_methods)
         end
         restored = register_restored_modules(sv, pkg, path)
 

src/aotcompile.cpp

@@ -92,7 +92,7 @@ void jl_get_function_id_impl(void *native_code, jl_code_instance_t *codeinst,
 }
 
 extern "C" JL_DLLEXPORT_CODEGEN void
-jl_get_llvm_mis_impl(void *native_code, size_t *num_elements, jl_method_instance_t **data)
+jl_get_llvm_cis_impl(void *native_code, size_t *num_elements, jl_code_instance_t **data)
 {
     jl_native_code_desc_t *desc = (jl_native_code_desc_t *)native_code;
     auto &map = desc->jl_fvar_map;
@@ -105,7 +105,7 @@ jl_get_llvm_mis_impl(void *native_code, size_t *num_elements, jl_method_instance
     assert(*num_elements == map.size());
     size_t i = 0;
     for (auto &ci : map) {
-        data[i++] = jl_get_ci_mi(ci.first);
+        data[i++] = ci.first;
     }
 }
 

src/codegen-stubs.c

@@ -13,10 +13,10 @@
 JL_DLLEXPORT void jl_dump_native_fallback(void *native_code,
         const char *bc_fname, const char *unopt_bc_fname, const char *obj_fname, const char *asm_fname,
         ios_t *z, ios_t *s) UNAVAILABLE
-JL_DLLEXPORT void jl_get_llvm_gvs_fallback(void *native_code, arraylist_t *gvs) UNAVAILABLE
-JL_DLLEXPORT void jl_get_llvm_gv_inits_fallback(void *native_code, arraylist_t *inits) UNAVAILABLE
-JL_DLLEXPORT void jl_get_llvm_external_fns_fallback(void *native_code, arraylist_t *gvs) UNAVAILABLE
-JL_DLLEXPORT void jl_get_llvm_mis_fallback(void *native_code, arraylist_t* MIs) UNAVAILABLE
+JL_DLLEXPORT void jl_get_llvm_gvs_fallback(void *native_code, size_t *num, void **gvs) UNAVAILABLE
+JL_DLLEXPORT void jl_get_llvm_gv_inits_fallback(void *native_code, size_t *num, void **inits) UNAVAILABLE
+JL_DLLEXPORT void jl_get_llvm_external_fns_fallback(void *native_code, size_t *num, void **gvs) UNAVAILABLE
+JL_DLLEXPORT void jl_get_llvm_cis_fallback(void *native_code, size_t *num, void **CIs) UNAVAILABLE
 
 JL_DLLEXPORT jl_value_t *jl_dump_method_asm_fallback(jl_method_instance_t *linfo, size_t world,
         char emit_mc, char getwrapper, const char* asm_variant, const char *debuginfo, char binary) UNAVAILABLE

src/jl_exported_funcs.inc

@@ -517,7 +517,7 @@
     YY(jl_get_llvm_gvs) \
     YY(jl_get_llvm_gv_inits) \
     YY(jl_get_llvm_external_fns) \
-    YY(jl_get_llvm_mis) \
+    YY(jl_get_llvm_cis) \
     YY(jl_dump_function_asm) \
     YY(jl_LLVMCreateDisasm) \
     YY(jl_LLVMDisasmInstruction) \

src/julia_internal.h

@@ -2067,8 +2067,8 @@ JL_DLLIMPORT void jl_get_function_id(void *native_code, jl_code_instance_t *ncod
         int32_t *func_idx, int32_t *specfunc_idx);
 JL_DLLIMPORT void jl_register_fptrs(uint64_t image_base, const struct _jl_image_fptrs_t *fptrs,
                                     jl_method_instance_t **linfos, size_t n);
-JL_DLLIMPORT void jl_get_llvm_mis(void *native_code, size_t *num_els,
-                                  jl_method_instance_t *MIs);
+JL_DLLIMPORT void jl_get_llvm_cis(void *native_code, size_t *num_els,
+                                  jl_code_instance_t **CIs);
 JL_DLLIMPORT void jl_init_codegen(void);
 JL_DLLIMPORT void jl_teardown_codegen(void) JL_NOTSAFEPOINT;
 JL_DLLIMPORT int jl_getFunctionInfo(jl_frame_t **frames, uintptr_t pointer, int skipC, int noInline) JL_NOTSAFEPOINT;