per_device_resource.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <rmm/cuda_device.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/mr/cuda_memory_resource.hpp>
#include <rmm/mr/device_memory_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <cuda/memory_resource>
 
#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, cuda::mr::any_resource<cuda::mr::device_accessible>>
    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;
}
 
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)
{
  using any_device_resource = cuda::mr::any_resource<cuda::mr::device_accessible>;
  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 ref to initial_resource
  if (old_itr == map.end()) {
    map.emplace(device_id.value(), static_cast<any_device_resource>(new_resource_ref));
    return device_async_resource_ref{*detail::initial_resource()};
  }
 
  device_async_resource_ref old_resource_ref{old_itr->second};
  old_itr->second = static_cast<any_device_resource>(new_resource_ref);  // reify and store
  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.
  if (new_mr != nullptr) { 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)
{
  using any_device_resource = cuda::mr::any_resource<cuda::mr::device_accessible>;
  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()) {
    // Create a resource_ref from the initial resource, then reify it to any_resource
    device_async_resource_ref initial_ref{*detail::initial_resource()};
    auto item = map.emplace(device_id.value(), static_cast<any_device_resource>(initial_ref));
    return device_async_resource_ref{item.first->second};
  }
  return device_async_resource_ref{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