Home | History | Annotate | Download | only in include 
      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
    115