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/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