buffer.hpp

 
#pragma once
 
#include <array>
#include <atomic>
#include <cstddef>
#include <functional>
#include <memory>
#include <variant>
#include <vector>
 
#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/utils.hpp>
 
namespace rapidsmpf {
 
enum class MemoryType : int {
    DEVICE = 0,  
    HOST = 1  
};
 
constexpr MemoryType LowestSpillType = MemoryType::HOST;
 
constexpr std::array<MemoryType, 2> MEMORY_TYPES{{MemoryType::DEVICE, MemoryType::HOST}};
 
class Buffer {
    friend class BufferResource;
 
  public:
    using DeviceStorageT = std::unique_ptr<rmm::device_buffer>;
 
    using HostStorageT = std::unique_ptr<std::vector<uint8_t>>;
 
    using StorageT = std::variant<DeviceStorageT, HostStorageT>;
 
    [[nodiscard]] HostStorageT const& host() const;
 
    [[nodiscard]] DeviceStorageT const& device() const;
 
    [[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 std::visit(
            overloaded{
                [](HostStorageT const&) -> MemoryType { return MemoryType::HOST; },
                [](DeviceStorageT const&) -> MemoryType { return MemoryType::DEVICE; }
            },
            storage_
        );
    }
 
    [[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<std::vector<uint8_t>> host_buffer, rmm::cuda_stream_view stream
    );
 
    Buffer(std::unique_ptr<rmm::device_buffer> device_buffer);
 
    void throw_if_locked() const;
 
    [[nodiscard]] HostStorageT& host();
 
    [[nodiscard]] DeviceStorageT& device();
 
    [[nodiscard]] DeviceStorageT release_device();
 
    [[nodiscard]] HostStorageT release_host();
 
  public:
    std::size_t const size;  
 
  private:
    StorageT 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