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

* enable using external code in pkgimages This was unintentionally disabled (and incomplete) in the original PR for pkgimages. Co-authored-by: Valentin Churavy <[email protected]>
JuliaLang · Feb 25, 2023 · c136b7e · c136b7e
1 parent 129d451
commit c136b7e
Show file tree

Hide file tree

Showing 6 changed files with 80 additions and 74 deletions.
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;