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/STLExtras.h"
  17 #include "llvm/ExecutionEngine/ExecutionEngine.h"
  18 #include "llvm/ExecutionEngine/JITSymbol.h"
  19 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
  20 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
  21 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
  22 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
  23 #include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
  24 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
  25 #include "llvm/ExecutionEngine/RTDyldMemoryManager.h"
  26 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
  27 #include "llvm/IR/DataLayout.h"
  28 #include "llvm/IR/LegacyPassManager.h"
  29 #include "llvm/IR/Mangler.h"
  30 #include "llvm/Support/DynamicLibrary.h"
  31 #include "llvm/Support/Error.h"
  32 #include "llvm/Support/raw_ostream.h"
  33 #include "llvm/Target/TargetMachine.h"
  34 #include "llvm/Transforms/InstCombine/InstCombine.h"
  35 #include "llvm/Transforms/Scalar.h"
  36 #include "llvm/Transforms/Scalar/GVN.h"
  37 #include <algorithm>
  38 #include <cassert>
  39 #include <cstdlib>
  40 #include <map>
  41 #include <memory>
  42 #include <string>
  43 #include <vector>
  44
  45 class PrototypeAST;
  46 class ExprAST;
  47
  48 /// FunctionAST - This class represents a function definition itself.
  49 class FunctionAST {
  50   std::unique_ptr<PrototypeAST> Proto;
  51   std::unique_ptr<ExprAST> Body;
  52
  53 public:
  54   FunctionAST(std::unique_ptr<PrototypeAST> Proto,
  55               std::unique_ptr<ExprAST> Body)
  56       : Proto(std::move(Proto)), Body(std::move(Body)) {}
  57
  58   const PrototypeAST& getProto() const;
  59   const std::string& getName() const;
  60   llvm::Function *codegen();
  61 };
  62
  63 /// This will compile FnAST to IR, rename the function to add the given
  64 /// suffix (needed to prevent a name-clash with the function's stub),
  65 /// and then take ownership of the module that the function was compiled
  66 /// into.
  67 std::unique_ptr<llvm::Module>
  68 irgenAndTakeOwnership(FunctionAST &FnAST, const std::string &Suffix);
  69
  70 namespace llvm {
  71 namespace orc {
  72
  73 class KaleidoscopeJIT {
  74 private:
  75   ExecutionSession ES;
  76   std::shared_ptr<SymbolResolver> Resolver;
  77   std::unique_ptr<TargetMachine> TM;
  78   const DataLayout DL;
  79   LegacyRTDyldObjectLinkingLayer ObjectLayer;
  80   LegacyIRCompileLayer<decltype(ObjectLayer), SimpleCompiler> CompileLayer;
  81
  82   using OptimizeFunction =
  83       std::function<std::unique_ptr<Module>(std::unique_ptr<Module>)>;
  84
  85   LegacyIRTransformLayer<decltype(CompileLayer), OptimizeFunction> OptimizeLayer;
  86
  87   std::unique_ptr<JITCompileCallbackManager> CompileCallbackMgr;
  88   std::unique_ptr<IndirectStubsManager> IndirectStubsMgr;
  89
  90 public:
  91   KaleidoscopeJIT()
  92       : Resolver(createLegacyLookupResolver(
  93             ES,
  94             [this](const std::string &Name) -> JITSymbol {
  95               if (auto Sym = IndirectStubsMgr->findStub(Name, false))
  96                 return Sym;
  97               if (auto Sym = OptimizeLayer.findSymbol(Name, false))
  98                 return Sym;
  99               else if (auto Err = Sym.takeError())
 100                 return std::move(Err);
 101               if (auto SymAddr =
 102                       RTDyldMemoryManager::getSymbolAddressInProcess(Name))
 103                 return JITSymbol(SymAddr, JITSymbolFlags::Exported);
 104               return nullptr;
 105             },
 106             [](Error Err) { cantFail(std::move(Err), "lookupFlags failed"); })),
 107         TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
 108         ObjectLayer(AcknowledgeORCv1Deprecation, ES,
 109                     [this](VModuleKey K) {
 110                       return LegacyRTDyldObjectLinkingLayer::Resources{
 111                           std::make_shared<SectionMemoryManager>(), Resolver};
 112                     }),
 113         CompileLayer(AcknowledgeORCv1Deprecation, ObjectLayer,
 114                      SimpleCompiler(*TM)),
 115         OptimizeLayer(AcknowledgeORCv1Deprecation, CompileLayer,
 116                       [this](std::unique_ptr<Module> M) {
 117                         return optimizeModule(std::move(M));
 118                       }),
 119         CompileCallbackMgr(cantFail(orc::createLocalCompileCallbackManager(
 120             TM->getTargetTriple(), ES, 0))) {
 121     auto IndirectStubsMgrBuilder =
 122       orc::createLocalIndirectStubsManagerBuilder(TM->getTargetTriple());
 123     IndirectStubsMgr = IndirectStubsMgrBuilder();
 124     llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
 125   }
 126
 127   TargetMachine &getTargetMachine() { return *TM; }
 128
 129   VModuleKey addModule(std::unique_ptr<Module> M) {
 130     // Add the module to the JIT with a new VModuleKey.
 131     auto K = ES.allocateVModule();
 132     cantFail(OptimizeLayer.addModule(K, std::move(M)));
 133     return K;
 134   }
 135
 136   Error addFunctionAST(std::unique_ptr<FunctionAST> FnAST) {
 137     // Move ownership of FnAST to a shared pointer - C++11 lambdas don't support
 138     // capture-by-move, which is be required for unique_ptr.
 139     auto SharedFnAST = std::shared_ptr<FunctionAST>(std::move(FnAST));
 140
 141     // Set the action to compile our AST. This lambda will be run if/when
 142     // execution hits the compile callback (via the stub).
 143     //
 144     // The steps to compile are:
 145     // (1) IRGen the function.
 146     // (2) Add the IR module to the JIT to make it executable like any other
 147     //     module.
 148     // (3) Use findSymbol to get the address of the compiled function.
 149     // (4) Update the stub pointer to point at the implementation so that
 150     ///    subsequent calls go directly to it and bypass the compiler.
 151     // (5) Return the address of the implementation: this lambda will actually
 152     //     be run inside an attempted call to the function, and we need to
 153     //     continue on to the implementation to complete the attempted call.
 154     //     The JIT runtime (the resolver block) will use the return address of
 155     //     this function as the address to continue at once it has reset the
 156     //     CPU state to what it was immediately before the call.
 157     auto CompileAction = [this, SharedFnAST]() {
 158       auto M = irgenAndTakeOwnership(*SharedFnAST, "$impl");
 159       addModule(std::move(M));
 160       auto Sym = findSymbol(SharedFnAST->getName() + "$impl");
 161       assert(Sym && "Couldn't find compiled function?");
 162       JITTargetAddress SymAddr = cantFail(Sym.getAddress());
 163       if (auto Err = IndirectStubsMgr->updatePointer(
 164               mangle(SharedFnAST->getName()), SymAddr)) {
 165         logAllUnhandledErrors(std::move(Err), errs(),
 166                               "Error updating function pointer: ");
 167         exit(1);
 168       }
 169
 170       return SymAddr;
 171     };
 172
 173     // Create a CompileCallback using the CompileAction - this is the re-entry
 174     // point into the compiler for functions that haven't been compiled yet.
 175     auto CCAddr = cantFail(
 176         CompileCallbackMgr->getCompileCallback(std::move(CompileAction)));
 177
 178     // Create an indirect stub. This serves as the functions "canonical
 179     // definition" - an unchanging (constant address) entry point to the
 180     // function implementation.
 181     // Initially we point the stub's function-pointer at the compile callback
 182     // that we just created. When the compile action for the callback is run we
 183     // will update the stub's function pointer to point at the function
 184     // implementation that we just implemented.
 185     if (auto Err = IndirectStubsMgr->createStub(
 186             mangle(SharedFnAST->getName()), CCAddr, JITSymbolFlags::Exported))
 187       return Err;
 188
 189     return Error::success();
 190   }
 191
 192   JITSymbol findSymbol(const std::string Name) {
 193     return OptimizeLayer.findSymbol(mangle(Name), true);
 194   }
 195
 196   void removeModule(VModuleKey K) {
 197     cantFail(OptimizeLayer.removeModule(K));
 198   }
 199
 200 private:
 201   std::string mangle(const std::string &Name) {
 202     std::string MangledName;
 203     raw_string_ostream MangledNameStream(MangledName);
 204     Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
 205     return MangledNameStream.str();
 206   }
 207
 208   std::unique_ptr<Module> optimizeModule(std::unique_ptr<Module> M) {
 209     // Create a function pass manager.
 210     auto FPM = std::make_unique<legacy::FunctionPassManager>(M.get());
 211
 212     // Add some optimizations.
 213     FPM->add(createInstructionCombiningPass());
 214     FPM->add(createReassociatePass());
 215     FPM->add(createGVNPass());
 216     FPM->add(createCFGSimplificationPass());
 217     FPM->doInitialization();
 218
 219     // Run the optimizations over all functions in the module being added to
 220     // the JIT.
 221     for (auto &F : *M)
 222       FPM->run(F);
 223
 224     return M;
 225   }
 226 };
 227
 228 } // end namespace orc
 229 } // end namespace llvm
 230
 231 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H