Home | History | Annotate | Line # | Download | only in Chapter4 
      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-exception
      6 //
      7 //===----------------------------------------------------------------------===//
      8 //
      9 // Contains a simple JIT definition for use in the kaleidoscope tutorials.
     10 //
     11 //===----------------------------------------------------------------------===//
     12 
     13 #ifndef LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
     14 #define LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
     15 
     16 #include "llvm/ADT/StringRef.h"
     17 #include "llvm/ExecutionEngine/JITSymbol.h"
     18 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
     19 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
     20 #include "llvm/ExecutionEngine/Orc/Core.h"
     21 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
     22 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
     23 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
     24 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
     25 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
     26 #include "llvm/ExecutionEngine/Orc/TPCIndirectionUtils.h"
     27 #include "llvm/ExecutionEngine/Orc/TargetProcessControl.h"
     28 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
     29 #include "llvm/IR/DataLayout.h"
     30 #include "llvm/IR/LLVMContext.h"
     31 #include "llvm/IR/LegacyPassManager.h"
     32 #include "llvm/Transforms/InstCombine/InstCombine.h"
     33 #include "llvm/Transforms/Scalar.h"
     34 #include "llvm/Transforms/Scalar/GVN.h"
     35 #include <memory>
     36 
     37 class PrototypeAST;
     38 class ExprAST;
     39 
     40 /// FunctionAST - This class represents a function definition itself.
     41 class FunctionAST {
     42   std::unique_ptr<PrototypeAST> Proto;
     43   std::unique_ptr<ExprAST> Body;
     44 
     45 public:
     46   FunctionAST(std::unique_ptr<PrototypeAST> Proto,
     47               std::unique_ptr<ExprAST> Body)
     48       : Proto(std::move(Proto)), Body(std::move(Body)) {}
     49 
     50   const PrototypeAST& getProto() const;
     51   const std::string& getName() const;
     52   llvm::Function *codegen();
     53 };
     54 
     55 /// This will compile FnAST to IR, rename the function to add the given
     56 /// suffix (needed to prevent a name-clash with the function's stub),
     57 /// and then take ownership of the module that the function was compiled
     58 /// into.
     59 llvm::orc::ThreadSafeModule irgenAndTakeOwnership(FunctionAST &FnAST,
     60                                                   const std::string &Suffix);
     61 
     62 namespace llvm {
     63 namespace orc {
     64 
     65 class KaleidoscopeASTLayer;
     66 class KaleidoscopeJIT;
     67 
     68 class KaleidoscopeASTMaterializationUnit : public MaterializationUnit {
     69 public:
     70   KaleidoscopeASTMaterializationUnit(KaleidoscopeASTLayer &L,
     71                                      std::unique_ptr<FunctionAST> F);
     72 
     73   StringRef getName() const override {
     74     return "KaleidoscopeASTMaterializationUnit";
     75   }
     76 
     77   void materialize(std::unique_ptr<MaterializationResponsibility> R) override;
     78 
     79 private:
     80   void discard(const JITDylib &JD, const SymbolStringPtr &Sym) override {
     81     llvm_unreachable("Kaleidoscope functions are not overridable");
     82   }
     83 
     84   KaleidoscopeASTLayer &L;
     85   std::unique_ptr<FunctionAST> F;
     86 };
     87 
     88 class KaleidoscopeASTLayer {
     89 public:
     90   KaleidoscopeASTLayer(IRLayer &BaseLayer, const DataLayout &DL)
     91       : BaseLayer(BaseLayer), DL(DL) {}
     92 
     93   Error add(ResourceTrackerSP RT, std::unique_ptr<FunctionAST> F) {
     94     return RT->getJITDylib().define(
     95         std::make_unique<KaleidoscopeASTMaterializationUnit>(*this,
     96                                                              std::move(F)),
     97         RT);
     98   }
     99 
    100   void emit(std::unique_ptr<MaterializationResponsibility> MR,
    101             std::unique_ptr<FunctionAST> F) {
    102     BaseLayer.emit(std::move(MR), irgenAndTakeOwnership(*F, ""));
    103   }
    104 
    105   SymbolFlagsMap getInterface(FunctionAST &F) {
    106     MangleAndInterner Mangle(BaseLayer.getExecutionSession(), DL);
    107     SymbolFlagsMap Symbols;
    108     Symbols[Mangle(F.getName())] =
    109         JITSymbolFlags(JITSymbolFlags::Exported | JITSymbolFlags::Callable);
    110     return Symbols;
    111   }
    112 
    113 private:
    114   IRLayer &BaseLayer;
    115   const DataLayout &DL;
    116 };
    117 
    118 KaleidoscopeASTMaterializationUnit::KaleidoscopeASTMaterializationUnit(
    119     KaleidoscopeASTLayer &L, std::unique_ptr<FunctionAST> F)
    120     : MaterializationUnit(L.getInterface(*F), nullptr), L(L), F(std::move(F)) {}
    121 
    122 void KaleidoscopeASTMaterializationUnit::materialize(
    123     std::unique_ptr<MaterializationResponsibility> R) {
    124   L.emit(std::move(R), std::move(F));
    125 }
    126 
    127 class KaleidoscopeJIT {
    128 private:
    129   std::unique_ptr<TargetProcessControl> TPC;
    130   std::unique_ptr<ExecutionSession> ES;
    131   std::unique_ptr<TPCIndirectionUtils> TPCIU;
    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 
    148 public:
    149   KaleidoscopeJIT(std::unique_ptr<TargetProcessControl> TPC,
    150                   std::unique_ptr<ExecutionSession> ES,
    151                   std::unique_ptr<TPCIndirectionUtils> TPCIU,
    152                   JITTargetMachineBuilder JTMB, DataLayout DL)
    153       : TPC(std::move(TPC)), ES(std::move(ES)), TPCIU(std::move(TPCIU)),
    154         DL(std::move(DL)), Mangle(*this->ES, this->DL),
    155         ObjectLayer(*this->ES,
    156                     []() { return std::make_unique<SectionMemoryManager>(); }),
    157         CompileLayer(*this->ES, ObjectLayer,
    158                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
    159         OptimizeLayer(*this->ES, CompileLayer, optimizeModule),
    160         ASTLayer(OptimizeLayer, this->DL),
    161         MainJD(this->ES->createBareJITDylib("<main>")) {
    162     MainJD.addGenerator(
    163         cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
    164             DL.getGlobalPrefix())));
    165   }
    166 
    167   ~KaleidoscopeJIT() {
    168     if (auto Err = ES->endSession())
    169       ES->reportError(std::move(Err));
    170     if (auto Err = TPCIU->cleanup())
    171       ES->reportError(std::move(Err));
    172   }
    173 
    174   static Expected<std::unique_ptr<KaleidoscopeJIT>> Create() {
    175     auto SSP = std::make_shared<SymbolStringPool>();
    176     auto TPC = SelfTargetProcessControl::Create(SSP);
    177     if (!TPC)
    178       return TPC.takeError();
    179 
    180     auto ES = std::make_unique<ExecutionSession>(std::move(SSP));
    181 
    182     auto TPCIU = TPCIndirectionUtils::Create(**TPC);
    183     if (!TPCIU)
    184       return TPCIU.takeError();
    185 
    186     (*TPCIU)->createLazyCallThroughManager(
    187         *ES, pointerToJITTargetAddress(&handleLazyCallThroughError));
    188 
    189     if (auto Err = setUpInProcessLCTMReentryViaTPCIU(**TPCIU))
    190       return std::move(Err);
    191 
    192     JITTargetMachineBuilder JTMB((*TPC)->getTargetTriple());
    193 
    194     auto DL = JTMB.getDefaultDataLayoutForTarget();
    195     if (!DL)
    196       return DL.takeError();
    197 
    198     return std::make_unique<KaleidoscopeJIT>(std::move(*TPC), std::move(ES),
    199                                              std::move(*TPCIU), std::move(JTMB),
    200                                              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<JITEvaluatedSymbol> lookup(StringRef Name) {
    221     return ES->lookup({&MainJD}, Mangle(Name.str()));
    222   }
    223 
    224 private:
    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 to
    239       // the JIT.
    240       for (auto &F : M)
    241         FPM->run(F);
    242     });
    243 
    244     return std::move(TSM);
    245   }
    246 };
    247 
    248 } // end namespace orc
    249 } // end namespace llvm
    250 
    251 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
    252