host_buffer.hpp

 
#pragma once
 
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <memory>
#include <span>
#include <vector>
 
#include <rmm/cuda_stream_view.hpp>
#include <rmm/device_buffer.hpp>
#include <rmm/resource_ref.hpp>
 
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/memory/host_memory_resource.hpp>
#include <rapidsmpf/memory/pinned_memory_resource.hpp>
 
namespace rapidsmpf {
 
 
class HostBuffer {
  public:
    using OwnedStorageDeleter = std::function<void(void*)>;
 
    HostBuffer(
        std::size_t size, rmm::cuda_stream_view stream, rmm::host_async_resource_ref mr
    );
 
    ~HostBuffer() noexcept;
 
    // Non copyable
    HostBuffer(HostBuffer const&) = delete;
    HostBuffer& operator=(HostBuffer const&) = delete;
 
    HostBuffer(HostBuffer&& other) noexcept;
 
    HostBuffer& operator=(HostBuffer&& other);
 
    void deallocate_async() noexcept;
 
    [[nodiscard]] rmm::cuda_stream_view stream() const noexcept;
 
    [[nodiscard]] std::size_t size() const noexcept;
 
    [[nodiscard]] bool empty() const noexcept;
 
    [[nodiscard]] std::byte* data() noexcept;
 
    [[nodiscard]] std::byte const* data() const noexcept;
 
    [[nodiscard]] std::vector<std::uint8_t> copy_to_uint8_vector() const;
 
    void set_stream(rmm::cuda_stream_view stream);
 
    static HostBuffer from_uint8_vector(
        std::vector<std::uint8_t> const& data,
        rmm::cuda_stream_view stream,
        rmm::host_async_resource_ref mr
    );
 
    static HostBuffer from_owned_vector(
        std::vector<std::uint8_t>&& data,
        rmm::cuda_stream_view stream,
        rmm::host_async_resource_ref mr
    );
 
    static HostBuffer from_rmm_device_buffer(
        std::unique_ptr<rmm::device_buffer> pinned_host_buffer,
        rmm::cuda_stream_view stream,
        PinnedMemoryResource& mr
    );
 
  private:
    HostBuffer(
        std::span<std::byte> span,
        rmm::cuda_stream_view stream,
        rmm::host_async_resource_ref mr,
        std::unique_ptr<void, OwnedStorageDeleter> owned_storage
    );
 
    rmm::cuda_stream_view stream_;
    rmm::host_async_resource_ref mr_;
    std::span<std::byte> span_{};
    std::unique_ptr<void, OwnedStorageDeleter> owned_storage_{nullptr, [](void*) {}};
};
 
}  // namespace rapidsmpf