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