Add API for non-native compiler caching of external methods

base/loading.jl

@@ -3254,6 +3254,7 @@ end
 
 # Const global for GC root
 const newly_inferred = CodeInstance[]
+const newly_inferred_external = CodeInstance[]
 
 # this is called in the external process that generates precompiled package files
 function include_package_for_output(pkg::PkgId, input::String, syntax_version::VersionNumber, depot_path::Vector{String}, dl_load_path::Vector{String}, load_path::Vector{String},
@@ -3280,6 +3281,7 @@ function include_package_for_output(pkg::PkgId, input::String, syntax_version::V
     end
 
     ccall(:jl_set_newly_inferred, Cvoid, (Any,), newly_inferred)
+    ccall(:jl_set_newly_inferred_external, Cvoid, (Any,), newly_inferred_external)
     # This one changes the parser behavior
     __toplevel__.var"#_internal_julia_parse" = VersionedParse(syntax_version)
     # This one is the compatibility marker for cache loading
@@ -3292,6 +3294,7 @@ function include_package_for_output(pkg::PkgId, input::String, syntax_version::V
         exit(125) # we define status = 125 means PrecompileableError
     finally
         ccall(:jl_set_newly_inferred, Cvoid, (Any,), nothing)
+        ccall(:jl_set_newly_inferred_external, Cvoid, (Any,), nothing)
     end
     # check that the package defined the expected module so we can give a nice error message if not
     m = maybe_root_module(pkg)
@@ -3301,6 +3304,7 @@ function include_package_for_output(pkg::PkgId, input::String, syntax_version::V
     # in the output. We removed it above to avoid including any code we may
     # have compiled for error handling and validation.
     ccall(:jl_set_newly_inferred, Cvoid, (Any,), newly_inferred)
+    ccall(:jl_set_newly_inferred_external, Cvoid, (Any,), newly_inferred_external)
     @lock require_lock end_loading(pkg, m)
     # insert_extension_triggers(pkg)
     # run_package_callbacks(pkg)
@@ -4522,6 +4526,16 @@ end
 precompile(mi::MethodInstance, world::UInt=get_world_counter()) =
     (ccall(:jl_compile_method_instance, Cvoid, (Any, Ptr{Cvoid}, UInt), mi, C_NULL, world); return true)
 
+"""
+    precompile(ci::CodeInstance)
+
+Mark a `CodeInstance` for inclusion in the precompilation cache. This is intended
+for use by non-native code caches (e.g. GPU compilers) to persist `CodeInstance`s
+whose `owner` is not `jl_nothing` into the package image.
+"""
+precompile(ci::CodeInstance) =
+    ccall(:jl_push_newly_inferred_external, Cvoid, (Any,), ci)
+
 """
     precompile(f, argtypes::Tuple{Vararg{Any}}, m::Method)
 

src/julia.h

@@ -2183,6 +2183,8 @@ JL_DLLEXPORT jl_value_t *jl_object_top_module(jl_value_t* v) JL_NOTSAFEPOINT;
 JL_DLLEXPORT void jl_set_newly_inferred(jl_value_t *newly_inferred);
 JL_DLLEXPORT jl_array_t* jl_compute_new_ext_cis(void);
 JL_DLLEXPORT void jl_push_newly_inferred(jl_value_t *ci);
+JL_DLLEXPORT void jl_set_newly_inferred_external(jl_value_t *newly_inferred_external);
+JL_DLLEXPORT void jl_push_newly_inferred_external(jl_value_t *ci);
 JL_DLLEXPORT void jl_set_inference_entrance_backtraces(jl_value_t *inference_entrance_backtraces);
 JL_DLLEXPORT void jl_push_inference_entrance_backtraces(jl_value_t *ci);
 JL_DLLEXPORT void jl_write_compiler_output(void);

src/staticdata.c

@@ -3357,11 +3357,11 @@ 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_precomp_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_precomp_cis);
 
     mod_array = jl_get_loaded_modules();  // __toplevel__ modules loaded in this session (from Base.loaded_modules_array)
     if (worklist) {
@@ -3420,6 +3420,16 @@ JL_DLLEXPORT void jl_create_system_image(void **_native_data, jl_array_t *workli
             new_ext_cis = jl_compute_new_ext_cis();
         }
 
+        // Merge explicitly-marked external CIs (from non-native compilers)
+        if (newly_inferred_external != NULL) {
+            ext_precomp_cis = newly_inferred_external; // root it
+            size_t n_ext = jl_array_nrows(ext_precomp_cis);
+            for (size_t i = 0; i < n_ext; i++) {
+                jl_array_ptr_1d_push(new_ext_cis, jl_array_ptr_ref(ext_precomp_cis, i));
+            }
+            ext_precomp_cis = NULL;
+        }
+
         // Collect method extensions
         extext_methods = jl_alloc_vec_any(0);
         jl_collect_extext_methods(extext_methods, mod_array);

src/staticdata_utils.c

@@ -83,6 +83,7 @@ static uint64_t jl_worklist_key(jl_array_t *worklist) JL_NOTSAFEPOINT
 }
 
 static jl_array_t *newly_inferred JL_GLOBALLY_ROOTED /*FIXME*/;
+static jl_array_t *newly_inferred_external JL_GLOBALLY_ROOTED /*FIXME*/ = NULL;
 // Mutex for newly_inferred
 jl_mutex_t newly_inferred_mutex;
 extern jl_mutex_t world_counter_lock;
@@ -144,6 +145,27 @@ JL_DLLEXPORT void jl_push_newly_inferred(jl_value_t* ci)
     JL_UNLOCK(&newly_inferred_mutex);
 }
 
+// Register array of CodeInstances from non-native compilers to be cached during precompilation
+JL_DLLEXPORT void jl_set_newly_inferred_external(jl_value_t *_newly_inferred_external)
+{
+    assert(_newly_inferred_external == NULL || _newly_inferred_external == jl_nothing ||
+           jl_is_array(_newly_inferred_external));
+    if (_newly_inferred_external == jl_nothing)
+        _newly_inferred_external = NULL;
+    newly_inferred_external = (jl_array_t*)_newly_inferred_external;
+}
+
+JL_DLLEXPORT void jl_push_newly_inferred_external(jl_value_t *ci)
+{
+    if (!newly_inferred_external)
+        return;
+    JL_LOCK(&newly_inferred_mutex);
+    size_t end = jl_array_nrows(newly_inferred_external);
+    jl_array_grow_end(newly_inferred_external, 1);
+    jl_array_ptr_set(newly_inferred_external, end, ci);
+    JL_UNLOCK(&newly_inferred_mutex);
+}
+
 
 static jl_array_t *inference_entrance_backtraces JL_GLOBALLY_ROOTED /*FIXME*/ = NULL;
 // Mutex for inference_entrance_backtraces