Skip to content

Commit

Permalink
Refactor main implementation to improve reusability and customization.
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@alliepiper
alliepiper committed Apr 7, 2024
1 parent a0f2fab commit 4613437
Showing 1 changed file with 199 additions and 46 deletions.
245 changes: 199 additions & 46 deletions nvbench/main.cuh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,14 +28,132 @@
#include <cstdlib>
#include <iostream>

// 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<std::string> 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"; \
Expand All @@ -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<std::string> main_convert_args(int argc, char const *const *argv)
{
std::vector<std::string> 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

0 comments on commit 4613437

Please sign in to comment.