utils.hpp

 
#pragma once
 
#include <cstdint>
#include <limits>
#include <memory>
#include <vector>
 
#include <rapidsmpf/communicator/communicator.hpp>
#include <rapidsmpf/memory/buffer.hpp>
#include <rapidsmpf/memory/buffer_resource.hpp>
#include <rapidsmpf/memory/packed_data.hpp>
 
namespace rapidsmpf::coll::detail {
 
using ChunkID = std::uint64_t;
 
class Chunk {
  private:
    ChunkID id_;  
    Rank destination_;  
    std::unique_ptr<std::vector<std::uint8_t>> metadata_;  
    std::unique_ptr<Buffer> data_;  
    std::uint64_t
        data_size_;  
 
    Chunk(
        ChunkID id,
        Rank destination,
        std::unique_ptr<std::vector<std::uint8_t>> metadata,
        std::unique_ptr<Buffer> data
    );
 
    Chunk(ChunkID id, Rank destination);
 
  public:
    static constexpr Rank INVALID_RANK = std::numeric_limits<Rank>::max();
 
    [[nodiscard]] bool is_ready() const noexcept;
 
    [[nodiscard]] MemoryType memory_type() const noexcept;
 
    [[nodiscard]] bool is_finish() const noexcept;
 
    [[nodiscard]] ChunkID id() const noexcept;
 
    [[nodiscard]] ChunkID sequence() const noexcept;
 
    [[nodiscard]] Rank origin() const noexcept;
 
    [[nodiscard]] Rank destination() const noexcept;
 
    [[nodiscard]] std::uint64_t data_size() const noexcept;
 
    [[nodiscard]] std::uint64_t metadata_size() const noexcept;
 
    [[nodiscard]] static std::unique_ptr<Chunk> from_packed_data(
        std::uint64_t sequence, Rank origin, Rank destination, PackedData&& packed_data
    );
 
    [[nodiscard]] static std::unique_ptr<Chunk> from_empty(
        std::uint64_t num_local_insertions, Rank origin, Rank destination
    );
 
    [[nodiscard]] PackedData release();
 
    static constexpr std::uint64_t ID_BITS = 38;
    static constexpr std::uint64_t RANK_BITS =
        sizeof(ChunkID) * std::numeric_limits<unsigned char>::digits - ID_BITS;
 
    static constexpr ChunkID chunk_id(std::uint64_t sequence, Rank origin);
 
    [[nodiscard]] std::unique_ptr<std::vector<std::uint8_t>> serialize() const;
 
    [[nodiscard]] static std::unique_ptr<Chunk> deserialize(
        std::vector<std::uint8_t>& data, BufferResource* br
    );
 
    [[nodiscard]] std::unique_ptr<Buffer> release_data_buffer() noexcept;
 
    void attach_data_buffer(std::unique_ptr<Buffer> data);
 
    ~Chunk() = default;
    Chunk(Chunk&&) = default;
    Chunk& operator=(Chunk&&) = default;
    Chunk(Chunk const&) = delete;
    Chunk& operator=(Chunk const&) = delete;
};
 
class PostBox {
  public:
    PostBox() = default;
    ~PostBox() = default;
    PostBox(PostBox const&) = delete;
    PostBox& operator=(PostBox const&) = delete;
    PostBox(PostBox&&) = delete;
    PostBox& operator=(PostBox&&) = delete;
 
    void insert(std::unique_ptr<Chunk> chunk);
 
    void insert(std::vector<std::unique_ptr<Chunk>>&& chunks);
 
    [[nodiscard]] std::vector<std::unique_ptr<Chunk>> extract_ready();
 
    [[nodiscard]] std::vector<std::unique_ptr<Chunk>> extract();
 
    [[nodiscard]] std::size_t size() const noexcept;
 
    [[nodiscard]] bool empty() const noexcept;
 
    [[nodiscard]] std::size_t spill(BufferResource* br, std::size_t amount);
 
  private:
    mutable std::mutex mutex_{};  
    std::vector<std::unique_ptr<Chunk>> chunks_{};  
};
 
[[nodiscard]] std::vector<std::unique_ptr<Chunk>> test_some(
    std::vector<std::unique_ptr<Chunk>>& chunks,
    std::vector<std::unique_ptr<Communicator::Future>>& futures,
    Communicator* comm
);
 
}  // namespace rapidsmpf::coll::detail