Skip to content

Commit

Permalink
enable re-using external code in pkgimages (#48723)
Browse files Browse the repository at this point in the history 
* 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]>
  • Loading branch information
@vtjnash @vchuravy @pull
2 people authored and pull[bot] committed Dec 20, 2023
1 parent 9b73701 commit 834dc26
Show file tree
Hide file tree
Showing 6 changed files with 80 additions and 74 deletions.
4 changes: 2 additions & 2 deletions base/linking.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
24 changes: 3 additions & 21 deletions src/aotcompile.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
111 changes: 66 additions & 45 deletions src/codegen.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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++) {
Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}

Expand Down Expand Up @@ -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;
}
}
Expand All @@ -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?
Expand All @@ -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));
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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);
Expand Down
2 changes: 1 addition & 1 deletion src/jitlayers.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
4 changes: 3 additions & 1 deletion src/julia.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
9 changes: 5 additions & 4 deletions src/staticdata.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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;
Expand All @@ -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
Expand Down Expand Up @@ -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;
Expand Down

0 comments on commit 834dc26

Please sign in to comment.