EPCGenericJITLinkMemoryManager.cpp

Go to the documentation of this file.

//===---- EPCGenericJITLinkMemoryManager.cpp -- Mem management via EPC ----===//
//
// 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/EPCGenericJITLinkMemoryManager.h"
 
#include "llvm/ExecutionEngine/JITLink/JITLink.h"
#include "llvm/ExecutionEngine/Orc/LookupAndRecordAddrs.h"
#include "llvm/ExecutionEngine/Orc/RTBridge/SPS/GenericMemoryManagerProxySpecs.h"
#include "llvm/ExecutionEngine/Orc/Shared/SPSCI/SimpleNativeMemoryMapSPSCI.h"
 
#include <limits>
 
using namespace llvm::jitlink;
 
namespace llvm {
namespace orc {
 
class EPCGenericJITLinkMemoryManager::InFlightAlloc
    : public jitlink::JITLinkMemoryManager::InFlightAlloc {
public:
 
  // FIXME: The C++98 initializer is an attempt to work around compile failures
  // due to http://www.open-std.org/jtc1/sc22/wg21/docs/cwg_defects.html#1397.
  // We should be able to switch this back to member initialization once that
  // issue is fixed.
  struct SegInfo {
    SegInfo() : WorkingMem(nullptr), ContentSize(0), ZeroFillSize(0) {}
 
    char *WorkingMem;
    ExecutorAddr Addr;
    uint64_t ContentSize;
    uint64_t ZeroFillSize;
  };
 
  using SegInfoMap = AllocGroupSmallMap<SegInfo>;
 
  InFlightAlloc(EPCGenericJITLinkMemoryManager &Parent, LinkGraph &G,
                ExecutorAddr AllocAddr, SegInfoMap Segs)
      : Parent(Parent), G(G), AllocAddr(AllocAddr), Segs(std::move(Segs)) {}
 
  void finalize(OnFinalizedFunction OnFinalize) override {
    tpctypes::FinalizeRequest FR;
    for (auto &KV : Segs) {
      assert(KV.second.ContentSize <= std::numeric_limits<size_t>::max());
      FR.Segments.push_back(tpctypes::SegFinalizeRequest{
          KV.first,
          KV.second.Addr,
          alignTo(KV.second.ContentSize + KV.second.ZeroFillSize,
                  Parent.ES.getExecutorProcessControl().getPageSize()),
          {KV.second.WorkingMem, static_cast<size_t>(KV.second.ContentSize)}});
    }
 
    // Transfer allocation actions.
    std::swap(FR.Actions, G.allocActions());
 
    Parent.B.Initialize(
        [OnFinalize = std::move(OnFinalize), AllocAddr = this->AllocAddr](
            Expected<ExecutorAddr> InitializeKey) mutable {
          // FIXME: Release abandoned alloc.
          if (!InitializeKey)
            return OnFinalize(InitializeKey.takeError());
          OnFinalize(FinalizedAlloc(AllocAddr));
        },
        Parent.ES, Parent.B.Instance, FR);
  }
 
  void abandon(OnAbandonedFunction OnAbandoned) override {
    // FIXME: Return memory to pool instead.
    Parent.B.Release(std::move(OnAbandoned), Parent.ES, Parent.B.Instance,
                     ArrayRef<ExecutorAddr>(AllocAddr));
  }
 
private:
  EPCGenericJITLinkMemoryManager &Parent;
  LinkGraph &G;
  ExecutorAddr AllocAddr;
  SegInfoMap Segs;
};
 
Expected<std::unique_ptr<EPCGenericJITLinkMemoryManager>>
EPCGenericJITLinkMemoryManager::Create(JITDylib &JD) {
  namespace sps = rt::sps;
  namespace sps_ci = rt::sps_ci;
  auto &ES = JD.getExecutionSession();
  Bindings B;
  // Instance is the executor-side allocator object -- a data symbol passed as
  // the first argument to each call, not a wrapper to proxy.
  if (auto Err = lookupAndRecordAddrs(
          ES, LookupKind::Static, makeJITDylibSearchOrder({&JD}),
          {{ES.intern(sps_ci::SimpleNativeMemoryMapInstanceName),
            &B.Instance}}))
    return std::move(Err);
  // The proxies resolve to the specs' default (SimpleNativeMemoryMap) names.
  if (auto Err = rt::buildProxies(
          JD, rt::proxyInit<sps::MemMgrReserveProxySpec>(&B.Reserve),
          rt::proxyInit<sps::MemMgrInitializeProxySpec>(&B.Initialize),
          rt::proxyInit<sps::MemMgrDeinitializeProxySpec>(&B.Deinitialize),
          rt::proxyInit<sps::MemMgrReleaseProxySpec>(&B.Release)))
    return std::move(Err);
  return std::make_unique<EPCGenericJITLinkMemoryManager>(ES, std::move(B));
}
 
Expected<std::unique_ptr<EPCGenericJITLinkMemoryManager>>
EPCGenericJITLinkMemoryManager::Create(ExecutionSession &ES) {
  return Create(ES.getBootstrapJITDylib());
}
 
void EPCGenericJITLinkMemoryManager::allocate(const JITLinkDylib *JD,
                                              LinkGraph &G,
                                              OnAllocatedFunction OnAllocated) {
  BasicLayout BL(G);
 
  auto Pages = BL.getContiguousPageBasedLayoutSizes(
      ES.getExecutorProcessControl().getPageSize());
  if (!Pages)
    return OnAllocated(Pages.takeError());
 
  B.Reserve(
      [this, BL = std::move(BL), OnAllocated = std::move(OnAllocated)](
          Expected<ExecutorAddr> AllocAddr) mutable {
        if (!AllocAddr)
          return OnAllocated(AllocAddr.takeError());
        completeAllocation(*AllocAddr, std::move(BL), std::move(OnAllocated));
      },
      ES, B.Instance, Pages->total());
}
 
void EPCGenericJITLinkMemoryManager::deallocate(
    std::vector<FinalizedAlloc> Allocs, OnDeallocatedFunction OnDeallocated) {
  std::vector<ExecutorAddr> Bases;
  Bases.reserve(Allocs.size());
  for (auto &A : Allocs)
    Bases.push_back(ExecutorAddr(A.getAddress()));
  B.Release(std::move(OnDeallocated), ES, B.Instance, Bases);
  for (auto &A : Allocs)
    A.release();
}
 
void EPCGenericJITLinkMemoryManager::completeAllocation(
    ExecutorAddr AllocAddr, BasicLayout BL, OnAllocatedFunction OnAllocated) {
 
  InFlightAlloc::SegInfoMap SegInfos;
 
  ExecutorAddr NextSegAddr = AllocAddr;
  for (auto &KV : BL.segments()) {
    const auto &AG = KV.first;
    auto &Seg = KV.second;
 
    Seg.Addr = NextSegAddr;
    KV.second.WorkingMem = BL.getGraph().allocateBuffer(Seg.ContentSize).data();
    NextSegAddr +=
        ExecutorAddrDiff(alignTo(Seg.ContentSize + Seg.ZeroFillSize,
                                 ES.getExecutorProcessControl().getPageSize()));
 
    auto &SegInfo = SegInfos[AG];
    SegInfo.ContentSize = Seg.ContentSize;
    SegInfo.ZeroFillSize = Seg.ZeroFillSize;
    SegInfo.Addr = Seg.Addr;
    SegInfo.WorkingMem = Seg.WorkingMem;
  }
 
  if (auto Err = BL.apply())
    return OnAllocated(std::move(Err));
 
  OnAllocated(std::make_unique<InFlightAlloc>(*this, BL.getGraph(), AllocAddr,
                                              std::move(SegInfos)));
}
 
} // end namespace orc
} // end namespace llvm