device_scalar.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2019-2025, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <rmm/cuda_stream_view.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/device_uvector.hpp>
#include <rmm/mr/device/per_device_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <type_traits>
 
namespace RMM_NAMESPACE {
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_view stream,
                         device_async_resource_ref mr = mr::get_current_device_resource_ref())
    : _storage{1, stream, mr}
  {
  }
 
  explicit device_scalar(value_type const& initial_value,
                         cuda_stream_view stream,
                         device_async_resource_ref mr = mr::get_current_device_resource_ref())
    : _storage{1, stream, mr}
  {
    set_value_async(initial_value, stream);
  }
 
  device_scalar(device_scalar const& other,
                cuda_stream_view stream,
                device_async_resource_ref mr = mr::get_current_device_resource_ref())
    : _storage{other._storage, stream, mr}
  {
  }
 
  [[nodiscard]] value_type value(cuda_stream_view stream) const
  {
    return _storage.front_element(stream);
  }
 
  void set_value_async(value_type const& value, cuda_stream_view 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&&, cuda_stream_view) = delete;
 
  void set_value_to_zero_async(cuda_stream_view stream)
  {
    _storage.set_element_to_zero_async(value_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_view stream() const noexcept { return _storage.stream(); }
 
  void set_stream(cuda_stream_view stream) noexcept { _storage.set_stream(stream); }
 
 private:
  rmm::device_uvector<T> _storage;
};
  // end of group
}  // namespace RMM_NAMESPACE