Skip to content

Pool more JIT resources to reduce memory usage/contention #44912

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Apr 12, 2022
Merged

Pool more JIT resources to reduce memory usage/contention #44912

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
e178003
Pool pass managers
pchintalapudi Apr 8, 2022
5da1b14
Condense IR pipeline
pchintalapudi Apr 8, 2022
4f86643
Scope resources better
pchintalapudi Apr 8, 2022
643b549
Pool compiler TargetMachines
pchintalapudi Apr 8, 2022
f56d75d
Rename acquire/release
pchintalapudi Apr 8, 2022
1bdf920
Lock note
pchintalapudi Apr 8, 2022
762ce48
Fix trailing whitespace
pchintalapudi Apr 8, 2022
b1b92fb
Remove old note about PM_mutex
pchintalapudi Apr 8, 2022
2627aec
Address review
pchintalapudi Apr 11, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions doc/src/devdocs/locks.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,11 +28,10 @@ The following are definitely leaf locks (level 1), and must not try to acquire a
> * gc_perm_lock
> * flisp
> * jl_in_stackwalk (Win32)
> * PM_mutex[i]
> * ContextPool::mutex
> * ResourcePool<?>::mutex
>
> > flisp itself is already threadsafe, this lock only protects the `jl_ast_context_list_t` pool
> > likewise, orc::ThreadSafeContexts carry their own lock, the ContextPool::mutex just protects the pool
> > likewise, the ResourcePool<?>::mutexes just protect the associated resource pool

The following is a leaf lock (level 2), and only acquires level 1 locks (safepoint) internally:

Expand Down
185 changes: 112 additions & 73 deletions src/jitlayers.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -494,58 +494,6 @@ static auto countBasicBlocks(const Function &F)
return std::distance(F.begin(), F.end());
}

OptimizerResultT JuliaOJIT::OptimizerT::operator()(orc::ThreadSafeModule TSM, orc::MaterializationResponsibility &R) {
TSM.withModuleDo([&](Module &M) {
uint64_t start_time = 0;
if (dump_llvm_opt_stream != NULL) {
// Print LLVM function statistics _before_ optimization
// Print all the information about this invocation as a YAML object
jl_printf(dump_llvm_opt_stream, "- \n");
// We print the name and some statistics for each function in the module, both
// before optimization and again afterwards.
jl_printf(dump_llvm_opt_stream, " before: \n");
for (auto &F : M.functions()) {
if (F.isDeclaration() || F.getName().startswith("jfptr_")) {
continue;
}
// Each function is printed as a YAML object with several attributes
jl_printf(dump_llvm_opt_stream, " \"%s\":\n", F.getName().str().c_str());
jl_printf(dump_llvm_opt_stream, " instructions: %u\n", F.getInstructionCount());
jl_printf(dump_llvm_opt_stream, " basicblocks: %lu\n", countBasicBlocks(F));
}

start_time = jl_hrtime();
}

JL_TIMING(LLVM_OPT);

{
//Lock around our pass manager
std::lock_guard<std::mutex> lock(this->mutex);
PM.run(M);
}

uint64_t end_time = 0;
if (dump_llvm_opt_stream != NULL) {
end_time = jl_hrtime();
jl_printf(dump_llvm_opt_stream, " time_ns: %" PRIu64 "\n", end_time - start_time);
jl_printf(dump_llvm_opt_stream, " optlevel: %d\n", optlevel);

// Print LLVM function statistics _after_ optimization
jl_printf(dump_llvm_opt_stream, " after: \n");
for (auto &F : M.functions()) {
if (F.isDeclaration() || F.getName().startswith("jfptr_")) {
continue;
}
jl_printf(dump_llvm_opt_stream, " \"%s\":\n", F.getName().str().c_str());
jl_printf(dump_llvm_opt_stream, " instructions: %u\n", F.getInstructionCount());
jl_printf(dump_llvm_opt_stream, " basicblocks: %lu\n", countBasicBlocks(F));
}
}
});
return Expected<orc::ThreadSafeModule>{std::move(TSM)};
}

void JuliaOJIT::OptSelLayerT::emit(std::unique_ptr<orc::MaterializationResponsibility> R, orc::ThreadSafeModule TSM) {
size_t optlevel = ~0ull;
TSM.withModuleDo([&](Module &M) {
Expand All @@ -570,7 +518,7 @@ void JuliaOJIT::OptSelLayerT::emit(std::unique_ptr<orc::MaterializationResponsib
}
});
assert(optlevel != ~0ull && "Failed to select a valid optimization level!");
this->optimizers[optlevel].emit(std::move(R), std::move(TSM));
this->optimizers[optlevel]->OptimizeLayer.emit(std::move(R), std::move(TSM));
}

