Skip to content

Commit 8d76c71

Browse files
committed
[ORC] Add ThreadSafeModule and ThreadSafeContext wrappers to support concurrent
compilation of IR in the JIT. ThreadSafeContext is a pair of an LLVMContext and a mutex that can be used to lock that context when it needs to be accessed from multiple threads. ThreadSafeModule is a pair of a unique_ptr<Module> and a shared_ptr<ThreadSafeContext>. This allows the lifetime of a ThreadSafeContext to be managed automatically in terms of the ThreadSafeModules that refer to it: Once all modules using a ThreadSafeContext are destructed, and providing the client has not held on to a copy of shared context pointer, the context will be automatically destructed. This scheme is necessary due to the following constraits: (1) We need multiple contexts for multithreaded compilation (at least one per compile thread plus one to store any IR not currently being compiled, though one context per module is simpler). (2) We need to free contexts that are no longer being used so that the JIT does not leak memory over time. (3) Module lifetimes are not predictable (modules are compiled as needed depending on the flow of JIT'd code) so there is no single point where contexts could be reclaimed. JIT clients not using concurrency can safely use one ThreadSafeContext for all ThreadSafeModules. JIT clients who want to be able to compile concurrently should use a different ThreadSafeContext for each module, or call setCloneToNewContextOnEmit on their top-level IRLayer. The former reduces compile latency (since no clone step is needed) at the cost of additional memory overhead for uncompiled modules (as every uncompiled module will duplicate the LLVM types, constants and metadata that have been shared). llvm-svn: 343055
1 parent e06831a commit 8d76c71

File tree

15 files changed

+364
-181
lines changed

15 files changed

+364
-181
lines changed

llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h

Lines changed: 4 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -68,17 +68,14 @@ class CompileOnDemandLayer2 : public IRLayer {
6868
using IndirectStubsManagerBuilder =
6969
std::function<std::unique_ptr<IndirectStubsManager>()>;
7070

71-
using GetAvailableContextFunction = std::function<LLVMContext &()>;
72-
7371
CompileOnDemandLayer2(ExecutionSession &ES, IRLayer &BaseLayer,
7472
JITCompileCallbackManager &CCMgr,
75-
IndirectStubsManagerBuilder BuildIndirectStubsManager,
76-
GetAvailableContextFunction GetAvailableContext);
73+
IndirectStubsManagerBuilder BuildIndirectStubsManager);
7774

78-
Error add(JITDylib &V, VModuleKey K, std::unique_ptr<Module> M) override;
75+
Error add(JITDylib &V, VModuleKey K, ThreadSafeModule TSM) override;
7976

8077
void emit(MaterializationResponsibility R, VModuleKey K,
81-
std::unique_ptr<Module> M) override;
78+
ThreadSafeModule TSM) override;
8279

8380
private:
8481
using StubManagersMap =
@@ -87,15 +84,14 @@ class CompileOnDemandLayer2 : public IRLayer {
8784
IndirectStubsManager &getStubsManager(const JITDylib &JD);
8885

8986
void emitExtractedFunctionsModule(MaterializationResponsibility R,
90-
std::unique_ptr<Module> M);
87+
ThreadSafeModule TSM);
9188

9289
mutable std::mutex CODLayerMutex;
9390

9491
IRLayer &BaseLayer;
9592
JITCompileCallbackManager &CCMgr;
9693
IndirectStubsManagerBuilder BuildIndirectStubsManager;
9794
StubManagersMap StubsMgrs;
98-
GetAvailableContextFunction GetAvailableContext;
9995
};
10096

10197
/// Compile-on-demand layer.

llvm/include/llvm/ExecutionEngine/Orc/IRCompileLayer.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,15 +34,15 @@ class IRCompileLayer2 : public IRLayer {
3434
std::function<Expected<std::unique_ptr<MemoryBuffer>>(Module &)>;
3535

3636
using NotifyCompiledFunction =
37-
std::function<void(VModuleKey K, std::unique_ptr<Module>)>;
37+
std::function<void(VModuleKey K, ThreadSafeModule TSM)>;
3838

3939
IRCompileLayer2(ExecutionSession &ES, ObjectLayer &BaseLayer,
4040
CompileFunction Compile);
4141

4242
void setNotifyCompiled(NotifyCompiledFunction NotifyCompiled);
4343

4444
void emit(MaterializationResponsibility R, VModuleKey K,
45-
std::unique_ptr<Module> M) override;
45+
ThreadSafeModule TSM) override;
4646

