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