Refactor mr_ref_tests to not depend on MR base classes (#1589)

Fixes #1444

Refactors mr_ref_tests.hpp and the .cpp files that include it to use a hierarchy of factory classes in which the base class contains a `resource_ref` and the derived classes are templated on the actual resource type and own the MR under test to maintain its lifetime. This way the tests can still be value parameterized (by name) and the test machinery can be type erased. In the future, if CCCL adds an `any_resource` type-erased owning resource container, we can use that to simplify this machinery.

Note that until #1598 merges, this PR necessarily disables (comments out) tests of `set_current_device_resource()`. This is because the test only has a resource_ref, so it can't `set_current_device_resource(ref).` #1598 adds `set_current_device_resource_ref()` which reenables this testing.

Authors:
  - Mark Harris (https://github.com/harrism)

Approvers:
  - Vyas Ramasubramani (https://github.com/vyasr)
  - Bradley Dice (https://github.com/bdice)

URL: #1589

tests/mr/device/mr_ref_multithreaded_tests.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -37,15 +37,15 @@ struct mr_ref_test_mt : public mr_ref_test {};
 
 INSTANTIATE_TEST_CASE_P(MultiThreadResourceTests,
  mr_ref_test_mt,
- ::testing::Values(mr_factory{"CUDA", &make_cuda},
+ ::testing::Values("CUDA",
 #ifdef RMM_CUDA_MALLOC_ASYNC_SUPPORT
- mr_factory{"CUDA_Async", &make_cuda_async},
+ "CUDA_Async",
 #endif
- mr_factory{"Managed", &make_managed},
- mr_factory{"Pool", &make_pool},
- mr_factory{"Arena", &make_arena},
- mr_factory{"Binning", &make_binning}),
- [](auto const& info) { return info.param.name; });
+ "Managed",
+ "Pool",
+ "Arena",
+ "Binning"),
+ [](auto const& info) { return info.param; });
 
 template <typename Task, typename... Arguments>
 void spawn_n(std::size_t num_threads, Task task, Arguments&&... args)
@@ -86,7 +86,8 @@ TEST(DefaultTest, GetCurrentDeviceResource_mt)
  });
 }
 
-TEST_P(mr_ref_test_mt, SetCurrentDeviceResource_mt)
+// Disable until we support resource_ref with set_current_device_resource
+/*TEST_P(mr_ref_test_mt, SetCurrentDeviceResource_mt)
 {
  // single thread changes default resource, then multiple threads use it
 
@@ -101,9 +102,9 @@ TEST_P(mr_ref_test_mt, SetCurrentDeviceResource_mt)
  // setting default resource w/ nullptr should reset to initial
  rmm::mr::set_current_device_resource(nullptr);
  EXPECT_TRUE(old->is_equal(*rmm::mr::get_current_device_resource()));
-}
+}*/
 
-TEST_P(mr_ref_test_mt, SetCurrentDeviceResourcePerThread_mt)
+/*TEST_P(mr_ref_test_mt, SetCurrentDeviceResourcePerThread_mt)
 {
  int num_devices{};
  RMM_CUDA_TRY(cudaGetDeviceCount(&num_devices));
@@ -135,7 +136,7 @@ TEST_P(mr_ref_test_mt, SetCurrentDeviceResourcePerThread_mt)
  for (auto& thread : threads) {
  thread.join();
  }
-}
+}*/
 
 TEST_P(mr_ref_test_mt, Allocate) { spawn(test_various_allocations, this->ref); }
 

tests/mr/device/mr_ref_test.hpp

@@ -27,7 +27,6 @@
 #include <rmm/mr/device/binning_memory_resource.hpp>
 #include <rmm/mr/device/cuda_async_memory_resource.hpp>
 #include <rmm/mr/device/cuda_memory_resource.hpp>
