JITLinkReentryTrampolines.cpp

Go to the documentation of this file.

//===----- JITLinkReentryTrampolines.cpp -- JITLink-based trampoline- -----===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
 
#include "llvm/ExecutionEngine/Orc/JITLinkReentryTrampolines.h"
 
#include "llvm/ExecutionEngine/JITLink/aarch64.h"
#include "llvm/ExecutionEngine/JITLink/x86_64.h"
#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
 
#include <memory>
 
#define DEBUG_TYPE "orc"
 
using namespace llvm;
using namespace llvm::jitlink;
 
namespace {
constexpr StringRef ReentryFnName = "__orc_rt_reenter";
constexpr StringRef ReentrySectionName = "__orc_stubs";
} // namespace
 
namespace llvm::orc {
 
class JITLinkReentryTrampolines::TrampolineAddrScraperPlugin
    : public ObjectLinkingLayer::Plugin {
public:
  void modifyPassConfig(MaterializationResponsibility &MR,
                        jitlink::LinkGraph &G,
                        jitlink::PassConfiguration &Config) override {
    Config.PreFixupPasses.push_back(
        [this](LinkGraph &G) { return recordTrampolineAddrs(G); });
  }
 
  Error notifyFailed(MaterializationResponsibility &MR) override {
    return Error::success();
  }
 
  Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override {
    return Error::success();
  }
 
  void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
                                   ResourceKey SrcKey) override {}
 
  void registerGraph(LinkGraph &G,
                     std::shared_ptr<std::vector<ExecutorSymbolDef>> Addrs) {
    std::lock_guard<std::mutex> Lock(M);
    assert(!PendingAddrs.count(&G) && "Duplicate registration");
    PendingAddrs[&G] = std::move(Addrs);
  }
 
  Error recordTrampolineAddrs(LinkGraph &G) {
    std::shared_ptr<std::vector<ExecutorSymbolDef>> Addrs;
    {
      std::lock_guard<std::mutex> Lock(M);
      auto I = PendingAddrs.find(&G);
      if (I == PendingAddrs.end())
        return Error::success();
      Addrs = std::move(I->second);
      PendingAddrs.erase(I);
    }
 
    auto *Sec = G.findSectionByName(ReentrySectionName);
    assert(Sec && "Reentry graph missing reentry section");
    assert(!Sec->empty() && "Reentry graph is empty");
 
    for (auto *Sym : Sec->symbols())
      if (!Sym->hasName())
        Addrs->push_back({Sym->getAddress(), JITSymbolFlags()});
 
    return Error::success();
  }
 
private:
  std::mutex M;
  DenseMap<LinkGraph *, std::shared_ptr<std::vector<ExecutorSymbolDef>>>
      PendingAddrs;
};
 
Expected<std::unique_ptr<JITLinkReentryTrampolines>>
JITLinkReentryTrampolines::Create(ObjectLinkingLayer &ObjLinkingLayer) {
 
  EmitTrampolineFn EmitTrampoline;
 
  const auto &TT = ObjLinkingLayer.getExecutionSession().getTargetTriple();
  switch (TT.getArch()) {
  case Triple::aarch64:
    EmitTrampoline = aarch64::createAnonymousReentryTrampoline;
    break;
  case Triple::x86_64:
    EmitTrampoline = x86_64::createAnonymousReentryTrampoline;
    break;
  default:
    return make_error<StringError>("JITLinkReentryTrampolines: architecture " +
           TT.getArchName() + " not supported",
                                   inconvertibleErrorCode());
  }
 
  return std::make_unique<JITLinkReentryTrampolines>(ObjLinkingLayer,
                                                     std::move(EmitTrampoline));
}
 
JITLinkReentryTrampolines::JITLinkReentryTrampolines(
    ObjectLinkingLayer &ObjLinkingLayer, EmitTrampolineFn EmitTrampoline)
    : ObjLinkingLayer(ObjLinkingLayer),
      EmitTrampoline(std::move(EmitTrampoline)) {
  auto TAS = std::make_shared<TrampolineAddrScraperPlugin>();
  TrampolineAddrScraper = TAS.get();
  ObjLinkingLayer.addPlugin(std::move(TAS));
}
 
void JITLinkReentryTrampolines::emit(ResourceTrackerSP RT,
                                     size_t NumTrampolines,
                                     OnTrampolinesReadyFn OnTrampolinesReady) {
 
  if (NumTrampolines == 0)
    return OnTrampolinesReady(std::vector<ExecutorSymbolDef>());
 
  JITDylibSP JD(&RT->getJITDylib());
  auto &ES = ObjLinkingLayer.getExecutionSession();
 
  auto ReentryGraphSym =
      ES.intern(("__orc_reentry_graph_#" + Twine(++ReentryGraphIdx)).str());
 
  auto G = std::make_unique<jitlink::LinkGraph>(
      (*ReentryGraphSym).str(), ES.getSymbolStringPool(), ES.getTargetTriple(),
      SubtargetFeatures(), jitlink::getGenericEdgeKindName);
 
  auto &ReentryFnSym = G->addExternalSymbol(ReentryFnName, 0, false);
 
  auto &ReentrySection =
      G->createSection(ReentrySectionName, MemProt::Exec | MemProt::Read);
 
  for (size_t I = 0; I != NumTrampolines; ++I)
    EmitTrampoline(*G, ReentrySection, ReentryFnSym).setLive(true);
 
  auto &FirstBlock = **ReentrySection.blocks().begin();
  G->addDefinedSymbol(FirstBlock, 0, *ReentryGraphSym, FirstBlock.getSize(),
                      Linkage::Strong, Scope::SideEffectsOnly, true, true);
 
  auto TrampolineAddrs = std::make_shared<std::vector<ExecutorSymbolDef>>();
  TrampolineAddrScraper->registerGraph(*G, TrampolineAddrs);
 
  // Add Graph via object linking layer.
  if (auto Err = ObjLinkingLayer.add(std::move(RT), std::move(G)))
    return OnTrampolinesReady(std::move(Err));
 
  // Trigger graph emission.
  ES.lookup(
      LookupKind::Static, {{JD.get(), JITDylibLookupFlags::MatchAllSymbols}},
      SymbolLookupSet(ReentryGraphSym,
                      SymbolLookupFlags::WeaklyReferencedSymbol),
      SymbolState::Ready,
      [OnTrampolinesReady = std::move(OnTrampolinesReady),
       TrampolineAddrs =
           std::move(TrampolineAddrs)](Expected<SymbolMap> Result) mutable {
        if (Result)
          OnTrampolinesReady(std::move(*TrampolineAddrs));
        else
          OnTrampolinesReady(Result.takeError());
      },
      NoDependenciesToRegister);
}
 
Expected<std::unique_ptr<LazyReexportsManager>>
createJITLinkLazyReexportsManager(ObjectLinkingLayer &ObjLinkingLayer,
                                  RedirectableSymbolManager &RSMgr,
                                  JITDylib &PlatformJD,
                                  LazyReexportsManager::Listener *L) {
  auto JLT = JITLinkReentryTrampolines::Create(ObjLinkingLayer);
  if (!JLT)
    return JLT.takeError();
 
  return LazyReexportsManager::Create(
      [JLT = std::move(*JLT)](ResourceTrackerSP RT, size_t NumTrampolines,
                              LazyReexportsManager::OnTrampolinesReadyFn
                                  OnTrampolinesReady) mutable {
        JLT->emit(std::move(RT), NumTrampolines, std::move(OnTrampolinesReady));
      },
      RSMgr, PlatformJD, L);
}
 
} // namespace llvm::orc