xref: /llvm-project/llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp (revision 04af63b267c391a4b0a0fb61060f724f8b5bc2be)
1 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
2 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
3 #include "llvm/ExecutionEngine/Orc/Core.h"
4 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
5 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
6 #include "llvm/ExecutionEngine/Orc/IRPartitionLayer.h"
7 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
8 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
9 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
10 #include "llvm/ExecutionEngine/Orc/SpeculateAnalyses.h"
11 #include "llvm/ExecutionEngine/Orc/Speculation.h"
12 #include "llvm/ExecutionEngine/Orc/TargetProcess/TargetExecutionUtils.h"
13 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
14 #include "llvm/IRReader/IRReader.h"
15 #include "llvm/Support/CommandLine.h"
16 #include "llvm/Support/Debug.h"
17 #include "llvm/Support/InitLLVM.h"
18 #include "llvm/Support/SourceMgr.h"
19 #include "llvm/Support/TargetSelect.h"
20 #include "llvm/Support/ThreadPool.h"
21 
22 #include <list>
23 #include <string>
24 
25 using namespace llvm;
26 using namespace llvm::orc;
27 
28 static cl::list<std::string> InputFiles(cl::Positional, cl::OneOrMore,
29                                         cl::desc("input files"));
30 
31 static cl::list<std::string> InputArgv("args", cl::Positional,
32                                        cl::desc("<program arguments>..."),
33                                        cl::PositionalEatsArgs);
34 
35 static cl::opt<unsigned> NumThreads("num-threads", cl::Optional,
36                                     cl::desc("Number of compile threads"),
37                                     cl::init(4));
38 
39 ExitOnError ExitOnErr;
40 
41 // Add Layers
42 class SpeculativeJIT {
43 public:
44   static Expected<std::unique_ptr<SpeculativeJIT>> Create() {
45     auto JTMB = orc::JITTargetMachineBuilder::detectHost();
46     if (!JTMB)
47       return JTMB.takeError();
48 
49     auto DL = JTMB->getDefaultDataLayoutForTarget();
50     if (!DL)
51       return DL.takeError();
52 
53     auto EPC = SelfExecutorProcessControl::Create(
54         nullptr,
55         std::make_unique<DynamicThreadPoolTaskDispatcher>(std::nullopt));
56     if (!EPC)
57       return EPC.takeError();
58 
59     auto ES = std::make_unique<ExecutionSession>(std::move(*EPC));
60 
61     auto LCTMgr = createLocalLazyCallThroughManager(
62         JTMB->getTargetTriple(), *ES,
63         ExecutorAddr::fromPtr(explodeOnLazyCompileFailure));
64     if (!LCTMgr)
65       return LCTMgr.takeError();
66 
67     auto ISMBuilder =
68         createLocalIndirectStubsManagerBuilder(JTMB->getTargetTriple());
69     if (!ISMBuilder)
70       return make_error<StringError>("No indirect stubs manager for target",
71                                      inconvertibleErrorCode());
72 
73     auto ProcessSymbolsSearchGenerator =
74         DynamicLibrarySearchGenerator::GetForCurrentProcess(
75             DL->getGlobalPrefix());
76     if (!ProcessSymbolsSearchGenerator)
77       return ProcessSymbolsSearchGenerator.takeError();
78 
79     std::unique_ptr<SpeculativeJIT> SJ(new SpeculativeJIT(
80         std::move(ES), std::move(*DL), std::move(*JTMB), std::move(*LCTMgr),
81         std::move(ISMBuilder), std::move(*ProcessSymbolsSearchGenerator)));
82     return std::move(SJ);
83   }
84 
85   ExecutionSession &getES() { return *ES; }
86 
87   Error addModule(ThreadSafeModule TSM) {
88     return CODLayer.add(MainJD, std::move(TSM));
89   }
90 
91   Expected<ExecutorSymbolDef> lookup(StringRef UnmangledName) {
92     return ES->lookup({&MainJD}, Mangle(UnmangledName));
93   }
94 
95   ~SpeculativeJIT() { CompileThreads.wait(); }
96 
97 private:
98   using IndirectStubsManagerBuilderFunction =
99       std::function<std::unique_ptr<IndirectStubsManager>()>;
100 
101   static void explodeOnLazyCompileFailure() {
102     errs() << "Lazy compilation failed, Symbol Implmentation not found!\n";
103     exit(1);
104   }
105 
106   SpeculativeJIT(
107       std::unique_ptr<ExecutionSession> ES, DataLayout DL,
108       orc::JITTargetMachineBuilder JTMB,
109       std::unique_ptr<LazyCallThroughManager> LCTMgr,
110       IndirectStubsManagerBuilderFunction ISMBuilder,
111       std::unique_ptr<DynamicLibrarySearchGenerator> ProcessSymbolsGenerator)
112       : ES(std::move(ES)), DL(std::move(DL)),
113         MainJD(this->ES->createBareJITDylib("<main>")),
114         LCTMgr(std::move(LCTMgr)),
115         CompileLayer(*this->ES, ObjLayer,
116                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
117         S(Imps, *this->ES),
118         SpeculateLayer(*this->ES, CompileLayer, S, Mangle, BlockFreqQuery()),
119         IPLayer(*this->ES, SpeculateLayer),
120         CODLayer(*this->ES, IPLayer, *this->LCTMgr, std::move(ISMBuilder)) {
121     MainJD.addGenerator(std::move(ProcessSymbolsGenerator));
122     this->CODLayer.setImplMap(&Imps);
123     ExitOnErr(S.addSpeculationRuntime(MainJD, Mangle));
124     LocalCXXRuntimeOverrides CXXRuntimeoverrides;
125     ExitOnErr(CXXRuntimeoverrides.enable(MainJD, Mangle));
126   }
127 
128   static std::unique_ptr<SectionMemoryManager> createMemMgr() {
129     return std::make_unique<SectionMemoryManager>();
130   }
131 
132   std::unique_ptr<ExecutionSession> ES;
133   DataLayout DL;
134   MangleAndInterner Mangle{*ES, DL};
135   DefaultThreadPool CompileThreads{llvm::hardware_concurrency(NumThreads)};
136 
137   JITDylib &MainJD;
138 
139   Triple TT;
140   std::unique_ptr<LazyCallThroughManager> LCTMgr;
141   IRCompileLayer CompileLayer;
142   ImplSymbolMap Imps;
143   Speculator S;
144   RTDyldObjectLinkingLayer ObjLayer{*ES, createMemMgr};
145   IRSpeculationLayer SpeculateLayer;
146   IRPartitionLayer IPLayer;
147   CompileOnDemandLayer CODLayer;
148 };
149 
150 int main(int argc, char *argv[]) {
151   // Initialize LLVM.
152   InitLLVM X(argc, argv);
153 
154   InitializeNativeTarget();
155   InitializeNativeTargetAsmPrinter();
156 
157   cl::ParseCommandLineOptions(argc, argv, "SpeculativeJIT");
158   ExitOnErr.setBanner(std::string(argv[0]) + ": ");
159 
160   if (NumThreads < 1) {
161     errs() << "Speculative compilation requires one or more dedicated compile "
162               "threads\n";
163     return 1;
164   }
165 
166   // Create a JIT instance.
167   auto SJ = ExitOnErr(SpeculativeJIT::Create());
168 
169   // Load the IR inputs.
170   for (const auto &InputFile : InputFiles) {
171     SMDiagnostic Err;
172     auto Ctx = std::make_unique<LLVMContext>();
173     auto M = parseIRFile(InputFile, Err, *Ctx);
174     if (!M) {
175       Err.print(argv[0], errs());
176       return 1;
177     }
178 
179     ExitOnErr(SJ->addModule(ThreadSafeModule(std::move(M), std::move(Ctx))));
180   }
181 
182   auto MainSym = ExitOnErr(SJ->lookup("main"));
183   auto Main = MainSym.getAddress().toPtr<int (*)(int, char *[])>();
184 
185   return runAsMain(Main, InputArgv, StringRef(InputFiles.front()));
186 
187   return 0;
188 }
189