Draft of julia initialization refactor. This is mostly focused on making embedding julia as a shared library behave better

gbaraldi · gbaraldi · commit 34fe416c5af3 · 2024-10-29T16:31:58.000-03:00
diff --git a/contrib/juliac.jl b/contrib/juliac.jl
@@ -9,6 +9,7 @@ outname = nothing
 file = nothing
 add_ccallables = false
 verbose = false
+output_llvm = nothing
 
 help = findfirst(x->x == "--help", ARGS)
 if help !== nothing
@@ -18,6 +19,7 @@ if help !== nothing
         --trim=<no,safe,unsafe,unsafe-warn>  Only output code statically determined to be reachable
         --compile-ccallable  Include all methods marked `@ccallable` in output
         --verbose            Request verbose output
+        --output-llvm=<opt,unopt>  Output LLVM bitcode
         """)
     exit(0)
 end
@@ -42,6 +44,17 @@ let i = 1
             global add_ccallables = true
         elseif arg == "--verbose"
             global verbose = true
+        elseif startswith(arg, "--output-llvm")
+            arg = split(arg, '=')
+            if length(arg) == 1
+                global output_llvm = "--output-bc"
+            elseif arg[2] == "opt"
+                global output_llvm = "--output-bc"
+            elseif arg[2] == "unopt"
+                global output_llvm = "--output-unopt-bc"
+            else
+                error("Invalid argument to --output-llvm")
+            end
         else
             if arg[1] == '-' || !isnothing(file)
                 println("Unexpected argument `$arg`")
@@ -65,7 +78,6 @@ tmpdir = mktempdir(cleanup=false)
 initsrc_path = joinpath(tmpdir, "init.c")
 init_path = joinpath(tmpdir, "init.a")
 img_path = joinpath(tmpdir, "img.a")
-bc_path = joinpath(tmpdir, "img-bc.a")
 
 open(initsrc_path, "w") do io
     print(io, """
@@ -80,6 +92,17 @@ open(initsrc_path, "w") do io
 end
 
 static_call_graph_arg() = isnothing(trim) ?  `` : `--trim=$(trim)`
+
+if !isnothing(output_llvm)
+    cmd = addenv(`$cmd --project=$(Base.active_project()) $output_llvm $(outname * ".a") --output-incremental=no --strip-ir --strip-metadata $(static_call_graph_arg()) $(joinpath(@__DIR__,"juliac-buildscript.jl")) $absfile $output_type $add_ccallables`, "OPENBLAS_NUM_THREADS" => 1, "JULIA_NUM_THREADS" => 1)
+    verbose && println("Running: $cmd")
+    if !success(pipeline(cmd; stdout, stderr))
+        println(stderr, "\nFailed to compile $file")
+        exit(1)
+    end
+    exit(0)
+end
+
 cmd = addenv(`$cmd --project=$(Base.active_project()) --output-o $img_path --output-incremental=no --strip-ir --strip-metadata $(static_call_graph_arg()) $(joinpath(@__DIR__,"juliac-buildscript.jl")) $absfile $output_type $add_ccallables`, "OPENBLAS_NUM_THREADS" => 1, "JULIA_NUM_THREADS" => 1)
 verbose && println("Running: $cmd")
 if !success(pipeline(cmd; stdout, stderr))
@@ -100,11 +123,11 @@ end
 julia_libs = Base.shell_split(Base.isdebugbuild() ? "-ljulia-debug -ljulia-internal-debug" : "-ljulia -ljulia-internal")
 try
     if output_type == "--output-lib"
-        run(`cc $(allflags) -o $outname -shared -Wl,$(Base.Linking.WHOLE_ARCHIVE) $img_path  -Wl,$(Base.Linking.NO_WHOLE_ARCHIVE) $init_path  $(julia_libs)`)
+        run(`cc $(allflags) -o $outname -shared -Wl,$(Base.Linking.WHOLE_ARCHIVE) $img_path  -Wl,$(Base.Linking.NO_WHOLE_ARCHIVE)  $(julia_libs)`)
     elseif output_type == "--output-sysimage"
         run(`cc $(allflags) -o $outname -shared -Wl,$(Base.Linking.WHOLE_ARCHIVE) $img_path  -Wl,$(Base.Linking.NO_WHOLE_ARCHIVE)             $(julia_libs)`)
     else