-#include <rmm/mr/device/device_memory_resource.hpp>
 #include <rmm/mr/device/fixed_size_memory_resource.hpp>
 #include <rmm/mr/device/managed_memory_resource.hpp>
 #include <rmm/mr/device/owning_wrapper.hpp>
@@ -42,12 +41,11 @@
 #include <gtest/gtest.h>
 
 #include <cstddef>
-#include <cstdint>
-#include <functional>
 #include <random>
+#include <string>
 #include <utility>
 
-using resource_ref = cuda::mr::resource_ref<cuda::mr::device_accessible>;
+using resource_ref = rmm::device_async_resource_ref;
 
 namespace rmm::test {
 
@@ -333,40 +331,6 @@ inline void test_mixed_random_async_allocation_free(rmm::device_async_resource_r
  EXPECT_EQ(allocations.size(), active_allocations);
 }
 
-using MRFactoryFunc = std::function<std::shared_ptr<rmm::mr::device_memory_resource>()>;
-
-/// Encapsulates a `device_memory_resource` factory function and associated name
-struct mr_factory {
- mr_factory(std::string name, MRFactoryFunc factory)
- : name{std::move(name)}, factory{std::move(factory)}
- {
- }
-
- std::string name; ///< Name to associate with tests that use this factory
- MRFactoryFunc factory; ///< Factory function that returns shared_ptr to `device_memory_resource`
- ///< instance to use in test
-};
-
-/// Test fixture class value-parameterized on different `mr_factory`s
-struct mr_ref_test : public ::testing::TestWithParam<mr_factory> {
- void SetUp() override
- {
- auto factory = GetParam().factory;
- mr = factory();
- if (mr == nullptr) {
- GTEST_SKIP() << "Skipping tests since the memory resource is not supported with this CUDA "
- << "driver/runtime version";
- }
- ref = rmm::device_async_resource_ref{*mr};
- }
-
- std::shared_ptr<rmm::mr::device_memory_resource> mr; ///< Pointer to resource to use in tests
- rmm::device_async_resource_ref ref{*mr};
- rmm::cuda_stream stream{};
-};
-
-struct mr_ref_allocation_test : public mr_ref_test {};
-
 /// MR factory functions
 inline auto make_cuda() { return std::make_shared<rmm::mr::cuda_memory_resource>(); }
 
@@ -426,4 +390,88 @@ inline auto make_binning()
  return mr;
 }
 
