Skip to content

Commit e9e26c0

Browse files
lhameslhames
committed
[ORC] Simplify use of lazyReexports with LLJIT.
This patch makes the target triple available via the LLJIT interface, and moves the IRTransformLayer from LLLazyJIT down into LLJIT. Together these changes make it easier to use the lazyReexports utility with LLJIT, and to apply IR transforms to code as it is compiled in LLJIT (rather than requiring transforms to be applied manually before code is added). An code example is added in llvm/examples/LLJITExamples/LLJITWithLazyReexports
1 parent d2fabd7 commit e9e26c0

File tree

6 files changed

+200
-22
lines changed

6 files changed

+200
-22
lines changed

‎llvm/examples/LLJITExamples/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
11
add_subdirectory(LLJITDumpObjects)
22
add_subdirectory(LLJITWithObjectCache)
33
add_subdirectory(LLJITWithCustomObjectLinkingLayer)
4+
add_subdirectory(LLJITWithLazyReexports)

‎llvm/examples/LLJITExamples/LLJITWithLazyReexports/CMakeLists.txt

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
set(LLVM_LINK_COMPONENTS
2+
Core
3+
ExecutionEngine
4+
IRReader
5+
OrcJIT
6+
Support
7+
nativecodegen
8+
)
9+
10+
add_llvm_example(LLJITWithLazyReexports
11+
LLJITWithLazyReexports.cpp
12+
)

‎llvm/examples/LLJITExamples/LLJITWithLazyReexports/LLJITWithLazyReexports.cpp

