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/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/IRTransformLayer.h" 23 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h" 24 #include "llvm/ExecutionEngine/Orc/LambdaResolver.h" 25 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" 26 #include "llvm/ExecutionEngine/RTDyldMemoryManager.h" 27 #include "llvm/ExecutionEngine/SectionMemoryManager.h" 28 #include "llvm/IR/DataLayout.h" 29 #include "llvm/IR/LegacyPassManager.h" 30 #include "llvm/IR/Mangler.h" 31 #include "llvm/Support/DynamicLibrary.h" 32 #include "llvm/Support/Error.h" 33 #include "llvm/Support/raw_ostream.h" 34 #include "llvm/Target/TargetMachine.h" 35 #include "llvm/Transforms/InstCombine/InstCombine.h" 36 #include "llvm/Transforms/Scalar.h" 37 #include "llvm/Transforms/Scalar/GVN.h" 38 #include <algorithm> 39 #include <cassert> 40 #include <cstdlib> 41 #include <map> 42 #include <memory> 43 #include <string> 44 #include <vector> 45 46 class PrototypeAST; 47 class ExprAST; 48 49 /// FunctionAST - This class represents a function definition itself. 50 class FunctionAST { 51 std::unique_ptr<PrototypeAST> Proto; 52 std::unique_ptr<ExprAST> Body; 53 54 public: FunctionAST(std::unique_ptr<PrototypeAST> Proto,std::unique_ptr<ExprAST> Body)55 FunctionAST(std::unique_ptr<PrototypeAST> Proto, 56 std::unique_ptr<ExprAST> Body) 57 : Proto(std::move(Proto)), Body(std::move(Body)) {} 58 59 const PrototypeAST& getProto() const; 60 const std::string& getName() const; 61 llvm::Function *codegen(); 62 }; 63 64 /// This will compile FnAST to IR, rename the function to add the given 65 /// suffix (needed to prevent a name-clash with the function's stub), 66 /// and then take ownership of the module that the function was compiled 67 /// into. 68 std::unique_ptr<llvm::Module> 69 irgenAndTakeOwnership(FunctionAST &FnAST, const std::string &Suffix); 70 71 namespace llvm { 72 namespace orc { 73 74 class KaleidoscopeJIT { 75 private: 76 ExecutionSession ES; 77 std::shared_ptr<SymbolResolver> Resolver; 78 std::unique_ptr<TargetMachine> TM; 79 const DataLayout DL; 80 RTDyldObjectLinkingLayer ObjectLayer; 81 IRCompileLayer<decltype(ObjectLayer), SimpleCompiler> CompileLayer; 82 83 using OptimizeFunction = 84 std::function<std::unique_ptr<Module>(std::unique_ptr<Module>)>; 85 86 IRTransformLayer<decltype(CompileLayer), OptimizeFunction> OptimizeLayer; 87 88 std::unique_ptr<JITCompileCallbackManager> CompileCallbackMgr; 89 std::unique_ptr<IndirectStubsManager> IndirectStubsMgr; 90 91 public: KaleidoscopeJIT()92 KaleidoscopeJIT() 93 : Resolver(createLegacyLookupResolver( 94 ES, 95 [this](const std::string &Name) -> JITSymbol { 96 if (auto Sym = IndirectStubsMgr->findStub(Name, false)) 97 return Sym; 98 if (auto Sym = OptimizeLayer.findSymbol(Name, false)) 99 return Sym; 100 else if (auto Err = Sym.takeError()) 101 return std::move(Err); 102 if (auto SymAddr = 103 RTDyldMemoryManager::getSymbolAddressInProcess(Name)) 104 return JITSymbol(SymAddr, JITSymbolFlags::Exported); 105 return nullptr; 106 }, 107 [](Error Err) { cantFail(std::move(Err), "lookupFlags failed"); })), 108 TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()), 109 ObjectLayer(ES, 110 [this](VModuleKey K) { 111 return RTDyldObjectLinkingLayer::Resources{ 112 std::make_shared<SectionMemoryManager>(), Resolver}; 113 }), 114 CompileLayer(ObjectLayer, SimpleCompiler(*TM)), 115 OptimizeLayer(CompileLayer, 116 [this](std::unique_ptr<Module> M) { 117 return optimizeModule(std::move(M)); 118 }), 119 CompileCallbackMgr(orc::createLocalCompileCallbackManager( 120 TM->getTargetTriple(), ES, 0)) { 121 auto IndirectStubsMgrBuilder = 122 orc::createLocalIndirectStubsManagerBuilder(TM->getTargetTriple()); 123 IndirectStubsMgr = IndirectStubsMgrBuilder(); 124 llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr); 125 } 126 getTargetMachine()127 TargetMachine &getTargetMachine() { return *TM; } 128 addModule(std::unique_ptr<Module> M)129 VModuleKey addModule(std::unique_ptr<Module> M) { 130 // Add the module to the JIT with a new VModuleKey. 131 auto K = ES.allocateVModule(); 132 cantFail(OptimizeLayer.addModule(K, std::move(M))); 133 return K; 134 } 135 addFunctionAST(std::unique_ptr<FunctionAST> FnAST)136 Error addFunctionAST(std::unique_ptr<FunctionAST> FnAST) { 137 // Move ownership of FnAST to a shared pointer - C++11 lambdas don't support 138 // capture-by-move, which is be required for unique_ptr. 139 auto SharedFnAST = std::shared_ptr<FunctionAST>(std::move(FnAST)); 140 141 // Set the action to compile our AST. This lambda will be run if/when 142 // execution hits the compile callback (via the stub). 143 // 144 // The steps to compile are: 145 // (1) IRGen the function. 146 // (2) Add the IR module to the JIT to make it executable like any other 147 // module. 148 // (3) Use findSymbol to get the address of the compiled function. 149 // (4) Update the stub pointer to point at the implementation so that 150 /// subsequent calls go directly to it and bypass the compiler. 151 // (5) Return the address of the implementation: this lambda will actually 152 // be run inside an attempted call to the function, and we need to 153 // continue on to the implementation to complete the attempted call. 154 // The JIT runtime (the resolver block) will use the return address of 155 // this function as the address to continue at once it has reset the 156 // CPU state to what it was immediately before the call. 157 auto CompileAction = [this, SharedFnAST]() { 158 auto M = irgenAndTakeOwnership(*SharedFnAST, "$impl"); 159 addModule(std::move(M)); 160 auto Sym = findSymbol(SharedFnAST->getName() + "$impl"); 161 assert(Sym && "Couldn't find compiled function?"); 162 JITTargetAddress SymAddr = cantFail(Sym.getAddress()); 163 if (auto Err = IndirectStubsMgr->updatePointer( 164 mangle(SharedFnAST->getName()), SymAddr)) { 165 logAllUnhandledErrors(std::move(Err), errs(), 166 "Error updating function pointer: "); 167 exit(1); 168 } 169 170 return SymAddr; 171 }; 172 173 // Create a CompileCallback using the CompileAction - this is the re-entry 174 // point into the compiler for functions that haven't been compiled yet. 175 auto CCAddr = cantFail( 176 CompileCallbackMgr->getCompileCallback(std::move(CompileAction))); 177 178 // Create an indirect stub. This serves as the functions "canonical 179 // definition" - an unchanging (constant address) entry point to the 180 // function implementation. 181 // Initially we point the stub's function-pointer at the compile callback 182 // that we just created. When the compile action for the callback is run we 183 // will update the stub's function pointer to point at the function 184 // implementation that we just implemented. 185 if (auto Err = IndirectStubsMgr->createStub( 186 mangle(SharedFnAST->getName()), CCAddr, JITSymbolFlags::Exported)) 187 return Err; 188 189 return Error::success(); 190 } 191 findSymbol(const std::string Name)192 JITSymbol findSymbol(const std::string Name) { 193 return OptimizeLayer.findSymbol(mangle(Name), true); 194 } 195 removeModule(VModuleKey K)196 void removeModule(VModuleKey K) { 197 cantFail(OptimizeLayer.removeModule(K)); 198 } 199 200 private: mangle(const std::string & Name)201 std::string mangle(const std::string &Name) { 202 std::string MangledName; 203 raw_string_ostream MangledNameStream(MangledName); 204 Mangler::getNameWithPrefix(MangledNameStream, Name, DL); 205 return MangledNameStream.str(); 206 } 207 optimizeModule(std::unique_ptr<Module> M)208 std::unique_ptr<Module> optimizeModule(std::unique_ptr<Module> M) { 209 // Create a function pass manager. 210 auto FPM = llvm::make_unique<legacy::FunctionPassManager>(M.get()); 211 212 // Add some optimizations. 213 FPM->add(createInstructionCombiningPass()); 214 FPM->add(createReassociatePass()); 215 FPM->add(createGVNPass()); 216 FPM->add(createCFGSimplificationPass()); 217 FPM->doInitialization(); 218 219 // Run the optimizations over all functions in the module being added to 220 // the JIT. 221 for (auto &F : *M) 222 FPM->run(F); 223 224 return M; 225 } 226 }; 227 228 } // end namespace orc 229 } // end namespace llvm 230 231 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H 232