+struct mr_factory_base {
+ std::string name{}; ///< Name to associate with tests that use this factory
+ resource_ref mr{rmm::mr::get_current_device_resource()};
+ bool skip_test{false};
+};
+
+/// Encapsulates a memory resource factory function and associated name
+template <class Resource, typename MRFactoryFunc>
+struct mr_factory : mr_factory_base {
+ mr_factory(std::string_view name, MRFactoryFunc factory)
+ : mr_factory_base{std::string{name}}, owned_mr{std::move(factory())}
+ {
+ if (owned_mr == nullptr) { skip_test = true; }
+
+ mr = *owned_mr;
+ }
+
+ // Owned resource to use in tests, type determined by the type of factory function
+ std::invoke_result_t<MRFactoryFunc> owned_mr;
+};
+
+using cuda_mr = rmm::mr::cuda_memory_resource;
+using pinned_mr = rmm::mr::pinned_host_memory_resource;
+using cuda_async_mr = rmm::mr::cuda_async_memory_resource;
+using managed_mr = rmm::mr::managed_memory_resource;
+using system_mr = rmm::mr::system_memory_resource;
+using pool_mr = rmm::mr::pool_memory_resource<cuda_mr>;
+using pinned_pool_mr = rmm::mr::pool_memory_resource<pinned_mr>;
+using arena_mr = rmm::mr::arena_memory_resource<cuda_mr>;
+using fixed_mr = rmm::mr::fixed_size_memory_resource<cuda_mr>;
+using binning_mr = rmm::mr::binning_memory_resource<pool_mr>;
+
+inline std::shared_ptr<mr_factory_base> mr_factory_dispatch(std::string name)
+{
+ if (name == "CUDA") {
+ return std::make_shared<mr_factory<cuda_mr, decltype(make_cuda)>>("CUDA", make_cuda);
+ } else if (name == "Host_Pinned") {
+ return std::make_shared<mr_factory<pinned_mr, decltype(make_host_pinned)>>("Host_Pinned",
+ make_host_pinned);
+ } else if (name == "CUDA_Async") {
+ return std::make_shared<mr_factory<cuda_async_mr, decltype(make_cuda_async)>>("CUDA_Async",
+ make_cuda_async);
+ } else if (name == "Managed") {
+ return std::make_shared<mr_factory<managed_mr, decltype(make_managed)>>("Managed",
+ make_managed);
+ } else if (name == "System") {
+ return std::make_shared<mr_factory<system_mr, decltype(make_system)>>("System", make_system);
+ } else if (name == "Pool") {
+ return std::make_shared<mr_factory<pool_mr, decltype(make_pool)>>("Pool", make_pool);
+ } else if (name == "Host_Pinned_Pool") {
+ return std::make_shared<mr_factory<pinned_pool_mr, decltype(make_host_pinned_pool)>>(
+ "Host_Pinned_Pool", make_host_pinned_pool);
+ } else if (name == "Arena") {
+ return std::make_shared<mr_factory<arena_mr, decltype(make_arena)>>("Arena", make_arena);
+ } else if (name == "Binning") {
+ return std::make_shared<mr_factory<binning_mr, decltype(make_binning)>>("Binning",
+ make_binning);
+ } else if (name == "Fixed_Size") {
+ return std::make_shared<mr_factory<fixed_mr, decltype(make_fixed_size)>>("Fixed_Size",
+ make_fixed_size);
+ } else {
+ return std::make_shared<mr_factory<cuda_mr, decltype(make_cuda)>>("Error", make_cuda);
+ }
+}
+
+/// Test fixture class value-parameterized on different `mr_factory`s
+struct mr_ref_test : public ::testing::TestWithParam<std::string> {
+ void SetUp() override
+ {
+ factory_obj = mr_factory_dispatch(GetParam());
+ if (factory_obj->skip_test) {
+ GTEST_SKIP() << "Skipping tests since the memory resource is not supported with this CUDA "
+ << "driver/runtime version";
+ }
+ ref = factory_obj->mr;
+ }
+
+ std::shared_ptr<mr_factory_base> factory_obj{};
+ resource_ref ref{rmm::mr::get_current_device_resource()};
+ rmm::cuda_stream stream{};
+};
+
+struct mr_ref_allocation_test : public mr_ref_test {};
+
 } // namespace rmm::test

tests/mr/device/mr_ref_tests.cpp

@@ -27,33 +27,33 @@ namespace {
 
 INSTANTIATE_TEST_SUITE_P(ResourceTests,
  mr_ref_test,
- ::testing::Values(mr_factory{"CUDA", &make_cuda},
+ ::testing::Values("CUDA",
 #ifdef RMM_CUDA_MALLOC_ASYNC_SUPPORT
- mr_factory{"CUDA_Async", &make_cuda_async},
+ "CUDA_Async",
 #endif
- mr_factory{"Managed", &make_managed},
- mr_factory{"System", &make_system},
- mr_factory{"Pool", &make_pool},
- mr_factory{"HostPinnedPool", &make_host_pinned_pool},
- mr_factory{"Arena", &make_arena},
- mr_factory{"Binning", &make_binning},
- mr_factory{"Fixed_Size", &make_fixed_size}),
- [](auto const& info) { return info.param.name; });
+ "Managed",
+ "System",
+ "Pool",
+ "HostPinnedPool",
+ "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(mr_factory{"CUDA", &make_cuda},
+ ::testing::Values("CUDA",
 #ifdef RMM_CUDA_MALLOC_ASYNC_SUPPORT
- mr_factory{"CUDA_Async", &make_cuda_async},
+ "CUDA_Async",
 #endif
- mr_factory{"Managed", &make_managed},
- mr_factory{"System", &make_system},
- mr_factory{"Pool", &make_pool},
- mr_factory{"HostPinnedPool", &make_host_pinned_pool},
- mr_factory{"Arena", &make_arena},
- mr_factory{"Binning", &make_binning}),
- [](auto const& info) { return info.param.name; });
+ "Managed",
+ "System"
+ "Pool",
+ "HostPinnedPool",
+ "Arena",
+ "Binning"),
+ [](auto const& info) { return info.param; });
 
 TEST(DefaultTest, CurrentDeviceResourceIsCUDA)
 {
@@ -70,7 +70,8 @@ TEST(DefaultTest, GetCurrentDeviceResource)
  EXPECT_TRUE(mr->is_equal(rmm::mr::cuda_memory_resource{}));
 }
 