-        run(`cc $(allflags) -o $outname -Wl,$(Base.Linking.WHOLE_ARCHIVE) $img_path -Wl,$(Base.Linking.NO_WHOLE_ARCHIVE) $init_path $(julia_libs)`)
+        run(`cc $(allflags) -o $outname -Wl,$(Base.Linking.WHOLE_ARCHIVE) $img_path -Wl,$(Base.Linking.NO_WHOLE_ARCHIVE) $(julia_libs)`)
     end
 catch
     println("\nCompilation failed.")
diff --git a/src/dlload.c b/src/dlload.c
@@ -254,7 +254,10 @@ void *jl_find_dynamic_library_by_addr(void *symbol) {
         jl_error("could not load base module");
     }
     handle = dlopen(info.dli_fname, RTLD_NOW | RTLD_NOLOAD | RTLD_LOCAL);
-    dlclose(handle); // Undo ref count increment from `dlopen`
+    if (handle == NULL && dlerror() == NULL) // We loaded the executable but got RTLD_DEFAULT back, give a real handle instead
+        handle = dlopen("", RTLD_NOW | RTLD_NOLOAD | RTLD_LOCAL);
+    if (handle) // We may get a null handle so don't segfault
+        dlclose(handle); // Undo ref count increment from `dlopen`
 #endif
     return handle;
 }
