content_description.hpp

 
#pragma once
#include <numeric>
#include <type_traits>
#include <utility>
 
#include <rapidsmpf/memory/buffer.hpp>
#include <rapidsmpf/memory/memory_type.hpp>
 
namespace rapidsmpf {
 
class ContentDescription {
  public:
    enum class Spillable : bool {
        NO,
        YES
    };
 
    template <
        std::ranges::input_range Range =
            std::initializer_list<std::pair<MemoryType, std::size_t>>>
        requires std::convertible_to<
            std::ranges::range_value_t<Range>,
            std::pair<MemoryType, std::size_t>>
    constexpr explicit ContentDescription(Range&& sizes, Spillable spillable)
        : spillable_(spillable == Spillable::YES) {
        content_sizes_.fill(0);
        for (auto&& [mem_type, size] : sizes) {
            auto idx = static_cast<std::size_t>(mem_type);
            if (idx < content_sizes_.size()) {
                content_sizes_[idx] = size;
            }
        }
    }
 
    constexpr ContentDescription(Spillable spillable = Spillable::NO)
        : ContentDescription({{}}, spillable) {}
 
    [[nodiscard]] constexpr std::size_t& content_size(MemoryType mem_type) noexcept {
        return content_sizes_[static_cast<std::size_t>(mem_type)];
    }
 
    [[nodiscard]] constexpr std::size_t content_size(MemoryType mem_type) const noexcept {
        return content_sizes_[static_cast<std::size_t>(mem_type)];
    }
 
    [[nodiscard]] constexpr MemoryType principal_memory_type() const noexcept {
        for (auto mem_type : MEMORY_TYPES) {
            if (content_size(mem_type) > 0) {
                return mem_type;
            }
        }
        return MemoryType::DEVICE;
    }
 
    [[nodiscard]] constexpr std::size_t content_size() const noexcept {
        return std::accumulate(
            content_sizes_.begin(), content_sizes_.end(), std::size_t{0}
        );
    }
 
    [[nodiscard]] constexpr bool spillable() const noexcept {
        return spillable_;
    }
 
    [[nodiscard]] constexpr bool operator==(
        ContentDescription const& other
    ) const noexcept {
        return spillable_ == other.spillable_ && content_sizes_ == other.content_sizes_;
    }
 
  private:
    std::array<std::size_t, MEMORY_TYPES.size()> content_sizes_ = {};
    bool spillable_;
};
 
static_assert(
    std::is_trivially_copyable_v<ContentDescription>,
    "ContentDescription must be trivially copyable"
);
 
}  // namespace rapidsmpf