memory.cpp

Go to the documentation of this file.

/**
 * @file
 * @ingroup numapp_mem
 * @copyright ESO 2023 - European Southern Observatory
 *
 * @brief Definition of memory functions from <numapp/memory.hpp>
 */
#include <numapp/memory.hpp>
 
#include <bit>
#include <cassert>
#include <numa.h>
#include <type_traits>
#include <utility>
#include <x86intrin.h>
 
#include <numapp/mempolicy.hpp>
 
namespace numapp {
namespace {
/**
 * Rounds allocation size up to nearest page size. `page_size` must be a power of two.
 */
constexpr auto AlignedCeil(std::size_t size, std::size_t page_size) noexcept -> std::size_t {
    if (size == 0) {
        std::unreachable();
    }
    if (page_size == 0) {
        std::unreachable();
    }
    if ((page_size & (page_size - 1)) != 0) {
        // Page size is power of 2
        std::unreachable();
    }
 
    return ((size + page_size - 1) / page_size) * page_size;
}
}  // namespace
 
std::size_t GetPageSize() noexcept {
    return numa_pagesize();
}
 
int GetNumNodes() noexcept {
    return numa_num_configured_nodes();
}
 
std::optional<int> GetNodeDistance(int node1, int node2) noexcept {
    int distance = numa_distance(node1, node2);
    return distance == 0 ? std::nullopt : std::optional<int>{distance};
}
 
std::optional<int> GetNodeOfCpu(int cpu) noexcept {
    int node = numa_node_of_cpu(cpu);
    return node == -1 ? std::nullopt : std::optional<int>{node};
}
 
std::error_code MemLock(void const* addr, std::size_t len, LockFlag flag) noexcept {
    if (mlock2(addr, len, static_cast<std::underlying_type<LockFlag>::type>(flag)) != 0) {
        return std::make_error_code(static_cast<std::errc>(errno));
    }
    return {};
}
 
std::error_code MemUnlock(void const* addr, std::size_t len) noexcept {
    if (munlock(addr, len) != 0) {
        return std::make_error_code(static_cast<std::errc>(errno));
    }
    return {};
}
 
std::error_code MemLockAll(LockAllFlag flags) noexcept {
    if (mlockall(static_cast<std::underlying_type<LockAllFlag>::type>(flags)) != 0) {
        return std::make_error_code(static_cast<std::errc>(errno));
    }
    return {};
}
 
std::error_code MemUnlockAll() noexcept {
    if (munlockall() != 0) {
        return std::make_error_code(static_cast<std::errc>(errno));
    }
    return {};
}
 
[[nodiscard]] void*
Allocate(std::size_t size, MemPolicy const& policy, std::error_code& ec) noexcept {
    return Allocate(size, policy, MemPolicyFlag::None, ec);
}
 
[[nodiscard]] void* Allocate(std::size_t size,
                             MemPolicy const& policy,
                             MemPolicyFlag flags,
                             std::error_code& ec) noexcept {
    auto ptr = mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, 0, 0);
    if (ptr == MAP_FAILED) {  // NOLINT
        ec = std::make_error_code(static_cast<std::errc>(errno));
        return nullptr;
    }
 
    ec = Apply(ptr, size, policy, flags);
    if (ec) {
        (void)munmap(ptr, size);
        return nullptr;
    }
    return ptr;
}
 
[[nodiscard]] void* Allocate(std::size_t size, MemPolicy const& policy) {
    std::error_code ec;
    auto ptr = Allocate(size, policy, ec);
    if (ec) {
        assert(ptr == nullptr);
        throw std::system_error(ec);
    }
    return ptr;
}
 
[[nodiscard]] void* Allocate(std::size_t size, MemPolicy const& policy, MemPolicyFlag flags) {
    std::error_code ec;
    auto ptr = Allocate(size, policy, flags, ec);
    if (ec) {
        assert(ptr == nullptr);
        throw std::system_error(ec);
    }
    return ptr;
}
 
void Free(void* ptr, std::size_t size, std::error_code& ec) noexcept {
    if (munmap(ptr, size) == -1) {
        ec = std::make_error_code(static_cast<std::errc>(errno));
    } else {
        ec.clear();
    }
}
 
void Free(void* ptr, std::size_t size) {
    std::error_code ec;
    Free(ptr, size, ec);
    if (ec) {
        throw std::system_error(ec);
    }
}
 
HugePageSize::HugePageSize(HugePagePreset preset) noexcept
    : m_size(static_cast<std::size_t>(preset)) {
}
 
HugePageSize::HugePageSize(std::size_t bytes) : m_size(bytes) {
    // Maintain invariant that `bytes` is an integer power of two and that base-2 log is not 0,
    // other than that you can technically request surprising page sizes with `mmap`.
    if (m_size < 2u) {
        throw std::invalid_argument("Invalid page size");
    }
    if (!std::has_single_bit(m_size)) {
        throw std::invalid_argument("Invalid page size");
    }
}
 
auto EncodeMmapFlags(HugePageSize page_size) noexcept -> int {
    // HugePageSize maintain the invariant that the size is a round base-2 log we can simplify to:
    return std::countr_zero(page_size.Size()) << MAP_HUGE_SHIFT;
}
 
