• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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/CompileUtils.h"
21 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
22 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.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/raw_ostream.h"
32 #include "llvm/Target/TargetMachine.h"
33 #include "llvm/Transforms/InstCombine/InstCombine.h"
34 #include "llvm/Transforms/Scalar.h"
35 #include "llvm/Transforms/Scalar/GVN.h"
36 #include <algorithm>
37 #include <memory>
38 #include <string>
39 #include <vector>
40 
41 namespace llvm {
42 namespace orc {
43 
44 class KaleidoscopeJIT {
45 private:
46   ExecutionSession ES;
47   std::shared_ptr<SymbolResolver> Resolver;
48   std::unique_ptr<TargetMachine> TM;
49   const DataLayout DL;
50   RTDyldObjectLinkingLayer ObjectLayer;
51   IRCompileLayer<decltype(ObjectLayer), SimpleCompiler> CompileLayer;
52 
53   using OptimizeFunction =
54       std::function<std::unique_ptr<Module>(std::unique_ptr<Module>)>;
55 
56   IRTransformLayer<decltype(CompileLayer), OptimizeFunction> OptimizeLayer;
57 
58 public:
KaleidoscopeJIT()59   KaleidoscopeJIT()
60       : Resolver(createLegacyLookupResolver(
61             ES,
62             [this](const std::string &Name) -> JITSymbol {
63               if (auto Sym = OptimizeLayer.findSymbol(Name, false))
64                 return Sym;
65               else if (auto Err = Sym.takeError())
66                 return std::move(Err);
67               if (auto SymAddr =
68                       RTDyldMemoryManager::getSymbolAddressInProcess(Name))
69                 return JITSymbol(SymAddr, JITSymbolFlags::Exported);
70               return nullptr;
71             },
72             [](Error Err) { cantFail(std::move(Err), "lookupFlags failed"); })),
73         TM(EngineBuilder().selectTarget()), DL(TM->createDataLayout()),
74         ObjectLayer(ES,
75                     [this](VModuleKey) {
76                       return RTDyldObjectLinkingLayer::Resources{
77                           std::make_shared<SectionMemoryManager>(), Resolver};
78                     }),
79         CompileLayer(ObjectLayer, SimpleCompiler(*TM)),
80         OptimizeLayer(CompileLayer, [this](std::unique_ptr<Module> M) {
81           return optimizeModule(std::move(M));
82         }) {
83     llvm::sys::DynamicLibrary::LoadLibraryPermanently(nullptr);
84   }
85 
getTargetMachine()86   TargetMachine &getTargetMachine() { return *TM; }
87 
addModule(std::unique_ptr<Module> M)88   VModuleKey addModule(std::unique_ptr<Module> M) {
89     // Add the module to the JIT with a new VModuleKey.
90     auto K = ES.allocateVModule();
91     cantFail(OptimizeLayer.addModule(K, std::move(M)));
92     return K;
93   }
94 
findSymbol(const std::string Name)95   JITSymbol findSymbol(const std::string Name) {
96     std::string MangledName;
97     raw_string_ostream MangledNameStream(MangledName);
98     Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
99     return OptimizeLayer.findSymbol(MangledNameStream.str(), true);
100   }
101 
removeModule(VModuleKey K)102   void removeModule(VModuleKey K) {
103     cantFail(OptimizeLayer.removeModule(K));
104   }
105 
106 private:
optimizeModule(std::unique_ptr<Module> M)107   std::unique_ptr<Module> optimizeModule(std::unique_ptr<Module> M) {
108     // Create a function pass manager.
109     auto FPM = llvm::make_unique<legacy::FunctionPassManager>(M.get());
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     return M;
124   }
125 };
126 
127 } // end namespace orc
128 } // end namespace llvm
129 
130 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
131