llvm/examples/Kaleidoscope/BuildingAJIT/Chapter2/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/CompileUtils.h"
  19 #include "llvm/ExecutionEngine/Orc/Core.h"
  20 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
  21 #include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.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/SectionMemoryManager.h"
  27 #include "llvm/IR/DataLayout.h"
  28 #include "llvm/IR/LLVMContext.h"
  29 #include "llvm/IR/LegacyPassManager.h"
  30 #include "llvm/Transforms/InstCombine/InstCombine.h"
  31 #include "llvm/Transforms/Scalar.h"
  32 #include "llvm/Transforms/Scalar/GVN.h"
  33 #include <memory>
  34
  35 namespace llvm {
  36 namespace orc {
  37
  38 class KaleidoscopeJIT {
  39 private:
  40   std::unique_ptr<ExecutionSession> ES;
  41
  42   DataLayout DL;
  43   MangleAndInterner Mangle;
  44
  45   RTDyldObjectLinkingLayer ObjectLayer;
  46   IRCompileLayer CompileLayer;
  47   IRTransformLayer OptimizeLayer;
  48
  49   JITDylib &MainJD;
  50
  51 public:
  52   KaleidoscopeJIT(std::unique_ptr<ExecutionSession> ES,
  53                   JITTargetMachineBuilder JTMB, DataLayout DL)
  54       : ES(std::move(ES)), DL(std::move(DL)), Mangle(*this->ES, this->DL),
  55         ObjectLayer(*this->ES,
  56                     []() { return std::make_unique<SectionMemoryManager>(); }),
  57         CompileLayer(*this->ES, ObjectLayer,
  58                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
  59         OptimizeLayer(*this->ES, CompileLayer, optimizeModule),
  60         MainJD(this->ES->createBareJITDylib("<main>")) {
  61     MainJD.addGenerator(
  62         cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
  63             DL.getGlobalPrefix())));
  64   }
  65
  66   ~KaleidoscopeJIT() {
  67     if (auto Err = ES->endSession())
  68       ES->reportError(std::move(Err));
  69   }
  70
  71   static Expected<std::unique_ptr<KaleidoscopeJIT>> Create() {
  72     auto EPC = SelfExecutorProcessControl::Create();
  73     if (!EPC)
  74       return EPC.takeError();
  75
  76     auto ES = std::make_unique<ExecutionSession>(std::move(*EPC));
  77
  78     JITTargetMachineBuilder JTMB(
  79         ES->getExecutorProcessControl().getTargetTriple());
  80
  81     auto DL = JTMB.getDefaultDataLayoutForTarget();
  82     if (!DL)
  83       return DL.takeError();
  84
  85     return std::make_unique<KaleidoscopeJIT>(std::move(ES), std::move(JTMB),
  86                                              std::move(*DL));
  87   }
  88
  89   const DataLayout &getDataLayout() const { return DL; }
  90
  91   JITDylib &getMainJITDylib() { return MainJD; }
  92
  93   Error addModule(ThreadSafeModule TSM, ResourceTrackerSP RT = nullptr) {
  94     if (!RT)
  95       RT = MainJD.getDefaultResourceTracker();
  96
  97     return OptimizeLayer.add(RT, std::move(TSM));
  98   }
  99
 100   Expected<JITEvaluatedSymbol> lookup(StringRef Name) {
 101     return ES->lookup({&MainJD}, Mangle(Name.str()));
 102   }
 103
 104 private:
 105   static Expected<ThreadSafeModule>
 106   optimizeModule(ThreadSafeModule TSM, const MaterializationResponsibility &R) {
 107     TSM.withModuleDo([](Module &M) {
 108       // Create a function pass manager.
 109       auto FPM = std::make_unique<legacy::FunctionPassManager>(&M);
 110
 111       // Add some optimizations.
 112       FPM->add(createInstructionCombiningPass());
 113       FPM->add(createReassociatePass());
 114       FPM->add(createGVNPass());
 115       FPM->add(createCFGSimplificationPass());
 116       FPM->doInitialization();
 117
 118       // Run the optimizations over all functions in the module being added to
 119       // the JIT.
 120       for (auto &F : M)
 121         FPM->run(F);
 122     });
 123
 124     return std::move(TSM);
 125   }
 126 };
 127
 128 } // end namespace orc
 129 } // end namespace llvm
 130
 131 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H