host_buffer.hpp

 
#pragma once
 
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <memory>
#include <span>
#include <vector>
 
#include <cuda/memory_resource>
 
#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>
 
namespace rapidsmpf {
 
 
class HostBuffer {
  public:
    HostBuffer(
        std::size_t size,
        rmm::cuda_stream_view stream,
        cuda::mr::any_resource<cuda::mr::host_accessible> 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
    );
 
    static HostBuffer from_rmm_device_buffer(
        std::unique_ptr<rmm::device_buffer> pinned_host_buffer,
        rmm::cuda_stream_view stream
    );
 
  private:
    HostBuffer(
        std::span<std::byte> span,
        rmm::cuda_stream_view stream,
        std::function<void(rmm::cuda_stream_view)> deallocate_fn
    );
 
    rmm::cuda_stream_view stream_;
    std::span<std::byte> span_{};
    std::function<void(rmm::cuda_stream_view)> deallocate_fn_{};
};
 
}  // namespace rapidsmpf