1*da58b97aSjoerg //===--- LLJITWithLazyReexports.cpp - LLJIT example with custom laziness --===//
2*da58b97aSjoerg //
3*da58b97aSjoerg // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4*da58b97aSjoerg // See https://llvm.org/LICENSE.txt for license information.
5*da58b97aSjoerg // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6*da58b97aSjoerg //
7*da58b97aSjoerg //===----------------------------------------------------------------------===//
8*da58b97aSjoerg //
9*da58b97aSjoerg // In this example we will use the lazy re-exports utility to lazily compile
10*da58b97aSjoerg // IR modules. We will do this in seven steps:
11*da58b97aSjoerg //
12*da58b97aSjoerg // 1. Create an LLJIT instance.
13*da58b97aSjoerg // 2. Install a transform so that we can see what is being compiled.
14*da58b97aSjoerg // 3. Create an indirect stubs manager and lazy call-through manager.
15*da58b97aSjoerg // 4. Add two modules that will be conditionally compiled, plus a main module.
16*da58b97aSjoerg // 5. Add lazy-rexports of the symbols in the conditionally compiled modules.
17*da58b97aSjoerg // 6. Dump the ExecutionSession state to see the symbol table prior to
18*da58b97aSjoerg //    executing any code.
19*da58b97aSjoerg // 7. Verify that only modules containing executed code are compiled.
20*da58b97aSjoerg //
21*da58b97aSjoerg //===----------------------------------------------------------------------===//
22*da58b97aSjoerg 
23*da58b97aSjoerg #include "llvm/ADT/StringMap.h"
24*da58b97aSjoerg #include "llvm/ExecutionEngine/JITLink/JITLinkMemoryManager.h"
25*da58b97aSjoerg #include "llvm/ExecutionEngine/Orc/LLJIT.h"
26*da58b97aSjoerg #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
27*da58b97aSjoerg #include "llvm/ExecutionEngine/Orc/OrcABISupport.h"
28*da58b97aSjoerg #include "llvm/ExecutionEngine/Orc/TPCDynamicLibrarySearchGenerator.h"
29*da58b97aSjoerg #include "llvm/ExecutionEngine/Orc/TPCIndirectionUtils.h"
30*da58b97aSjoerg #include "llvm/ExecutionEngine/Orc/TargetProcessControl.h"
31*da58b97aSjoerg #include "llvm/Support/InitLLVM.h"
32*da58b97aSjoerg #include "llvm/Support/TargetSelect.h"
33*da58b97aSjoerg #include "llvm/Support/raw_ostream.h"
34*da58b97aSjoerg 
35*da58b97aSjoerg #include "../ExampleModules.h"
36*da58b97aSjoerg 
37*da58b97aSjoerg #include <future>
38*da58b97aSjoerg 
39*da58b97aSjoerg using namespace llvm;
40*da58b97aSjoerg using namespace llvm::orc;
41*da58b97aSjoerg 
42*da58b97aSjoerg ExitOnError ExitOnErr;
43*da58b97aSjoerg 
44*da58b97aSjoerg // Example IR modules.
45*da58b97aSjoerg //
46*da58b97aSjoerg // Note that in the conditionally compiled modules, FooMod and BarMod, functions
47*da58b97aSjoerg // have been given an _body suffix. This is to ensure that their names do not
48*da58b97aSjoerg // clash with their lazy-reexports.
49*da58b97aSjoerg // For clients who do not wish to rename function bodies (e.g. because they want
50*da58b97aSjoerg // to re-use cached objects between static and JIT compiles) techniques exist to
51*da58b97aSjoerg // avoid renaming. See the lazy-reexports section of the ORCv2 design doc.
52*da58b97aSjoerg 
53*da58b97aSjoerg const llvm::StringRef FooMod =
54*da58b97aSjoerg     R"(
55*da58b97aSjoerg   declare i32 @return1()
56*da58b97aSjoerg 
57*da58b97aSjoerg   define i32 @foo_body() {
58*da58b97aSjoerg   entry:
59*da58b97aSjoerg     %0 = call i32 @return1()
60*da58b97aSjoerg     ret i32 %0
61*da58b97aSjoerg   }
62*da58b97aSjoerg )";
63*da58b97aSjoerg 
64*da58b97aSjoerg const llvm::StringRef BarMod =
65*da58b97aSjoerg     R"(
66*da58b97aSjoerg   declare i32 @return2()
67*da58b97aSjoerg 
68*da58b97aSjoerg   define i32 @bar_body() {
69*da58b97aSjoerg   entry:
70*da58b97aSjoerg     %0 = call i32 @return2()
71*da58b97aSjoerg     ret i32 %0
72*da58b97aSjoerg   }
73*da58b97aSjoerg )";
74*da58b97aSjoerg 
75*da58b97aSjoerg const llvm::StringRef MainMod =
76*da58b97aSjoerg     R"(
77*da58b97aSjoerg 
78*da58b97aSjoerg   define i32 @entry(i32 %argc) {
79*da58b97aSjoerg   entry:
80*da58b97aSjoerg     %and = and i32 %argc, 1
81*da58b97aSjoerg     %tobool = icmp eq i32 %and, 0
82*da58b97aSjoerg     br i1 %tobool, label %if.end, label %if.then
83*da58b97aSjoerg 
84*da58b97aSjoerg   if.then:                                          ; preds = %entry
85*da58b97aSjoerg     %call = tail call i32 @foo() #2
86*da58b97aSjoerg     br label %return
87*da58b97aSjoerg 
88*da58b97aSjoerg   if.end:                                           ; preds = %entry
89*da58b97aSjoerg     %call1 = tail call i32 @bar() #2
90*da58b97aSjoerg     br label %return
91*da58b97aSjoerg 
92*da58b97aSjoerg   return:                                           ; preds = %if.end, %if.then
93*da58b97aSjoerg     %retval.0 = phi i32 [ %call, %if.then ], [ %call1, %if.end ]
94*da58b97aSjoerg     ret i32 %retval.0
95*da58b97aSjoerg   }
96*da58b97aSjoerg 
97*da58b97aSjoerg   declare i32 @foo()
98*da58b97aSjoerg   declare i32 @bar()
99*da58b97aSjoerg )";
100*da58b97aSjoerg 
return1()101*da58b97aSjoerg extern "C" int32_t return1() { return 1; }
return2()102*da58b97aSjoerg extern "C" int32_t return2() { return 2; }
103*da58b97aSjoerg 
reenter(void * Ctx,void * TrampolineAddr)104*da58b97aSjoerg static void *reenter(void *Ctx, void *TrampolineAddr) {
105*da58b97aSjoerg   std::promise<void *> LandingAddressP;
106*da58b97aSjoerg   auto LandingAddressF = LandingAddressP.get_future();
107*da58b97aSjoerg 
108*da58b97aSjoerg   auto *TPCIU = static_cast<TPCIndirectionUtils *>(Ctx);
109*da58b97aSjoerg   TPCIU->getLazyCallThroughManager().resolveTrampolineLandingAddress(
110*da58b97aSjoerg       pointerToJITTargetAddress(TrampolineAddr),
111*da58b97aSjoerg       [&](JITTargetAddress LandingAddress) {
112*da58b97aSjoerg         LandingAddressP.set_value(
113*da58b97aSjoerg             jitTargetAddressToPointer<void *>(LandingAddress));
114*da58b97aSjoerg       });
115*da58b97aSjoerg   return LandingAddressF.get();
116*da58b97aSjoerg }
117*da58b97aSjoerg 
reportErrorAndExit()118*da58b97aSjoerg static void reportErrorAndExit() {
119*da58b97aSjoerg   errs() << "Unable to lazily compile function. Exiting.\n";
120*da58b97aSjoerg   exit(1);
121*da58b97aSjoerg }
122*da58b97aSjoerg 
123*da58b97aSjoerg cl::list<std::string> InputArgv(cl::Positional,
124*da58b97aSjoerg                                 cl::desc("<program arguments>..."));
125*da58b97aSjoerg 
main(int argc,char * argv[])126*da58b97aSjoerg int main(int argc, char *argv[]) {
127*da58b97aSjoerg   // Initialize LLVM.
128*da58b97aSjoerg   InitLLVM X(argc, argv);
129*da58b97aSjoerg 
130*da58b97aSjoerg   InitializeNativeTarget();
131*da58b97aSjoerg   InitializeNativeTargetAsmPrinter();
132*da58b97aSjoerg 
133*da58b97aSjoerg   cl::ParseCommandLineOptions(argc, argv, "LLJITWithLazyReexports");
134*da58b97aSjoerg   ExitOnErr.setBanner(std::string(argv[0]) + ": ");
135*da58b97aSjoerg 
136*da58b97aSjoerg   // (1) Create LLJIT instance.
137*da58b97aSjoerg   auto SSP = std::make_shared<SymbolStringPool>();
138*da58b97aSjoerg   auto TPC = ExitOnErr(SelfTargetProcessControl::Create(std::move(SSP)));
139*da58b97aSjoerg   auto J = ExitOnErr(LLJITBuilder().setTargetProcessControl(*TPC).create());
140*da58b97aSjoerg 
141*da58b97aSjoerg   // (2) Install transform to print modules as they are compiled:
142*da58b97aSjoerg   J->getIRTransformLayer().setTransform(
143*da58b97aSjoerg       [](ThreadSafeModule TSM,
144*da58b97aSjoerg          const MaterializationResponsibility &R) -> Expected<ThreadSafeModule> {
145*da58b97aSjoerg         TSM.withModuleDo([](Module &M) { dbgs() << "---Compiling---\n" << M; });
146*da58b97aSjoerg         return std::move(TSM); // Not a redundant move: fix build on gcc-7.5
147*da58b97aSjoerg       });
148*da58b97aSjoerg 
149*da58b97aSjoerg   // (3) Create stubs and call-through managers:
150*da58b97aSjoerg   auto TPCIU = ExitOnErr(TPCIndirectionUtils::Create(*TPC));
151*da58b97aSjoerg   ExitOnErr(TPCIU->writeResolverBlock(pointerToJITTargetAddress(&reenter),
152*da58b97aSjoerg                                       pointerToJITTargetAddress(TPCIU.get())));
153*da58b97aSjoerg   TPCIU->createLazyCallThroughManager(
154*da58b97aSjoerg       J->getExecutionSession(), pointerToJITTargetAddress(&reportErrorAndExit));
155*da58b97aSjoerg   auto ISM = TPCIU->createIndirectStubsManager();
156*da58b97aSjoerg   J->getMainJITDylib().addGenerator(
157*da58b97aSjoerg       ExitOnErr(TPCDynamicLibrarySearchGenerator::GetForTargetProcess(*TPC)));
158*da58b97aSjoerg 
159*da58b97aSjoerg   // (4) Add modules.
160*da58b97aSjoerg   ExitOnErr(J->addIRModule(ExitOnErr(parseExampleModule(FooMod, "foo-mod"))));
161*da58b97aSjoerg   ExitOnErr(J->addIRModule(ExitOnErr(parseExampleModule(BarMod, "bar-mod"))));
162*da58b97aSjoerg   ExitOnErr(J->addIRModule(ExitOnErr(parseExampleModule(MainMod, "main-mod"))));
163*da58b97aSjoerg 
164*da58b97aSjoerg   // (5) Add lazy reexports.
165*da58b97aSjoerg   MangleAndInterner Mangle(J->getExecutionSession(), J->getDataLayout());
166*da58b97aSjoerg   SymbolAliasMap ReExports(
167*da58b97aSjoerg       {{Mangle("foo"),
168*da58b97aSjoerg         {Mangle("foo_body"),
169*da58b97aSjoerg          JITSymbolFlags::Exported | JITSymbolFlags::Callable}},
170*da58b97aSjoerg        {Mangle("bar"),
171*da58b97aSjoerg         {Mangle("bar_body"),
172*da58b97aSjoerg          JITSymbolFlags::Exported | JITSymbolFlags::Callable}}});
173*da58b97aSjoerg   ExitOnErr(J->getMainJITDylib().define(
174*da58b97aSjoerg       lazyReexports(TPCIU->getLazyCallThroughManager(), *ISM,
175*da58b97aSjoerg                     J->getMainJITDylib(), std::move(ReExports))));
176*da58b97aSjoerg 
177*da58b97aSjoerg   // (6) Dump the ExecutionSession state.
178*da58b97aSjoerg   dbgs() << "---Session state---\n";
179*da58b97aSjoerg   J->getExecutionSession().dump(dbgs());
180*da58b97aSjoerg   dbgs() << "\n";
181*da58b97aSjoerg 
182*da58b97aSjoerg   // (7) Execute the JIT'd main function and pass the example's command line
183*da58b97aSjoerg   // arguments unmodified. This should cause either ExampleMod1 or ExampleMod2
184*da58b97aSjoerg   // to be compiled, and either "1" or "2" returned depending on the number of
185*da58b97aSjoerg   // arguments passed.
186*da58b97aSjoerg 
187*da58b97aSjoerg   // Look up the JIT'd function, cast it to a function pointer, then call it.
188*da58b97aSjoerg   auto EntrySym = ExitOnErr(J->lookup("entry"));
189*da58b97aSjoerg   auto *Entry = (int (*)(int))EntrySym.getAddress();
190*da58b97aSjoerg 
191*da58b97aSjoerg   int Result = Entry(argc);
192*da58b97aSjoerg   outs() << "---Result---\n"
193*da58b97aSjoerg          << "entry(" << argc << ") = " << Result << "\n";
194*da58b97aSjoerg 
195*da58b97aSjoerg   return 0;
196*da58b97aSjoerg }
197