1//===----- KaleidoscopeJIT.h - A simple JIT for Kaleidoscope ----*- C++ -*-===//
2//
3//                     The LLVM Compiler Infrastructure
4//
5// This file is distributed under the University of Illinois Open Source
6// License. See LICENSE.TXT for details.
7//
8//===----------------------------------------------------------------------===//
9//
10// Contains a simple JIT definition for use in the kaleidoscope tutorials.
11//
12//===----------------------------------------------------------------------===//
13
14#ifndef LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
15#define LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
16
17#include "llvm/ADT/STLExtras.h"
18#include "llvm/ExecutionEngine/ExecutionEngine.h"
19#include "llvm/ExecutionEngine/RuntimeDyld.h"
20#include "llvm/ExecutionEngine/SectionMemoryManager.h"
21#include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
22#include "llvm/ExecutionEngine/Orc/CompileUtils.h"
23#include "llvm/ExecutionEngine/Orc/JITSymbol.h"
24#include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
25#include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
26#include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
27#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
28#include "llvm/IR/DataLayout.h"
29#include "llvm/IR/Mangler.h"
30#include "llvm/Support/DynamicLibrary.h"
31#include "llvm/Support/raw_ostream.h"
32#include "llvm/Target/TargetMachine.h"
33#include <algorithm>
34#include <memory>
35#include <string>
36#include <vector>
37
38namespace llvm {
39namespace orc {
40
41class KaleidoscopeJIT {
42private:
43  std::unique_ptr<TargetMachine> TM;
44  const DataLayout DL;
45  std::unique_ptr<JITCompileCallbackManager> CompileCallbackManager;
46  ObjectLinkingLayer<> ObjectLayer;
47  IRCompileLayer<decltype(ObjectLayer)> CompileLayer;
48
49  typedef std::function<std::unique_ptr<Module>(std::unique_ptr<Module>)>
50    OptimizeFunction;
51
52  IRTransformLayer<decltype(CompileLayer), OptimizeFunction> OptimizeLayer;
53  CompileOnDemandLayer<decltype(OptimizeLayer)> CODLayer;
54
55public:
56  typedef decltype(CODLayer)::ModuleSetHandleT ModuleHandle;
57
58  KaleidoscopeJIT()
59      : TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
60        CompileCallbackManager(
61            orc::createLocalCompileCallbackManager(TM->getTargetTriple(), 0)),
62        CompileLayer(ObjectLayer, SimpleCompiler(*TM)),
63        OptimizeLayer(CompileLayer,
64                      [this](std::unique_ptr<Module> M) {
65                        return optimizeModule(std::move(M));
66                      }),
67        CODLayer(OptimizeLayer,
68                 [this](Function &F) { return std::set<Function*>({&F}); },
69                 *CompileCallbackManager,
70                 orc::createLocalIndirectStubsManagerBuilder(
71                   TM->getTargetTriple())) {
72    llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
73  }
74
75  TargetMachine &getTargetMachine() { return *TM; }
76
77  ModuleHandle addModule(std::unique_ptr<Module> M) {
78    // Build our symbol resolver:
79    // Lambda 1: Look back into the JIT itself to find symbols that are part of
80    //           the same "logical dylib".
81    // Lambda 2: Search for external symbols in the host process.
82    auto Resolver = createLambdaResolver(
83        [&](const std::string &Name) {
84          if (auto Sym = CODLayer.findSymbol(Name, false))
85            return Sym.toRuntimeDyldSymbol();
86          return RuntimeDyld::SymbolInfo(nullptr);
87        },
88        [](const std::string &Name) {
89          if (auto SymAddr =
90                RTDyldMemoryManager::getSymbolAddressInProcess(Name))
91            return RuntimeDyld::SymbolInfo(SymAddr, JITSymbolFlags::Exported);
92          return RuntimeDyld::SymbolInfo(nullptr);
93        });
94
95    // Build a singlton module set to hold our module.
96    std::vector<std::unique_ptr<Module>> Ms;
97    Ms.push_back(std::move(M));
98
99    // Add the set to the JIT with the resolver we created above and a newly
100    // created SectionMemoryManager.
101    return CODLayer.addModuleSet(std::move(Ms),
102                                 make_unique<SectionMemoryManager>(),
103                                 std::move(Resolver));
104  }
105
106  JITSymbol findSymbol(const std::string Name) {
107    std::string MangledName;
108    raw_string_ostream MangledNameStream(MangledName);
109    Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
110    return CODLayer.findSymbol(MangledNameStream.str(), true);
111  }
112
113  void removeModule(ModuleHandle H) {
114    CODLayer.removeModuleSet(H);
115  }
116
117private:
118
119  std::unique_ptr<Module> optimizeModule(std::unique_ptr<Module> M) {
120    // Create a function pass manager.
121    auto FPM = llvm::make_unique<legacy::FunctionPassManager>(M.get());
122
123    // Add some optimizations.
124    FPM->add(createInstructionCombiningPass());
125    FPM->add(createReassociatePass());
126    FPM->add(createGVNPass());
127    FPM->add(createCFGSimplificationPass());
128    FPM->doInitialization();
129
130    // Run the optimizations over all functions in the module being added to
131    // the JIT.
132    for (auto &F : *M)
133      FPM->run(F);
134
135    return M;
136  }
137
138};
139
140} // end namespace orc
141} // end namespace llvm
142
143#endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
144