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