memory_reservation.hpp

 
#pragma once
 
#include <rmm/cuda_stream_pool.hpp>
 
#include <rapidsmpf/memory/buffer.hpp>
#include <rapidsmpf/memory/memory_reservation.hpp>
 
namespace rapidsmpf {
 
class BufferResource;
 
class MemoryReservation {
    friend class BufferResource;
 
  public:
    ~MemoryReservation() noexcept;
 
    void clear() noexcept;
 
    MemoryReservation(MemoryReservation&& o);
 
    MemoryReservation& operator=(MemoryReservation&& o) noexcept;
 
    MemoryReservation(MemoryReservation const&) = delete;
    MemoryReservation& operator=(MemoryReservation const&) = delete;
 
    [[nodiscard]] constexpr std::size_t size() const noexcept {
        return size_;
    }
 
    [[nodiscard]] constexpr MemoryType mem_type() const noexcept {
        return mem_type_;
    }
 
    [[nodiscard]] constexpr BufferResource* br() const noexcept {
        return br_;
    }
 
  private:
    constexpr MemoryReservation(MemoryType mem_type, BufferResource* br, std::size_t size)
        : mem_type_{mem_type}, br_{br}, size_{size} {}
 
  private:
    MemoryType mem_type_;  
    BufferResource* br_;  
    std::size_t size_;  
};
}  // namespace rapidsmpf