Skip to content

Commit

Permalink
Refactor main implementation to improve reusability and customization. (
Browse files Browse the repository at this point in the history 
#165)

* 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.

* Add tests that validate common main customizations.
  • Loading branch information
@alliepiper
alliepiper committed Apr 9, 2024
1 parent 9e8efa2 commit 165cf92
Show file tree
Hide file tree
Showing 4 changed files with 403 additions and 49 deletions.
241 changes: 194 additions & 47 deletions nvbench/main.cuh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,14 +28,130 @@
#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.
// argc/argv are the usual command line arguments types. The ARGS version of this
// macro is a bit more convenient.
#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.
// Args is a vector of strings, each element is an argument.
#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, args) []() {}()
#endif

// Customization point, called after NVBench parsing.
#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, args); \
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 +162,87 @@
std::cerr << "\nNVBench encountered an unknown error.\n"; \
return 1; \
}

#ifdef NVBENCH_HAS_CUPTI
#define NVBENCH_INITIALIZE_DRIVER_API NVBENCH_DRIVER_API_CALL(cuInit(0))
#else
// clang-format off
#define NVBENCH_INITIALIZE_DRIVER_API do {} while (false)
// clang-format on
#endif

#define NVBENCH_MAIN_PARSE(argc, argv) \
nvbench::option_parser parser; \
parser.parse(argc, argv)
namespace nvbench::detail
{

// See NVIDIA/NVBench#136 for CUDA_MODULE_LOADING
inline void set_env(const char *name, const char *value)
{
#ifdef _MSC_VER
#define NVBENCH_INITIALIZE_CUDA_ENV _putenv_s("CUDA_MODULE_LOADING", "EAGER")
_putenv_s(name, value);
#else
#define NVBENCH_INITIALIZE_CUDA_ENV setenv("CUDA_MODULE_LOADING", "EAGER", 1)
setenv(name, value, 1);
#endif
}

#define NVBENCH_INITIALIZE_BENCHMARKS() \
nvbench::benchmark_manager::get().initialize()
inline void main_initialize(int, char **)
{
// See NVIDIA/NVBench#136 for CUDA_MODULE_LOADING
set_env("CUDA_MODULE_LOADING", "EAGER");

#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 CUDA driver API if needed:
#ifdef NVBENCH_HAS_CUPTI
NVBENCH_DRIVER_API_CALL(cuInit(0));
#endif

// 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
15 changes: 13 additions & 2 deletions testing/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@ set(test_srcs
cpu_timer.cu
criterion_manager.cu
criterion_params.cu
custom_main_custom_args.cu
custom_main_custom_exceptions.cu
enum_type_list.cu
entropy_criterion.cu
float64_axis.cu
Expand All @@ -24,7 +26,12 @@ set(test_srcs
type_list.cu
)

# Metatarget for all examples:
# Custom arguments:
# CTest commands+args can't be modified after creation, so we need to rely on substitution.
set(NVBench_TEST_ARGS_nvbench.test.custom_main_custom_args "--quiet" "--my-custom-arg" "--run-once" "-d" "0")
set(NVBench_TEST_ARGS_nvbench.test.custom_main_custom_exceptions "--quiet" "--run-once" "-d" "0")

# Metatarget for all tests:
add_custom_target(nvbench.test.all)
add_dependencies(nvbench.all nvbench.test.all)

Expand All @@ -36,10 +43,14 @@ foreach(test_src IN LISTS test_srcs)
target_link_libraries(${test_name} PRIVATE nvbench::nvbench fmt)
set_target_properties(${test_name} PROPERTIES COMPILE_FEATURES cuda_std_17)
nvbench_config_target(${test_name})
add_test(NAME ${test_name} COMMAND "$<TARGET_FILE:${test_name}>")
add_test(NAME ${test_name} COMMAND "$<TARGET_FILE:${test_name}>" ${NVBench_TEST_ARGS_${test_name}})

add_dependencies(nvbench.test.all ${test_name})
endforeach()

set_tests_properties(nvbench.test.custom_main_custom_exceptions PROPERTIES
PASS_REGULAR_EXPRESSION "Custom error detected: Expected exception thrown."
)

add_subdirectory(cmake)
add_subdirectory(device)
Loading

0 comments on commit 165cf92

Please sign in to comment.