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/CompileUtils.h" |
18 | #include "llvm/ExecutionEngine/Orc/Core.h" |
19 | #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h" |
20 | #include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h" |
21 | #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h" |
22 | #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h" |
23 | #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h" |
24 | #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h" |
25 | #include "llvm/ExecutionEngine/Orc/Shared/ExecutorSymbolDef.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: |
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>(args: std::move(JTMB))), |
59 | OptimizeLayer(*this->ES, CompileLayer, optimizeModule), |
60 | MainJD(this->ES->createBareJITDylib(Name: "<main>" )) { |
61 | MainJD.addGenerator( |
62 | DefGenerator: cantFail(ValOrErr: DynamicLibrarySearchGenerator::GetForCurrentProcess( |
63 | GlobalPrefix: DL.getGlobalPrefix()))); |
64 | } |
65 | |
66 | ~KaleidoscopeJIT() { |
67 | if (auto Err = ES->endSession()) |
68 | ES->reportError(Err: std::move(Err)); |
69 | } |
70 | |
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>(args: 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>(args: std::move(ES), args: std::move(JTMB), |
86 | args: std::move(*DL)); |
87 | } |
88 | |
89 | const DataLayout &getDataLayout() const { return DL; } |
90 | |
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, TSM: std::move(TSM)); |
98 | } |
99 | |
100 | Expected<ExecutorSymbolDef> lookup(StringRef Name) { |
101 | return ES->lookup(SearchOrder: {&MainJD}, Symbol: Mangle(Name.str())); |
102 | } |
103 | |
104 | private: |
105 | static Expected<ThreadSafeModule> |
106 | optimizeModule(ThreadSafeModule TSM, const MaterializationResponsibility &R) { |
107 | TSM.withModuleDo(F: [](Module &M) { |
108 | // Create a function pass manager. |
109 | auto FPM = std::make_unique<legacy::FunctionPassManager>(args: &M); |
110 | |
111 | // Add some optimizations. |
112 | FPM->add(P: createInstructionCombiningPass()); |
113 | FPM->add(P: createReassociatePass()); |
114 | FPM->add(P: createGVNPass()); |
115 | FPM->add(P: 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 | |