[Orc] Refactor JITCompileCallbackManagerBase and CompileOnDemandLayer to support
[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     BaseLayerModuleSetHandleListT &BaseLayerHandles = H->second;
197     for (auto &BH : BaseLayerHandles) {
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                             CallbackInfo.getAddress());
262         makeStub(*Proto, *FunctionBodyPointer);
263
264         F.setName(Name + BodySuffix);
265         F.setVisibility(GlobalValue::HiddenVisibility);
266
267         auto KV = std::make_pair(std::move(Name), std::move(CallbackInfo));
268         NewStubInfos.push_back(StubInfos.insert(StubInfos.begin(), KV));
269       }
270
271       auto H = addModule(std::move(SubM), MSI, LogicalModule, FallbackLookup);
272
273       // Set the compile actions for this module:
274       for (auto &KVPair : NewStubInfos) {
275         std::string BodyName = Mangle(KVPair->first + BodySuffix,
276                                       M.getDataLayout());
277         auto &CCInfo = KVPair->second;
278         CCInfo.setCompileAction(
279           [=](){
280             return BaseLayer.findSymbolIn(H, BodyName, false).getAddress();
281           });
282       }
283
284     }
285
286     // Ok - we've processed all the partitioned modules. Now add the
287     // stubs/globals module and set the update actions.
288     auto StubsH =
289       addModule(std::move(StubsModule), MSI, LogicalModule, FallbackLookup);
290
291     for (auto &KVPair : StubInfos) {
292       std::string AddrName = Mangle(KVPair.first + AddrSuffix,
293                                     M.getDataLayout());
294       auto &CCInfo = KVPair.second;
295       CCInfo.setUpdateAction(
296         getLocalFPUpdater(BaseLayer, StubsH, AddrName));
297     }
298   }
299
300   // Add the given Module to the base layer using a memory manager that will
301   // perform the appropriate scoped lookup (i.e. will look first with in the
302   // module from which it was extracted, then into the set to which that module
303   // belonged, and finally externally).
304   BaseLayerModuleSetHandleT addModule(
305                                std::unique_ptr<Module> M,
306                                ModuleSetInfo &MSI,
307                                typename CODScopedLookup::LMHandle LogicalModule,
308                                LookupFtor FallbackLookup) {
309
310     // Add this module to the JIT with a memory manager that uses the
311     // DylibLookup to resolve symbols.
312     std::vector<std::unique_ptr<Module>> MSet;
313     MSet.push_back(std::move(M));
314
315     auto DylibLookup = MSI.Lookup;
316     auto MM =
317       createLookasideRTDyldMM<SectionMemoryManager>(
318         [=](const std::string &Name) {
319           if (auto Symbol = DylibLookup->findSymbol(LogicalModule, Name))
320             return Symbol.getAddress();
321           return FallbackLookup(Name);
322         },
323         [=](const std::string &Name) {
324           return DylibLookup->findSymbol(LogicalModule, Name).getAddress();
325         });
326
327     BaseLayerModuleSetHandleT H =
328       BaseLayer.addModuleSet(std::move(MSet), std::move(MM));
329     // Add this module to the logical module lookup.
330     DylibLookup->addToLogicalModule(LogicalModule, H);
331     MSI.BaseLayerModuleSetHandles.push_back(H);
332
333     return H;
334   }
335
336   static std::string Mangle(StringRef Name, const DataLayout &DL) {
337     Mangler M(&DL);
338     std::string MangledName;
339     {
340       raw_string_ostream MangledNameStream(MangledName);
341       M.getNameWithPrefix(MangledNameStream, Name);
342     }
343     return MangledName;
344   }
345
346   BaseLayerT &BaseLayer;
347   CompileCallbackMgrT &CompileCallbackMgr;
348   ModuleSetInfoListT ModuleSetInfos;
349 };
350
351 } // End namespace orc.
352 } // End namespace llvm.
353
354 #endif // LLVM_EXECUTIONENGINE_ORC_COMPILEONDEMANDLAYER_H