per_device_resource.hpp

查看此文件的文档。

/*
 * Copyright (c) 2020-2024, 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
 *
 * https://apache.ac.cn/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 <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