diff --git a/cpp/benchmarks/fixture/nvbench_fixture.hpp b/cpp/benchmarks/fixture/nvbench_fixture.hpp index 4e4eec3547f..ac0cab4071b 100644 --- a/cpp/benchmarks/fixture/nvbench_fixture.hpp +++ b/cpp/benchmarks/fixture/nvbench_fixture.hpp @@ -45,6 +45,8 @@ static std::string cuio_host_mem_param{ * Initializes the default memory resource to use the RMM pool device resource. */ struct nvbench_base_fixture { + using host_pooled_mr_t = rmm::mr::pool_memory_resource; + inline auto make_cuda() { return std::make_shared(); } inline auto make_pool() @@ -90,12 +92,14 @@ struct nvbench_base_fixture { inline rmm::host_async_resource_ref make_cuio_host_pinned_pool() { - using host_pooled_mr = rmm::mr::pool_memory_resource; - static std::shared_ptr mr = std::make_shared( - std::make_shared().get(), - size_t{1} * 1024 * 1024 * 1024); + if (!this->host_pooled_mr) { + // Don't store in static, as the CUDA context may be destroyed before static destruction + this->host_pooled_mr = std::make_shared( + std::make_shared().get(), + size_t{1} * 1024 * 1024 * 1024); + } - return *mr; + return *this->host_pooled_mr; } inline rmm::host_async_resource_ref create_cuio_host_memory_resource(std::string const& mode) @@ -126,9 +130,16 @@ struct nvbench_base_fixture { std::cout << "CUIO host memory resource = " << cuio_host_mode << "\n"; } + ~nvbench_base_fixture() + { + // Ensure the the pool is freed before the CUDA context is destroyed: + cudf::io::set_host_memory_resource(this->make_cuio_host_pinned()); + } + std::shared_ptr mr; std::string rmm_mode{"pool"}; + std::shared_ptr host_pooled_mr; std::string cuio_host_mode{"pinned"}; }; diff --git a/cpp/benchmarks/fixture/nvbench_main.cpp b/cpp/benchmarks/fixture/nvbench_main.cpp index f46cb11a6c3..5dfd67b1c54 100644 --- a/cpp/benchmarks/fixture/nvbench_main.cpp +++ b/cpp/benchmarks/fixture/nvbench_main.cpp @@ -15,29 +15,44 @@ */ #include -#define NVBENCH_ENVIRONMENT cudf::nvbench_base_fixture #include +#include #include +namespace cudf { + // strip off the rmm_mode and cuio_host_mem parameters before passing the // remaining arguments to nvbench::option_parser -#undef NVBENCH_MAIN_PARSE -#define NVBENCH_MAIN_PARSE(argc, argv) \ - nvbench::option_parser parser; \ - std::vector m_args; \ - for (int i = 0; i < argc; ++i) { \ - std::string arg = argv[i]; \ - if (arg == cudf::detail::rmm_mode_param) { \ - i += 2; \ - } else if (arg == cudf::detail::cuio_host_mem_param) { \ - i += 2; \ - } else { \ - m_args.push_back(arg); \ - } \ - } \ - parser.parse(m_args) +void benchmark_arg_handler(std::vector& args) +{ + std::vector _cudf_tmp_args; + + for (std::size_t i = 0; i < args.size(); ++i) { + std::string arg = args[i]; + if (arg == cudf::detail::rmm_mode_param) { + i++; // skip the next argument + } else if (arg == cudf::detail::cuio_host_mem_param) { + i++; // skip the next argument + } else { + _cudf_tmp_args.push_back(arg); + } + } + + args = _cudf_tmp_args; +} + +} // namespace cudf + +// Install arg handler +#undef NVBENCH_MAIN_CUSTOM_ARGS_HANDLER +#define NVBENCH_MAIN_CUSTOM_ARGS_HANDLER(args) cudf::benchmark_arg_handler(args) + +// Global fixture setup: +#undef NVBENCH_MAIN_INITIALIZE_CUSTOM_POST +#define NVBENCH_MAIN_INITIALIZE_CUSTOM_POST(argc, argv) \ + [[maybe_unused]] auto env_state = cudf::nvbench_base_fixture(argc, argv); // this declares/defines the main() function using the definitions above NVBENCH_MAIN diff --git a/cpp/cmake/thirdparty/patches/nvbench_global_setup.diff b/cpp/cmake/thirdparty/patches/nvbench_global_setup.diff deleted file mode 100644 index 04f96f49b48..00000000000 --- a/cpp/cmake/thirdparty/patches/nvbench_global_setup.diff +++ /dev/null @@ -1,29 +0,0 @@ -diff --git a/nvbench/main.cuh b/nvbench/main.cuh -index 0ba82d7..cca5273 100644 ---- a/nvbench/main.cuh -+++ b/nvbench/main.cuh -@@ -54,6 +54,16 @@ - // clang-format on - #endif - -+#ifndef NVBENCH_ENVIRONMENT -+namespace nvbench { -+struct no_environment -+{ -+ no_environment(int, char const *const *) {} -+}; -+} -+#define NVBENCH_ENVIRONMENT nvbench::no_environment -+#endif -+ - #define NVBENCH_MAIN_PARSE(argc, argv) \ - nvbench::option_parser parser; \ - parser.parse(argc, argv) -@@ -77,6 +87,7 @@ - printer.set_total_state_count(total_states); \ - \ - printer.set_completed_state_count(0); \ -+ [[maybe_unused]] auto env_state = NVBENCH_ENVIRONMENT(argc, argv); \ - for (auto &bench_ptr : benchmarks) \ - { \ - bench_ptr->set_printer(printer); \ diff --git a/cpp/cmake/thirdparty/patches/nvbench_override.json b/cpp/cmake/thirdparty/patches/nvbench_override.json index ad9b19c29c1..ef0deb4c1e9 100644 --- a/cpp/cmake/thirdparty/patches/nvbench_override.json +++ b/cpp/cmake/thirdparty/patches/nvbench_override.json @@ -2,13 +2,8 @@ { "packages" : { "nvbench" : { - "patches" : [ - { - "file" : "${current_json_dir}/nvbench_global_setup.diff", - "issue" : "Fix add support for global setup to initialize RMM in nvbench [https://github.com/NVIDIA/nvbench/pull/123]", - "fixed_in" : "" - } - ] + "git_url": "https://github.com/NVIDIA/nvbench.git", + "git_tag": "555d628e9b250868c9da003e4407087ff1982e8e" } } }