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), 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