void* AllocateHuge(std::size_t size,
                   HugePageSize page_size,
                   MemPolicy const& policy,
                   MemPolicyFlag flags,
                   std::error_code& ec) noexcept {
    auto const tlb_bits = EncodeMmapFlags(page_size);
    auto alloc_size = AlignedCeil(size, page_size.Size());
    // note: Although freeing huge pages with `munmap()` requires length to be multiple of page
    // size, `mmap()` does not.
    auto ptr = mmap(nullptr,
                    alloc_size,
                    PROT_READ | PROT_WRITE,
                    MAP_PRIVATE | MAP_ANONYMOUS | MAP_HUGETLB | tlb_bits,
                    0,
                    0);
    if (ptr == MAP_FAILED) {  // NOLINT
        ec = std::make_error_code(static_cast<std::errc>(errno));
        return nullptr;
    }
 
    ec = Apply(ptr, alloc_size, policy, flags);
    if (ec) {
        (void)munmap(ptr, alloc_size);
        return nullptr;
    }
    return ptr;
}
 
void* AllocateHuge(std::size_t size,
                   HugePageSize page_size,
                   MemPolicy const& policy,
                   MemPolicyFlag flags) {
    std::error_code ec;
    auto* ptr = AllocateHuge(size, page_size, policy, flags, ec);
    if (ec) {
        throw std::system_error(ec, "AllocateHuge failed");
    }
    return ptr;
}
 
void FreeHuge(void* ptr, std::size_t size, HugePageSize page_size, std::error_code& ec) noexcept {
    // note: `munmap()` documents that for huge pages the size must be a multiple of the page size.
    auto const alloc_size = AlignedCeil(size, page_size.Size());
    if (munmap(ptr, alloc_size) == -1) {
        ec = std::make_error_code(static_cast<std::errc>(errno));
    } else {
        ec.clear();
    }
}
 
void FreeHuge(void* ptr, std::size_t size, HugePageSize page_size) {
    std::error_code ec;
    FreeHuge(ptr, size, page_size, ec);
    if (ec) {
        throw std::system_error(ec, "FreeHuge failed");
    }
}
 
PageResource::PageResource(MemPolicy policy, MemPolicyFlag flags)
    : m_policy(std::move(policy)), m_flags(flags) {
}
 
auto PageResource::GetPolicy() const noexcept -> MemPolicy const& {
    return m_policy;
}
 
auto PageResource::GetFlags() const noexcept -> MemPolicyFlag {
    return m_flags;
}
 
void* PageResource::do_allocate(std::size_t bytes, std::size_t alignment) {
    return Allocate(bytes, m_policy, m_flags);
}
 
void PageResource::do_deallocate(void* p, std::size_t bytes, std::size_t alignment) {
    return Free(p, bytes);
}
 
bool PageResource::do_is_equal(const std::pmr::memory_resource& other) const noexcept {
    if (auto const* o = dynamic_cast<PageResource const*>(&other); o != nullptr) {
        return GetFlags() == o->GetFlags() && GetPolicy() == o->GetPolicy();
    }
    return false;
}
 
HugePageResource::HugePageResource(HugePageSize page_size, MemPolicy policy, MemPolicyFlag flags)
    : m_page_size(page_size), m_policy(std::move(policy)), m_flags(flags) {
}
 
auto HugePageResource::GetPolicy() const noexcept -> MemPolicy const& {
    return m_policy;
}
 
auto HugePageResource::GetFlags() const noexcept -> MemPolicyFlag {
    return m_flags;
}
 
auto HugePageResource::GetPageSize() const noexcept -> HugePageSize {
    return m_page_size;
}
 
void* HugePageResource::do_allocate(std::size_t bytes, std::size_t alignment) {
    return AllocateHuge(bytes, m_page_size, m_policy, m_flags);
}
 
void HugePageResource::do_deallocate(void* p, std::size_t bytes, std::size_t alignment) {
    return FreeHuge(p, bytes, m_page_size);
}
 
bool HugePageResource::do_is_equal(const std::pmr::memory_resource& other) const noexcept {
    if (auto const* o = dynamic_cast<HugePageResource const*>(&other); o != nullptr) {
        return GetFlags() == o->GetFlags() && GetPolicy() == o->GetPolicy() &&
               GetPageSize() == o->GetPageSize();
    }
    return false;
}
 
LockResource::LockResource(LockFlag flag) noexcept
    : m_upstream(std::pmr::get_default_resource()), m_flag(flag) {
}
 
LockResource::LockResource(std::pmr::memory_resource* upstream) noexcept
    : m_upstream(upstream), m_flag(LockFlag::PreFault) {
    assert(upstream != nullptr);
}
 
LockResource::LockResource(LockFlag flag, std::pmr::memory_resource* upstream) noexcept
    : m_upstream(upstream), m_flag(flag) {
    assert(upstream != nullptr);
}
 
auto LockResource::do_allocate(std::size_t bytes, std::size_t alignment) -> void* {
    auto* mem = m_upstream->allocate(bytes, alignment);
    assert(mem != nullptr);
    auto ec = MemLock(mem, bytes, m_flag);
    if (ec) {
        m_upstream->deallocate(mem, bytes, alignment);
        throw std::system_error(ec, "MemLock failed");
    }
    return mem;
}
 
void LockResource::do_deallocate(void* p, std::size_t bytes, std::size_t alignment) {
    m_upstream->deallocate(p, bytes, alignment);
}
 
auto LockResource::do_is_equal(const std::pmr::memory_resource& other) const noexcept -> bool {
    if (auto const* o = dynamic_cast<LockResource const*>(&other); o != nullptr) {
        return GetFlags() == o->GetFlags() && GetUpstream() == o->GetUpstream();
    }
    return false;
}
 
auto LockResource::GetFlags() const noexcept -> LockFlag {
    return m_flag;
}
 
auto LockResource::GetUpstream() const noexcept -> std::pmr::memory_resource* {
    return m_upstream;
}
 
}  // namespace numapp