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