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-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 #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 #include <cstddef>
25 #include <utility>
26 
27 namespace RMM_NAMESPACE {
28 
29 struct cuda_device_id;
30 inline cuda_device_id get_current_cuda_device();
31 
42  using value_type = int;
43 
47  cuda_device_id() noexcept : id_{get_current_cuda_device().value()} {}
48 
54  explicit constexpr cuda_device_id(value_type dev_id) noexcept : id_{dev_id} {}
55 
57  [[nodiscard]] constexpr value_type value() const noexcept { return id_; }
58 
59  // TODO re-add doxygen comment specifier /** for these hidden friend operators once this Breathe
60  // bug is fixed: https://github.com/breathe-doc/breathe/issues/916
62 
69  [[nodiscard]] constexpr friend bool operator==(cuda_device_id const& lhs,
70  cuda_device_id const& rhs) noexcept
71  {
72  return lhs.value() == rhs.value();
73  }
74 
82  [[nodiscard]] constexpr friend bool operator!=(cuda_device_id const& lhs,
83  cuda_device_id const& rhs) noexcept
84  {
85  return lhs.value() != rhs.value();
86  }
88  private:
89  value_type id_;
90 };
91 
100 {
101  cuda_device_id::value_type dev_id{-1};
102  RMM_ASSERT_CUDA_SUCCESS(cudaGetDevice(&dev_id));
103  return cuda_device_id{dev_id};
104 }
105 
112 {
113  cuda_device_id::value_type num_dev{-1};
114  RMM_ASSERT_CUDA_SUCCESS(cudaGetDeviceCount(&num_dev));
115  return num_dev;
116 }
117 
123 inline std::pair<std::size_t, std::size_t> available_device_memory()
124 {
125  std::size_t free{};
126  std::size_t total{};
127  RMM_CUDA_TRY(cudaMemGetInfo(&free, &total));
128  return {free, total};
129 }
130 
139 inline std::size_t percent_of_free_device_memory(int percent)
140 {
141  [[maybe_unused]] auto const [free, total] = rmm::available_device_memory();
142  auto fraction = static_cast<double>(percent) / 100.0;
143  return rmm::align_down(static_cast<std::size_t>(static_cast<double>(free) * fraction),
145 }
146 
158  : old_device_{get_current_cuda_device()},
159  needs_reset_{dev_id.value() >= 0 && old_device_ != dev_id}
160  {
161  if (needs_reset_) { RMM_ASSERT_CUDA_SUCCESS(cudaSetDevice(dev_id.value())); }
162  }
167  {
168  if (needs_reset_) { RMM_ASSERT_CUDA_SUCCESS(cudaSetDevice(old_device_.value())); }
169  }
170 
172  cuda_set_device_raii& operator=(cuda_set_device_raii const&) = delete;
174  cuda_set_device_raii& operator=(cuda_set_device_raii&&) = delete;
175 
176  private:
177  cuda_device_id old_device_;
178  bool needs_reset_;
179 };
180  // end of group
182 } // 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:123
cuda_device_id get_current_cuda_device()
Returns a cuda_device_id for the current device.
Definition: cuda_device.hpp:99
int get_num_cuda_devices()
Returns the number of CUDA devices in the system.
Definition: cuda_device.hpp:111
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:139
bool operator==(cuda_stream_view lhs, cuda_stream_view rhs)
Equality comparison operator for streams.
Definition: cuda_stream_view.hpp:175
bool operator!=(cuda_stream_view lhs, cuda_stream_view rhs)
Inequality comparison operator for streams.
Definition: cuda_stream_view.hpp:187
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:41
constexpr cuda_device_id(value_type dev_id) noexcept
Construct a cuda_device_id from the specified integer value.
Definition: cuda_device.hpp:54
cuda_device_id() noexcept
Construct a cuda_device_id from the current device.
Definition: cuda_device.hpp:47
constexpr value_type value() const noexcept
The wrapped integer value.
Definition: cuda_device.hpp:57
int value_type
Integer type used for device identifier.
Definition: cuda_device.hpp:42
RAII class that sets the current CUDA device to the specified device on construction and restores the...
Definition: cuda_device.hpp:151
~cuda_set_device_raii() noexcept
Reactivates the previous CUDA device.
Definition: cuda_device.hpp:166
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:157