Skip to content

Commit

Permalink
Refactor mr_ref_tests to not depend on MR base classes (#1589)
Browse files Browse the repository at this point in the history 
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
  • Loading branch information
@harrism
harrism authored Jul 2, 2024
1 parent a727a03 commit d71f9e1
Show file tree
Hide file tree
Showing 5 changed files with 131 additions and 81 deletions.
25 changes: 13 additions & 12 deletions tests/mr/device/mr_ref_multithreaded_tests.cpp
View file
Original file line number Diff line number Diff line change
@@ -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.
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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));
Expand Down Expand Up @@ -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); }

Expand Down
124 changes: 86 additions & 38 deletions tests/mr/device/mr_ref_test.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand All @@ -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 {

Expand Down Expand Up @@ -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>(); }

Expand Down Expand Up @@ -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
43 changes: 22 additions & 21 deletions tests/mr/device/mr_ref_tests.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
{
Expand All @@ -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());
Expand All @@ -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); }

Expand Down
19 changes: 10 additions & 9 deletions tests/mr/device/thrust_allocator_tests.cu
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand All @@ -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)
{
Expand All @@ -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
1 change: 0 additions & 1 deletion tests/mr/host/mr_ref_tests.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down

0 comments on commit d71f9e1

Please sign in to comment.