per_device_resource.hpp

Go to the documentation of this file.

/*
 * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */
 
#pragma once
 
#include <rmm/cuda_device.hpp>
#include <rmm/detail/export.hpp>
#include <rmm/detail/runtime_shutdown.hpp>
#include <rmm/mr/cuda_memory_resource.hpp>
#include <rmm/resource_ref.hpp>
 
#include <cuda/memory_resource>
 
#include <map>
#include <mutex>
#include <utility>

 
RMM_NAMESPACE_BEGIN
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 cuda_memory_resource& initial_resource()
{
  static cuda_memory_resource mr{};
  return mr;
}

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;
  // Register the process-exit hook immediately after constructing the map.
  rmm::detail::register_process_exit_hook();
  return device_id_to_resource;
}
 
}  // namespace detail

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.emplace(device_id.value(), detail::initial_resource());
    return device_async_resource_ref{item.first->second};
  }
  return device_async_resource_ref{found->second};
}

inline cuda::mr::any_resource<cuda::mr::device_accessible> set_per_device_resource(
  cuda_device_id device_id, cuda::mr::any_resource<cuda::mr::device_accessible> new_resource)
{
  std::lock_guard<std::mutex> lock{detail::ref_map_lock()};
  auto& map          = detail::get_ref_map();
  auto const old_itr = map.find(device_id.value());
  if (old_itr == map.end()) {
    map.emplace(device_id.value(), std::move(new_resource));
    return {detail::initial_resource()};
  }
  return std::exchange(old_itr->second, std::move(new_resource));
}

inline device_async_resource_ref get_current_device_resource_ref()
{
  return get_per_device_resource_ref(rmm::get_current_cuda_device());
}

inline cuda::mr::any_resource<cuda::mr::device_accessible> set_current_device_resource(
  cuda::mr::any_resource<cuda::mr::device_accessible> new_resource)
{
  return set_per_device_resource(rmm::get_current_cuda_device(), std::move(new_resource));
}

inline cuda::mr::any_resource<cuda::mr::device_accessible> reset_per_device_resource(
  cuda_device_id device_id)
{
  return set_per_device_resource(device_id, {detail::initial_resource()});
}

inline cuda::mr::any_resource<cuda::mr::device_accessible> reset_current_device_resource()
{
  return reset_per_device_resource(rmm::get_current_cuda_device());
}
  // end of group
}  // namespace mr
RMM_NAMESPACE_END