device_buffer.h

Go to the documentation of this file.

 
#pragma once
 
#include <cstddef>
#include <utility>
 
namespace fz {
 
struct DeviceSpan {
    void*  data  = nullptr;
    size_t bytes = 0;
 
    DeviceSpan() = default;
    DeviceSpan(void* d, size_t n) : data(d), bytes(n) {}
 
    bool empty() const { return bytes == 0; }
};
 
struct ConstDeviceSpan {
    const void* data  = nullptr;
    size_t      bytes = 0;
 
    ConstDeviceSpan() = default;
    ConstDeviceSpan(const void* d, size_t n) : data(d), bytes(n) {}
    ConstDeviceSpan(const DeviceSpan& s) : data(s.data), bytes(s.bytes) {}
 
    bool empty() const { return bytes == 0; }
};
 
class BorrowedDeviceBuffer {
public:
    BorrowedDeviceBuffer() = default;
    BorrowedDeviceBuffer(void* d, size_t n) : data_(d), bytes_(n) {}
 
    void*  data()  const { return data_; }
    size_t bytes() const { return bytes_; }
    bool   empty() const { return bytes_ == 0; }
 
    DeviceSpan      span()  const { return DeviceSpan(data_, bytes_); }
    ConstDeviceSpan cspan() const { return ConstDeviceSpan(data_, bytes_); }
 
private:
    void*  data_  = nullptr;
    size_t bytes_ = 0;
};
 
namespace detail {
void freeDeviceBuffer(void* ptr, int device) noexcept;
}  // namespace detail
 
class OwnedDeviceBuffer {
public:
    OwnedDeviceBuffer() = default;
    OwnedDeviceBuffer(void* d, size_t n, int device)
        : data_(d), bytes_(n), device_(device) {}
 
    OwnedDeviceBuffer(const OwnedDeviceBuffer&)            = delete;
    OwnedDeviceBuffer& operator=(const OwnedDeviceBuffer&) = delete;
 
    OwnedDeviceBuffer(OwnedDeviceBuffer&& other) noexcept
        : data_(other.data_), bytes_(other.bytes_), device_(other.device_) {
        other.data_  = nullptr;
        other.bytes_ = 0;
    }
 
    OwnedDeviceBuffer& operator=(OwnedDeviceBuffer&& other) noexcept {
        if (this != &other) {
            reset();
            data_        = other.data_;
            bytes_       = other.bytes_;
            device_      = other.device_;
            other.data_  = nullptr;
            other.bytes_ = 0;
        }
        return *this;
    }
 
    ~OwnedDeviceBuffer() { reset(); }
 
    void*  data()   const { return data_; }
    size_t bytes()  const { return bytes_; }
    int    device() const { return device_; }
    bool   empty()  const { return bytes_ == 0; }
    explicit operator bool() const { return data_ != nullptr; }
 
    DeviceSpan      span()  const { return DeviceSpan(data_, bytes_); }
    ConstDeviceSpan cspan() const { return ConstDeviceSpan(data_, bytes_); }
 
    void* release() {
        void* p = data_;
        data_   = nullptr;
        bytes_  = 0;
        return p;
    }
 
    void reset() {
        if (data_) detail::freeDeviceBuffer(data_, device_);
        data_  = nullptr;
        bytes_ = 0;
    }
 
private:
    void*  data_   = nullptr;
    size_t bytes_  = 0;
    int    device_ = 0;
};
 
}  // namespace fz