Convert union-find benchmark to use Google benchmark

benchmarks/union_find/CMakeLists.txt

@@ -1,3 +1,7 @@
+# We require version 1.5.0 or higher but the format used by Google benchmark is
+# wrong and thus, we cannot check the version during the configuration step.
+find_package(benchmark REQUIRED)
+
 add_executable(ArborX_Benchmark_UnionFind.exe union_find.cpp)
-target_link_libraries(ArborX_Benchmark_UnionFind.exe ArborX::ArborX Boost::program_options)
-add_test(NAME ArborX_Benchmark_UnionFind COMMAND ./ArborX_Benchmark_UnionFind.exe)
+target_link_libraries(ArborX_Benchmark_UnionFind.exe ArborX::ArborX benchmark::benchmark Boost::program_options)
+add_test(NAME ArborX_Benchmark_UnionFind COMMAND ./ArborX_Benchmark_UnionFind.exe --benchmark_color=true)

benchmarks/union_find/union_find.cpp

@@ -20,6 +20,14 @@
 
 #include <boost/program_options.hpp>
 
+#include <benchmark/benchmark.h>
+
+struct Spec
+{
+  int num_edges;
+  bool allow_loops;
+};
+
 struct UnweightedEdge
 {
   unsigned int source;
@@ -100,25 +108,84 @@ auto buildUnionFind(ExecutionSpace const &exec_space, int n)
     return ArborX::Details::UnionFind<MemorySpace, /*DoSerial*/ false>(labels);
 }
 
-template <typename ExecutionSpace, typename Edges, typename UnionFind>
-double benchmark(ExecutionSpace const &exec_space, Edges edges,
-                 UnionFind union_find)
+template <typename ExecutionSpace>
+void BM_union_find(benchmark::State &state, Spec const &spec)
 {
-  Kokkos::fence();
-  Kokkos::Timer timer;
-  Kokkos::parallel_for(
-      "ArborX::Bechmark::union-find",
-      Kokkos::RangePolicy<ExecutionSpace>(exec_space, 0, edges.size()),
-      KOKKOS_LAMBDA(int e) {
-        int i = edges(e).source;
-        int j = edges(e).target;
-
-        union_find.merge(i, j);
-      });
-  Kokkos::fence();
-  return timer.seconds();
+  ExecutionSpace exec_space;
+
+  auto const n = spec.num_edges + 1;
+
+  auto edges = buildEdges(exec_space, spec.num_edges, spec.allow_loops);
+  auto union_find = buildUnionFind(exec_space, n);
+
+  for (auto _ : state)
+  {
+    exec_space.fence();
+    auto const start = std::chrono::high_resolution_clock::now();
+
+    Kokkos::parallel_for(
+        "ArborX::Bechmark::union-find",
+        Kokkos::RangePolicy<ExecutionSpace>(exec_space, 0, edges.size()),
+        KOKKOS_LAMBDA(int e) {
+          int i = edges(e).source;
+          int j = edges(e).target;
+
+          union_find.merge(i, j);
+        });
+
+    exec_space.fence();
+    auto const end = std::chrono::high_resolution_clock::now();
+    std::chrono::duration<double> elapsed_seconds = end - start;
+    state.SetIterationTime(elapsed_seconds.count());
+  }
+  state.counters["rate"] = benchmark::Counter(
+      spec.num_edges, benchmark::Counter::kIsIterationInvariantRate);
 }
 
