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/CompileOnDemandLayer.h"
19 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
20 #include "llvm/ExecutionEngine/Orc/Core.h"
21 #include "llvm/ExecutionEngine/Orc/EPCIndirectionUtils.h"
22 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
23 #include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
24 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
25 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
26 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
27 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.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 namespace llvm {
38 namespace orc {
39 
40 class KaleidoscopeJIT {
41 private:
42   std::unique_ptr<ExecutionSession> ES;
43   std::unique_ptr<EPCIndirectionUtils> EPCIU;
44 
45   DataLayout DL;
46   MangleAndInterner Mangle;
47 
48   RTDyldObjectLinkingLayer ObjectLayer;
49   IRCompileLayer CompileLayer;
50   IRTransformLayer OptimizeLayer;
51   CompileOnDemandLayer CODLayer;
52 
53   JITDylib &MainJD;
54 
55   static void handleLazyCallThroughError() {
56     errs() << "LazyCallThrough error: Could not find function body";
57     exit(1);
58   }
59 
60 public:
61   KaleidoscopeJIT(std::unique_ptr<ExecutionSession> ES,
62                   std::unique_ptr<EPCIndirectionUtils> EPCIU,
63                   JITTargetMachineBuilder JTMB, DataLayout DL)
64       : ES(std::move(ES)), EPCIU(std::move(EPCIU)), DL(std::move(DL)),
65         Mangle(*this->ES, this->DL),
66         ObjectLayer(*this->ES,
67                     []() { return std::make_unique<SectionMemoryManager>(); }),
68         CompileLayer(*this->ES, ObjectLayer,
69                      std::make_unique<ConcurrentIRCompiler>(std::move(JTMB))),
70         OptimizeLayer(*this->ES, CompileLayer, optimizeModule),
71         CODLayer(*this->ES, OptimizeLayer,
72                  this->EPCIU->getLazyCallThroughManager(),
73                  [this] { return this->EPCIU->createIndirectStubsManager(); }),
74         MainJD(this->ES->createBareJITDylib("<main>")) {
75     MainJD.addGenerator(
76         cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
77             DL.getGlobalPrefix())));
78   }
79 
80   ~KaleidoscopeJIT() {
81     if (auto Err = ES->endSession())
82       ES->reportError(std::move(Err));
83     if (auto Err = EPCIU->cleanup())
84       ES->reportError(std::move(Err));
85   }
86 
87   static Expected<std::unique_ptr<KaleidoscopeJIT>> Create() {
88     auto EPC = SelfExecutorProcessControl::Create();
89     if (!EPC)
90       return EPC.takeError();
91 
92     auto ES = std::make_unique<ExecutionSession>(std::move(*EPC));
93 
94     auto EPCIU = EPCIndirectionUtils::Create(ES->getExecutorProcessControl());
95     if (!EPCIU)
96       return EPCIU.takeError();
97 
98     (*EPCIU)->createLazyCallThroughManager(
99         *ES, pointerToJITTargetAddress(&handleLazyCallThroughError));
100 
101     if (auto Err = setUpInProcessLCTMReentryViaEPCIU(**EPCIU))
102       return std::move(Err);
103 
104     JITTargetMachineBuilder JTMB(
105         ES->getExecutorProcessControl().getTargetTriple());
106 
107     auto DL = JTMB.getDefaultDataLayoutForTarget();
108     if (!DL)
109       return DL.takeError();
110 
111     return std::make_unique<KaleidoscopeJIT>(std::move(ES), std::move(*EPCIU),
112                                              std::move(JTMB), std::move(*DL));
113   }
114 
115   const DataLayout &getDataLayout() const { return DL; }
116 
117   JITDylib &getMainJITDylib() { return MainJD; }
118 
119   Error addModule(ThreadSafeModule TSM, ResourceTrackerSP RT = nullptr) {
120     if (!RT)
121       RT = MainJD.getDefaultResourceTracker();
122 
123     return CODLayer.add(RT, std::move(TSM));
124   }
125 
126   Expected<JITEvaluatedSymbol> lookup(StringRef Name) {
127     return ES->lookup({&MainJD}, Mangle(Name.str()));
128   }
129 
130 private:
131   static Expected<ThreadSafeModule>
132   optimizeModule(ThreadSafeModule TSM, const MaterializationResponsibility &R) {
133     TSM.withModuleDo([](Module &M) {
134       // Create a function pass manager.
135       auto FPM = std::make_unique<legacy::FunctionPassManager>(&M);
136 
137       // Add some optimizations.
138       FPM->add(createInstructionCombiningPass());
139       FPM->add(createReassociatePass());
140       FPM->add(createGVNPass());
141       FPM->add(createCFGSimplificationPass());
142       FPM->doInitialization();
143 
144       // Run the optimizations over all functions in the module being added to
145       // the JIT.
146       for (auto &F : M)
147         FPM->run(F);
148     });
149 
150     return std::move(TSM);
151   }
152 };
153 
154 } // end namespace orc
155 } // end namespace llvm
156 
157 #endif // LLVM_EXECUTIONENGINE_ORC_KALEIDOSCOPEJIT_H
158