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/StringRef.h"
17 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
18 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
19 #include "llvm/ExecutionEngine/Orc/Core.h"
20 #include "llvm/ExecutionEngine/Orc/EPCIndirectionUtils.h"
21 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
22 #include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
23 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
24 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
25 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
26 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
27 #include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.h"
28 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
29 #include "llvm/IR/DataLayout.h"
30 #include "llvm/IR/LLVMContext.h"
31 #include "llvm/IR/LegacyPassManager.h"
32 #include "llvm/Transforms/InstCombine/InstCombine.h"
33 #include "llvm/Transforms/Scalar.h"
34 #include "llvm/Transforms/Scalar/GVN.h"
35 #include <memory>
36
37 class PrototypeAST;
38 class ExprAST;
39
40 /// FunctionAST - This class represents a function definition itself.
41 class FunctionAST {
42 std::unique_ptr<PrototypeAST> Proto;
43 std::unique_ptr<ExprAST> Body;
44
45 public:
FunctionAST(std::unique_ptr<PrototypeAST> Proto,std::unique_ptr<ExprAST> Body)46 FunctionAST(std::unique_ptr<PrototypeAST> Proto,
47 std::unique_ptr<ExprAST> Body)
48 : Proto(std::move(Proto)), Body(std::move(Body)) {}
49
50 const PrototypeAST& getProto() const;
51 const std::string& getName() const;
52 llvm::Function *codegen();
53 };
54
55 /// This will compile FnAST to IR, rename the function to add the given
56 /// suffix (needed to prevent a name-clash with the function's stub),
57 /// and then take ownership of the module that the function was compiled
58 /// into.
59 llvm::orc::ThreadSafeModule irgenAndTakeOwnership(FunctionAST &FnAST,
60 const std::string &Suffix);
61
62 namespace llvm {
63 namespace orc {
64
65 class KaleidoscopeASTLayer;
66 class KaleidoscopeJIT;
67
68 class KaleidoscopeASTMaterializationUnit : public MaterializationUnit {
69 public:
70 KaleidoscopeASTMaterializationUnit(KaleidoscopeASTLayer &L,
71 std::unique_ptr<FunctionAST> F);
72
getName()73 StringRef getName() const override {
74 return "KaleidoscopeASTMaterializationUnit";
75 }
76
77 void materialize(std::unique_ptr<MaterializationResponsibility> R) override;
78
79 private:
discard(const JITDylib & JD,const SymbolStringPtr & Sym)80 void discard(const JITDylib &JD, const SymbolStringPtr &Sym) override {
81 llvm_unreachable("Kaleidoscope functions are not overridable");
82 }
83
84 KaleidoscopeASTLayer &L;
85 std::unique_ptr<FunctionAST> F;
86 };
87
88 class KaleidoscopeASTLayer {
89 public:
KaleidoscopeASTLayer(IRLayer & BaseLayer,const DataLayout & DL)90 KaleidoscopeASTLayer(IRLayer &BaseLayer, const DataLayout &DL)
91 : BaseLayer(BaseLayer), DL(DL) {}
92
add(ResourceTrackerSP RT,std::unique_ptr<FunctionAST> F)93 Error add(ResourceTrackerSP RT, std::unique_ptr<FunctionAST> F) {
94 return RT->getJITDylib().define(
95 std::make_unique<KaleidoscopeASTMaterializationUnit>(*this,
96 std::move(F)),
97 RT);
98 }
99
emit(std::unique_ptr<MaterializationResponsibility> MR,std::unique_ptr<FunctionAST> F)100 void emit(std::unique_ptr<MaterializationResponsibility> MR,
101 std::unique_ptr<FunctionAST> F) {
102 BaseLayer.emit(std::move(MR), irgenAndTakeOwnership(*F, ""));
103 }
104
getInterface(FunctionAST & F)105 MaterializationUnit::Interface getInterface(FunctionAST &F) {
106 MangleAndInterner Mangle(BaseLayer.getExecutionSession(), DL);
107 SymbolFlagsMap Symbols;
108 Symbols[Mangle(F.getName())] =
109 JITSymbolFlags(JITSymbolFlags::Exported | JITSymbolFlags::Callable);
110 return MaterializationUnit::Interface(std::move(Symbols), nullptr);
111 }
112
113 private:
114 IRLayer &BaseLayer;
115 const DataLayout &DL;
116 };
117
KaleidoscopeASTMaterializationUnit(KaleidoscopeASTLayer & L,std::unique_ptr<FunctionAST> F)118 KaleidoscopeASTMaterializationUnit::KaleidoscopeASTMaterializationUnit(
119 KaleidoscopeASTLayer &L, std::unique_ptr<FunctionAST> F)
120 : MaterializationUnit(L.getInterface(*F)), L(L), F(std::move(F)) {}
121
materialize(std::unique_ptr<MaterializationResponsibility> R)122 void KaleidoscopeASTMaterializationUnit::materialize(
123 std::unique_ptr<MaterializationResponsibility> R) {
124 L.emit(std::move(R), std::move(F));
125 }
126
127 class KaleidoscopeJIT {
128 private:
129 std::unique_ptr<ExecutionSession> ES;
130 std::unique_ptr<EPCIndirectionUtils> EPCIU;
131
132 DataLayout DL;
133 MangleAndInterner Mangle;
134
135 RTDyldObjectLinkingLayer ObjectLayer;
136 IRCompileLayer CompileLayer;
137 IRTransformLayer OptimizeLayer;
138 KaleidoscopeASTLayer ASTLayer;
139
140 JITDylib &MainJD;
141
handleLazyCallThroughError()142 static void handleLazyCallThroughError() {
143 errs() << "LazyCallThrough error: Could not find function body";
144 exit(1);
145 }
146
147 public:
KaleidoscopeJIT(std::unique_ptr<ExecutionSession> ES,std::unique_ptr<EPCIndirectionUtils> EPCIU,JITTargetMachineBuilder JTMB,DataLayout DL)148 KaleidoscopeJIT(std::unique_ptr<ExecutionSession> ES,
149 std::unique_ptr<EPCIndirectionUtils> EPCIU,
150 JITTargetMachineBuilder JTMB, DataLayout DL)
151 : ES(std::move(ES)), EPCIU(std::move(EPCIU)), DL(std::move(DL)),
152 Mangle(*this->ES, this->DL),
153 ObjectLayer(*this->ES,
154 []() { return std::make_unique<SectionMemoryManager>(); }),
155 CompileLayer(*this->ES, ObjectLayer,
156 std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
157 OptimizeLayer(*this->ES, CompileLayer, optimizeModule),
158 ASTLayer(OptimizeLayer, this->DL),
159 MainJD(this->ES->createBareJITDylib("<main>")) {
160 MainJD.addGenerator(
161 cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
162 DL.getGlobalPrefix())));
163 }
164
~KaleidoscopeJIT()165 ~KaleidoscopeJIT() {
166 if (auto Err = ES->endSession())
167 ES->reportError(std::move(Err));
168 if (auto Err = EPCIU->cleanup())
169 ES->reportError(std::move(Err));
170 }
171
Create()172 static Expected<std::unique_ptr<KaleidoscopeJIT>> Create() {
173 auto EPC = SelfExecutorProcessControl::Create();
174 if (!EPC)
175 return EPC.takeError();
176
177 auto ES = std::make_unique<ExecutionSession>(std::move(*EPC));
178
179 auto EPCIU = EPCIndirectionUtils::Create(*ES);
180 if (!EPCIU)
181 return EPCIU.takeError();
182
183 (*EPCIU)->createLazyCallThroughManager(
184 *ES, ExecutorAddr::fromPtr(&handleLazyCallThroughError));
185
186 if (auto Err = setUpInProcessLCTMReentryViaEPCIU(**EPCIU))
187 return std::move(Err);
188
189 JITTargetMachineBuilder JTMB(
190 ES->getExecutorProcessControl().getTargetTriple());
191
192 auto DL = JTMB.getDefaultDataLayoutForTarget();
193 if (!DL)
194 return DL.takeError();
195
196 return std::make_unique<KaleidoscopeJIT>(std::move(ES), std::move(*EPCIU),
197 std::move(JTMB), std::move(*DL));
198 }
199
getDataLayout()200 const DataLayout &getDataLayout() const { return DL; }
201
getMainJITDylib()202 JITDylib &getMainJITDylib() { return MainJD; }
203
204 Error addModule(ThreadSafeModule TSM, ResourceTrackerSP RT = nullptr) {
205 if (!RT)
206 RT = MainJD.getDefaultResourceTracker();
207
208 return OptimizeLayer.add(RT, std::move(TSM));
209 }
210
211 Error addAST(std::unique_ptr<FunctionAST> F, ResourceTrackerSP RT = nullptr) {
212 if (!RT)
213 RT = MainJD.getDefaultResourceTracker();
214 return ASTLayer.add(RT, std::move(F));
215 }
216
lookup(StringRef Name)217 Expected<ExecutorSymbolDef> lookup(StringRef Name) {
218 return ES->lookup({&MainJD}, Mangle(Name.str()));
219 }
220
221 private:
222 static Expected<ThreadSafeModule>
optimizeModule(ThreadSafeModule TSM,const MaterializationResponsibility & R)223 optimizeModule(ThreadSafeModule TSM, const MaterializationResponsibility &R) {
224 TSM.withModuleDo([](Module &M) {
225 // Create a function pass manager.
226 auto FPM = std::make_unique<legacy::FunctionPassManager>(&M);
227
228 // Add some optimizations.
229 FPM->add(createInstructionCombiningPass());
230 FPM->add(createReassociatePass());
231 FPM->add(createGVNPass());
232 FPM->add(createCFGSimplificationPass());
233 FPM->doInitialization();
234
235 // Run the optimizations over all functions in the module being added to
236 // the JIT.
237 for (auto &F : M)
238 FPM->run(F);
239 });
240
241 return std::move(TSM);
242 }
243 };
244
245 } // end namespace orc
246 } // end namespace llvm
247
248 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
249