device_scalar.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <rmm/detail/export.hpp>
#include <rmm/device_uvector.hpp>
#include <rmm/mr/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <cuda/stream>
 
#include <type_traits>
 
RMM_NAMESPACE_BEGIN

template <typename T>
class device_scalar {
 public:
  static_assert(std::is_trivially_copyable_v<T>, "Scalar type must be trivially copyable");
 
  using value_type = typename device_uvector<T>::value_type;  
  using size_type  = typename device_uvector<T>::size_type;   
  using reference  = typename device_uvector<T>::reference;   
  using const_reference = typename device_uvector<T>::const_reference;  
  using pointer =
    typename device_uvector<T>::pointer;  
  using const_pointer = typename device_uvector<T>::const_pointer;  
 
  RMM_EXEC_CHECK_DISABLE
  ~device_scalar() = default;
 
  RMM_EXEC_CHECK_DISABLE
  device_scalar(device_scalar&&) noexcept = default;  

  device_scalar& operator=(device_scalar&&) noexcept = default;

  device_scalar(device_scalar const&) = delete;

  device_scalar& operator=(device_scalar const&) = delete;

  device_scalar() = delete;

  explicit device_scalar(
    cuda::stream_ref stream,
    cuda::mr::any_resource<cuda::mr::device_accessible> mr = mr::get_current_device_resource_ref())
    : _storage{1, stream, std::move(mr)}
  {
  }

  explicit device_scalar(
    value_type const& initial_value,
    cuda::stream_ref stream,
    cuda::mr::any_resource<cuda::mr::device_accessible> mr = mr::get_current_device_resource_ref())
    : _storage{1, stream, std::move(mr)}
  {
    set_value_async(initial_value, stream);
  }
 
  // Disallow passing literals to the constructor to avoid race conditions where the
  // memory holding the literal can be freed before the async memcpy / memset executes.
  device_scalar(value_type const&&,
                cuda::stream_ref stream,
                cuda::mr::any_resource<cuda::mr::device_accessible> mr =
                  mr::get_current_device_resource_ref()) = delete;
  device_scalar(
    device_scalar 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]] value_type value(cuda::stream_ref stream) const
  {
    return _storage.front_element(stream);
  }

  void set_value_async(value_type const& value, cuda::stream_ref stream)
  {
    _storage.set_element_async(0, value, stream);
  }
 
  // Disallow passing literals to set_value to avoid race conditions where the memory holding the
  // literal can be freed before the async memcpy / memset executes.
  void set_value_async(value_type const&&, cuda::stream_ref) = delete;

  void set_value_to_zero_async(cuda::stream_ref stream)
  {
    _storage.set_element_to_zero_async(size_type{0}, stream);
  }

  [[nodiscard]] pointer data() noexcept { return static_cast<pointer>(_storage.data()); }

  [[nodiscard]] const_pointer data() const noexcept
  {
    return static_cast<const_pointer>(_storage.data());
  }

  [[nodiscard]] constexpr size_type size() const noexcept { return 1; }

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

  void set_stream(cuda::stream_ref stream) noexcept { _storage.set_stream(stream); }
 
 private:
  rmm::device_uvector<T> _storage;
};
 
static_assert(std::is_constructible_v<device_scalar<int>, int const&, cuda::stream_ref>);
static_assert(std::is_constructible_v<device_scalar<int>,
                                      int const&,
                                      cuda::stream_ref,
                                      cuda::mr::any_resource<cuda::mr::device_accessible>>);
static_assert(!std::is_constructible_v<device_scalar<int>, int, cuda::stream_ref>);
static_assert(!std::is_constructible_v<device_scalar<int>, int const, cuda::stream_ref>);
static_assert(!std::is_constructible_v<device_scalar<int>,
                                       int,
                                       cuda::stream_ref,
                                       cuda::mr::any_resource<cuda::mr::device_accessible>>);
static_assert(!std::is_constructible_v<device_scalar<int>,
                                       int const,
                                       cuda::stream_ref,
                                       cuda::mr::any_resource<cuda::mr::device_accessible>>);
static_assert([]<typename Scalar>(Scalar*) {
  return requires(Scalar& scalar, int value, cuda::stream_ref stream) {
    scalar.set_value_async(value, stream);
  } && requires(Scalar& scalar, int const value, cuda::stream_ref stream) {
    scalar.set_value_async(value, stream);
  } && !requires(Scalar& scalar, int value, cuda::stream_ref stream) {
    scalar.set_value_async(std::move(value), stream);
  } && !requires(Scalar& scalar, int const value, cuda::stream_ref stream) {
    scalar.set_value_async(std::move(value), stream);
  };
}(static_cast<device_scalar<int>*>(nullptr)));
  // end of group
RMM_NAMESPACE_END