From 75fd465d5aacff6b12e32c71c33b2c59d5d0b728 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Tue, 2 Dec 2025 23:53:48 -0600 Subject: [PATCH 1/2] Split MR_REF tests --- cpp/tests/CMakeLists.txt | 24 +- cpp/tests/mr/mr_ref_arena_tests.cpp | 37 +++ cpp/tests/mr/mr_ref_binning_tests.cpp | 37 +++ cpp/tests/mr/mr_ref_cuda_async_tests.cpp | 37 +++ cpp/tests/mr/mr_ref_cuda_tests.cpp | 36 +++ cpp/tests/mr/mr_ref_default_tests.cpp | 111 +++++++ cpp/tests/mr/mr_ref_fixed_size_tests.cpp | 30 ++ cpp/tests/mr/mr_ref_managed_tests.cpp | 36 +++ cpp/tests/mr/mr_ref_multithreaded_tests.cpp | 317 ------------------ cpp/tests/mr/mr_ref_pinned_pool_tests.cpp | 36 +++ cpp/tests/mr/mr_ref_pinned_tests.cpp | 36 +++ cpp/tests/mr/mr_ref_pool_tests.cpp | 37 +++ cpp/tests/mr/mr_ref_system_tests.cpp | 36 +++ cpp/tests/mr/mr_ref_test.hpp | 335 ++++++++++++++++++++ cpp/tests/mr/mr_ref_tests.cpp | 140 -------- 15 files changed, 825 insertions(+), 460 deletions(-) create mode 100644 cpp/tests/mr/mr_ref_arena_tests.cpp create mode 100644 cpp/tests/mr/mr_ref_binning_tests.cpp create mode 100644 cpp/tests/mr/mr_ref_cuda_async_tests.cpp create mode 100644 cpp/tests/mr/mr_ref_cuda_tests.cpp create mode 100644 cpp/tests/mr/mr_ref_default_tests.cpp create mode 100644 cpp/tests/mr/mr_ref_fixed_size_tests.cpp create mode 100644 cpp/tests/mr/mr_ref_managed_tests.cpp delete mode 100644 cpp/tests/mr/mr_ref_multithreaded_tests.cpp create mode 100644 cpp/tests/mr/mr_ref_pinned_pool_tests.cpp create mode 100644 cpp/tests/mr/mr_ref_pinned_tests.cpp create mode 100644 cpp/tests/mr/mr_ref_pool_tests.cpp create mode 100644 cpp/tests/mr/mr_ref_system_tests.cpp delete mode 100644 cpp/tests/mr/mr_ref_tests.cpp diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index cdd8657bb..2b1773640 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -133,9 +133,27 @@ endfunction() # test sources -# device mr_ref tests -ConfigureTest(DEVICE_MR_REF_TEST mr/mr_ref_tests.cpp mr/mr_ref_multithreaded_tests.cpp GPUS 1 - PERCENT 100) +# device mr_ref tests - split by memory resource type for parallel execution PERCENT values based on +# maximum memory allocation per test (assuming 16GB GPU) + +# Full test suites with multi-threading (~2GB max allocation = 15%) +ConfigureTest(CUDA_MR_REF_TEST mr/mr_ref_cuda_tests.cpp GPUS 1 PERCENT 15) +ConfigureTest(CUDA_ASYNC_MR_REF_TEST mr/mr_ref_cuda_async_tests.cpp GPUS 1 PERCENT 15) +ConfigureTest(MANAGED_MR_REF_TEST mr/mr_ref_managed_tests.cpp GPUS 1 PERCENT 15) + +# Pre-allocating resources (50% of free device memory = 50%) +ConfigureTest(POOL_MR_REF_TEST mr/mr_ref_pool_tests.cpp GPUS 1 PERCENT 50) +ConfigureTest(ARENA_MR_REF_TEST mr/mr_ref_arena_tests.cpp GPUS 1 PERCENT 50) +ConfigureTest(BINNING_MR_REF_TEST mr/mr_ref_binning_tests.cpp GPUS 1 PERCENT 50) + +# Partial test suites without multi-threading (~500MB max = 5%) +ConfigureTest(SYSTEM_MR_REF_TEST mr/mr_ref_system_tests.cpp GPUS 1 PERCENT 5) +ConfigureTest(PINNED_MR_REF_TEST mr/mr_ref_pinned_tests.cpp GPUS 1 PERCENT 5) +ConfigureTest(PINNED_POOL_MR_REF_TEST mr/mr_ref_pinned_pool_tests.cpp GPUS 1 PERCENT 5) + +# Small test suites (minimal allocations = 5%) +ConfigureTest(FIXED_SIZE_MR_REF_TEST mr/mr_ref_fixed_size_tests.cpp GPUS 1 PERCENT 5) +ConfigureTest(DEFAULT_MR_REF_TEST mr/mr_ref_default_tests.cpp GPUS 1 PERCENT 5) # general adaptor tests ConfigureTest(ADAPTOR_TEST mr/adaptor_tests.cpp) diff --git a/cpp/tests/mr/mr_ref_arena_tests.cpp b/cpp/tests/mr/mr_ref_arena_tests.cpp new file mode 100644 index 000000000..33a86b954 --- /dev/null +++ b/cpp/tests/mr/mr_ref_arena_tests.cpp @@ -0,0 +1,37 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test.hpp" + +#include +#include +#include +#include + +#include + +namespace rmm::test { +namespace { + +// Single-threaded basic tests (5 tests) +INSTANTIATE_TEST_SUITE_P(ArenaResourceTests, + mr_ref_test, + ::testing::Values("Arena"), + [](auto const& info) { return info.param; }); + +// Single-threaded allocation tests (9 tests) +INSTANTIATE_TEST_SUITE_P(ArenaResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("Arena"), + [](auto const& info) { return info.param; }); + +// Multi-threaded tests (15 tests) +INSTANTIATE_TEST_SUITE_P(ArenaMultiThreadResourceTests, + mr_ref_test_mt, + ::testing::Values("Arena"), + [](auto const& info) { return info.param; }); + +} // namespace +} // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_binning_tests.cpp b/cpp/tests/mr/mr_ref_binning_tests.cpp new file mode 100644 index 000000000..14df36d8b --- /dev/null +++ b/cpp/tests/mr/mr_ref_binning_tests.cpp @@ -0,0 +1,37 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test.hpp" + +#include +#include +#include +#include + +#include + +namespace rmm::test { +namespace { + +// Single-threaded basic tests (5 tests) +INSTANTIATE_TEST_SUITE_P(BinningResourceTests, + mr_ref_test, + ::testing::Values("Binning"), + [](auto const& info) { return info.param; }); + +// Single-threaded allocation tests (9 tests) +INSTANTIATE_TEST_SUITE_P(BinningResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("Binning"), + [](auto const& info) { return info.param; }); + +// Multi-threaded tests (15 tests) +INSTANTIATE_TEST_SUITE_P(BinningMultiThreadResourceTests, + mr_ref_test_mt, + ::testing::Values("Binning"), + [](auto const& info) { return info.param; }); + +} // namespace +} // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_cuda_async_tests.cpp b/cpp/tests/mr/mr_ref_cuda_async_tests.cpp new file mode 100644 index 000000000..6e19a3f97 --- /dev/null +++ b/cpp/tests/mr/mr_ref_cuda_async_tests.cpp @@ -0,0 +1,37 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test.hpp" + +#include +#include +#include +#include + +#include + +namespace rmm::test { +namespace { + +// Single-threaded basic tests (5 tests) +INSTANTIATE_TEST_SUITE_P(CudaAsyncResourceTests, + mr_ref_test, + ::testing::Values("CUDA_Async"), + [](auto const& info) { return info.param; }); + +// Single-threaded allocation tests (9 tests) +INSTANTIATE_TEST_SUITE_P(CudaAsyncResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("CUDA_Async"), + [](auto const& info) { return info.param; }); + +// Multi-threaded tests (15 tests) +INSTANTIATE_TEST_SUITE_P(CudaAsyncMultiThreadResourceTests, + mr_ref_test_mt, + ::testing::Values("CUDA_Async"), + [](auto const& info) { return info.param; }); + +} // namespace +} // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_cuda_tests.cpp b/cpp/tests/mr/mr_ref_cuda_tests.cpp new file mode 100644 index 000000000..ca159c535 --- /dev/null +++ b/cpp/tests/mr/mr_ref_cuda_tests.cpp @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test.hpp" + +#include +#include +#include + +#include + +namespace rmm::test { +namespace { + +// Single-threaded basic tests (5 tests) +INSTANTIATE_TEST_SUITE_P(CudaResourceTests, + mr_ref_test, + ::testing::Values("CUDA"), + [](auto const& info) { return info.param; }); + +// Single-threaded allocation tests (9 tests) +INSTANTIATE_TEST_SUITE_P(CudaResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("CUDA"), + [](auto const& info) { return info.param; }); + +// Multi-threaded tests (15 tests) +INSTANTIATE_TEST_SUITE_P(CudaMultiThreadResourceTests, + mr_ref_test_mt, + ::testing::Values("CUDA"), + [](auto const& info) { return info.param; }); + +} // namespace +} // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_default_tests.cpp b/cpp/tests/mr/mr_ref_default_tests.cpp new file mode 100644 index 000000000..d8d9ca815 --- /dev/null +++ b/cpp/tests/mr/mr_ref_default_tests.cpp @@ -0,0 +1,111 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test.hpp" + +#include +#include +#include + +#include + +#include +#include + +namespace rmm::test { +namespace { + +// Suppress warnings about uninstantiated parameterized tests in this file +GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_test); +GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_allocation_test); +GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_test_mt); + +// Helper functions for multi-threaded tests +template +void spawn_n(std::size_t num_threads, Task task, Arguments&&... args) +{ + std::vector threads; + threads.reserve(num_threads); + for (std::size_t i = 0; i < num_threads; ++i) { + threads.emplace_back(std::thread(task, std::forward(args)...)); + } + + for (auto& thread : threads) { + thread.join(); + } +} + +template +void spawn(Task task, Arguments&&... args) +{ + spawn_n(4, task, std::forward(args)...); +} + +// Single-threaded default resource tests + +TEST(DefaultTest, CurrentDeviceResourceIsCUDA) +{ + EXPECT_NE(nullptr, rmm::mr::get_current_device_resource()); + EXPECT_TRUE(rmm::mr::get_current_device_resource()->is_equal(rmm::mr::cuda_memory_resource{})); +} + +TEST(DefaultTest, UseCurrentDeviceResource) { test_get_current_device_resource(); } + +TEST(DefaultTest, UseCurrentDeviceResourceRef) { test_get_current_device_resource_ref(); } + +TEST(DefaultTest, GetCurrentDeviceResource) +{ + auto* mr = rmm::mr::get_current_device_resource(); + EXPECT_NE(nullptr, mr); + EXPECT_TRUE(mr->is_equal(rmm::mr::cuda_memory_resource{})); +} + +TEST(DefaultTest, GetCurrentDeviceResourceRef) +{ + auto mr = rmm::mr::get_current_device_resource_ref(); + EXPECT_EQ(mr, rmm::device_async_resource_ref{rmm::mr::detail::initial_resource()}); +} + +// Multi-threaded default resource tests + +TEST(DefaultTest, UseCurrentDeviceResource_mt) { spawn(test_get_current_device_resource); } + +TEST(DefaultTest, UseCurrentDeviceResourceRef_mt) { spawn(test_get_current_device_resource_ref); } + +TEST(DefaultTest, CurrentDeviceResourceIsCUDA_mt) +{ + spawn([]() { + EXPECT_NE(nullptr, rmm::mr::get_current_device_resource()); + EXPECT_TRUE(rmm::mr::get_current_device_resource()->is_equal(rmm::mr::cuda_memory_resource{})); + }); +} + +TEST(DefaultTest, CurrentDeviceResourceRefIsCUDA_mt) +{ + spawn([]() { + EXPECT_EQ(rmm::mr::get_current_device_resource_ref(), + rmm::device_async_resource_ref{rmm::mr::detail::initial_resource()}); + }); +} + +TEST(DefaultTest, GetCurrentDeviceResource_mt) +{ + spawn([]() { + rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource(); + EXPECT_NE(nullptr, mr); + EXPECT_TRUE(mr->is_equal(rmm::mr::cuda_memory_resource{})); + }); +} + +TEST(DefaultTest, GetCurrentDeviceResourceRef_mt) +{ + spawn([]() { + auto mr = rmm::mr::get_current_device_resource_ref(); + EXPECT_EQ(mr, rmm::device_async_resource_ref{rmm::mr::detail::initial_resource()}); + }); +} + +} // namespace +} // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_fixed_size_tests.cpp b/cpp/tests/mr/mr_ref_fixed_size_tests.cpp new file mode 100644 index 000000000..6d45b117d --- /dev/null +++ b/cpp/tests/mr/mr_ref_fixed_size_tests.cpp @@ -0,0 +1,30 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test.hpp" + +#include +#include +#include +#include + +#include + +namespace rmm::test { +namespace { + +// Suppress warnings about uninstantiated tests (Fixed_Size only has basic tests) +GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_allocation_test); +GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_test_mt); + +// Single-threaded basic tests (5 tests) +// Note: Fixed_Size MR cannot handle dynamic allocation sizes, so only basic tests are included +INSTANTIATE_TEST_SUITE_P(FixedSizeResourceTests, + mr_ref_test, + ::testing::Values("Fixed_Size"), + [](auto const& info) { return info.param; }); + +} // namespace +} // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_managed_tests.cpp b/cpp/tests/mr/mr_ref_managed_tests.cpp new file mode 100644 index 000000000..e4179b079 --- /dev/null +++ b/cpp/tests/mr/mr_ref_managed_tests.cpp @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test.hpp" + +#include +#include +#include + +#include + +namespace rmm::test { +namespace { + +// Single-threaded basic tests (5 tests) +INSTANTIATE_TEST_SUITE_P(ManagedResourceTests, + mr_ref_test, + ::testing::Values("Managed"), + [](auto const& info) { return info.param; }); + +// Single-threaded allocation tests (9 tests) +INSTANTIATE_TEST_SUITE_P(ManagedResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("Managed"), + [](auto const& info) { return info.param; }); + +// Multi-threaded tests (15 tests) +INSTANTIATE_TEST_SUITE_P(ManagedMultiThreadResourceTests, + mr_ref_test_mt, + ::testing::Values("Managed"), + [](auto const& info) { return info.param; }); + +} // namespace +} // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_multithreaded_tests.cpp b/cpp/tests/mr/mr_ref_multithreaded_tests.cpp deleted file mode 100644 index 6f097ba6c..000000000 --- a/cpp/tests/mr/mr_ref_multithreaded_tests.cpp +++ /dev/null @@ -1,317 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "mr_ref_test.hpp" - -#include -#include -#include -#include -#include -#include -#include - -#include - -#include -#include - -namespace rmm::test { -namespace { - -struct mr_ref_test_mt : public mr_ref_test {}; - -INSTANTIATE_TEST_SUITE_P( - MultiThreadResourceTests, - mr_ref_test_mt, - ::testing::Values("CUDA", "CUDA_Async", "Managed", "Pool", "Arena", "Binning"), - [](auto const& info) { return info.param; }); - -template -void spawn_n(std::size_t num_threads, Task task, Arguments&&... args) -{ - std::vector threads; - threads.reserve(num_threads); - for (std::size_t i = 0; i < num_threads; ++i) { - threads.emplace_back(std::thread(task, std::forward(args)...)); - } - - for (auto& thread : threads) { - thread.join(); - } -} - -template -void spawn(Task task, Arguments&&... args) -{ - spawn_n(4, task, std::forward(args)...); -} - -TEST(DefaultTest, UseCurrentDeviceResource_mt) { spawn(test_get_current_device_resource); } - -TEST(DefaultTest, UseCurrentDeviceResourceRef_mt) { spawn(test_get_current_device_resource_ref); } - -TEST(DefaultTest, CurrentDeviceResourceIsCUDA_mt) -{ - spawn([]() { - EXPECT_NE(nullptr, rmm::mr::get_current_device_resource()); - EXPECT_TRUE(rmm::mr::get_current_device_resource()->is_equal(rmm::mr::cuda_memory_resource{})); - }); -} - -TEST(DefaultTest, CurrentDeviceResourceRefIsCUDA_mt) -{ - spawn([]() { - EXPECT_EQ(rmm::mr::get_current_device_resource_ref(), - rmm::device_async_resource_ref{rmm::mr::detail::initial_resource()}); - }); -} - -TEST(DefaultTest, GetCurrentDeviceResource_mt) -{ - spawn([]() { - rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource(); - EXPECT_NE(nullptr, mr); - EXPECT_TRUE(mr->is_equal(rmm::mr::cuda_memory_resource{})); - }); -} - -TEST(DefaultTest, GetCurrentDeviceResourceRef_mt) -{ - spawn([]() { - auto mr = rmm::mr::get_current_device_resource_ref(); - EXPECT_EQ(mr, rmm::device_async_resource_ref{rmm::mr::detail::initial_resource()}); - }); -} - -TEST_P(mr_ref_test_mt, SetCurrentDeviceResourceRef_mt) -{ - // single thread changes default resource, then multiple threads use it - auto old = rmm::mr::set_current_device_resource_ref(this->ref); - test_get_current_device_resource_ref(); - - int device; - RMM_CUDA_TRY(cudaGetDevice(&device)); - - spawn([device, mr = this->ref]() { - RMM_CUDA_TRY(cudaSetDevice(device)); - EXPECT_EQ(mr, rmm::mr::get_current_device_resource_ref()); - test_get_current_device_resource_ref(); // test allocating with the new default resource - }); - - // resetting default resource should reset to initial - rmm::mr::reset_current_device_resource_ref(); - EXPECT_EQ(old, rmm::mr::get_current_device_resource_ref()); -} - -TEST_P(mr_ref_test_mt, SetCurrentDeviceResourceRefPerThread_mt) -{ - int num_devices{}; - RMM_CUDA_TRY(cudaGetDeviceCount(&num_devices)); - - std::vector threads; - threads.reserve(num_devices); - - auto mr = this->ref; - - for (int i = 0; i < num_devices; ++i) { - threads.emplace_back( - [mr](auto dev_id) { - RMM_CUDA_TRY(cudaSetDevice(dev_id)); - auto cuda_ref = rmm::mr::get_current_device_resource_ref(); - auto old = rmm::mr::set_current_device_resource_ref(mr); - - // initial resource for this device should be CUDA mr - EXPECT_EQ(old, cuda_ref); - // get_current_device_resource_ref should equal the resource we - // just set - EXPECT_EQ(mr, rmm::mr::get_current_device_resource_ref()); - // Resetting current dev resource ref should make it - // cuda MR and return the MR we previously set - old = rmm::mr::reset_current_device_resource_ref(); - EXPECT_EQ(old, mr); - EXPECT_EQ(cuda_ref, rmm::mr::get_current_device_resource_ref()); - }, - i); - } - - for (auto& thread : threads) { - thread.join(); - } -} - -TEST_P(mr_ref_test_mt, Allocate) -{ - int device; - RMM_CUDA_TRY(cudaGetDevice(&device)); - - auto mr = this->ref; - spawn([device, mr]() { - RMM_CUDA_TRY(cudaSetDevice(device)); - test_various_allocations(mr); - }); -} - -TEST_P(mr_ref_test_mt, AllocateDefaultStream) -{ - spawn(test_various_async_allocations, this->ref, rmm::cuda_stream_view{}); -} - -TEST_P(mr_ref_test_mt, AllocateOnStream) -{ - spawn(test_various_async_allocations, this->ref, this->stream.view()); -} - -TEST_P(mr_ref_test_mt, RandomAllocations) -{ - spawn(test_random_allocations, this->ref, default_num_allocations, default_max_size); -} - -TEST_P(mr_ref_test_mt, RandomAllocationsDefaultStream) -{ - spawn(test_random_async_allocations, - this->ref, - default_num_allocations, - default_max_size, - rmm::cuda_stream_view{}); -} - -TEST_P(mr_ref_test_mt, RandomAllocationsStream) -{ - spawn(test_random_async_allocations, - this->ref, - default_num_allocations, - default_max_size, - this->stream.view()); -} - -TEST_P(mr_ref_test_mt, MixedRandomAllocationFree) -{ - spawn(test_mixed_random_allocation_free, this->ref, default_max_size); -} - -TEST_P(mr_ref_test_mt, MixedRandomAllocationFreeDefaultStream) -{ - spawn( - test_mixed_random_async_allocation_free, this->ref, default_max_size, rmm::cuda_stream_view{}); -} - -TEST_P(mr_ref_test_mt, MixedRandomAllocationFreeStream) -{ - spawn(test_mixed_random_async_allocation_free, this->ref, default_max_size, this->stream.view()); -} - -void async_allocate_loop(rmm::device_async_resource_ref ref, - std::size_t num_allocations, - std::list& allocations, - std::mutex& mtx, - std::condition_variable& allocations_ready, - cudaEvent_t& event, - rmm::cuda_stream_view stream) -{ - constexpr std::size_t max_size{1_MiB}; - - std::default_random_engine generator; - std::uniform_int_distribution size_distribution(1, max_size); - - for (std::size_t i = 0; i < num_allocations; ++i) { - std::size_t size = size_distribution(generator); - void* ptr = ref.allocate(stream, size); - { - std::lock_guard lock(mtx); - RMM_CUDA_TRY(cudaEventRecord(event, stream.value())); - allocations.emplace_back(ptr, size); - } - allocations_ready.notify_one(); - } - - // Work around for threads going away before cudaEvent has finished async processing - cudaEventSynchronize(event); -} - -void async_deallocate_loop(rmm::device_async_resource_ref ref, - std::size_t num_allocations, - std::list& allocations, - std::mutex& mtx, - std::condition_variable& allocations_ready, - cudaEvent_t& event, - rmm::cuda_stream_view stream) -{ - for (std::size_t i = 0; i < num_allocations; i++) { - std::unique_lock lock(mtx); - allocations_ready.wait(lock, [&allocations] { return !allocations.empty(); }); - RMM_CUDA_TRY(cudaStreamWaitEvent(stream.value(), event)); - allocation alloc = allocations.front(); - allocations.pop_front(); - ref.deallocate(stream, alloc.ptr, alloc.size); - } - - // Work around for threads going away before cudaEvent has finished async processing - cudaEventSynchronize(event); -} - -void test_async_allocate_free_different_threads(rmm::device_async_resource_ref ref, - rmm::cuda_stream_view streamA, - rmm::cuda_stream_view streamB) -{ - constexpr std::size_t num_allocations{100}; - - std::mutex mtx; - std::condition_variable allocations_ready; - std::list allocations; - cudaEvent_t event{}; - - RMM_CUDA_TRY(cudaEventCreate(&event)); - - std::thread producer(async_allocate_loop, - ref, - num_allocations, - std::ref(allocations), - std::ref(mtx), - std::ref(allocations_ready), - std::ref(event), - streamA); - - std::thread consumer(async_deallocate_loop, - ref, - num_allocations, - std::ref(allocations), - std::ref(mtx), - std::ref(allocations_ready), - std::ref(event), - streamB); - - producer.join(); - consumer.join(); - - RMM_CUDA_TRY(cudaEventDestroy(event)); -} - -TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsDefaultStream) -{ - test_async_allocate_free_different_threads( - this->ref, rmm::cuda_stream_default, rmm::cuda_stream_default); -} - -TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsPerThreadDefaultStream) -{ - test_async_allocate_free_different_threads( - this->ref, rmm::cuda_stream_per_thread, rmm::cuda_stream_per_thread); -} - -TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsSameStream) -{ - test_async_allocate_free_different_threads(this->ref, this->stream, this->stream); -} - -TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsDifferentStream) -{ - rmm::cuda_stream streamB; - test_async_allocate_free_different_threads(this->ref, this->stream, streamB); - streamB.synchronize(); -} - -} // namespace -} // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_pinned_pool_tests.cpp b/cpp/tests/mr/mr_ref_pinned_pool_tests.cpp new file mode 100644 index 000000000..f2874ab5f --- /dev/null +++ b/cpp/tests/mr/mr_ref_pinned_pool_tests.cpp @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test.hpp" + +#include +#include +#include +#include + +#include + +namespace rmm::test { +namespace { + +// Suppress warning about uninstantiated multi-threaded tests (PinnedPool doesn't support MT tests) +GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_test_mt); + +// Single-threaded basic tests (5 tests) +INSTANTIATE_TEST_SUITE_P(PinnedPoolResourceTests, + mr_ref_test, + ::testing::Values("PinnedPool"), + [](auto const& info) { return info.param; }); + +// Single-threaded allocation tests (9 tests) +INSTANTIATE_TEST_SUITE_P(PinnedPoolResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("PinnedPool"), + [](auto const& info) { return info.param; }); + +// Note: No multi-threaded tests for PinnedPool memory resource + +} // namespace +} // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_pinned_tests.cpp b/cpp/tests/mr/mr_ref_pinned_tests.cpp new file mode 100644 index 000000000..5aedd3918 --- /dev/null +++ b/cpp/tests/mr/mr_ref_pinned_tests.cpp @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test.hpp" + +#include +#include +#include +#include + +#include + +namespace rmm::test { +namespace { + +// Suppress warning about uninstantiated multi-threaded tests (Pinned doesn't support MT tests) +GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_test_mt); + +// Single-threaded basic tests (5 tests) +INSTANTIATE_TEST_SUITE_P(PinnedResourceTests, + mr_ref_test, + ::testing::Values("Pinned"), + [](auto const& info) { return info.param; }); + +// Single-threaded allocation tests (9 tests) +INSTANTIATE_TEST_SUITE_P(PinnedResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("Pinned"), + [](auto const& info) { return info.param; }); + +// Note: No multi-threaded tests for Pinned memory resource + +} // namespace +} // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_pool_tests.cpp b/cpp/tests/mr/mr_ref_pool_tests.cpp new file mode 100644 index 000000000..9920b4df6 --- /dev/null +++ b/cpp/tests/mr/mr_ref_pool_tests.cpp @@ -0,0 +1,37 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test.hpp" + +#include +#include +#include +#include + +#include + +namespace rmm::test { +namespace { + +// Single-threaded basic tests (5 tests) +INSTANTIATE_TEST_SUITE_P(PoolResourceTests, + mr_ref_test, + ::testing::Values("Pool"), + [](auto const& info) { return info.param; }); + +// Single-threaded allocation tests (9 tests) +INSTANTIATE_TEST_SUITE_P(PoolResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("Pool"), + [](auto const& info) { return info.param; }); + +// Multi-threaded tests (15 tests) +INSTANTIATE_TEST_SUITE_P(PoolMultiThreadResourceTests, + mr_ref_test_mt, + ::testing::Values("Pool"), + [](auto const& info) { return info.param; }); + +} // namespace +} // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_system_tests.cpp b/cpp/tests/mr/mr_ref_system_tests.cpp new file mode 100644 index 000000000..ae9dc20ac --- /dev/null +++ b/cpp/tests/mr/mr_ref_system_tests.cpp @@ -0,0 +1,36 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test.hpp" + +#include +#include +#include +#include + +#include + +namespace rmm::test { +namespace { + +// Suppress warning about uninstantiated multi-threaded tests (System doesn't support MT tests) +GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_test_mt); + +// Single-threaded basic tests (5 tests) +INSTANTIATE_TEST_SUITE_P(SystemResourceTests, + mr_ref_test, + ::testing::Values("System"), + [](auto const& info) { return info.param; }); + +// Single-threaded allocation tests (9 tests) +INSTANTIATE_TEST_SUITE_P(SystemResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("System"), + [](auto const& info) { return info.param; }); + +// Note: No multi-threaded tests for System memory resource + +} // namespace +} // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_test.hpp b/cpp/tests/mr/mr_ref_test.hpp index 3d25434e6..702eee194 100644 --- a/cpp/tests/mr/mr_ref_test.hpp +++ b/cpp/tests/mr/mr_ref_test.hpp @@ -517,4 +517,339 @@ struct mr_ref_test : public ::testing::TestWithParam { struct mr_ref_allocation_test : public mr_ref_test {}; +// Multi-threaded test fixture +struct mr_ref_test_mt : public mr_ref_test {}; + +// Parameterized test definitions for mr_ref_test + +TEST_P(mr_ref_test, SetCurrentDeviceResourceRef) +{ + rmm::mr::cuda_memory_resource cuda_mr{}; + auto cuda_ref = rmm::device_async_resource_ref{cuda_mr}; + + rmm::mr::set_current_device_resource_ref(cuda_ref); + auto old = rmm::mr::set_current_device_resource_ref(this->ref); + + // old mr should equal a cuda mr + EXPECT_EQ(old, cuda_ref); + + // current dev resource should equal this resource + EXPECT_EQ(this->ref, rmm::mr::get_current_device_resource_ref()); + + test_get_current_device_resource_ref(); + + // Resetting should reset to initial cuda resource + rmm::mr::reset_current_device_resource_ref(); + EXPECT_EQ(rmm::device_async_resource_ref{rmm::mr::detail::initial_resource()}, + rmm::mr::get_current_device_resource_ref()); +} + +TEST_P(mr_ref_test, SelfEquality) { EXPECT_TRUE(this->ref == this->ref); } + +// Simple reproducer for https://github.com/rapidsai/rmm/issues/861 +TEST_P(mr_ref_test, AllocationsAreDifferent) { concurrent_allocations_are_different(this->ref); } + +TEST_P(mr_ref_test, AsyncAllocationsAreDifferentDefaultStream) +{ + concurrent_async_allocations_are_different(this->ref, cuda_stream_view{}); +} + +TEST_P(mr_ref_test, AsyncAllocationsAreDifferent) +{ + concurrent_async_allocations_are_different(this->ref, this->stream); +} + +// Parameterized test definitions for mr_ref_allocation_test + +TEST_P(mr_ref_allocation_test, AllocateDefault) { test_various_allocations(this->ref); } + +TEST_P(mr_ref_allocation_test, AllocateDefaultStream) +{ + test_various_async_allocations(this->ref, cuda_stream_view{}); +} + +TEST_P(mr_ref_allocation_test, AllocateOnStream) +{ + test_various_async_allocations(this->ref, this->stream); +} + +TEST_P(mr_ref_allocation_test, RandomAllocations) { test_random_allocations(this->ref); } + +TEST_P(mr_ref_allocation_test, RandomAllocationsDefaultStream) +{ + test_random_async_allocations( + this->ref, default_num_allocations, default_max_size, cuda_stream_view{}); +} + +TEST_P(mr_ref_allocation_test, RandomAllocationsStream) +{ + test_random_async_allocations(this->ref, default_num_allocations, default_max_size, this->stream); +} + +TEST_P(mr_ref_allocation_test, MixedRandomAllocationFree) +{ + test_mixed_random_allocation_free(this->ref, default_max_size); +} + +TEST_P(mr_ref_allocation_test, MixedRandomAllocationFreeDefaultStream) +{ + test_mixed_random_async_allocation_free(this->ref, default_max_size, cuda_stream_view{}); +} + +TEST_P(mr_ref_allocation_test, MixedRandomAllocationFreeStream) +{ + test_mixed_random_async_allocation_free(this->ref, default_max_size, this->stream); +} + +// Helper functions for multi-threaded tests + +template +void spawn_n(std::size_t num_threads, Task task, Arguments&&... args) +{ + std::vector threads; + threads.reserve(num_threads); + for (std::size_t i = 0; i < num_threads; ++i) { + threads.emplace_back(std::thread(task, std::forward(args)...)); + } + + for (auto& thread : threads) { + thread.join(); + } +} + +template +void spawn(Task task, Arguments&&... args) +{ + spawn_n(4, task, std::forward(args)...); +} + +inline void async_allocate_loop(rmm::device_async_resource_ref ref, + std::size_t num_allocations, + std::list& allocations, + std::mutex& mtx, + std::condition_variable& allocations_ready, + cudaEvent_t& event, + rmm::cuda_stream_view stream) +{ + constexpr std::size_t max_size{1_MiB}; + + std::default_random_engine generator; + std::uniform_int_distribution size_distribution(1, max_size); + + for (std::size_t i = 0; i < num_allocations; ++i) { + std::size_t size = size_distribution(generator); + void* ptr = ref.allocate(stream, size); + { + std::lock_guard lock(mtx); + RMM_CUDA_TRY(cudaEventRecord(event, stream.value())); + allocations.emplace_back(ptr, size); + } + allocations_ready.notify_one(); + } + + // Work around for threads going away before cudaEvent has finished async processing + cudaEventSynchronize(event); +} + +inline void async_deallocate_loop(rmm::device_async_resource_ref ref, + std::size_t num_allocations, + std::list& allocations, + std::mutex& mtx, + std::condition_variable& allocations_ready, + cudaEvent_t& event, + rmm::cuda_stream_view stream) +{ + for (std::size_t i = 0; i < num_allocations; i++) { + std::unique_lock lock(mtx); + allocations_ready.wait(lock, [&allocations] { return !allocations.empty(); }); + RMM_CUDA_TRY(cudaStreamWaitEvent(stream.value(), event)); + allocation alloc = allocations.front(); + allocations.pop_front(); + ref.deallocate(stream, alloc.ptr, alloc.size); + } + + // Work around for threads going away before cudaEvent has finished async processing + cudaEventSynchronize(event); +} + +inline void test_async_allocate_free_different_threads(rmm::device_async_resource_ref ref, + rmm::cuda_stream_view streamA, + rmm::cuda_stream_view streamB) +{ + constexpr std::size_t num_allocations{100}; + + std::mutex mtx; + std::condition_variable allocations_ready; + std::list allocations; + cudaEvent_t event{}; + + RMM_CUDA_TRY(cudaEventCreate(&event)); + + std::thread producer(async_allocate_loop, + ref, + num_allocations, + std::ref(allocations), + std::ref(mtx), + std::ref(allocations_ready), + std::ref(event), + streamA); + + std::thread consumer(async_deallocate_loop, + ref, + num_allocations, + std::ref(allocations), + std::ref(mtx), + std::ref(allocations_ready), + std::ref(event), + streamB); + + producer.join(); + consumer.join(); + + RMM_CUDA_TRY(cudaEventDestroy(event)); +} + +// Parameterized test definitions for mr_ref_test_mt + +TEST_P(mr_ref_test_mt, SetCurrentDeviceResourceRef_mt) +{ + // single thread changes default resource, then multiple threads use it + auto old = rmm::mr::set_current_device_resource_ref(this->ref); + test_get_current_device_resource_ref(); + + int device; + RMM_CUDA_TRY(cudaGetDevice(&device)); + + spawn([device, mr = this->ref]() { + RMM_CUDA_TRY(cudaSetDevice(device)); + EXPECT_EQ(mr, rmm::mr::get_current_device_resource_ref()); + test_get_current_device_resource_ref(); // test allocating with the new default resource + }); + + // resetting default resource should reset to initial + rmm::mr::reset_current_device_resource_ref(); + EXPECT_EQ(old, rmm::mr::get_current_device_resource_ref()); +} + +TEST_P(mr_ref_test_mt, SetCurrentDeviceResourceRefPerThread_mt) +{ + int num_devices{}; + RMM_CUDA_TRY(cudaGetDeviceCount(&num_devices)); + + std::vector threads; + threads.reserve(num_devices); + + auto mr = this->ref; + + for (int i = 0; i < num_devices; ++i) { + threads.emplace_back( + [mr](auto dev_id) { + RMM_CUDA_TRY(cudaSetDevice(dev_id)); + auto cuda_ref = rmm::mr::get_current_device_resource_ref(); + auto old = rmm::mr::set_current_device_resource_ref(mr); + + // initial resource for this device should be CUDA mr + EXPECT_EQ(old, cuda_ref); + // get_current_device_resource_ref should equal the resource we + // just set + EXPECT_EQ(mr, rmm::mr::get_current_device_resource_ref()); + // Resetting current dev resource ref should make it + // cuda MR and return the MR we previously set + old = rmm::mr::reset_current_device_resource_ref(); + EXPECT_EQ(old, mr); + EXPECT_EQ(cuda_ref, rmm::mr::get_current_device_resource_ref()); + }, + i); + } + + for (auto& thread : threads) { + thread.join(); + } +} + +TEST_P(mr_ref_test_mt, Allocate) +{ + int device; + RMM_CUDA_TRY(cudaGetDevice(&device)); + + auto mr = this->ref; + spawn([device, mr]() { + RMM_CUDA_TRY(cudaSetDevice(device)); + test_various_allocations(mr); + }); +} + +TEST_P(mr_ref_test_mt, AllocateDefaultStream) +{ + spawn(test_various_async_allocations, this->ref, rmm::cuda_stream_view{}); +} + +TEST_P(mr_ref_test_mt, AllocateOnStream) +{ + spawn(test_various_async_allocations, this->ref, this->stream.view()); +} + +TEST_P(mr_ref_test_mt, RandomAllocations) +{ + spawn(test_random_allocations, this->ref, default_num_allocations, default_max_size); +} + +TEST_P(mr_ref_test_mt, RandomAllocationsDefaultStream) +{ + spawn(test_random_async_allocations, + this->ref, + default_num_allocations, + default_max_size, + rmm::cuda_stream_view{}); +} + +TEST_P(mr_ref_test_mt, RandomAllocationsStream) +{ + spawn(test_random_async_allocations, + this->ref, + default_num_allocations, + default_max_size, + this->stream.view()); +} + +TEST_P(mr_ref_test_mt, MixedRandomAllocationFree) +{ + spawn(test_mixed_random_allocation_free, this->ref, default_max_size); +} + +TEST_P(mr_ref_test_mt, MixedRandomAllocationFreeDefaultStream) +{ + spawn( + test_mixed_random_async_allocation_free, this->ref, default_max_size, rmm::cuda_stream_view{}); +} + +TEST_P(mr_ref_test_mt, MixedRandomAllocationFreeStream) +{ + spawn(test_mixed_random_async_allocation_free, this->ref, default_max_size, this->stream.view()); +} + +TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsDefaultStream) +{ + test_async_allocate_free_different_threads( + this->ref, rmm::cuda_stream_default, rmm::cuda_stream_default); +} + +TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsPerThreadDefaultStream) +{ + test_async_allocate_free_different_threads( + this->ref, rmm::cuda_stream_per_thread, rmm::cuda_stream_per_thread); +} + +TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsSameStream) +{ + test_async_allocate_free_different_threads(this->ref, this->stream, this->stream); +} + +TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsDifferentStream) +{ + rmm::cuda_stream streamB; + test_async_allocate_free_different_threads(this->ref, this->stream, streamB); + streamB.synchronize(); +} + } // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_tests.cpp b/cpp/tests/mr/mr_ref_tests.cpp deleted file mode 100644 index 7ba0442db..000000000 --- a/cpp/tests/mr/mr_ref_tests.cpp +++ /dev/null @@ -1,140 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "mr_ref_test.hpp" - -#include -#include -#include - -#include - -namespace rmm::test { -namespace { - -INSTANTIATE_TEST_SUITE_P(ResourceTests, - mr_ref_test, - ::testing::Values("CUDA", - "CUDA_Async", - "Managed", - "System", - "Pool", - "Pinned", - "PinnedPool", - "Arena", - "Binning", - "Fixed_Size"), - [](auto const& info) { return info.param; }); - -// Leave out fixed-size MR here because it can't handle the dynamic allocation sizes -INSTANTIATE_TEST_SUITE_P( - ResourceAllocationTests, - mr_ref_allocation_test, - ::testing::Values( - "CUDA", "CUDA_Async", "Managed", "System", "Pool", "Pinned", "PinnedPool", "Arena", "Binning"), - [](auto const& info) { return info.param; }); - -TEST(DefaultTest, CurrentDeviceResourceIsCUDA) -{ - EXPECT_NE(nullptr, rmm::mr::get_current_device_resource()); - EXPECT_TRUE(rmm::mr::get_current_device_resource()->is_equal(rmm::mr::cuda_memory_resource{})); -} - -TEST(DefaultTest, UseCurrentDeviceResource) { test_get_current_device_resource(); } - -TEST(DefaultTest, UseCurrentDeviceResourceRef) { test_get_current_device_resource_ref(); } - -TEST(DefaultTest, GetCurrentDeviceResource) -{ - auto* mr = rmm::mr::get_current_device_resource(); - EXPECT_NE(nullptr, mr); - EXPECT_TRUE(mr->is_equal(rmm::mr::cuda_memory_resource{})); -} - -TEST(DefaultTest, GetCurrentDeviceResourceRef) -{ - auto mr = rmm::mr::get_current_device_resource_ref(); - EXPECT_EQ(mr, rmm::device_async_resource_ref{rmm::mr::detail::initial_resource()}); -} - -TEST_P(mr_ref_test, SetCurrentDeviceResourceRef) -{ - rmm::mr::cuda_memory_resource cuda_mr{}; - auto cuda_ref = rmm::device_async_resource_ref{cuda_mr}; - - rmm::mr::set_current_device_resource_ref(cuda_ref); - auto old = rmm::mr::set_current_device_resource_ref(this->ref); - - // old mr should equal a cuda mr - EXPECT_EQ(old, cuda_ref); - - // current dev resource should equal this resource - EXPECT_EQ(this->ref, rmm::mr::get_current_device_resource_ref()); - - test_get_current_device_resource_ref(); - - // Resetting should reset to initial cuda resource - rmm::mr::reset_current_device_resource_ref(); - EXPECT_EQ(rmm::device_async_resource_ref{rmm::mr::detail::initial_resource()}, - rmm::mr::get_current_device_resource_ref()); -} - -TEST_P(mr_ref_test, SelfEquality) { EXPECT_TRUE(this->ref == this->ref); } - -// Simple reproducer for https://github.com/rapidsai/rmm/issues/861 -TEST_P(mr_ref_test, AllocationsAreDifferent) { concurrent_allocations_are_different(this->ref); } - -TEST_P(mr_ref_test, AsyncAllocationsAreDifferentDefaultStream) -{ - concurrent_async_allocations_are_different(this->ref, cuda_stream_view{}); -} - -TEST_P(mr_ref_test, AsyncAllocationsAreDifferent) -{ - concurrent_async_allocations_are_different(this->ref, this->stream); -} - -TEST_P(mr_ref_allocation_test, AllocateDefault) { test_various_allocations(this->ref); } - -TEST_P(mr_ref_allocation_test, AllocateDefaultStream) -{ - test_various_async_allocations(this->ref, cuda_stream_view{}); -} - -TEST_P(mr_ref_allocation_test, AllocateOnStream) -{ - test_various_async_allocations(this->ref, this->stream); -} - -TEST_P(mr_ref_allocation_test, RandomAllocations) { test_random_allocations(this->ref); } - -TEST_P(mr_ref_allocation_test, RandomAllocationsDefaultStream) -{ - test_random_async_allocations( - this->ref, default_num_allocations, default_max_size, cuda_stream_view{}); -} - -TEST_P(mr_ref_allocation_test, RandomAllocationsStream) -{ - test_random_async_allocations(this->ref, default_num_allocations, default_max_size, this->stream); -} - -TEST_P(mr_ref_allocation_test, MixedRandomAllocationFree) -{ - test_mixed_random_allocation_free(this->ref, default_max_size); -} - -TEST_P(mr_ref_allocation_test, MixedRandomAllocationFreeDefaultStream) -{ - test_mixed_random_async_allocation_free(this->ref, default_max_size, cuda_stream_view{}); -} - -TEST_P(mr_ref_allocation_test, MixedRandomAllocationFreeStream) -{ - test_mixed_random_async_allocation_free(this->ref, default_max_size, this->stream); -} - -} // namespace -} // namespace rmm::test From cad95211dd02b965a2ab06840419e383557a1834 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Mon, 8 Dec 2025 14:44:08 -0600 Subject: [PATCH 2/2] Clean up mr_ref test splits --- cpp/tests/mr/mr_ref_arena_tests.cpp | 14 +- cpp/tests/mr/mr_ref_binning_tests.cpp | 14 +- cpp/tests/mr/mr_ref_cuda_async_tests.cpp | 14 +- cpp/tests/mr/mr_ref_cuda_tests.cpp | 13 +- cpp/tests/mr/mr_ref_default_tests.cpp | 5 - cpp/tests/mr/mr_ref_fixed_size_tests.cpp | 14 +- cpp/tests/mr/mr_ref_managed_tests.cpp | 13 +- cpp/tests/mr/mr_ref_pinned_pool_tests.cpp | 17 +- cpp/tests/mr/mr_ref_pinned_tests.cpp | 17 +- cpp/tests/mr/mr_ref_pool_tests.cpp | 14 +- cpp/tests/mr/mr_ref_system_tests.cpp | 17 +- cpp/tests/mr/mr_ref_test.hpp | 335 ---------------------- cpp/tests/mr/mr_ref_test_allocation.hpp | 54 ++++ cpp/tests/mr/mr_ref_test_basic.hpp | 51 ++++ cpp/tests/mr/mr_ref_test_mt.hpp | 272 ++++++++++++++++++ 15 files changed, 402 insertions(+), 462 deletions(-) create mode 100644 cpp/tests/mr/mr_ref_test_allocation.hpp create mode 100644 cpp/tests/mr/mr_ref_test_basic.hpp create mode 100644 cpp/tests/mr/mr_ref_test_mt.hpp diff --git a/cpp/tests/mr/mr_ref_arena_tests.cpp b/cpp/tests/mr/mr_ref_arena_tests.cpp index 33a86b954..3a73c78cd 100644 --- a/cpp/tests/mr/mr_ref_arena_tests.cpp +++ b/cpp/tests/mr/mr_ref_arena_tests.cpp @@ -3,31 +3,23 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "mr_ref_test.hpp" - -#include -#include -#include -#include - -#include +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" +#include "mr_ref_test_mt.hpp" namespace rmm::test { namespace { -// Single-threaded basic tests (5 tests) INSTANTIATE_TEST_SUITE_P(ArenaResourceTests, mr_ref_test, ::testing::Values("Arena"), [](auto const& info) { return info.param; }); -// Single-threaded allocation tests (9 tests) INSTANTIATE_TEST_SUITE_P(ArenaResourceAllocationTests, mr_ref_allocation_test, ::testing::Values("Arena"), [](auto const& info) { return info.param; }); -// Multi-threaded tests (15 tests) INSTANTIATE_TEST_SUITE_P(ArenaMultiThreadResourceTests, mr_ref_test_mt, ::testing::Values("Arena"), diff --git a/cpp/tests/mr/mr_ref_binning_tests.cpp b/cpp/tests/mr/mr_ref_binning_tests.cpp index 14df36d8b..0c6308e05 100644 --- a/cpp/tests/mr/mr_ref_binning_tests.cpp +++ b/cpp/tests/mr/mr_ref_binning_tests.cpp @@ -3,31 +3,23 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "mr_ref_test.hpp" - -#include -#include -#include -#include - -#include +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" +#include "mr_ref_test_mt.hpp" namespace rmm::test { namespace { -// Single-threaded basic tests (5 tests) INSTANTIATE_TEST_SUITE_P(BinningResourceTests, mr_ref_test, ::testing::Values("Binning"), [](auto const& info) { return info.param; }); -// Single-threaded allocation tests (9 tests) INSTANTIATE_TEST_SUITE_P(BinningResourceAllocationTests, mr_ref_allocation_test, ::testing::Values("Binning"), [](auto const& info) { return info.param; }); -// Multi-threaded tests (15 tests) INSTANTIATE_TEST_SUITE_P(BinningMultiThreadResourceTests, mr_ref_test_mt, ::testing::Values("Binning"), diff --git a/cpp/tests/mr/mr_ref_cuda_async_tests.cpp b/cpp/tests/mr/mr_ref_cuda_async_tests.cpp index 6e19a3f97..ce08183ba 100644 --- a/cpp/tests/mr/mr_ref_cuda_async_tests.cpp +++ b/cpp/tests/mr/mr_ref_cuda_async_tests.cpp @@ -3,31 +3,23 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "mr_ref_test.hpp" - -#include -#include -#include -#include - -#include +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" +#include "mr_ref_test_mt.hpp" namespace rmm::test { namespace { -// Single-threaded basic tests (5 tests) INSTANTIATE_TEST_SUITE_P(CudaAsyncResourceTests, mr_ref_test, ::testing::Values("CUDA_Async"), [](auto const& info) { return info.param; }); -// Single-threaded allocation tests (9 tests) INSTANTIATE_TEST_SUITE_P(CudaAsyncResourceAllocationTests, mr_ref_allocation_test, ::testing::Values("CUDA_Async"), [](auto const& info) { return info.param; }); -// Multi-threaded tests (15 tests) INSTANTIATE_TEST_SUITE_P(CudaAsyncMultiThreadResourceTests, mr_ref_test_mt, ::testing::Values("CUDA_Async"), diff --git a/cpp/tests/mr/mr_ref_cuda_tests.cpp b/cpp/tests/mr/mr_ref_cuda_tests.cpp index ca159c535..5c9616998 100644 --- a/cpp/tests/mr/mr_ref_cuda_tests.cpp +++ b/cpp/tests/mr/mr_ref_cuda_tests.cpp @@ -3,30 +3,23 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "mr_ref_test.hpp" - -#include -#include -#include - -#include +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" +#include "mr_ref_test_mt.hpp" namespace rmm::test { namespace { -// Single-threaded basic tests (5 tests) INSTANTIATE_TEST_SUITE_P(CudaResourceTests, mr_ref_test, ::testing::Values("CUDA"), [](auto const& info) { return info.param; }); -// Single-threaded allocation tests (9 tests) INSTANTIATE_TEST_SUITE_P(CudaResourceAllocationTests, mr_ref_allocation_test, ::testing::Values("CUDA"), [](auto const& info) { return info.param; }); -// Multi-threaded tests (15 tests) INSTANTIATE_TEST_SUITE_P(CudaMultiThreadResourceTests, mr_ref_test_mt, ::testing::Values("CUDA"), diff --git a/cpp/tests/mr/mr_ref_default_tests.cpp b/cpp/tests/mr/mr_ref_default_tests.cpp index d8d9ca815..b9326d631 100644 --- a/cpp/tests/mr/mr_ref_default_tests.cpp +++ b/cpp/tests/mr/mr_ref_default_tests.cpp @@ -17,11 +17,6 @@ namespace rmm::test { namespace { -// Suppress warnings about uninstantiated parameterized tests in this file -GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_test); -GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_allocation_test); -GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_test_mt); - // Helper functions for multi-threaded tests template void spawn_n(std::size_t num_threads, Task task, Arguments&&... args) diff --git a/cpp/tests/mr/mr_ref_fixed_size_tests.cpp b/cpp/tests/mr/mr_ref_fixed_size_tests.cpp index 6d45b117d..1c2e001dc 100644 --- a/cpp/tests/mr/mr_ref_fixed_size_tests.cpp +++ b/cpp/tests/mr/mr_ref_fixed_size_tests.cpp @@ -3,23 +3,11 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "mr_ref_test.hpp" - -#include -#include -#include -#include - -#include +#include "mr_ref_test_basic.hpp" namespace rmm::test { namespace { -// Suppress warnings about uninstantiated tests (Fixed_Size only has basic tests) -GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_allocation_test); -GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_test_mt); - -// Single-threaded basic tests (5 tests) // Note: Fixed_Size MR cannot handle dynamic allocation sizes, so only basic tests are included INSTANTIATE_TEST_SUITE_P(FixedSizeResourceTests, mr_ref_test, diff --git a/cpp/tests/mr/mr_ref_managed_tests.cpp b/cpp/tests/mr/mr_ref_managed_tests.cpp index e4179b079..5a57486e6 100644 --- a/cpp/tests/mr/mr_ref_managed_tests.cpp +++ b/cpp/tests/mr/mr_ref_managed_tests.cpp @@ -3,30 +3,23 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "mr_ref_test.hpp" - -#include -#include -#include - -#include +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" +#include "mr_ref_test_mt.hpp" namespace rmm::test { namespace { -// Single-threaded basic tests (5 tests) INSTANTIATE_TEST_SUITE_P(ManagedResourceTests, mr_ref_test, ::testing::Values("Managed"), [](auto const& info) { return info.param; }); -// Single-threaded allocation tests (9 tests) INSTANTIATE_TEST_SUITE_P(ManagedResourceAllocationTests, mr_ref_allocation_test, ::testing::Values("Managed"), [](auto const& info) { return info.param; }); -// Multi-threaded tests (15 tests) INSTANTIATE_TEST_SUITE_P(ManagedMultiThreadResourceTests, mr_ref_test_mt, ::testing::Values("Managed"), diff --git a/cpp/tests/mr/mr_ref_pinned_pool_tests.cpp b/cpp/tests/mr/mr_ref_pinned_pool_tests.cpp index f2874ab5f..d6e1910a0 100644 --- a/cpp/tests/mr/mr_ref_pinned_pool_tests.cpp +++ b/cpp/tests/mr/mr_ref_pinned_pool_tests.cpp @@ -3,34 +3,21 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "mr_ref_test.hpp" - -#include -#include -#include -#include - -#include +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" namespace rmm::test { namespace { -// Suppress warning about uninstantiated multi-threaded tests (PinnedPool doesn't support MT tests) -GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_test_mt); - -// Single-threaded basic tests (5 tests) INSTANTIATE_TEST_SUITE_P(PinnedPoolResourceTests, mr_ref_test, ::testing::Values("PinnedPool"), [](auto const& info) { return info.param; }); -// Single-threaded allocation tests (9 tests) INSTANTIATE_TEST_SUITE_P(PinnedPoolResourceAllocationTests, mr_ref_allocation_test, ::testing::Values("PinnedPool"), [](auto const& info) { return info.param; }); -// Note: No multi-threaded tests for PinnedPool memory resource - } // namespace } // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_pinned_tests.cpp b/cpp/tests/mr/mr_ref_pinned_tests.cpp index 5aedd3918..f073a5e9b 100644 --- a/cpp/tests/mr/mr_ref_pinned_tests.cpp +++ b/cpp/tests/mr/mr_ref_pinned_tests.cpp @@ -3,34 +3,21 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "mr_ref_test.hpp" - -#include -#include -#include -#include - -#include +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" namespace rmm::test { namespace { -// Suppress warning about uninstantiated multi-threaded tests (Pinned doesn't support MT tests) -GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_test_mt); - -// Single-threaded basic tests (5 tests) INSTANTIATE_TEST_SUITE_P(PinnedResourceTests, mr_ref_test, ::testing::Values("Pinned"), [](auto const& info) { return info.param; }); -// Single-threaded allocation tests (9 tests) INSTANTIATE_TEST_SUITE_P(PinnedResourceAllocationTests, mr_ref_allocation_test, ::testing::Values("Pinned"), [](auto const& info) { return info.param; }); -// Note: No multi-threaded tests for Pinned memory resource - } // namespace } // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_pool_tests.cpp b/cpp/tests/mr/mr_ref_pool_tests.cpp index 9920b4df6..176792869 100644 --- a/cpp/tests/mr/mr_ref_pool_tests.cpp +++ b/cpp/tests/mr/mr_ref_pool_tests.cpp @@ -3,31 +3,23 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "mr_ref_test.hpp" - -#include -#include -#include -#include - -#include +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" +#include "mr_ref_test_mt.hpp" namespace rmm::test { namespace { -// Single-threaded basic tests (5 tests) INSTANTIATE_TEST_SUITE_P(PoolResourceTests, mr_ref_test, ::testing::Values("Pool"), [](auto const& info) { return info.param; }); -// Single-threaded allocation tests (9 tests) INSTANTIATE_TEST_SUITE_P(PoolResourceAllocationTests, mr_ref_allocation_test, ::testing::Values("Pool"), [](auto const& info) { return info.param; }); -// Multi-threaded tests (15 tests) INSTANTIATE_TEST_SUITE_P(PoolMultiThreadResourceTests, mr_ref_test_mt, ::testing::Values("Pool"), diff --git a/cpp/tests/mr/mr_ref_system_tests.cpp b/cpp/tests/mr/mr_ref_system_tests.cpp index ae9dc20ac..ca640eb97 100644 --- a/cpp/tests/mr/mr_ref_system_tests.cpp +++ b/cpp/tests/mr/mr_ref_system_tests.cpp @@ -3,34 +3,21 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "mr_ref_test.hpp" - -#include -#include -#include -#include - -#include +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" namespace rmm::test { namespace { -// Suppress warning about uninstantiated multi-threaded tests (System doesn't support MT tests) -GTEST_ALLOW_UNINSTANTIATED_PARAMETERIZED_TEST(mr_ref_test_mt); - -// Single-threaded basic tests (5 tests) INSTANTIATE_TEST_SUITE_P(SystemResourceTests, mr_ref_test, ::testing::Values("System"), [](auto const& info) { return info.param; }); -// Single-threaded allocation tests (9 tests) INSTANTIATE_TEST_SUITE_P(SystemResourceAllocationTests, mr_ref_allocation_test, ::testing::Values("System"), [](auto const& info) { return info.param; }); -// Note: No multi-threaded tests for System memory resource - } // namespace } // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_test.hpp b/cpp/tests/mr/mr_ref_test.hpp index 702eee194..3d25434e6 100644 --- a/cpp/tests/mr/mr_ref_test.hpp +++ b/cpp/tests/mr/mr_ref_test.hpp @@ -517,339 +517,4 @@ struct mr_ref_test : public ::testing::TestWithParam { struct mr_ref_allocation_test : public mr_ref_test {}; -// Multi-threaded test fixture -struct mr_ref_test_mt : public mr_ref_test {}; - -// Parameterized test definitions for mr_ref_test - -TEST_P(mr_ref_test, SetCurrentDeviceResourceRef) -{ - rmm::mr::cuda_memory_resource cuda_mr{}; - auto cuda_ref = rmm::device_async_resource_ref{cuda_mr}; - - rmm::mr::set_current_device_resource_ref(cuda_ref); - auto old = rmm::mr::set_current_device_resource_ref(this->ref); - - // old mr should equal a cuda mr - EXPECT_EQ(old, cuda_ref); - - // current dev resource should equal this resource - EXPECT_EQ(this->ref, rmm::mr::get_current_device_resource_ref()); - - test_get_current_device_resource_ref(); - - // Resetting should reset to initial cuda resource - rmm::mr::reset_current_device_resource_ref(); - EXPECT_EQ(rmm::device_async_resource_ref{rmm::mr::detail::initial_resource()}, - rmm::mr::get_current_device_resource_ref()); -} - -TEST_P(mr_ref_test, SelfEquality) { EXPECT_TRUE(this->ref == this->ref); } - -// Simple reproducer for https://github.com/rapidsai/rmm/issues/861 -TEST_P(mr_ref_test, AllocationsAreDifferent) { concurrent_allocations_are_different(this->ref); } - -TEST_P(mr_ref_test, AsyncAllocationsAreDifferentDefaultStream) -{ - concurrent_async_allocations_are_different(this->ref, cuda_stream_view{}); -} - -TEST_P(mr_ref_test, AsyncAllocationsAreDifferent) -{ - concurrent_async_allocations_are_different(this->ref, this->stream); -} - -// Parameterized test definitions for mr_ref_allocation_test - -TEST_P(mr_ref_allocation_test, AllocateDefault) { test_various_allocations(this->ref); } - -TEST_P(mr_ref_allocation_test, AllocateDefaultStream) -{ - test_various_async_allocations(this->ref, cuda_stream_view{}); -} - -TEST_P(mr_ref_allocation_test, AllocateOnStream) -{ - test_various_async_allocations(this->ref, this->stream); -} - -TEST_P(mr_ref_allocation_test, RandomAllocations) { test_random_allocations(this->ref); } - -TEST_P(mr_ref_allocation_test, RandomAllocationsDefaultStream) -{ - test_random_async_allocations( - this->ref, default_num_allocations, default_max_size, cuda_stream_view{}); -} - -TEST_P(mr_ref_allocation_test, RandomAllocationsStream) -{ - test_random_async_allocations(this->ref, default_num_allocations, default_max_size, this->stream); -} - -TEST_P(mr_ref_allocation_test, MixedRandomAllocationFree) -{ - test_mixed_random_allocation_free(this->ref, default_max_size); -} - -TEST_P(mr_ref_allocation_test, MixedRandomAllocationFreeDefaultStream) -{ - test_mixed_random_async_allocation_free(this->ref, default_max_size, cuda_stream_view{}); -} - -TEST_P(mr_ref_allocation_test, MixedRandomAllocationFreeStream) -{ - test_mixed_random_async_allocation_free(this->ref, default_max_size, this->stream); -} - -// Helper functions for multi-threaded tests - -template -void spawn_n(std::size_t num_threads, Task task, Arguments&&... args) -{ - std::vector threads; - threads.reserve(num_threads); - for (std::size_t i = 0; i < num_threads; ++i) { - threads.emplace_back(std::thread(task, std::forward(args)...)); - } - - for (auto& thread : threads) { - thread.join(); - } -} - -template -void spawn(Task task, Arguments&&... args) -{ - spawn_n(4, task, std::forward(args)...); -} - -inline void async_allocate_loop(rmm::device_async_resource_ref ref, - std::size_t num_allocations, - std::list& allocations, - std::mutex& mtx, - std::condition_variable& allocations_ready, - cudaEvent_t& event, - rmm::cuda_stream_view stream) -{ - constexpr std::size_t max_size{1_MiB}; - - std::default_random_engine generator; - std::uniform_int_distribution size_distribution(1, max_size); - - for (std::size_t i = 0; i < num_allocations; ++i) { - std::size_t size = size_distribution(generator); - void* ptr = ref.allocate(stream, size); - { - std::lock_guard lock(mtx); - RMM_CUDA_TRY(cudaEventRecord(event, stream.value())); - allocations.emplace_back(ptr, size); - } - allocations_ready.notify_one(); - } - - // Work around for threads going away before cudaEvent has finished async processing - cudaEventSynchronize(event); -} - -inline void async_deallocate_loop(rmm::device_async_resource_ref ref, - std::size_t num_allocations, - std::list& allocations, - std::mutex& mtx, - std::condition_variable& allocations_ready, - cudaEvent_t& event, - rmm::cuda_stream_view stream) -{ - for (std::size_t i = 0; i < num_allocations; i++) { - std::unique_lock lock(mtx); - allocations_ready.wait(lock, [&allocations] { return !allocations.empty(); }); - RMM_CUDA_TRY(cudaStreamWaitEvent(stream.value(), event)); - allocation alloc = allocations.front(); - allocations.pop_front(); - ref.deallocate(stream, alloc.ptr, alloc.size); - } - - // Work around for threads going away before cudaEvent has finished async processing - cudaEventSynchronize(event); -} - -inline void test_async_allocate_free_different_threads(rmm::device_async_resource_ref ref, - rmm::cuda_stream_view streamA, - rmm::cuda_stream_view streamB) -{ - constexpr std::size_t num_allocations{100}; - - std::mutex mtx; - std::condition_variable allocations_ready; - std::list allocations; - cudaEvent_t event{}; - - RMM_CUDA_TRY(cudaEventCreate(&event)); - - std::thread producer(async_allocate_loop, - ref, - num_allocations, - std::ref(allocations), - std::ref(mtx), - std::ref(allocations_ready), - std::ref(event), - streamA); - - std::thread consumer(async_deallocate_loop, - ref, - num_allocations, - std::ref(allocations), - std::ref(mtx), - std::ref(allocations_ready), - std::ref(event), - streamB); - - producer.join(); - consumer.join(); - - RMM_CUDA_TRY(cudaEventDestroy(event)); -} - -// Parameterized test definitions for mr_ref_test_mt - -TEST_P(mr_ref_test_mt, SetCurrentDeviceResourceRef_mt) -{ - // single thread changes default resource, then multiple threads use it - auto old = rmm::mr::set_current_device_resource_ref(this->ref); - test_get_current_device_resource_ref(); - - int device; - RMM_CUDA_TRY(cudaGetDevice(&device)); - - spawn([device, mr = this->ref]() { - RMM_CUDA_TRY(cudaSetDevice(device)); - EXPECT_EQ(mr, rmm::mr::get_current_device_resource_ref()); - test_get_current_device_resource_ref(); // test allocating with the new default resource - }); - - // resetting default resource should reset to initial - rmm::mr::reset_current_device_resource_ref(); - EXPECT_EQ(old, rmm::mr::get_current_device_resource_ref()); -} - -TEST_P(mr_ref_test_mt, SetCurrentDeviceResourceRefPerThread_mt) -{ - int num_devices{}; - RMM_CUDA_TRY(cudaGetDeviceCount(&num_devices)); - - std::vector threads; - threads.reserve(num_devices); - - auto mr = this->ref; - - for (int i = 0; i < num_devices; ++i) { - threads.emplace_back( - [mr](auto dev_id) { - RMM_CUDA_TRY(cudaSetDevice(dev_id)); - auto cuda_ref = rmm::mr::get_current_device_resource_ref(); - auto old = rmm::mr::set_current_device_resource_ref(mr); - - // initial resource for this device should be CUDA mr - EXPECT_EQ(old, cuda_ref); - // get_current_device_resource_ref should equal the resource we - // just set - EXPECT_EQ(mr, rmm::mr::get_current_device_resource_ref()); - // Resetting current dev resource ref should make it - // cuda MR and return the MR we previously set - old = rmm::mr::reset_current_device_resource_ref(); - EXPECT_EQ(old, mr); - EXPECT_EQ(cuda_ref, rmm::mr::get_current_device_resource_ref()); - }, - i); - } - - for (auto& thread : threads) { - thread.join(); - } -} - -TEST_P(mr_ref_test_mt, Allocate) -{ - int device; - RMM_CUDA_TRY(cudaGetDevice(&device)); - - auto mr = this->ref; - spawn([device, mr]() { - RMM_CUDA_TRY(cudaSetDevice(device)); - test_various_allocations(mr); - }); -} - -TEST_P(mr_ref_test_mt, AllocateDefaultStream) -{ - spawn(test_various_async_allocations, this->ref, rmm::cuda_stream_view{}); -} - -TEST_P(mr_ref_test_mt, AllocateOnStream) -{ - spawn(test_various_async_allocations, this->ref, this->stream.view()); -} - -TEST_P(mr_ref_test_mt, RandomAllocations) -{ - spawn(test_random_allocations, this->ref, default_num_allocations, default_max_size); -} - -TEST_P(mr_ref_test_mt, RandomAllocationsDefaultStream) -{ - spawn(test_random_async_allocations, - this->ref, - default_num_allocations, - default_max_size, - rmm::cuda_stream_view{}); -} - -TEST_P(mr_ref_test_mt, RandomAllocationsStream) -{ - spawn(test_random_async_allocations, - this->ref, - default_num_allocations, - default_max_size, - this->stream.view()); -} - -TEST_P(mr_ref_test_mt, MixedRandomAllocationFree) -{ - spawn(test_mixed_random_allocation_free, this->ref, default_max_size); -} - -TEST_P(mr_ref_test_mt, MixedRandomAllocationFreeDefaultStream) -{ - spawn( - test_mixed_random_async_allocation_free, this->ref, default_max_size, rmm::cuda_stream_view{}); -} - -TEST_P(mr_ref_test_mt, MixedRandomAllocationFreeStream) -{ - spawn(test_mixed_random_async_allocation_free, this->ref, default_max_size, this->stream.view()); -} - -TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsDefaultStream) -{ - test_async_allocate_free_different_threads( - this->ref, rmm::cuda_stream_default, rmm::cuda_stream_default); -} - -TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsPerThreadDefaultStream) -{ - test_async_allocate_free_different_threads( - this->ref, rmm::cuda_stream_per_thread, rmm::cuda_stream_per_thread); -} - -TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsSameStream) -{ - test_async_allocate_free_different_threads(this->ref, this->stream, this->stream); -} - -TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsDifferentStream) -{ - rmm::cuda_stream streamB; - test_async_allocate_free_different_threads(this->ref, this->stream, streamB); - streamB.synchronize(); -} - } // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_test_allocation.hpp b/cpp/tests/mr/mr_ref_test_allocation.hpp new file mode 100644 index 000000000..f3f97c6cd --- /dev/null +++ b/cpp/tests/mr/mr_ref_test_allocation.hpp @@ -0,0 +1,54 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "mr_ref_test.hpp" + +namespace rmm::test { + +// Parameterized test definitions for mr_ref_allocation_test + +TEST_P(mr_ref_allocation_test, AllocateDefault) { test_various_allocations(this->ref); } + +TEST_P(mr_ref_allocation_test, AllocateDefaultStream) +{ + test_various_async_allocations(this->ref, cuda_stream_view{}); +} + +TEST_P(mr_ref_allocation_test, AllocateOnStream) +{ + test_various_async_allocations(this->ref, this->stream); +} + +TEST_P(mr_ref_allocation_test, RandomAllocations) { test_random_allocations(this->ref); } + +TEST_P(mr_ref_allocation_test, RandomAllocationsDefaultStream) +{ + test_random_async_allocations( + this->ref, default_num_allocations, default_max_size, cuda_stream_view{}); +} + +TEST_P(mr_ref_allocation_test, RandomAllocationsStream) +{ + test_random_async_allocations(this->ref, default_num_allocations, default_max_size, this->stream); +} + +TEST_P(mr_ref_allocation_test, MixedRandomAllocationFree) +{ + test_mixed_random_allocation_free(this->ref, default_max_size); +} + +TEST_P(mr_ref_allocation_test, MixedRandomAllocationFreeDefaultStream) +{ + test_mixed_random_async_allocation_free(this->ref, default_max_size, cuda_stream_view{}); +} + +TEST_P(mr_ref_allocation_test, MixedRandomAllocationFreeStream) +{ + test_mixed_random_async_allocation_free(this->ref, default_max_size, this->stream); +} + +} // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_test_basic.hpp b/cpp/tests/mr/mr_ref_test_basic.hpp new file mode 100644 index 000000000..9925776fb --- /dev/null +++ b/cpp/tests/mr/mr_ref_test_basic.hpp @@ -0,0 +1,51 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "mr_ref_test.hpp" + +namespace rmm::test { + +// Parameterized test definitions for mr_ref_test (basic tests) + +TEST_P(mr_ref_test, SetCurrentDeviceResourceRef) +{ + rmm::mr::cuda_memory_resource cuda_mr{}; + auto cuda_ref = rmm::device_async_resource_ref{cuda_mr}; + + rmm::mr::set_current_device_resource_ref(cuda_ref); + auto old = rmm::mr::set_current_device_resource_ref(this->ref); + + // old mr should equal a cuda mr + EXPECT_EQ(old, cuda_ref); + + // current dev resource should equal this resource + EXPECT_EQ(this->ref, rmm::mr::get_current_device_resource_ref()); + + test_get_current_device_resource_ref(); + + // Resetting should reset to initial cuda resource + rmm::mr::reset_current_device_resource_ref(); + EXPECT_EQ(rmm::device_async_resource_ref{rmm::mr::detail::initial_resource()}, + rmm::mr::get_current_device_resource_ref()); +} + +TEST_P(mr_ref_test, SelfEquality) { EXPECT_TRUE(this->ref == this->ref); } + +// Simple reproducer for https://github.com/rapidsai/rmm/issues/861 +TEST_P(mr_ref_test, AllocationsAreDifferent) { concurrent_allocations_are_different(this->ref); } + +TEST_P(mr_ref_test, AsyncAllocationsAreDifferentDefaultStream) +{ + concurrent_async_allocations_are_different(this->ref, cuda_stream_view{}); +} + +TEST_P(mr_ref_test, AsyncAllocationsAreDifferent) +{ + concurrent_async_allocations_are_different(this->ref, this->stream); +} + +} // namespace rmm::test diff --git a/cpp/tests/mr/mr_ref_test_mt.hpp b/cpp/tests/mr/mr_ref_test_mt.hpp new file mode 100644 index 000000000..3d08278b9 --- /dev/null +++ b/cpp/tests/mr/mr_ref_test_mt.hpp @@ -0,0 +1,272 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "mr_ref_test.hpp" + +#include +#include +#include +#include +#include + +namespace rmm::test { + +// Multi-threaded test fixture +struct mr_ref_test_mt : public mr_ref_test {}; + +// Helper functions for multi-threaded tests + +template +void spawn_n(std::size_t num_threads, Task task, Arguments&&... args) +{ + std::vector threads; + threads.reserve(num_threads); + for (std::size_t i = 0; i < num_threads; ++i) { + threads.emplace_back(std::thread(task, std::forward(args)...)); + } + + for (auto& thread : threads) { + thread.join(); + } +} + +template +void spawn(Task task, Arguments&&... args) +{ + spawn_n(4, task, std::forward(args)...); +} + +inline void async_allocate_loop(rmm::device_async_resource_ref ref, + std::size_t num_allocations, + std::list& allocations, + std::mutex& mtx, + std::condition_variable& allocations_ready, + cudaEvent_t& event, + rmm::cuda_stream_view stream) +{ + constexpr std::size_t max_size{1_MiB}; + + std::default_random_engine generator; + std::uniform_int_distribution size_distribution(1, max_size); + + for (std::size_t i = 0; i < num_allocations; ++i) { + std::size_t size = size_distribution(generator); + void* ptr = ref.allocate(stream, size); + { + std::lock_guard lock(mtx); + RMM_CUDA_TRY(cudaEventRecord(event, stream.value())); + allocations.emplace_back(ptr, size); + } + allocations_ready.notify_one(); + } + + // Work around for threads going away before cudaEvent has finished async processing + cudaEventSynchronize(event); +} + +inline void async_deallocate_loop(rmm::device_async_resource_ref ref, + std::size_t num_allocations, + std::list& allocations, + std::mutex& mtx, + std::condition_variable& allocations_ready, + cudaEvent_t& event, + rmm::cuda_stream_view stream) +{ + for (std::size_t i = 0; i < num_allocations; i++) { + std::unique_lock lock(mtx); + allocations_ready.wait(lock, [&allocations] { return !allocations.empty(); }); + RMM_CUDA_TRY(cudaStreamWaitEvent(stream.value(), event)); + allocation alloc = allocations.front(); + allocations.pop_front(); + ref.deallocate(stream, alloc.ptr, alloc.size); + } + + // Work around for threads going away before cudaEvent has finished async processing + cudaEventSynchronize(event); +} + +inline void test_async_allocate_free_different_threads(rmm::device_async_resource_ref ref, + rmm::cuda_stream_view streamA, + rmm::cuda_stream_view streamB) +{ + constexpr std::size_t num_allocations{100}; + + std::mutex mtx; + std::condition_variable allocations_ready; + std::list allocations; + cudaEvent_t event{}; + + RMM_CUDA_TRY(cudaEventCreate(&event)); + + std::thread producer(async_allocate_loop, + ref, + num_allocations, + std::ref(allocations), + std::ref(mtx), + std::ref(allocations_ready), + std::ref(event), + streamA); + + std::thread consumer(async_deallocate_loop, + ref, + num_allocations, + std::ref(allocations), + std::ref(mtx), + std::ref(allocations_ready), + std::ref(event), + streamB); + + producer.join(); + consumer.join(); + + RMM_CUDA_TRY(cudaEventDestroy(event)); +} + +// Parameterized test definitions for mr_ref_test_mt + +TEST_P(mr_ref_test_mt, SetCurrentDeviceResourceRef_mt) +{ + // single thread changes default resource, then multiple threads use it + auto old = rmm::mr::set_current_device_resource_ref(this->ref); + test_get_current_device_resource_ref(); + + int device; + RMM_CUDA_TRY(cudaGetDevice(&device)); + + spawn([device, mr = this->ref]() { + RMM_CUDA_TRY(cudaSetDevice(device)); + EXPECT_EQ(mr, rmm::mr::get_current_device_resource_ref()); + test_get_current_device_resource_ref(); // test allocating with the new default resource + }); + + // resetting default resource should reset to initial + rmm::mr::reset_current_device_resource_ref(); + EXPECT_EQ(old, rmm::mr::get_current_device_resource_ref()); +} + +TEST_P(mr_ref_test_mt, SetCurrentDeviceResourceRefPerThread_mt) +{ + int num_devices{}; + RMM_CUDA_TRY(cudaGetDeviceCount(&num_devices)); + + std::vector threads; + threads.reserve(num_devices); + + auto mr = this->ref; + + for (int i = 0; i < num_devices; ++i) { + threads.emplace_back( + [mr](auto dev_id) { + RMM_CUDA_TRY(cudaSetDevice(dev_id)); + auto cuda_ref = rmm::mr::get_current_device_resource_ref(); + auto old = rmm::mr::set_current_device_resource_ref(mr); + + // initial resource for this device should be CUDA mr + EXPECT_EQ(old, cuda_ref); + // get_current_device_resource_ref should equal the resource we + // just set + EXPECT_EQ(mr, rmm::mr::get_current_device_resource_ref()); + // Resetting current dev resource ref should make it + // cuda MR and return the MR we previously set + old = rmm::mr::reset_current_device_resource_ref(); + EXPECT_EQ(old, mr); + EXPECT_EQ(cuda_ref, rmm::mr::get_current_device_resource_ref()); + }, + i); + } + + for (auto& thread : threads) { + thread.join(); + } +} + +TEST_P(mr_ref_test_mt, Allocate) +{ + int device; + RMM_CUDA_TRY(cudaGetDevice(&device)); + + auto mr = this->ref; + spawn([device, mr]() { + RMM_CUDA_TRY(cudaSetDevice(device)); + test_various_allocations(mr); + }); +} + +TEST_P(mr_ref_test_mt, AllocateDefaultStream) +{ + spawn(test_various_async_allocations, this->ref, rmm::cuda_stream_view{}); +} + +TEST_P(mr_ref_test_mt, AllocateOnStream) +{ + spawn(test_various_async_allocations, this->ref, this->stream.view()); +} + +TEST_P(mr_ref_test_mt, RandomAllocations) +{ + spawn(test_random_allocations, this->ref, default_num_allocations, default_max_size); +} + +TEST_P(mr_ref_test_mt, RandomAllocationsDefaultStream) +{ + spawn(test_random_async_allocations, + this->ref, + default_num_allocations, + default_max_size, + rmm::cuda_stream_view{}); +} + +TEST_P(mr_ref_test_mt, RandomAllocationsStream) +{ + spawn(test_random_async_allocations, + this->ref, + default_num_allocations, + default_max_size, + this->stream.view()); +} + +TEST_P(mr_ref_test_mt, MixedRandomAllocationFree) +{ + spawn(test_mixed_random_allocation_free, this->ref, default_max_size); +} + +TEST_P(mr_ref_test_mt, MixedRandomAllocationFreeDefaultStream) +{ + spawn( + test_mixed_random_async_allocation_free, this->ref, default_max_size, rmm::cuda_stream_view{}); +} + +TEST_P(mr_ref_test_mt, MixedRandomAllocationFreeStream) +{ + spawn(test_mixed_random_async_allocation_free, this->ref, default_max_size, this->stream.view()); +} + +TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsDefaultStream) +{ + test_async_allocate_free_different_threads( + this->ref, rmm::cuda_stream_default, rmm::cuda_stream_default); +} + +TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsPerThreadDefaultStream) +{ + test_async_allocate_free_different_threads( + this->ref, rmm::cuda_stream_per_thread, rmm::cuda_stream_per_thread); +} + +TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsSameStream) +{ + test_async_allocate_free_different_threads(this->ref, this->stream, this->stream); +} + +TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsDifferentStream) +{ + rmm::cuda_stream streamB; + test_async_allocate_free_different_threads(this->ref, this->stream, streamB); + streamB.synchronize(); +} + +} // namespace rmm::test