Re-work jl_adopt_thread runtime auto-initialization

This should be a dramatically simpler implementation, by just setting up
the `jl_pgcstack_func_slot` to have a usable function pointer by default
and then performing library introspection based on
`__builtin_return_address` in the runtime.

This avoids needing to modify codegen for `julia.get_pgcstack` at all.

Author: topolarity

src/aotcompile.cpp

@@ -975,17 +975,19 @@ static GlobalVariable *emit_shard_table(Module &M, Type *T_size, Type *T_psize,
 }
 
 // See src/processor.h for documentation about this table. Corresponds to jl_image_ptls_t.
-static GlobalVariable *emit_ptls_table(Module &M, Type *T_size, Type *T_psize) {
+static GlobalVariable *emit_ptls_table(Module &M, Type *T_size, Type *T_ptr) {
+    Function *jl_pgcstack_default_func = Function::Create(
+        FunctionType::get(T_ptr, false), GlobalValue::ExternalLinkage, "jl_pgcstack_default_func", &M);
     std::array<Constant *, 3> ptls_table{
-        new GlobalVariable(M, T_size, false, GlobalValue::ExternalLinkage, Constant::getNullValue(T_size), "jl_pgcstack_func_slot"),
+        new GlobalVariable(M, T_ptr, false, GlobalValue::ExternalLinkage, jl_pgcstack_default_func, "jl_pgcstack_func_slot"),
         new GlobalVariable(M, T_size, false, GlobalValue::ExternalLinkage, Constant::getNullValue(T_size), "jl_pgcstack_key_slot"),
         new GlobalVariable(M, T_size, false, GlobalValue::ExternalLinkage, Constant::getNullValue(T_size), "jl_tls_offset"),
     };
     for (auto &gv : ptls_table) {
         cast<GlobalVariable>(gv)->setVisibility(GlobalValue::HiddenVisibility);
         cast<GlobalVariable>(gv)->setDSOLocal(true);
     }
-    auto ptls_table_arr = ConstantArray::get(ArrayType::get(T_psize, ptls_table.size()), ptls_table);
+    auto ptls_table_arr = ConstantArray::get(ArrayType::get(T_ptr, ptls_table.size()), ptls_table);
     auto ptls_table_gv = new GlobalVariable(M, ptls_table_arr->getType(), false,
                                             GlobalValue::ExternalLinkage, ptls_table_arr, "jl_ptls_table");
     ptls_table_gv->setVisibility(GlobalValue::HiddenVisibility);
@@ -2184,6 +2186,7 @@ void jl_dump_native_impl(void *native_code,
 
         Type *T_size = DL.getIntPtrType(Context);
         Type *T_psize = T_size->getPointerTo();
+        Type *T_ptr = PointerType::get(Context, 0);
 
         auto FT = FunctionType::get(Type::getInt8Ty(Context)->getPointerTo()->getPointerTo(), {}, false);
         auto F = Function::Create(FT, Function::ExternalLinkage, "get_jl_RTLD_DEFAULT_handle_addr", metadataM);
@@ -2226,7 +2229,7 @@ void jl_dump_native_impl(void *native_code,
                                         GlobalVariable::InternalLinkage,
                                         value, "jl_dispatch_target_ids");
             auto shards = emit_shard_table(metadataM, T_size, T_psize, threads);
-            auto ptls = emit_ptls_table(metadataM, T_size, T_psize);
+            auto ptls = emit_ptls_table(metadataM, T_size, T_ptr);
             auto header = emit_image_header(metadataM, threads, nfvars, ngvars);
             auto AT = ArrayType::get(T_size, sizeof(jl_small_typeof) / sizeof(void*));
             auto jl_small_typeof_copy = new GlobalVariable(metadataM, AT, false,

src/dlload.c

@@ -240,22 +240,30 @@ JL_DLLEXPORT int jl_dlclose(void *handle) JL_NOTSAFEPOINT
 #endif
 }
 
-void *jl_find_dynamic_library_by_addr(void *symbol) {
+void *jl_find_dynamic_library_by_addr(void *symbol, int throw_err) {
     void *handle;
 #ifdef _OS_WINDOWS_
     if (!GetModuleHandleExW(GET_MODULE_HANDLE_EX_FLAG_FROM_ADDRESS | GET_MODULE_HANDLE_EX_FLAG_UNCHANGED_REFCOUNT,
                             (LPCWSTR)symbol,
                             (HMODULE*)&handle)) {
-        jl_error("could not load base module");
+        if (throw_err)
+            jl_error("could not load base module");
+        else
+            return NULL;
     }
 #else
     Dl_info info;
     if (!dladdr(symbol, &info) || !info.dli_fname) {
-        jl_error("could not load base module");
+        if (throw_err)
+            jl_error("could not load base module");
+        else
+            return NULL;
     }
     handle = dlopen(info.dli_fname, RTLD_NOW | RTLD_NOLOAD | RTLD_LOCAL);
-    if (handle == NULL && dlerror() == NULL) // We loaded the executable but got RTLD_DEFAULT back, give a real handle instead
+    if (handle == RTLD_DEFAULT && (RTLD_DEFAULT != NULL || dlerror() == NULL)) {
+        // We loaded the executable but got RTLD_DEFAULT back, ask for 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
@@ -280,7 +288,7 @@ JL_DLLEXPORT void *jl_load_dynamic_library(const char *modname, unsigned flags,
 
     // modname == NULL is a sentinel value requesting the handle of libjulia-internal
     if (modname == NULL)
-        return jl_find_dynamic_library_by_addr(&jl_load_dynamic_library);
+        return jl_find_dynamic_library_by_addr(&jl_load_dynamic_library, throw_err);
 
     abspath = jl_isabspath(modname);
     is_atpath = 0;

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(NULL), jl_task_t, gcstack);
+        ct = container_of(jl_adopt_thread(), jl_task_t, gcstack);
     }
     else if (ct != NULL) {
         // we are about to start tearing everything down, so lets try not to get
@@ -808,8 +808,8 @@ JL_DLLEXPORT void julia_init(JL_IMAGE_SEARCH rel)
     void *stack_lo, *stack_hi;
     jl_init_stack_limits(1, &stack_lo, &stack_hi);
 
-    jl_libjulia_internal_handle = jl_find_dynamic_library_by_addr(&jl_load_dynamic_library);
-    jl_libjulia_handle = jl_find_dynamic_library_by_addr(&jl_any_type);
+    jl_libjulia_internal_handle = jl_find_dynamic_library_by_addr(&jl_load_dynamic_library, /* throw_err */ 1);
+    jl_libjulia_handle = jl_find_dynamic_library_by_addr(&jl_any_type, /* throw_err */ 1);
 #ifdef _OS_WINDOWS_
     jl_exe_handle = GetModuleHandleA(NULL);
     jl_RTLD_DEFAULT_handle = jl_libjulia_internal_handle;

src/jl_exported_funcs.inc

@@ -40,6 +40,7 @@
     XX(jl_atomic_store_bits) \
     XX(jl_atomic_storeonce_bits) \
     XX(jl_atomic_swap_bits) \
+    XX(jl_autoinit_and_adopt_thread) \
     XX(jl_backtrace_from_here) \
     XX(jl_base_relative_to) \
     XX(jl_bitcast) \
@@ -292,7 +293,6 @@
     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) \
@@ -362,6 +362,7 @@
     XX(jl_pathname_for_handle) \
     XX(jl_pchar_to_array) \
     XX(jl_pchar_to_string) \
+    XX(jl_pgcstack_default_func) \
     XX(jl_pointerref) \
     XX(jl_pointerset) \
     XX(jl_pop_handler) \

src/jlapi.c

@@ -1125,13 +1125,6 @@ JL_DLLEXPORT int jl_repl_entrypoint(int argc, char *argv[])
     jl_atexit_hook(ret);
     return ret;
 }
-int jl_init_runtime_adopt_thread(void* sysimg_handle)
-{
-    //Unimplemented
-    abort();
-    return 0;
-}
-
 
 #ifdef __cplusplus
 }

src/julia.h

@@ -2241,7 +2241,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* sysimage_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);

src/julia_internal.h

@@ -1099,7 +1099,6 @@ 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);
 
@@ -1514,7 +1513,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);
-JL_DLLEXPORT void *jl_find_dynamic_library_by_addr(void *symbol);
+void *jl_find_dynamic_library_by_addr(void *symbol, int throw_err);
 #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);

