Class Memory

include/base/memory.h:34–250

 */
template<class _T, class _Index = size_t>
class Memory {
public:
    typedef _T Data;
    typedef _Index Index;

    int device_id;
    Index count = 0, capacity = 0;
    cudaStream_t stream;
    int *refer_count = nullptr;
    Data *host_ptr = nullptr, *device_ptr = nullptr;

    /**
     * @brief Construct a memory space
     * @param _device_id GPU id, -1 for CPU
     * @param _count number of data
     * @param _stream CUDA stream
     */
    Memory(int _device_id, Index _count = 0, cudaStream_t _stream = 0) :
            device_id(_device_id), stream(_stream) {
        resize(_count);
    }

    /** Shallow copy constructor */
    Memory(const Memory &m) :
            device_id(m.device_id), count(m.count), capacity(m.capacity), stream(m.stream), refer_count(m.refer_count),
            host_ptr(m.host_ptr), device_ptr(m.device_ptr) {
        if (capacity)
            (*refer_count)++;
    }

    Memory &operator=(const Memory &) = delete;

    ~Memory() { reallocate(0); }

    /** Swap two memory spaces */
    void swap(Memory &m) {
        std::swap(device_id, m.device_id);
        std::swap(count, m.count);
        std::swap(capacity, m.capacity);
        std::swap(stream, m.stream);
        std::swap(refer_count, m.refer_count);
        std::swap(host_ptr, m.host_ptr);
        std::swap(device_ptr, m.device_ptr);
    }

    __host__ __device__ Data &operator[](Index index) {
#ifdef __CUDA_ARCH__
        return device_ptr[index];
#else
        return host_ptr[index];
#endif
    }

    __host__ __device__ Data &operator[](Index index) const {
#ifdef __CUDA_ARCH__
        return device_ptr[index];
#else
        return host_ptr[index];