Support external linkage in "sysimages"

This takes the first steps towards saving native code in package precompilation: - extend staticdata.c to support external linkage - add linkage to LLVM module-for-sysimage creation - add a stub implementation of a new LLD_jll stdlib This has no user-visible consequences (yet), because the machinery is only invoked by command-line arguments `--output-o $output --output-incremental=yes` which is not a previously-supported combination. Co-authored-by: Valentin Churavy <[email protected]>
JuliaLang · Oct 3, 2022 · 92ebe50 · 92ebe50
1 parent fcdc5bc
commit 92ebe50
Show file tree

Hide file tree

Showing 23 changed files with 758 additions and 158 deletions.
diff --git a/deps/Makefile b/deps/Makefile
@@ -85,7 +85,7 @@ DEP_LIBS += dsfmt
 endif
 
 ifeq ($(USE_SYSTEM_LLVM), 0)
-DEP_LIBS += llvm
+DEP_LIBS += llvm lld
 endif
 
 ifeq ($(USE_SYSTEM_LLD), 0)

diff --git a/deps/llvm.mk b/deps/llvm.mk
@@ -308,8 +308,8 @@ LLVM_TOOLS_JLL_TAGS := -llvm_version+$(LLVM_VER_MAJ)
 endif
 
 $(eval $(call bb-install,llvm,LLVM,false,true))
-$(eval $(call bb-install,clang,CLANG,false,true))
 $(eval $(call bb-install,lld,LLD,false,true))
+$(eval $(call bb-install,clang,CLANG,false,true))
 $(eval $(call bb-install,llvm-tools,LLVM_TOOLS,false,true))
 
 endif # USE_BINARYBUILDER_LLVM

diff --git a/src/aotcompile.cpp b/src/aotcompile.cpp
@@ -123,6 +123,17 @@ int32_t jl_get_llvm_gv_impl(void *native_code, jl_value_t *p)
     return 0;
 }
 
