buffer.h

 
#pragma once
 
#include <memory>
#include <utility>
 
#include <ucxx/log.h>
 
namespace rmm {
// Forward declaration to prevent symbols from being added to symbol table unnecessarily.
class device_buffer;
}  // namespace rmm
 
namespace ucxx {
 
enum class BufferType {
  Host = 0,
  RMM,
  Invalid,
};
 
class Buffer {
 protected:
  BufferType _bufferType{BufferType::Invalid};  
  size_t _size;                                 
 
  Buffer(const BufferType bufferType, const size_t size);
 
 public:
  Buffer()                         = delete;
  Buffer(const Buffer&)            = delete;
  Buffer& operator=(Buffer const&) = delete;
  Buffer(Buffer&& o)               = delete;
  Buffer& operator=(Buffer&& o)    = delete;
 
  virtual ~Buffer();
 
  [[nodiscard]] BufferType getType() const noexcept;
 
  [[nodiscard]] size_t getSize() const noexcept;
 
  [[nodiscard]] virtual void* data() = 0;
};
 
class HostBuffer : public Buffer {
 private:
  void* _buffer;  
 
 public:
  HostBuffer()                             = delete;
  HostBuffer(const HostBuffer&)            = delete;
  HostBuffer& operator=(HostBuffer const&) = delete;
  HostBuffer(HostBuffer&& o)               = delete;
  HostBuffer& operator=(HostBuffer&& o)    = delete;
 
  explicit HostBuffer(const size_t size);
 
  HostBuffer(const void* buffer, const size_t size);
 
  ~HostBuffer();
 
  [[nodiscard]] void* release();
 
  [[nodiscard]] virtual void* data();
};
 
#if UCXX_ENABLE_RMM
class RMMBuffer : public Buffer {
 private:
  std::unique_ptr<rmm::device_buffer> _buffer;  
 
 public:
  RMMBuffer()                            = delete;
  RMMBuffer(const RMMBuffer&)            = delete;
  RMMBuffer& operator=(RMMBuffer const&) = delete;
  RMMBuffer(RMMBuffer&& o)               = delete;
  RMMBuffer& operator=(RMMBuffer&& o)    = delete;
 
  ~RMMBuffer() override;
 
  explicit RMMBuffer(const size_t size);
 
  explicit RMMBuffer(std::unique_ptr<rmm::device_buffer> rmm_buffer);
 
  [[nodiscard]] std::unique_ptr<rmm::device_buffer> release();
 
  [[nodiscard]] virtual void* data();
};
#endif
 
[[nodiscard]] std::shared_ptr<Buffer> allocateBuffer(BufferType bufferType, const size_t size);
 
}  // namespace ucxx