examples/Kaleidoscope/include/KaleidoscopeJIT.h

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