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..3a73c78cd --- /dev/null +++ b/cpp/tests/mr/mr_ref_arena_tests.cpp @@ -0,0 +1,29 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" +#include "mr_ref_test_mt.hpp" + +namespace rmm::test { +namespace { + +INSTANTIATE_TEST_SUITE_P(ArenaResourceTests, + mr_ref_test, + ::testing::Values("Arena"), + [](auto const& info) { return info.param; }); + +INSTANTIATE_TEST_SUITE_P(ArenaResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("Arena"), + [](auto const& info) { return info.param; }); + +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..0c6308e05 --- /dev/null +++ b/cpp/tests/mr/mr_ref_binning_tests.cpp @@ -0,0 +1,29 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" +#include "mr_ref_test_mt.hpp" + +namespace rmm::test { +namespace { + +INSTANTIATE_TEST_SUITE_P(BinningResourceTests, + mr_ref_test, + ::testing::Values("Binning"), + [](auto const& info) { return info.param; }); + +INSTANTIATE_TEST_SUITE_P(BinningResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("Binning"), + [](auto const& info) { return info.param; }); + +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..ce08183ba --- /dev/null +++ b/cpp/tests/mr/mr_ref_cuda_async_tests.cpp @@ -0,0 +1,29 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" +#include "mr_ref_test_mt.hpp" + +namespace rmm::test { +namespace { + +INSTANTIATE_TEST_SUITE_P(CudaAsyncResourceTests, + mr_ref_test, + ::testing::Values("CUDA_Async"), + [](auto const& info) { return info.param; }); + +INSTANTIATE_TEST_SUITE_P(CudaAsyncResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("CUDA_Async"), + [](auto const& info) { return info.param; }); + +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..5c9616998 --- /dev/null +++ b/cpp/tests/mr/mr_ref_cuda_tests.cpp @@ -0,0 +1,29 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" +#include "mr_ref_test_mt.hpp" + +namespace rmm::test { +namespace { + +INSTANTIATE_TEST_SUITE_P(CudaResourceTests, + mr_ref_test, + ::testing::Values("CUDA"), + [](auto const& info) { return info.param; }); + +INSTANTIATE_TEST_SUITE_P(CudaResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("CUDA"), + [](auto const& info) { return info.param; }); + +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..b9326d631 --- /dev/null +++ b/cpp/tests/mr/mr_ref_default_tests.cpp @@ -0,0 +1,106 @@ +/* + * 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 { + +// 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..1c2e001dc --- /dev/null +++ b/cpp/tests/mr/mr_ref_fixed_size_tests.cpp @@ -0,0 +1,18 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test_basic.hpp" + +namespace rmm::test { +namespace { + +// 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..5a57486e6 --- /dev/null +++ b/cpp/tests/mr/mr_ref_managed_tests.cpp @@ -0,0 +1,29 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" +#include "mr_ref_test_mt.hpp" + +namespace rmm::test { +namespace { + +INSTANTIATE_TEST_SUITE_P(ManagedResourceTests, + mr_ref_test, + ::testing::Values("Managed"), + [](auto const& info) { return info.param; }); + +INSTANTIATE_TEST_SUITE_P(ManagedResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("Managed"), + [](auto const& info) { return info.param; }); + +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_pinned_pool_tests.cpp b/cpp/tests/mr/mr_ref_pinned_pool_tests.cpp new file mode 100644 index 000000000..d6e1910a0 --- /dev/null +++ b/cpp/tests/mr/mr_ref_pinned_pool_tests.cpp @@ -0,0 +1,23 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" + +namespace rmm::test { +namespace { + +INSTANTIATE_TEST_SUITE_P(PinnedPoolResourceTests, + mr_ref_test, + ::testing::Values("PinnedPool"), + [](auto const& info) { return info.param; }); + +INSTANTIATE_TEST_SUITE_P(PinnedPoolResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("PinnedPool"), + [](auto const& info) { return info.param; }); + +} // 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..f073a5e9b --- /dev/null +++ b/cpp/tests/mr/mr_ref_pinned_tests.cpp @@ -0,0 +1,23 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" + +namespace rmm::test { +namespace { + +INSTANTIATE_TEST_SUITE_P(PinnedResourceTests, + mr_ref_test, + ::testing::Values("Pinned"), + [](auto const& info) { return info.param; }); + +INSTANTIATE_TEST_SUITE_P(PinnedResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("Pinned"), + [](auto const& info) { return info.param; }); + +} // 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..176792869 --- /dev/null +++ b/cpp/tests/mr/mr_ref_pool_tests.cpp @@ -0,0 +1,29 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" +#include "mr_ref_test_mt.hpp" + +namespace rmm::test { +namespace { + +INSTANTIATE_TEST_SUITE_P(PoolResourceTests, + mr_ref_test, + ::testing::Values("Pool"), + [](auto const& info) { return info.param; }); + +INSTANTIATE_TEST_SUITE_P(PoolResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("Pool"), + [](auto const& info) { return info.param; }); + +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..ca640eb97 --- /dev/null +++ b/cpp/tests/mr/mr_ref_system_tests.cpp @@ -0,0 +1,23 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "mr_ref_test_allocation.hpp" +#include "mr_ref_test_basic.hpp" + +namespace rmm::test { +namespace { + +INSTANTIATE_TEST_SUITE_P(SystemResourceTests, + mr_ref_test, + ::testing::Values("System"), + [](auto const& info) { return info.param; }); + +INSTANTIATE_TEST_SUITE_P(SystemResourceAllocationTests, + mr_ref_allocation_test, + ::testing::Values("System"), + [](auto const& info) { return info.param; }); + +} // namespace +} // 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_multithreaded_tests.cpp b/cpp/tests/mr/mr_ref_test_mt.hpp similarity index 72% rename from cpp/tests/mr/mr_ref_multithreaded_tests.cpp rename to cpp/tests/mr/mr_ref_test_mt.hpp index 6f097ba6c..3d08278b9 100644 --- a/cpp/tests/mr/mr_ref_multithreaded_tests.cpp +++ b/cpp/tests/mr/mr_ref_test_mt.hpp @@ -3,31 +3,22 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include "mr_ref_test.hpp" - -#include -#include -#include -#include -#include -#include -#include +#pragma once -#include +#include "mr_ref_test.hpp" +#include +#include +#include #include #include namespace rmm::test { -namespace { +// Multi-threaded test fixture 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; }); +// Helper functions for multi-threaded tests template void spawn_n(std::size_t num_threads, Task task, Arguments&&... args) @@ -49,43 +40,94 @@ void spawn(Task task, Arguments&&... args) spawn_n(4, task, std::forward(args)...); } -TEST(DefaultTest, UseCurrentDeviceResource_mt) { spawn(test_get_current_device_resource); } +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}; -TEST(DefaultTest, UseCurrentDeviceResourceRef_mt) { spawn(test_get_current_device_resource_ref); } + std::default_random_engine generator; + std::uniform_int_distribution size_distribution(1, max_size); -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{})); - }); -} + 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(); + } -TEST(DefaultTest, CurrentDeviceResourceRefIsCUDA_mt) -{ - spawn([]() { - EXPECT_EQ(rmm::mr::get_current_device_resource_ref(), - rmm::device_async_resource_ref{rmm::mr::detail::initial_resource()}); - }); + // Work around for threads going away before cudaEvent has finished async processing + cudaEventSynchronize(event); } -TEST(DefaultTest, GetCurrentDeviceResource_mt) +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) { - 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{})); - }); + 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); } -TEST(DefaultTest, GetCurrentDeviceResourceRef_mt) +inline void test_async_allocate_free_different_threads(rmm::device_async_resource_ref ref, + rmm::cuda_stream_view streamA, + rmm::cuda_stream_view streamB) { - spawn([]() { - auto mr = rmm::mr::get_current_device_resource_ref(); - EXPECT_EQ(mr, rmm::device_async_resource_ref{rmm::mr::detail::initial_resource()}); - }); + 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 @@ -203,92 +245,6 @@ 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( @@ -313,5 +269,4 @@ TEST_P(mr_ref_test_mt, AllocFreeDifferentThreadsDifferentStream) streamB.synchronize(); } -} // namespace } // 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