per_device_resource.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2020-2024, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#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 <rmm/resource_ref.hpp>
 
#include <map>
#include <mutex>
 
namespace RMM_NAMESPACE {
namespace mr {
namespace detail {
 
// These symbols must have default visibility so that when they are
// referenced in multiple different DSOs the linker correctly
// determines that there is only a single unique reference to the
// function symbols (and hence they return unique static references
// across different DSOs). See also
// https://github.com/rapidsai/rmm/issues/826
// Although currently the entire RMM namespace is RMM_EXPORT, we
// explicitly mark these functions as exported in case the namespace
// export changes.
RMM_EXPORT inline device_memory_resource* initial_resource()
{
  static cuda_memory_resource mr{};
  return &mr;
}
 
RMM_EXPORT inline std::mutex& map_lock()
{
  static std::mutex map_lock;
  return map_lock;
}
 
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;
}
 
RMM_EXPORT inline std::mutex& ref_map_lock()
{
  static std::mutex ref_map_lock;
  return ref_map_lock;
}
 
// This symbol must have default visibility, see: https://github.com/rapidsai/rmm/issues/826
RMM_EXPORT inline auto& get_ref_map()
{
  static std::map<cuda_device_id::value_type, device_async_resource_ref> device_id_to_resource_ref;
  return device_id_to_resource_ref;
}
 
}  // 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;
}
 
namespace detail {
 
// The non-thread-safe implementation of `set_per_device_resource_ref`. This exists because
// we need to call this function from two places: the thread-safe version of
// `set_per_device_resource_ref` and the thread-safe version of `set_per_device_resource`,
// both of which take the lock, so we need an implementation that doesn't take the lock.
inline device_async_resource_ref set_per_device_resource_ref_unsafe(
  cuda_device_id device_id, device_async_resource_ref new_resource_ref)
{
  auto& map          = detail::get_ref_map();
  auto const old_itr = map.find(device_id.value());
  // If a resource didn't previously exist for `device_id`, return pointer to initial_resource
  // Note: because resource_ref is not default-constructible, we can't use std::map::operator[]
  if (old_itr == map.end()) {
    map.insert({device_id.value(), new_resource_ref});
    return device_async_resource_ref{detail::initial_resource()};
  }
 
  auto old_resource_ref = old_itr->second;
  old_itr->second       = new_resource_ref;  // update map directly via iterator
  return old_resource_ref;
}
}  // namespace detail
 
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()};
 
  // Note: even though set_per_device_resource() and set_per_device_resource_ref() are not
  // interchangeable, we call the latter from the former to maintain resource_ref
  // state consistent with the resource pointer state. This is necessary because the
  // Python API still uses the raw pointer API. Once the Python API is updated to use
  // resource_ref, this call can be removed.
  detail::set_per_device_resource_ref_unsafe(device_id, new_mr);
 
  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);
}
 
inline device_async_resource_ref get_per_device_resource_ref(cuda_device_id device_id)
{
  std::lock_guard<std::mutex> lock{detail::ref_map_lock()};
  auto& map = detail::get_ref_map();
  // If a resource was never set for `id`, set to the initial resource
  auto const found = map.find(device_id.value());
  if (found == map.end()) {
    auto item = map.insert({device_id.value(), detail::initial_resource()});
    return item.first->second;
  }
  return found->second;
}
 
inline device_async_resource_ref set_per_device_resource_ref(
  cuda_device_id device_id, device_async_resource_ref new_resource_ref)
{
  std::lock_guard<std::mutex> lock{detail::ref_map_lock()};
  return detail::set_per_device_resource_ref_unsafe(device_id, new_resource_ref);
}
 
inline device_async_resource_ref get_current_device_resource_ref()
{
  return get_per_device_resource_ref(rmm::get_current_cuda_device());
}
 
inline device_async_resource_ref set_current_device_resource_ref(
  device_async_resource_ref new_resource_ref)
{
  return set_per_device_resource_ref(rmm::get_current_cuda_device(), new_resource_ref);
}
 
inline device_async_resource_ref reset_per_device_resource_ref(cuda_device_id device_id)
{
  return set_per_device_resource_ref(device_id, detail::initial_resource());
}
 
inline device_async_resource_ref reset_current_device_resource_ref()
{
  return reset_per_device_resource_ref(rmm::get_current_cuda_device());
}  // end of group
}  // namespace mr
}  // namespace RMM_NAMESPACE