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 KaleidoscopeJIT()48 KaleidoscopeJIT() 49 : TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()), 50 CompileLayer(ObjectLayer, SimpleCompiler(*TM)) { 51 llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr); 52 } 53 getTargetMachine()54 TargetMachine &getTargetMachine() { return *TM; } 55 addModule(std::unique_ptr<Module> M)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 removeModule(ModuleHandleT H)75 void removeModule(ModuleHandleT H) { 76 ModuleHandles.erase( 77 std::find(ModuleHandles.begin(), ModuleHandles.end(), H)); 78 CompileLayer.removeModuleSet(H); 79 } 80 findSymbol(const std::string Name)81 JITSymbol findSymbol(const std::string Name) { 82 return findMangledSymbol(mangle(Name)); 83 } 84 85 private: mangle(const std::string & Name)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 singletonSet(T t)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 findMangledSymbol(const std::string & Name)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