[Orc] Use raw TargetAddresses for callback trampoline addresses, rather than IR.
[oota-llvm.git] / include / llvm / ExecutionEngine / Orc / CompileOnDemandLayer.h
1 //===- CompileOnDemandLayer.h - Compile each function on demand -*- C++ -*-===//
2 //
3 //                     The LLVM Compiler Infrastructure
4 //
5 // This file is distributed under the University of Illinois Open Source
6 // License. See LICENSE.TXT for details.
7 //
8 //===----------------------------------------------------------------------===//
9 //
10 // JIT layer for breaking up modules and inserting callbacks to allow
11 // individual functions to be compiled on demand.
12 //
13 //===----------------------------------------------------------------------===//
14
15 #ifndef LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
16 #define LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H
17
18 #include "IndirectionUtils.h"
19 #include "LookasideRTDyldMM.h"
20 #include "llvm/ADT/STLExtras.h"
21 #include "llvm/ExecutionEngine/SectionMemoryManager.h"
22 #include <list>
23
24 namespace llvm {
25 namespace orc {
26
27 /// @brief Compile-on-demand layer.
28 ///
29 ///   Modules added to this layer have their calls indirected, and are then
30 /// broken up into a set of single-function modules, each of which is added
31 /// to the layer below in a singleton set. The lower layer can be any layer that
32 /// accepts IR module sets.
33 ///
34 /// It is expected that this layer will frequently be used on top of a
35 /// LazyEmittingLayer. The combination of the two ensures that each function is
36 /// compiled only when it is first called.
37 template <typename BaseLayerT, typename CompileCallbackMgrT>
38 class CompileOnDemandLayer {
39 public:
40   /// @brief Lookup helper that provides compatibility with the classic
41   ///        static-compilation symbol resolution process.
42   ///
43   ///   The CompileOnDemand (COD) layer splits modules up into multiple
44   /// sub-modules, each held in its own llvm::Module instance, in order to
45   /// support lazy compilation. When a module that contains private symbols is
46   /// broken up symbol linkage changes may be required to enable access to
47   /// "private" data that now resides in a different llvm::Module instance. To
48   /// retain expected symbol resolution behavior for clients of the COD layer,
49   /// the CODScopedLookup class uses a two-tiered lookup system to resolve
50   /// symbols. Lookup first scans sibling modules that were split from the same
51   /// original module (logical-module scoped lookup), then scans all other
52   /// modules that have been added to the lookup scope (logical-dylib scoped
53   /// lookup).
54   class CODScopedLookup {
55   private:
56     typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
57     typedef std::vector<BaseLayerModuleSetHandleT> SiblingHandlesList;
58     typedef std::list<SiblingHandlesList> PseudoDylibModuleSetHandlesList;
59
60   public:
61     /// @brief Handle for a logical module.
62     typedef typename PseudoDylibModuleSetHandlesList::iterator LMHandle;
63
64     /// @brief Construct a scoped lookup.
65     CODScopedLookup(BaseLayerT &BaseLayer) : BaseLayer(BaseLayer) {}
66
67     /// @brief Start a new context for a single logical module.
68     LMHandle createLogicalModule() {
69       Handles.push_back(SiblingHandlesList());
70       return std::prev(Handles.end());
71     }
72
73     /// @brief Add a concrete Module's handle to the given logical Module's
74     ///        lookup scope.
75     void addToLogicalModule(LMHandle LMH, BaseLayerModuleSetHandleT H) {
76       LMH->push_back(H);
77     }
78
79     /// @brief Remove a logical Module from the CODScopedLookup entirely.
80     void removeLogicalModule(LMHandle LMH) { Handles.erase(LMH); }
81
82     /// @brief Look up a symbol in this context.
83     JITSymbol findSymbol(LMHandle LMH, const std::string &Name) {
84       if (auto Symbol = findSymbolIn(LMH, Name))
85         return Symbol;
86
87       for (auto I = Handles.begin(), E = Handles.end(); I != E; ++I)
88         if (I != LMH)
89           if (auto Symbol = findSymbolIn(I, Name))
90             return Symbol;
91
92       return nullptr;
93     }
94
95   private:
96
97     JITSymbol findSymbolIn(LMHandle LMH, const std::string &Name) {
98       for (auto H : *LMH)
99         if (auto Symbol = BaseLayer.findSymbolIn(H, Name, false))
100           return Symbol;
101       return nullptr;
102     }
103
104     BaseLayerT &BaseLayer;
105     PseudoDylibModuleSetHandlesList Handles;
106   };
107
108 private:
109   typedef typename BaseLayerT::ModuleSetHandleT BaseLayerModuleSetHandleT;
110   typedef std::vector<BaseLayerModuleSetHandleT> BaseLayerModuleSetHandleListT;
111
112   struct ModuleSetInfo {
113     // Symbol lookup - just one for the whole module set.
114     std::shared_ptr<CODScopedLookup> Lookup;
115
116     // Logical module handles.
117     std::vector<typename CODScopedLookup::LMHandle> LMHandles;
118
119     // List of vectors of module set handles:
120     // One vector per logical module - each vector holds the handles for the
121     // exploded modules for that logical module in the base layer.
122     BaseLayerModuleSetHandleListT BaseLayerModuleSetHandles;
123
124     ModuleSetInfo(std::shared_ptr<CODScopedLookup> Lookup)
125         : Lookup(std::move(Lookup)) {}
126
127     void releaseResources(BaseLayerT &BaseLayer) {
128       for (auto LMH : LMHandles)
129         Lookup->removeLogicalModule(LMH);
130       for (auto H : BaseLayerModuleSetHandles)
131         BaseLayer.removeModuleSet(H);
132     }
133   };
134
135   typedef std::list<ModuleSetInfo> ModuleSetInfoListT;
136
137 public:
138   /// @brief Handle to a set of loaded modules.
139   typedef typename ModuleSetInfoListT::iterator ModuleSetHandleT;
140
141   // @brief Fallback lookup functor.
142   typedef std::function<uint64_t(const std::string &)> LookupFtor;
143
144   /// @brief Construct a compile-on-demand layer instance.
145   CompileOnDemandLayer(BaseLayerT &BaseLayer, CompileCallbackMgrT &CallbackMgr)
146       : BaseLayer(BaseLayer), CompileCallbackMgr(CallbackMgr) {}
147
148   /// @brief Add a module to the compile-on-demand layer.
149   template <typename ModuleSetT>
150   ModuleSetHandleT addModuleSet(ModuleSetT Ms,
151                                 LookupFtor FallbackLookup = nullptr) {
152
153     // If the user didn't supply a fallback lookup then just use
154     // getSymbolAddress.
155     if (!FallbackLookup)
156       FallbackLookup = [=](const std::string &Name) {
157                          return findSymbol(Name, true).getAddress();
158                        };
159
160     // Create a lookup context and ModuleSetInfo for this module set.
161     // For the purposes of symbol resolution the set Ms will be treated as if
162     // the modules it contained had been linked together as a dylib.
163     auto DylibLookup = std::make_shared<CODScopedLookup>(BaseLayer);
164     ModuleSetHandleT H =
165         ModuleSetInfos.insert(ModuleSetInfos.end(), ModuleSetInfo(DylibLookup));
166     ModuleSetInfo &MSI = ModuleSetInfos.back();
167
168     // Process each of the modules in this module set.
169     for (auto &M : Ms)
170       partitionAndAdd(*M, MSI, FallbackLookup);
171
172     return H;
173   }
174
175   /// @brief Remove the module represented by the given handle.
176   ///
177   ///   This will remove all modules in the layers below that were derived from
178   /// the module represented by H.
179   void removeModuleSet(ModuleSetHandleT H) {
180     H->releaseResources(BaseLayer);
181     ModuleSetInfos.erase(H);
182   }
183
184   /// @brief Search for the given named symbol.
185   /// @param Name The name of the symbol to search for.
186   /// @param ExportedSymbolsOnly If true, search only for exported symbols.
187   /// @return A handle for the given named symbol, if it exists.
188   JITSymbol findSymbol(StringRef Name, bool ExportedSymbolsOnly) {
189     return BaseLayer.findSymbol(Name, ExportedSymbolsOnly);
190   }
191
192   /// @brief Get the address of a symbol provided by this layer, or some layer
193   ///        below this one.
194   JITSymbol findSymbolIn(ModuleSetHandleT H, const std::string &Name,
195                          bool ExportedSymbolsOnly) {
196
197     for (auto &BH : H->BaseLayerModuleSetHandles) {
198       if (auto Symbol = BaseLayer.findSymbolIn(BH, Name, ExportedSymbolsOnly))
199         return Symbol;
200     }
201     return nullptr;
202   }
203
204 private:
205
206   void partitionAndAdd(Module &M, ModuleSetInfo &MSI,
207                        LookupFtor FallbackLookup) {
208     const char *AddrSuffix = "$orc_addr";
209     const char *BodySuffix = "$orc_body";
210
211     // We're going to break M up into a bunch of sub-modules, but we want
212     // internal linkage symbols to still resolve sensibly. CODScopedLookup
213     // provides the "logical module" concept to make this work, so create a
214     // new logical module for M.
215     auto DylibLookup = MSI.Lookup;
216     auto LogicalModule = DylibLookup->createLogicalModule();
217     MSI.LMHandles.push_back(LogicalModule);
218
219     // Partition M into a "globals and stubs" module, a "common symbols" module,
220     // and a list of single-function modules.
221     auto PartitionedModule = fullyPartition(M);
222     auto StubsModule = std::move(PartitionedModule.GlobalVars);
223     auto CommonsModule = std::move(PartitionedModule.Commons);
224     auto FunctionModules = std::move(PartitionedModule.Functions);
225
226     // Emit the commons stright away.
227     auto CommonHandle = addModule(std::move(CommonsModule), MSI, LogicalModule,
228                                   FallbackLookup);
229     BaseLayer.emitAndFinalize(CommonHandle);
230
231     // Map of definition names to callback-info data structures. We'll use
232     // this to build the compile actions for the stubs below.
233     typedef std::map<std::string,
234                      typename CompileCallbackMgrT::CompileCallbackInfo>
235       StubInfoMap;
236     StubInfoMap StubInfos;
237
238     // Now we need to take each of the extracted Modules and add them to
239     // base layer. Each Module will be added individually to make sure they
240     // can be compiled separately, and each will get its own lookaside
241     // memory manager that will resolve within this logical module first.
242     for (auto &SubM : FunctionModules) {
243
244       // Keep track of the stubs we create for this module so that we can set
245       // their compile actions.
246       std::vector<typename StubInfoMap::iterator> NewStubInfos;
247
248       // Search for function definitions and insert stubs into the stubs
249       // module.
250       for (auto &F : *SubM) {
251         if (F.isDeclaration())
252           continue;
253
254         std::string Name = F.getName();
255         Function *Proto = StubsModule->getFunction(Name);
256         assert(Proto && "Failed to clone function decl into stubs module.");
257         auto CallbackInfo =
258           CompileCallbackMgr.getCompileCallback(*Proto->getFunctionType());
259         GlobalVariable *FunctionBodyPointer =
260           createImplPointer(*Proto, Name + AddrSuffix,
261                             createIRTypedAddress(*Proto->getFunctionType(),
262                                                  CallbackInfo.getAddress()));
263         makeStub(*Proto, *FunctionBodyPointer);
264
265         F.setName(Name + BodySuffix);
266         F.setVisibility(GlobalValue::HiddenVisibility);
267
268         auto KV = std::make_pair(std::move(Name), std::move(CallbackInfo));
269         NewStubInfos.push_back(StubInfos.insert(StubInfos.begin(), KV));
270       }
271
272       auto H = addModule(std::move(SubM), MSI, LogicalModule, FallbackLookup);
273
274       // Set the compile actions for this module:
275       for (auto &KVPair : NewStubInfos) {
276         std::string BodyName = Mangle(KVPair->first + BodySuffix,
277                                       M.getDataLayout());
278         auto &CCInfo = KVPair->second;
279         CCInfo.setCompileAction(
280           [=](){
281             return BaseLayer.findSymbolIn(H, BodyName, false).getAddress();
282           });
283       }
284
285     }
286
287     // Ok - we've processed all the partitioned modules. Now add the
288     // stubs/globals module and set the update actions.
289     auto StubsH =
290       addModule(std::move(StubsModule), MSI, LogicalModule, FallbackLookup);
291
292     for (auto &KVPair : StubInfos) {
293       std::string AddrName = Mangle(KVPair.first + AddrSuffix,
294                                     M.getDataLayout());
295       auto &CCInfo = KVPair.second;
296       CCInfo.setUpdateAction(
297         getLocalFPUpdater(BaseLayer, StubsH, AddrName));
298     }
299   }
300
301   // Add the given Module to the base layer using a memory manager that will
302   // perform the appropriate scoped lookup (i.e. will look first with in the
303   // module from which it was extracted, then into the set to which that module
304   // belonged, and finally externally).
305   BaseLayerModuleSetHandleT addModule(
306                                std::unique_ptr<Module> M,
307                                ModuleSetInfo &MSI,
308                                typename CODScopedLookup::LMHandle LogicalModule,
309                                LookupFtor FallbackLookup) {
310
311     // Add this module to the JIT with a memory manager that uses the
312     // DylibLookup to resolve symbols.
313     std::vector<std::unique_ptr<Module>> MSet;
314     MSet.push_back(std::move(M));
315
316     auto DylibLookup = MSI.Lookup;
317     auto MM =
318       createLookasideRTDyldMM<SectionMemoryManager>(
319         [=](const std::string &Name) {
320           if (auto Symbol = DylibLookup->findSymbol(LogicalModule, Name))
321             return Symbol.getAddress();
322           return FallbackLookup(Name);
323         },
324         [=](const std::string &Name) {
325           return DylibLookup->findSymbol(LogicalModule, Name).getAddress();
326         });
327
328     BaseLayerModuleSetHandleT H =
329       BaseLayer.addModuleSet(std::move(MSet), std::move(MM));
330     // Add this module to the logical module lookup.
331     DylibLookup->addToLogicalModule(LogicalModule, H);
332     MSI.BaseLayerModuleSetHandles.push_back(H);
333
334     return H;
335   }
336
337   static std::string Mangle(StringRef Name, const DataLayout &DL) {
338     Mangler M(&DL);
339     std::string MangledName;
340     {
341       raw_string_ostream MangledNameStream(MangledName);
342       M.getNameWithPrefix(MangledNameStream, Name);
343     }
344     return MangledName;
345   }
346
347   BaseLayerT &BaseLayer;
348   CompileCallbackMgrT &CompileCallbackMgr;
349   ModuleSetInfoListT ModuleSetInfos;
350 };
351
352 } // End namespace orc.
353 } // End namespace llvm.
354
355 #endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H