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 
DEFINE_SIMPLE_CONVERSION_FUNCTIONS(OrcCBindingsStack,LLVMOrcJITStackRef)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), nullptr,
225                        std::move(ExternalResolver), ExternalResolverCtx);
226   }
227 
228   void removeModule(ModuleHandleT H) {
229     GenericHandles[H]->removeModule();
230     GenericHandles[H] = nullptr;
231     FreeHandleIndexes.push_back(H);
232   }
233 
234   orc::JITSymbol findSymbol(const std::string &Name, bool ExportedSymbolsOnly) {
235     if (auto Sym = IndirectStubsMgr->findStub(Name, ExportedSymbolsOnly))
236       return Sym;
237     return CODLayer.findSymbol(mangle(Name), ExportedSymbolsOnly);
238   }
239 
240   orc::JITSymbol findSymbolIn(ModuleHandleT H, const std::string &Name,
241                               bool ExportedSymbolsOnly) {
242     return GenericHandles[H]->findSymbolIn(Name, ExportedSymbolsOnly);
243   }
244 
245 private:
246 
247   template <typename LayerT>
248   unsigned createHandle(LayerT &Layer,
249                         typename LayerT::ModuleSetHandleT Handle) {
250     unsigned NewHandle;
251     if (!FreeHandleIndexes.empty()) {
252       NewHandle = FreeHandleIndexes.back();
253       FreeHandleIndexes.pop_back();
254       GenericHandles[NewHandle] = createGenericHandle(Layer, std::move(Handle));
255       return NewHandle;
256     } else {
257       NewHandle = GenericHandles.size();
258       GenericHandles.push_back(createGenericHandle(Layer, std::move(Handle)));
259     }
260     return NewHandle;
261   }
262 
263   DataLayout DL;
264   SectionMemoryManager CCMgrMemMgr;
265 
266   std::unique_ptr<CompileCallbackMgr> CCMgr;
267   ObjLayerT ObjectLayer;
268   CompileLayerT CompileLayer;
269   CODLayerT CODLayer;
270 
271   std::unique_ptr<orc::IndirectStubsManager> IndirectStubsMgr;
272 
273   std::vector<std::unique_ptr<GenericHandle>> GenericHandles;
274   std::vector<unsigned> FreeHandleIndexes;
275 
276   orc::LocalCXXRuntimeOverrides CXXRuntimeOverrides;
277   std::vector<orc::CtorDtorRunner<OrcCBindingsStack>> IRStaticDestructorRunners;
278 };
279 
280 } // end namespace llvm
281 
282 #endif // LLVM_LIB_EXECUTIONENGINE_ORC_ORCCBINDINGSSTACK_H
283