per_device_resource.hpp
Go to the documentation of this file.
1 /*
2  * Copyright (c) 2020, 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/mr/device/cuda_memory_resource.hpp>
20 #include <rmm/mr/device/device_memory_resource.hpp>
21 
22 #include <map>
23 #include <mutex>
24 
72 namespace rmm {
73 
79  using value_type = int;
80 
86  explicit constexpr cuda_device_id(value_type id) noexcept : id_{id} {}
87 
89  constexpr value_type value() const noexcept { return id_; }
90 
91  private:
92  value_type id_;
93 };
94 
95 namespace mr {
96 
97 namespace detail {
98 
107 {
108  static cuda_memory_resource mr{};
109  return &mr;
110 }
111 
112 inline std::mutex& map_lock()
113 {
114  static std::mutex map_lock;
115  return map_lock;
116 }
117 
118 inline auto& get_map()
119 {
120  static std::map<cuda_device_id::value_type, device_memory_resource*> device_id_to_resource;
121  return device_id_to_resource;
122 }
123 
132 {
133  int dev_id;
134  RMM_CUDA_TRY(cudaGetDevice(&dev_id));
135  return cuda_device_id{dev_id};
136 }
137 } // namespace detail
138 
161 {
162  std::lock_guard<std::mutex> lock{detail::map_lock()};
163  auto& map = detail::get_map();
164  // If a resource was never set for `id`, set to the initial resource
165  auto const found = map.find(id.value());
166  return (found == map.end()) ? (map[id.value()] = detail::initial_resource()) : found->second;
167 }
168 
197  device_memory_resource* new_mr)
198 {
199  std::lock_guard<std::mutex> lock{detail::map_lock()};
200  auto& map = detail::get_map();
201  auto const old_itr = map.find(id.value());
202  // If a resource didn't previously exist for `id`, return pointer to initial_resource
203  auto old_mr = (old_itr == map.end()) ? detail::initial_resource() : old_itr->second;
204  map[id.value()] = (new_mr == nullptr) ? detail::initial_resource() : new_mr;
205  return old_mr;
206 }
207 
230 {
231  return get_per_device_resource(detail::current_device());
232 }
233 
259 {
260  return set_per_device_resource(detail::current_device(), new_mr);
261 }
262 } // namespace mr
263 } // namespace rmm
rmm::mr::set_current_device_resource
device_memory_resource * set_current_device_resource(device_memory_resource *new_mr)
Set the memory resource for the current device.
Definition: per_device_resource.hpp:258
rmm::mr::set_per_device_resource
device_memory_resource * set_per_device_resource(cuda_device_id id, device_memory_resource *new_mr)
Set the device_memory_resource for the specified device.
Definition: per_device_resource.hpp:196
rmm::mr::detail::initial_resource
device_memory_resource * initial_resource()
Returns a pointer to the initial resource.
Definition: per_device_resource.hpp:106
rmm::mr::detail::current_device
cuda_device_id current_device()
Returns a cuda_device_id for the current device.
Definition: per_device_resource.hpp:131
rmm::cuda_device_id::cuda_device_id
constexpr cuda_device_id(value_type id) noexcept
Construct a cuda_device_id from the specified integer value.
Definition: per_device_resource.hpp:86
rmm::mr::get_per_device_resource
device_memory_resource * get_per_device_resource(cuda_device_id id)
Get the resource for the specified device.
Definition: per_device_resource.hpp:160
rmm::mr::device_memory_resource
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:83
rmm::cuda_device_id
Strong type for a CUDA device identifier.
Definition: per_device_resource.hpp:78
rmm::mr::get_current_device_resource
device_memory_resource * get_current_device_resource()
Get the memory resource for the current device.
Definition: per_device_resource.hpp:229
rmm::cuda_device_id::value
constexpr value_type value() const noexcept
Returns the wrapped integer value.
Definition: per_device_resource.hpp:89
rmm::mr::cuda_memory_resource
device_memory_resource derived class that uses cudaMalloc/Free for allocation/deallocation.
Definition: cuda_memory_resource.hpp:29