All Classes Files Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
device_scalar.hpp
Go to the documentation of this file.
1 /*
2  * Copyright (c) 2019-2025, NVIDIA CORPORATION.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #pragma once
18 
19 #include <rmm/cuda_stream_view.hpp>
20 #include <rmm/detail/export.hpp>
21 #include <rmm/device_uvector.hpp>
23 #include <rmm/resource_ref.hpp>
24 
25 #include <type_traits>
26 
27 namespace RMM_NAMESPACE {
41 template <typename T>
43  public:
44  static_assert(std::is_trivially_copyable_v<T>, "Scalar type must be trivially copyable");
45 
50  using pointer =
54 
55  RMM_EXEC_CHECK_DISABLE
56  ~device_scalar() = default;
57 
58  RMM_EXEC_CHECK_DISABLE
59  device_scalar(device_scalar&&) noexcept = default;
60 
66  device_scalar& operator=(device_scalar&&) noexcept = default;
67 
71  device_scalar(device_scalar const&) = delete;
72 
76  device_scalar& operator=(device_scalar const&) = delete;
77 
81  device_scalar() = delete;
82 
99  : _storage{1, stream, mr}
100  {
101  }
102 
119  explicit device_scalar(value_type const& initial_value,
120  cuda_stream_view stream,
122  : _storage{1, stream, mr}
123  {
124  set_value_async(initial_value, stream);
125  }
126 
140  cuda_stream_view stream,
142  : _storage{other._storage, stream, mr}
143  {
144  }
145 
162  [[nodiscard]] value_type value(cuda_stream_view stream) const
163  {
164  return _storage.front_element(stream);
165  }
166 
204  void set_value_async(value_type const& value, cuda_stream_view stream)
205  {
206  _storage.set_element_async(0, value, stream);
207  }
208 
209  // Disallow passing literals to set_value to avoid race conditions where the memory holding the
210  // literal can be freed before the async memcpy / memset executes.
211  void set_value_async(value_type&&, cuda_stream_view) = delete;
212 
228  {
229  _storage.set_element_to_zero_async(value_type{0}, stream);
230  }
231 
242  [[nodiscard]] pointer data() noexcept { return static_cast<pointer>(_storage.data()); }
243 
254  [[nodiscard]] const_pointer data() const noexcept
255  {
256  return static_cast<const_pointer>(_storage.data());
257  }
258 
262  [[nodiscard]] constexpr size_type size() const noexcept { return 1; }
263 
267  [[nodiscard]] cuda_stream_view stream() const noexcept { return _storage.stream(); }
268 
274  void set_stream(cuda_stream_view stream) noexcept { _storage.set_stream(stream); }
275 
276  private:
277  rmm::device_uvector<T> _storage;
278 };
279  // end of group
281 } // namespace RMM_NAMESPACE
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:39
Container for a single object of type T in device memory.
Definition: device_scalar.hpp:42
typename device_uvector< T >::value_type value_type
T, the type of the scalar element.
Definition: device_scalar.hpp:46
constexpr size_type size() const noexcept
The size of the scalar: always 1.
Definition: device_scalar.hpp:262
const_pointer data() const noexcept
Returns const pointer to object in device memory.
Definition: device_scalar.hpp:254
typename device_uvector< T >::const_reference const_reference
const value_type&
Definition: device_scalar.hpp:49
typename device_uvector< T >::size_type size_type
The type used for the size.
Definition: device_scalar.hpp:47
device_scalar(device_scalar &&) noexcept=default
Default move constructor.
void set_stream(cuda_stream_view stream) noexcept
Sets the stream to be used for deallocation.
Definition: device_scalar.hpp:274
device_scalar(value_type const &initial_value, cuda_stream_view stream, device_async_resource_ref mr=mr::get_current_device_resource_ref())
Construct a new device_scalar with an initial value.
Definition: device_scalar.hpp:119
device_scalar(device_scalar const &other, cuda_stream_view stream, device_async_resource_ref mr=mr::get_current_device_resource_ref())
Construct a new device_scalar by deep copying the contents of another device_scalar,...
Definition: device_scalar.hpp:139
typename device_uvector< T >::const_pointer const_pointer
Definition: device_scalar.hpp:53
cuda_stream_view stream() const noexcept
Stream associated with the device memory allocation.
Definition: device_scalar.hpp:267
pointer data() noexcept
Returns pointer to object in device memory.
Definition: device_scalar.hpp:242
value_type value(cuda_stream_view stream) const
Copies the value from device to host, synchronizes, and returns the value.
Definition: device_scalar.hpp:162
void set_value_to_zero_async(cuda_stream_view stream)
Sets the value of the device_scalar to zero on the specified stream.
Definition: device_scalar.hpp:227
void set_value_async(value_type const &value, cuda_stream_view stream)
Sets the value of the device_scalar to the value of v.
Definition: device_scalar.hpp:204
typename device_uvector< T >::pointer pointer
The type of the pointer returned by data()
Definition: device_scalar.hpp:51
typename device_uvector< T >::reference reference
value_type&
Definition: device_scalar.hpp:48
An uninitialized vector of elements in device memory.
Definition: device_uvector.hpp:78
value_type * pointer
The type of the pointer returned by data()
Definition: device_uvector.hpp:88
std::size_t size_type
The type used for the size of the vector.
Definition: device_uvector.hpp:84
T value_type
T; stored value type.
Definition: device_uvector.hpp:83
value_type & reference
value_type&; reference type returned by operator[](size_type)
Definition: device_uvector.hpp:85
value_type const * const_pointer
The type of the pointer returned by data() const.
Definition: device_uvector.hpp:89
value_type const & const_reference
Definition: device_uvector.hpp:87
cuda::mr::async_resource_ref< cuda::mr::device_accessible > device_async_resource_ref
Alias for a cuda::mr::async_resource_ref with the property cuda::mr::device_accessible.
Definition: resource_ref.hpp:41
device_async_resource_ref get_current_device_resource_ref()
Get the device_async_resource_ref for the current device.
Definition: per_device_resource.hpp:411
Management of per-device device_memory_resources.