All Classes Files Functions Variables Typedefs Enumerations Enumerator Friends Modules Pages
cuda_device.hpp
Go to the documentation of this file.
1 /*
2  * Copyright (c) 2021-2024, 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 #pragma once
17 
18 #include <rmm/aligned.hpp>
19 #include <rmm/detail/error.hpp>
20 #include <rmm/detail/export.hpp>
21 
22 #include <cuda_runtime_api.h>
23 
24 namespace RMM_NAMESPACE {
25 
26 struct cuda_device_id;
27 inline cuda_device_id get_current_cuda_device();
28 
39  using value_type = int;
40 
44  cuda_device_id() noexcept : id_{get_current_cuda_device().value()} {}
45 
51  explicit constexpr cuda_device_id(value_type dev_id) noexcept : id_{dev_id} {}
52 
54  [[nodiscard]] constexpr value_type value() const noexcept { return id_; }
55 
56  // TODO re-add doxygen comment specifier /** for these hidden friend operators once this Breathe
57  // bug is fixed: https://github.com/breathe-doc/breathe/issues/916
59 
66  [[nodiscard]] constexpr friend bool operator==(cuda_device_id const& lhs,
67  cuda_device_id const& rhs) noexcept
68  {
69  return lhs.value() == rhs.value();
70  }
71 
79  [[nodiscard]] constexpr friend bool operator!=(cuda_device_id const& lhs,
80  cuda_device_id const& rhs) noexcept
81  {
82  return lhs.value() != rhs.value();
83  }
85  private:
86  value_type id_;
87 };
88 
97 {
98  cuda_device_id::value_type dev_id{-1};
99  RMM_ASSERT_CUDA_SUCCESS(cudaGetDevice(&dev_id));
100  return cuda_device_id{dev_id};
101 }
102 
109 {
110  cuda_device_id::value_type num_dev{-1};
111  RMM_ASSERT_CUDA_SUCCESS(cudaGetDeviceCount(&num_dev));
112  return num_dev;
113 }
114 
120 inline std::pair<std::size_t, std::size_t> available_device_memory()
121 {
122  std::size_t free{};
123  std::size_t total{};
124  RMM_CUDA_TRY(cudaMemGetInfo(&free, &total));
125  return {free, total};
126 }
127 
136 inline std::size_t percent_of_free_device_memory(int percent)
137 {
138  [[maybe_unused]] auto const [free, total] = rmm::available_device_memory();
139  auto fraction = static_cast<double>(percent) / 100.0;
140  return rmm::align_down(static_cast<std::size_t>(static_cast<double>(free) * fraction),
142 }
143 
155  : old_device_{get_current_cuda_device()},
156  needs_reset_{dev_id.value() >= 0 && old_device_ != dev_id}
157  {
158  if (needs_reset_) { RMM_ASSERT_CUDA_SUCCESS(cudaSetDevice(dev_id.value())); }
159  }
164  {
165  if (needs_reset_) { RMM_ASSERT_CUDA_SUCCESS(cudaSetDevice(old_device_.value())); }
166  }
167 
169  cuda_set_device_raii& operator=(cuda_set_device_raii const&) = delete;
171  cuda_set_device_raii& operator=(cuda_set_device_raii&&) = delete;
172 
173  private:
174  cuda_device_id old_device_;
175  bool needs_reset_;
176 };
177  // end of group
179 } // namespace RMM_NAMESPACE
std::pair< std::size_t, std::size_t > available_device_memory()
Returns the available and total device memory in bytes for the current device.
Definition: cuda_device.hpp:120
cuda_device_id get_current_cuda_device()
Returns a cuda_device_id for the current device.
Definition: cuda_device.hpp:96
int get_num_cuda_devices()
Returns the number of CUDA devices in the system.
Definition: cuda_device.hpp:108
std::size_t percent_of_free_device_memory(int percent)
Returns the approximate specified percent of available device memory on the current CUDA device,...
Definition: cuda_device.hpp:136
bool operator==(cuda_stream_view lhs, cuda_stream_view rhs)
Equality comparison operator for streams.
Definition: cuda_stream_view.hpp:177
bool operator!=(cuda_stream_view lhs, cuda_stream_view rhs)
Inequality comparison operator for streams.
Definition: cuda_stream_view.hpp:189
static constexpr std::size_t CUDA_ALLOCATION_ALIGNMENT
Default alignment used for CUDA memory allocation.
Definition: aligned.hpp:43
constexpr std::size_t align_down(std::size_t value, std::size_t alignment) noexcept
Align down to the nearest multiple of specified power of 2.
Definition: aligned.hpp:91
Strong type for a CUDA device identifier.
Definition: cuda_device.hpp:38
constexpr cuda_device_id(value_type dev_id) noexcept
Construct a cuda_device_id from the specified integer value.
Definition: cuda_device.hpp:51
cuda_device_id() noexcept
Construct a cuda_device_id from the current device.
Definition: cuda_device.hpp:44
constexpr value_type value() const noexcept
The wrapped integer value.
Definition: cuda_device.hpp:54
int value_type
Integer type used for device identifier.
Definition: cuda_device.hpp:39
RAII class that sets the current CUDA device to the specified device on construction and restores the...
Definition: cuda_device.hpp:148
~cuda_set_device_raii() noexcept
Reactivates the previous CUDA device.
Definition: cuda_device.hpp:163
cuda_set_device_raii(cuda_device_id dev_id)
Construct a new cuda_set_device_raii object and sets the current CUDA device to dev_id
Definition: cuda_device.hpp:154