llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp

   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 }