enable re-using external code in pkgimages (#48723)

vtjnash · vchuravy · pull[bot] · commit 834dc26a430c · 2023-12-20T14:17:42.000Z
* enable using external code in pkgimages

This was unintentionally disabled (and incomplete) in the original PR
for pkgimages.

Co-authored-by: Valentin Churavy &lt;v.churavy@gmail.com&gt;
diff --git a/base/linking.jl b/base/linking.jl
@@ -162,8 +162,8 @@ function link_image_cmd(path, out)
     `$(ld()) $V $SHARED -o $out $WHOLE_ARCHIVE $path $NO_WHOLE_ARCHIVE $LIBDIR $PRIVATE_LIBDIR $SHLIBDIR $LIBS`
 end
 
-function link_image(path, out, internal_stderr::IO = stderr, internal_stdout::IO = stdout)
-    run(link_image_cmd(path, out), Base.DevNull(), stderr, stdout)
+function link_image(path, out, internal_stderr::IO=stderr, internal_stdout::IO=stdout)
+    run(link_image_cmd(path, out), Base.DevNull(), internal_stderr, internal_stdout)
 end
 
 end # module Linking
diff --git a/src/aotcompile.cpp b/src/aotcompile.cpp
@@ -369,34 +369,16 @@ void *jl_create_native_impl(jl_array_t *methods, LLVMOrcThreadSafeModuleRef llvm
     size_t offset = gvars.size();
     data->jl_external_to_llvm.resize(params.external_fns.size());
 
-    auto tbaa_const = tbaa_make_child_with_context(*ctxt.getContext(), "jtbaa_const", nullptr, true).first;
     for (auto &extern_fn : params.external_fns) {
         jl_code_instance_t *this_code = std::get<0>(extern_fn.first);
         bool specsig = std::get<1>(extern_fn.first);
         assert(specsig && "Error external_fns doesn't handle non-specsig yet");
-        (void)specsig;
-        Function *F = extern_fn.second;
-        Module *M = F->getParent();
-
-        Type *T_funcp = F->getFunctionType()->getPointerTo();
-        // Can't create a GC with type FunctionType. Alias also doesn't work
-        GlobalVariable *GV = new GlobalVariable(*M, T_funcp, false,
-                                                GlobalVariable::ExternalLinkage,
-                                                Constant::getNullValue(T_funcp),
-                                                F->getName());
-
-
-        // Need to insert load instruction, thus we can't use replace all uses with
-        replaceUsesWithLoad(*F, [GV](Instruction &) { return GV; }, tbaa_const);
-
-        assert(F->getNumUses() == 0); // declaration counts as use
-        GV->takeName(F);
-        F->eraseFromParent();
-
+        (void) specsig;
+        GlobalVariable *F = extern_fn.second;
         size_t idx = gvars.size() - offset;
         assert(idx >= 0);
         data->jl_external_to_llvm.at(idx) = this_code;
-        gvars.push_back(std::string(GV->getName()));
+        gvars.push_back(std::string(F->getName()));
     }
 
     // clones the contents of the module `m` to the shadow_output collector
diff --git a/src/codegen.cpp b/src/codegen.cpp
@@ -1542,7 +1542,7 @@ class jl_codectx_t {
     jl_codegen_params_t &emission_context;
     llvm::MapVector<jl_code_instance_t*, jl_codegen_call_target_t> call_targets;
     std::map<void*, GlobalVariable*> &global_targets;
-    std::map<std::tuple<jl_code_instance_t*, bool>, Function*> &external_calls;
+    std::map<std::tuple<jl_code_instance_t*, bool>, GlobalVariable*> &external_calls;
     Function *f = NULL;
     // local var info. globals are not in here.
     std::vector<jl_varinfo_t> slots;
@@ -1704,7 +1704,7 @@ static Value *get_current_task(jl_codectx_t &ctx);
 static Value *get_current_ptls(jl_codectx_t &ctx);
 static Value *get_last_age_field(jl_codectx_t &ctx);
 static void CreateTrap(IRBuilder<> &irbuilder, bool create_new_block = true);
-static CallInst *emit_jlcall(jl_codectx_t &ctx, Function *theFptr, Value *theF,
+static CallInst *emit_jlcall(jl_codectx_t &ctx, FunctionCallee theFptr, Value *theF,
                              const jl_cgval_t *args, size_t nargs, JuliaFunction *trampoline);
 static CallInst *emit_jlcall(jl_codectx_t &ctx, JuliaFunction *theFptr, Value *theF,
                              const jl_cgval_t *args, size_t nargs, JuliaFunction *trampoline);
@@ -4039,14 +4039,14 @@ static bool emit_builtin_call(jl_codectx_t &ctx, jl_cgval_t *ret, jl_value_t *f,
 }
 
 // Returns ctx.types().T_prjlvalue
-static CallInst *emit_jlcall(jl_codectx_t &ctx, Function *theFptr, Value *theF,
+static CallInst *emit_jlcall(jl_codectx_t &ctx, FunctionCallee theFptr, Value *theF,
                              const jl_cgval_t *argv, size_t nargs, JuliaFunction *trampoline)
 {
     ++EmittedJLCalls;
     Function *TheTrampoline = prepare_call(trampoline);
     // emit arguments
     SmallVector<Value*, 4> theArgs;
-    theArgs.push_back(theFptr);
+    theArgs.push_back(theFptr.getCallee());
     if (theF)
         theArgs.push_back(theF);
     for (size_t i = 0; i < nargs; i++) {
@@ -4067,7 +4067,7 @@ static CallInst *emit_jlcall(jl_codectx_t &ctx, JuliaFunction *theFptr, Value *t
 }
 
 
-static jl_cgval_t emit_call_specfun_other(jl_codectx_t &ctx, jl_method_instance_t *mi, jl_value_t *jlretty, StringRef specFunctionObject,
+static jl_cgval_t emit_call_specfun_other(jl_codectx_t &ctx, jl_method_instance_t *mi, jl_value_t *jlretty, StringRef specFunctionObject, jl_code_instance_t *fromexternal,
                                           const jl_cgval_t *argv, size_t nargs, jl_returninfo_t::CallingConv *cc, unsigned *return_roots, jl_value_t *inferred_retty)
 {
     ++EmittedSpecfunCalls;
@@ -4143,7 +4143,22 @@ static jl_cgval_t emit_call_specfun_other(jl_codectx_t &ctx, jl_method_instance_
         idx++;
     }
     assert(idx == nfargs);
-    CallInst *call = ctx.builder.CreateCall(returninfo.decl, ArrayRef<Value*>(&argvals[0], nfargs));
+    Value *callee = returninfo.decl;
+    if (fromexternal) {
+        std::string namep("p");
+        namep += returninfo.decl->getName();
+        GlobalVariable *GV = cast_or_null<GlobalVariable>(jl_Module->getNamedValue(namep));
+        if (GV == nullptr) {
+            GV = new GlobalVariable(*jl_Module, callee->getType(), false,
+                                    GlobalVariable::ExternalLinkage,
+                                    Constant::getNullValue(callee->getType()),
+                                    namep);
+            ctx.external_calls[std::make_tuple(fromexternal, true)] = GV;
+        }
+        jl_aliasinfo_t ai = jl_aliasinfo_t::fromTBAA(ctx, ctx.tbaa().tbaa_const);
+        callee = ai.decorateInst(ctx.builder.CreateAlignedLoad(callee->getType(), GV, Align(sizeof(void*))));
+    }
+    CallInst *call = ctx.builder.CreateCall(cft, callee, ArrayRef<Value*>(&argvals[0], nfargs));
     call->setAttributes(returninfo.decl->getAttributes());
 
     jl_cgval_t retval;
@@ -4182,13 +4197,30 @@ static jl_cgval_t emit_call_specfun_other(jl_codectx_t &ctx, jl_method_instance_
     return update_julia_type(ctx, retval, inferred_retty);
 }
 
-static jl_cgval_t emit_call_specfun_boxed(jl_codectx_t &ctx, jl_value_t *jlretty, StringRef specFunctionObject,
+static jl_cgval_t emit_call_specfun_boxed(jl_codectx_t &ctx, jl_value_t *jlretty, StringRef specFunctionObject, jl_code_instance_t *fromexternal,
                                           const jl_cgval_t *argv, size_t nargs, jl_value_t *inferred_retty)
 {
-    auto theFptr = cast<Function>(
-        jl_Module->getOrInsertFunction(specFunctionObject, ctx.types().T_jlfunc).getCallee());
-    addRetAttr(theFptr, Attribute::NonNull);
-    Value *ret = emit_jlcall(ctx, theFptr, nullptr, argv, nargs, julia_call);
+    Value *theFptr;
+    if (fromexternal) {
+        std::string namep("p");
+        namep += specFunctionObject;
+        GlobalVariable *GV = cast_or_null<GlobalVariable>(jl_Module->getNamedValue(namep));
+        Type *pfunc = ctx.types().T_jlfunc->getPointerTo();
+        if (GV == nullptr) {
+            GV = new GlobalVariable(*jl_Module, pfunc, false,
+                                    GlobalVariable::ExternalLinkage,
+                                    Constant::getNullValue(pfunc),
+                                    namep);
+            ctx.external_calls[std::make_tuple(fromexternal, false)] = GV;
+        }
+        jl_aliasinfo_t ai = jl_aliasinfo_t::fromTBAA(ctx, ctx.tbaa().tbaa_const);
+        theFptr = ai.decorateInst(ctx.builder.CreateAlignedLoad(pfunc, GV, Align(sizeof(void*))));
+    }
+    else {
+        theFptr = jl_Module->getOrInsertFunction(specFunctionObject, ctx.types().T_jlfunc).getCallee();
+        addRetAttr(cast<Function>(theFptr), Attribute::NonNull);
+    }
+    Value *ret = emit_jlcall(ctx, FunctionCallee(ctx.types().T_jlfunc, theFptr), nullptr, argv, nargs, julia_call);
     return update_julia_type(ctx, mark_julia_type(ctx, ret, true, jlretty), inferred_retty);
 }
 
@@ -4223,12 +4255,12 @@ static jl_cgval_t emit_invoke(jl_codectx_t &ctx, const jl_cgval_t &lival, const
             FunctionType *ft = ctx.f->getFunctionType();
             StringRef protoname = ctx.f->getName();
             if (ft == ctx.types().T_jlfunc) {
-                result = emit_call_specfun_boxed(ctx, ctx.rettype, protoname, argv, nargs, rt);
+                result = emit_call_specfun_boxed(ctx, ctx.rettype, protoname, nullptr, argv, nargs, rt);
                 handled = true;
             }
             else if (ft != ctx.types().T_jlfuncparams) {
                 unsigned return_roots = 0;
-                result = emit_call_specfun_other(ctx, mi, ctx.rettype, protoname, argv, nargs, &cc, &return_roots, rt);
+                result = emit_call_specfun_other(ctx, mi, ctx.rettype, protoname, nullptr, argv, nargs, &cc, &return_roots, rt);
                 handled = true;
             }
         }
@@ -4248,16 +4280,17 @@ 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;
-                    bool cache_valid = ctx.use_cache;
+                    bool cache_valid = ctx.use_cache || ctx.external_linkage;
                     bool external = false;
-                    if (ctx.external_linkage) {
-                       if (0 && jl_object_in_image((jl_value_t*)codeinst)) {
-                           // Target is present in another pkgimage
-                           cache_valid = true;
-                           external = true;
-                       }
+
+                    // Check if we already queued this up
+                    auto it = ctx.call_targets.find(codeinst);
+                    if (need_to_emit && it != ctx.call_targets.end()) {
+                        protoname = std::get<2>(it->second)->getName();
+                        need_to_emit = cache_valid = false;
                     }
 
+                    // Check if it is already compiled (either JIT or externally)
                     if (cache_valid) {
                         // optimization: emit the correct name immediately, if we know it
                         // TODO: use `emitted` map here too to try to consolidate names?
@@ -4270,32 +4303,30 @@ static jl_cgval_t emit_invoke(jl_codectx_t &ctx, const jl_cgval_t &lival, const
                             invoke = jl_atomic_load_relaxed(&codeinst->invoke);
                             if (specsig ? jl_atomic_load_relaxed(&codeinst->specsigflags) & 0b1 : invoke == jl_fptr_args_addr) {
                                 protoname = jl_ExecutionEngine->getFunctionAtAddress((uintptr_t)fptr, codeinst);
-                                need_to_emit = false;
+                                if (ctx.external_linkage) {
+                                    // TODO: Add !specsig support to aotcompile.cpp
+                                    // Check that the codeinst is containing native code
+                                    if (specsig && jl_atomic_load_relaxed(&codeinst->specsigflags) & 0b100) {
+                                        external = true;
+                                        need_to_emit = false;
+                                    }
+                                }
+                                else { // ctx.use_cache
+                                    need_to_emit = false;
+                                }
                             }
                         }
                     }
-                    auto it = ctx.call_targets.find(codeinst);
-                    if (need_to_emit && it != ctx.call_targets.end()) {
-                        protoname = std::get<2>(it->second)->getName();
-                        need_to_emit = false;
-                    }
                     if (need_to_emit) {
                         raw_string_ostream(name) << (specsig ? "j_" : "j1_") << name_from_method_instance(mi) << "_" << jl_atomic_fetch_add(&globalUniqueGeneratedNames, 1);
                         protoname = StringRef(name);
                     }
                     jl_returninfo_t::CallingConv cc = jl_returninfo_t::CallingConv::Boxed;
                     unsigned return_roots = 0;
                     if (specsig)
-                        result = emit_call_specfun_other(ctx, mi, codeinst->rettype, protoname, argv, nargs, &cc, &return_roots, rt);
+                        result = emit_call_specfun_other(ctx, mi, codeinst->rettype, protoname, external ? codeinst : nullptr, argv, nargs, &cc, &return_roots, rt);
                     else
-                        result = emit_call_specfun_boxed(ctx, codeinst->rettype, protoname, argv, nargs, rt);
-                    if (external) {
-                        assert(!need_to_emit);
-                        auto calledF = jl_Module->getFunction(protoname);
-                        assert(calledF);
-                        // TODO: Check if already present?
-                        ctx.external_calls[std::make_tuple(codeinst, specsig)] = calledF;
-                    }
+                        result = emit_call_specfun_boxed(ctx, codeinst->rettype, protoname, external ? codeinst : nullptr, argv, nargs, rt);
                     handled = true;
                     if (need_to_emit) {
                         Function *trampoline_decl = cast<Function>(jl_Module->getNamedValue(protoname));
@@ -5617,14 +5648,7 @@ static Function *emit_tojlinvoke(jl_code_instance_t *codeinst, Module *M, jl_cod
     Function *theFunc;
     Value *theFarg;
     auto invoke = jl_atomic_load_relaxed(&codeinst->invoke);
-
     bool cache_valid = params.cache;
-    if (params.external_linkage) {
-        if (0 && jl_object_in_image((jl_value_t*)codeinst)) {
-            // Target is present in another pkgimage
-            cache_valid = true;
-        }
-    }
 
     if (cache_valid && invoke != NULL) {
         StringRef theFptrName = jl_ExecutionEngine->getFunctionAtAddress((uintptr_t)invoke, codeinst);
@@ -8537,9 +8561,6 @@ void jl_compile_workqueue(
         bool preal_specsig = false;
         auto invoke = jl_atomic_load_acquire(&codeinst->invoke);
         bool cache_valid = params.cache;
-        if (params.external_linkage) {
-            cache_valid = 0 && jl_object_in_image((jl_value_t*)codeinst);
-        }
         // WARNING: isspecsig is protected by the codegen-lock. If that lock is removed, then the isspecsig load needs to be properly atomically sequenced with this.
         if (cache_valid && invoke != NULL) {
             auto fptr = jl_atomic_load_relaxed(&codeinst->specptr.fptr);
diff --git a/src/jitlayers.h b/src/jitlayers.h
@@ -190,7 +190,7 @@ typedef struct _jl_codegen_params_t {
     // outputs
     std::vector<std::pair<jl_code_instance_t*, jl_codegen_call_target_t>> workqueue;
     std::map<void*, GlobalVariable*> globals;
-    std::map<std::tuple<jl_code_instance_t*,bool>, Function*> external_fns;
+    std::map<std::tuple<jl_code_instance_t*,bool>, GlobalVariable*> external_fns;
     std::map<jl_datatype_t*, DIType*> ditypes;
     std::map<jl_datatype_t*, Type*> llvmtypes;
     DenseMap<Constant*, GlobalVariable*> mergedConstants;
diff --git a/src/julia.h b/src/julia.h
@@ -424,7 +424,9 @@ typedef struct _jl_code_instance_t {
     jl_value_t *argescapes; // escape information of call arguments
 
     // compilation state cache
-    _Atomic(uint8_t) specsigflags; // & 0b1 == specptr is a specialized function signature for specTypes->rettype, &0b10 == invokeptr matches specptr
+    _Atomic(uint8_t) specsigflags; // & 0b001 == specptr is a specialized function signature for specTypes->rettype
+                                   // & 0b010 == invokeptr matches specptr
+                                   // & 0b100 == From image
     _Atomic(uint8_t) precompile;  // if set, this will be added to the output system image
     uint8_t relocatability;  // nonzero if all roots are built into sysimg or tagged by module key
     _Atomic(jl_callptr_t) invoke; // jlcall entry point
diff --git a/src/staticdata.c b/src/staticdata.c
@@ -1077,7 +1077,7 @@ static void record_external_fns(jl_serializer_state *s, arraylist_t *external_fn
 #ifndef JL_NDEBUG
     for (size_t i = 0; i < external_fns->len; i++) {
         jl_code_instance_t *ci = (jl_code_instance_t*)external_fns->items[i];
-        assert(jl_object_in_image((jl_value_t*)ci));
+        assert(jl_atomic_load_relaxed(&ci->specsigflags) & 0b100);
     }
 #endif
 }
@@ -1889,7 +1889,7 @@ static void jl_update_all_fptrs(jl_serializer_state *s, jl_image_t *image)
             void *fptr = (void*)(base + offset);
             if (specfunc) {
                 codeinst->specptr.fptr = fptr;
-                codeinst->specsigflags = 0b11; // TODO: set only if confirmed to be true
+                codeinst->specsigflags = 0b111; // TODO: set only if confirmed to be true
             }
             else {
                 codeinst->invoke = (jl_callptr_t)fptr;
@@ -1913,11 +1913,12 @@ static uint32_t write_gvars(jl_serializer_state *s, arraylist_t *globals, arrayl
     }
     for (size_t i = 0; i < external_fns->len; i++) {
         jl_code_instance_t *ci = (jl_code_instance_t*)external_fns->items[i];
+        assert(ci && (jl_atomic_load_relaxed(&ci->specsigflags) & 0b001));
         uintptr_t item = backref_id(s, (void*)ci, s->link_ids_external_fnvars);
         uintptr_t reloc = get_reloc_for_item(item, 0);
         write_reloc_t(s->gvar_record, reloc);
     }
-    return globals->len + 1;
+    return globals->len;
 }
 
 // Pointer relocation for native-code referenced global variables
@@ -1962,7 +1963,7 @@ static void jl_root_new_gvars(jl_serializer_state *s, jl_image_t *image, uint32_
                 v = (uintptr_t)jl_as_global_root((jl_value_t*)v);
         } else {
             jl_code_instance_t *codeinst = (jl_code_instance_t*) v;
-            assert(codeinst && (codeinst->specsigflags & 0b01));
+            assert(codeinst && (codeinst->specsigflags & 0b01) && codeinst->specptr.fptr);
             v = (uintptr_t)codeinst->specptr.fptr;
         }
         *gv = v;

Original file line number	Diff line number	Diff line change
`@@ -1077,7 +1077,7 @@ static void record_external_fns(jl_serializer_state s, arraylist_t external_fn`
`1077`	`1077`	`#ifndef JL_NDEBUG`
`1078`	`1078`	`for (size_t i = 0; i < external_fns->len; i++) {`
`1079`	`1079`	`jl_code_instance_t ci = (jl_code_instance_t)external_fns->items[i];`
`1080`		`- assert(jl_object_in_image((jl_value_t*)ci));`
	`1080`	`+ assert(jl_atomic_load_relaxed(&ci->specsigflags) & 0b100);`
`1081`	`1081`	`}`
`1082`	`1082`	`#endif`
`1083`	`1083`	`}`
`@@ -1889,7 +1889,7 @@ static void jl_update_all_fptrs(jl_serializer_state s, jl_image_t image)`
`1889`	`1889`	`void fptr = (void)(base + offset);`
`1890`	`1890`	`if (specfunc) {`
`1891`	`1891`	`codeinst->specptr.fptr = fptr;`
`1892`		`- codeinst->specsigflags = 0b11; // TODO: set only if confirmed to be true`
	`1892`	`+ codeinst->specsigflags = 0b111; // TODO: set only if confirmed to be true`
`1893`	`1893`	`}`
`1894`	`1894`	`else {`
`1895`	`1895`	`codeinst->invoke = (jl_callptr_t)fptr;`
`@@ -1913,11 +1913,12 @@ static uint32_t write_gvars(jl_serializer_state s, arraylist_t globals, arrayl`
`1913`	`1913`	`}`
`1914`	`1914`	`for (size_t i = 0; i < external_fns->len; i++) {`
`1915`	`1915`	`jl_code_instance_t ci = (jl_code_instance_t)external_fns->items[i];`
	`1916`	`+ assert(ci && (jl_atomic_load_relaxed(&ci->specsigflags) & 0b001));`
`1916`	`1917`	`uintptr_t item = backref_id(s, (void*)ci, s->link_ids_external_fnvars);`
`1917`	`1918`	`uintptr_t reloc = get_reloc_for_item(item, 0);`
`1918`	`1919`	`write_reloc_t(s->gvar_record, reloc);`
`1919`	`1920`	`}`
`1920`		`- return globals->len + 1;`
	`1921`	`+ return globals->len;`
`1921`	`1922`	`}`
`1922`	`1923`
`1923`	`1924`	`// Pointer relocation for native-code referenced global variables`
`@@ -1962,7 +1963,7 @@ static void jl_root_new_gvars(jl_serializer_state s, jl_image_t image, uint32_`
`1962`	`1963`	`v = (uintptr_t)jl_as_global_root((jl_value_t*)v);`
`1963`	`1964`	`} else {`
`1964`	`1965`	`jl_code_instance_t codeinst = (jl_code_instance_t) v;`
`1965`		`- assert(codeinst && (codeinst->specsigflags & 0b01));`
	`1966`	`+ assert(codeinst && (codeinst->specsigflags & 0b01) && codeinst->specptr.fptr);`
`1966`	`1967`	`v = (uintptr_t)codeinst->specptr.fptr;`
`1967`	`1968`	`}`
`1968`	`1969`	`*gv = v;`