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