device_uvector.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/error.hpp>
#include <rmm/detail/exec_check_disable.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/device_buffer.hpp>
#include <rmm/mr/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <cuda/std/span>
#include <thrust/iterator/reverse_iterator.h>
 
#include <cstddef>
#include <type_traits>
#include <utility>
 
namespace RMM_NAMESPACE {
template <typename T>
class device_uvector {
  static_assert(std::is_trivially_copyable_v<T>,
                "device_uvector only supports types that are trivially copyable.");
 
 public:
  using value_type = T;            
  using size_type  = std::size_t;  
  using reference  = value_type&;  
  using const_reference =
    value_type const&;  
  using pointer        = value_type*;        
  using const_pointer  = value_type const*;  
  using iterator       = pointer;            
  using const_iterator = const_pointer;  
  using reverse_iterator =
    thrust::reverse_iterator<iterator>;  
  using const_reverse_iterator =
    thrust::reverse_iterator<const_iterator>;  
 
  RMM_EXEC_CHECK_DISABLE
  ~device_uvector() = default;
 
  RMM_EXEC_CHECK_DISABLE
  device_uvector(device_uvector&&) noexcept = default;  
 
  RMM_EXEC_CHECK_DISABLE
  device_uvector& operator=(device_uvector&&) noexcept =
    default;  
 
  device_uvector(device_uvector const&) = delete;
 
  device_uvector& operator=(device_uvector const&) = delete;
 
  device_uvector() = delete;
 
  explicit device_uvector(size_type size,
                          cuda_stream_view stream,
                          device_async_resource_ref mr = mr::get_current_device_resource_ref())
    : _storage{elements_to_bytes(size), stream, mr}
  {
  }
 
  explicit device_uvector(device_uvector const& other,
                          cuda_stream_view stream,
                          device_async_resource_ref mr = mr::get_current_device_resource_ref())
    : _storage{other._storage, stream, mr}
  {
  }
 
  [[nodiscard]] pointer element_ptr(size_type element_index) noexcept
  {
    assert(element_index < size());
    return data() + element_index;
  }
 
  [[nodiscard]] const_pointer element_ptr(size_type element_index) const noexcept
  {
    assert(element_index < size());
    return data() + element_index;
  }
 
  void set_element_async(size_type element_index, value_type const& value, cuda_stream_view stream)
  {
    RMM_EXPECTS(
      element_index < size(), "Attempt to access out of bounds element.", rmm::out_of_range);
 
    if constexpr (std::is_same_v<value_type, bool>) {
      RMM_CUDA_TRY(
        cudaMemsetAsync(element_ptr(element_index), value, sizeof(value), stream.value()));
      return;
    }
 
    if constexpr (std::is_fundamental_v<value_type>) {
      if (value == value_type{0}) {
        set_element_to_zero_async(element_index, stream);
        return;
      }
    }
 
    RMM_CUDA_TRY(cudaMemcpyAsync(
      element_ptr(element_index), &value, sizeof(value), cudaMemcpyDefault, stream.value()));
  }
 
  // We delete the r-value reference overload to prevent asynchronously copying from a literal or
  // implicit temporary value after it is deleted or goes out of scope.
  void set_element_async(size_type, value_type const&&, cuda_stream_view) = delete;
 
  void set_element_to_zero_async(size_type element_index, cuda_stream_view stream)
  {
    RMM_EXPECTS(
      element_index < size(), "Attempt to access out of bounds element.", rmm::out_of_range);
    RMM_CUDA_TRY(
      cudaMemsetAsync(element_ptr(element_index), 0, sizeof(value_type), stream.value()));
  }
 
  void set_element(size_type element_index, T const& value, cuda_stream_view stream)
  {
    set_element_async(element_index, value, stream);
    stream.synchronize_no_throw();
  }
 
