Skip to content

Commit 04af63b

Browse files
committed
[ORC] Introduce IRPartitionLayer for common partition functionality.
1 parent 2c01b27 commit 04af63b

File tree

10 files changed

+417
-322
lines changed

10 files changed

+417
-322
lines changed

llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
2222
#include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
2323
#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
24+
#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
2425
#include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
2526
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
2627
#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
@@ -48,6 +49,7 @@ class KaleidoscopeJIT {
4849
RTDyldObjectLinkingLayer ObjectLayer;
4950
IRCompileLayer CompileLayer;
5051
IRTransformLayer OptimizeLayer;
52+
IRPartitionLayer IPLayer;
5153
CompileOnDemandLayer CODLayer;
5254

5355
JITDylib &MainJD;
@@ -68,8 +70,8 @@ class KaleidoscopeJIT {
6870
CompileLayer(*this->ES, ObjectLayer,
6971
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
7072
OptimizeLayer(*this->ES, CompileLayer, optimizeModule),
71-
CODLayer(*this->ES, OptimizeLayer,
72-
this->EPCIU->getLazyCallThroughManager(),
73+
IPLayer(*this->ES, OptimizeLayer),
74+
CODLayer(*this->ES, IPLayer, this->EPCIU->getLazyCallThroughManager(),
7375
[this] { return this->EPCIU->createIndirectStubsManager(); }),
7476
MainJD(this->ES->createBareJITDylib("<main>")) {
7577
MainJD.addGenerator(

llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
#include "llvm/ExecutionEngine/Orc/Core.h"
44
#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
55
#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
6+
#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
67
#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
78
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
89
#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
@@ -109,13 +110,14 @@ class SpeculativeJIT {
109110
IndirectStubsManagerBuilderFunction ISMBuilder,
110111
std::unique_ptr<DynamicLibrarySearchGenerator> ProcessSymbolsGenerator)
111112
: ES(std::move(ES)), DL(std::move(DL)),
112-
MainJD(this->ES->createBareJITDylib("<main>")), LCTMgr(std::move(LCTMgr)),
113+
MainJD(this->ES->createBareJITDylib("<main>")),
114+
LCTMgr(std::move(LCTMgr)),
113115
CompileLayer(*this->ES, ObjLayer,
114116
std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
115117
S(Imps, *this->ES),
116118
SpeculateLayer(*this->ES, CompileLayer, S, Mangle, BlockFreqQuery()),
117-
CODLayer(*this->ES, SpeculateLayer, *this->LCTMgr,
118-
std::move(ISMBuilder)) {
119+
IPLayer(*this->ES, SpeculateLayer),
120+
CODLayer(*this->ES, IPLayer, *this->LCTMgr, std::move(ISMBuilder)) {
119121
MainJD.addGenerator(std::move(ProcessSymbolsGenerator));
120122
this->CODLayer.setImplMap(&Imps);
121123
ExitOnErr(S.addSpeculationRuntime(MainJD, Mangle));
@@ -141,6 +143,7 @@ class SpeculativeJIT {
141143
Speculator S;
142144
RTDyldObjectLinkingLayer ObjLayer{*ES, createMemMgr};
143145
IRSpeculationLayer SpeculateLayer;
146+
IRPartitionLayer IPLayer;
144147
CompileOnDemandLayer CODLayer;
145148
};
146149

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

Lines changed: 2 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -53,37 +53,15 @@ namespace llvm {
5353
namespace orc {
5454

5555
class CompileOnDemandLayer : public IRLayer {
56-
friend class PartitioningIRMaterializationUnit;
57-
5856
public:
5957
/// Builder for IndirectStubsManagers.
6058
using IndirectStubsManagerBuilder =
6159
std::function<std::unique_ptr<IndirectStubsManager>()>;
6260

63-
using GlobalValueSet = std::set<const GlobalValue *>;
64-
65-
/// Partitioning function.
66-
using PartitionFunction =
67-
std::function<std::optional<GlobalValueSet>(GlobalValueSet Requested)>;
68-
69-
/// Off-the-shelf partitioning which compiles all requested symbols (usually
70-
/// a single function at a time).
71-
static std::optional<GlobalValueSet>
72-
compileRequested(GlobalValueSet Requested);
73-
74-
/// Off-the-shelf partitioning which compiles whole modules whenever any
75-
/// symbol in them is requested.
76-
static std::optional<GlobalValueSet>
77-
compileWholeModule(GlobalValueSet Requested);
78-
7961
/// Construct a CompileOnDemandLayer.
8062
CompileOnDemandLayer(ExecutionSession &ES, IRLayer &BaseLayer,
81-
LazyCallThroughManager &LCTMgr,
82-
IndirectStubsManagerBuilder BuildIndirectStubsManager);
83-
84-
/// Sets the partition function.
85-
void setPartitionFunction(PartitionFunction Partition);
86-
63+
LazyCallThroughManager &LCTMgr,
64+
IndirectStubsManagerBuilder BuildIndirectStubsManager);
8765
/// Sets the ImplSymbolMap
8866
void setImplMap(ImplSymbolMap *Imp);
8967

@@ -110,22 +88,12 @@ class CompileOnDemandLayer : public IRLayer {
11088

11189
PerDylibResources &getPerDylibResources(JITDylib &TargetD);
11290

113-
void cleanUpModule(Module &M);
114-
115-
void expandPartition(GlobalValueSet &Partition);
116-
117-
void emitPartition(std::unique_ptr<MaterializationResponsibility> R,
118-
ThreadSafeModule TSM,
119-
IRMaterializationUnit::SymbolNameToDefinitionMap Defs);
120-
12191
mutable std::mutex CODLayerMutex;
12292

12393
IRLayer &BaseLayer;
12494
LazyCallThroughManager &LCTMgr;
12595
IndirectStubsManagerBuilder BuildIndirectStubsManager;
12696
PerDylibResourcesMap DylibResources;
127-
PartitionFunction Partition = compileRequested;
128-
SymbolLinkagePromoter PromoteSymbols;
12997
ImplSymbolMap *AliaseeImpls = nullptr;
13098
};
13199

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
//===- IRPartitionLayer.h - Partition IR module on lookup -------*- C++ -*-===//
2+
//
3+
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4+
// See https://llvm.org/LICENSE.txt for license information.
5+
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6+
//
7+
//===----------------------------------------------------------------------===//
8+
//
9+
// JIT layer for breaking up modules into smaller submodules that only contains
10+
// looked up symbols.
11+
//
12+
//===----------------------------------------------------------------------===//
13+
14+
#ifndef LLVM_EXECUTIONENGINE_ORC_IRPARTITIONLAYER_H
15+
#define LLVM_EXECUTIONENGINE_ORC_IRPARTITIONLAYER_H
16+
17+
#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
18+
#include "llvm/ExecutionEngine/Orc/Layer.h"
19+
#include "llvm/IR/Attributes.h"
20+
#include "llvm/IR/Constant.h"
21+
#include "llvm/IR/Constants.h"
22+
#include "llvm/IR/DataLayout.h"
23+
#include "llvm/IR/Function.h"
24+
#include "llvm/IR/GlobalAlias.h"
25+
#include "llvm/IR/GlobalValue.h"
26+
#include "llvm/IR/GlobalVariable.h"
27+
#include "llvm/IR/Instruction.h"
28+
#include "llvm/IR/Mangler.h"
29+
#include "llvm/IR/Module.h"
30+
#include "llvm/IR/Type.h"
31+
32+
namespace llvm {
33+
namespace orc {
34+
35+
/// A layer that breaks up IR modules into smaller submodules that only contains
36+
/// looked up symbols.
37+
class IRPartitionLayer : public IRLayer {
38+
friend class PartitioningIRMaterializationUnit;
39+
40+
public:
41+
using GlobalValueSet = std::set<const GlobalValue *>;
42+
43+
/// Partitioning function.
44+
using PartitionFunction =
45+
std::function<std::optional<GlobalValueSet>(GlobalValueSet Requested)>;
46+
47+
/// Construct a IRPartitionLayer.
48+
IRPartitionLayer(ExecutionSession &ES, IRLayer &BaseLayer);
49+
50+
/// Off-the-shelf partitioning which compiles all requested symbols (usually
51+
/// a single function at a time).
52+
static std::optional<GlobalValueSet>
53+
compileRequested(GlobalValueSet Requested);
54+
55+
/// Off-the-shelf partitioning which compiles whole modules whenever any
56+
/// symbol in them is requested.
57+
static std::optional<GlobalValueSet>
58+
compileWholeModule(GlobalValueSet Requested);
59+
60+
/// Sets the partition function.
61+
void setPartitionFunction(PartitionFunction Partition);
62+
63+
/// Emits the given module. This should not be called by clients: it will be
64+
/// called by the JIT when a definition added via the add method is requested.
65+
void emit(std::unique_ptr<MaterializationResponsibility> R,
66+
ThreadSafeModule TSM) override;
67+
68+
private:
69+
void cleanUpModule(Module &M);
70+
71+
void expandPartition(GlobalValueSet &Partition);
72+
73+
void emitPartition(std::unique_ptr<MaterializationResponsibility> R,
74+
ThreadSafeModule TSM,
75+
IRMaterializationUnit::SymbolNameToDefinitionMap Defs);
76+
77+
IRLayer &BaseLayer;
78+
PartitionFunction Partition = compileRequested;
79+
SymbolLinkagePromoter PromoteSymbols;
80+
};
81+
82+
} // namespace orc
83+
} // namespace llvm
84+
85+
#endif

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

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
1919
#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
2020
#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
21+
#include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
2122
#include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
2223
#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
2324
#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
@@ -271,9 +272,8 @@ class LLLazyJIT : public LLJIT {
271272
public:
272273

273274
/// Sets the partition function.
274-
void
275-
setPartitionFunction(CompileOnDemandLayer::PartitionFunction Partition) {
276-
CODLayer->setPartitionFunction(std::move(Partition));
275+
void setPartitionFunction(IRPartitionLayer::PartitionFunction Partition) {
276+
IPLayer->setPartitionFunction(std::move(Partition));
277277
}
278278

279279
/// Returns a reference to the on-demand layer.
@@ -293,6 +293,7 @@ class LLLazyJIT : public LLJIT {
293293
LLLazyJIT(LLLazyJITBuilderState &S, Error &Err);
294294

295295
std::unique_ptr<LazyCallThroughManager> LCTMgr;
296+
std::unique_ptr<IRPartitionLayer> IPLayer;
296297
std::unique_ptr<CompileOnDemandLayer> CODLayer;
297298
};
298299

llvm/lib/ExecutionEngine/Orc/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ add_llvm_component_library(LLVMOrcJIT
2626
IndirectionUtils.cpp
2727
IRCompileLayer.cpp
2828
IRTransformLayer.cpp
29+
IRPartitionLayer.cpp
2930
JITTargetMachineBuilder.cpp
3031
LazyReexports.cpp
3132
Layer.cpp

0 commit comments

Comments
 (0)