Class SharedArray

grape/cuda/utils/shared_array.h:25–136

namespace cuda {
template <typename T>
class SharedArray {
  static_assert(std::is_pod<T>::value, "Unsupported datatype");

 public:
  using device_t = thrust::device_vector<T>;
  using host_t = pinned_vector<T>;

  SharedArray() = default;

  explicit SharedArray(size_t size) {
    d_buffer_.resize(size);
    h_buffer_.resize(size);
  }

  void resize(size_t size) {
    d_buffer_.resize(size);
    h_buffer_.resize(size);
  }

  void set(size_t idx, const T& t) { d_buffer_[idx] = t; }

  void set(size_t idx, const T& t, const Stream& stream) {
    h_buffer_[idx] = t;
    CHECK_CUDA(cudaMemcpyAsync(thrust::raw_pointer_cast(d_buffer_.data()),
                               thrust::raw_pointer_cast(h_buffer_.data()),
                               sizeof(T), cudaMemcpyHostToDevice,
                               stream.cuda_stream()));
  }

  void fill(const T& t) {
    auto size = h_buffer_.size();
    thrust::fill_n(h_buffer_.data(), size, t);
    d_buffer_ = h_buffer_;
  }

  void fill(const T& t, const Stream& stream) {
    auto size = h_buffer_.size();

    thrust::fill_n(h_buffer_.data(), size, t);
    CHECK_CUDA(cudaMemcpyAsync(thrust::raw_pointer_cast(d_buffer_.data()),
                               thrust::raw_pointer_cast(h_buffer_.data()),
                               sizeof(T) * size, cudaMemcpyHostToDevice,
                               stream.cuda_stream()));
  }

  typename thrust::device_vector<T>::reference get(size_t idx) {
    return d_buffer_[idx];
  }

  typename thrust::device_vector<T>::const_reference get(size_t idx) const {
    return d_buffer_[idx];
  }

  T get(size_t idx, const Stream& stream) const {
    CHECK_CUDA(cudaMemcpyAsync(
        const_cast<T*>((thrust::raw_pointer_cast(h_buffer_.data()) + idx)),
        thrust::raw_pointer_cast(d_buffer_.data()) + idx, sizeof(T),
        cudaMemcpyDeviceToHost, stream.cuda_stream()));