1 //===----- LLJIT.h -- An ORC-based JIT for compiling LLVM IR ----*- 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 // An ORC-based JIT for compiling LLVM IR.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #ifndef LLVM_EXECUTIONENGINE_ORC_LLJIT_H
14 #define LLVM_EXECUTIONENGINE_ORC_LLJIT_H
15 
16 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
17 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
18 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
19 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
20 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
21 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
22 #include "llvm/ExecutionEngine/Orc/ObjectTransformLayer.h"
23 #include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
24 #include "llvm/Support/ThreadPool.h"
25 
26 namespace llvm {
27 namespace orc {
28 
29 class LLJITBuilderState;
30 class LLLazyJITBuilderState;
31 
32 /// A pre-fabricated ORC JIT stack that can serve as an alternative to MCJIT.
33 ///
34 /// Create instances using LLJITBuilder.
35 class LLJIT {
36   template <typename, typename, typename> friend class LLJITBuilderSetters;
37 
38 public:
39   static Expected<std::unique_ptr<LLJIT>> Create(LLJITBuilderState &S);
40 
41   /// Destruct this instance. If a multi-threaded instance, waits for all
42   /// compile threads to complete.
43   ~LLJIT();
44 
45   /// Returns the ExecutionSession for this instance.
46   ExecutionSession &getExecutionSession() { return *ES; }
47 
48   /// Returns a reference to the DataLayout for this instance.
49   const DataLayout &getDataLayout() const { return DL; }
50 
51   /// Returns a reference to the JITDylib representing the JIT'd main program.
52   JITDylib &getMainJITDylib() { return Main; }
53 
54   /// Returns the JITDylib with the given name, or nullptr if no JITDylib with
55   /// that name exists.
56   JITDylib *getJITDylibByName(StringRef Name) {
57     return ES->getJITDylibByName(Name);
58   }
59 
60   /// Create a new JITDylib with the given name and return a reference to it.
61   ///
62   /// JITDylib names must be unique. If the given name is derived from user
63   /// input or elsewhere in the environment then the client should check
64   /// (e.g. by calling getJITDylibByName) that the given name is not already in
65   /// use.
66   JITDylib &createJITDylib(std::string Name) {
67     return ES->createJITDylib(std::move(Name));
68   }
69 
70   /// Convenience method for defining an absolute symbol.
71   Error defineAbsolute(StringRef Name, JITEvaluatedSymbol Address);
72 
73   /// Adds an IR module to the given JITDylib.
74   Error addIRModule(JITDylib &JD, ThreadSafeModule TSM);
75 
76   /// Adds an IR module to the Main JITDylib.
77   Error addIRModule(ThreadSafeModule TSM) {
78     return addIRModule(Main, std::move(TSM));
79   }
80 
81   /// Adds an object file to the given JITDylib.
82   Error addObjectFile(JITDylib &JD, std::unique_ptr<MemoryBuffer> Obj);
83 
84   /// Adds an object file to the given JITDylib.
85   Error addObjectFile(std::unique_ptr<MemoryBuffer> Obj) {
86     return addObjectFile(Main, std::move(Obj));
87   }
88 
89   /// Look up a symbol in JITDylib JD by the symbol's linker-mangled name (to
90   /// look up symbols based on their IR name use the lookup function instead).
91   Expected<JITEvaluatedSymbol> lookupLinkerMangled(JITDylib &JD,
92                                                    StringRef Name);
93 
94   /// Look up a symbol in the main JITDylib by the symbol's linker-mangled name
95   /// (to look up symbols based on their IR name use the lookup function
96   /// instead).
97   Expected<JITEvaluatedSymbol> lookupLinkerMangled(StringRef Name) {
98     return lookupLinkerMangled(Main, Name);
99   }
100 
101   /// Look up a symbol in JITDylib JD based on its IR symbol name.
102   Expected<JITEvaluatedSymbol> lookup(JITDylib &JD, StringRef UnmangledName) {
103     return lookupLinkerMangled(JD, mangle(UnmangledName));
104   }
105 
106   /// Look up a symbol in the main JITDylib based on its IR symbol name.
107   Expected<JITEvaluatedSymbol> lookup(StringRef UnmangledName) {
108     return lookup(Main, UnmangledName);
109   }
110 
111   /// Runs all not-yet-run static constructors.
112   Error runConstructors() { return CtorRunner.run(); }
113 
114   /// Runs all not-yet-run static destructors.
115   Error runDestructors() { return DtorRunner.run(); }
116 
117   /// Returns a reference to the ObjLinkingLayer
118   ObjectLayer &getObjLinkingLayer() { return *ObjLinkingLayer; }
119 
120   /// Returns a reference to the object transform layer.
121   ObjectTransformLayer &getObjTransformLayer() { return ObjTransformLayer; }
122 
123 protected:
124   static std::unique_ptr<ObjectLayer>
125   createObjectLinkingLayer(LLJITBuilderState &S, ExecutionSession &ES);
126 
127   static Expected<std::unique_ptr<IRCompileLayer::IRCompiler>>
128   createCompileFunction(LLJITBuilderState &S, JITTargetMachineBuilder JTMB);
129 
130   /// Create an LLJIT instance with a single compile thread.
131   LLJIT(LLJITBuilderState &S, Error &Err);
132 
133   std::string mangle(StringRef UnmangledName);
134 
135   Error applyDataLayout(Module &M);
136 
137   void recordCtorDtors(Module &M);
138 
139   std::unique_ptr<ExecutionSession> ES;
140   JITDylib &Main;
141 
142   DataLayout DL;
143   std::unique_ptr<ThreadPool> CompileThreads;
144 
145   std::unique_ptr<ObjectLayer> ObjLinkingLayer;
146   ObjectTransformLayer ObjTransformLayer;
147   std::unique_ptr<IRCompileLayer> CompileLayer;
148 
149   CtorDtorRunner CtorRunner, DtorRunner;
150 };
151 
152 /// An extended version of LLJIT that supports lazy function-at-a-time
153 /// compilation of LLVM IR.
154 class LLLazyJIT : public LLJIT {
155   template <typename, typename, typename> friend class LLJITBuilderSetters;
156 
157 public:
158 
159   /// Set an IR transform (e.g. pass manager pipeline) to run on each function
160   /// when it is compiled.
161   void setLazyCompileTransform(IRTransformLayer::TransformFunction Transform) {
162     TransformLayer->setTransform(std::move(Transform));
163   }
164 
165   /// Sets the partition function.
166   void
167   setPartitionFunction(CompileOnDemandLayer::PartitionFunction Partition) {
168     CODLayer->setPartitionFunction(std::move(Partition));
169   }
170 
171   /// Add a module to be lazily compiled to JITDylib JD.
172   Error addLazyIRModule(JITDylib &JD, ThreadSafeModule M);
173 
174   /// Add a module to be lazily compiled to the main JITDylib.
175   Error addLazyIRModule(ThreadSafeModule M) {
176     return addLazyIRModule(Main, std::move(M));
177   }
178 
179 private:
180 
181   // Create a single-threaded LLLazyJIT instance.
182   LLLazyJIT(LLLazyJITBuilderState &S, Error &Err);
183 
184   std::unique_ptr<LazyCallThroughManager> LCTMgr;
185   std::unique_ptr<IRTransformLayer> TransformLayer;
186   std::unique_ptr<CompileOnDemandLayer> CODLayer;
187 };
188 
189 class LLJITBuilderState {
190 public:
191   using ObjectLinkingLayerCreator = std::function<std::unique_ptr<ObjectLayer>(
192       ExecutionSession &, const Triple &TT)>;
193 
194   using CompileFunctionCreator =
195       std::function<Expected<std::unique_ptr<IRCompileLayer::IRCompiler>>(
196           JITTargetMachineBuilder JTMB)>;
197 
198   std::unique_ptr<ExecutionSession> ES;
199   Optional<JITTargetMachineBuilder> JTMB;
200   ObjectLinkingLayerCreator CreateObjectLinkingLayer;
201   CompileFunctionCreator CreateCompileFunction;
202   unsigned NumCompileThreads = 0;
203 
204   /// Called prior to JIT class construcion to fix up defaults.
205   Error prepareForConstruction();
206 };
207 
208 template <typename JITType, typename SetterImpl, typename State>
209 class LLJITBuilderSetters {
210 public:
211   /// Set the JITTargetMachineBuilder for this instance.
212   ///
213   /// If this method is not called, JITTargetMachineBuilder::detectHost will be
214   /// used to construct a default target machine builder for the host platform.
215   SetterImpl &setJITTargetMachineBuilder(JITTargetMachineBuilder JTMB) {
216     impl().JTMB = std::move(JTMB);
217     return impl();
218   }
219 
220   /// Return a reference to the JITTargetMachineBuilder.
221   ///
222   Optional<JITTargetMachineBuilder> &getJITTargetMachineBuilder() {
223     return impl().JTMB;
224   }
225 
226   /// Set an ObjectLinkingLayer creation function.
227   ///
228   /// If this method is not called, a default creation function will be used
229   /// that will construct an RTDyldObjectLinkingLayer.
230   SetterImpl &setObjectLinkingLayerCreator(
231       LLJITBuilderState::ObjectLinkingLayerCreator CreateObjectLinkingLayer) {
232     impl().CreateObjectLinkingLayer = std::move(CreateObjectLinkingLayer);
233     return impl();
234   }
235 
236   /// Set a CompileFunctionCreator.
237   ///
238   /// If this method is not called, a default creation function wil be used
239   /// that will construct a basic IR compile function that is compatible with
240   /// the selected number of threads (SimpleCompiler for '0' compile threads,
241   /// ConcurrentIRCompiler otherwise).
242   SetterImpl &setCompileFunctionCreator(
243       LLJITBuilderState::CompileFunctionCreator CreateCompileFunction) {
244     impl().CreateCompileFunction = std::move(CreateCompileFunction);
245     return impl();
246   }
247 
248   /// Set the number of compile threads to use.
249   ///
250   /// If set to zero, compilation will be performed on the execution thread when
251   /// JITing in-process. If set to any other number N, a thread pool of N
252   /// threads will be created for compilation.
253   ///
254   /// If this method is not called, behavior will be as if it were called with
255   /// a zero argument.
256   SetterImpl &setNumCompileThreads(unsigned NumCompileThreads) {
257     impl().NumCompileThreads = NumCompileThreads;
258     return impl();
259   }
260 
261   /// Create an instance of the JIT.
262   Expected<std::unique_ptr<JITType>> create() {
263     if (auto Err = impl().prepareForConstruction())
264       return std::move(Err);
265 
266     Error Err = Error::success();
267     std::unique_ptr<JITType> J(new JITType(impl(), Err));
268     if (Err)
269       return std::move(Err);
270     return std::move(J);
271   }
272 
273 protected:
274   SetterImpl &impl() { return static_cast<SetterImpl &>(*this); }
275 };
276 
277 /// Constructs LLJIT instances.
278 class LLJITBuilder
279     : public LLJITBuilderState,
280       public LLJITBuilderSetters<LLJIT, LLJITBuilder, LLJITBuilderState> {};
281 
282 class LLLazyJITBuilderState : public LLJITBuilderState {
283   friend class LLLazyJIT;
284 
285 public:
286   using IndirectStubsManagerBuilderFunction =
287       std::function<std::unique_ptr<IndirectStubsManager>()>;
288 
289   Triple TT;
290   JITTargetAddress LazyCompileFailureAddr = 0;
291   std::unique_ptr<LazyCallThroughManager> LCTMgr;
292   IndirectStubsManagerBuilderFunction ISMBuilder;
293 
294   Error prepareForConstruction();
295 };
296 
297 template <typename JITType, typename SetterImpl, typename State>
298 class LLLazyJITBuilderSetters
299     : public LLJITBuilderSetters<JITType, SetterImpl, State> {
300 public:
301   /// Set the address in the target address to call if a lazy compile fails.
302   ///
303   /// If this method is not called then the value will default to 0.
304   SetterImpl &setLazyCompileFailureAddr(JITTargetAddress Addr) {
305     this->impl().LazyCompileFailureAddr = Addr;
306     return this->impl();
307   }
308 
309   /// Set the lazy-callthrough manager.
310   ///
311   /// If this method is not called then a default, in-process lazy callthrough
312   /// manager for the host platform will be used.
313   SetterImpl &
314   setLazyCallthroughManager(std::unique_ptr<LazyCallThroughManager> LCTMgr) {
315     this->impl().LCTMgr = std::move(LCTMgr);
316     return this->impl();
317   }
318 
319   /// Set the IndirectStubsManager builder function.
320   ///
321   /// If this method is not called then a default, in-process
322   /// IndirectStubsManager builder for the host platform will be used.
323   SetterImpl &setIndirectStubsManagerBuilder(
324       LLLazyJITBuilderState::IndirectStubsManagerBuilderFunction ISMBuilder) {
325     this->impl().ISMBuilder = std::move(ISMBuilder);
326     return this->impl();
327   }
328 };
329 
330 /// Constructs LLLazyJIT instances.
331 class LLLazyJITBuilder
332     : public LLLazyJITBuilderState,
333       public LLLazyJITBuilderSetters<LLLazyJIT, LLLazyJITBuilder,
334                                      LLLazyJITBuilderState> {};
335 
336 } // End namespace orc
337 } // End namespace llvm
338 
339 #endif // LLVM_EXECUTIONENGINE_ORC_LLJIT_H
340