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