cuda_buffer.hpp

// SPDX-License-Identifier: MIT
// Copyright (c) 2021  Kenji Koide (k.koide@aist.go.jp)
 
#pragma once
 
#include <string>
 
struct CUstream_st;
 
namespace gtsam_points {
 
class CUDABuffer {
public:
  CUDABuffer(bool use_pinned_buffer = true);
  ~CUDABuffer();
 
  void resize(size_t size, CUstream_st* stream);
 
  void upload(CUstream_st* stream);
 
  void upload(size_t size, CUstream_st* stream);
 
  void upload(const void* buffer, size_t size, CUstream_st* stream);
 
  void download(CUstream_st* stream);
 
  void download(void* buffer, size_t size, CUstream_st* stream);
 
  size_t size() const;
 
  void* host_buffer();
 
  void* device_buffer();
 
  template <typename T>
  void upload(const T* buffer, size_t size, CUstream_st* stream) {
    upload(reinterpret_cast<const void*>(buffer), size, stream);
  }
  template <typename T>
  void download(T* buffer, size_t size, CUstream_st* stream) {
    download(reinterpret_cast<void*>(buffer), size, stream);
  }
  template <typename T>
  T* host_buffer() {
    return reinterpret_cast<T*>(host_buffer());
  }
  template <typename T>
  T* device_buffer() {
    return reinterpret_cast<T*>(device_buffer());
  }
 
private:
  const bool use_pinned_buffer;
 
  size_t buffer_size;
  void* h_buffer;
  void* d_buffer;
};
 
}  // namespace gtsam_points