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/JITSymbol.h"
18 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
19 #include "llvm/ExecutionEngine/Orc/Core.h"
20 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
21 #include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
22 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
23 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
24 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
25 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
26 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
27 #include "llvm/IR/DataLayout.h"
28 #include "llvm/IR/LLVMContext.h"
29 #include "llvm/IR/LegacyPassManager.h"
30 #include "llvm/Transforms/InstCombine/InstCombine.h"
31 #include "llvm/Transforms/Scalar.h"
32 #include "llvm/Transforms/Scalar/GVN.h"
33 #include <memory>
34 
35 namespace llvm {
36 namespace orc {
37 
38 class KaleidoscopeJIT {
39 private:
40   std::unique_ptr<ExecutionSession> ES;
41 
42   DataLayout DL;
43   MangleAndInterner Mangle;
44 
45   RTDyldObjectLinkingLayer ObjectLayer;
46   IRCompileLayer CompileLayer;
47   IRTransformLayer OptimizeLayer;
48 
49   JITDylib &MainJD;
50 
51 public:
KaleidoscopeJIT(std::unique_ptr<ExecutionSession> ES,JITTargetMachineBuilder JTMB,DataLayout DL)52   KaleidoscopeJIT(std::unique_ptr<ExecutionSession> ES,
53                   JITTargetMachineBuilder JTMB, DataLayout DL)
54       : ES(std::move(ES)), DL(std::move(DL)), Mangle(*this->ES, this->DL),
55         ObjectLayer(*this->ES,
56                     []() { return std::make_unique<SectionMemoryManager>(); }),
57         CompileLayer(*this->ES, ObjectLayer,
58                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
59         OptimizeLayer(*this->ES, CompileLayer, optimizeModule),
60         MainJD(this->ES->createBareJITDylib("<main>")) {
61     MainJD.addGenerator(
62         cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
63             DL.getGlobalPrefix())));
64   }
65 
~KaleidoscopeJIT()66   ~KaleidoscopeJIT() {
67     if (auto Err = ES->endSession())
68       ES->reportError(std::move(Err));
69   }
70 
Create()71   static Expected<std::unique_ptr<KaleidoscopeJIT>> Create() {
72     auto EPC = SelfExecutorProcessControl::Create();
73     if (!EPC)
74       return EPC.takeError();
75 
76     auto ES = std::make_unique<ExecutionSession>(std::move(*EPC));
77 
78     JITTargetMachineBuilder JTMB(
79         ES->getExecutorProcessControl().getTargetTriple());
80 
81     auto DL = JTMB.getDefaultDataLayoutForTarget();
82     if (!DL)
83       return DL.takeError();
84 
85     return std::make_unique<KaleidoscopeJIT>(std::move(ES), std::move(JTMB),
86                                              std::move(*DL));
87   }
88 
getDataLayout()89   const DataLayout &getDataLayout() const { return DL; }
90 
getMainJITDylib()91   JITDylib &getMainJITDylib() { return MainJD; }
92 
93   Error addModule(ThreadSafeModule TSM, ResourceTrackerSP RT = nullptr) {
94     if (!RT)
95       RT = MainJD.getDefaultResourceTracker();
96 
97     return OptimizeLayer.add(RT, std::move(TSM));
98   }
99 
lookup(StringRef Name)100   Expected<JITEvaluatedSymbol> lookup(StringRef Name) {
101     return ES->lookup({&MainJD}, Mangle(Name.str()));
102   }
103 
104 private:
105   static Expected<ThreadSafeModule>
optimizeModule(ThreadSafeModule TSM,const MaterializationResponsibility & R)106   optimizeModule(ThreadSafeModule TSM, const MaterializationResponsibility &R) {
107     TSM.withModuleDo([](Module &M) {
108       // Create a function pass manager.
109       auto FPM = std::make_unique<legacy::FunctionPassManager>(&M);
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 
124     return std::move(TSM);
125   }
126 };
127 
128 } // end namespace orc
129 } // end namespace llvm
130 
131 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
132