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/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/Shared/ExecutorSymbolDef.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   MaterializationUnit::Interface 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 MaterializationUnit::Interface(std::move(Symbols), nullptr);
 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)), 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<ExecutionSession> ES;
 130   std::unique_ptr<EPCIndirectionUtils> EPCIU;
 131
 132   DataLayout DL;
 133   MangleAndInterner Mangle;
 134
 135   RTDyldObjectLinkingLayer ObjectLayer;
 136   IRCompileLayer CompileLayer;
 137   IRTransformLayer OptimizeLayer;
 138   KaleidoscopeASTLayer ASTLayer;
 139
 140   JITDylib &MainJD;
 141
 142   static void handleLazyCallThroughError() {
 143     errs() << "LazyCallThrough error: Could not find function body";
 144     exit(1);
 145   }
 146
 147 public:
 148   KaleidoscopeJIT(std::unique_ptr<ExecutionSession> ES,
 149                   std::unique_ptr<EPCIndirectionUtils> EPCIU,
 150                   JITTargetMachineBuilder JTMB, DataLayout DL)
 151       : ES(std::move(ES)), EPCIU(std::move(EPCIU)), DL(std::move(DL)),
 152         Mangle(*this->ES, this->DL),
 153         ObjectLayer(*this->ES,
 154                     []() { return std::make_unique<SectionMemoryManager>(); }),
 155         CompileLayer(*this->ES, ObjectLayer,
 156                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
 157         OptimizeLayer(*this->ES, CompileLayer, optimizeModule),
 158         ASTLayer(OptimizeLayer, this->DL),
 159         MainJD(this->ES->createBareJITDylib("<main>")) {
 160     MainJD.addGenerator(
 161         cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
 162             DL.getGlobalPrefix())));
 163   }
 164
 165   ~KaleidoscopeJIT() {
 166     if (auto Err = ES->endSession())
 167       ES->reportError(std::move(Err));
 168     if (auto Err = EPCIU->cleanup())
 169       ES->reportError(std::move(Err));
 170   }
 171
 172   static Expected<std::unique_ptr<KaleidoscopeJIT>> Create() {
 173     auto EPC = SelfExecutorProcessControl::Create();
 174     if (!EPC)
 175       return EPC.takeError();
 176
 177     auto ES = std::make_unique<ExecutionSession>(std::move(*EPC));
 178
 179     auto EPCIU = EPCIndirectionUtils::Create(*ES);
 180     if (!EPCIU)
 181       return EPCIU.takeError();
 182
 183     (*EPCIU)->createLazyCallThroughManager(
 184         *ES, ExecutorAddr::fromPtr(&handleLazyCallThroughError));
 185
 186     if (auto Err = setUpInProcessLCTMReentryViaEPCIU(**EPCIU))
 187       return std::move(Err);
 188
 189     JITTargetMachineBuilder JTMB(
 190         ES->getExecutorProcessControl().getTargetTriple());
 191
 192     auto DL = JTMB.getDefaultDataLayoutForTarget();
 193     if (!DL)
 194       return DL.takeError();
 195
 196     return std::make_unique<KaleidoscopeJIT>(std::move(ES), std::move(*EPCIU),
 197                                              std::move(JTMB), std::move(*DL));
 198   }
 199
 200   const DataLayout &getDataLayout() const { return DL; }
 201
 202   JITDylib &getMainJITDylib() { return MainJD; }
 203
 204   Error addModule(ThreadSafeModule TSM, ResourceTrackerSP RT = nullptr) {
 205     if (!RT)
 206       RT = MainJD.getDefaultResourceTracker();
 207
 208     return OptimizeLayer.add(RT, std::move(TSM));
 209   }
 210
 211   Error addAST(std::unique_ptr<FunctionAST> F, ResourceTrackerSP RT = nullptr) {
 212     if (!RT)
 213       RT = MainJD.getDefaultResourceTracker();
 214     return ASTLayer.add(RT, std::move(F));
 215   }
 216
 217   Expected<ExecutorSymbolDef> lookup(StringRef Name) {
 218     return ES->lookup({&MainJD}, Mangle(Name.str()));
 219   }
 220
 221 private:
 222   static Expected<ThreadSafeModule>
 223   optimizeModule(ThreadSafeModule TSM, const MaterializationResponsibility &R) {
 224     TSM.withModuleDo([](Module &M) {
 225       // Create a function pass manager.
 226       auto FPM = std::make_unique<legacy::FunctionPassManager>(&M);
 227
 228       // Add some optimizations.
 229       FPM->add(createInstructionCombiningPass());
 230       FPM->add(createReassociatePass());
 231       FPM->add(createGVNPass());
 232       FPM->add(createCFGSimplificationPass());
 233       FPM->doInitialization();
 234
 235       // Run the optimizations over all functions in the module being added to
 236       // the JIT.
 237       for (auto &F : M)
 238         FPM->run(F);
 239     });
 240
 241     return std::move(TSM);
 242   }
 243 };
 244
 245 } // end namespace orc
 246 } // end namespace llvm
 247
 248 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H