src/llvm-ptls.cpp

@@ -5,9 +5,6 @@
 #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>
 
@@ -146,17 +143,7 @@ void LowerPTLS::fix_pgcstack_use(CallInst *pgcstack, Function *pgcstack_getter,
         return;
     }
     if (or_new) {
-        // 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();
+        // pgcstack();
         // if (pgcstack != nullptr)
         //     last_gc_state = emit_gc_unsafe_enter(ctx);
         //     phi = pgcstack;        // fast
@@ -182,29 +169,17 @@ void LowerPTLS::fix_pgcstack_use(CallInst *pgcstack, Function *pgcstack_getter,
         if (CFGModified)
             *CFGModified = true;
         // emit slow branch code
-        Function *adoptFunc = M->getFunction(XSTR(jl_adopt_thread));
+        CallInst *adopt = cast<CallInst>(pgcstack->clone());
+        Function *adoptFunc = M->getFunction(XSTR(jl_autoinit_and_adopt_thread));
         if (adoptFunc == NULL) {
-            adoptFunc = Function::Create(FunctionType::get(builder.getPtrTy(), { builder.getPtrTy()}, false),
+            adoptFunc = Function::Create(pgcstack_getter->getFunctionType(),
                 pgcstack_getter->getLinkage(), pgcstack_getter->getAddressSpace(),
-                XSTR(jl_adopt_thread), M);
+                XSTR(jl_autoinit_and_adopt_thread), M);
             adoptFunc->copyAttributesFrom(pgcstack_getter);
             adoptFunc->copyMetadata(pgcstack_getter, 0);
         }
-        builder.SetInsertPoint(slowTerm);
-        Value* handle = Constant::getNullValue(builder.getPtrTy());
-        if (imaging_mode) {
-            // 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);
-            }
-            auto this_func = builder.GetInsertBlock()->getParent();
-            handle = builder.CreateCall(dladdr, {ConstantExpr::getBitCast(this_func, builder.getPtrTy())});
-        }
-
-        auto adopt = builder.CreateCall(adoptFunc, {handle});
+        adopt->setCalledFunction(adoptFunc);
+        adopt->insertBefore(slowTerm);
         phi->addIncoming(adopt, slowTerm->getParent());
         // emit fast branch code
         builder.SetInsertPoint(fastTerm->getParent());
@@ -235,20 +210,17 @@ 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
-            //     if pgcstack_getter == null
-            //        pgcstack = null;    // slow
-            //     else
-            //        pgcstack = pgcstack_getter();    // slow
+            //     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,
@@ -265,52 +237,21 @@ 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;
-                builder.SetInsertPoint(pgcstack);
-                // 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(phi_value, phi_value->getParent());
+            phi->addIncoming(pgcstack, pgcstack->getParent());
+
             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);
