Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Refactor stream mode setup for gtests #15337

Merged
merged 4 commits into from
Apr 1, 2024
Merged

Refactor stream mode setup for gtests #15337

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8e21c58
Refactor stream mode setup for gtests
davidwendt Mar 19, 2024
87ff441
fix merge conflict
davidwendt Mar 19, 2024
8ab99b2
remove unneeded blank lines
davidwendt Mar 19, 2024
5d44699
Merge branch 'branch-24.06' into ref-stream-mode-setup
davidwendt Mar 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 36 additions & 21 deletions cpp/include/cudf_test/testing_main.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,32 @@ inline auto parse_cudf_test_opts(int argc, char** argv)
}
}

/**
* @brief Sets up stream mode memory resource adaptor
*
* The resource adaptor is only set as the current device resource if the
* stream mode is enabled.
*
* The caller must keep the return object alive for the life of the test runs.
*
* @param cmd_opts Command line options returned by parse_cudf_test_opts
* @return Memory resource adaptor
*/
inline auto make_stream_mode_adaptor(cxxopts::ParseResult const& cmd_opts)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I have a mild preference for actually declaring the return type here instead of using auto. Gives the caller a little more context with what's going on.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Normally I prefer the return types as well. The type here is something created through templated RMM classes only to be passed back to an RMM function. This means RMM could change the type/name on us and would then require this code to be updated even though we don't actually rely on the type for anything--we don't make calls to it, we just pass it to set_current_device_resource. The returned object then just needs to stay alive for the life of the tests. Leaving it as auto here helps future proofs this code a bit.

{
auto resource = rmm::mr::get_current_device_resource();
auto const stream_mode = cmd_opts["stream_mode"].as<std::string>();
auto const stream_error_mode = cmd_opts["stream_error_mode"].as<std::string>();
auto const error_on_invalid_stream = (stream_error_mode == "error");
auto const check_default_stream = (stream_mode == "new_cudf_default");
auto adaptor =
make_stream_checking_resource_adaptor(resource, error_on_invalid_stream, check_default_stream);
if ((stream_mode == "new_cudf_default") || (stream_mode == "new_testing_default")) {
rmm::mr::set_current_device_resource(&adaptor);
}
return adaptor;
}

/**
* @brief Macro that defines main function for gtest programs that use rmm
*
Expand All @@ -155,25 +181,14 @@ inline auto parse_cudf_test_opts(int argc, char** argv)
* function parses the command line to customize test behavior, like the
* allocation mode used for creating the default memory resource.
*/
#define CUDF_TEST_PROGRAM_MAIN() \
int main(int argc, char** argv) \
{ \
::testing::InitGoogleTest(&argc, argv); \
auto const cmd_opts = parse_cudf_test_opts(argc, argv); \
auto const rmm_mode = cmd_opts["rmm_mode"].as<std::string>(); \
auto resource = cudf::test::create_memory_resource(rmm_mode); \
rmm::mr::set_current_device_resource(resource.get()); \
\
auto const stream_mode = cmd_opts["stream_mode"].as<std::string>(); \
if ((stream_mode == "new_cudf_default") || (stream_mode == "new_testing_default")) { \
auto const stream_error_mode = cmd_opts["stream_error_mode"].as<std::string>(); \
auto const error_on_invalid_stream = (stream_error_mode == "error"); \
auto const check_default_stream = (stream_mode == "new_cudf_default"); \
auto adaptor = make_stream_checking_resource_adaptor( \
resource.get(), error_on_invalid_stream, check_default_stream); \
rmm::mr::set_current_device_resource(&adaptor); \
return RUN_ALL_TESTS(); \
} \
\
return RUN_ALL_TESTS(); \
#define CUDF_TEST_PROGRAM_MAIN() \
int main(int argc, char** argv) \
{ \
::testing::InitGoogleTest(&argc, argv); \
auto const cmd_opts = parse_cudf_test_opts(argc, argv); \
auto const rmm_mode = cmd_opts["rmm_mode"].as<std::string>(); \
auto resource = cudf::test::create_memory_resource(rmm_mode); \
rmm::mr::set_current_device_resource(resource.get()); \
auto adaptor = make_stream_mode_adaptor(cmd_opts); \
return RUN_ALL_TESTS(); \
}
14 changes: 2 additions & 12 deletions cpp/tests/error/error_handling_test.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -128,17 +128,7 @@ TEST(DebugAssert, cudf_assert_true)
int main(int argc, char** argv)
{
::testing::InitGoogleTest(&argc, argv);
auto const cmd_opts = parse_cudf_test_opts(argc, argv);
auto const stream_mode = cmd_opts["stream_mode"].as<std::string>();
if ((stream_mode == "new_cudf_default") || (stream_mode == "new_testing_default")) {
auto resource = rmm::mr::get_current_device_resource();
auto const stream_error_mode = cmd_opts["stream_error_mode"].as<std::string>();
auto const error_on_invalid_stream = (stream_error_mode == "error");
auto const check_default_stream = (stream_mode == "new_cudf_default");
auto adaptor = make_stream_checking_resource_adaptor(
resource, error_on_invalid_stream, check_default_stream);
rmm::mr::set_current_device_resource(&adaptor);
return RUN_ALL_TESTS();
}
auto const cmd_opts = parse_cudf_test_opts(argc, argv);
auto adaptor = make_stream_mode_adaptor(cmd_opts);
return RUN_ALL_TESTS();
}
Loading