Lines changed: 163 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,163 @@
1+
//===--- LLJITWithLazyReexports.cpp - LLJIT example with custom laziness --===//
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+
// In this example we will use the lazy re-exports utility to lazily compile
10+
// IR modules. We will do this in seven steps:
11+
//
12+
// 1. Create an LLJIT instance.
13+
// 2. Install a transform so that we is being compiled.
14+
// 3. Create an indirect stubs manager and lazy call-through manager.
15+
// 4. Add two modules that will be conditionally compiled, plus a main module.
16+
// 5. Add lazy-rexports of the symbols in the conditionally compiled modules.
17+
// 6. Dump the ExecutionSession state to see the symbol table prior to
18+
// executing any code.
19+
// 7. Verify that only modules containing executed code are compiled.
20+
//
21+
//===----------------------------------------------------------------------===//
22+
23+
#include "llvm/ADT/StringMap.h"
24+
#include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
25+
#include "llvm/ExecutionEngine/Orc/LLJIT.h"
26+
#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
27+
#include "llvm/Support/InitLLVM.h"
28+
#include "llvm/Support/TargetSelect.h"
29+
#include "llvm/Support/raw_ostream.h"
30+
31+
#include "../ExampleModules.h"
32+
33+
using namespace llvm;
34+
using namespace llvm::orc;
35+
36+
ExitOnError ExitOnErr;
37+
38+
// Example IR modules.
39+
//
40+
// Note that in the conditionally compiled modules, FooMod and BarMod, functions
41+
// have been given an _body suffix. This is to ensure that their names do not
42+
// clash with their lazy-reexports.
43+
// For clients who do not wish to rename function bodies (e.g. because they want
44+
// to re-use cached objects between static and JIT compiles) techniques exist to
45+
// avoid renaming. See the lazy-reexports section of the ORCv2 design doc.
46+
47+
const llvm::StringRef FooMod =
48+
R"(
49+
define i32 @foo_body() {
50+
entry:
51+
ret i32 1
52+
}
53+
)";
54+
55+
const llvm::StringRef BarMod =
56+
R"(
57+
define i32 @bar_body(i32 %x) {
58+
entry:
59+
ret i32 2
60+
}
61+
)";
62+
63+
const llvm::StringRef MainMod =
64+
R"(
65+
66+
define i32 @entry(i32 %argc) {
67+
entry:
68+
%and = and i32 %argc, 1
69+
%tobool = icmp eq i32 %and, 0
70+
br i1 %tobool, label %if.end, label %if.then
71+
72+
if.then: ; preds = %entry
73+
%call = tail call i32 @foo() #2
74+
br label %return
75+
76+
if.end: ; preds = %entry
77+
%call1 = tail call i32 @bar() #2
78+
br label %return
79+
80+
return: ; preds = %if.end, %if.then
81+
%retval.0 = phi i32 [ %call, %if.then ], [ %call1, %if.end ]
82+
ret i32 %retval.0
83+
}
84+
85+
declare i32 @foo()
86+
declare i32 @bar()
87+
)";
88+
89+
cl::list<std::string> InputArgv(cl::Positional,
90+
cl::desc("<program arguments>..."));
91+
92+
int main(int argc, char *argv[]) {
93+
// Initialize LLVM.
94+
InitLLVM X(argc, argv);
95+
96+
InitializeNativeTarget();
97+
InitializeNativeTargetAsmPrinter();
98+
99+
cl::ParseCommandLineOptions(argc, argv, "LLJITWithLazyReexports");
100+
ExitOnErr.setBanner(std::string(argv[0]) + ": ");
101+
102+
// (1) Create LLJIT instance.
103+
auto J = ExitOnErr(LLJITBuilder().create());
104+
105+
// (2) Install transform to print modules as they are compiled:
106+
J->getIRTransformLayer().setTransform(
107+
[](ThreadSafeModule TSM,
108+
const MaterializationResponsibility &R) -> Expected<ThreadSafeModule> {
109+
TSM.withModuleDo([](Module &M) { dbgs() << "---Compiling---\n" << M; });
110+
return TSM;
111+
});
112+
113+
// (3) Create stubs and call-through managers:
114+
std::unique_ptr<IndirectStubsManager> ISM;
115+
{
116+
auto ISMBuilder =
117+
createLocalIndirectStubsManagerBuilder(J->getTargetTriple());
118+
if (!ISMBuilder())
119+
ExitOnErr(make_error<StringError>("Could not create stubs manager for " +
120+
J->getTargetTriple().str(),
121+
inconvertibleErrorCode()));
122+
ISM = ISMBuilder();
123+
}
124+
auto LCTM = ExitOnErr(createLocalLazyCallThroughManager(
125+
J->getTargetTriple(), J->getExecutionSession(), 0));
126+
127+
// (4) Add modules.
128+
ExitOnErr(J->addIRModule(ExitOnErr(parseExampleModule(FooMod, "foo-mod"))));
129+
ExitOnErr(J->addIRModule(ExitOnErr(parseExampleModule(BarMod, "bar-mod"))));
130+
ExitOnErr(J->addIRModule(ExitOnErr(parseExampleModule(MainMod, "main-mod"))));
131+
132+
// (5) Add lazy reexports.
133+
MangleAndInterner Mangle(J->getExecutionSession(), J->getDataLayout());
134+
SymbolAliasMap ReExports(
135+
{{Mangle("foo"),
136+
{Mangle("foo_body"),
137+
JITSymbolFlags::Exported | JITSymbolFlags::Callable}},
138+
{Mangle("bar"),
139+
{Mangle("bar_body"),
140+
JITSymbolFlags::Exported | JITSymbolFlags::Callable}}});
141+
ExitOnErr(J->getMainJITDylib().define(
142+
lazyReexports(*LCTM, *ISM, J->getMainJITDylib(), std::move(ReExports))));
143+
144+
// (6) Dump the ExecutionSession state.
145+
dbgs() << "---Session state---\n";
146+
J->getExecutionSession().dump(dbgs());
147+
dbgs() << "\n";
148+
149+
// (7) Execute the JIT'd main function and pass the example's command line
150+
// arguments unmodified. This should cause either ExampleMod1 or ExampleMod2
151+
// to be compiled, and either "1" or "2" returned depending on the number of
152+
// arguments passed.
153+
154+
// Look up the JIT'd function, cast it to a function pointer, then call it.
155+
auto EntrySym = ExitOnErr(J->lookup("entry"));
156+
auto *Entry = (int (*)(int))EntrySym.getAddress();
157+
158+
int Result = Entry(argc);
159+
outs() << "---Result---\n"
160+
<< "entry(" << argc << ") = " << Result << "\n";
161+
162+
return 0;
163+
}

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

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,9 @@ class LLJIT {
4545
/// Returns the ExecutionSession for this instance.
4646
ExecutionSession &getExecutionSession() { return *ES; }
4747

48+
/// Returns a reference to the triple for this instance.
49+
const Triple &getTargetTriple() const { return TT; }
50+
4851
/// Returns a reference to the DataLayout for this instance.
4952
const DataLayout &getDataLayout() const { return DL; }
5053

@@ -120,6 +123,9 @@ class LLJIT {
120123
/// Returns a reference to the object transform layer.
121124
ObjectTransformLayer &getObjTransformLayer() { return ObjTransformLayer; }
122125

126+
/// Returns a reference to the IR transform layer.
127+
IRTransformLayer &getIRTransformLayer() { return *TransformLayer; }
128+
123129
protected:
124130
static std::unique_ptr<ObjectLayer>
125131
createObjectLinkingLayer(LLJITBuilderState &S, ExecutionSession &ES);
@@ -140,11 +146,13 @@ class LLJIT {
140146
JITDylib &Main;
141147

142148
DataLayout DL;
149+
Triple TT;
143150
std::unique_ptr<ThreadPool> CompileThreads;
144151

145152
std::unique_ptr<ObjectLayer> ObjLinkingLayer;
146153
ObjectTransformLayer ObjTransformLayer;
147154
std::unique_ptr<IRCompileLayer> CompileLayer;
155+
std::unique_ptr<IRTransformLayer> TransformLayer;
148156

149157
CtorDtorRunner CtorRunner, DtorRunner;
150158
};
@@ -156,12 +164,6 @@ class LLLazyJIT : public LLJIT {
156164

157165
public:
158166

159-
/// Set an IR transform (e.g. pass manager pipeline) to run on each function
160-
/// when it is compiled.
161-
void setLazyCompileTransform(IRTransformLayer::TransformFunction Transform) {
162-
TransformLayer->setTransform(std::move(Transform));
163-
}
164-
165167
/// Sets the partition function.
166168
void
167169
setPartitionFunction(CompileOnDemandLayer::PartitionFunction Partition) {
@@ -182,7 +184,6 @@ class LLLazyJIT : public LLJIT {
182184
LLLazyJIT(LLLazyJITBuilderState &S, Error &Err);
183185

184186
std::unique_ptr<LazyCallThroughManager> LCTMgr;
185-
std::unique_ptr<IRTransformLayer> TransformLayer;
186187
std::unique_ptr<CompileOnDemandLayer> CODLayer;
187188
};
188189

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,7 @@ Error LLJIT::addIRModule(JITDylib &JD, ThreadSafeModule TSM) {
6767
TSM.withModuleDo([&](Module &M) { return applyDataLayout(M); }))
6868
return Err;
6969

70-
return CompileLayer->add(JD, std::move(TSM), ES->allocateVModule());
70+
return TransformLayer->add(JD, std::move(TSM), ES->allocateVModule());
7171
}
7272

7373
Error LLJIT::addObjectFile(JITDylib &JD, std::unique_ptr<MemoryBuffer> Obj) {
@@ -128,6 +128,7 @@ LLJIT::createCompileFunction(LLJITBuilderState &S,
128128
LLJIT::LLJIT(LLJITBuilderState &S, Error &Err)
129129
: ES(S.ES ? std::move(S.ES) : std::make_unique<ExecutionSession>()),
130130
Main(this->ES->createJITDylib("<main>")), DL(""),
131+
TT(S.JTMB->getTargetTriple()),
131132
ObjLinkingLayer(createObjectLinkingLayer(S, *ES)),
132133
ObjTransformLayer(*this->ES, *ObjLinkingLayer), CtorRunner(Main),
133134
DtorRunner(Main) {
@@ -162,6 +163,8 @@ LLJIT::LLJIT(LLJITBuilderState &S, Error &Err)
162163
CompileThreads->async(std::move(Work));
163164
});
164165
}
166+
167+
TransformLayer = std::make_unique<IRTransformLayer>(*ES, *CompileLayer);
165168
}
166169

167170
std::string LLJIT::mangle(StringRef UnmangledName) {
@@ -249,9 +252,6 @@ LLLazyJIT::LLLazyJIT(LLLazyJITBuilderState &S, Error &Err) : LLJIT(S, Err) {
249252
return;
250253
}
251254

252-
// Create the transform layer.
253-
TransformLayer = std::make_unique<IRTransformLayer>(*ES, *CompileLayer);
254-
255255
// Create the COD layer.
256256
CODLayer = std::make_unique<CompileOnDemandLayer>(
257257
*ES, *TransformLayer, *LCTMgr, std::move(ISMBuilder));

‎llvm/tools/lli/lli.cpp

Lines changed: 12 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -781,17 +781,18 @@ int runOrcLazyJIT(const char *ProgName) {
781781

782782
auto Dump = createDebugDumper();
783783

784-
J->setLazyCompileTransform([&](orc::ThreadSafeModule TSM,
785-
const orc::MaterializationResponsibility &R) {
786-
TSM.withModuleDo([&](Module &M) {
787-
if (verifyModule(M, &dbgs())) {
788-
dbgs() << "Bad module: " << &M << "\n";
789-
exit(1);
790-
}
791-
Dump(M);
792-
});
793-
return TSM;
794-
});
784+
J->getIRTransformLayer().setTransform(
785+
[&](orc::ThreadSafeModule TSM,
786+
const orc::MaterializationResponsibility &R) {
787+
TSM.withModuleDo([&](Module &M) {
788+
if (verifyModule(M, &dbgs())) {
789+
dbgs() << "Bad module: " << &M << "\n";
790+
exit(1);
791+
}
792+
Dump(M);
793+
});
794+
return TSM;
795+
});
795796

796797
orc::MangleAndInterner Mangle(J->getExecutionSession(), J->getDataLayout());
797798
J->getMainJITDylib().addGenerator(

0 commit comments

Comments
 (0)