device_scalar.hpp
1 /*
2  * Copyright (c) 2019-2021, 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/device_uvector.hpp>
21 #include <rmm/mr/device/device_memory_resource.hpp>
23 
24 #include <type_traits>
25 
26 namespace rmm {
27 
35 template <typename T>
37  public:
38  static_assert(std::is_trivially_copyable<T>::value, "Scalar type must be trivially copyable");
39 
40  using value_type = typename device_uvector<T>::value_type;
41  using reference = typename device_uvector<T>::reference;
42  using const_reference = typename device_uvector<T>::const_reference;
43  using pointer = typename device_uvector<T>::pointer;
44  using const_pointer = typename device_uvector<T>::const_pointer;
45 
46  RMM_EXEC_CHECK_DISABLE
47  ~device_scalar() = default;
48 
49  RMM_EXEC_CHECK_DISABLE
50  device_scalar(device_scalar&&) noexcept = default;
51 
52  device_scalar& operator=(device_scalar&&) noexcept = default;
53 
57  device_scalar(device_scalar const&) = delete;
58 
63 
67  device_scalar() = delete;
68 
83  explicit device_scalar(
85  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource())
86  : _storage{1, stream, mr}
87  {
88  }
89 
106  explicit device_scalar(
107  value_type const& initial_value,
109  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource())
110  : _storage{1, stream, mr}
111  {
112  set_value_async(initial_value, stream);
113  }
114 
129  rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource())
130  : _storage{other._storage, stream, mr}
131  {
132  }
133 
150  [[nodiscard]] value_type value(cuda_stream_view stream) const
151  {
152  return _storage.front_element(stream);
153  }
154 
192  void set_value_async(value_type const& value, cuda_stream_view stream)
193  {
194  _storage.set_element_async(0, value, stream);
195  }
196 
197  // Disallow passing literals to set_value to avoid race conditions where the memory holding the
198  // literal can be freed before the async memcpy / memset executes.
199  void set_value_async(value_type&&, cuda_stream_view) = delete;
200 
216  {
217  _storage.set_element_to_zero_async(value_type{0}, stream);
218  }
219 
228  [[nodiscard]] pointer data() noexcept { return static_cast<pointer>(_storage.data()); }
229 
238  [[nodiscard]] const_pointer data() const noexcept
239  {
240  return static_cast<const_pointer>(_storage.data());
241  }
242 
246  [[nodiscard]] cuda_stream_view stream() const noexcept { return _storage.stream(); }
247 
251  void set_stream(cuda_stream_view stream) noexcept { _storage.set_stream(stream); }
252 
253  private:
254  rmm::device_uvector<T> _storage;
255 };
256 } // namespace rmm
rmm::device_scalar::device_scalar
device_scalar(device_scalar const &other, cuda_stream_view stream, rmm::mr::device_memory_resource *mr=rmm::mr::get_current_device_resource())
Construct a new device_scalar by deep copying the contents of another device_scalar,...
Definition: device_scalar.hpp:127
per_device_resource.hpp
Management of per-device device_memory_resources.
rmm::device_scalar::set_stream
void set_stream(cuda_stream_view stream) noexcept
Sets the stream to be used for deallocation.
Definition: device_scalar.hpp:251
rmm::cuda_stream_view
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:34
rmm::device_uvector
An uninitialized vector of elements in device memory.
Definition: device_uvector.hpp:69
rmm::device_scalar::operator=
device_scalar & operator=(device_scalar const &)=delete
Copy assignment is deleted as it doesn't allow a stream argument.
rmm::device_scalar::device_scalar
device_scalar(cuda_stream_view stream, rmm::mr::device_memory_resource *mr=rmm::mr::get_current_device_resource())
Construct a new uninitialized device_scalar.
Definition: device_scalar.hpp:83
rmm::device_scalar::data
pointer data() noexcept
Returns pointer to object in device memory.
Definition: device_scalar.hpp:228
rmm::device_scalar::data
const_pointer data() const noexcept
Returns const pointer to object in device memory.
Definition: device_scalar.hpp:238
rmm::device_scalar::device_scalar
device_scalar()=delete
Default constructor is deleted as it doesn't allow a stream argument.
rmm::device_scalar::set_value_async
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:192
rmm::device_scalar::set_value_to_zero_async
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:215
rmm::device_scalar
Container for a single object of type T in device memory.
Definition: device_scalar.hpp:36
rmm::device_scalar::value
value_type value(cuda_stream_view stream) const
Copies the value from device to host, synchronizes, and returns the value.
Definition: device_scalar.hpp:150
rmm::device_scalar::stream
cuda_stream_view stream() const noexcept
Returns stream most recently specified for allocation/deallocation.
Definition: device_scalar.hpp:246
rmm::mr::device_memory_resource
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:82
rmm::device_scalar::device_scalar
device_scalar(device_scalar const &)=delete
Copy ctor is deleted as it doesn't allow a stream argument.
rmm::device_scalar::device_scalar
device_scalar(value_type const &initial_value, cuda_stream_view stream, rmm::mr::device_memory_resource *mr=rmm::mr::get_current_device_resource())
Construct a new device_scalar with an initial value.
Definition: device_scalar.hpp:106