+// NOTE Motivation for this class that stores the argument count and values is
+// I could not figure out how to make the parser consume arguments with
+// Boost.Program_options
+// Benchmark removes its own arguments from the command line arguments. This
+// means, that by virtue of returning references to internal data members in
+// argc() and argv() function, it will necessarily modify the members. It will
+// decrease _argc, and "reduce" _argv data. Hence, we must keep a copy of _argv
+// that is not modified from the outside to release memory in the destructor
+// correctly.
+class CmdLineArgs
+{
+private:
+  int _argc;
+  std::vector<char *> _argv;
+  std::vector<char *> _owner_ptrs;
+
+public:
+  CmdLineArgs(std::vector<std::string> const &args, char const *exe)
+      : _argc(args.size() + 1)
+      , _owner_ptrs{new char[std::strlen(exe) + 1]}
+  {
+    std::strcpy(_owner_ptrs[0], exe);
+    _owner_ptrs.reserve(_argc);
+    for (auto const &s : args)
+    {
+      _owner_ptrs.push_back(new char[s.size() + 1]);
+      std::strcpy(_owner_ptrs.back(), s.c_str());
+    }
+    _argv = _owner_ptrs;
+  }
+
+  ~CmdLineArgs()
+  {
+    for (auto *p : _owner_ptrs)
+    {
+      delete[] p;
+    }
+  }
+
+  int &argc() { return _argc; }
+
+  char **argv() { return _argv.data(); }
+};
+
 int main(int argc, char *argv[])
 {
   Kokkos::ScopeGuard guard(argc, argv);
@@ -129,37 +196,64 @@ int main(int argc, char *argv[])
 
   namespace bpo = boost::program_options;
 
-  int num_edges;
+  Spec spec;
   bool disallow_loops;
 
   bpo::options_description desc("Allowed options");
   // clang-format off
   desc.add_options()
       ( "help", "help message" )
       ( "no-loops", bpo::bool_switch(&disallow_loops)->default_value(false), "disallow loops in the graph")
-      ( "n", bpo::value<int>(&num_edges)->default_value(50000000), "number of edges in the graph" )
+      ( "n", bpo::value<int>(&spec.num_edges)->default_value(50000000), "number of edges in the graph" )
       ;
   // clang-format on
   bpo::variables_map vm;
-  bpo::store(bpo::command_line_parser(argc, argv).options(desc).run(), vm);
+  bpo::parsed_options parsed = bpo::command_line_parser(argc, argv)
+                                   .options(desc)
+                                   .allow_unregistered()
+                                   .run();
+  bpo::store(parsed, vm);
+  CmdLineArgs pass_further{
+      bpo::collect_unrecognized(parsed.options, bpo::include_positional),
+      argv[0]};
   bpo::notify(vm);
 
   if (vm.count("help") > 0)
   {
-    std::cout << desc << '\n';
+    // Full list of options consists of Kokkos + Boost.Program_options +
+    // Google Benchmark and we still need to call benchmark::Initialize() to
+    // get those printed to the standard output.
+    std::cout << desc << "\n";
+    int ac = 2;
+    char *av[] = {(char *)"ignored", (char *)"--help"};
+    // benchmark::Initialize() calls exit(0) when `--help` so register
+    // Kokkos::finalize() to be called on normal program termination.
+    std::atexit(Kokkos::finalize);
+    benchmark::Initialize(&ac, av);
     return 1;
   }
-  bool allow_loops = !disallow_loops;
 
-  printf("allow loops       : %s\n", (allow_loops ? "yes" : "no"));
-  printf("n                 : %d\n", num_edges);
-  auto const n = num_edges + 1;
+  benchmark::Initialize(&pass_further.argc(), pass_further.argv());
+  // Throw if some of the arguments have not been recognized.
+  std::ignore =
+      bpo::command_line_parser(pass_further.argc(), pass_further.argv())
+          .options(bpo::options_description(""))
+          .run();
+
+  spec.allow_loops = !disallow_loops;
+
+  printf("allow loops       : %s\n", (spec.allow_loops ? "yes" : "no"));
+  printf("n                 : %d\n", spec.num_edges);
 
-  Kokkos::DefaultExecutionSpace exec_space;
+  using ExecutionSpace = Kokkos::DefaultExecutionSpace;
+  benchmark::RegisterBenchmark("union_find",
+                               [=](benchmark::State &state) {
+                                 BM_union_find<ExecutionSpace>(state, spec);
+                               })
+      ->UseManualTime()
+      ->Unit(benchmark::kMicrosecond);
 
-  auto edges = buildEdges(exec_space, num_edges, allow_loops);
-  printf("Rate : %.2lf M/sec\n",
-         n / 1e6 / benchmark(exec_space, edges, buildUnionFind(exec_space, n)));
+  benchmark::RunSpecifiedBenchmarks();
 
   return EXIT_SUCCESS;
 }