From 4dc46a0da3ed5f86125f06efe4eb8f7c9e042942 Mon Sep 17 00:00:00 2001
From: Robert Maynard <rmaynard@nvidia.com>
Date: Thu, 20 Apr 2023 11:24:19 -0400
Subject: [PATCH] nvbench executables now build with a custom nvbench main

This ensures that the `nvbench_base_fixture` we need, is constructed
before any benchmark is run.
---
 cpp/benchmarks/CMakeLists.txt                 |  4 ++--
 cpp/benchmarks/fixture/nvbench_main.cu        | 22 +++++++++++++++++++
 cpp/benchmarks/fixture/rmm_pool_raii.hpp      |  2 --
 .../patches/nvbench_global_setup.diff         |  2 +-
 4 files changed, 25 insertions(+), 5 deletions(-)
 create mode 100644 cpp/benchmarks/fixture/nvbench_main.cu

diff --git a/cpp/benchmarks/CMakeLists.txt b/cpp/benchmarks/CMakeLists.txt
index 5d3be71286d..a6a78bb977b 100644
--- a/cpp/benchmarks/CMakeLists.txt
+++ b/cpp/benchmarks/CMakeLists.txt
@@ -91,14 +91,14 @@ endfunction()
 # This function takes in a benchmark name and benchmark source for nvbench benchmarks and handles
 # setting all of the associated properties and linking to build the benchmark
 function(ConfigureNVBench CMAKE_BENCH_NAME)
-  add_executable(${CMAKE_BENCH_NAME} ${ARGN})
+  add_executable(${CMAKE_BENCH_NAME} ${ARGN} fixture/nvbench_main.cu)
   set_target_properties(
     ${CMAKE_BENCH_NAME}
     PROPERTIES RUNTIME_OUTPUT_DIRECTORY "$<BUILD_INTERFACE:${CUDF_BINARY_DIR}/benchmarks>"
                INSTALL_RPATH "\$ORIGIN/../../../lib"
   )
   target_link_libraries(
-    ${CMAKE_BENCH_NAME} PRIVATE cudf_benchmark_common cudf_datagen nvbench::main
+    ${CMAKE_BENCH_NAME} PRIVATE cudf_benchmark_common cudf_datagen nvbench::nvbench
                                 $<TARGET_NAME_IF_EXISTS:conda_env>
   )
   install(
diff --git a/cpp/benchmarks/fixture/nvbench_main.cu b/cpp/benchmarks/fixture/nvbench_main.cu
new file mode 100644
index 00000000000..f58eae62372
--- /dev/null
+++ b/cpp/benchmarks/fixture/nvbench_main.cu
@@ -0,0 +1,22 @@
+/*
+ * Copyright (c) 2023, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <benchmarks/fixture/rmm_pool_raii.hpp>
+#define NVBENCH_ENVIRONMENT cudf::nvbench_base_fixture
+
+#include <nvbench/main.cuh>
+
+NVBENCH_MAIN
diff --git a/cpp/benchmarks/fixture/rmm_pool_raii.hpp b/cpp/benchmarks/fixture/rmm_pool_raii.hpp
index 465c53a91ea..23f49735855 100644
--- a/cpp/benchmarks/fixture/rmm_pool_raii.hpp
+++ b/cpp/benchmarks/fixture/rmm_pool_raii.hpp
@@ -79,5 +79,3 @@ struct nvbench_base_fixture {
 };
 
 }  // namespace cudf
-
-#define NVBENCH_ENVIRONMENT cudf::nvbench_base_fixture
diff --git a/cpp/cmake/thirdparty/patches/nvbench_global_setup.diff b/cpp/cmake/thirdparty/patches/nvbench_global_setup.diff
index 0487b0a1ac3..aa61262fdcb 100644
--- a/cpp/cmake/thirdparty/patches/nvbench_global_setup.diff
+++ b/cpp/cmake/thirdparty/patches/nvbench_global_setup.diff
@@ -21,7 +21,7 @@ index 0ba82d7..7ab02c1 100644
      printer.set_total_state_count(total_states);                                                   \
                                                                                                     \
      printer.set_completed_state_count(0);                                                          \
-+    NVBENCH_ENVIRONMENT();                                                                         \
++    auto env_state = NVBENCH_ENVIRONMENT();                                                        \
      for (auto &bench_ptr : benchmarks)                                                             \
      {                                                                                              \
        bench_ptr->set_printer(printer);                                                             \