-TEST_P(mr_ref_test, SetCurrentDeviceResource)
+// Disable until we support resource_ref with set_current_device_resource
+/*TEST_P(mr_ref_test, SetCurrentDeviceResource)
 {
  rmm::mr::device_memory_resource* old{};
  old = rmm::mr::set_current_device_resource(this->mr.get());
@@ -87,7 +88,7 @@ TEST_P(mr_ref_test, SetCurrentDeviceResource)
  // setting to `nullptr` should reset to initial cuda resource
  rmm::mr::set_current_device_resource(nullptr);
  EXPECT_TRUE(rmm::mr::get_current_device_resource()->is_equal(rmm::mr::cuda_memory_resource{}));
-}
+}*/
 
 TEST_P(mr_ref_test, SelfEquality) { EXPECT_TRUE(this->ref == this->ref); }
 

tests/mr/device/thrust_allocator_tests.cu

@@ -36,7 +36,8 @@ namespace {
 
 struct allocator_test : public mr_ref_test {};
 
-TEST_P(allocator_test, first)
+// Disable until we support resource_ref with set_current_device_resource
+/*TEST_P(allocator_test, first)
 {
  rmm::mr::set_current_device_resource(this->mr.get());
  auto const num_ints{100};
@@ -51,7 +52,7 @@ TEST_P(allocator_test, defaults)
  EXPECT_EQ(allocator.stream(), rmm::cuda_stream_default);
  EXPECT_EQ(allocator.get_upstream_resource(),
  rmm::device_async_resource_ref{rmm::mr::get_current_device_resource()});
-}
+}*/
 
 TEST_P(allocator_test, multi_device)
 {
@@ -70,15 +71,15 @@ TEST_P(allocator_test, multi_device)
 
 INSTANTIATE_TEST_CASE_P(ThrustAllocatorTests,
  allocator_test,
- ::testing::Values(mr_factory{"CUDA", &make_cuda},
+ ::testing::Values("CUDA",
 #ifdef RMM_CUDA_MALLOC_ASYNC_SUPPORT
- mr_factory{"CUDA_Async", &make_cuda_async},
+ "CUDA_Async",
 #endif
- mr_factory{"Managed", &make_managed},
- mr_factory{"Pool", &make_pool},
- mr_factory{"Arena", &make_arena},
- mr_factory{"Binning", &make_binning}),
- [](auto const& info) { return info.param.name; });
+ "Managed",
+ "Pool",
+ "Arena",
+ "Binning"),
+ [](auto const& info) { return info.param; });
 
 } // namespace
 } // namespace rmm::test

tests/mr/host/mr_ref_tests.cpp

@@ -17,7 +17,6 @@
 #include "../../byte_literals.hpp"
 
 #include <rmm/aligned.hpp>
-#include <rmm/mr/host/host_memory_resource.hpp>
 #include <rmm/mr/host/new_delete_resource.hpp>
 #include <rmm/mr/host/pinned_memory_resource.hpp>
 #include <rmm/resource_ref.hpp>