252 lines · c
1//===- KaleidoscopeJIT.h - A simple JIT for Kaleidoscope --------*- C++ -*-===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// Contains a simple JIT definition for use in the kaleidoscope tutorials.10//11//===----------------------------------------------------------------------===//12 13#ifndef LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H14#define LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H15 16#include "llvm/ADT/StringRef.h"17#include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"18#include "llvm/ExecutionEngine/Orc/CompileUtils.h"19#include "llvm/ExecutionEngine/Orc/Core.h"20#include "llvm/ExecutionEngine/Orc/EPCIndirectionUtils.h"21#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"22#include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"23#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"24#include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"25#include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"26#include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"27#include "llvm/ExecutionEngine/Orc/SelfExecutorProcessControl.h"28#include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"29#include "llvm/ExecutionEngine/SectionMemoryManager.h"30#include "llvm/IR/DataLayout.h"31#include "llvm/IR/LLVMContext.h"32#include "llvm/IR/LegacyPassManager.h"33#include "llvm/Transforms/InstCombine/InstCombine.h"34#include "llvm/Transforms/Scalar.h"35#include "llvm/Transforms/Scalar/GVN.h"36#include <memory>37 38class PrototypeAST;39class ExprAST;40 41/// FunctionAST - This class represents a function definition itself.42class FunctionAST {43 std::unique_ptr<PrototypeAST> Proto;44 std::unique_ptr<ExprAST> Body;45 46public:47 FunctionAST(std::unique_ptr<PrototypeAST> Proto,48 std::unique_ptr<ExprAST> Body)49 : Proto(std::move(Proto)), Body(std::move(Body)) {}50 51 const PrototypeAST& getProto() const;52 const std::string& getName() const;53 llvm::Function *codegen();54};55 56/// This will compile FnAST to IR, rename the function to add the given57/// suffix (needed to prevent a name-clash with the function's stub),58/// and then take ownership of the module that the function was compiled59/// into.60llvm::orc::ThreadSafeModule irgenAndTakeOwnership(FunctionAST &FnAST,61 const std::string &Suffix);62 63namespace llvm {64namespace orc {65 66class KaleidoscopeASTLayer;67class KaleidoscopeJIT;68 69class KaleidoscopeASTMaterializationUnit : public MaterializationUnit {70public:71 KaleidoscopeASTMaterializationUnit(KaleidoscopeASTLayer &L,72 std::unique_ptr<FunctionAST> F);73 74 StringRef getName() const override {75 return "KaleidoscopeASTMaterializationUnit";76 }77 78 void materialize(std::unique_ptr<MaterializationResponsibility> R) override;79 80private:81 void discard(const JITDylib &JD, const SymbolStringPtr &Sym) override {82 llvm_unreachable("Kaleidoscope functions are not overridable");83 }84 85 KaleidoscopeASTLayer &L;86 std::unique_ptr<FunctionAST> F;87};88 89class KaleidoscopeASTLayer {90public:91 KaleidoscopeASTLayer(IRLayer &BaseLayer, const DataLayout &DL)92 : BaseLayer(BaseLayer), DL(DL) {}93 94 Error add(ResourceTrackerSP RT, std::unique_ptr<FunctionAST> F) {95 return RT->getJITDylib().define(96 std::make_unique<KaleidoscopeASTMaterializationUnit>(*this,97 std::move(F)),98 RT);99 }100 101 void emit(std::unique_ptr<MaterializationResponsibility> MR,102 std::unique_ptr<FunctionAST> F) {103 BaseLayer.emit(std::move(MR), irgenAndTakeOwnership(*F, ""));104 }105 106 MaterializationUnit::Interface getInterface(FunctionAST &F) {107 MangleAndInterner Mangle(BaseLayer.getExecutionSession(), DL);108 SymbolFlagsMap Symbols;109 Symbols[Mangle(F.getName())] =110 JITSymbolFlags(JITSymbolFlags::Exported | JITSymbolFlags::Callable);111 return MaterializationUnit::Interface(std::move(Symbols), nullptr);112 }113 114private:115 IRLayer &BaseLayer;116 const DataLayout &DL;117};118 119KaleidoscopeASTMaterializationUnit::KaleidoscopeASTMaterializationUnit(120 KaleidoscopeASTLayer &L, std::unique_ptr<FunctionAST> F)121 : MaterializationUnit(L.getInterface(*F)), L(L), F(std::move(F)) {}122 123void KaleidoscopeASTMaterializationUnit::materialize(124 std::unique_ptr<MaterializationResponsibility> R) {125 L.emit(std::move(R), std::move(F));126}127 128class KaleidoscopeJIT {129private:130 std::unique_ptr<ExecutionSession> ES;131 std::unique_ptr<EPCIndirectionUtils> EPCIU;132 133 DataLayout DL;134 MangleAndInterner Mangle;135 136 RTDyldObjectLinkingLayer ObjectLayer;137 IRCompileLayer CompileLayer;138 IRTransformLayer OptimizeLayer;139 KaleidoscopeASTLayer ASTLayer;140 141 JITDylib &MainJD;142 143 static void handleLazyCallThroughError() {144 errs() << "LazyCallThrough error: Could not find function body";145 exit(1);146 }147 148public:149 KaleidoscopeJIT(std::unique_ptr<ExecutionSession> ES,150 std::unique_ptr<EPCIndirectionUtils> EPCIU,151 JITTargetMachineBuilder JTMB, DataLayout DL)152 : ES(std::move(ES)), EPCIU(std::move(EPCIU)), DL(std::move(DL)),153 Mangle(*this->ES, this->DL),154 ObjectLayer(*this->ES,155 [](const MemoryBuffer &) {156 return std::make_unique<SectionMemoryManager>();157 }),158 CompileLayer(*this->ES, ObjectLayer,159 std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),160 OptimizeLayer(*this->ES, CompileLayer, optimizeModule),161 ASTLayer(OptimizeLayer, this->DL),162 MainJD(this->ES->createBareJITDylib("<main>")) {163 MainJD.addGenerator(164 cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(165 DL.getGlobalPrefix())));166 }167 168 ~KaleidoscopeJIT() {169 if (auto Err = ES->endSession())170 ES->reportError(std::move(Err));171 if (auto Err = EPCIU->cleanup())172 ES->reportError(std::move(Err));173 }174 175 static Expected<std::unique_ptr<KaleidoscopeJIT>> Create() {176 auto EPC = SelfExecutorProcessControl::Create();177 if (!EPC)178 return EPC.takeError();179 180 auto ES = std::make_unique<ExecutionSession>(std::move(*EPC));181 182 auto EPCIU = EPCIndirectionUtils::Create(*ES);183 if (!EPCIU)184 return EPCIU.takeError();185 186 (*EPCIU)->createLazyCallThroughManager(187 *ES, ExecutorAddr::fromPtr(&handleLazyCallThroughError));188 189 if (auto Err = setUpInProcessLCTMReentryViaEPCIU(**EPCIU))190 return std::move(Err);191 192 JITTargetMachineBuilder JTMB(193 ES->getExecutorProcessControl().getTargetTriple());194 195 auto DL = JTMB.getDefaultDataLayoutForTarget();196 if (!DL)197 return DL.takeError();198 199 return std::make_unique<KaleidoscopeJIT>(std::move(ES), std::move(*EPCIU),200 std::move(JTMB), std::move(*DL));201 }202 203 const DataLayout &getDataLayout() const { return DL; }204 205 JITDylib &getMainJITDylib() { return MainJD; }206 207 Error addModule(ThreadSafeModule TSM, ResourceTrackerSP RT = nullptr) {208 if (!RT)209 RT = MainJD.getDefaultResourceTracker();210 211 return OptimizeLayer.add(RT, std::move(TSM));212 }213 214 Error addAST(std::unique_ptr<FunctionAST> F, ResourceTrackerSP RT = nullptr) {215 if (!RT)216 RT = MainJD.getDefaultResourceTracker();217 return ASTLayer.add(RT, std::move(F));218 }219 220 Expected<ExecutorSymbolDef> lookup(StringRef Name) {221 return ES->lookup({&MainJD}, Mangle(Name.str()));222 }223 224private:225 static Expected<ThreadSafeModule>226 optimizeModule(ThreadSafeModule TSM, const MaterializationResponsibility &R) {227 TSM.withModuleDo([](Module &M) {228 // Create a function pass manager.229 auto FPM = std::make_unique<legacy::FunctionPassManager>(&M);230 231 // Add some optimizations.232 FPM->add(createInstructionCombiningPass());233 FPM->add(createReassociatePass());234 FPM->add(createGVNPass());235 FPM->add(createCFGSimplificationPass());236 FPM->doInitialization();237 238 // Run the optimizations over all functions in the module being added to239 // the JIT.240 for (auto &F : M)241 FPM->run(F);242 });243 244 return std::move(TSM);245 }246};247 248} // end namespace orc249} // end namespace llvm250 251#endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H252