memory_type.hpp

 
#pragma once
 
#include <algorithm>
#include <array>
#include <iostream>
#include <ranges>
#include <span>
 
namespace rapidsmpf {
 
enum class MemoryType : int {
    DEVICE = 0,  
    PINNED_HOST = 1,  
    HOST = 2  
};
 
constexpr std::array<MemoryType, 3> MEMORY_TYPES{
    {MemoryType::DEVICE, MemoryType::PINNED_HOST, MemoryType::HOST}
};
 
constexpr std::array<char const*, MEMORY_TYPES.size()> MEMORY_TYPE_NAMES{
    {"DEVICE", "PINNED_HOST", "HOST"}
};
 
constexpr std::array<MemoryType, 2> SPILL_TARGET_MEMORY_TYPES{
    {MemoryType::PINNED_HOST, MemoryType::HOST}
};
 
constexpr std::span<MemoryType const> leq_memory_types(MemoryType mem_type) noexcept {
    return std::views::drop_while(MEMORY_TYPES, [&](MemoryType const& mt) {
        return mt != mem_type;
    });
}
 
static_assert(std::ranges::equal(leq_memory_types(MemoryType::DEVICE), MEMORY_TYPES));
static_assert(std::ranges::equal(
    leq_memory_types(MemoryType::HOST), std::ranges::single_view{MemoryType::HOST}
));
// unknown memory type should return an empty view
static_assert(std::ranges::equal(
    leq_memory_types(static_cast<MemoryType>(-1)), std::ranges::empty_view<MemoryType>{}
));
 
constexpr char const* to_string(MemoryType mem_type) {
    return MEMORY_TYPE_NAMES[static_cast<std::size_t>(mem_type)];
}
 
std::ostream& operator<<(std::ostream& os, MemoryType mem_type);
 
std::istream& operator>>(std::istream& is, MemoryType& out);
 
}  // namespace rapidsmpf