void jl_register_jit_object(const object::ObjectFile &debugObj,
Expand Down Expand Up @@ -911,6 +859,106 @@ namespace {
.setCodeModel(TM.getCodeModel())
.setCodeGenOptLevel(CodeGenOptLevelFor(optlevel));
}

struct TMCreator {
orc::JITTargetMachineBuilder JTMB;

TMCreator(TargetMachine &TM, int optlevel) : JTMB(createJTMBFromTM(TM, optlevel)) {}

std::unique_ptr<TargetMachine> operator()() {
return cantFail(JTMB.createTargetMachine());
}
};

struct PMCreator {
std::unique_ptr<TargetMachine> TM;
int optlevel;
PMCreator(TargetMachine &TM, int optlevel) : TM(cantFail(createJTMBFromTM(TM, optlevel).createTargetMachine())), optlevel(optlevel) {}
PMCreator(const PMCreator &other) : PMCreator(*other.TM, other.optlevel) {}
PMCreator(PMCreator &&other) : TM(std::move(other.TM)), optlevel(other.optlevel) {}
friend void swap(PMCreator &self, PMCreator &other) {
using std::swap;
swap(self.TM, other.TM);
swap(self.optlevel, other.optlevel);
}
PMCreator &operator=(PMCreator other) {
swap(*this, other);
return *this;
}
std::unique_ptr<legacy::PassManager> operator()() {
auto PM = std::make_unique<legacy::PassManager>();
addPassesForOptLevel(*PM, *TM, optlevel);
return PM;
}
};

struct OptimizerT {
OptimizerT(TargetMachine &TM, int optlevel) : optlevel(optlevel), PMs(PMCreator(TM, optlevel)) {}

OptimizerResultT operator()(orc::ThreadSafeModule TSM, orc::MaterializationResponsibility &R) {
TSM.withModuleDo([&](Module &M) {
uint64_t start_time = 0;
if (dump_llvm_opt_stream != NULL) {
Copy link
Member

Choose a reason for hiding this comment

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

It appears this might need some locking later (for dump_llvm_opt_stream)

Copy link
Member Author

Choose a reason for hiding this comment

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

This should be addressed in the latest commit in #44914 to lock around bundles of stream operations.

// Print LLVM function statistics _before_ optimization
// Print all the information about this invocation as a YAML object
jl_printf(dump_llvm_opt_stream, "- \n");
// We print the name and some statistics for each function in the module, both
// before optimization and again afterwards.
jl_printf(dump_llvm_opt_stream, " before: \n");
for (auto &F : M.functions()) {
if (F.isDeclaration() || F.getName().startswith("jfptr_")) {
continue;
}
// Each function is printed as a YAML object with several attributes
jl_printf(dump_llvm_opt_stream, " \"%s\":\n", F.getName().str().c_str());
jl_printf(dump_llvm_opt_stream, " instructions: %u\n", F.getInstructionCount());
jl_printf(dump_llvm_opt_stream, " basicblocks: %lu\n", countBasicBlocks(F));
}

start_time = jl_hrtime();
}

JL_TIMING(LLVM_OPT);

//Run the optimization
(***PMs).run(M);

uint64_t end_time = 0;
if (dump_llvm_opt_stream != NULL) {
end_time = jl_hrtime();
jl_printf(dump_llvm_opt_stream, " time_ns: %" PRIu64 "\n", end_time - start_time);
jl_printf(dump_llvm_opt_stream, " optlevel: %d\n", optlevel);

// Print LLVM function statistics _after_ optimization
jl_printf(dump_llvm_opt_stream, " after: \n");
for (auto &F : M.functions()) {
if (F.isDeclaration() || F.getName().startswith("jfptr_")) {
continue;
}
jl_printf(dump_llvm_opt_stream, " \"%s\":\n", F.getName().str().c_str());
jl_printf(dump_llvm_opt_stream, " instructions: %u\n", F.getInstructionCount());
jl_printf(dump_llvm_opt_stream, " basicblocks: %lu\n", countBasicBlocks(F));
}
}
});
return Expected<orc::ThreadSafeModule>{std::move(TSM)};
}
private:
int optlevel;
JuliaOJIT::ResourcePool<std::unique_ptr<legacy::PassManager>> PMs;
};

struct CompilerT : orc::IRCompileLayer::IRCompiler {

CompilerT(orc::IRSymbolMapper::ManglingOptions MO, TargetMachine &TM, int optlevel)
: orc::IRCompileLayer::IRCompiler(MO), TMs(TMCreator(TM, optlevel)) {}

Expected<std::unique_ptr<MemoryBuffer>> operator()(Module &M) override {
return orc::SimpleCompiler(***TMs)(M);
}

JuliaOJIT::ResourcePool<std::unique_ptr<TargetMachine>> TMs;
};
}

