[lit] Improve lit.Run class
[llvm-complete.git] / examples / SpeculativeJIT / SpeculativeJIT.cpp
blob1fd1fc92a73f9225898272e0db9f74878a8704a8
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/IndirectionUtils.h"
7 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
8 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
9 #include "llvm/ExecutionEngine/Orc/SpeculateAnalyses.h"
10 #include "llvm/ExecutionEngine/Orc/Speculation.h"
11 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
12 #include "llvm/IRReader/IRReader.h"
13 #include "llvm/Support/CommandLine.h"
14 #include "llvm/Support/Debug.h"
15 #include "llvm/Support/InitLLVM.h"
16 #include "llvm/Support/SourceMgr.h"
17 #include "llvm/Support/TargetSelect.h"
18 #include "llvm/Support/ThreadPool.h"
20 #include <list>
21 #include <string>
23 using namespace llvm;
24 using namespace llvm::orc;
26 static cl::list<std::string> InputFiles(cl::Positional, cl::OneOrMore,
27 cl::desc("input files"));
29 static cl::list<std::string> InputArgv("args", cl::Positional,
30 cl::desc("<program arguments>..."),
31 cl::ZeroOrMore, cl::PositionalEatsArgs);
33 static cl::opt<unsigned> NumThreads("num-threads", cl::Optional,
34 cl::desc("Number of compile threads"),
35 cl::init(4));
37 ExitOnError ExitOnErr;
39 // Add Layers
40 class SpeculativeJIT {
41 public:
42 static Expected<std::unique_ptr<SpeculativeJIT>> Create() {
43 auto JTMB = orc::JITTargetMachineBuilder::detectHost();
44 if (!JTMB)
45 return JTMB.takeError();
47 auto DL = JTMB->getDefaultDataLayoutForTarget();
48 if (!DL)
49 return DL.takeError();
51 auto ES = std::make_unique<ExecutionSession>();
53 auto LCTMgr = createLocalLazyCallThroughManager(
54 JTMB->getTargetTriple(), *ES,
55 pointerToJITTargetAddress(explodeOnLazyCompileFailure));
56 if (!LCTMgr)
57 return LCTMgr.takeError();
59 auto ISMBuilder =
60 createLocalIndirectStubsManagerBuilder(JTMB->getTargetTriple());
61 if (!ISMBuilder)
62 return make_error<StringError>("No indirect stubs manager for target",
63 inconvertibleErrorCode());
65 auto ProcessSymbolsSearchGenerator =
66 DynamicLibrarySearchGenerator::GetForCurrentProcess(
67 DL->getGlobalPrefix());
68 if (!ProcessSymbolsSearchGenerator)
69 return ProcessSymbolsSearchGenerator.takeError();
71 std::unique_ptr<SpeculativeJIT> SJ(new SpeculativeJIT(
72 std::move(ES), std::move(*DL), std::move(*JTMB), std::move(*LCTMgr),
73 std::move(ISMBuilder), std::move(*ProcessSymbolsSearchGenerator)));
74 return std::move(SJ);
77 ExecutionSession &getES() { return *ES; }
79 Error addModule(JITDylib &JD, ThreadSafeModule TSM) {
80 return CODLayer.add(JD, std::move(TSM));
83 Expected<JITEvaluatedSymbol> lookup(StringRef UnmangledName) {
84 return ES->lookup({&ES->getMainJITDylib()}, Mangle(UnmangledName));
87 ~SpeculativeJIT() { CompileThreads.wait(); }
89 private:
90 using IndirectStubsManagerBuilderFunction =
91 std::function<std::unique_ptr<IndirectStubsManager>()>;
93 static void explodeOnLazyCompileFailure() {
94 errs() << "Lazy compilation failed, Symbol Implmentation not found!\n";
95 exit(1);
98 SpeculativeJIT(
99 std::unique_ptr<ExecutionSession> ES, DataLayout DL,
100 orc::JITTargetMachineBuilder JTMB,
101 std::unique_ptr<LazyCallThroughManager> LCTMgr,
102 IndirectStubsManagerBuilderFunction ISMBuilder,
103 std::unique_ptr<DynamicLibrarySearchGenerator> ProcessSymbolsGenerator)
104 : ES(std::move(ES)), DL(std::move(DL)), LCTMgr(std::move(LCTMgr)),
105 CompileLayer(*this->ES, ObjLayer,
106 ConcurrentIRCompiler(std::move(JTMB))),
107 S(Imps, *this->ES),
108 SpeculateLayer(*this->ES, CompileLayer, S, Mangle, BlockFreqQuery()),
109 CODLayer(*this->ES, SpeculateLayer, *this->LCTMgr,
110 std::move(ISMBuilder)) {
111 this->ES->getMainJITDylib().addGenerator(
112 std::move(ProcessSymbolsGenerator));
113 this->CODLayer.setImplMap(&Imps);
114 this->ES->setDispatchMaterialization(
116 [this](JITDylib &JD, std::unique_ptr<MaterializationUnit> MU) {
117 // FIXME: Switch to move capture once we have C++14.
118 auto SharedMU = std::shared_ptr<MaterializationUnit>(std::move(MU));
119 auto Work = [SharedMU, &JD]() { SharedMU->doMaterialize(JD); };
120 CompileThreads.async(std::move(Work));
122 ExitOnErr(S.addSpeculationRuntime(this->ES->getMainJITDylib(), Mangle));
123 LocalCXXRuntimeOverrides CXXRuntimeoverrides;
124 ExitOnErr(CXXRuntimeoverrides.enable(this->ES->getMainJITDylib(), Mangle));
127 static std::unique_ptr<SectionMemoryManager> createMemMgr() {
128 return std::make_unique<SectionMemoryManager>();
131 std::unique_ptr<ExecutionSession> ES;
132 DataLayout DL;
133 MangleAndInterner Mangle{*ES, DL};
134 ThreadPool CompileThreads{NumThreads};
136 Triple TT;
137 std::unique_ptr<LazyCallThroughManager> LCTMgr;
138 IRCompileLayer CompileLayer;
139 ImplSymbolMap Imps;
140 Speculator S;
141 RTDyldObjectLinkingLayer ObjLayer{*ES, createMemMgr};
142 IRSpeculationLayer SpeculateLayer;
143 CompileOnDemandLayer CODLayer;
146 int main(int argc, char *argv[]) {
147 // Initialize LLVM.
148 InitLLVM X(argc, argv);
150 InitializeNativeTarget();
151 InitializeNativeTargetAsmPrinter();
153 cl::ParseCommandLineOptions(argc, argv, "SpeculativeJIT");
154 ExitOnErr.setBanner(std::string(argv[0]) + ": ");
156 if (NumThreads < 1) {
157 errs() << "Speculative compilation requires one or more dedicated compile "
158 "threads\n";
159 return 1;
162 // Create a JIT instance.
163 auto SJ = ExitOnErr(SpeculativeJIT::Create());
165 // Load the IR inputs.
166 for (const auto &InputFile : InputFiles) {
167 SMDiagnostic Err;
168 auto Ctx = std::make_unique<LLVMContext>();
169 auto M = parseIRFile(InputFile, Err, *Ctx);
170 if (!M) {
171 Err.print(argv[0], errs());
172 return 1;
175 ExitOnErr(SJ->addModule(SJ->getES().getMainJITDylib(),
176 ThreadSafeModule(std::move(M), std::move(Ctx))));
179 // Build an argv array for the JIT'd main.
180 std::vector<const char *> ArgV;
181 ArgV.push_back(argv[0]);
182 for (const auto &InputArg : InputArgv)
183 ArgV.push_back(InputArg.data());
184 ArgV.push_back(nullptr);
186 // Look up the JIT'd main, cast it to a function pointer, then call it.
188 auto MainSym = ExitOnErr(SJ->lookup("main"));
189 int (*Main)(int, const char *[]) =
190 (int (*)(int, const char *[]))MainSym.getAddress();
192 Main(ArgV.size() - 1, ArgV.data());
194 return 0;