Skip to content

Commit

Permalink
Reuse incremental JIT compilation for --image-codegen (#50649)
Browse files Browse the repository at this point in the history
We don't need to merge all of the workqueue modules when performing
compilation with `--image-codegen` set, we just need the global variable
initializers to be defined before they're used in one of the modules.
Therefore we can do this by compiling all of the global variable
initializers upfront, so that later references will link them properly.
  • Loading branch information
pchintalapudi authored Jul 26, 2023
1 parent 8c3452f commit ff14eaf
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 28 deletions.
79 changes: 52 additions & 27 deletions src/jitlayers.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -141,11 +141,12 @@ void jl_link_global(GlobalVariable *GV, void *addr) JL_NOTSAFEPOINT
++LinkedGlobals;
Constant *P = literal_static_pointer_val(addr, GV->getValueType());
GV->setInitializer(P);
GV->setDSOLocal(true);
if (jl_options.image_codegen) {
// If we are forcing imaging mode codegen for debugging,
// emit external non-const symbol to avoid LLVM optimizing the code
// similar to non-imaging mode.
GV->setLinkage(GlobalValue::ExternalLinkage);
assert(GV->hasExternalLinkage());
}
else {
GV->setConstant(true);
Expand All @@ -162,6 +163,23 @@ void jl_jit_globals(std::map<void *, GlobalVariable*> &globals) JL_NOTSAFEPOINT
}
}

// used for image_codegen, where we keep all the gvs external
// so we can't jit them directly into each module
static orc::ThreadSafeModule jl_get_globals_module(orc::ThreadSafeContext &ctx, bool imaging_mode, const DataLayout &DL, const Triple &T, std::map<void *, GlobalVariable*> &globals) JL_NOTSAFEPOINT
{
auto lock = ctx.getLock();
auto GTSM = jl_create_ts_module("globals", ctx, imaging_mode, DL, T);
auto GM = GTSM.getModuleUnlocked();
for (auto &global : globals) {
auto GV = global.second;
auto GV2 = new GlobalVariable(*GM, GV->getValueType(), GV->isConstant(), GlobalValue::ExternalLinkage, literal_static_pointer_val(global.first, GV->getValueType()), GV->getName(), nullptr, GV->getThreadLocalMode(), GV->getAddressSpace(), false);
GV2->copyAttributesFrom(GV);
GV2->setDSOLocal(true);
GV2->setAlignment(GV->getAlign());
}
return GTSM;
}

// this generates llvm code for the lambda info
// and adds the result to the jitlayers
// (and the shadow module),
Expand Down Expand Up @@ -211,46 +229,53 @@ static jl_callptr_t _jl_compile_codeinst(

if (params._shared_module)
jl_ExecutionEngine->addModule(orc::ThreadSafeModule(std::move(params._shared_module), params.tsctx));
if (!params.imaging) {
StringMap<orc::ThreadSafeModule*> NewExports;

// In imaging mode, we can't inline global variable initializers in order to preserve
// the fiction that we don't know what loads from the global will return. Thus, we
// need to emit a separate module for the globals before any functions are compiled,
// to ensure that the globals are defined when they are compiled.
if (params.imaging) {
jl_ExecutionEngine->addModule(jl_get_globals_module(params.tsctx, params.imaging, params.DL, params.TargetTriple, params.global_targets));
} else {
StringMap<void*> NewGlobals;
for (auto &global : params.global_targets) {
NewGlobals[global.second->getName()] = global.first;
}
for (auto &def : emitted) {
orc::ThreadSafeModule &TSM = std::get<0>(def.second);
//The underlying context object is still locked because params is not destroyed yet
auto M = TSM.getModuleUnlocked();
for (auto &F : M->global_objects()) {
if (!F.isDeclaration() && F.getLinkage() == GlobalValue::ExternalLinkage) {
NewExports[F.getName()] = &TSM;
}
}
// Let's link all globals here also (for now)
auto M = std::get<0>(def.second).getModuleUnlocked();
for (auto &GV : M->globals()) {
auto InitValue = NewGlobals.find(GV.getName());
if (InitValue != NewGlobals.end()) {
jl_link_global(&GV, InitValue->second);
}
}
}
DenseMap<orc::ThreadSafeModule*, int> Queued;
std::vector<orc::ThreadSafeModule*> Stack;
for (auto &def : emitted) {
// Add the results to the execution engine now
orc::ThreadSafeModule &M = std::get<0>(def.second);
jl_add_to_ee(M, NewExports, Queued, Stack);
assert(Queued.empty() && Stack.empty() && !M);
}
} else {
jl_jit_globals(params.global_targets);
auto main = std::move(emitted[codeinst].first);
for (auto &def : emitted) {
if (def.first != codeinst) {
jl_merge_module(main, std::move(def.second.first));
}

// Collect the exported functions from the emitted modules,
// which form dependencies on which functions need to be
// compiled first. Cycles of functions are compiled together.
// (essentially we compile a DAG of SCCs in reverse topological order,
// if we treat declarations of external functions as edges from declaration
// to definition)
StringMap<orc::ThreadSafeModule*> NewExports;
for (auto &def : emitted) {
orc::ThreadSafeModule &TSM = std::get<0>(def.second);
//The underlying context object is still locked because params is not destroyed yet
auto M = TSM.getModuleUnlocked();
for (auto &F : M->global_objects()) {
if (!F.isDeclaration() && F.getLinkage() == GlobalValue::ExternalLinkage) {
NewExports[F.getName()] = &TSM;
}
}
jl_ExecutionEngine->addModule(std::move(main));
}
DenseMap<orc::ThreadSafeModule*, int> Queued;
std::vector<orc::ThreadSafeModule*> Stack;
for (auto &def : emitted) {
// Add the results to the execution engine now
orc::ThreadSafeModule &M = std::get<0>(def.second);
jl_add_to_ee(M, NewExports, Queued, Stack);
assert(Queued.empty() && Stack.empty() && !M);
}
++CompiledCodeinsts;
MaxWorkqueueSize.updateMax(emitted.size());
Expand Down
3 changes: 2 additions & 1 deletion test/llvmpasses/image-codegen.jl
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,8 @@
# CHECK-NOT: internal global
# CHECK-NOT: private global
# CHECK: jl_global
# CHECK-SAME: = global
# COM: we emit both declarations and definitions, so we may see either style in the IR
# CHECK-SAME: = {{(external )?}}global
# CHECK: julia_f_
# CHECK-NOT: internal global
# CHECK-NOT: private global
Expand Down

4 comments on commit ff14eaf

@nanosoldier
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Executing the daily package evaluation, I will reply here when finished:

@nanosoldier runtests(isdaily = true)

@nanosoldier
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The package evaluation job you requested has completed - possible new issues were detected.
The full report is available.

@vtjnash
Copy link
Member

@vtjnash vtjnash commented on ff14eaf Aug 1, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@nanosoldier runbenchmarks(ALL, isdaily = true)

@nanosoldier
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Your benchmark job has completed - possible performance regressions were detected. A full report can be found here.

Please sign in to comment.