[Orc] Enable user-supplied memory managers in the CompileOnDemand layer.
[oota-llvm.git] / lib / ExecutionEngine / Orc / OrcCBindingsStack.h
1 //===--- OrcCBindingsStack.h - Orc JIT stack for C bindings ---*- 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 #ifndef LLVM_LIB_EXECUTIONENGINE_ORC_ORCCBINDINGSSTACK_H
11 #define LLVM_LIB_EXECUTIONENGINE_ORC_ORCCBINDINGSSTACK_H
12
13 #include "llvm/ADT/Triple.h"
14 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
15 #include "llvm/ExecutionEngine/Orc/CompileUtils.h"
16 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
17 #include "llvm/ExecutionEngine/Orc/IRCompileLayer.h"
18 #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
19 #include "llvm/IR/LLVMContext.h"
20 #include "llvm-c/OrcBindings.h"
21
22 namespace llvm {
23
24 class OrcCBindingsStack;
25
26 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(OrcCBindingsStack, LLVMOrcJITStackRef)
27 DEFINE_SIMPLE_CONVERSION_FUNCTIONS(TargetMachine, LLVMTargetMachineRef)
28
29 class OrcCBindingsStack {
30 public:
31
32   typedef orc::JITCompileCallbackManager CompileCallbackMgr;
33   typedef orc::ObjectLinkingLayer<> ObjLayerT;
34   typedef orc::IRCompileLayer<ObjLayerT> CompileLayerT;
35   typedef orc::CompileOnDemandLayer<CompileLayerT, CompileCallbackMgr> CODLayerT;
36
37   typedef std::function<std::unique_ptr<CompileCallbackMgr>()>
38     CallbackManagerBuilder;
39
40   typedef CODLayerT::IndirectStubsManagerBuilderT IndirectStubsManagerBuilder;
41
42 private:
43
44   class GenericHandle {
45   public:
46     virtual ~GenericHandle() {}
47     virtual orc::JITSymbol findSymbolIn(const std::string &Name,
48                                         bool ExportedSymbolsOnly) = 0;
49     virtual void removeModule() = 0;
50   };
51
52   template <typename LayerT>
53   class GenericHandleImpl : public GenericHandle {
54   public:
55     GenericHandleImpl(LayerT &Layer, typename LayerT::ModuleSetHandleT Handle)
56       : Layer(Layer), Handle(std::move(Handle)) {}
57
58     orc::JITSymbol findSymbolIn(const std::string &Name,
59                                 bool ExportedSymbolsOnly) override {
60       return Layer.findSymbolIn(Handle, Name, ExportedSymbolsOnly);
61     }
62
63     void removeModule() override {
64       return Layer.removeModuleSet(Handle);
65     }
66
67   private:
68     LayerT &Layer;
69     typename LayerT::ModuleSetHandleT Handle;
70   };
71
72   template <typename LayerT>
73   std::unique_ptr<GenericHandleImpl<LayerT>>
74   createGenericHandle(LayerT &Layer, typename LayerT::ModuleSetHandleT Handle) {
75     return llvm::make_unique<GenericHandleImpl<LayerT>>(Layer,
76                                                         std::move(Handle));
77   }
78
79 public:
80
81   // We need a 'ModuleSetHandleT' to conform to the layer concept.
82   typedef unsigned ModuleSetHandleT;
83
84   typedef unsigned ModuleHandleT;
85
86   static std::unique_ptr<CompileCallbackMgr> createCompileCallbackMgr(Triple T);
87   static IndirectStubsManagerBuilder createIndirectStubsMgrBuilder(Triple T);
88
89   OrcCBindingsStack(TargetMachine &TM,
90                     std::unique_ptr<CompileCallbackMgr> CCMgr, 
91                     IndirectStubsManagerBuilder IndirectStubsMgrBuilder)
92     : DL(TM.createDataLayout()), CCMgr(std::move(CCMgr)),
93       ObjectLayer(),
94       CompileLayer(ObjectLayer, orc::SimpleCompiler(TM)),
95       CODLayer(CompileLayer,
96                [](Function &F) { std::set<Function*> S; S.insert(&F); return S; },
97                *this->CCMgr, std::move(IndirectStubsMgrBuilder), false),
98       IndirectStubsMgr(IndirectStubsMgrBuilder()),
99       CXXRuntimeOverrides([this](const std::string &S) { return mangle(S); }) {}
100
101   ~OrcCBindingsStack() {
102     // Run any destructors registered with __cxa_atexit.
103     CXXRuntimeOverrides.runDestructors();
104     // Run any IR destructors.
105     for (auto &DtorRunner : IRStaticDestructorRunners)
106       DtorRunner.runViaLayer(*this);
107   }
108
109   std::string mangle(StringRef Name) {
110     std::string MangledName;
111     {
112       raw_string_ostream MangledNameStream(MangledName);
113       Mangler::getNameWithPrefix(MangledNameStream, Name, DL);
114     }
115     return MangledName;
116   }
117
118   template <typename PtrTy>
119   static PtrTy fromTargetAddress(orc::TargetAddress Addr) {
120     return reinterpret_cast<PtrTy>(static_cast<uintptr_t>(Addr));
121   }
122
123   orc::TargetAddress
124   createLazyCompileCallback(LLVMOrcLazyCompileCallbackFn Callback,
125                             void *CallbackCtx) {
126     auto CCInfo = CCMgr->getCompileCallback();
127     CCInfo.setCompileAction(
128       [=]() -> orc::TargetAddress {
129         return Callback(wrap(this), CallbackCtx);
130       });
131     return CCInfo.getAddress();
132   }
133
134   void createIndirectStub(StringRef StubName, orc::TargetAddress Addr) {
135     IndirectStubsMgr->createStub(StubName, Addr, JITSymbolFlags::Exported);
136   }
137
138   void setIndirectStubPointer(StringRef Name, orc::TargetAddress Addr) {
139     IndirectStubsMgr->updatePointer(Name, Addr);
140   }
141
142   std::shared_ptr<RuntimeDyld::SymbolResolver>
143   createResolver(LLVMOrcSymbolResolverFn ExternalResolver,
144                  void *ExternalResolverCtx) {
145     auto Resolver = orc::createLambdaResolver(
146       [this, ExternalResolver, ExternalResolverCtx](const std::string &Name) {
147         // Search order:
148         // 1. JIT'd symbols.
149         // 2. Runtime overrides.
150         // 3. External resolver (if present).
151
152         if (auto Sym = CODLayer.findSymbol(Name, true))
153           return RuntimeDyld::SymbolInfo(Sym.getAddress(),
154                                          Sym.getFlags());
155         if (auto Sym = CXXRuntimeOverrides.searchOverrides(Name))
156           return Sym;
157
158         if (ExternalResolver)
159           return RuntimeDyld::SymbolInfo(ExternalResolver(Name.c_str(),
160                                                           ExternalResolverCtx),
161                                          llvm::JITSymbolFlags::Exported);
162
163         return RuntimeDyld::SymbolInfo(nullptr);
164       },
165       [](const std::string &Name) {
166         return RuntimeDyld::SymbolInfo(nullptr);
167       }
168     );
169
170     return std::shared_ptr<RuntimeDyld::SymbolResolver>(std::move(Resolver));
171   }
172
173   template <typename LayerT>
174   ModuleHandleT addIRModule(LayerT &Layer,
175                             Module *M,
176                             std::unique_ptr<RuntimeDyld::MemoryManager> MemMgr,
177                             LLVMOrcSymbolResolverFn ExternalResolver,
178                             void *ExternalResolverCtx) {
179
180     // Attach a data-layout if one isn't already present.
181     if (M->getDataLayout().isDefault())
182       M->setDataLayout(DL);
183
184     // Record the static constructors and destructors. We have to do this before
185     // we hand over ownership of the module to the JIT.
186     std::vector<std::string> CtorNames, DtorNames;
187     for (auto Ctor : orc::getConstructors(*M))
188       CtorNames.push_back(mangle(Ctor.Func->getName()));
189     for (auto Dtor : orc::getDestructors(*M))
190       DtorNames.push_back(mangle(Dtor.Func->getName()));
191
192     // Create the resolver.
193     auto Resolver = createResolver(ExternalResolver, ExternalResolverCtx);
194
195     // Add the module to the JIT.
196     std::vector<Module*> S;
197     S.push_back(std::move(M));
198
199     auto LH = Layer.addModuleSet(std::move(S), std::move(MemMgr),
200                                  std::move(Resolver));
201     ModuleHandleT H = createHandle(Layer, LH);
202
203     // Run the static constructors, and save the static destructor runner for
204     // execution when the JIT is torn down.
205     orc::CtorDtorRunner<OrcCBindingsStack> CtorRunner(std::move(CtorNames), H);
206     CtorRunner.runViaLayer(*this);
207
208     IRStaticDestructorRunners.emplace_back(std::move(DtorNames), H);
209
210     return H;
211   }
212
213   ModuleHandleT addIRModuleEager(Module* M,
214                                  LLVMOrcSymbolResolverFn ExternalResolver,
215                                  void *ExternalResolverCtx) {
216     return addIRModule(CompileLayer, std::move(M),
217                        llvm::make_unique<SectionMemoryManager>(),
218                        std::move(ExternalResolver), ExternalResolverCtx);
219   }
220
221   ModuleHandleT addIRModuleLazy(Module* M,
222                                 LLVMOrcSymbolResolverFn ExternalResolver,
223                                 void *ExternalResolverCtx) {
224     return addIRModule(CODLayer, std::move(M),
225                        llvm::make_unique<SectionMemoryManager>(),
226                        std::move(ExternalResolver), ExternalResolverCtx);
227   }
228
229   void removeModule(ModuleHandleT H) {
230     GenericHandles[H]->removeModule();
231     GenericHandles[H] = nullptr;
232     FreeHandleIndexes.push_back(H);
233   }
234
235   orc::JITSymbol findSymbol(const std::string &Name, bool ExportedSymbolsOnly) {
236     if (auto Sym = IndirectStubsMgr->findStub(Name, ExportedSymbolsOnly))
237       return Sym;
238     return CODLayer.findSymbol(mangle(Name), ExportedSymbolsOnly);
239   }
240
241   orc::JITSymbol findSymbolIn(ModuleHandleT H, const std::string &Name,
242                               bool ExportedSymbolsOnly) {
243     return GenericHandles[H]->findSymbolIn(Name, ExportedSymbolsOnly);
244   }
245
246 private:
247
248   template <typename LayerT>
249   unsigned createHandle(LayerT &Layer,
250                         typename LayerT::ModuleSetHandleT Handle) {
251     unsigned NewHandle;
252     if (!FreeHandleIndexes.empty()) {
253       NewHandle = FreeHandleIndexes.back();
254       FreeHandleIndexes.pop_back();
255       GenericHandles[NewHandle] = createGenericHandle(Layer, std::move(Handle));
256       return NewHandle;
257     } else {
258       NewHandle = GenericHandles.size();
259       GenericHandles.push_back(createGenericHandle(Layer, std::move(Handle)));
260     }
261     return NewHandle;
262   }
263
264   DataLayout DL;
265   SectionMemoryManager CCMgrMemMgr;
266
267   std::unique_ptr<CompileCallbackMgr> CCMgr;
268   ObjLayerT ObjectLayer;
269   CompileLayerT CompileLayer;
270   CODLayerT CODLayer;
271
272   std::unique_ptr<orc::IndirectStubsManager> IndirectStubsMgr;
273
274   std::vector<std::unique_ptr<GenericHandle>> GenericHandles;
275   std::vector<unsigned> FreeHandleIndexes;
276
277   orc::LocalCXXRuntimeOverrides CXXRuntimeOverrides;
278   std::vector<orc::CtorDtorRunner<OrcCBindingsStack>> IRStaticDestructorRunners;
279 };
280
281 } // end namespace llvm
282
283 #endif // LLVM_LIB_EXECUTIONENGINE_ORC_ORCCBINDINGSSTACK_H