per_device_resource.hpp

Go to the documentation of this file.

/*
 * Copyright (c) 2020-2021, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <rmm/cuda_device.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/mr/device/cuda_memory_resource.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
 
#include <map>
#include <mutex>
 
namespace rmm::mr {
namespace detail {
 
inline device_memory_resource* initial_resource()
{
  static cuda_memory_resource mr{};
  return &mr;
}
 
inline std::mutex& map_lock()
{
  static std::mutex map_lock;
  return map_lock;
}
 
// This symbol must have default visibility, see: https://github.com/rapidsai/rmm/issues/826
RMM_EXPORT inline auto& get_map()
{
  static std::map<cuda_device_id::value_type, device_memory_resource*> device_id_to_resource;
  return device_id_to_resource;
}
 
}  // namespace detail
 
inline device_memory_resource* get_per_device_resource(cuda_device_id device_id)
{
  std::lock_guard<std::mutex> lock{detail::map_lock()};
  auto& map = detail::get_map();
  // If a resource was never set for `id`, set to the initial resource
  auto const found = map.find(device_id.value());
  return (found == map.end()) ? (map[device_id.value()] = detail::initial_resource())
                              : found->second;
}
 
inline device_memory_resource* set_per_device_resource(cuda_device_id device_id,
                                                       device_memory_resource* new_mr)
{
  std::lock_guard<std::mutex> lock{detail::map_lock()};
  auto& map          = detail::get_map();
  auto const old_itr = map.find(device_id.value());
  // If a resource didn't previously exist for `id`, return pointer to initial_resource
  auto* old_mr           = (old_itr == map.end()) ? detail::initial_resource() : old_itr->second;
  map[device_id.value()] = (new_mr == nullptr) ? detail::initial_resource() : new_mr;
  return old_mr;
}
 
inline device_memory_resource* get_current_device_resource()
{
  return get_per_device_resource(rmm::get_current_cuda_device());
}
 
inline device_memory_resource* set_current_device_resource(device_memory_resource* new_mr)
{
  return set_per_device_resource(rmm::get_current_cuda_device(), new_mr);
}  // end of group
}  // namespace rmm::mr