llvm/examples/Kaleidoscope/BuildingAJIT/Chapter4/KaleidoscopeJIT.h

   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 TPC = SelfTargetProcessControl::Create();
 176     if (!TPC)
 177       return TPC.takeError();
 178
 179     auto ES = std::make_unique<ExecutionSession>();
 180
 181     auto TPCIU = TPCIndirectionUtils::Create(**TPC);
 182     if (!TPCIU)
 183       return TPCIU.takeError();
 184
 185     (*TPCIU)->createLazyCallThroughManager(
 186         *ES, pointerToJITTargetAddress(&handleLazyCallThroughError));
 187
 188     if (auto Err = setUpInProcessLCTMReentryViaTPCIU(**TPCIU))
 189       return std::move(Err);
 190
 191     JITTargetMachineBuilder JTMB((*TPC)->getTargetTriple());
 192
 193     auto DL = JTMB.getDefaultDataLayoutForTarget();
 194     if (!DL)
 195       return DL.takeError();
 196
 197     return std::make_unique<KaleidoscopeJIT>(std::move(*TPC), std::move(ES),
 198                                              std::move(*TPCIU), std::move(JTMB),
 199                                              std::move(*DL));
 200   }
 201
 202   const DataLayout &getDataLayout() const { return DL; }
 203
 204   JITDylib &getMainJITDylib() { return MainJD; }
 205
 206   Error addModule(ThreadSafeModule TSM, ResourceTrackerSP RT = nullptr) {
 207     if (!RT)
 208       RT = MainJD.getDefaultResourceTracker();
 209
 210     return OptimizeLayer.add(RT, std::move(TSM));
 211   }
 212
 213   Error addAST(std::unique_ptr<FunctionAST> F, ResourceTrackerSP RT = nullptr) {
 214     if (!RT)
 215       RT = MainJD.getDefaultResourceTracker();
 216     return ASTLayer.add(RT, std::move(F));
 217   }
 218
 219   Expected<JITEvaluatedSymbol> lookup(StringRef Name) {
 220     return ES->lookup({&MainJD}, Mangle(Name.str()));
 221   }
 222
 223 private:
 224   static Expected<ThreadSafeModule>
 225   optimizeModule(ThreadSafeModule TSM, const MaterializationResponsibility &R) {
 226     TSM.withModuleDo([](Module &M) {
 227       // Create a function pass manager.
 228       auto FPM = std::make_unique<legacy::FunctionPassManager>(&M);
 229
 230       // Add some optimizations.
 231       FPM->add(createInstructionCombiningPass());
 232       FPM->add(createReassociatePass());
 233       FPM->add(createGVNPass());
 234       FPM->add(createCFGSimplificationPass());
 235       FPM->doInitialization();
 236
 237       // Run the optimizations over all functions in the module being added to
 238       // the JIT.
 239       for (auto &F : M)
 240         FPM->run(F);
 241     });
 242
 243     return std::move(TSM);
 244   }
 245 };
 246
 247 } // end namespace orc
 248 } // end namespace llvm
 249
 250 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H