examples/Kaleidoscope/include/KaleidoscopeJIT.h

   1 //===----- KaleidoscopeJIT.h - A simple JIT for Kaleidoscope ----*- C++ -*-===//
   2 //
   3 //                     The LLVM Compiler Infrastructure
   4 //
   5 // This file is distributed under the University of Illinois Open Source
   6 // License. See LICENSE.TXT for details.
   7 //
   8 //===----------------------------------------------------------------------===//
   9 //
  10 // Contains a simple JIT definition for use in the kaleidoscope tutorials.
  11 //
  12 //===----------------------------------------------------------------------===//
  13
  14 #ifndef LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
  15 #define LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
  16
  17 #include "llvm/ExecutionEngine/ExecutionEngine.h"
  18 #include "llvm/ExecutionEngine/RTDyldMemoryManager.h"
  19 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
  20 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
  21 #include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
  22 #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
  23 #include "llvm/IR/Mangler.h"
  24 #include "llvm/Support/DynamicLibrary.h"
  25
  26 namespace llvm {
  27 namespace orc {
  28
  29 class KaleidoscopeJIT {
  30 public:
  31   typedef ObjectLinkingLayer<> ObjLayerT;
  32   typedef IRCompileLayer<ObjLayerT> CompileLayerT;
  33   typedef CompileLayerT::ModuleSetHandleT ModuleHandleT;
  34
  35   KaleidoscopeJIT()
  36       : TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
  37         CompileLayer(ObjectLayer, SimpleCompiler(*TM)) {
  38     llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
  39   }
  40
  41   TargetMachine &getTargetMachine() { return *TM; }
  42
  43   ModuleHandleT addModule(std::unique_ptr<Module> M) {
  44     // We need a memory manager to allocate memory and resolve symbols for this
  45     // new module. Create one that resolves symbols by looking back into the
  46     // JIT.
  47     auto Resolver = createLambdaResolver(
  48         [&](const std::string &Name) {
  49           if (auto Sym = findMangledSymbol(Name))
  50             return RuntimeDyld::SymbolInfo(Sym.getAddress(), Sym.getFlags());
  51           return RuntimeDyld::SymbolInfo(nullptr);
  52         },
  53         [](const std::string &S) { return nullptr; });
  54     auto H = CompileLayer.addModuleSet(singletonSet(std::move(M)),
  55                                        make_unique<SectionMemoryManager>(),
  56                                        std::move(Resolver));
  57
  58     ModuleHandles.push_back(H);
  59     return H;
  60   }
  61
  62   void removeModule(ModuleHandleT H) {
  63     ModuleHandles.erase(
  64         std::find(ModuleHandles.begin(), ModuleHandles.end(), H));
  65     CompileLayer.removeModuleSet(H);
  66   }
  67
  68   JITSymbol findSymbol(const std::string Name) {
  69     return findMangledSymbol(mangle(Name));
  70   }
  71
  72 private:
  73
  74   std::string mangle(const std::string &Name) {
  75     std::string MangledName;
  76     {
  77       raw_string_ostream MangledNameStream(MangledName);
  78       Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
  79     }
  80     return MangledName;
  81   }
  82
  83   template <typename T> static std::vector<T> singletonSet(T t) {
  84     std::vector<T> Vec;
  85     Vec.push_back(std::move(t));
  86     return Vec;
  87   }
  88
  89   JITSymbol findMangledSymbol(const std::string &Name) {
  90     // Search modules in reverse order: from last added to first added.
  91     // This is the opposite of the usual search order for dlsym, but makes more
  92     // sense in a REPL where we want to bind to the newest available definition.
  93     for (auto H : make_range(ModuleHandles.rbegin(), ModuleHandles.rend()))
  94       if (auto Sym = CompileLayer.findSymbolIn(H, Name, true))
  95         return Sym;
  96
  97     // If we can't find the symbol in the JIT, try looking in the host process.
  98     if (auto SymAddr = RTDyldMemoryManager::getSymbolAddressInProcess(Name))
  99       return JITSymbol(SymAddr, JITSymbolFlags::Exported);
 100
 101     return nullptr;
 102   }
 103
 104   std::unique_ptr<TargetMachine> TM;
 105   const DataLayout DL;
 106   ObjLayerT ObjectLayer;
 107   CompileLayerT CompileLayer;
 108   std::vector<ModuleHandleT> ModuleHandles;
 109 };
 110
 111 } // End namespace orc.
 112 } // End namespace llvm
 113
 114 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H