From 4613437c63658e7121443c2cb6e199946243b79c Mon Sep 17 00:00:00 2001 From: Allison Piper Date: Sat, 6 Apr 2024 15:35:08 +0000 Subject: [PATCH] Refactor main implementation to improve reusability and customization. Move the implementation of `main` out of macros and into separate functions. This allows for easier reuse and customization of the macros. Existing macro usage should still work as expected, and new customization points will simplify common tasks like argument parsing going forward. --- nvbench/main.cuh | 245 ++++++++++++++++++++++++++++++++++++++--------- 1 file changed, 199 insertions(+), 46 deletions(-) diff --git a/nvbench/main.cuh b/nvbench/main.cuh index bcdced5..21ad7eb 100644 --- a/nvbench/main.cuh +++ b/nvbench/main.cuh @@ -28,14 +28,132 @@ #include #include +// Advanced users can rebuild NVBench's `main` function using the macros in this file, or replace +// them with customized implementations. + +// Customization point, called before NVBench initialization. +#ifndef NVBENCH_MAIN_INITIALIZE_CUSTOM_PRE +#define NVBENCH_MAIN_INITIALIZE_CUSTOM_PRE(argc, argv) []() {}() +#endif + +// Customization point, called after NVBench initialization. +#ifndef NVBENCH_MAIN_INITIALIZE_CUSTOM_POST +#define NVBENCH_MAIN_INITIALIZE_CUSTOM_POST(argc, argv) []() {}() +#endif + +// Customization point, called before NVBench parsing. Update argc/argv if needed. +#ifndef NVBENCH_MAIN_CUSTOM_ARGC_ARGV_HANDLER +#define NVBENCH_MAIN_CUSTOM_ARGC_ARGV_HANDLER(argc, argv) []() {}() +#endif + +// Customization point, called before NVBench parsing. Update args if needed. +#ifndef NVBENCH_MAIN_CUSTOM_ARGS_HANDLER +#define NVBENCH_MAIN_CUSTOM_ARGS_HANDLER(args) []() {}() +#endif + +// Customization point, called before NVBench parsing. +#ifndef NVBENCH_MAIN_PARSE_CUSTOM_PRE +#define NVBENCH_MAIN_PARSE_CUSTOM_PRE(parser) []() {}() +#endif + +// Customization point, called after NVBench parsing. +#ifndef NVBENCH_MAIN_PARSE_CUSTOM_POST +#define NVBENCH_MAIN_PARSE_CUSTOM_POST(parser) []() {}() +#endif + +// Customization point, called after NVBench parser is created. +#ifndef NVBENCH_MAIN_PARSE_CUSTOM_POST +#define NVBENCH_MAIN_PARSE_CUSTOM_POST(parser) []() {}() +#endif + +// Customization point, called before NVBench finalization. +#ifndef NVBENCH_MAIN_FINALIZE_CUSTOM_PRE +#define NVBENCH_MAIN_FINALIZE_CUSTOM_PRE() []() {}() +#endif + +// Customization point, called after NVBench finalization. +#ifndef NVBENCH_MAIN_FINALIZE_CUSTOM_POST +#define NVBENCH_MAIN_FINALIZE_CUSTOM_POST() []() {}() +#endif + +// Customization point, use to catch addition exceptions. +#ifndef NVBENCH_MAIN_CATCH_EXCEPTIONS_CUSTOM +#define NVBENCH_MAIN_CATCH_EXCEPTIONS_CUSTOM +#endif + +/************************************ Default implementation **************************************/ + +#ifndef NVBENCH_MAIN #define NVBENCH_MAIN \ - int main(int argc, char const *const *argv) \ + int main(int argc, char **argv) \ try \ { \ NVBENCH_MAIN_BODY(argc, argv); \ - NVBENCH_CUDA_CALL(cudaDeviceReset()); \ return 0; \ } \ + NVBENCH_MAIN_CATCH_EXCEPTIONS_CUSTOM \ + NVBENCH_MAIN_CATCH_EXCEPTIONS +#endif + +#ifndef NVBENCH_MAIN_BODY +#define NVBENCH_MAIN_BODY(argc, argv) \ + NVBENCH_MAIN_INITIALIZE(argc, argv); \ + { \ + NVBENCH_MAIN_PARSE(argc, argv); \ + \ + NVBENCH_MAIN_PRINT_PREAMBLE(parser); \ + NVBENCH_MAIN_RUN_BENCHMARKS(parser); \ + NVBENCH_MAIN_PRINT_EPILOGUE(parser); \ + \ + NVBENCH_MAIN_PRINT_RESULTS(parser); \ + } /* Tear down parser before finalization */ \ + NVBENCH_MAIN_FINALIZE(); \ + return 0; +#endif + +#ifndef NVBENCH_MAIN_INITIALIZE +#define NVBENCH_MAIN_INITIALIZE(argc, argv) \ + NVBENCH_MAIN_INITIALIZE_CUSTOM_PRE(argc, argv); \ + nvbench::detail::main_initialize(argc, argv); \ + NVBENCH_MAIN_INITIALIZE_CUSTOM_POST(argc, argv) +#endif + +#ifndef NVBENCH_MAIN_PARSE +#define NVBENCH_MAIN_PARSE(argc, argv) \ + NVBENCH_MAIN_CUSTOM_ARGC_ARGV_HANDLER(argc, argv); \ + std::vector args = nvbench::detail::main_convert_args(argc, argv); \ + NVBENCH_MAIN_CUSTOM_ARGS_HANDLER(args); \ + nvbench::option_parser parser; \ + NVBENCH_MAIN_PARSE_CUSTOM_PRE(parser); \ + parser.parse(args); \ + NVBENCH_MAIN_PARSE_CUSTOM_POST(parser) +#endif + +#ifndef NVBENCH_MAIN_PRINT_PREAMBLE +#define NVBENCH_MAIN_PRINT_PREAMBLE(parser) nvbench::detail::main_print_preamble(parser) +#endif + +#ifndef NVBENCH_MAIN_RUN_BENCHMARKS +#define NVBENCH_MAIN_RUN_BENCHMARKS(parser) nvbench::detail::main_run_benchmarks(parser) +#endif + +#ifndef NVBENCH_MAIN_PRINT_EPILOGUE +#define NVBENCH_MAIN_PRINT_EPILOGUE(parser) nvbench::detail::main_print_epilogue(parser) +#endif + +#ifndef NVBENCH_MAIN_PRINT_RESULTS +#define NVBENCH_MAIN_PRINT_RESULTS(parser) nvbench::detail::main_print_results(parser) +#endif + +#ifndef NVBENCH_MAIN_FINALIZE +#define NVBENCH_MAIN_FINALIZE() \ + NVBENCH_MAIN_FINALIZE_CUSTOM_PRE(); \ + nvbench::detail::main_finalize(); \ + NVBENCH_MAIN_FINALIZE_CUSTOM_POST() +#endif + +#ifndef NVBENCH_MAIN_CATCH_EXCEPTIONS +#define NVBENCH_MAIN_CATCH_EXCEPTIONS \ catch (std::exception & e) \ { \ std::cerr << "\nNVBench encountered an error:\n\n" << e.what() << "\n"; \ @@ -46,56 +164,91 @@ std::cerr << "\nNVBench encountered an unknown error.\n"; \ return 1; \ } +#endif -#ifdef NVBENCH_HAS_CUPTI -#define NVBENCH_INITIALIZE_DRIVER_API NVBENCH_DRIVER_API_CALL(cuInit(0)) +namespace nvbench::detail +{ + +inline void set_env(const char *name, const char *value) +{ +#ifdef _MSC_VER + _putenv_s(name, value); #else -// clang-format off -#define NVBENCH_INITIALIZE_DRIVER_API do {} while (false) -// clang-format on + setenv(name, value, 1); #endif +} -#define NVBENCH_MAIN_PARSE(argc, argv) \ - nvbench::option_parser parser; \ - parser.parse(argc, argv) - -// See NVIDIA/NVBench#136 for CUDA_MODULE_LOADING +inline void main_initialize(int, char **) +{ + // See NVIDIA/NVBench#136 for CUDA_MODULE_LOADING #ifdef _MSC_VER -#define NVBENCH_INITIALIZE_CUDA_ENV _putenv_s("CUDA_MODULE_LOADING", "EAGER") + _putenv_s("CUDA_MODULE_LOADING", "EAGER"); #else -#define NVBENCH_INITIALIZE_CUDA_ENV setenv("CUDA_MODULE_LOADING", "EAGER", 1) + setenv("CUDA_MODULE_LOADING", "EAGER", 1); #endif -#define NVBENCH_INITIALIZE_BENCHMARKS() \ - nvbench::benchmark_manager::get().initialize() + // Initialize CUDA driver API if needed: +#ifdef NVBENCH_HAS_CUPTI + NVBENCH_DRIVER_API_CALL(cuInit(0)); +#endif -#define NVBENCH_MAIN_BODY(argc, argv) \ - do \ - { \ - NVBENCH_INITIALIZE_CUDA_ENV; \ - NVBENCH_INITIALIZE_DRIVER_API; \ - NVBENCH_INITIALIZE_BENCHMARKS(); \ - NVBENCH_MAIN_PARSE(argc, argv); \ - auto &printer = parser.get_printer(); \ - \ - printer.print_device_info(); \ - printer.print_log_preamble(); \ - auto &benchmarks = parser.get_benchmarks(); \ - \ - std::size_t total_states = 0; \ - for (auto &bench_ptr : benchmarks) \ - { \ - total_states += bench_ptr->get_config_count(); \ - } \ - printer.set_total_state_count(total_states); \ - \ - printer.set_completed_state_count(0); \ - for (auto &bench_ptr : benchmarks) \ - { \ - bench_ptr->set_printer(printer); \ - bench_ptr->run(); \ - bench_ptr->clear_printer(); \ - } \ - printer.print_log_epilogue(); \ - printer.print_benchmark_results(benchmarks); \ - } while (false) + // Initialize the benchmarks *after* setting up the CUDA environment: + nvbench::benchmark_manager::get().initialize(); +} + +inline std::vector main_convert_args(int argc, char const *const *argv) +{ + std::vector args; + for (int i = 0; i < argc; ++i) + { + args.push_back(argv[i]); + } + return args; +} + +inline void main_print_preamble(option_parser &parser) +{ + auto &printer = parser.get_printer(); + + printer.print_device_info(); + printer.print_log_preamble(); +} + +inline void main_run_benchmarks(option_parser &parser) +{ + auto &printer = parser.get_printer(); + auto &benchmarks = parser.get_benchmarks(); + + std::size_t total_states = 0; + for (auto &bench_ptr : benchmarks) + { + total_states += bench_ptr->get_config_count(); + } + + printer.set_completed_state_count(0); + printer.set_total_state_count(total_states); + + for (auto &bench_ptr : benchmarks) + { + bench_ptr->set_printer(printer); + bench_ptr->run(); + bench_ptr->clear_printer(); + } +} + +inline void main_print_epilogue(option_parser &parser) +{ + auto &printer = parser.get_printer(); + printer.print_log_epilogue(); +} + +inline void main_print_results(option_parser &parser) +{ + auto &printer = parser.get_printer(); + auto &benchmarks = parser.get_benchmarks(); + printer.print_benchmark_results(benchmarks); +} + +inline void main_finalize() { NVBENCH_CUDA_CALL(cudaDeviceReset()); } + +} // namespace nvbench::detail