Ensure CI from external MI and foreign abs interpreter survive precompilation

Compiler/src/precompile.jl

@@ -311,8 +311,7 @@ function enqueue_specialization!(all::Bool, worklist, mi::Core.MethodInstance)
     while codeinst !== nothing
         do_compile = false
         if codeinst.owner !== nothing
-            # TODO(vchuravy) native code caching for foreign interpreters
-            # Skip foreign code instances
+            # This code instance is from a foreign interpreter, so we skip it
         elseif use_const_api(codeinst) # Check if invoke is jl_fptr_const_return
             do_compile = true
         elseif codeinst.invoke != C_NULL || codeinst.precompile
@@ -344,7 +343,8 @@ function compile_and_emit_native(worlds::Vector{UInt},
                                  newmodules, # Vector{Module} or Nothing
                                  mod_array, # Vector{Module} or Nothing
                                  all::Bool,
-                                 module_init_order::Vector{Any}) # Vector{Module}
+                                 module_init_order::Vector{Any}, # Vector{Module}
+                                 ext_foreign_cis::Vector{Any}) # Vector{CodeInstance}
     latestworld = worlds[end]
 
     # Step 1: Precompile all __init__ methods that will be required
@@ -375,7 +375,13 @@ function compile_and_emit_native(worlds::Vector{UInt},
             if new_ext_cis !== nothing
                 for i in 1:length(new_ext_cis::Vector{Any})
                     ci = new_ext_cis[i]::CodeInstance
-                    enqueue_specialization!(all, specialization_worklist, get_ci_mi(ci))
+                    if ci.owner !== nothing
+                        # enqueue_specialization will skip over CIs from foreign interpreters
+                        # and currently will visit at most one (do_compile) CI per method instance
+                        push!(ext_foreign_cis, ci)
+                    else
+                        enqueue_specialization!(all, specialization_worklist, get_ci_mi(ci))
+                    end
                 end
             end
         end

Compiler/src/typeinfer.jl

@@ -164,7 +164,7 @@ function finish!(interp::AbstractInterpreter, caller::InferenceState, validation
                 end
                 inferred_result = maybe_compress_codeinfo(interp, mi, inferred_result)
             elseif ci.owner === nothing
-                # The global cache can only handle objects that codegen understands
+                # The global cache can only handle objects that codegen understands (nothing or CodeInfo)
                 inferred_result = nothing
             end
         else

src/aotcompile.cpp

@@ -531,7 +531,7 @@ static bool canPartition(const Function &F)
 // `external_linkage` create linkages between pkgimages.
 extern "C" JL_DLLEXPORT_CODEGEN
 void *jl_create_native_impl(LLVMOrcThreadSafeModuleRef llvmmod, int trim, int external_linkage, size_t world,
-                           jl_array_t *mod_array, jl_array_t *worklist, int all, jl_array_t *module_init_order)
+                           jl_array_t *mod_array, jl_array_t *worklist, int all, jl_array_t *module_init_order, jl_array_t *ext_foreign_cis)
 {
     JL_TIMING(INFERENCE, INFERENCE);
     auto ct = jl_current_task;
@@ -547,7 +547,7 @@ void *jl_create_native_impl(LLVMOrcThreadSafeModuleRef llvmmod, int trim, int ex
         compiler_start_time = jl_hrtime();
 
     jl_value_t **fargs;
-    JL_GC_PUSHARGS(fargs, 8);
+    JL_GC_PUSHARGS(fargs, 9);
 #ifdef _P64
     jl_value_t *jl_array_ulong_type = jl_array_uint64_type;
 #else
@@ -568,10 +568,11 @@ void *jl_create_native_impl(LLVMOrcThreadSafeModuleRef llvmmod, int trim, int ex
     fargs[5] = mod_array ? (jl_value_t*)mod_array : jl_nothing; // mod_array (or nothing)
     fargs[6] = jl_box_bool(all);
     fargs[7] = module_init_order ? (jl_value_t*)module_init_order : jl_nothing; // module_init_order (or nothing)
+    fargs[8] = ext_foreign_cis ? (jl_value_t*)ext_foreign_cis : jl_nothing; // ext_foreign_cis (or nothing)
     size_t last_age = ct->world_age;
     ct->world_age = jl_typeinf_world;
-    fargs[0] = jl_apply(fargs, 8);
-    fargs[1] = fargs[2] = fargs[3] = fargs[4] = fargs[5] = fargs[6] = fargs[7] = NULL;
+    fargs[0] = jl_apply(fargs, 9);
+    fargs[1] = fargs[2] = fargs[3] = fargs[4] = fargs[5] = fargs[6] = fargs[7] = fargs[8] = NULL;
     ct->world_age = last_age;
     jl_value_t *codeinfos = fargs[0];
     JL_TYPECHK(jl_create_native, array_any, codeinfos);

src/julia_internal.h

@@ -2178,7 +2178,7 @@ JL_DLLIMPORT jl_value_t *jl_dump_function_ir(jl_llvmf_dump_t *dump, char strip_i
 JL_DLLIMPORT jl_value_t *jl_dump_function_asm(jl_llvmf_dump_t *dump, char emit_mc, const char* asm_variant, const char *debuginfo, char binary, char raw);
 
 typedef jl_value_t *(*jl_codeinstance_lookup_t)(jl_method_instance_t *mi JL_PROPAGATES_ROOT, size_t min_world, size_t max_world);
-JL_DLLIMPORT void *jl_create_native(LLVMOrcThreadSafeModuleRef llvmmod, int trim, int cache, size_t world, jl_array_t *mod_array JL_MAYBE_UNROOTED, jl_array_t *worklist JL_MAYBE_UNROOTED, int all, jl_array_t *module_init_order JL_MAYBE_UNROOTED);
+JL_DLLIMPORT void *jl_create_native(LLVMOrcThreadSafeModuleRef llvmmod, int trim, int cache, size_t world, jl_array_t *mod_array JL_MAYBE_UNROOTED, jl_array_t *worklist JL_MAYBE_UNROOTED, int all, jl_array_t *module_init_order JL_MAYBE_UNROOTED, jl_array_t *ext_foreign_cis JL_MAYBE_UNROOTED);
 JL_DLLIMPORT void *jl_emit_native(jl_array_t *codeinfos, LLVMOrcThreadSafeModuleRef llvmmod, const jl_cgparams_t *cgparams, int _external_linkage);
 JL_DLLIMPORT void jl_dump_native(void *native_code,
         const char *bc_fname, const char *unopt_bc_fname, const char *obj_fname, const char *asm_fname,

src/staticdata.c

@@ -3350,18 +3350,20 @@ JL_DLLEXPORT void jl_create_system_image(void **_native_data, jl_array_t *workli
         ff = f;
     }
 
-    jl_array_t *mod_array = NULL, *extext_methods = NULL, *new_ext_cis = NULL;
+    jl_array_t *mod_array = NULL, *extext_methods = NULL, *new_ext_cis = NULL, *ext_foreign_cis = NULL;
     int64_t checksumpos = 0;
     int64_t checksumpos_ff = 0;
     int64_t datastartpos = 0;
-    JL_GC_PUSH3(&mod_array, &extext_methods, &new_ext_cis);
+    JL_GC_PUSH4(&mod_array, &extext_methods, &new_ext_cis, &ext_foreign_cis);
+
+    ext_foreign_cis = jl_alloc_vec_any(0);
 
     mod_array = jl_get_loaded_modules();  // __toplevel__ modules loaded in this session (from Base.loaded_modules_array)
     if (worklist) {
         if (_native_data != NULL) {
             if (suppress_precompile)
                 newly_inferred = NULL;
-            *_native_data = jl_create_native(NULL, 0, 1, jl_atomic_load_acquire(&jl_world_counter), NULL, suppress_precompile ? (jl_array_t*)jl_an_empty_vec_any : worklist, 0, module_init_order);
+            *_native_data = jl_create_native(NULL, 0, 1, jl_atomic_load_acquire(&jl_world_counter), NULL, suppress_precompile ? (jl_array_t*)jl_an_empty_vec_any : worklist, 0, module_init_order, ext_foreign_cis);
         }
         jl_write_header_for_incremental(f, worklist, mod_array, udeps, srctextpos, &checksumpos);
         if (emit_split) {
@@ -3375,7 +3377,7 @@ JL_DLLEXPORT void jl_create_system_image(void **_native_data, jl_array_t *workli
         }
     }
     else if (_native_data != NULL) {
-        *_native_data = jl_create_native(NULL, jl_options.trim, 0, jl_atomic_load_acquire(&jl_world_counter), mod_array, NULL, jl_options.compile_enabled == JL_OPTIONS_COMPILE_ALL, module_init_order);
+        *_native_data = jl_create_native(NULL, jl_options.trim, 0, jl_atomic_load_acquire(&jl_world_counter), mod_array, NULL, jl_options.compile_enabled == JL_OPTIONS_COMPILE_ALL, module_init_order, ext_foreign_cis);
     }
     if (_native_data != NULL)
         native_functions = *_native_data;
@@ -3413,6 +3415,15 @@ JL_DLLEXPORT void jl_create_system_image(void **_native_data, jl_array_t *workli
             new_ext_cis = jl_compute_new_ext_cis();
         }
 
+        // Merge foreign & external CIs
+        if (ext_foreign_cis) {
+            size_t n_ext = jl_array_nrows(ext_foreign_cis);
+            for (size_t i = 0; i < n_ext; i++) {
+                jl_array_ptr_1d_push(new_ext_cis, jl_array_ptr_ref(ext_foreign_cis, i));
+            }
+        }
+        ext_foreign_cis = NULL; // not needed anymore, free it
+
         // Collect method extensions
         extext_methods = jl_alloc_vec_any(0);
         jl_collect_extext_methods(extext_methods, mod_array);

src/staticdata_utils.c

@@ -502,7 +502,7 @@ static jl_array_t *queue_external_cis(jl_array_t *list, jl_query_cache *query_ca
         int dispatch_status = jl_atomic_load_relaxed(&m->dispatch_status);
         if (!(dispatch_status & METHOD_SIG_LATEST_WHICH))
             continue; // ignore replaced methods
-        if (ci->owner == jl_nothing && jl_atomic_load_relaxed(&ci->inferred) && jl_is_method(m) && jl_object_in_image((jl_value_t*)m->module)) {
+        if (jl_atomic_load_relaxed(&ci->inferred) && jl_is_method(m) && jl_object_in_image((jl_value_t*)m->module)) {
             int found = has_backedge_to_worklist(mi, &visited, &stack, query_cache);
             assert(found == 0 || found == 1 || found == 2);
             assert(stack.len == 0);

test/precompile.jl

@@ -1872,6 +1872,7 @@ end
     dir = @__DIR__
     @test success(pipeline(Cmd(`$(Base.julia_cmd()) --startup-file=no precompile_absint1.jl`; dir); stdout, stderr))
     @test success(pipeline(Cmd(`$(Base.julia_cmd()) --startup-file=no precompile_absint2.jl`; dir); stdout, stderr))
+    @test success(pipeline(Cmd(`$(Base.julia_cmd()) --startup-file=no precompile_extmi.jl`; dir); stdout, stderr))
 end
 
 precompile_test_harness("Recursive types") do load_path

test/precompile_absint1.jl

@@ -43,39 +43,35 @@ precompile_test_harness() do load_path
         let m = only(methods(TestAbsIntPrecompile1.basic_callee))
             mi = only(Base.specializations(m))
             ci = mi.cache
-            @test isdefined(ci, :next)
-            @test ci.owner === cache_owner
+            ci = check_presence(mi, nothing)
+            @test ci !== nothing
+            @test ci.owner === nothing
             @test ci.max_world == typemax(UInt)
             @test Base.module_build_id(TestAbsIntPrecompile1) ==
                 Base.object_build_id(ci)
-            ci = ci.next
-            @test !isdefined(ci, :next)
-            @test ci.owner === nothing
+            ci = check_presence(mi, cache_owner)
+            @test ci !== nothing
+            @test ci.owner === cache_owner
             @test ci.max_world == typemax(UInt)
             @test Base.module_build_id(TestAbsIntPrecompile1) ==
                 Base.object_build_id(ci)
         end
         let m = only(methods(sum, (Vector{Float64},)))
-            found = false
             for mi in Base.specializations(m)
                 if mi isa Core.MethodInstance && mi.specTypes == Tuple{typeof(sum),Vector{Float64}}
-                    ci = mi.cache
-                    @test isdefined(ci, :next)
-                    @test ci.owner === cache_owner
-                    @test ci.max_world == typemax(UInt)
-                    @test Base.module_build_id(TestAbsIntPrecompile1) ==
-                        Base.object_build_id(ci)
-                    ci = ci.next
-                    @test !isdefined(ci, :next)
+                    ci = check_presence(mi, nothing)
+                    @test ci !== nothing
                     @test ci.owner === nothing
                     @test ci.max_world == typemax(UInt)
+                    @test Base.module_build_id(TestAbsIntPrecompile1) == Base.object_build_id(ci)
+                    ci = check_presence(mi, cache_owner)
+                    @test ci !== nothing
+                    @test ci.owner === cache_owner
+                    @test ci.max_world == typemax(UInt)
                     @test Base.module_build_id(TestAbsIntPrecompile1) ==
                         Base.object_build_id(ci)
-                    found = true
-                    break
                 end
             end
-            @test found
         end
     end
 end

test/precompile_absint2.jl

@@ -62,40 +62,34 @@ precompile_test_harness() do load_path
             TestAbsIntPrecompile2.Custom.PrecompileInterpreter())
         let m = only(methods(TestAbsIntPrecompile2.basic_callee))
             mi = only(Base.specializations(m))
-            ci = mi.cache
-            @test isdefined(ci, :next)
-            @test ci.owner === cache_owner
-            @test ci.max_world == typemax(UInt)
-            @test Base.module_build_id(TestAbsIntPrecompile2) ==
-                Base.object_build_id(ci)
-            ci = ci.next
-            @test !isdefined(ci, :next)
+
+            ci = check_presence(mi, nothing)
+            @test ci !== nothing
             @test ci.owner === nothing
             @test ci.max_world == typemax(UInt)
             @test Base.module_build_id(TestAbsIntPrecompile2) ==
                 Base.object_build_id(ci)
+            ci = check_presence(mi, cache_owner)
+            @test ci !== nothing
+            @test ci.owner === cache_owner
+            @test ci.max_world == typemax(UInt)
+            @test Base.module_build_id(TestAbsIntPrecompile2) == Base.object_build_id(ci)
         end
         let m = only(methods(sum, (Vector{Float64},)))
-            found = false
             for mi = Base.specializations(m)
                 if mi isa Core.MethodInstance && mi.specTypes == Tuple{typeof(sum),Vector{Float64}}
-                    ci = mi.cache
-                    @test isdefined(ci, :next)
-                    @test ci.owner === cache_owner
-                    @test ci.max_world == typemax(UInt)
-                    @test Base.module_build_id(TestAbsIntPrecompile2) ==
-                        Base.object_build_id(ci)
-                    ci = ci.next
-                    @test !isdefined(ci, :next)
+                    ci = check_presence(mi, nothing)
+                    @test ci !== nothing
                     @test ci.owner === nothing
                     @test ci.max_world == typemax(UInt)
-                    @test Base.module_build_id(TestAbsIntPrecompile2) ==
-                        Base.object_build_id(ci)
-                    found = true
-                    break
+                    @test Base.module_build_id(TestAbsIntPrecompile2) == Base.object_build_id(ci)
+                    ci = check_presence(mi, cache_owner)
+                    @test ci !== nothing
+                    @test ci.owner === cache_owner
+                    @test ci.max_world == typemax(UInt)
+                    @test Base.module_build_id(TestAbsIntPrecompile2) == Base.object_build_id(ci)
                 end
             end
-            @test found
         end
     end
 end

test/precompile_extmi.jl

@@ -0,0 +1,75 @@
+# This file is a part of Julia. License is MIT: https://julialang.org/license
+
+using Test
+
+include("precompile_utils.jl")
+
+precompile_test_harness() do load_path
+    write(joinpath(load_path, "ExampleCompiler.jl"), :(
+    module ExampleCompiler
+        const CC = Core.Compiler
+
+        struct ExampleInterpreter <: CC.AbstractInterpreter
+            world::UInt
+            inf_cache::CC.InferenceCache
+        end
+        ExampleInterpreter(world::UInt) =
+            ExampleInterpreter(world, CC.InferenceCache())
+
+        CC.InferenceParams(::ExampleInterpreter) = CC.InferenceParams()
+        CC.OptimizationParams(::ExampleInterpreter) = CC.OptimizationParams()
+        CC.get_inference_cache(interp::ExampleInterpreter) = interp.inf_cache
+        CC.cache_owner(interp::ExampleInterpreter) = :ExampleInterpreter
+
+        CC.get_inference_world(interp::ExampleInterpreter) = interp.world
+        CC.lock_mi_inference(::ExampleInterpreter, ::Core.MethodInstance) = nothing
+        CC.unlock_mi_inference(::ExampleInterpreter, ::Core.MethodInstance) = nothing
+
+        function infer(mi, world)
+            interp = ExampleInterpreter(world)
+            ci = CC.typeinf_ext(interp, mi, CC.SOURCE_MODE_GET_SOURCE)
+            @assert ci !== nothing "Inference of $mi failed"
+
+            return ci
+        end
+
+        function precompile(f, tt; world = Base.get_world_counter())
+            mi = Base.method_instance(f, tt; world)
+            infer(mi, world)
+        end
+    end) |> string)
+    Base.compilecache(Base.PkgId("ExampleCompiler"))
+
+    write(joinpath(load_path, "ExampleUser.jl"), :(
+    module ExampleUser
+        import ExampleCompiler
+
+        function square(x)
+            return x * x
+        end
+        ExampleCompiler.precompile(square, (Float64,))
+
+        # Stubbed together version of PrecompileTools
+        ccall(:jl_tag_newly_inferred_enable, Cvoid, ())
+        try
+            # Important `identity(::Any) mi does not belong to us`
+            ExampleCompiler.precompile(identity, (Float64,))
+        finally
+            ccall(:jl_tag_newly_inferred_disable, Cvoid, ())
+        end
+    end) |> string)
+    Base.compilecache(Base.PkgId("ExampleUser"))
+
+    @eval let
+        using ExampleUser
+        cache_owner = :ExampleInterpreter
+
+        mi_square = Base.method_instance(ExampleUser.square, (Float64,))
+        @test check_presence(mi_square, :ExampleInterpreter) !== nothing
+
+        mi_identity = Base.method_instance(identity, (Float64,))
+        @test check_presence(mi_identity, :ExampleInterpreter) !== nothing
+    end
+end
+
+finish_precompile_test!()

test/precompile_utils.jl

@@ -29,3 +29,18 @@ let original_depot_path = copy(Base.DEPOT_PATH)
         append!(Base.LOAD_PATH, original_load_path)
     end
 end
+
+function check_presence(mi, token)
+    ci = isdefined(mi, :cache) ? mi.cache : nothing
+    while ci !== nothing
+        # CI should have been validated
+        @test ci.max_world != Base.ReinferUtils.WORLD_AGE_REVALIDATION_SENTINEL
+        @test ci.min_world != ~zero(UInt)
+        # Chose a CI with the right owner and current validity.
+        if ci.owner === token && ci.max_world == typemax(UInt)
+            return ci
+        end
+        ci = isdefined(ci, :next) ? ci.next : nothing
+    end
+    return nothing
+end