diff --git a/src/init.c b/src/init.c
@@ -238,7 +238,7 @@ JL_DLLEXPORT void jl_atexit_hook(int exitcode) JL_NOTSAFEPOINT_ENTER
     jl_task_t *ct = jl_get_current_task();
 
     if (ct == NULL && jl_base_module) {
-        ct = container_of(jl_adopt_thread(), jl_task_t, gcstack);
+        ct = container_of(jl_adopt_thread(NULL, NULL), jl_task_t, gcstack);
     }
     else if (ct != NULL) {
         // we are about to start tearing everything down, so lets try not to get
@@ -379,6 +379,8 @@ JL_DLLEXPORT void *jl_libjulia_internal_handle;
 JL_DLLEXPORT void *jl_libjulia_handle;
 JL_DLLEXPORT void *jl_RTLD_DEFAULT_handle;
 JL_DLLEXPORT void *jl_exe_handle;
+JL_DLLEXPORT void *jl_base_image_handle; // Used when the sysimage is bundled as an executable or a standalone dylib
+
 #ifdef _OS_WINDOWS_
 void *jl_ntdll_handle;
 void *jl_kernel32_handle;
@@ -761,7 +763,6 @@ JL_DLLEXPORT void julia_init(JL_IMAGE_SEARCH rel)
     rec_backtrace(NULL, 0, 0);
 #endif
 
-    libsupport_init();
     jl_safepoint_init();
     jl_page_size = jl_getpagesize();
     htable_new(&jl_current_modules, 0);
@@ -843,9 +844,10 @@ 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
+    // If the image is in memory then jl_base_image_handle must be set before calling init
     if (rel == JL_IMAGE_IN_MEMORY) {
-        jl_set_sysimg_so(jl_exe_handle);
-        jl_options.image_file = jl_options.julia_bin;
+        jl_set_sysimg_so(jl_base_image_handle);
+        jl_options.image_file = jl_options.julia_bin; // XXX: Use dladdr on base_image_handle?
     }
     else if (jl_options.image_file)
         jl_preload_sysimg_so(jl_options.image_file);
diff --git a/src/jl_exported_funcs.inc b/src/jl_exported_funcs.inc
@@ -291,6 +291,7 @@
     XX(jl_load_) \
     XX(jl_load_and_lookup) \
     XX(jl_load_dynamic_library) \
+    XX(jl_find_dynamic_library_by_addr) \
     XX(jl_load_file_string) \
     XX(jl_lookup_code_address) \
     XX(jl_lseek) \
diff --git a/src/jlapi.c b/src/jlapi.c
@@ -25,7 +25,15 @@ extern "C" {
 #else
 #include <fenv.h>
 #endif
+_Atomic(int) jl_runtime_is_initialized = 0;
+uv_mutex_t initialization_lock;
 
+// Should be called in an attribute constructor function to make initialization threadsafe when embedding as an SO.
+__attribute__((constructor)) void jl_preinit_runtime(void)
+{
+    libsupport_init();
+    uv_mutex_init(&initialization_lock);
+}
 /**
  * @brief Check if Julia is already initialized.
  *
@@ -36,7 +44,7 @@ extern "C" {
  */
 JL_DLLEXPORT int jl_is_initialized(void)
 {
-    return jl_main_module != NULL;
+    return jl_atomic_load_acquire(&jl_runtime_is_initialized);
 }
 
 /**
@@ -87,7 +95,7 @@ JL_DLLEXPORT void jl_init_with_image(const char *julia_bindir,
 {
     if (jl_is_initialized())
         return;
-    libsupport_init();
+    uv_mutex_lock(&initialization_lock);
     jl_options.julia_bindir = julia_bindir;
     if (image_path != NULL)
         jl_options.image_file = image_path;
@@ -1029,7 +1037,7 @@ JL_DLLEXPORT int jl_repl_entrypoint(int argc, char *argv[])
     // No-op on non-windows
     lock_low32();
 
-    libsupport_init();
+    jl_preinit_runtime();
     int lisp_prompt = (argc >= 2 && strcmp((char*)argv[1],"--lisp") == 0);
     if (lisp_prompt) {
         memmove(&argv[1], &argv[2], (argc-2)*sizeof(void*));
@@ -1061,6 +1069,23 @@ JL_DLLEXPORT int jl_repl_entrypoint(int argc, char *argv[])
     return ret;
 }
 
+int jl_init_runtime_adopt_thread(void* sysimg_handle) {
+    // Check if we are the main thread
+    uv_mutex_lock(&initialization_lock);
+    if (jl_is_initialized()) {
+        uv_mutex_unlock(&initialization_lock); // Someone won the race
+        return 0;
+    }
+    //TODO: Implement option parsing
+    assert(sysimg_handle);
+    // This code path assumes that you are either an executable or dylib with the sysimage linked in
+    jl_base_image_handle = sysimg_handle;
+    julia_init(JL_IMAGE_IN_MEMORY);
+    jl_enter_threaded_region();
+    uv_mutex_unlock(&initialization_lock);
+    return 1;
+}
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/src/julia.h b/src/julia.h
@@ -2120,7 +2120,7 @@ JL_DLLEXPORT void jl_postoutput_hook(void);
 JL_DLLEXPORT void JL_NORETURN jl_exit(int status);
 JL_DLLEXPORT void JL_NORETURN jl_raise(int signo);
 JL_DLLEXPORT const char *jl_pathname_for_handle(void *handle);
-JL_DLLEXPORT jl_gcframe_t **jl_adopt_thread(void);
+JL_DLLEXPORT jl_gcframe_t **jl_adopt_thread(void* sysimage_handle);
 
 JL_DLLEXPORT int jl_deserialize_verify_header(ios_t *s);
 JL_DLLEXPORT void jl_preload_sysimg_so(const char *fname);
diff --git a/src/julia_internal.h b/src/julia_internal.h
@@ -1058,13 +1058,15 @@ void jl_init_serializer(void);
 void jl_init_uv(void);
 void jl_init_int32_int64_cache(void);
 JL_DLLEXPORT void jl_init_options(void);
+int jl_init_runtime_adopt_thread(void* sysimg_handle);
 
 void jl_set_base_ctx(char *__stk);
 
 extern JL_DLLEXPORT ssize_t jl_tls_offset;
 extern JL_DLLEXPORT const int jl_tls_elf_support;
 void jl_init_threading(void);
 void jl_start_threads(void);
+JL_DLLEXPORT void jl_enter_threaded_region(void);
 
 // Whether the GC is running
 extern uv_mutex_t safepoint_lock;
@@ -1454,6 +1456,7 @@ JL_DLLEXPORT void jl_srand(uint64_t) JL_NOTSAFEPOINT;
 JL_DLLEXPORT void jl_init_rand(void);
 
 JL_DLLEXPORT extern void *jl_exe_handle;
+JL_DLLEXPORT extern void *jl_base_image_handle;
 JL_DLLEXPORT extern void *jl_libjulia_handle;
 JL_DLLEXPORT extern void *jl_libjulia_internal_handle;
 JL_DLLEXPORT extern void *jl_RTLD_DEFAULT_handle;
@@ -1468,7 +1471,7 @@ 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);
-void *jl_find_dynamic_library_by_addr(void *symbol);
+JL_DLLEXPORT void *jl_find_dynamic_library_by_addr(void *symbol);
 #define jl_get_library(f_lib) jl_get_library_(f_lib, 1)
 JL_DLLEXPORT void *jl_load_and_lookup(const char *f_lib, const char *f_name, _Atomic(void*) *hnd);
 JL_DLLEXPORT void *jl_lazy_load_and_lookup(jl_value_t *lib_val, const char *f_name);
diff --git a/src/llvm-ptls.cpp b/src/llvm-ptls.cpp
@@ -5,6 +5,9 @@
 #include "support/dtypes.h"
 #include "passes.h"
 
+#include "llvm/IR/BasicBlock.h"
+#include "llvm/IR/Instruction.h"
+#include "llvm/Support/Debug.h"
 #include <llvm-c/Core.h>
 #include <llvm-c/Types.h>
 
@@ -151,7 +154,17 @@ void LowerPTLS::fix_pgcstack_use(CallInst *pgcstack, Function *pgcstack_getter,
         return;
     }
     if (or_new) {
-        // pgcstack();
+        // pgcstack = pgstack_intrinsic()
+
+
+        // if (offset != 0)
+        //     pgcstack = tp + offset; // fast
+        // else
+        // pgcstack_getter = load pgcstack_func_slot
+        // if pgcstack_getter == nullptr // Runtime not initialized
+        //      pgcstack = nullptr
+        // else
+        //      pgcstack = pgcstack_getter();
         // if (pgcstack != nullptr)
         //     last_gc_state = emit_gc_unsafe_enter(ctx);
         //     phi = pgcstack;        // fast
@@ -177,17 +190,25 @@ void LowerPTLS::fix_pgcstack_use(CallInst *pgcstack, Function *pgcstack_getter,
         if (CFGModified)
             *CFGModified = true;
         // emit slow branch code
-        CallInst *adopt = cast<CallInst>(pgcstack->clone());
         Function *adoptFunc = M->getFunction(XSTR(jl_adopt_thread));
         if (adoptFunc == NULL) {
-            adoptFunc = Function::Create(pgcstack_getter->getFunctionType(),
+            adoptFunc = Function::Create(FunctionType::get(builder.getPtrTy(), { builder.getPtrTy()}, false),
                 pgcstack_getter->getLinkage(), pgcstack_getter->getAddressSpace(),
                 XSTR(jl_adopt_thread), M);
             adoptFunc->copyAttributesFrom(pgcstack_getter);
             adoptFunc->copyMetadata(pgcstack_getter, 0);
         }
-        adopt->setCalledFunction(adoptFunc);
-        adopt->insertBefore(slowTerm);
+        // Adopt thread takes in a handle to the sysimage and this is the easiest way to get it.
+        Function *dladdr = M->getFunction(XSTR(jl_find_dynamic_library_by_addr)); // gets handle to sysimage
+        if (dladdr == NULL) {
+            dladdr = Function::Create(FunctionType::get(builder.getPtrTy(), { builder.getPtrTy()}, false),
+                pgcstack_getter->getLinkage(), pgcstack_getter->getAddressSpace(),
+                XSTR(jl_find_dynamic_library_by_addr), M);
+        }
+        builder.SetInsertPoint(slowTerm);
+        auto this_func = builder.GetInsertBlock()->getParent();
+        auto handle = builder.CreateCall(dladdr, {ConstantExpr::getBitCast(this_func, builder.getPtrTy())});
+        auto adopt = builder.CreateCall(adoptFunc, {handle});
         phi->addIncoming(adopt, slowTerm->getParent());
         // emit fast branch code
         builder.SetInsertPoint(fastTerm->getParent());
@@ -213,17 +234,20 @@ void LowerPTLS::fix_pgcstack_use(CallInst *pgcstack, Function *pgcstack_getter,
 
     if (imaging_mode) {
         IRBuilder<> builder(pgcstack);
+        SmallVector<uint32_t, 2> Weights{9, 1};
+        MDBuilder MDB(pgcstack->getContext());
         if (jl_tls_elf_support) {
             // if (offset != 0)
             //     pgcstack = tp + offset; // fast
             // else
-            //     pgcstack = getter();    // slow
+            //     if pgcstack_getter == null
+            //        pgcstack = null;    // slow
+            //     else
+            //        pgcstack = pgcstack_getter();    // slow
             auto offset = builder.CreateLoad(T_size, pgcstack_offset);
             offset->setMetadata(llvm::LLVMContext::MD_tbaa, tbaa_const);
             offset->setMetadata(llvm::LLVMContext::MD_invariant_load, MDNode::get(pgcstack->getContext(), None));
             auto cmp = builder.CreateICmpNE(offset, Constant::getNullValue(offset->getType()));
-            MDBuilder MDB(pgcstack->getContext());
-            SmallVector<uint32_t, 2> Weights{9, 1};
             TerminatorInst *fastTerm;
             TerminatorInst *slowTerm;
             SplitBlockAndInsertIfThenElse(cmp, pgcstack, &fastTerm, &slowTerm,
@@ -240,21 +264,53 @@ void LowerPTLS::fix_pgcstack_use(CallInst *pgcstack, Function *pgcstack_getter,
             // refresh the basic block in the builder
             builder.SetInsertPoint(pgcstack);
             auto getter = builder.CreateLoad(T_pgcstack_getter, pgcstack_func_slot);
+            auto phi_value = cast<Instruction>(pgcstack);
+            if (or_new) {
+                // if pgcstack_func_slot is not initialized we set pgcstack to null to trigger the slow path
+                TerminatorInst *nonNullTerm;
+                TerminatorInst *nullTerm;
+                auto is_null = builder.CreateICmpEQ(getter, Constant::getNullValue(builder.getPtrTy()));
+                SplitBlockAndInsertIfThenElse(is_null, pgcstack, &nullTerm, &nonNullTerm,
+                                    MDB.createBranchWeights(Weights));
+                builder.SetInsertPoint(pgcstack);
+                auto phi2 = builder.CreatePHI(T_pppjlvalue, 2, "pgcstack");
+                pgcstack->moveBefore(nonNullTerm);
+                phi2->addIncoming(pgcstack, nonNullTerm->getParent());
+                phi2->addIncoming(Constant::getNullValue(T_pppjlvalue), nullTerm->getParent());
+                phi_value = phi2;
+                // Check if pgcstack_func_slot is initialized
+            }
             getter->setMetadata(llvm::LLVMContext::MD_tbaa, tbaa_const);
             getter->setMetadata(llvm::LLVMContext::MD_invariant_load, MDNode::get(pgcstack->getContext(), None));
             pgcstack->setCalledFunction(pgcstack->getFunctionType(), getter);
             set_pgcstack_attrs(pgcstack);
 
             phi->addIncoming(fastTLS, fastTLS->getParent());
-            phi->addIncoming(pgcstack, pgcstack->getParent());
-
+            phi->addIncoming(phi_value, phi_value->getParent());
+            if (or_new) {
+                pgcstack->getParent()->getParent()->print(dbgs());
+            }
             return;
         }
         // In imaging mode, we emit the function address as a load of a static
         // variable to be filled (in `staticdata.c`) at initialization time of the sysimg.
         // This way we can bypass the extra indirection in `jl_get_pgcstack`
         // since we may not know which getter function to use ahead of time.
         auto getter = builder.CreateLoad(T_pgcstack_getter, pgcstack_func_slot);
+        if (or_new) {
+            // if pgcstack_func_slot is not initialized we set pgcstack to null to trigger the slow path
+            TerminatorInst *nonNullTerm;
+            TerminatorInst *nullTerm;
+            auto is_null = builder.CreateICmpEQ(getter, Constant::getNullValue(builder.getPtrTy()));
+            SplitBlockAndInsertIfThenElse(is_null, pgcstack, &nullTerm, &nonNullTerm,
+                                MDB.createBranchWeights(Weights));
+            builder.SetInsertPoint(pgcstack);
+            auto phi2 = builder.CreatePHI(T_pppjlvalue, 2, "pgcstack");
+            pgcstack->moveBefore(nonNullTerm);
+            phi2->addIncoming(pgcstack, nonNullTerm->getParent());
+            phi2->addIncoming(Constant::getNullValue(T_pppjlvalue), nullTerm->getParent());
+            pgcstack->replaceAllUsesWith(phi2);
+        }
         getter->setMetadata(llvm::LLVMContext::MD_tbaa, tbaa_const);
         getter->setMetadata(llvm::LLVMContext::MD_invariant_load, MDNode::get(pgcstack->getContext(), None));
         if (TargetTriple.isOSDarwin()) {
diff --git a/src/threading.c b/src/threading.c
@@ -401,8 +401,14 @@ jl_ptls_t jl_init_threadtls(int16_t tid)
     return ptls;
 }
 
-JL_DLLEXPORT jl_gcframe_t **jl_adopt_thread(void)
+// Pass in the handle to the system image. This is used to initialize the runtime correctly in case we are a shared library
+JL_DLLEXPORT jl_gcframe_t **jl_adopt_thread(void* sysimg_handle)
 {
+    if (!jl_is_initialized()) {
+        if (jl_init_runtime_adopt_thread(sysimg_handle) == 1)
+            return &jl_get_current_task()->gcstack;
+        // We lost the race and need to be initialized as usual
+    }
     // `jl_init_threadtls` puts us in a GC unsafe region, so ensure GC isn't running.
     // we can't use a normal safepoint because we don't have signal handlers yet.
     // we also can't use jl_safepoint_wait_gc because that assumes we're in a task.

Original file line number	Diff line number	Diff line change
`@@ -254,7 +254,10 @@ void jl_find_dynamic_library_by_addr(void symbol) {`
`254`	`254`	`jl_error("could not load base module");`
`255`	`255`	`}`
`256`	`256`	`handle = dlopen(info.dli_fname, RTLD_NOW \| RTLD_NOLOAD \| RTLD_LOCAL);`
`257`		- dlclose(handle); // Undo ref count increment from `dlopen`
	`257`	`+ if (handle == NULL && dlerror() == NULL) // We loaded the executable but got RTLD_DEFAULT back, give a real handle instead`
	`258`	`+ handle = dlopen("", RTLD_NOW \| RTLD_NOLOAD \| RTLD_LOCAL);`
	`259`	`+ if (handle) // We may get a null handle so don't segfault`
	`260`	+ dlclose(handle); // Undo ref count increment from `dlopen`
`258`	`261`	`#endif`
`259`	`262`	`return handle;`
`260`	`263`	`}`