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/CompileOnDemandLayer.h" 21 #include "llvm/ExecutionEngine/Orc/CompileUtils.h" 22 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" 23 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.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/RuntimeDyld.h" 28 #include "llvm/ExecutionEngine/SectionMemoryManager.h" 29 #include "llvm/IR/DataLayout.h" 30 #include "llvm/IR/LegacyPassManager.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 "llvm/Transforms/InstCombine/InstCombine.h" 36 #include "llvm/Transforms/Scalar.h" 37 #include "llvm/Transforms/Scalar/GVN.h" 38 #include <algorithm> 39 #include <map> 40 #include <memory> 41 #include <set> 42 #include <string> 43 #include <vector> 44 45 namespace llvm { 46 namespace orc { 47 48 class KaleidoscopeJIT { 49 private: 50 ExecutionSession ES; 51 std::map<VModuleKey, std::shared_ptr<SymbolResolver>> Resolvers; 52 std::unique_ptr<TargetMachine> TM; 53 const DataLayout DL; 54 RTDyldObjectLinkingLayer ObjectLayer; 55 IRCompileLayer<decltype(ObjectLayer), SimpleCompiler> CompileLayer; 56 57 using OptimizeFunction = 58 std::function<std::unique_ptr<Module>(std::unique_ptr<Module>)>; 59 60 IRTransformLayer<decltype(CompileLayer), OptimizeFunction> OptimizeLayer; 61 62 std::unique_ptr<JITCompileCallbackManager> CompileCallbackManager; 63 CompileOnDemandLayer<decltype(OptimizeLayer)> CODLayer; 64 65 public: 66 KaleidoscopeJIT() 67 : TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()), 68 ObjectLayer(ES, 69 [this](VModuleKey K) { 70 return RTDyldObjectLinkingLayer::Resources{ 71 std::make_shared<SectionMemoryManager>(), 72 Resolvers[K]}; 73 }), 74 CompileLayer(ObjectLayer, SimpleCompiler(*TM)), 75 OptimizeLayer(CompileLayer, 76 [this](std::unique_ptr<Module> M) { 77 return optimizeModule(std::move(M)); 78 }), 79 CompileCallbackManager(orc::createLocalCompileCallbackManager( 80 TM->getTargetTriple(), ES, 0)), 81 CODLayer(ES, OptimizeLayer, 82 [&](orc::VModuleKey K) { return Resolvers[K]; }, 83 [&](orc::VModuleKey K, std::shared_ptr<SymbolResolver> R) { 84 Resolvers[K] = std::move(R); 85 }, 86 [](Function &F) { return std::set<Function *>({&F}); }, 87 *CompileCallbackManager, 88 orc::createLocalIndirectStubsManagerBuilder( 89 TM->getTargetTriple())) { 90 llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr); 91 } 92 93 TargetMachine &getTargetMachine() { return *TM; } 94 95 VModuleKey addModule(std::unique_ptr<Module> M) { 96 // Create a new VModuleKey. 97 VModuleKey K = ES.allocateVModule(); 98 99 // Build a resolver and associate it with the new key. 100 Resolvers[K] = createLegacyLookupResolver( 101 ES, 102 [this](const std::string &Name) -> JITSymbol { 103 if (auto Sym = CompileLayer.findSymbol(Name, false)) 104 return Sym; 105 else if (auto Err = Sym.takeError()) 106 return std::move(Err); 107 if (auto SymAddr = 108 RTDyldMemoryManager::getSymbolAddressInProcess(Name)) 109 return JITSymbol(SymAddr, JITSymbolFlags::Exported); 110 return nullptr; 111 }, 112 [](Error Err) { cantFail(std::move(Err), "lookupFlags failed"); }); 113 114 // Add the module to the JIT with the new key. 115 cantFail(CODLayer.addModule(K, std::move(M))); 116 return K; 117 } 118 119 JITSymbol findSymbol(const std::string Name) { 120 std::string MangledName; 121 raw_string_ostream MangledNameStream(MangledName); 122 Mangler::getNameWithPrefix(MangledNameStream, Name, DL); 123 return CODLayer.findSymbol(MangledNameStream.str(), true); 124 } 125 126 void removeModule(VModuleKey K) { 127 cantFail(CODLayer.removeModule(K)); 128 } 129 130 private: 131 std::unique_ptr<Module> optimizeModule(std::unique_ptr<Module> M) { 132 // Create a function pass manager. 133 auto FPM = llvm::make_unique<legacy::FunctionPassManager>(M.get()); 134 135 // Add some optimizations. 136 FPM->add(createInstructionCombiningPass()); 137 FPM->add(createReassociatePass()); 138 FPM->add(createGVNPass()); 139 FPM->add(createCFGSimplificationPass()); 140 FPM->doInitialization(); 141 142 // Run the optimizations over all functions in the module being added to 143 // the JIT. 144 for (auto &F : *M) 145 FPM->run(F); 146 147 return M; 148 } 149 }; 150 151 } // end namespace orc 152 } // end namespace llvm 153 154 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H 155
