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 "RemoteJITUtils.h" 18 #include "llvm/ADT/STLExtras.h" 19 #include "llvm/ExecutionEngine/ExecutionEngine.h" 20 #include "llvm/ExecutionEngine/RuntimeDyld.h" 21 #include "llvm/ExecutionEngine/SectionMemoryManager.h" 22 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h" 23 #include "llvm/ExecutionEngine/Orc/CompileUtils.h" 24 #include "llvm/ExecutionEngine/Orc/JITSymbol.h" 25 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" 26 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h" 27 #include "llvm/ExecutionEngine/Orc/LambdaResolver.h" 28 #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h" 29 #include "llvm/ExecutionEngine/Orc/OrcRemoteTargetClient.h" 30 #include "llvm/IR/DataLayout.h" 31 #include "llvm/IR/Mangler.h" 32 #include "llvm/Support/DynamicLibrary.h" 33 #include "llvm/Support/raw_ostream.h" 34 #include "llvm/Target/TargetMachine.h" 35 #include <algorithm> 36 #include <memory> 37 #include <string> 38 #include <vector> 39 40 class PrototypeAST; 41 class ExprAST; 42 43 /// FunctionAST - This class represents a function definition itself. 44 class FunctionAST { 45 std::unique_ptr<PrototypeAST> Proto; 46 std::unique_ptr<ExprAST> Body; 47 48 public: 49 FunctionAST(std::unique_ptr<PrototypeAST> Proto, 50 std::unique_ptr<ExprAST> Body) 51 : Proto(std::move(Proto)), Body(std::move(Body)) {} 52 const PrototypeAST& getProto() const; 53 const std::string& getName() const; 54 llvm::Function *codegen(); 55 }; 56 57 /// This will compile FnAST to IR, rename the function to add the given 58 /// suffix (needed to prevent a name-clash with the function's stub), 59 /// and then take ownership of the module that the function was compiled 60 /// into. 61 std::unique_ptr<llvm::Module> 62 irgenAndTakeOwnership(FunctionAST &FnAST, const std::string &Suffix); 63 64 namespace llvm { 65 namespace orc { 66 67 // Typedef the remote-client API. 68 typedef remote::OrcRemoteTargetClient<FDRPCChannel> MyRemote; 69 70 class KaleidoscopeJIT { 71 private: 72 MyRemote &Remote; 73 std::unique_ptr<TargetMachine> TM; 74 const DataLayout DL; 75 JITCompileCallbackManager *CompileCallbackMgr; 76 std::unique_ptr<IndirectStubsManager> IndirectStubsMgr; 77 ObjectLinkingLayer<> ObjectLayer; 78 IRCompileLayer<decltype(ObjectLayer)> CompileLayer; 79 80 typedef std::function<std::unique_ptr<Module>(std::unique_ptr<Module>)> 81 OptimizeFunction; 82 83 IRTransformLayer<decltype(CompileLayer), OptimizeFunction> OptimizeLayer; 84 85 public: 86 typedef decltype(OptimizeLayer)::ModuleSetHandleT ModuleHandle; 87 88 KaleidoscopeJIT(MyRemote &Remote) 89 : Remote(Remote), 90 TM(EngineBuilder().selectTarget()), 91 DL(TM->createDataLayout()), 92 CompileLayer(ObjectLayer, SimpleCompiler(*TM)), 93 OptimizeLayer(CompileLayer, 94 [this](std::unique_ptr<Module> M) { 95 return optimizeModule(std::move(M)); 96 }) { 97 auto CCMgrOrErr = Remote.enableCompileCallbacks(0); 98 if (!CCMgrOrErr) { 99 logAllUnhandledErrors(CCMgrOrErr.takeError(), errs(), 100 "Error enabling remote compile callbacks:"); 101 exit(1); 102 } 103 CompileCallbackMgr = &*CCMgrOrErr; 104 std::unique_ptr<MyRemote::RCIndirectStubsManager> ISM; 105 if (auto Err = Remote.createIndirectStubsManager(ISM)) { 106 logAllUnhandledErrors(std::move(Err), errs(), 107 "Error creating indirect stubs manager:"); 108 exit(1); 109 } 110 IndirectStubsMgr = std::move(ISM); 111 llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr); 112 } 113 114 TargetMachine &getTargetMachine() { return *TM; } 115 116 ModuleHandle addModule(std::unique_ptr<Module> M) { 117 118 // Build our symbol resolver: 119 // Lambda 1: Look back into the JIT itself to find symbols that are part of 120 // the same "logical dylib". 121 // Lambda 2: Search for external symbols in the host process. 122 auto Resolver = createLambdaResolver( 123 [&](const std::string &Name) { 124 if (auto Sym = IndirectStubsMgr->findStub(Name, false)) 125 return Sym.toRuntimeDyldSymbol(); 126 if (auto Sym = OptimizeLayer.findSymbol(Name, false)) 127 return Sym.toRuntimeDyldSymbol(); 128 return RuntimeDyld::SymbolInfo(nullptr); 129 }, 130 [&](const std::string &Name) { 131 if (auto AddrOrErr = Remote.getSymbolAddress(Name)) 132 return RuntimeDyld::SymbolInfo(*AddrOrErr, 133 JITSymbolFlags::Exported); 134 else { 135 logAllUnhandledErrors(AddrOrErr.takeError(), errs(), 136 "Error resolving remote symbol:"); 137 exit(1); 138 } 139 return RuntimeDyld::SymbolInfo(nullptr); 140 }); 141 142 std::unique_ptr<MyRemote::RCMemoryManager> MemMgr; 143 if (auto Err = Remote.createRemoteMemoryManager(MemMgr)) { 144 logAllUnhandledErrors(std::move(Err), errs(), 145 "Error creating remote memory manager:"); 146 exit(1); 147 } 148 149 // Build a singlton module set to hold our module. 150 std::vector<std::unique_ptr<Module>> Ms; 151 Ms.push_back(std::move(M)); 152 153 // Add the set to the JIT with the resolver we created above and a newly 154 // created SectionMemoryManager. 155 return OptimizeLayer.addModuleSet(std::move(Ms), 156 std::move(MemMgr), 157 std::move(Resolver)); 158 } 159 160 Error addFunctionAST(std::unique_ptr<FunctionAST> FnAST) { 161 // Create a CompileCallback - this is the re-entry point into the compiler 162 // for functions that haven't been compiled yet. 163 auto CCInfo = CompileCallbackMgr->getCompileCallback(); 164 165 // Create an indirect stub. This serves as the functions "canonical 166 // definition" - an unchanging (constant address) entry point to the 167 // function implementation. 168 // Initially we point the stub's function-pointer at the compile callback 169 // that we just created. In the compile action for the callback (see below) 170 // we will update the stub's function pointer to point at the function 171 // implementation that we just implemented. 172 if (auto Err = IndirectStubsMgr->createStub(mangle(FnAST->getName()), 173 CCInfo.getAddress(), 174 JITSymbolFlags::Exported)) 175 return Err; 176 177 // Move ownership of FnAST to a shared pointer - C++11 lambdas don't support 178 // capture-by-move, which is be required for unique_ptr. 179 auto SharedFnAST = std::shared_ptr<FunctionAST>(std::move(FnAST)); 180 181 // Set the action to compile our AST. This lambda will be run if/when 182 // execution hits the compile callback (via the stub). 183 // 184 // The steps to compile are: 185 // (1) IRGen the function. 186 // (2) Add the IR module to the JIT to make it executable like any other 187 // module. 188 // (3) Use findSymbol to get the address of the compiled function. 189 // (4) Update the stub pointer to point at the implementation so that 190 /// subsequent calls go directly to it and bypass the compiler. 191 // (5) Return the address of the implementation: this lambda will actually 192 // be run inside an attempted call to the function, and we need to 193 // continue on to the implementation to complete the attempted call. 194 // The JIT runtime (the resolver block) will use the return address of 195 // this function as the address to continue at once it has reset the 196 // CPU state to what it was immediately before the call. 197 CCInfo.setCompileAction( 198 [this, SharedFnAST]() { 199 auto M = irgenAndTakeOwnership(*SharedFnAST, "$impl"); 200 addModule(std::move(M)); 201 auto Sym = findSymbol(SharedFnAST->getName() + "$impl"); 202 assert(Sym && "Couldn't find compiled function?"); 203 TargetAddress SymAddr = Sym.getAddress(); 204 if (auto Err = 205 IndirectStubsMgr->updatePointer(mangle(SharedFnAST->getName()), 206 SymAddr)) { 207 logAllUnhandledErrors(std::move(Err), errs(), 208 "Error updating function pointer: "); 209 exit(1); 210 } 211 212 return SymAddr; 213 }); 214 215 return Error::success(); 216 } 217 218 Error executeRemoteExpr(TargetAddress ExprAddr) { 219 return Remote.callVoidVoid(ExprAddr); 220 } 221 222 JITSymbol findSymbol(const std::string Name) { 223 return OptimizeLayer.findSymbol(mangle(Name), true); 224 } 225 226 void removeModule(ModuleHandle H) { 227 OptimizeLayer.removeModuleSet(H); 228 } 229 230 private: 231 232 std::string mangle(const std::string &Name) { 233 std::string MangledName; 234 raw_string_ostream MangledNameStream(MangledName); 235 Mangler::getNameWithPrefix(MangledNameStream, Name, DL); 236 return MangledNameStream.str(); 237 } 238 239 std::unique_ptr<Module> optimizeModule(std::unique_ptr<Module> M) { 240 // Create a function pass manager. 241 auto FPM = llvm::make_unique<legacy::FunctionPassManager>(M.get()); 242 243 // Add some optimizations. 244 FPM->add(createInstructionCombiningPass()); 245 FPM->add(createReassociatePass()); 246 FPM->add(createGVNPass()); 247 FPM->add(createCFGSimplificationPass()); 248 FPM->doInitialization(); 249 250 // Run the optimizations over all functions in the module being added to 251 // the JIT. 252 for (auto &F : *M) 253 FPM->run(F); 254 255 return M; 256 } 257 258 }; 259 260 } // end namespace orc 261 } // end namespace llvm 262 263 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H 264
