staticdata: Refactor sysimage loading (#57542)

Introduce `jl_image_buf_t` to represent the in-memory details of an
unparsed system image and clean-up several global variables and split
loading logic in staticdata.c so that we get (almost) everything we need
from the sysimage handle at once.

This allows sysimage loading to be separated into three phases:
  1. Lookup the raw sysimage buffer as a `jl_image_buf_t`
2. For .so sysimages, parse the sysimage and initialize JIT targets,
etc.
  3. Finally load the sysimage into a `jl_image_t`

Care was taken to preserve the existing behavior of calling
`jl_set_sysimg_so` to configure the sysimage before initializing Julia,
although this is likely to be next on the chopping block in a follow-up.

Dependent on #57523.

Author: gbaraldi

src/aotcompile.cpp

@@ -2201,7 +2201,7 @@ void jl_dump_native_impl(void *native_code,
             builder.CreateRet(ConstantInt::get(T_int32, 1));
         }
         if (imaging_mode) {
-            auto specs = jl_get_llvm_clone_targets();
+            auto specs = jl_get_llvm_clone_targets(jl_options.cpu_target);
             const uint32_t base_flags = has_veccall ? JL_TARGET_VEC_CALL : 0;
             SmallVector<uint8_t, 0> data;
             auto push_i32 = [&] (uint32_t v) {

src/init.c

@@ -868,21 +868,33 @@ static NOINLINE void _finish_julia_init(JL_IMAGE_SEARCH rel, jl_ptls_t ptls, jl_
 {
     JL_TIMING(JULIA_INIT, JULIA_INIT);
     jl_resolve_sysimg_location(rel);
+
     // loads sysimg if available, and conditionally sets jl_options.cpu_target
+    jl_image_buf_t sysimage = { JL_IMAGE_KIND_NONE };
     if (rel == JL_IMAGE_IN_MEMORY) {
-        jl_set_sysimg_so(jl_exe_handle);
+        sysimage = jl_set_sysimg_so(jl_exe_handle);
         jl_options.image_file = jl_options.julia_bin;
     }
     else if (jl_options.image_file)
-        jl_preload_sysimg_so(jl_options.image_file);
+        sysimage = jl_preload_sysimg(jl_options.image_file);
+
+    if (sysimage.kind == JL_IMAGE_KIND_SO)
+        jl_gc_notify_image_load(sysimage.data, sysimage.size);
+
     if (jl_options.cpu_target == NULL)
         jl_options.cpu_target = "native";
-    jl_init_codegen();
 
+    // Parse image, perform relocations, and init JIT targets, etc.
+    jl_image_t parsed_image = jl_init_processor_sysimg(sysimage, jl_options.cpu_target);
+
+    jl_init_codegen();
     jl_init_common_symbols();
-    if (jl_options.image_file) {
-        jl_restore_system_image(jl_options.image_file);
+
+    if (sysimage.kind != JL_IMAGE_KIND_NONE) {
+        // Load the .ji or .so sysimage
+        jl_restore_system_image(&parsed_image, sysimage);
     } else {
+        // No sysimage provided, init a minimal environment
         jl_init_types();
         jl_global_roots_list = (jl_genericmemory_t*)jl_an_empty_memory_any;
         jl_global_roots_keyset = (jl_genericmemory_t*)jl_an_empty_memory_any;
@@ -891,7 +903,7 @@ static NOINLINE void _finish_julia_init(JL_IMAGE_SEARCH rel, jl_ptls_t ptls, jl_
     jl_init_flisp();
     jl_init_serializer();
 
-    if (!jl_options.image_file) {
+    if (sysimage.kind == JL_IMAGE_KIND_NONE) {
         jl_top_module = jl_core_module;
         jl_init_intrinsic_functions();
         jl_init_primitives();
@@ -919,7 +931,8 @@ static NOINLINE void _finish_julia_init(JL_IMAGE_SEARCH rel, jl_ptls_t ptls, jl_
 
     jl_gc_enable(1);
 
-    if (jl_options.image_file && (!jl_generating_output() || jl_options.incremental) && jl_module_init_order) {
+    if ((sysimage.kind != JL_IMAGE_KIND_NONE) &&
+            (!jl_generating_output() || jl_options.incremental) && jl_module_init_order) {
         jl_array_t *init_order = jl_module_init_order;
         JL_GC_PUSH1(&init_order);
         jl_module_init_order = NULL;

src/jitlayers.cpp

@@ -1281,7 +1281,7 @@ namespace {
 #endif
 #endif
         uint32_t target_flags = 0;
-        auto target = jl_get_llvm_target(jl_generating_output(), target_flags);
+        auto target = jl_get_llvm_target(jl_options.cpu_target, jl_generating_output(), target_flags);
         auto &TheCPU = target.first;
         SmallVector<std::string, 10> targetFeatures(target.second.begin(), target.second.end());
         std::string errorstr;

src/jl_exported_funcs.inc

@@ -358,7 +358,7 @@
     XX(jl_pointerset) \
     XX(jl_pop_handler) \
     XX(jl_pop_handler_noexcept) \
-    XX(jl_preload_sysimg_so) \
+    XX(jl_preload_sysimg) \
     XX(jl_prepend_cwd) \
     XX(jl_printf) \
     XX(jl_print_backtrace) \
@@ -388,7 +388,6 @@
     XX(jl_restore_incremental) \
     XX(jl_restore_package_image_from_file) \
     XX(jl_restore_system_image) \
-    XX(jl_restore_system_image_data) \
     XX(jl_rethrow) \
     XX(jl_rethrow_other) \
     XX(jl_running_on_valgrind) \

src/julia.h

@@ -2219,6 +2219,25 @@ typedef enum {
     JL_IMAGE_IN_MEMORY = 2
 } JL_IMAGE_SEARCH;
 
+typedef enum {
+    JL_IMAGE_KIND_NONE = 0,
+    JL_IMAGE_KIND_JI,
+    JL_IMAGE_KIND_SO,
+} jl_image_kind_t;
+
+// A loaded, but unparsed .ji or .so image file
+typedef struct {
+    jl_image_kind_t kind;
+    void *handle;
+    const void *pointers; // jl_image_pointers_t *
+    const char *data;
+    size_t size;
+    uint64_t base;
+} jl_image_buf_t;
+
+struct _jl_image_t;
+typedef struct _jl_image_t jl_image_t;
+
 JL_DLLIMPORT const char *jl_get_libdir(void);
 JL_DLLEXPORT void julia_init(JL_IMAGE_SEARCH rel);
 JL_DLLEXPORT void jl_init(void);
@@ -2235,11 +2254,10 @@ JL_DLLEXPORT const char *jl_pathname_for_handle(void *handle);
 JL_DLLEXPORT jl_gcframe_t **jl_adopt_thread(void);
 
 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 jl_image_buf_t jl_preload_sysimg(const char *fname);
+JL_DLLEXPORT jl_image_buf_t jl_set_sysimg_so(void *handle);
 JL_DLLEXPORT void jl_create_system_image(void **, jl_array_t *worklist, bool_t emit_split, ios_t **s, ios_t **z, jl_array_t **udeps, int64_t *srctextpos);
-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 void jl_restore_system_image(jl_image_t *image, jl_image_buf_t buf);
 JL_DLLEXPORT jl_value_t *jl_restore_incremental(const char *fname, jl_array_t *depmods, int complete, const char *pkgimage);
 JL_DLLEXPORT jl_value_t *jl_object_top_module(jl_value_t* v) JL_NOTSAFEPOINT;
 

src/llvm-multiversioning.cpp

@@ -216,7 +216,7 @@ static void annotate_module_clones(Module &M) {
     if (auto maybe_specs = get_target_specs(M)) {
         specs = std::move(*maybe_specs);
     } else {
-        auto full_specs = jl_get_llvm_clone_targets();
+        auto full_specs = jl_get_llvm_clone_targets(jl_options.cpu_target);
         specs.reserve(full_specs.size());
         for (auto &spec: full_specs) {
             specs.push_back(TargetSpec::fromSpec(spec));

src/processor.cpp

@@ -504,7 +504,8 @@ static inline llvm::SmallVector<TargetData<n>, 0>
 parse_cmdline(const char *option, F &&feature_cb)
 {
     if (!option)
-        option = "native";
+        abort();
+
     llvm::SmallVector<TargetData<n>, 0> res;
     TargetData<n> arg{};
     auto reset_arg = [&] {
@@ -612,31 +613,29 @@ parse_cmdline(const char *option, F &&feature_cb)
 
 // Cached version of command line parsing
 template<size_t n, typename F>
-static inline llvm::SmallVector<TargetData<n>, 0> &get_cmdline_targets(F &&feature_cb)
+static inline llvm::SmallVector<TargetData<n>, 0> &get_cmdline_targets(const char *cpu_target, F &&feature_cb)
 {
     static llvm::SmallVector<TargetData<n>, 0> targets =
-        parse_cmdline<n>(jl_options.cpu_target, std::forward<F>(feature_cb));
+        parse_cmdline<n>(cpu_target, std::forward<F>(feature_cb));
     return targets;
 }
 
 // Load sysimg, use the `callback` for dispatch and perform all relocations
 // for the selected target.
 template<typename F>
-static inline jl_image_t parse_sysimg(void *hdl, F &&callback)
+static inline jl_image_t parse_sysimg(jl_image_buf_t image, F &&callback, void *ctx)
 {
     JL_TIMING(LOAD_IMAGE, LOAD_Processor);
     jl_image_t res{};
 
-    const jl_image_pointers_t *pointers;
-    if (jl_system_image_size == 0)
-        jl_dlsym(hdl, "jl_image_pointers", (void**)&pointers, 1);
-    else
-        pointers = &jl_image_pointers; // libjulia-internal and sysimage statically linked
+    if (image.kind != JL_IMAGE_KIND_SO)
+        return res;
 
+    const jl_image_pointers_t *pointers = (const jl_image_pointers_t *)image.pointers;
     const void *ids = pointers->target_data;
     jl_value_t* rejection_reason = nullptr;
     JL_GC_PUSH1(&rejection_reason);
-    uint32_t target_idx = callback(ids, &rejection_reason);
+    uint32_t target_idx = callback(ctx, ids, &rejection_reason);
     if (target_idx == UINT32_MAX) {
         jl_error(jl_string_ptr(rejection_reason));
     }
@@ -794,17 +793,7 @@ static inline jl_image_t parse_sysimg(void *hdl, F &&callback)
         res.fptrs.nclones = clones.size();
     }
 
-#ifdef _OS_WINDOWS_
-    res.base = (intptr_t)hdl;
-#else
-    Dl_info dlinfo;
-    if (dladdr((void*)pointers, &dlinfo) != 0) {
-        res.base = (intptr_t)dlinfo.dli_fbase;
-    }
-    else {
-        res.base = 0;
-    }
-#endif
+    res.base = image.base;
 
     {
         void *pgcstack_func_slot = pointers->ptls->pgcstack_func_slot;
@@ -1024,7 +1013,7 @@ JL_DLLEXPORT jl_value_t *jl_get_cpu_features(void)
 }
 
 extern "C" JL_DLLEXPORT jl_value_t* jl_reflect_clone_targets() {
-    auto specs = jl_get_llvm_clone_targets();
+    auto specs = jl_get_llvm_clone_targets(jl_options.cpu_target);
     const uint32_t base_flags = 0;
     llvm::SmallVector<uint8_t, 0> data;
     auto push_i32 = [&] (uint32_t v) {

src/processor.h

@@ -64,6 +64,7 @@ JL_DLLEXPORT int jl_test_cpu_feature(jl_cpu_feature_t feature);
 static const uint32_t jl_sysimg_tag_mask = 0x80000000u;
 static const uint32_t jl_sysimg_val_mask = ~((uint32_t)0x80000000u);
 
+// A parsed image file
 typedef struct _jl_image_fptrs_t {
     // number of functions
     uint32_t nptrs;
@@ -82,14 +83,14 @@ typedef struct _jl_image_fptrs_t {
     const uint32_t *clone_idxs;
 } jl_image_fptrs_t;
 
-typedef struct {
+struct _jl_image_t {
     uint64_t base;
     const char *gvars_base;
     const int32_t *gvars_offsets;
     uint32_t ngvars;
     jl_image_fptrs_t fptrs;
     void **jl_small_typeof;
-} jl_image_t;
+};
 
 // The header for each image
 // Details important counts about the image
@@ -206,8 +207,8 @@ typedef struct {
  *
  * Return the data about the function pointers selected.
  */
-jl_image_t jl_init_processor_sysimg(void *hdl);
-jl_image_t jl_init_processor_pkgimg(void *hdl);
+jl_image_t jl_init_processor_sysimg(jl_image_buf_t image, const char *cpu_target);
+jl_image_t jl_init_processor_pkgimg(jl_image_buf_t image);
 
 // Return the name of the host CPU as a julia string.
 JL_DLLEXPORT jl_value_t *jl_get_cpu_name(void);
@@ -251,7 +252,7 @@ extern JL_DLLEXPORT bool jl_processor_print_help;
  * If the detected/specified CPU name is not available on the LLVM version specified,
  * a fallback CPU name will be used. Unsupported features will be ignored.
  */
-extern "C" JL_DLLEXPORT std::pair<std::string,llvm::SmallVector<std::string, 0>> jl_get_llvm_target(bool imaging, uint32_t &flags) JL_NOTSAFEPOINT;
+extern "C" JL_DLLEXPORT std::pair<std::string,llvm::SmallVector<std::string, 0>> jl_get_llvm_target(const char *cpu_target, bool imaging, uint32_t &flags) JL_NOTSAFEPOINT;
 
 /**
  * Returns the CPU name and feature string to be used by LLVM disassembler.
@@ -275,7 +276,7 @@ struct jl_target_spec_t {
 /**
  * Return the list of targets to clone
  */
-extern "C" JL_DLLEXPORT llvm::SmallVector<jl_target_spec_t, 0> jl_get_llvm_clone_targets(void) JL_NOTSAFEPOINT;
+extern "C" JL_DLLEXPORT llvm::SmallVector<jl_target_spec_t, 0> jl_get_llvm_clone_targets(const char *cpu_target) JL_NOTSAFEPOINT;
 // NOLINTEND(clang-diagnostic-return-type-c-linkage)
 struct FeatureName {
     const char *name;

src/processor_arm.cpp

@@ -1519,7 +1519,7 @@ static inline void disable_depends(FeatureList<n> &features)
     ::disable_depends(features, Feature::deps, sizeof(Feature::deps) / sizeof(FeatureDep));
 }
 
-static const llvm::SmallVector<TargetData<feature_sz>, 0> &get_cmdline_targets(void)
+static const llvm::SmallVector<TargetData<feature_sz>, 0> &get_cmdline_targets(const char *cpu_target)
 {
     auto feature_cb = [] (const char *str, size_t len, FeatureList<feature_sz> &list) {
 #ifdef _CPU_AARCH64_
@@ -1536,7 +1536,7 @@ static const llvm::SmallVector<TargetData<feature_sz>, 0> &get_cmdline_targets(v
         set_bit(list, fbit, true);
         return true;
     };
-    auto &targets = ::get_cmdline_targets<feature_sz>(feature_cb);
+    auto &targets = ::get_cmdline_targets<feature_sz>(cpu_target, feature_cb);
     for (auto &t: targets) {
         if (auto nname = normalize_cpu_name(t.name)) {
             t.name = nname;
@@ -1599,10 +1599,11 @@ static int max_vector_size(const FeatureList<feature_sz> &features)
 #endif
 }
 
-static uint32_t sysimg_init_cb(const void *id, jl_value_t **rejection_reason)
+static uint32_t sysimg_init_cb(void *ctx, const void *id, jl_value_t **rejection_reason)
 {
     // First see what target is requested for the JIT.
-    auto &cmdline = get_cmdline_targets();
+    const char *cpu_target = (const char *)ctx;
+    auto &cmdline = get_cmdline_targets(cpu_target);
     TargetData<feature_sz> target = arg_target_data(cmdline[0], true);
     // Then find the best match in the sysimg
     auto sysimg = deserialize_target_data<feature_sz>((const uint8_t*)id);
@@ -1626,7 +1627,7 @@ static uint32_t sysimg_init_cb(const void *id, jl_value_t **rejection_reason)
     return match.best_idx;
 }
 
-static uint32_t pkgimg_init_cb(const void *id, jl_value_t **rejection_reason JL_REQUIRE_ROOTED_SLOT)
+static uint32_t pkgimg_init_cb(void *ctx, const void *id, jl_value_t **rejection_reason JL_REQUIRE_ROOTED_SLOT)
 {
     TargetData<feature_sz> target = jit_targets.front();
     auto pkgimg = deserialize_target_data<feature_sz>((const uint8_t*)id);
@@ -1639,9 +1640,9 @@ static uint32_t pkgimg_init_cb(const void *id, jl_value_t **rejection_reason JL_
     return match.best_idx;
 }
 
-static void ensure_jit_target(bool imaging)
+static void ensure_jit_target(const char *cpu_target, bool imaging)
 {
-    auto &cmdline = get_cmdline_targets();
+    auto &cmdline = get_cmdline_targets(cpu_target);
     check_cmdline(cmdline, imaging);
     if (!jit_targets.empty())
         return;
@@ -1852,36 +1853,36 @@ JL_DLLEXPORT jl_value_t *jl_cpu_has_fma(int bits)
 #endif
 }
 
-jl_image_t jl_init_processor_sysimg(void *hdl)
+jl_image_t jl_init_processor_sysimg(jl_image_buf_t image, const char *cpu_target)
 {
     if (!jit_targets.empty())
         jl_error("JIT targets already initialized");
-    return parse_sysimg(hdl, sysimg_init_cb);
+    return parse_sysimg(image, sysimg_init_cb, (void *)cpu_target);
 }
 
-jl_image_t jl_init_processor_pkgimg(void *hdl)
+jl_image_t jl_init_processor_pkgimg(jl_image_buf_t image)
 {
     if (jit_targets.empty())
         jl_error("JIT targets not initialized");
     if (jit_targets.size() > 1)
         jl_error("Expected only one JIT target");
-    return parse_sysimg(hdl, pkgimg_init_cb);
+    return parse_sysimg(image, pkgimg_init_cb, NULL);
 }
 
 JL_DLLEXPORT jl_value_t* jl_check_pkgimage_clones(char *data)
 {
     jl_value_t *rejection_reason = NULL;
     JL_GC_PUSH1(&rejection_reason);
-    uint32_t match_idx = pkgimg_init_cb(data, &rejection_reason);
+    uint32_t match_idx = pkgimg_init_cb(NULL, data, &rejection_reason);
     JL_GC_POP();
     if (match_idx == UINT32_MAX)
         return rejection_reason;
     return jl_nothing;
 }
 
-std::pair<std::string,llvm::SmallVector<std::string, 0>> jl_get_llvm_target(bool imaging, uint32_t &flags)
+std::pair<std::string,llvm::SmallVector<std::string, 0>> jl_get_llvm_target(const char *cpu_target, bool imaging, uint32_t &flags)
 {
-    ensure_jit_target(imaging);
+    ensure_jit_target(cpu_target, imaging);
     flags = jit_targets[0].en.flags;
     return get_llvm_target_vec(jit_targets[0]);
 }
@@ -1900,10 +1901,10 @@ const std::pair<std::string,std::string> &jl_get_llvm_disasm_target(void)
 }
 
 #ifndef __clang_gcanalyzer__
-llvm::SmallVector<jl_target_spec_t, 0> jl_get_llvm_clone_targets(void)
+llvm::SmallVector<jl_target_spec_t, 0> jl_get_llvm_clone_targets(const char *cpu_target)
 {
 
-    auto &cmdline = get_cmdline_targets();
+    auto &cmdline = get_cmdline_targets(cpu_target);
     check_cmdline(cmdline, true);
     llvm::SmallVector<TargetData<feature_sz>, 0> image_targets;
     for (auto &arg: cmdline) {