llvm::DataLayout jl_create_datalayout(TargetMachine &TM) {
Expand All @@ -920,15 +968,14 @@ llvm::DataLayout jl_create_datalayout(TargetMachine &TM) {
return jl_data_layout;
}

JuliaOJIT::PipelineT::PipelineT(orc::ObjectLayer &BaseLayer, TargetMachine &TM, int optlevel)
: CompileLayer(BaseLayer.getExecutionSession(), BaseLayer,
std::make_unique<CompilerT>(orc::irManglingOptionsFromTargetOptions(TM.Options), TM, optlevel)),
OptimizeLayer(CompileLayer.getExecutionSession(), CompileLayer, OptimizerT(TM, optlevel)) {}

JuliaOJIT::JuliaOJIT()
: TM(createTargetMachine()),
DL(jl_create_datalayout(*TM)),
TMs{
cantFail(createJTMBFromTM(*TM, 0).createTargetMachine()),
cantFail(createJTMBFromTM(*TM, 1).createTargetMachine()),
cantFail(createJTMBFromTM(*TM, 2).createTargetMachine()),
cantFail(createJTMBFromTM(*TM, 3).createTargetMachine())
},
#if JL_LLVM_VERSION >= 130000
ES(cantFail(orc::SelfExecutorProcessControl::Create())),
#else
Expand All @@ -955,17 +1002,13 @@ JuliaOJIT::JuliaOJIT()
}
),
#endif
CompileLayer0(ES, ObjectLayer, std::make_unique<orc::ConcurrentIRCompiler>(createJTMBFromTM(*TM, 0))),
CompileLayer1(ES, ObjectLayer, std::make_unique<orc::ConcurrentIRCompiler>(createJTMBFromTM(*TM, 1))),
CompileLayer2(ES, ObjectLayer, std::make_unique<orc::ConcurrentIRCompiler>(createJTMBFromTM(*TM, 2))),
CompileLayer3(ES, ObjectLayer, std::make_unique<orc::ConcurrentIRCompiler>(createJTMBFromTM(*TM, 3))),
OptimizeLayers{
{ES, CompileLayer0, OptimizerT(PM0, PM_mutexes[0], 0)},
{ES, CompileLayer1, OptimizerT(PM1, PM_mutexes[1], 1)},
{ES, CompileLayer2, OptimizerT(PM2, PM_mutexes[2], 2)},
{ES, CompileLayer3, OptimizerT(PM3, PM_mutexes[3], 3)},
Pipelines{
std::make_unique<PipelineT>(ObjectLayer, *TM, 0),
std::make_unique<PipelineT>(ObjectLayer, *TM, 1),
std::make_unique<PipelineT>(ObjectLayer, *TM, 2),
std::make_unique<PipelineT>(ObjectLayer, *TM, 3),
},
OptSelLayer(OptimizeLayers)
OptSelLayer(Pipelines)
{
#ifdef JL_USE_JITLINK
# if defined(_OS_DARWIN_) && defined(LLVM_SHLIB)
Expand All @@ -987,10 +1030,6 @@ JuliaOJIT::JuliaOJIT()
registerRTDyldJITObject(Object, LO, MemMgr);
});
#endif
addPassesForOptLevel(PM0, *TMs[0], 0);
addPassesForOptLevel(PM1, *TMs[1], 1);
addPassesForOptLevel(PM2, *TMs[2], 2);
addPassesForOptLevel(PM3, *TMs[3], 3);

