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/STLExtras.h"
     18 #include "llvm/ADT/iterator_range.h"
     19 #include "llvm/ExecutionEngine/ExecutionEngine.h"
     20 #include "llvm/ExecutionEngine/JITSymbol.h"
     21 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
     22 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
     23 #include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
     24 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
     25 #include "llvm/ExecutionEngine/RTDyldMemoryManager.h"
     26 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
     27 #include "llvm/IR/DataLayout.h"
     28 #include "llvm/IR/Mangler.h"
     29 #include "llvm/Support/DynamicLibrary.h"
     30 #include "llvm/Support/raw_ostream.h"
     31 #include "llvm/Target/TargetMachine.h"
     32 #include <algorithm>
     33 #include <map>
     34 #include <memory>
     35 #include <string>
     36 #include <vector>
     37 
     38 namespace llvm {
     39 namespace orc {
     40 
     41 class KaleidoscopeJIT {
     42 public:
     43   using ObjLayerT = RTDyldObjectLinkingLayer;
     44   using CompileLayerT = IRCompileLayer<ObjLayerT, SimpleCompiler>;
     45 
     46   KaleidoscopeJIT()
     47       : Resolver(createLegacyLookupResolver(
     48             ES,
     49             [this](const std::string &Name) {
     50               return ObjectLayer.findSymbol(Name, true);
     51             },
     52             [](Error Err) { cantFail(std::move(Err), "lookupFlags failed"); })),
     53         TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
     54         ObjectLayer(ES,
     55                     [this](VModuleKey) {
     56                       return ObjLayerT::Resources{
     57                           std::make_shared<SectionMemoryManager>(), Resolver};
     58                     }),
     59         CompileLayer(ObjectLayer, SimpleCompiler(*TM)) {
     60     llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
     61   }
     62 
     63   TargetMachine &getTargetMachine() { return *TM; }
     64 
     65   VModuleKey addModule(std::unique_ptr<Module> M) {
     66     auto K = ES.allocateVModule();
     67     cantFail(CompileLayer.addModule(K, std::move(M)));
     68     ModuleKeys.push_back(K);
     69     return K;
     70   }
     71 
     72   void removeModule(VModuleKey K) {
     73     ModuleKeys.erase(find(ModuleKeys, K));
     74     cantFail(CompileLayer.removeModule(K));
     75   }
     76 
     77   JITSymbol findSymbol(const std::string Name) {
     78     return findMangledSymbol(mangle(Name));
     79   }
     80 
     81 private:
     82   std::string mangle(const std::string &Name) {
     83     std::string MangledName;
     84     {
     85       raw_string_ostream MangledNameStream(MangledName);
     86       Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
     87     }
     88     return MangledName;
     89   }
     90 
     91   JITSymbol findMangledSymbol(const std::string &Name) {
     92 #ifdef _WIN32
     93     // The symbol lookup of ObjectLinkingLayer uses the SymbolRef::SF_Exported
     94     // flag to decide whether a symbol will be visible or not, when we call
     95     // IRCompileLayer::findSymbolIn with ExportedSymbolsOnly set to true.
     96     //
     97     // But for Windows COFF objects, this flag is currently never set.
     98     // For a potential solution see: https://reviews.llvm.org/rL258665
     99     // For now, we allow non-exported symbols on Windows as a workaround.
    100     const bool ExportedSymbolsOnly = false;
    101 #else
    102     const bool ExportedSymbolsOnly = true;
    103 #endif
    104 
    105     // Search modules in reverse order: from last added to first added.
    106     // This is the opposite of the usual search order for dlsym, but makes more
    107     // sense in a REPL where we want to bind to the newest available definition.
    108     for (auto H : make_range(ModuleKeys.rbegin(), ModuleKeys.rend()))
    109       if (auto Sym = CompileLayer.findSymbolIn(H, Name, ExportedSymbolsOnly))
    110         return Sym;
    111 
    112     // If we can't find the symbol in the JIT, try looking in the host process.
    113     if (auto SymAddr = RTDyldMemoryManager::getSymbolAddressInProcess(Name))
    114       return JITSymbol(SymAddr, JITSymbolFlags::Exported);
    115 
    116 #ifdef _WIN32
    117     // For Windows retry without "_" at beginning, as RTDyldMemoryManager uses
    118     // GetProcAddress and standard libraries like msvcrt.dll use names
    119     // with and without "_" (for example "_itoa" but "sin").
    120     if (Name.length() > 2 && Name[0] == '_')
    121       if (auto SymAddr =
    122               RTDyldMemoryManager::getSymbolAddressInProcess(Name.substr(1)))
    123         return JITSymbol(SymAddr, JITSymbolFlags::Exported);
    124 #endif
    125 
    126     return nullptr;
    127   }
    128 
    129   ExecutionSession ES;
    130   std::shared_ptr<SymbolResolver> Resolver;
    131   std::unique_ptr<TargetMachine> TM;
    132   const DataLayout DL;
    133   ObjLayerT ObjectLayer;
    134   CompileLayerT CompileLayer;
    135   std::vector<VModuleKey> ModuleKeys;
    136 };
    137 
    138 } // end namespace orc
    139 } // end namespace llvm
    140 
    141 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
    142