llvm/examples/Kaleidoscope/BuildingAJIT/Chapter3/KaleidoscopeJIT.h

   1 //===- KaleidoscopeJIT.h - A simple JIT for Kaleidoscope --------*- C++ -*-===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // Contains a simple JIT definition for use in the kaleidoscope tutorials.
  10 //
  11 //===----------------------------------------------------------------------===//
  12
  13 #ifndef LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
  14 #define LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
  15
  16 #include "llvm/ADT/StringRef.h"
  17 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
  18 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
  19 #include "llvm/ExecutionEngine/Orc/Core.h"
  20 #include "llvm/ExecutionEngine/Orc/EPCIndirectionUtils.h"
  21 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
  22 #include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
  23 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
  24 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
  25 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
  26 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
  27 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
  28 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
  29 #include "llvm/IR/DataLayout.h"
  30 #include "llvm/IR/LLVMContext.h"
  31 #include "llvm/IR/LegacyPassManager.h"
  32 #include "llvm/Transforms/InstCombine/InstCombine.h"
  33 #include "llvm/Transforms/Scalar.h"
  34 #include "llvm/Transforms/Scalar/GVN.h"
  35 #include <memory>
  36
  37 namespace llvm {
  38 namespace orc {
  39
  40 class KaleidoscopeJIT {
  41 private:
  42   std::unique_ptr<ExecutionSession> ES;
  43   std::unique_ptr<EPCIndirectionUtils> EPCIU;
  44
  45   DataLayout DL;
  46   MangleAndInterner Mangle;
  47
  48   RTDyldObjectLinkingLayer ObjectLayer;
  49   IRCompileLayer CompileLayer;
  50   IRTransformLayer OptimizeLayer;
  51   CompileOnDemandLayer CODLayer;
  52
  53   JITDylib &MainJD;
  54
  55   static void handleLazyCallThroughError() {
  56     errs() << "LazyCallThrough error: Could not find function body";
  57     exit(1);
  58   }
  59
  60 public:
  61   KaleidoscopeJIT(std::unique_ptr<ExecutionSession> ES,
  62                   std::unique_ptr<EPCIndirectionUtils> EPCIU,
  63                   JITTargetMachineBuilder JTMB, DataLayout DL)
  64       : ES(std::move(ES)), EPCIU(std::move(EPCIU)), DL(std::move(DL)),
  65         Mangle(*this->ES, this->DL),
  66         ObjectLayer(*this->ES,
  67                     []() { return std::make_unique<SectionMemoryManager>(); }),
  68         CompileLayer(*this->ES, ObjectLayer,
  69                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
  70         OptimizeLayer(*this->ES, CompileLayer, optimizeModule),
  71         CODLayer(*this->ES, OptimizeLayer,
  72                  this->EPCIU->getLazyCallThroughManager(),
  73                  [this] { return this->EPCIU->createIndirectStubsManager(); }),
  74         MainJD(this->ES->createBareJITDylib("<main>")) {
  75     MainJD.addGenerator(
  76         cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
  77             DL.getGlobalPrefix())));
  78   }
  79
  80   ~KaleidoscopeJIT() {
  81     if (auto Err = ES->endSession())
  82       ES->reportError(std::move(Err));
  83     if (auto Err = EPCIU->cleanup())
  84       ES->reportError(std::move(Err));
  85   }
  86
  87   static Expected<std::unique_ptr<KaleidoscopeJIT>> Create() {
  88     auto EPC = SelfExecutorProcessControl::Create();
  89     if (!EPC)
  90       return EPC.takeError();
  91
  92     auto ES = std::make_unique<ExecutionSession>(std::move(*EPC));
  93
  94     auto EPCIU = EPCIndirectionUtils::Create(*ES);
  95     if (!EPCIU)
  96       return EPCIU.takeError();
  97
  98     (*EPCIU)->createLazyCallThroughManager(
  99         *ES, ExecutorAddr::fromPtr(&handleLazyCallThroughError));
 100
 101     if (auto Err = setUpInProcessLCTMReentryViaEPCIU(**EPCIU))
 102       return std::move(Err);
 103
 104     JITTargetMachineBuilder JTMB(
 105         ES->getExecutorProcessControl().getTargetTriple());
 106
 107     auto DL = JTMB.getDefaultDataLayoutForTarget();
 108     if (!DL)
 109       return DL.takeError();
 110
 111     return std::make_unique<KaleidoscopeJIT>(std::move(ES), std::move(*EPCIU),
 112                                              std::move(JTMB), std::move(*DL));
 113   }
 114
 115   const DataLayout &getDataLayout() const { return DL; }
 116
 117   JITDylib &getMainJITDylib() { return MainJD; }
 118
 119   Error addModule(ThreadSafeModule TSM, ResourceTrackerSP RT = nullptr) {
 120     if (!RT)
 121       RT = MainJD.getDefaultResourceTracker();
 122
 123     return CODLayer.add(RT, std::move(TSM));
 124   }
 125
 126   Expected<ExecutorSymbolDef> lookup(StringRef Name) {
 127     return ES->lookup({&MainJD}, Mangle(Name.str()));
 128   }
 129
 130 private:
 131   static Expected<ThreadSafeModule>
 132   optimizeModule(ThreadSafeModule TSM, const MaterializationResponsibility &R) {
 133     TSM.withModuleDo([](Module &M) {
 134       // Create a function pass manager.
 135       auto FPM = std::make_unique<legacy::FunctionPassManager>(&M);
 136
 137       // Add some optimizations.
 138       FPM->add(createInstructionCombiningPass());
 139       FPM->add(createReassociatePass());
 140       FPM->add(createGVNPass());
 141       FPM->add(createCFGSimplificationPass());
 142       FPM->doInitialization();
 143
 144       // Run the optimizations over all functions in the module being added to
 145       // the JIT.
 146       for (auto &F : M)
 147         FPM->run(F);
 148     });
 149
 150     return std::move(TSM);
 151   }
 152 };
 153
 154 } // end namespace orc
 155 } // end namespace llvm
 156
 157 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H