4747
private:
4848
mutable std::mutex IRLayerMutex;

llvm/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,8 @@ namespace orc {
2525

2626
class IRTransformLayer2 : public IRLayer {
2727
public:
28-
2928
using TransformFunction =
30-
std::function<Expected<std::unique_ptr<Module>>(std::unique_ptr<Module>)>;
29+
std::function<Expected<ThreadSafeModule>(ThreadSafeModule)>;
3130

3231
IRTransformLayer2(ExecutionSession &ES, IRLayer &BaseLayer,
3332
TransformFunction Transform = identityTransform);
@@ -37,10 +36,10 @@ class IRTransformLayer2 : public IRLayer {
3736
}
3837

3938
void emit(MaterializationResponsibility R, VModuleKey K,
40-
std::unique_ptr<Module> M) override;
39+
ThreadSafeModule TSM) override;
4140

42-
static std::unique_ptr<Module> identityTransform(std::unique_ptr<Module> M) {
43-
return M;
41+
static ThreadSafeModule identityTransform(ThreadSafeModule TSM) {
42+
return TSM;
4443
}
4544

4645
private:

llvm/include/llvm/ExecutionEngine/Orc/LLJIT.h

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
#include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
2222
#include "llvm/ExecutionEngine/Orc/ObjectTransformLayer.h"
2323
#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
24+
#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
2425
#include "llvm/Target/TargetMachine.h"
2526

2627
namespace llvm {
@@ -43,11 +44,11 @@ class LLJIT {
4344
Error defineAbsolute(StringRef Name, JITEvaluatedSymbol Address);
4445

4546
/// Adds an IR module to the given JITDylib.
46-
Error addIRModule(JITDylib &JD, std::unique_ptr<Module> M);
47+
Error addIRModule(JITDylib &JD, ThreadSafeModule TSM);
4748

4849
/// Adds an IR module to the Main JITDylib.
49-
Error addIRModule(std::unique_ptr<Module> M) {
50-
return addIRModule(Main, std::move(M));
50+
Error addIRModule(ThreadSafeModule TSM) {
51+
return addIRModule(Main, std::move(TSM));
5152
}
5253

5354
/// Adds an object file to the given JITDylib.
@@ -119,7 +120,7 @@ class LLLazyJIT : public LLJIT {
119120
public:
120121
/// Create an LLLazyJIT instance.
121122
static Expected<std::unique_ptr<LLLazyJIT>>
122-
Create(std::unique_ptr<TargetMachine> TM, DataLayout DL, LLVMContext &Ctx);
123+
Create(std::unique_ptr<TargetMachine> TM, DataLayout DL);
123124

124125
/// Set an IR transform (e.g. pass manager pipeline) to run on each function
125126
/// when it is compiled.
@@ -128,16 +129,16 @@ class LLLazyJIT : public LLJIT {
128129
}
129130

130131
/// Add a module to be lazily compiled to JITDylib JD.
131-
Error addLazyIRModule(JITDylib &JD, std::unique_ptr<Module> M);
132+
Error addLazyIRModule(JITDylib &JD, ThreadSafeModule M);
132133

133134
/// Add a module to be lazily compiled to the main JITDylib.
134-
Error addLazyIRModule(std::unique_ptr<Module> M) {
135+
Error addLazyIRModule(ThreadSafeModule M) {
135136
return addLazyIRModule(Main, std::move(M));
136137
}
137138

138139
private:
139140
LLLazyJIT(std::unique_ptr<ExecutionSession> ES,
140-
std::unique_ptr<TargetMachine> TM, DataLayout DL, LLVMContext &Ctx,
141+
std::unique_ptr<TargetMachine> TM, DataLayout DL,
141142
std::unique_ptr<JITCompileCallbackManager> CCMgr,
142143
std::function<std::unique_ptr<IndirectStubsManager>()> ISMBuilder);
143144

llvm/include/llvm/ExecutionEngine/Orc/Layer.h

Lines changed: 25 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
#define LLVM_EXECUTIONENGINE_ORC_LAYER_H
1616

1717
#include "llvm/ExecutionEngine/Orc/Core.h"
18+
#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
1819
#include "llvm/IR/Module.h"
1920
#include "llvm/Support/MemoryBuffer.h"
2021

@@ -30,21 +31,36 @@ class IRLayer {
3031
/// Returns the ExecutionSession for this layer.
3132
ExecutionSession &getExecutionSession() { return ES; }
3233

34+
/// Sets the CloneToNewContextOnEmit flag (false by default).
35+
///
36+
/// When set, IR modules added to this layer will be cloned on to a new
37+
/// context before emit is called. This can be used by clients who want
38+
/// to load all IR using one LLVMContext (to save memory via type and
39+
/// constant uniquing), but want to move Modules to fresh contexts before
40+
/// compiling them to enable concurrent compilation.
41+
/// Single threaded clients, or clients who load every module on a new
42+
/// context, need not set this.
43+
void setCloneToNewContextOnEmit(bool CloneToNewContextOnEmit);
44+
45+
/// Returns the current value of the CloneToNewContextOnEmit flag.
46+
bool getCloneToNewContextOnEmit() const { return CloneToNewContextOnEmit; }
47+
3348
/// Adds a MaterializationUnit representing the given IR to the given
3449
/// JITDylib.
35-
virtual Error add(JITDylib &JD, VModuleKey K, std::unique_ptr<Module> M);
50+
virtual Error add(JITDylib &JD, VModuleKey K, ThreadSafeModule TSM);
3651

3752
/// Adds a MaterializationUnit representing the given IR to the main
3853
/// JITDylib.
39-
Error add(VModuleKey K, std::unique_ptr<Module> M) {
40-
return add(ES.getMainJITDylib(), K, std::move(M));
54+
Error add(VModuleKey K, ThreadSafeModule TSM) {
55+
return add(ES.getMainJITDylib(), K, std::move(TSM));
4156
}
4257

4358
/// Emit should materialize the given IR.
4459
virtual void emit(MaterializationResponsibility R, VModuleKey K,
45-
std::unique_ptr<Module> M) = 0;
60+
ThreadSafeModule TSM) = 0;
4661

4762
private:
63+
bool CloneToNewContextOnEmit = false;
4864
ExecutionSession &ES;
4965
};
5066

@@ -58,18 +74,18 @@ class IRMaterializationUnit : public MaterializationUnit {
5874

5975
/// Create an IRMaterializationLayer. Scans the module to build the
6076
/// SymbolFlags and SymbolToDefinition maps.
61-
IRMaterializationUnit(ExecutionSession &ES, std::unique_ptr<Module> M);
77+
IRMaterializationUnit(ExecutionSession &ES, ThreadSafeModule TSM);
6278

6379
/// Create an IRMaterializationLayer from a module, and pre-existing
6480
/// SymbolFlags and SymbolToDefinition maps. The maps must provide
6581
/// entries for each definition in M.
6682
/// This constructor is useful for delegating work from one
6783
/// IRMaterializationUnit to another.
68-
IRMaterializationUnit(std::unique_ptr<Module> M, SymbolFlagsMap SymbolFlags,
84+
IRMaterializationUnit(ThreadSafeModule TSM, SymbolFlagsMap SymbolFlags,
6985
SymbolNameToDefinitionMap SymbolToDefinition);
7086

7187
protected:
72-
std::unique_ptr<Module> M;
88+
ThreadSafeModule TSM;
7389
SymbolNameToDefinitionMap SymbolToDefinition;
7490

7591
private:
@@ -81,7 +97,8 @@ class IRMaterializationUnit : public MaterializationUnit {
8197
class BasicIRLayerMaterializationUnit : public IRMaterializationUnit {
8298
public:
8399
BasicIRLayerMaterializationUnit(IRLayer &L, VModuleKey K,
84-
std::unique_ptr<Module> M);
100+
ThreadSafeModule TSM);
101+
85102
private:
86103

87104
void materialize(MaterializationResponsibility R) override;
Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
//===----------- ThreadSafeModule.h -- Layer interfaces ---------*- C++ -*-===//
2+
//
3+
// The LLVM Compiler Infrastructure
4+
//
5+
// This file is distributed under the University of Illinois Open Source
6+
// License. See LICENSE.TXT for details.
7+
//
8+
//===----------------------------------------------------------------------===//
9+
//
10+
// Thread safe wrappers and utilities for Module and LLVMContext.
11+
//
12+
//===----------------------------------------------------------------------===//
13+
14+
#ifndef LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULEWRAPPER_H
15+
#define LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULEWRAPPER_H
16+
17+
#include "llvm/IR/LLVMContext.h"
18+
#include "llvm/IR/Module.h"
19+
20+
namespace llvm {
21+
namespace orc {
22+
23+
/// An LLVMContext together with an associated mutex that can be used to lock
24+
/// the context to prevent concurrent access by other threads.
25+
class ThreadSafeContext {
26+
private:
27+
28+
struct State {
29+
State(std::unique_ptr<LLVMContext> Ctx)
30+
: Ctx(std::move(Ctx)) {}
31+
32+
std::unique_ptr<LLVMContext> Ctx;
33+
std::recursive_mutex Mutex;
34+
};
35+
36+
public:
37+
38+
// RAII based lock for ThreadSafeContext.
39+
class Lock {
40+
private:
41+
using UnderlyingLock = std::lock_guard<std::recursive_mutex>;
42+
public:
43+
44+
Lock(std::shared_ptr<State> S)
45+
: S(std::move(S)),
46+
L(llvm::make_unique<UnderlyingLock>(this->S->Mutex)) {}
47+
private:
48+
std::shared_ptr<State> S;
49+
std::unique_ptr<UnderlyingLock> L;
50+
};
51+
52+
/// Construct a null context.
53+
ThreadSafeContext() = default;
54+
55+
/// Construct a ThreadSafeContext from the given LLVMContext.
56+
ThreadSafeContext(std::unique_ptr<LLVMContext> NewCtx)
57+
: S(std::make_shared<State>(std::move(NewCtx))) {
58+
assert(S->Ctx != nullptr &&
59+
"Can not construct a ThreadSafeContext from a nullptr");
60+
}
61+
62+
/// Returns a pointer to the LLVMContext that was used to construct this
63+
/// instance, or null if the instance was default constructed.
64+
LLVMContext* getContext() {
65+
return S ? S->Ctx.get() : nullptr;
66+
}
67+
68+
Lock getLock() {
69+
assert(S && "Can not lock an empty ThreadSafeContext");
70+
return Lock(S);
71+
}
72+
73+
private:
74+
std::shared_ptr<State> S;
75+
};
76+
77+
/// An LLVM Module together with a shared ThreadSafeContext.
78+
class ThreadSafeModule {
79+
public:
80+
/// Default construct a ThreadSafeModule. This results in a null module and
81+
/// null context.
82+
ThreadSafeModule() = default;
83+
84+
/// Construct a ThreadSafeModule from a unique_ptr<Module> and a
85+
/// unique_ptr<LLVMContext>. This creates a new ThreadSafeContext from the
86+
/// given context.
87+
ThreadSafeModule(std::unique_ptr<Module> M,
88+
std::unique_ptr<LLVMContext> Ctx)
89+
: M(std::move(M)), TSCtx(std::move(Ctx)) {}
90+
91+
ThreadSafeModule(std::unique_ptr<Module> M,
92+
ThreadSafeContext TSCtx)
93+
: M(std::move(M)), TSCtx(std::move(TSCtx)) {}
94+
95+
Module* getModule() { return M.get(); }
96+
97+
ThreadSafeContext::Lock getContextLock() { return TSCtx.getLock(); }
98+
99+
explicit operator bool() {
100+
if (M) {
101+
assert(TSCtx.getContext() && "Non-null module must have non-null context");
102+
return true;
103+
}
104+
return false;
105+
}
106+
107+
private:
108+
std::unique_ptr<Module> M;
109+
ThreadSafeContext TSCtx;
110+
};
111+
112+
using GVPredicate = std::function<bool(const GlobalValue&)>;
113+
using GVModifier = std::function<void(GlobalValue&)>;
114+
115+
/// Clones the given module on to a new context.
116+
ThreadSafeModule
117+
cloneToNewContext(ThreadSafeModule &TSMW,
118+
GVPredicate ShouldCloneDef = GVPredicate(),
119+
GVModifier UpdateClonedDefSource = GVModifier());
120+
121+
} // End namespace orc
122+
} // End namespace llvm
123+
124+
#endif // LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULEWRAPPER_H

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ add_llvm_library(LLVMOrcJIT
1616
OrcMCJITReplacement.cpp
1717
RPCUtils.cpp
1818
RTDyldObjectLinkingLayer.cpp
19+
ThreadSafeModule.cpp
1920

2021
ADDITIONAL_HEADER_DIRS
2122
${LLVM_MAIN_INCLUDE_DIR}/llvm/ExecutionEngine/Orc

0 commit comments

Comments
 (0)