diff --git a/cpp/include/raft/core/memory_stats_resources.hpp b/cpp/include/raft/core/memory_stats_resources.hpp index 136b48beb6..38eaa12270 100644 --- a/cpp/include/raft/core/memory_stats_resources.hpp +++ b/cpp/include/raft/core/memory_stats_resources.hpp @@ -4,6 +4,7 @@ */ #pragma once +#include #include #include #include @@ -77,7 +78,8 @@ class memory_stats_resources : public resources { explicit memory_stats_resources(const resources& existing) : resources(existing), old_host_(mr::get_default_host_resource()), - old_device_(rmm::mr::get_current_device_resource_ref()) + old_device_(rmm::mr::get_per_device_resource_ref( + rmm::cuda_device_id{resource::get_device_id(*this)})) { init(); } @@ -85,7 +87,8 @@ class memory_stats_resources : public resources { ~memory_stats_resources() override { mr::set_default_host_resource(old_host_); - rmm::mr::set_current_device_resource(std::move(old_device_)); + rmm::mr::set_per_device_resource(rmm::cuda_device_id{resource::get_device_id(*this)}, + std::move(old_device_)); } memory_stats_resources(memory_stats_resources const&) = delete; @@ -215,7 +218,8 @@ class memory_stats_resources : public resources { device_stats_adaptor_t sa{rmm::device_async_resource_ref{old_device_}}; device_stats_ = sa.get_stats(); device_adaptor_ = std::make_unique(std::move(sa)); - rmm::mr::set_current_device_resource(*device_adaptor_); + rmm::mr::set_per_device_resource(rmm::cuda_device_id{resource::get_device_id(*this)}, + *device_adaptor_); } // --- Workspace --- { diff --git a/cpp/include/raft/core/memory_tracking_resources.hpp b/cpp/include/raft/core/memory_tracking_resources.hpp index d06e14ef0e..d3a36e4b66 100644 --- a/cpp/include/raft/core/memory_tracking_resources.hpp +++ b/cpp/include/raft/core/memory_tracking_resources.hpp @@ -5,6 +5,7 @@ #pragma once #include +#include #include #include #include @@ -108,7 +109,8 @@ class memory_tracking_resources : public resources { { report_.stop(); raft::mr::set_default_host_resource(old_host_); - rmm::mr::set_current_device_resource(old_device_); + rmm::mr::set_per_device_resource(rmm::cuda_device_id{resource::get_device_id(*this)}, + std::move(old_device_)); } memory_tracking_resources(memory_tracking_resources const&) = delete; @@ -128,7 +130,8 @@ class memory_tracking_resources : public resources { owned_stream_(std::move(owned_stream)), report_(out_override ? *out_override : *owned_stream_, sample_interval), old_host_(raft::mr::get_default_host_resource()), - old_device_(rmm::mr::get_current_device_resource_ref()) + old_device_(rmm::mr::get_per_device_resource_ref( + rmm::cuda_device_id{resource::get_device_id(*this)})) { init(); } @@ -217,7 +220,8 @@ class memory_tracking_resources : public resources { device_stats_t sa{rmm::device_async_resource_ref{old_device_}}; report_.register_source("device", sa.get_stats()); device_adaptor_ = std::make_unique(std::move(sa), report_.get_notifier()); - rmm::mr::set_current_device_resource(*device_adaptor_); + rmm::mr::set_per_device_resource(rmm::cuda_device_id{resource::get_device_id(*this)}, + *device_adaptor_); } // --- Workspace (track upstream to preserve limiting_resource_adaptor) --- diff --git a/cpp/tests/core/memory_stats_resources.cpp b/cpp/tests/core/memory_stats_resources.cpp index be615c9041..0d16f93b64 100644 --- a/cpp/tests/core/memory_stats_resources.cpp +++ b/cpp/tests/core/memory_stats_resources.cpp @@ -4,19 +4,49 @@ */ #include +#include #include #include +#include #include +#include #include +#include #include #include #include +#include namespace raft { +namespace { + +struct device_resource_restore_guard { + int device_id; + raft::mr::device_resource resource; + + ~device_resource_restore_guard() + { + rmm::mr::set_per_device_resource(rmm::cuda_device_id{device_id}, std::move(resource)); + } +}; + +auto current_device_uses_pool_resource() -> bool +{ + auto current_mr = rmm::mr::get_current_device_resource_ref(); + return cuda::mr::resource_cast(¤t_mr) != nullptr; +} + +auto current_device_uses_default_cuda_resource() -> bool +{ + auto current_mr = rmm::mr::get_current_device_resource_ref(); + return cuda::mr::resource_cast(¤t_mr) != nullptr; +} + +} // namespace TEST(MemoryStatsResources, IndependentCounting_DefaultWorkspace) { @@ -93,4 +123,105 @@ TEST(MemoryStatsResources, IndependentCounting_PoolWorkspace) dev_mr.deallocate(cuda::stream_ref{cudaStreamLegacy}, dev_ptr, kGlobalSize); } +TEST(MemoryStatsResources, RestoresDeviceResourceOnConstructionDevice) +{ + if (device_setter::get_device_count() < 2) { + GTEST_SKIP() << "Requires at least 2 CUDA devices"; + } + + auto device0 = 0; + auto device1 = 1; + + auto device0_guard = [&]() { + auto scoped_device = device_setter{device0}; + auto upstream = rmm::mr::get_current_device_resource_ref(); + return device_resource_restore_guard{ + device0, + rmm::mr::set_current_device_resource( + raft::mr::device_resource{rmm::mr::pool_memory_resource(upstream, 1 << 20, 2 << 20)})}; + }(); + + auto device1_guard = [&]() { + auto scoped_device = device_setter{device1}; + return device_resource_restore_guard{device1, rmm::mr::reset_current_device_resource()}; + }(); + + { + auto scoped_device = device_setter{device0}; + raft::resources res; + auto tracked = std::make_unique(res); + auto wrong_device = device_setter{device1}; + static_cast(wrong_device); + tracked.reset(); + } + + { + auto scoped_device = device_setter{device0}; + auto current_mr = rmm::mr::get_current_device_resource_ref(); + EXPECT_NE(cuda::mr::resource_cast(¤t_mr), nullptr); + } + + { + auto scoped_device = device_setter{device1}; + auto current_mr = rmm::mr::get_current_device_resource_ref(); + EXPECT_NE(cuda::mr::resource_cast(¤t_mr), nullptr); + } +} + +TEST(MemoryStatsResources, InstallsTrackedResourceOnHandleDevice) +{ + if (device_setter::get_device_count() < 2) { + GTEST_SKIP() << "Requires at least 2 CUDA devices"; + } + + auto device0 = 0; + auto device1 = 1; + + auto device0_guard = [&]() { + auto scoped_device = device_setter{device0}; + auto upstream = rmm::mr::get_current_device_resource_ref(); + return device_resource_restore_guard{ + device0, + rmm::mr::set_current_device_resource( + raft::mr::device_resource{rmm::mr::pool_memory_resource(upstream, 1 << 20, 2 << 20)})}; + }(); + + auto device1_guard = [&]() { + auto scoped_device = device_setter{device1}; + return device_resource_restore_guard{device1, rmm::mr::reset_current_device_resource()}; + }(); + + { + auto scoped_device = device_setter{device0}; + raft::resources res; + static_cast(resource::get_device_id(res)); + + auto wrong_device = device_setter{device1}; + static_cast(wrong_device); + auto tracked = std::make_unique(res); + + { + auto verify_device0 = device_setter{device0}; + EXPECT_FALSE(current_device_uses_pool_resource()); + } + + { + auto verify_device1 = device_setter{device1}; + EXPECT_TRUE(current_device_uses_default_cuda_resource()); + } + + tracked.reset(); + } + + { + auto scoped_device = device_setter{device0}; + EXPECT_TRUE(current_device_uses_pool_resource()); + } + + { + auto scoped_device = device_setter{device1}; + EXPECT_TRUE(current_device_uses_default_cuda_resource()); + } +} + } // namespace raft diff --git a/cpp/tests/core/monitor_resources.cu b/cpp/tests/core/monitor_resources.cu index c7be4bd285..94f7ff9904 100644 --- a/cpp/tests/core/monitor_resources.cu +++ b/cpp/tests/core/monitor_resources.cu @@ -4,19 +4,48 @@ */ #include +#include #include #include +#include +#include +#include + +#include #include #include #include +#include #include #include #include namespace { +struct device_resource_restore_guard { + int device_id; + raft::mr::device_resource resource; + + ~device_resource_restore_guard() + { + rmm::mr::set_per_device_resource(rmm::cuda_device_id{device_id}, std::move(resource)); + } +}; + +auto current_device_uses_pool_resource() -> bool +{ + auto current_mr = rmm::mr::get_current_device_resource_ref(); + return cuda::mr::resource_cast(¤t_mr) != nullptr; +} + +auto current_device_uses_default_cuda_resource() -> bool +{ + auto current_mr = rmm::mr::get_current_device_resource_ref(); + return cuda::mr::resource_cast(¤t_mr) != nullptr; +} + TEST(MemoryTrackingResources, TracksDeviceAllocations) { using namespace std::chrono_literals; @@ -49,4 +78,106 @@ TEST(MemoryTrackingResources, TracksDeviceAllocations) << output; } +TEST(MemoryTrackingResources, RestoresDeviceResourceOnConstructionDevice) +{ + if (raft::device_setter::get_device_count() < 2) { + GTEST_SKIP() << "Requires at least 2 CUDA devices"; + } + + auto device0 = 0; + auto device1 = 1; + + auto device0_guard = [&]() { + auto scoped_device = raft::device_setter{device0}; + auto upstream = rmm::mr::get_current_device_resource_ref(); + return device_resource_restore_guard{ + device0, + rmm::mr::set_current_device_resource( + raft::mr::device_resource{rmm::mr::pool_memory_resource(upstream, 1 << 20, 2 << 20)})}; + }(); + + auto device1_guard = [&]() { + auto scoped_device = raft::device_setter{device1}; + return device_resource_restore_guard{device1, rmm::mr::reset_current_device_resource()}; + }(); + + { + auto scoped_device = raft::device_setter{device0}; + std::ostringstream oss; + auto tracked = std::make_unique(oss); + auto wrong_device = raft::device_setter{device1}; + static_cast(wrong_device); + tracked.reset(); + } + + { + auto scoped_device = raft::device_setter{device0}; + auto current_mr = rmm::mr::get_current_device_resource_ref(); + EXPECT_NE(cuda::mr::resource_cast(¤t_mr), nullptr); + } + + { + auto scoped_device = raft::device_setter{device1}; + auto current_mr = rmm::mr::get_current_device_resource_ref(); + EXPECT_NE(cuda::mr::resource_cast(¤t_mr), nullptr); + } +} + +TEST(MemoryTrackingResources, InstallsTrackedResourceOnHandleDevice) +{ + if (raft::device_setter::get_device_count() < 2) { + GTEST_SKIP() << "Requires at least 2 CUDA devices"; + } + + auto device0 = 0; + auto device1 = 1; + + auto device0_guard = [&]() { + auto scoped_device = raft::device_setter{device0}; + auto upstream = rmm::mr::get_current_device_resource_ref(); + return device_resource_restore_guard{ + device0, + rmm::mr::set_current_device_resource( + raft::mr::device_resource{rmm::mr::pool_memory_resource(upstream, 1 << 20, 2 << 20)})}; + }(); + + auto device1_guard = [&]() { + auto scoped_device = raft::device_setter{device1}; + return device_resource_restore_guard{device1, rmm::mr::reset_current_device_resource()}; + }(); + + { + auto scoped_device = raft::device_setter{device0}; + raft::resources res; + static_cast(raft::resource::get_device_id(res)); + + auto wrong_device = raft::device_setter{device1}; + static_cast(wrong_device); + std::ostringstream oss; + auto tracked = std::make_unique(res, oss); + + { + auto verify_device0 = raft::device_setter{device0}; + EXPECT_FALSE(current_device_uses_pool_resource()); + } + + { + auto verify_device1 = raft::device_setter{device1}; + EXPECT_TRUE(current_device_uses_default_cuda_resource()); + } + + tracked.reset(); + } + + { + auto scoped_device = raft::device_setter{device0}; + EXPECT_TRUE(current_device_uses_pool_resource()); + } + + { + auto scoped_device = raft::device_setter{device1}; + EXPECT_TRUE(current_device_uses_default_cuda_resource()); + } +} + } // namespace