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/ADT/iterator_range.h"
     18 #include "llvm/ADT/STLExtras.h"
     19 #include "llvm/ExecutionEngine/ExecutionEngine.h"
     20 #include "llvm/ExecutionEngine/JITSymbolFlags.h"
     21 #include "llvm/ExecutionEngine/RTDyldMemoryManager.h"
     22 #include "llvm/ExecutionEngine/RuntimeDyld.h"
     23 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
     24 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
     25 #include "llvm/ExecutionEngine/Orc/JITSymbol.h"
     26 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
     27 #include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
     28 #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
     29 #include "llvm/IR/DataLayout.h"
     30 #include "llvm/IR/Mangler.h"
     31 #include "llvm/Support/DynamicLibrary.h"
     32 #include "llvm/Support/raw_ostream.h"
     33 #include "llvm/Target/TargetMachine.h"
     34 #include <algorithm>
     35 #include <memory>
     36 #include <string>
     37 #include <vector>
     38 
     39 namespace llvm {
     40 namespace orc {
     41 
     42 class KaleidoscopeJIT {
     43 public:
     44   typedef ObjectLinkingLayer<> ObjLayerT;
     45   typedef IRCompileLayer<ObjLayerT> CompileLayerT;
     46   typedef CompileLayerT::ModuleSetHandleT ModuleHandleT;
     47 
     48   KaleidoscopeJIT()
     49       : TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
     50         CompileLayer(ObjectLayer, SimpleCompiler(*TM)) {
     51     llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
     52   }
     53 
     54   TargetMachine &getTargetMachine() { return *TM; }
     55 
     56   ModuleHandleT addModule(std::unique_ptr<Module> M) {
     57     // We need a memory manager to allocate memory and resolve symbols for this
     58     // new module. Create one that resolves symbols by looking back into the
     59     // JIT.
     60     auto Resolver = createLambdaResolver(
     61         [&](const std::string &Name) {
     62           if (auto Sym = findMangledSymbol(Name))
     63             return Sym.toRuntimeDyldSymbol();
     64           return RuntimeDyld::SymbolInfo(nullptr);
     65         },
     66         [](const std::string &S) { return nullptr; });
     67     auto H = CompileLayer.addModuleSet(singletonSet(std::move(M)),
     68                                        make_unique<SectionMemoryManager>(),
     69                                        std::move(Resolver));
     70 
     71     ModuleHandles.push_back(H);
     72     return H;
     73   }
     74 
     75   void removeModule(ModuleHandleT H) {
     76     ModuleHandles.erase(
     77         std::find(ModuleHandles.begin(), ModuleHandles.end(), H));
     78     CompileLayer.removeModuleSet(H);
     79   }
     80 
     81   JITSymbol findSymbol(const std::string Name) {
     82     return findMangledSymbol(mangle(Name));
     83   }
     84 
     85 private:
     86   std::string mangle(const std::string &Name) {
     87     std::string MangledName;
     88     {
     89       raw_string_ostream MangledNameStream(MangledName);
     90       Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
     91     }
     92     return MangledName;
     93   }
     94 
     95   template <typename T> static std::vector<T> singletonSet(T t) {
     96     std::vector<T> Vec;
     97     Vec.push_back(std::move(t));
     98     return Vec;
     99   }
    100 
    101   JITSymbol findMangledSymbol(const std::string &Name) {
    102     // Search modules in reverse order: from last added to first added.
    103     // This is the opposite of the usual search order for dlsym, but makes more
    104     // sense in a REPL where we want to bind to the newest available definition.
    105     for (auto H : make_range(ModuleHandles.rbegin(), ModuleHandles.rend()))
    106       if (auto Sym = CompileLayer.findSymbolIn(H, Name, true))
    107         return Sym;
    108 
    109     // If we can't find the symbol in the JIT, try looking in the host process.
    110     if (auto SymAddr = RTDyldMemoryManager::getSymbolAddressInProcess(Name))
    111       return JITSymbol(SymAddr, JITSymbolFlags::Exported);
    112 
    113     return nullptr;
    114   }
    115 
    116   std::unique_ptr<TargetMachine> TM;
    117   const DataLayout DL;
    118   ObjLayerT ObjectLayer;
    119   CompileLayerT CompileLayer;
    120   std::vector<ModuleHandleT> ModuleHandles;
    121 };
    122 
    123 } // end namespace orc
    124 } // end namespace llvm
    125 
    126 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
    127