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/JITSymbol.h"
20 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
21 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
22 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
23 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
24 #include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
25 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
26 #include "llvm/ExecutionEngine/RTDyldMemoryManager.h"
27 #include "llvm/ExecutionEngine/RuntimeDyld.h"
28 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
29 #include "llvm/IR/DataLayout.h"
30 #include "llvm/IR/LegacyPassManager.h"
31 #include "llvm/IR/Mangler.h"
32 #include "llvm/Support/DynamicLibrary.h"
33 #include "llvm/Support/raw_ostream.h"
34 #include "llvm/Target/TargetMachine.h"
35 #include "llvm/Transforms/InstCombine/InstCombine.h"
36 #include "llvm/Transforms/Scalar.h"
37 #include "llvm/Transforms/Scalar/GVN.h"
38 #include <algorithm>
39 #include <map>
40 #include <memory>
41 #include <set>
42 #include <string>
43 #include <vector>
44 
45 namespace llvm {
46 namespace orc {
47 
48 class KaleidoscopeJIT {
49 private:
50   ExecutionSession ES;
51   std::map<VModuleKey, std::shared_ptr<SymbolResolver>> Resolvers;
52   std::unique_ptr<TargetMachine> TM;
53   const DataLayout DL;
54   RTDyldObjectLinkingLayer ObjectLayer;
55   IRCompileLayer<decltype(ObjectLayer), SimpleCompiler> CompileLayer;
56 
57   using OptimizeFunction =
58       std::function<std::unique_ptr<Module>(std::unique_ptr<Module>)>;
59 
60   IRTransformLayer<decltype(CompileLayer), OptimizeFunction> OptimizeLayer;
61 
62   std::unique_ptr<JITCompileCallbackManager> CompileCallbackManager;
63   CompileOnDemandLayer<decltype(OptimizeLayer)> CODLayer;
64 
65 public:
KaleidoscopeJIT()66   KaleidoscopeJIT()
67       : TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
68         ObjectLayer(ES,
69                     [this](VModuleKey K) {
70                       return RTDyldObjectLinkingLayer::Resources{
71                           std::make_shared<SectionMemoryManager>(),
72                           Resolvers[K]};
73                     }),
74         CompileLayer(ObjectLayer, SimpleCompiler(*TM)),
75         OptimizeLayer(CompileLayer,
76                       [this](std::unique_ptr<Module> M) {
77                         return optimizeModule(std::move(M));
78                       }),
79         CompileCallbackManager(orc::createLocalCompileCallbackManager(
80             TM->getTargetTriple(), ES, 0)),
81         CODLayer(ES, OptimizeLayer,
82                  [&](orc::VModuleKey K) { return Resolvers[K]; },
83                  [&](orc::VModuleKey K, std::shared_ptr<SymbolResolver> R) {
84                    Resolvers[K] = std::move(R);
85                  },
86                  [](Function &F) { return std::set<Function *>({&F}); },
87                  *CompileCallbackManager,
88                  orc::createLocalIndirectStubsManagerBuilder(
89                      TM->getTargetTriple())) {
90     llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
91   }
92 
getTargetMachine()93   TargetMachine &getTargetMachine() { return *TM; }
94 
addModule(std::unique_ptr<Module> M)95   VModuleKey addModule(std::unique_ptr<Module> M) {
96     // Create a new VModuleKey.
97     VModuleKey K = ES.allocateVModule();
98 
99     // Build a resolver and associate it with the new key.
100     Resolvers[K] = createLegacyLookupResolver(
101         ES,
102         [this](const std::string &Name) -> JITSymbol {
103           if (auto Sym = CompileLayer.findSymbol(Name, false))
104             return Sym;
105           else if (auto Err = Sym.takeError())
106             return std::move(Err);
107           if (auto SymAddr =
108                   RTDyldMemoryManager::getSymbolAddressInProcess(Name))
109             return JITSymbol(SymAddr, JITSymbolFlags::Exported);
110           return nullptr;
111         },
112         [](Error Err) { cantFail(std::move(Err), "lookupFlags failed"); });
113 
114     // Add the module to the JIT with the new key.
115     cantFail(CODLayer.addModule(K, std::move(M)));
116     return K;
117   }
118 
findSymbol(const std::string Name)119   JITSymbol findSymbol(const std::string Name) {
120     std::string MangledName;
121     raw_string_ostream MangledNameStream(MangledName);
122     Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
123     return CODLayer.findSymbol(MangledNameStream.str(), true);
124   }
125 
removeModule(VModuleKey K)126   void removeModule(VModuleKey K) {
127     cantFail(CODLayer.removeModule(K));
128   }
129 
130 private:
optimizeModule(std::unique_ptr<Module> M)131   std::unique_ptr<Module> optimizeModule(std::unique_ptr<Module> M) {
132     // Create a function pass manager.
133     auto FPM = llvm::make_unique<legacy::FunctionPassManager>(M.get());
134 
135     // Add some optimizations.
136     FPM->add(createInstructionCombiningPass());
137     FPM->add(createReassociatePass());
138     FPM->add(createGVNPass());
139     FPM->add(createCFGSimplificationPass());
140     FPM->doInitialization();
141 
142     // Run the optimizations over all functions in the module being added to
143     // the JIT.
144     for (auto &F : *M)
145       FPM->run(F);
146 
147     return M;
148   }
149 };
150 
151 } // end namespace orc
152 } // end namespace llvm
153 
154 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
155