buffer.hpp

 
#pragma once
 
#include <atomic>
#include <cstddef>
#include <functional>
#include <memory>
#include <variant>
 
#include <cuda_runtime.h>
 
#include <rmm/cuda_stream_view.hpp>
#include <rmm/device_buffer.hpp>
 
#include <rapidsmpf/cuda_event.hpp>
#include <rapidsmpf/error.hpp>
#include <rapidsmpf/memory/host_buffer.hpp>
#include <rapidsmpf/memory/memory_type.hpp>
#include <rapidsmpf/utils.hpp>
 
namespace rapidsmpf {
 
class Buffer {
    friend class BufferResource;
 
  public:
    using DeviceBufferT = std::unique_ptr<rmm::device_buffer>;
 
    using HostBufferT = std::unique_ptr<HostBuffer>;
 
    static constexpr std::array<MemoryType, 1> device_buffer_types{MemoryType::DEVICE};
 
    static constexpr std::array<MemoryType, 2> host_buffer_types{
        MemoryType::HOST, MemoryType::PINNED_HOST
    };
 
    [[nodiscard]] std::byte const* data() const;
 
    template <typename F>
    auto write_access(F&& f)
        -> std::invoke_result_t<F, std::byte*, rmm::cuda_stream_view> {
        using Fn = std::remove_reference_t<F>;
        static_assert(
            std::is_invocable_v<Fn, std::byte*, rmm::cuda_stream_view>,
            "write_access() expects callable R(std::byte*, rmm::cuda_stream_view)"
        );
        using R = std::invoke_result_t<Fn, std::byte*, rmm::cuda_stream_view>;
 
        auto* ptr = const_cast<std::byte*>(data());
        if constexpr (std::is_void_v<R>) {
            std::invoke(std::forward<F>(f), ptr, stream_);
            latest_write_event_.record(stream_);
        } else {
            auto ret = std::invoke(std::forward<F>(f), ptr, stream_);
            latest_write_event_.record(stream_);
            return ret;
        }
    }
 
    std::byte* exclusive_data_access();
 
    void unlock();
 
    [[nodiscard]] MemoryType constexpr mem_type() const {
        return mem_type_;
    }
 
    [[nodiscard]] constexpr rmm::cuda_stream_view stream() const noexcept {
        return stream_;
    }
 
    [[nodiscard]] bool is_latest_write_done() const;
 
    Buffer(Buffer&&) = delete;
    Buffer(Buffer const&) = delete;
    Buffer& operator=(Buffer& o) = delete;
    Buffer& operator=(Buffer&& o) = delete;
 
  private:
    Buffer(
        std::unique_ptr<HostBuffer> host_buffer,
        rmm::cuda_stream_view stream,
        MemoryType mem_type
    );
 
    Buffer(std::unique_ptr<rmm::device_buffer> device_buffer, MemoryType mem_type);
 
    void throw_if_locked() const;
 
    [[nodiscard]] DeviceBufferT release_device_buffer();
 
    [[nodiscard]] HostBufferT release_host_buffer();
 
  public:
    std::size_t const size;  
 
  private:
    MemoryType const mem_type_;
    std::variant<DeviceBufferT, HostBufferT> storage_;
    rmm::cuda_stream_view stream_;
    CudaEvent latest_write_event_;
    std::atomic<bool> lock_;
};
 
void buffer_copy(
    Buffer& dst,
    Buffer const& src,
    std::size_t size,
    std::ptrdiff_t dst_offset = 0,
    std::ptrdiff_t src_offset = 0
);
 
}  // namespace rapidsmpf