-            pgcstack->replaceAllUsesWith(phi2);
-            phi2->addIncoming(pgcstack, nonNullTerm->getParent());
-            phi2->addIncoming(Constant::getNullValue(T_pppjlvalue), nullTerm->getParent());
-            builder.SetInsertPoint(pgcstack);
-        }
         getter->setMetadata(llvm::LLVMContext::MD_tbaa, tbaa_const);
         getter->setMetadata(llvm::LLVMContext::MD_invariant_load, MDNode::get(pgcstack->getContext(), None));
         if (TargetTriple.isOSDarwin()) {

src/threading.c

@@ -276,6 +276,11 @@ JL_DLLEXPORT jl_gcframe_t **jl_get_pgcstack(void) JL_GLOBALLY_ROOTED
 #endif
 }
 
+JL_DLLEXPORT jl_gcframe_t **jl_pgcstack_default_func(void) JL_GLOBALLY_ROOTED
+{
+    return NULL;
+}
+
 void jl_pgcstack_getkey(jl_get_pgcstack_func **f, jl_pgcstack_key_t *k)
 {
 #ifndef __clang_gcanalyzer__
@@ -433,14 +438,9 @@ static void jl_init_task_lock(jl_task_t *ct)
         }
     }
 }
-// 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)
+
+JL_DLLEXPORT jl_gcframe_t **jl_adopt_thread(void)
 {
-    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.
     jl_atomic_fetch_add(&jl_gc_disable_counter, 1);
@@ -464,6 +464,23 @@ JL_DLLEXPORT jl_gcframe_t **jl_adopt_thread(void* sysimg_handle)
     return &ct->gcstack;
 }
 
+JL_DLLEXPORT jl_gcframe_t **jl_autoinit_and_adopt_thread(void)
+{
+    void *retaddr = __builtin_extract_return_addr(__builtin_return_address(0));
+    void *sysimg_handle = jl_find_dynamic_library_by_addr(retaddr, /* throw_err */ 0);
+
+    if (!jl_is_initialized()) {
+        if (sysimg_handle == NULL) {
+            fprintf(stderr, "error: runtime auto-initialization failed due to bad sysimage lookup\n"
+                            "       (this should not happen, please file a bug report)\n");
+            abort();
+        }
+
+        assert(0 && "TODO: implement auto-init");
+    }
+
+    return jl_adopt_thread();
+}
 
 void jl_safepoint_suspend_all_threads(jl_task_t *ct)
 {