  [[nodiscard]] value_type element(size_type element_index, cuda_stream_view stream) const
  {
    RMM_EXPECTS(
      element_index < size(), "Attempt to access out of bounds element.", rmm::out_of_range);
    value_type value;
    RMM_CUDA_TRY(cudaMemcpyAsync(
      &value, element_ptr(element_index), sizeof(value), cudaMemcpyDefault, stream.value()));
    stream.synchronize();
    return value;
  }
 
  [[nodiscard]] value_type front_element(cuda_stream_view stream) const
  {
    return element(0, stream);
  }
 
  [[nodiscard]] value_type back_element(cuda_stream_view stream) const
  {
    return element(size() - 1, stream);
  }
 
  void reserve(size_type new_capacity, cuda_stream_view stream)
  {
    _storage.reserve(elements_to_bytes(new_capacity), stream);
  }
 
  void resize(size_type new_size, cuda_stream_view stream)
  {
    _storage.resize(elements_to_bytes(new_size), stream);
  }
 
  void shrink_to_fit(cuda_stream_view stream) { _storage.shrink_to_fit(stream); }
 
  device_buffer release() noexcept { return std::move(_storage); }
 
  [[nodiscard]] size_type capacity() const noexcept
  {
    return bytes_to_elements(_storage.capacity());
  }
 
  [[nodiscard]] pointer data() noexcept { return static_cast<pointer>(_storage.data()); }
 
  [[nodiscard]] const_pointer data() const noexcept
  {
    return static_cast<const_pointer>(_storage.data());
  }
 
  [[nodiscard]] iterator begin() noexcept { return data(); }
 
  [[nodiscard]] const_iterator cbegin() const noexcept { return data(); }
 
  [[nodiscard]] const_iterator begin() const noexcept { return cbegin(); }
 
  [[nodiscard]] iterator end() noexcept { return data() + size(); }
 
  [[nodiscard]] const_iterator cend() const noexcept { return data() + size(); }
 
  [[nodiscard]] const_iterator end() const noexcept { return cend(); }
 
  [[nodiscard]] reverse_iterator rbegin() noexcept { return reverse_iterator(end()); }
 
  [[nodiscard]] const_reverse_iterator crbegin() const noexcept
  {
    return const_reverse_iterator(cend());
  }
 
  [[nodiscard]] const_reverse_iterator rbegin() const noexcept { return crbegin(); }
 
  [[nodiscard]] reverse_iterator rend() noexcept { return reverse_iterator(begin()); }
 
  [[nodiscard]] const_reverse_iterator crend() const noexcept
  {
    return const_reverse_iterator(begin());
  }
 
  [[nodiscard]] const_reverse_iterator rend() const noexcept { return crend(); }
 
  [[nodiscard]] size_type size() const noexcept { return bytes_to_elements(_storage.size()); }
 
  [[nodiscard]] std::int64_t ssize() const noexcept
  {
    assert(size() < static_cast<size_type>(std::numeric_limits<int64_t>::max()) &&
           "Size overflows signed integer");
    return static_cast<int64_t>(size());
  }
 
  [[nodiscard]] bool is_empty() const noexcept { return size() == 0; }
 
  [[nodiscard]] operator cuda::std::span<T const>() const noexcept
  {
    return cuda::std::span<T const>{data(), size()};
  }
 
  [[nodiscard]] operator cuda::std::span<T>() noexcept
  {
    return cuda::std::span<T>{data(), size()};
  }
 
  [[nodiscard]] rmm::device_async_resource_ref memory_resource() const noexcept
  {
    return _storage.memory_resource();
  }
 
  [[nodiscard]] cuda_stream_view stream() const noexcept { return _storage.stream(); }
 
  void set_stream(cuda_stream_view stream) noexcept { _storage.set_stream(stream); }
 
 private:
  device_buffer _storage{};  
 
  [[nodiscard]] size_type constexpr elements_to_bytes(size_type num_elements) const noexcept
  {
    return num_elements * sizeof(value_type);
  }
 
  [[nodiscard]] size_type constexpr bytes_to_elements(size_type num_bytes) const noexcept
  {
    return num_bytes / sizeof(value_type);
  }
};
  // end of group
}  // namespace RMM_NAMESPACE