device_uvector.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <rmm/detail/cuda_memcpy.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/iterator>
#include <cuda/std/span>
#include <cuda/stream>
 
#include <cstddef>
#include <limits>
#include <type_traits>
#include <utility>
 
RMM_NAMESPACE_BEGIN

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 =
    cuda::std::reverse_iterator<iterator>;  
  using const_reverse_iterator =
    cuda::std::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_ref stream,
    cuda::mr::any_resource<cuda::mr::device_accessible> mr = mr::get_current_device_resource_ref())
    : _storage{elements_to_bytes(size), std::alignment_of_v<T>, stream, std::move(mr)}
  {
  }

  explicit device_uvector(
    device_uvector const& other,
    cuda::stream_ref stream,
    cuda::mr::any_resource<cuda::mr::device_accessible> mr = mr::get_current_device_resource_ref())
    : _storage{other._storage, stream, std::move(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_ref stream)
  {
    RMM_EXPECTS(
      element_index < size(), "Attempt to access out of bounds element.", rmm::out_of_range);
    RMM_CUDA_TRY(
      rmm::detail::memcpy_async(element_ptr(element_index), &value, sizeof(value), stream));
  }
 
  // 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_ref) = delete;

  void set_element_to_zero_async(size_type element_index, cuda::stream_ref 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.get()));
  }

  void set_element(size_type element_index, T const& value, cuda::stream_ref stream)
  {
    set_element_async(element_index, value, stream);
    stream.sync();
  }

  [[nodiscard]] value_type element(size_type element_index, cuda::stream_ref stream) const
  {
    RMM_EXPECTS(
      element_index < size(), "Attempt to access out of bounds element.", rmm::out_of_range);
    value_type value;
    RMM_CUDA_TRY(
      rmm::detail::memcpy_async(&value, element_ptr(element_index), sizeof(value), stream));
    stream.sync();
    return value;
  }

  [[nodiscard]] value_type front_element(cuda::stream_ref stream) const
  {
    return element(0, stream);
  }

  [[nodiscard]] value_type back_element(cuda::stream_ref stream) const
  {
    return element(size() - 1, stream);
  }

  void reserve(size_type new_capacity, cuda::stream_ref stream)
  {
    _storage.reserve(elements_to_bytes(new_capacity), stream);
  }

  void resize(size_type new_size, cuda::stream_ref stream)
  {
    _storage.resize(elements_to_bytes(new_size), stream);
  }

  void shrink_to_fit(cuda::stream_ref 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() noexcept
  {
    return _storage.memory_resource();
  }


  [[nodiscard]] cuda::stream_ref stream() const noexcept { return _storage.stream(); }

  void set_stream(cuda::stream_ref stream) noexcept { _storage.set_stream(stream); }
 
 private:
  device_buffer _storage{};  
 
  [[nodiscard]] size_type elements_to_bytes(size_type num_elements) const
  {
    RMM_EXPECTS(num_elements <= std::numeric_limits<size_type>::max() / sizeof(value_type),
                "Requested size overflows device_uvector storage.",
                rmm::invalid_argument);
    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
RMM_NAMESPACE_END