+extern "C" JL_DLLEXPORT
+void jl_iterate_llvm_gv_impl(void *native_code, void (*callback)(void*, int32_t, void*), void* ctx)
+{
+    jl_native_code_desc_t *data = (jl_native_code_desc_t*)native_code;
+    if (data) {
+        for (std::pair<void*, int32_t> pair : data->jl_value_to_llvm) {
+            callback(ctx, pair.second, pair.first);
+        }
+    }
+}
+
 extern "C" JL_DLLEXPORT
 LLVMOrcThreadSafeModuleRef jl_get_llvm_module_impl(void *native_code)
 {
@@ -148,7 +159,6 @@ static void emit_offset_table(Module &mod, const std::vector<GlobalValue*> &vars
 {
     // Emit a global variable with all the variable addresses.
     // The cloning pass will convert them into offsets.
-    assert(!vars.empty());
     size_t nvars = vars.size();
     std::vector<Constant*> addrs(nvars);
     for (size_t i = 0; i < nvars; i++) {
@@ -258,17 +268,17 @@ static void jl_ci_cache_lookup(const jl_cgparams_t &cgparams, jl_method_instance
 // this builds the object file portion of the sysimage files for fast startup, and can
 // also be used be extern consumers like GPUCompiler.jl to obtain a module containing
 // all reachable & inferrrable functions. The `policy` flag switches between the default
-// mode `0`, the extern mode `1`, and imaging mode `2`.
+// mode `0`, the extern mode `1`.
 extern "C" JL_DLLEXPORT
-void *jl_create_native_impl(jl_array_t *methods, LLVMOrcThreadSafeModuleRef llvmmod, const jl_cgparams_t *cgparams, int _policy)
+void *jl_create_native_impl(jl_array_t *methods, LLVMOrcThreadSafeModuleRef llvmmod, const jl_cgparams_t *cgparams, int _policy, int _imaging_mode)
 {
     ++CreateNativeCalls;
     CreateNativeMax.updateMax(jl_array_len(methods));
     if (cgparams == NULL)
         cgparams = &jl_default_cgparams;
     jl_native_code_desc_t *data = new jl_native_code_desc_t;
     CompilationPolicy policy = (CompilationPolicy) _policy;
-    bool imaging = imaging_default() || policy == CompilationPolicy::ImagingMode;
+    bool imaging = imaging_default() || _imaging_mode == 1;
     jl_workqueue_t emitted;
     jl_method_instance_t *mi = NULL;
     jl_code_info_t *src = NULL;
@@ -572,7 +582,7 @@ void jl_dump_native_impl(void *native_code,
     Type *T_psize = T_size->getPointerTo();
 
     // add metadata information
-    if (imaging_default()) {
+    if (imaging_default() || jl_options.outputo) {
         emit_offset_table(*dataM, data->jl_sysimg_gvars, "jl_sysimg_gvars", T_psize);
         emit_offset_table(*dataM, data->jl_sysimg_fvars, "jl_sysimg_fvars", T_psize);
 

diff --git a/src/codegen-stubs.c b/src/codegen-stubs.c
@@ -14,6 +14,7 @@ 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,
         const char *sysimg_data, size_t sysimg_len) UNAVAILABLE
 JL_DLLEXPORT int32_t jl_get_llvm_gv_fallback(void *native_code, jl_value_t *p) UNAVAILABLE
+JL_DLLEXPORT void jl_iterate_llvm_gv_fallback(void *native_code, void (*callback)(void*, int32_t, void*), void* ctx) UNAVAILABLE
 
 JL_DLLEXPORT void jl_extern_c_fallback(jl_function_t *f, jl_value_t *rt, jl_value_t *argt, char *name) UNAVAILABLE
 JL_DLLEXPORT jl_value_t *jl_dump_method_asm_fallback(jl_method_instance_t *linfo, size_t world,
@@ -66,7 +67,7 @@ JL_DLLEXPORT size_t jl_jit_total_bytes_fallback(void)
     return 0;
 }
 
-JL_DLLEXPORT void *jl_create_native_fallback(jl_array_t *methods, LLVMOrcThreadSafeModuleRef llvmctxt, const jl_cgparams_t *cgparams, int _policy) UNAVAILABLE
+JL_DLLEXPORT void *jl_create_native_fallback(jl_array_t *methods, LLVMOrcThreadSafeModuleRef llvmmod, const jl_cgparams_t *cgparams, int _policy, int _imaging_mode) UNAVAILABLE
 
 JL_DLLEXPORT void jl_dump_compiles_fallback(void *s)
 {

diff --git a/src/codegen.cpp b/src/codegen.cpp
@@ -4015,6 +4015,8 @@ static jl_cgval_t emit_invoke(jl_codectx_t &ctx, const jl_cgval_t &lival, const
                     std::string name;
                     StringRef protoname;
                     bool need_to_emit = true;
+                    // TODO: We should check if the code is available externally
+                    //       and then emit a trampoline.
                     if (ctx.use_cache) {
                         // optimization: emit the correct name immediately, if we know it
                         // TODO: use `emitted` map here too to try to consolidate names?
@@ -6765,6 +6767,8 @@ static jl_llvm_functions_t
         funcName << "japi3_";
     else
         funcName << "japi1_";
+    if (jl_precompile_toplevel_module != NULL)
+        funcName << jl_precompile_toplevel_module->build_id << "_";
     const char* unadorned_name = ctx.name;
 #if defined(_OS_LINUX_)
     if (unadorned_name[0] == '@')
@@ -6812,7 +6816,11 @@ static jl_llvm_functions_t
         }();
 
         std::string wrapName;
-        raw_string_ostream(wrapName) << "jfptr_" << unadorned_name << "_" << globalUniqueGeneratedNames++;
+        raw_string_ostream wrapNameStream(wrapName);
+        wrapNameStream << "jfptr_" << unadorned_name;
+        if (jl_precompile_toplevel_module)
+            wrapNameStream << "_" << jl_precompile_toplevel_module->build_id;
+        wrapNameStream << "_"  << globalUniqueGeneratedNames++;
         declarations.functionObject = wrapName;
         (void)gen_invoke_wrapper(lam, jlrettype, returninfo, retarg, declarations.functionObject, M, ctx.emission_context);
         // TODO: add attributes: maybe_mark_argument_dereferenceable(Arg, argType)
@@ -8284,6 +8292,10 @@ void jl_compile_workqueue(
         StringRef preal_decl = "";
         bool preal_specsig = false;
         auto invoke = jl_atomic_load_relaxed(&codeinst->invoke);
+        // TODO: available_extern
+        // We need to emit a trampoline that loads the target address in an extern_module from a GV
+        // Right now we will unecessarily emit a function we have already compiled in a native module
+        // again in a calling module.
         if (params.cache && invoke != NULL) {
             auto fptr = jl_atomic_load_relaxed(&codeinst->specptr.fptr);
             if (invoke == jl_fptr_args_addr) {

diff --git a/src/datatype.c b/src/datatype.c
@@ -104,6 +104,7 @@ jl_datatype_t *jl_new_uninitialized_datatype(void)
     t->zeroinit = 0;
     t->has_concrete_subtype = 1;
     t->cached_by_hash = 0;
+    t->imgcache = 0;
     t->name = NULL;
     t->super = NULL;
     t->parameters = NULL;

diff --git a/src/dlload.c b/src/dlload.c
@@ -73,7 +73,7 @@ const char *jl_crtdll_name = CRTDLL_BASENAME ".dll";
 #define JL_RTLD(flags, FLAG) (flags & JL_RTLD_ ## FLAG ? RTLD_ ## FLAG : 0)
 
 #ifdef _OS_WINDOWS_
-static void win32_formatmessage(DWORD code, char *reason, int len) JL_NOTSAFEPOINT
+void win32_formatmessage(DWORD code, char *reason, int len) JL_NOTSAFEPOINT
 {
     DWORD res;
     LPWSTR errmsg;

diff --git a/src/dump.c b/src/dump.c
@@ -150,7 +150,7 @@ static arraylist_t reinit_list;
 // list of modules being serialized
 // This is not quite globally rooted, but we take care to only
 // ever assigned rooted values here.
-static jl_array_t *serializer_worklist JL_GLOBALLY_ROOTED;
+jl_array_t *serializer_worklist JL_GLOBALLY_ROOTED;
 // The set of external MethodInstances we want to serialize
 // (methods owned by other modules that were first inferred for a
 //  module currently being serialized)
@@ -237,7 +237,7 @@ static void jl_serialize_cnull(jl_serializer_state *s, jl_value_t *t)
     jl_serialize_value(s, t);
 }
 
-static int module_in_worklist(jl_module_t *mod) JL_NOTSAFEPOINT
+int module_in_worklist(jl_module_t *mod) JL_NOTSAFEPOINT
 {
     int i, l = jl_array_len(serializer_worklist);
     for (i = 0; i < l; i++) {
@@ -933,7 +933,7 @@ static void jl_serialize_value_(jl_serializer_state *s, jl_value_t *v, int as_li
                     rletable = (uint64_t*)jl_array_data(m->root_blocks);
                     nblocks2 = jl_array_len(m->root_blocks);
                 }
-                // this visits every item, if it becomes a bottlneck we could hop blocks
+                // this visits every item, if it becomes a bottleneck we could hop blocks
                 while (rle_iter_increment(&rootiter, nroots, rletable, nblocks2))
                     if (rootiter.key == key)
                         jl_serialize_value(s, jl_array_ptr_ref(m->roots, rootiter.i));

diff --git a/src/gc.c b/src/gc.c
@@ -2884,6 +2884,8 @@ static void mark_roots(jl_gc_mark_cache_t *gc_cache, jl_gc_mark_sp_t *sp)
         gc_mark_queue_obj(gc_cache, sp, jl_all_methods);
     if (_jl_debug_method_invalidation != NULL)
         gc_mark_queue_obj(gc_cache, sp, _jl_debug_method_invalidation);
+    if (jl_build_ids != NULL)
+        gc_mark_queue_obj(gc_cache, sp, jl_build_ids);
 
     // constants
     gc_mark_queue_obj(gc_cache, sp, jl_emptytuple_type);

diff --git a/src/gf.c b/src/gf.c
@@ -461,7 +461,7 @@ static int get_method_unspec_list(jl_typemap_entry_t *def, void *closure)
     return 1;
 }
 
-static int foreach_mtable_in_module(
+int foreach_mtable_in_module(
         jl_module_t *m,
         int (*visit)(jl_methtable_t *mt, void *env),
         void *env)

diff --git a/src/init.c b/src/init.c
@@ -719,6 +719,9 @@ JL_DLLEXPORT void julia_init(JL_IMAGE_SEARCH rel)
         jl_install_default_signal_handlers();
 
     jl_gc_init();
+
+    arraylist_new(&jl_linkage_blobs, 0);
+
     jl_ptls_t ptls = jl_init_threadtls(0);
 #pragma GCC diagnostic push
 #if defined(_COMPILER_GCC_) && __GNUC__ >= 12

diff --git a/src/jitlayers.h b/src/jitlayers.h
@@ -220,7 +220,6 @@ jl_llvm_functions_t jl_emit_codeinst(
 enum CompilationPolicy {
     Default = 0,
     Extern = 1,
-    ImagingMode = 2
 };
 
 typedef std::map<jl_code_instance_t*, std::pair<orc::ThreadSafeModule, jl_llvm_functions_t>> jl_workqueue_t;

diff --git a/src/jl_exported_funcs.inc b/src/jl_exported_funcs.inc
@@ -401,6 +401,7 @@
     XX(jl_restore_excstack) \
     XX(jl_restore_incremental) \
     XX(jl_restore_incremental_from_buf) \
+    XX(jl_restore_package_image_from_file) \
     XX(jl_restore_system_image) \
     XX(jl_restore_system_image_data) \
     XX(jl_rethrow) \
@@ -521,6 +522,7 @@
     XX(jl_vexceptionf) \
     XX(jl_vprintf) \
     XX(jl_wakeup_thread) \
+    XX(jl_write_compiler_output) \
     XX(jl_yield) \
 
 #define JL_RUNTIME_EXPORTED_FUNCS_WIN(XX) \
@@ -537,6 +539,7 @@
     YY(jl_get_LLVM_VERSION) \
     YY(jl_dump_native) \
     YY(jl_get_llvm_gv) \
+    YY(jl_iterate_llvm_gv) \
     YY(jl_dump_function_asm) \
     YY(jl_LLVMCreateDisasm) \
     YY(jl_LLVMDisasmInstruction) \

diff --git a/src/julia.h b/src/julia.h
@@ -552,6 +552,7 @@ typedef struct _jl_datatype_t {
     uint8_t zeroinit:1; // if one or more fields requires zero-initialization
     uint8_t has_concrete_subtype:1; // If clear, no value will have this datatype
     uint8_t cached_by_hash:1; // stored in hash-based set cache (instead of linear cache)
+    uint8_t imgcache:1; // true if this type was loaded from sysimg/pkgimg
 } jl_datatype_t;
 
 typedef struct _jl_vararg_t {
@@ -1759,15 +1760,18 @@ JL_DLLEXPORT const char *jl_pathname_for_handle(void *handle);
 JL_DLLEXPORT int jl_deserialize_verify_header(ios_t *s);
 JL_DLLEXPORT void jl_preload_sysimg_so(const char *fname);
 JL_DLLEXPORT void jl_set_sysimg_so(void *handle);
-JL_DLLEXPORT ios_t *jl_create_system_image(void *);
+JL_DLLEXPORT ios_t *jl_create_system_image(void *, jl_array_t *worklist);
 JL_DLLEXPORT void jl_save_system_image(const char *fname);
 JL_DLLEXPORT void jl_restore_system_image(const char *fname);
-JL_DLLEXPORT void jl_restore_system_image_data(const char *buf, size_t len);
+JL_DLLEXPORT jl_value_t *jl_restore_system_image_data(const char *buf, size_t len);
+JL_DLLEXPORT jl_value_t *jl_restore_package_image_from_file(const char *fname);
 JL_DLLEXPORT void jl_set_newly_inferred(jl_value_t *newly_inferred);
 JL_DLLEXPORT int jl_save_incremental(const char *fname, jl_array_t *worklist);
 JL_DLLEXPORT jl_value_t *jl_restore_incremental(const char *fname, jl_array_t *depmods);
 JL_DLLEXPORT jl_value_t *jl_restore_incremental_from_buf(const char *buf, size_t sz, jl_array_t *depmods);
 
+JL_DLLEXPORT void jl_write_compiler_output(void);
+
 // parsing
 JL_DLLEXPORT jl_value_t *jl_parse_all(const char *text, size_t text_len,
                                       const char *filename, size_t filename_len, size_t lineno);

diff --git a/src/julia_internal.h b/src/julia_internal.h
@@ -290,6 +290,9 @@ void jl_call_tracer(tracer_cb callback, jl_value_t *tracee);
 void print_func_loc(JL_STREAM *s, jl_method_t *m);
 extern jl_array_t *_jl_debug_method_invalidation JL_GLOBALLY_ROOTED;
 void invalidate_backedges(void (*f)(jl_code_instance_t*), jl_method_instance_t *replaced_mi, size_t max_world, const char *why);
+extern arraylist_t jl_linkage_blobs;                        // external linkage: sysimg/pkgimages
+extern jl_array_t *jl_build_ids JL_GLOBALLY_ROOTED;         // external linkage: corresponding build_ids
+extern uint64_t jl_worklist_key(jl_array_t *worklist);
 
 extern JL_DLLEXPORT size_t jl_page_size;
 extern jl_function_t *jl_typeinf_func;
@@ -617,6 +620,9 @@ int nroots_with_key(jl_method_t *m, uint64_t key);
 
 int jl_valid_type_param(jl_value_t *v);
 
+extern jl_array_t *serializer_worklist;
+int module_in_worklist(jl_module_t *mod);
+
 JL_DLLEXPORT jl_value_t *jl_apply_2va(jl_value_t *f, jl_value_t **args, uint32_t nargs);
 
 void JL_NORETURN jl_method_error(jl_function_t *f, jl_value_t **args, size_t na, size_t world);
@@ -693,6 +699,7 @@ jl_expr_t *jl_exprn(jl_sym_t *head, size_t n);
 jl_function_t *jl_new_generic_function(jl_sym_t *name, jl_module_t *module);
 jl_function_t *jl_new_generic_function_with_supertype(jl_sym_t *name, jl_module_t *module, jl_datatype_t *st);
 int jl_foreach_reachable_mtable(int (*visit)(jl_methtable_t *mt, void *env), void *env);
+int foreach_mtable_in_module(jl_module_t *m, int (*visit)(jl_methtable_t *mt, void *env), void *env);
 void jl_init_main_module(void);
 JL_DLLEXPORT int jl_is_submodule(jl_module_t *child, jl_module_t *parent) JL_NOTSAFEPOINT;
 jl_array_t *jl_get_loaded_modules(void);
@@ -931,11 +938,12 @@ JL_DLLEXPORT jl_value_t *jl_dump_fptr_asm(uint64_t fptr, char raw_mc, const char
 JL_DLLEXPORT jl_value_t *jl_dump_function_ir(jl_llvmf_dump_t *dump, char strip_ir_metadata, char dump_module, const char *debuginfo);
 JL_DLLEXPORT jl_value_t *jl_dump_function_asm(jl_llvmf_dump_t *dump, char raw_mc, const char* asm_variant, const char *debuginfo, char binary);
 
-void *jl_create_native(jl_array_t *methods, LLVMOrcThreadSafeModuleRef llvmmod, const jl_cgparams_t *cgparams, int policy);
+void *jl_create_native(jl_array_t *methods, LLVMOrcThreadSafeModuleRef llvmmod, const jl_cgparams_t *cgparams, int policy, int imaging_mode);
 void jl_dump_native(void *native_code,
         const char *bc_fname, const char *unopt_bc_fname, const char *obj_fname, const char *asm_fname,
         const char *sysimg_data, size_t sysimg_len);
 int32_t jl_get_llvm_gv(void *native_code, jl_value_t *p) JL_NOTSAFEPOINT;
+void jl_iterate_llvm_gv(void *native_code, void (*callback)(void*, int32_t, void*), void*) JL_NOTSAFEPOINT;
 JL_DLLEXPORT void jl_get_function_id(void *native_code, jl_code_instance_t *ncode,
         int32_t *func_idx, int32_t *specfunc_idx);
 
@@ -1220,6 +1228,7 @@ extern void *jl_ntdll_handle;
 extern void *jl_kernel32_handle;
 extern void *jl_crtdll_handle;
 extern void *jl_winsock_handle;
+void win32_formatmessage(DWORD code, char *reason, int len) JL_NOTSAFEPOINT;
 #endif
 
 JL_DLLEXPORT void *jl_get_library_(const char *f_lib, int throw_err);
@@ -1565,7 +1574,6 @@ void jl_register_fptrs(uint64_t sysimage_base, const struct _jl_sysimg_fptrs_t *
                        jl_method_instance_t **linfos, size_t n);
 void jl_write_coverage_data(const char*);
 void jl_write_malloc_log(void);
-void jl_write_compiler_output(void);
 
 #if jl_has_builtin(__builtin_unreachable) || defined(_COMPILER_GCC_) || defined(_COMPILER_INTEL_)
 #  define jl_unreachable() __builtin_unreachable()

diff --git a/src/llvm-multiversioning.cpp b/src/llvm-multiversioning.cpp
@@ -304,23 +304,31 @@ static inline std::vector<T*> consume_gv(Module &M, const char *name, bool allow
     // Strip them from the Module so that it's easier to handle the uses.
     GlobalVariable *gv = M.getGlobalVariable(name);
     assert(gv && gv->hasInitializer());
-    auto *ary = cast<ConstantArray>(gv->getInitializer());
-    unsigned nele = ary->getNumOperands();
+    ArrayType *Ty = cast<ArrayType>(gv->getInitializer()->getType());
+    unsigned nele = Ty->getArrayNumElements();
     std::vector<T*> res(nele);
-    unsigned i = 0;
-    while (i < nele) {
-        llvm::Value *val = ary->getOperand(i)->stripPointerCasts();
-        if (allow_bad_fvars && (!isa<T>(val) || (isa<Function>(val) && cast<Function>(val)->isDeclaration()))) {
-            // Shouldn't happen in regular use, but can happen in bugpoint.
-            nele--;
-            continue;
+    ConstantArray *ary = nullptr;
+    if (gv->getInitializer()->isNullValue()) {
+        for (unsigned i = 0; i < nele; ++i)
+            res[i] = cast<T>(Constant::getNullValue(Ty->getArrayElementType()));
+    }
+    else {
+        ary = cast<ConstantArray>(gv->getInitializer());
+        unsigned i = 0;
+        while (i < nele) {
+            llvm::Value *val = ary->getOperand(i)->stripPointerCasts();
+            if (allow_bad_fvars && (!isa<T>(val) || (isa<Function>(val) && cast<Function>(val)->isDeclaration()))) {
+                // Shouldn't happen in regular use, but can happen in bugpoint.
+                nele--;
+                continue;
+            }
+            res[i++] = cast<T>(val);
         }
-        res[i++] = cast<T>(val);
+        res.resize(nele);
     }
-    res.resize(nele);
     assert(gv->use_empty());
     gv->eraseFromParent();
-    if (ary->use_empty())
+    if (ary && ary->use_empty())
         ary->destroyConstant();
     return res;
 }
@@ -925,12 +933,17 @@ Constant *CloneCtx::emit_offset_table(const std::vector<T*> &vars, StringRef nam
 {
     auto T_int32 = Type::getInt32Ty(M.getContext());
     auto T_size = getSizeTy(M.getContext());
-    assert(!vars.empty());
-    add_comdat(GlobalAlias::create(T_size, 0, GlobalVariable::ExternalLinkage,
-                                   name + "_base",
-                                   ConstantExpr::getBitCast(vars[0], T_size->getPointerTo()), &M));
-    auto vbase = ConstantExpr::getPtrToInt(vars[0], T_size);
     uint32_t nvars = vars.size();
+    Constant *base = nullptr;
+    if (nvars > 0) {
+        base = ConstantExpr::getBitCast(vars[0], T_size->getPointerTo());
+        add_comdat(GlobalAlias::create(T_size, 0, GlobalVariable::ExternalLinkage,
+                                       name + "_base",
+                                       base, &M));
+    } else {
+        base = ConstantExpr::getNullValue(T_size->getPointerTo());
+    }
+    auto vbase = ConstantExpr::getPtrToInt(base, T_size);
     std::vector<Constant*> offsets(nvars + 1);
     offsets[0] = ConstantInt::get(T_int32, nvars);
     offsets[1] = ConstantInt::get(T_int32, 0);