device_scalar.hpp
1 /*
2  * Copyright (c) 2019, 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_buffer.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 
54  explicit device_scalar(
55  cuda_stream_view const &stream,
56  rmm::mr::device_memory_resource *mr = rmm::mr::get_current_device_resource())
57  : buffer{sizeof(T), stream, mr}
58  {
59  }
60 
77  explicit device_scalar(
78  T const &initial_value,
79  cuda_stream_view const &stream = cuda_stream_view{},
80  rmm::mr::device_memory_resource *mr = rmm::mr::get_current_device_resource())
81  : buffer{sizeof(T), stream, mr}
82  {
83  set_value(initial_value, stream);
84  }
85 
99  cuda_stream_view const &stream = {},
100  rmm::mr::device_memory_resource *mr = rmm::mr::get_current_device_resource())
101  : buffer{other.buffer, stream, mr}
102  {
103  }
104 
121  T value(cuda_stream_view const &stream = cuda_stream_view{}) const
122  {
123  T host_value{};
124  _memcpy(&host_value, buffer.data(), stream);
125  stream.synchronize();
126  return host_value;
127  }
128 
163  template <typename Placeholder = void>
164  auto set_value(T const &host_value, cuda_stream_view const &stream = cuda_stream_view{})
165  -> std::enable_if_t<std::is_fundamental<T>::value, Placeholder>
166  {
167  if (host_value == T{0}) {
168  RMM_CUDA_TRY(cudaMemsetAsync(buffer.data(), 0, sizeof(T), stream.value()));
169  } else {
170  _memcpy(buffer.data(), &host_value, stream);
171  }
172  }
173 
208  template <typename Placeholder = void>
209  auto set_value(T const &host_value, cuda_stream_view const &stream = cuda_stream_view{})
210  -> std::enable_if_t<not std::is_fundamental<T>::value, Placeholder>
211  {
212  _memcpy(buffer.data(), &host_value, stream);
213  }
214 
223  T *data() noexcept { return static_cast<T *>(buffer.data()); }
224 
233  T const *data() const noexcept { return static_cast<T const *>(buffer.data()); }
234 
235  device_scalar() = default;
236  ~device_scalar() = default;
237  device_scalar(device_scalar &&) = default;
238  device_scalar &operator=(device_scalar const &) = delete;
239  device_scalar &operator=(device_scalar &&) = delete;
240 
241  private:
242  rmm::device_buffer buffer{sizeof(T)};
243 
244  inline void _memcpy(void *dst, const void *src, cuda_stream_view const &stream) const
245  {
246  RMM_CUDA_TRY(cudaMemcpyAsync(dst, src, sizeof(T), cudaMemcpyDefault, stream.value()));
247  }
248 };
249 } // namespace rmm
device_buffer.hpp
RAII construct for device memory allocation.
rmm::device_scalar::set_value
auto set_value(T const &host_value, cuda_stream_view const &stream=cuda_stream_view{}) -> std::enable_if_t< std::is_fundamental< T >::value, Placeholder >
Sets the value of the device_scalar to the given host_value.
Definition: device_scalar.hpp:164
per_device_resource.hpp
Management of per-device device_memory_resources.
rmm::cuda_stream_view
Strongly-typed non-owning wrapper for CUDA streams with default constructor.
Definition: cuda_stream_view.hpp:34
rmm::device_scalar::device_scalar
device_scalar(cuda_stream_view const &stream, rmm::mr::device_memory_resource *mr=rmm::mr::get_current_device_resource())
Construct a new uninitialized device_scalar.
Definition: device_scalar.hpp:54
rmm::device_buffer
Definition: device_buffer.hpp:73
rmm::device_scalar::value
T value(cuda_stream_view const &stream=cuda_stream_view{}) const
Copies the value from device to host, synchronizes, and returns the value.
Definition: device_scalar.hpp:121
rmm::device_scalar::device_scalar
device_scalar(device_scalar const &other, cuda_stream_view const &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:98
rmm::device_scalar::data
T * data() noexcept
Returns pointer to object in device memory.
Definition: device_scalar.hpp:223
rmm::device_scalar
Container for a single object of type T in device memory.
Definition: device_scalar.hpp:36
rmm::device_scalar::device_scalar
device_scalar(T const &initial_value, cuda_stream_view const &stream=cuda_stream_view{}, 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:77
rmm::cuda_stream_view::synchronize
void synchronize() const
Synchronize the viewed CUDA stream.
Definition: cuda_stream_view.hpp:96
rmm::device_buffer::data
void const * data() const noexcept
Returns raw pointer to underlying device memory allocation.
Definition: device_buffer.hpp:335
rmm::cuda_stream_view::value
constexpr cudaStream_t value() const noexcept
Get the wrapped stream.
Definition: cuda_stream_view.hpp:58
rmm::mr::device_memory_resource
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:83
rmm::device_scalar::set_value
auto set_value(T const &host_value, cuda_stream_view const &stream=cuda_stream_view{}) -> std::enable_if_t< not std::is_fundamental< T >::value, Placeholder >
Sets the value of the device_scalar to the given host_value.
Definition: device_scalar.hpp:209
rmm::device_scalar::data
T const * data() const noexcept
Returns const pointer to object in device memory.
Definition: device_scalar.hpp:233