1de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar//===----- KaleidoscopeJIT.h - A simple JIT for Kaleidoscope ----*- C++ -*-===// 2de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar// 3de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar// The LLVM Compiler Infrastructure 4de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar// 5de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar// This file is distributed under the University of Illinois Open Source 6de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar// License. See LICENSE.TXT for details. 7de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar// 8de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar//===----------------------------------------------------------------------===// 9de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar// 10de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar// Contains a simple JIT definition for use in the kaleidoscope tutorials. 11de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar// 12de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar//===----------------------------------------------------------------------===// 13de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 14de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#ifndef LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H 15de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#define LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H 16de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 17de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/ADT/STLExtras.h" 18de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/ExecutionEngine/ExecutionEngine.h" 19de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/ExecutionEngine/RuntimeDyld.h" 20de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/ExecutionEngine/SectionMemoryManager.h" 21de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/ExecutionEngine/Orc/CompileUtils.h" 22de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/ExecutionEngine/Orc/JITSymbol.h" 23de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" 24de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/ExecutionEngine/Orc/IRTransformLayer.h" 25de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/ExecutionEngine/Orc/LambdaResolver.h" 26de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h" 27de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/IR/DataLayout.h" 28de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/IR/Mangler.h" 29de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/Support/DynamicLibrary.h" 30de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/Support/raw_ostream.h" 31de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include "llvm/Target/TargetMachine.h" 32de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include <algorithm> 33de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include <memory> 34de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include <string> 35de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#include <vector> 36de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 37de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainarnamespace llvm { 38de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainarnamespace orc { 39de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 40de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainarclass KaleidoscopeJIT { 41de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainarprivate: 42de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar std::unique_ptr<TargetMachine> TM; 43de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar const DataLayout DL; 44de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar ObjectLinkingLayer<> ObjectLayer; 45de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar IRCompileLayer<decltype(ObjectLayer)> CompileLayer; 46de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 47de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar typedef std::function<std::unique_ptr<Module>(std::unique_ptr<Module>)> 48de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar OptimizeFunction; 49de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 50de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar IRTransformLayer<decltype(CompileLayer), OptimizeFunction> OptimizeLayer; 51de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 52de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainarpublic: 53de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar typedef decltype(OptimizeLayer)::ModuleSetHandleT ModuleHandle; 54de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 55de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar KaleidoscopeJIT() 56de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar : TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()), 57de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar CompileLayer(ObjectLayer, SimpleCompiler(*TM)), 58de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar OptimizeLayer(CompileLayer, 59de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar [this](std::unique_ptr<Module> M) { 60de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar return optimizeModule(std::move(M)); 61de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar }) { 62de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr); 63de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar } 64de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 65de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar TargetMachine &getTargetMachine() { return *TM; } 66de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 67de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar ModuleHandle addModule(std::unique_ptr<Module> M) { 68de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar // Build our symbol resolver: 69de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar // Lambda 1: Look back into the JIT itself to find symbols that are part of 70de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar // the same "logical dylib". 71de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar // Lambda 2: Search for external symbols in the host process. 72de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar auto Resolver = createLambdaResolver( 73de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar [&](const std::string &Name) { 74de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar if (auto Sym = OptimizeLayer.findSymbol(Name, false)) 75de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar return Sym.toRuntimeDyldSymbol(); 76de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar return RuntimeDyld::SymbolInfo(nullptr); 77de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar }, 78de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar [](const std::string &Name) { 79de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar if (auto SymAddr = 80de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar RTDyldMemoryManager::getSymbolAddressInProcess(Name)) 81de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar return RuntimeDyld::SymbolInfo(SymAddr, JITSymbolFlags::Exported); 82de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar return RuntimeDyld::SymbolInfo(nullptr); 83de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar }); 84de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 85de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar // Build a singlton module set to hold our module. 86de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar std::vector<std::unique_ptr<Module>> Ms; 87de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar Ms.push_back(std::move(M)); 88de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 89de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar // Add the set to the JIT with the resolver we created above and a newly 90de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar // created SectionMemoryManager. 91de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar return OptimizeLayer.addModuleSet(std::move(Ms), 92de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar make_unique<SectionMemoryManager>(), 93de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar std::move(Resolver)); 94de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar } 95de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 96de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar JITSymbol findSymbol(const std::string Name) { 97de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar std::string MangledName; 98de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar raw_string_ostream MangledNameStream(MangledName); 99de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar Mangler::getNameWithPrefix(MangledNameStream, Name, DL); 100de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar return OptimizeLayer.findSymbol(MangledNameStream.str(), true); 101de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar } 102de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 103de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar void removeModule(ModuleHandle H) { 104de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar OptimizeLayer.removeModuleSet(H); 105de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar } 106de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 107de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainarprivate: 108de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 109de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar std::unique_ptr<Module> optimizeModule(std::unique_ptr<Module> M) { 110de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar // Create a function pass manager. 111de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar auto FPM = llvm::make_unique<legacy::FunctionPassManager>(M.get()); 112de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 113de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar // Add some optimizations. 114de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar FPM->add(createInstructionCombiningPass()); 115de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar FPM->add(createReassociatePass()); 116de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar FPM->add(createGVNPass()); 117de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar FPM->add(createCFGSimplificationPass()); 118de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar FPM->doInitialization(); 119de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 120de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar // Run the optimizations over all functions in the module being added to 121de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar // the JIT. 122de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar for (auto &F : *M) 123de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar FPM->run(F); 124de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 125de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar return M; 126de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar } 127de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 128de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar}; 129de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 130de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar} // end namespace orc 131de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar} // end namespace llvm 132de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar 133de2d8694e25a814696358e95141f4b1aa4d8847ePirama Arumuga Nainar#endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H 134