per_device_resource.hpp
Go to the documentation of this file.
1 /*
2  * Copyright (c) 2020-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_device.hpp>
20 #include <rmm/detail/export.hpp>
21 #include <rmm/mr/device/cuda_memory_resource.hpp>
22 #include <rmm/mr/device/device_memory_resource.hpp>
23 
24 #include <map>
25 #include <mutex>
26 
74 namespace rmm::mr {
75 
76 namespace detail {
77 
86 {
87  static cuda_memory_resource mr{};
88  return &mr;
89 }
90 
91 inline std::mutex& map_lock()
92 {
93  static std::mutex map_lock;
94  return map_lock;
95 }
96 
97 // Must have default visibility, see: https://github.com/rapidsai/rmm/issues/826
98 RMM_EXPORT inline auto& get_map()
99 {
100  static std::map<cuda_device_id::value_type, device_memory_resource*> device_id_to_resource;
101  return device_id_to_resource;
102 }
103 
104 } // namespace detail
105 
128 {
129  std::lock_guard<std::mutex> lock{detail::map_lock()};
130  auto& map = detail::get_map();
131  // If a resource was never set for `id`, set to the initial resource
132  auto const found = map.find(device_id.value());
133  return (found == map.end()) ? (map[device_id.value()] = detail::initial_resource())
134  : found->second;
135 }
136 
165  device_memory_resource* new_mr)
166 {
167  std::lock_guard<std::mutex> lock{detail::map_lock()};
168  auto& map = detail::get_map();
169  auto const old_itr = map.find(device_id.value());
170  // If a resource didn't previously exist for `id`, return pointer to initial_resource
171  auto* old_mr = (old_itr == map.end()) ? detail::initial_resource() : old_itr->second;
172  map[device_id.value()] = (new_mr == nullptr) ? detail::initial_resource() : new_mr;
173  return old_mr;
174 }
175 
198 {
199  return get_per_device_resource(rmm::detail::current_device());
200 }
201 
227 {
228  return set_per_device_resource(rmm::detail::current_device(), new_mr);
229 }
230 } // namespace rmm::mr
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:226
rmm::mr::detail::initial_resource
device_memory_resource * initial_resource()
Returns a pointer to the initial resource.
Definition: per_device_resource.hpp:85
rmm::mr::set_per_device_resource
device_memory_resource * set_per_device_resource(cuda_device_id device_id, device_memory_resource *new_mr)
Set the device_memory_resource for the specified device.
Definition: per_device_resource.hpp:164
rmm::mr::get_per_device_resource
device_memory_resource * get_per_device_resource(cuda_device_id device_id)
Get the resource for the specified device.
Definition: per_device_resource.hpp:127
rmm::mr::device_memory_resource
Base class for all libcudf device memory allocation.
Definition: device_memory_resource.hpp:82
rmm::cuda_device_id
Strong type for a CUDA device identifier.
Definition: cuda_device.hpp:28
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:197
rmm::cuda_device_id::value
constexpr value_type value() const noexcept
Returns the wrapped integer value.
Definition: cuda_device.hpp:39
rmm::mr::cuda_memory_resource
device_memory_resource derived class that uses cudaMalloc/Free for allocation/deallocation.
Definition: cuda_memory_resource.hpp:30