// Make sure SectionMemoryManager::getSymbolAddressInProcess can resolve
// symbols in the program as well. The nullptr argument to the function
Expand Down
63 changes: 26 additions & 37 deletions src/jitlayers.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -193,11 +193,10 @@ class JuliaOJIT {
typedef orc::IRCompileLayer CompileLayerT;
typedef orc::IRTransformLayer OptimizeLayerT;
typedef object::OwningBinary<object::ObjectFile> OwningObj;
private:
template<typename ResourceT, size_t max = 0>
struct ResourcePool {
public:
ResourcePool(function_ref<ResourceT()> creator) : creator(std::move(creator)), mutex(std::make_unique<WNMutex>()) {}
ResourcePool(std::function<ResourceT()> creator) : creator(std::move(creator)), mutex(std::make_unique<WNMutex>()) {}
class OwningResource {
public:
OwningResource(ResourcePool &pool, ResourceT resource) : pool(pool), resource(std::move(resource)) {}
Expand All @@ -206,7 +205,7 @@ class JuliaOJIT {
OwningResource(OwningResource &&) = default;
OwningResource &operator=(OwningResource &&) = default;
~OwningResource() {
if (resource) pool.release_(std::move(*resource));
if (resource) pool.release(std::move(*resource));
}
ResourceT release() {
ResourceT res(std::move(*resource));
Expand Down Expand Up @@ -242,11 +241,15 @@ class JuliaOJIT {
llvm::Optional<ResourceT> resource;
};

OwningResource acquire() {
return OwningResource(*this, acquire_());
OwningResource operator*() {
return OwningResource(*this, acquire());
}

OwningResource get() {
return **this;
}

ResourceT acquire_() {
ResourceT acquire() {
std::unique_lock<std::mutex> lock(mutex->mutex);
if (!pool.empty()) {
return pool.pop_back_val();
Expand All @@ -259,13 +262,13 @@ class JuliaOJIT {
assert(!pool.empty() && "Expected resource pool to have a value!");
return pool.pop_back_val();
}
void release_(ResourceT &&resource) {
void release(ResourceT &&resource) {
std::lock_guard<std::mutex> lock(mutex->mutex);
pool.push_back(std::move(resource));
mutex->empty.notify_one();
}
private:
llvm::function_ref<ResourceT()> creator;
std::function<ResourceT()> creator;
size_t created = 0;
llvm::SmallVector<ResourceT, max == 0 ? 8 : max> pool;
struct WNMutex {
Expand All @@ -275,33 +278,31 @@ class JuliaOJIT {

std::unique_ptr<WNMutex> mutex;
};
struct OptimizerT {
OptimizerT(legacy::PassManager &PM, std::mutex &mutex, int optlevel) : optlevel(optlevel), PM(PM), mutex(mutex) {}

OptimizerResultT operator()(orc::ThreadSafeModule M, orc::MaterializationResponsibility &R);
private:
int optlevel;
legacy::PassManager &PM;
std::mutex &mutex;
struct PipelineT {
PipelineT(orc::ObjectLayer &BaseLayer, TargetMachine &TM, int optlevel);
CompileLayerT CompileLayer;
OptimizeLayerT OptimizeLayer;
};
// Custom object emission notification handler for the JuliaOJIT
template <typename ObjT, typename LoadResult>
void registerObject(const ObjT &Obj, const LoadResult &LO);

struct OptSelLayerT : orc::IRLayer {

template<size_t N>
OptSelLayerT(OptimizeLayerT (&optimizers)[N]) : orc::IRLayer(optimizers[0].getExecutionSession(), optimizers[0].getManglingOptions()), optimizers(optimizers), count(N) {
OptSelLayerT(std::unique_ptr<PipelineT> (&optimizers)[N]) : orc::IRLayer(optimizers[0]->OptimizeLayer.getExecutionSession(), optimizers[0]->OptimizeLayer.getManglingOptions()), optimizers(optimizers), count(N) {
static_assert(N > 0, "Expected array with at least one optimizer!");
}

void emit(std::unique_ptr<orc::MaterializationResponsibility> R, orc::ThreadSafeModule TSM) override;

private:
OptimizeLayerT *optimizers;
std::unique_ptr<PipelineT> *optimizers;
size_t count;
};

private:
// Custom object emission notification handler for the JuliaOJIT
template <typename ObjT, typename LoadResult>
void registerObject(const ObjT &Obj, const LoadResult &LO);

public:

JuliaOJIT();
Expand All @@ -321,13 +322,13 @@ class JuliaOJIT {
uint64_t getFunctionAddress(StringRef Name);
StringRef getFunctionAtAddress(uint64_t Addr, jl_code_instance_t *codeinst);
auto getContext() {
return ContextPool.acquire();
return *ContextPool;
}
orc::ThreadSafeContext acquireContext() {
return ContextPool.acquire_();
return ContextPool.acquire();
}
void releaseContext(orc::ThreadSafeContext &&ctx) {
ContextPool.release_(std::move(ctx));
ContextPool.release(std::move(ctx));
}
const DataLayout& getDataLayout() const;
TargetMachine &getTargetMachine();
Expand All @@ -340,14 +341,6 @@ class JuliaOJIT {

std::unique_ptr<TargetMachine> TM;
DataLayout DL;
// Should be big enough that in the common case, The
// object fits in its entirety
legacy::PassManager PM0; // per-optlevel pass managers
legacy::PassManager PM1;
legacy::PassManager PM2;
legacy::PassManager PM3;
std::mutex PM_mutexes[4];
std::unique_ptr<TargetMachine> TMs[4];

orc::ExecutionSession ES;
orc::JITDylib &GlobalJD;
Expand All @@ -359,11 +352,7 @@ class JuliaOJIT {
std::shared_ptr<RTDyldMemoryManager> MemMgr;
#endif
ObjLayerT ObjectLayer;
CompileLayerT CompileLayer0;
CompileLayerT CompileLayer1;
CompileLayerT CompileLayer2;
CompileLayerT CompileLayer3;
OptimizeLayerT OptimizeLayers[4];
std::unique_ptr<PipelineT> Pipelines[4];
OptSelLayerT OptSelLayer;

DenseMap<void*, std::string> ReverseLocalSymbolTable;
Expand Down