Skip to content

Commit 8b94274

Browse files
lhameslhames
committed
[ORC] Make the VModuleKey optional, propagate it via MaterializationUnit and
MaterializationResponsibility. VModuleKeys are intended to enable selective removal of modules from a JIT session, however for a wide variety of use cases selective removal is not needed and introduces unnecessary overhead. As of this commit, the default constructed VModuleKey value is reserved as a "do not track" value, and becomes the default when adding a new module to the JIT. This commit also changes the propagation of VModuleKeys. They were passed alongside the MaterializationResponsibity instance in XXLayer::emit methods, but are now propagated as part of the MaterializationResponsibility instance itself (and as part of MaterializationUnit when stored in a JITDylib). Associating VModuleKeys with MaterializationUnits in this way should allow for a thread-safe module removal mechanism in the future, even when a module is in the process of being compiled, by having the MaterializationResponsibility object check in on its VModuleKey's state before commiting its results to the JITDylib. llvm-svn: 344643
1 parent 4084df0 commit 8b94274

21 files changed

+124
-117
lines changed

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -94,8 +94,7 @@ class CompileOnDemandLayer : public IRLayer {
9494

9595
/// Emits the given module. This should not be called by clients: it will be
9696
/// called by the JIT when a definition added via the add method is requested.
97-
void emit(MaterializationResponsibility R, VModuleKey K,
98-
ThreadSafeModule TSM) override;
97+
void emit(MaterializationResponsibility R, ThreadSafeModule TSM) override;
9998

10099
private:
101100
struct PerDylibResources {

‎llvm/include/llvm/ExecutionEngine/Orc/Core.h

Lines changed: 25 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,9 @@ class MaterializationResponsibility {
168168
/// into.
169169
JITDylib &getTargetJITDylib() const { return JD; }
170170

171+
/// Returns the VModuleKey for this instance.
172+
VModuleKey getVModuleKey() const { return K; }
173+
171174
/// Returns the symbol flags map for this responsibility instance.
172175
/// Note: The returned flags may have transient flags (Lazy, Materializing)
173176
/// set. These should be stripped with JITSymbolFlags::stripTransientFlags
@@ -218,7 +221,8 @@ class MaterializationResponsibility {
218221
/// Delegates responsibility for the given symbols to the returned
219222
/// materialization responsibility. Useful for breaking up work between
220223
/// threads, or different kinds of materialization processes.
221-
MaterializationResponsibility delegate(const SymbolNameSet &Symbols);
224+
MaterializationResponsibility delegate(const SymbolNameSet &Symbols,
225+
VModuleKey NewKey = VModuleKey());
222226

223227
void addDependencies(const SymbolStringPtr &Name,
224228
const SymbolDependenceMap &Dependencies);
@@ -229,10 +233,12 @@ class MaterializationResponsibility {
229233
private:
230234
/// Create a MaterializationResponsibility for the given JITDylib and
231235
/// initial symbols.
232-
MaterializationResponsibility(JITDylib &JD, SymbolFlagsMap SymbolFlags);
236+
MaterializationResponsibility(JITDylib &JD, SymbolFlagsMap SymbolFlags,
237+
VModuleKey K);
233238

234239
JITDylib &JD;
235240
SymbolFlagsMap SymbolFlags;
241+
VModuleKey K;
236242
};
237243

238244
/// A MaterializationUnit represents a set of symbol definitions that can
@@ -245,8 +251,8 @@ class MaterializationResponsibility {
245251
/// stronger definition is added or already present.
246252
class MaterializationUnit {
247253
public:
248-
MaterializationUnit(SymbolFlagsMap InitalSymbolFlags)
249-
: SymbolFlags(std::move(InitalSymbolFlags)) {}
254+
MaterializationUnit(SymbolFlagsMap InitalSymbolFlags, VModuleKey K)
255+
: SymbolFlags(std::move(InitalSymbolFlags)), K(std::move(K)) {}
250256

251257
virtual ~MaterializationUnit() {}
252258

@@ -261,7 +267,8 @@ class MaterializationUnit {
261267
/// ExecutionSession::DispatchMaterializationFunction) to trigger
262268
/// materialization of this MaterializationUnit.
263269
void doMaterialize(JITDylib &JD) {
264-
materialize(MaterializationResponsibility(JD, std::move(SymbolFlags)));
270+
materialize(MaterializationResponsibility(JD, std::move(SymbolFlags),
271+
std::move(K)));
265272
}
266273

267274
/// Called by JITDylibs to notify MaterializationUnits that the given symbol
@@ -273,6 +280,7 @@ class MaterializationUnit {
273280

274281
protected:
275282
SymbolFlagsMap SymbolFlags;
283+
VModuleKey K;
276284

277285
private:
278286
virtual void anchor();
@@ -298,7 +306,7 @@ using MaterializationUnitList =
298306
/// materialized.
299307
class AbsoluteSymbolsMaterializationUnit : public MaterializationUnit {
300308
public:
301-
AbsoluteSymbolsMaterializationUnit(SymbolMap Symbols);
309+
AbsoluteSymbolsMaterializationUnit(SymbolMap Symbols, VModuleKey K);
302310

303311
StringRef getName() const override;
304312

@@ -321,9 +329,9 @@ class AbsoluteSymbolsMaterializationUnit : public MaterializationUnit {
321329
/// \endcode
322330
///
323331
inline std::unique_ptr<AbsoluteSymbolsMaterializationUnit>
324-
absoluteSymbols(SymbolMap Symbols) {
332+
absoluteSymbols(SymbolMap Symbols, VModuleKey K = VModuleKey()) {
325333
return llvm::make_unique<AbsoluteSymbolsMaterializationUnit>(
326-
std::move(Symbols));
334+
std::move(Symbols), std::move(K));
327335
}
328336

329337
struct SymbolAliasMapEntry {
@@ -349,7 +357,8 @@ class ReExportsMaterializationUnit : public MaterializationUnit {
349357
/// Note: Care must be taken that no sets of aliases form a cycle, as such
350358
/// a cycle will result in a deadlock when any symbol in the cycle is
351359
/// resolved.
352-
ReExportsMaterializationUnit(JITDylib *SourceJD, SymbolAliasMap Aliases);
360+
ReExportsMaterializationUnit(JITDylib *SourceJD, SymbolAliasMap Aliases,
361+
VModuleKey K);
353362

354363
StringRef getName() const override;
355364

@@ -374,17 +383,18 @@ class ReExportsMaterializationUnit : public MaterializationUnit {
374383
/// return Err;
375384
/// \endcode
376385
inline std::unique_ptr<ReExportsMaterializationUnit>
377-
symbolAliases(SymbolAliasMap Aliases) {
378-
return llvm::make_unique<ReExportsMaterializationUnit>(nullptr,
379-
std::move(Aliases));
386+
symbolAliases(SymbolAliasMap Aliases, VModuleKey K = VModuleKey()) {
387+
return llvm::make_unique<ReExportsMaterializationUnit>(
388+
nullptr, std::move(Aliases), std::move(K));
380389
}
381390

382391
/// Create a materialization unit for re-exporting symbols from another JITDylib
383392
/// with alternative names/flags.
384393
inline std::unique_ptr<ReExportsMaterializationUnit>
385-
reexports(JITDylib &SourceJD, SymbolAliasMap Aliases) {
386-
return llvm::make_unique<ReExportsMaterializationUnit>(&SourceJD,
387-
std::move(Aliases));
394+
reexports(JITDylib &SourceJD, SymbolAliasMap Aliases,
395+
VModuleKey K = VModuleKey()) {
396+
return llvm::make_unique<ReExportsMaterializationUnit>(
397+
&SourceJD, std::move(Aliases), std::move(K));
388398
}
389399

390400
/// Build a SymbolAliasMap for the common case where you want to re-export

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -41,8 +41,7 @@ class IRCompileLayer : public IRLayer {
4141

4242
void setNotifyCompiled(NotifyCompiledFunction NotifyCompiled);
4343

44-
void emit(MaterializationResponsibility R, VModuleKey K,
45-
ThreadSafeModule TSM) override;
44+
void emit(MaterializationResponsibility R, ThreadSafeModule TSM) override;
4645

4746
private:
4847
mutable std::mutex IRLayerMutex;

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

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,7 @@ class IRTransformLayer : public IRLayer {
3535
this->Transform = std::move(Transform);
3636
}
3737

38-
void emit(MaterializationResponsibility R, VModuleKey K,
39-
ThreadSafeModule TSM) override;
38+
void emit(MaterializationResponsibility R, ThreadSafeModule TSM) override;
4039

4140
static ThreadSafeModule
4241
identityTransform(ThreadSafeModule TSM,

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

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -111,8 +111,6 @@ class LLJIT {
111111
LLJIT(std::unique_ptr<ExecutionSession> ES, JITTargetMachineBuilder JTMB,
112112
DataLayout DL, unsigned NumCompileThreads);
113113

114-
std::unique_ptr<RuntimeDyld::MemoryManager> getMemoryManager(VModuleKey K);
115-
116114
std::string mangle(StringRef UnmangledName);
117115

118116
Error applyDataLayout(Module &M);

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

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -49,11 +49,11 @@ class IRLayer {
4949

5050
/// Adds a MaterializationUnit representing the given IR to the given
5151
/// JITDylib.
52-
virtual Error add(JITDylib &JD, VModuleKey K, ThreadSafeModule TSM);
52+
virtual Error add(JITDylib &JD, ThreadSafeModule TSM,
53+
VModuleKey K = VModuleKey());
5354

5455
/// Emit should materialize the given IR.
55-
virtual void emit(MaterializationResponsibility R, VModuleKey K,
56-
ThreadSafeModule TSM) = 0;
56+
virtual void emit(MaterializationResponsibility R, ThreadSafeModule TSM) = 0;
5757

5858
private:
5959
bool CloneToNewContextOnEmit = false;
@@ -70,14 +70,16 @@ class IRMaterializationUnit : public MaterializationUnit {
7070

7171
/// Create an IRMaterializationLayer. Scans the module to build the
7272
/// SymbolFlags and SymbolToDefinition maps.
73-
IRMaterializationUnit(ExecutionSession &ES, ThreadSafeModule TSM);
73+
IRMaterializationUnit(ExecutionSession &ES, ThreadSafeModule TSM,
74+
VModuleKey K);
7475

7576
/// Create an IRMaterializationLayer from a module, and pre-existing
7677
/// SymbolFlags and SymbolToDefinition maps. The maps must provide
7778
/// entries for each definition in M.
7879
/// This constructor is useful for delegating work from one
7980
/// IRMaterializationUnit to another.
80-
IRMaterializationUnit(ThreadSafeModule TSM, SymbolFlagsMap SymbolFlags,
81+
IRMaterializationUnit(ThreadSafeModule TSM, VModuleKey K,
82+
SymbolFlagsMap SymbolFlags,
8183
SymbolNameToDefinitionMap SymbolToDefinition);
8284

8385
/// Return the ModuleIdentifier as the name for this MaterializationUnit.
@@ -119,10 +121,11 @@ class ObjectLayer {
119121

120122
/// Adds a MaterializationUnit representing the given IR to the given
121123
/// JITDylib.
122-
virtual Error add(JITDylib &JD, VModuleKey K, std::unique_ptr<MemoryBuffer> O);
124+
virtual Error add(JITDylib &JD, std::unique_ptr<MemoryBuffer> O,
125+
VModuleKey K = VModuleKey());
123126

124127
/// Emit should materialize the given IR.
125-
virtual void emit(MaterializationResponsibility R, VModuleKey K,
128+
virtual void emit(MaterializationResponsibility R,
126129
std::unique_ptr<MemoryBuffer> O) = 0;
127130

128131
private:
@@ -149,7 +152,6 @@ class BasicObjectLayerMaterializationUnit : public MaterializationUnit {
149152
void discard(const JITDylib &JD, const SymbolStringPtr &Name) override;
150153

151154
ObjectLayer &L;
152-
VModuleKey K;
153155
std::unique_ptr<MemoryBuffer> O;
154156
};
155157

‎llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -159,7 +159,8 @@ class LazyReexportsMaterializationUnit : public MaterializationUnit {
159159
LazyReexportsMaterializationUnit(LazyCallThroughManager &LCTManager,
160160
IndirectStubsManager &ISManager,
161161
JITDylib &SourceJD,
162-
SymbolAliasMap CallableAliases);
162+
SymbolAliasMap CallableAliases,
163+
VModuleKey K);
163164

164165
StringRef getName() const override;
165166

@@ -182,9 +183,10 @@ class LazyReexportsMaterializationUnit : public MaterializationUnit {
182183
inline std::unique_ptr<LazyReexportsMaterializationUnit>
183184
lazyReexports(LazyCallThroughManager &LCTManager,
184185
IndirectStubsManager &ISManager, JITDylib &SourceJD,
185-
SymbolAliasMap CallableAliases) {
186+
SymbolAliasMap CallableAliases, VModuleKey K = VModuleKey()) {
186187
return llvm::make_unique<LazyReexportsMaterializationUnit>(
187-
LCTManager, ISManager, SourceJD, std::move(CallableAliases));
188+
LCTManager, ISManager, SourceJD, std::move(CallableAliases),
189+
std::move(K));
188190
}
189191

190192
} // End namespace orc

‎llvm/include/llvm/ExecutionEngine/Orc/ObjectTransformLayer.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ class ObjectTransformLayer : public ObjectLayer {
3232
ObjectTransformLayer(ExecutionSession &ES, ObjectLayer &BaseLayer,
3333
TransformFunction Transform);
3434

35-
void emit(MaterializationResponsibility R, VModuleKey K,
35+
void emit(MaterializationResponsibility R,
3636
std::unique_ptr<MemoryBuffer> O) override;
3737

3838
private:

‎llvm/include/llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ class RTDyldObjectLinkingLayer : public ObjectLayer {
4747
using NotifyEmittedFunction = std::function<void(VModuleKey)>;
4848

4949
using GetMemoryManagerFunction =
50-
std::function<std::unique_ptr<RuntimeDyld::MemoryManager>(VModuleKey)>;
50+
std::function<std::unique_ptr<RuntimeDyld::MemoryManager>()>;
5151

5252
/// Construct an ObjectLinkingLayer with the given NotifyLoaded,
5353
/// and NotifyEmitted functors.
@@ -57,7 +57,7 @@ class RTDyldObjectLinkingLayer : public ObjectLayer {
5757
NotifyEmittedFunction NotifyEmitted = NotifyEmittedFunction());
5858

5959
/// Emit the object.
60-
void emit(MaterializationResponsibility R, VModuleKey K,
60+
void emit(MaterializationResponsibility R,
6161
std::unique_ptr<MemoryBuffer> O) override;
6262

6363
/// Set the 'ProcessAllSections' flag.
@@ -118,7 +118,7 @@ class RTDyldObjectLinkingLayer : public ObjectLayer {
118118
bool ProcessAllSections = false;
119119
bool OverrideObjectFlags = false;
120120
bool AutoClaimObjectSymbols = false;
121-
std::map<VModuleKey, std::shared_ptr<RuntimeDyld::MemoryManager>> MemMgrs;
121+
std::vector<std::unique_ptr<RuntimeDyld::MemoryManager>> MemMgrs;
122122
};
123123

124124
class LegacyRTDyldObjectLinkingLayerBase {

‎llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -68,14 +68,16 @@ namespace orc {
6868
class PartitioningIRMaterializationUnit : public IRMaterializationUnit {
6969
public:
7070
PartitioningIRMaterializationUnit(ExecutionSession &ES, ThreadSafeModule TSM,
71-
CompileOnDemandLayer &Parent)
72-
: IRMaterializationUnit(ES, std::move(TSM)), Parent(Parent) {}
71+
VModuleKey K, CompileOnDemandLayer &Parent)
72+
: IRMaterializationUnit(ES, std::move(TSM), std::move(K)),
73+
Parent(Parent) {}
7374

7475
PartitioningIRMaterializationUnit(
7576
ThreadSafeModule TSM, SymbolFlagsMap SymbolFlags,
7677
SymbolNameToDefinitionMap SymbolToDefinition,
7778
CompileOnDemandLayer &Parent)
78-
: IRMaterializationUnit(std::move(TSM), std::move(SymbolFlags),
79+
: IRMaterializationUnit(std::move(TSM), std::move(K),
80+
std::move(SymbolFlags),
7981
std::move(SymbolToDefinition)),
8082
Parent(Parent) {}
8183

@@ -116,8 +118,8 @@ void CompileOnDemandLayer::setPartitionFunction(PartitionFunction Partition) {
116118
this->Partition = std::move(Partition);
117119
}
118120

119-
void CompileOnDemandLayer::emit(MaterializationResponsibility R, VModuleKey K,
120-
ThreadSafeModule TSM) {
121+
void CompileOnDemandLayer::emit(MaterializationResponsibility R,
122+
ThreadSafeModule TSM) {
121123
assert(TSM.getModule() && "Null module");
122124

123125
auto &ES = getExecutionSession();
@@ -149,7 +151,7 @@ void CompileOnDemandLayer::emit(MaterializationResponsibility R, VModuleKey K,
149151
// implementation dylib.
150152
if (auto Err = PDR.getImplDylib().define(
151153
llvm::make_unique<PartitioningIRMaterializationUnit>(
152-
ES, std::move(TSM), *this))) {
154+
ES, std::move(TSM), R.getVModuleKey(), *this))) {
153155
ES.reportError(std::move(Err));
154156
R.failMaterialization();
155157
return;
@@ -245,7 +247,7 @@ void CompileOnDemandLayer::emitPartition(
245247
// unmodified to the base layer.
246248
if (GVsToExtract == None) {
247249
Defs.clear();
248-
BaseLayer.emit(std::move(R), ES.allocateVModule(), std::move(TSM));
250+
BaseLayer.emit(std::move(R), std::move(TSM));
249251
return;
250252
}
251253

@@ -285,9 +287,9 @@ void CompileOnDemandLayer::emitPartition(
285287

286288
auto ExtractedTSM = extractSubModule(TSM, ".submodule", ShouldExtract);
287289
R.replace(llvm::make_unique<PartitioningIRMaterializationUnit>(
288-
ES, std::move(TSM), *this));
290+
ES, std::move(TSM), R.getVModuleKey(), *this));
289291

290-
BaseLayer.emit(std::move(R), ES.allocateVModule(), std::move(ExtractedTSM));
292+
BaseLayer.emit(std::move(R), std::move(ExtractedTSM));
291293
}
292294

293295
} // end namespace orc

‎llvm/lib/ExecutionEngine/Orc/Core.cpp

Lines changed: 16 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -368,8 +368,8 @@ void AsynchronousSymbolQuery::detach() {
368368
}
369369

370370
MaterializationResponsibility::MaterializationResponsibility(
371-
JITDylib &JD, SymbolFlagsMap SymbolFlags)
372-
: JD(JD), SymbolFlags(std::move(SymbolFlags)) {
371+
JITDylib &JD, SymbolFlagsMap SymbolFlags, VModuleKey K)
372+
: JD(JD), SymbolFlags(std::move(SymbolFlags)), K(std::move(K)) {
373373
assert(!this->SymbolFlags.empty() && "Materializing nothing?");
374374

375375
#ifndef NDEBUG
@@ -461,7 +461,12 @@ void MaterializationResponsibility::replace(
461461
}
462462

463463
MaterializationResponsibility
464-
MaterializationResponsibility::delegate(const SymbolNameSet &Symbols) {
464+
MaterializationResponsibility::delegate(const SymbolNameSet &Symbols,
465+
VModuleKey NewKey) {
466+
467+
if (NewKey == VModuleKey())
468+
NewKey = K;
469+
465470
SymbolFlagsMap DelegatedFlags;
466471

467472
for (auto &Name : Symbols) {
@@ -474,7 +479,8 @@ MaterializationResponsibility::delegate(const SymbolNameSet &Symbols) {
474479
SymbolFlags.erase(I);
475480
}
476481

477-
return MaterializationResponsibility(JD, std::move(DelegatedFlags));
482+
return MaterializationResponsibility(JD, std::move(DelegatedFlags),
483+
std::move(NewKey));
478484
}
479485

480486
void MaterializationResponsibility::addDependencies(
@@ -491,8 +497,9 @@ void MaterializationResponsibility::addDependenciesForAll(
491497
}
492498

493499
AbsoluteSymbolsMaterializationUnit::AbsoluteSymbolsMaterializationUnit(
494-
SymbolMap Symbols)
495-
: MaterializationUnit(extractFlags(Symbols)), Symbols(std::move(Symbols)) {}
500+
SymbolMap Symbols, VModuleKey K)
501+
: MaterializationUnit(extractFlags(Symbols), std::move(K)),
502+
Symbols(std::move(Symbols)) {}
496503

497504
StringRef AbsoluteSymbolsMaterializationUnit::getName() const {
498505
return "<Absolute Symbols>";
@@ -519,9 +526,9 @@ AbsoluteSymbolsMaterializationUnit::extractFlags(const SymbolMap &Symbols) {
519526
}
520527

521528
ReExportsMaterializationUnit::ReExportsMaterializationUnit(
522-
JITDylib *SourceJD, SymbolAliasMap Aliases)
523-
: MaterializationUnit(extractFlags(Aliases)), SourceJD(SourceJD),
524-
Aliases(std::move(Aliases)) {}
529+
JITDylib *SourceJD, SymbolAliasMap Aliases, VModuleKey K)
530+
: MaterializationUnit(extractFlags(Aliases), std::move(K)),
531+
SourceJD(SourceJD), Aliases(std::move(Aliases)) {}
525532

526533
StringRef ReExportsMaterializationUnit::getName() const {
527534
return "<Reexports>";

0 commit comments

Comments
 (0)