buffer.h

 
#pragma once
 
#include <memory>
 
#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,
  CCCL,
  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]] void* data() override;
};
 
#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;
 
  [[deprecated(
    "RMMBuffer is deprecated and will be removed in a future release. Use CCCL buffers instead "
    "(UCXX_ENABLE_CCCL).")]]
  explicit RMMBuffer(const size_t size);
 
  [[deprecated(
    "RMMBuffer is deprecated and will be removed in a future release. Use CCCL buffers instead "
    "(UCXX_ENABLE_CCCL).")]]
  explicit RMMBuffer(std::unique_ptr<rmm::device_buffer> rmm_buffer);
 
  [[nodiscard]] std::unique_ptr<rmm::device_buffer> release();
 
  [[nodiscard]] void* data() override;
};
#endif
 
#if UCXX_ENABLE_CCCL
struct CCCLBufferImpl;
 
class CCCLBuffer : public Buffer {
 private:
  std::unique_ptr<CCCLBufferImpl> _impl;  
 
 public:
  CCCLBuffer()                             = delete;
  CCCLBuffer(const CCCLBuffer&)            = delete;
  CCCLBuffer& operator=(CCCLBuffer const&) = delete;
  CCCLBuffer(CCCLBuffer&& o)               = delete;
  CCCLBuffer& operator=(CCCLBuffer&& o)    = delete;
 
  ~CCCLBuffer() override;
 
  explicit CCCLBuffer(const size_t size);
 
  [[nodiscard]] void* data() override;
};
#endif
 
[[nodiscard]] std::shared_ptr<Buffer> allocateBuffer(BufferType bufferType, const size_t size);
 
}  // namespace ucxx