From bc2fa1165338faf9164fe7b03f09106ad124749f Mon Sep 17 00:00:00 2001
From: Robert Maynard <rmaynard@nvidia.com>
Date: Mon, 24 Apr 2023 10:06:01 -0400
Subject: [PATCH] Use custom nvbench entry point to ensure
 `cudf::nvbench_base_fixture` usage (#13183)

Previously we used the `nvbench::main` CMake target which would cause the same `main.o` to re-used across targets.
Since this object file was compiled without any knowledge of rmm or libcudf it wouldn't have our custom `NVBENCH_ENVIRONMENT` type. This caused our benchmarks to not execute with the rmm pool allocator.

Authors:
  - Robert Maynard (https://github.com/robertmaynard)
  - Bradley Dice (https://github.com/bdice)

Approvers:
  - David Wendt (https://github.com/davidwendt)
  - Nghia Truong (https://github.com/ttnghia)
  - Karthikeyan (https://github.com/karthikeyann)
  - Bradley Dice (https://github.com/bdice)

URL: https://github.com/rapidsai/cudf/pull/13183
---
 cpp/benchmarks/CMakeLists.txt                 |  8 +++----
 cpp/benchmarks/fixture/nvbench_main.cpp       | 22 +++++++++++++++++++
 cpp/benchmarks/fixture/rmm_pool_raii.hpp      |  2 --
 .../patches/nvbench_global_setup.diff         |  2 +-
 4 files changed, 27 insertions(+), 7 deletions(-)
 create mode 100644 cpp/benchmarks/fixture/nvbench_main.cpp

diff --git a/cpp/benchmarks/CMakeLists.txt b/cpp/benchmarks/CMakeLists.txt
index 8e230f2d5d8..565a396d913 100644
--- a/cpp/benchmarks/CMakeLists.txt
+++ b/cpp/benchmarks/CMakeLists.txt
@@ -24,8 +24,8 @@ target_compile_options(
 
 target_link_libraries(
   cudf_datagen
-  PUBLIC GTest::gmock GTest::gtest GTest::gmock_main GTest::gtest_main benchmark::benchmark
-         nvbench::nvbench Threads::Threads cudf cudftestutil
+  PUBLIC GTest::gmock GTest::gtest benchmark::benchmark nvbench::nvbench Threads::Threads cudf
+         cudftestutil
   PRIVATE $<TARGET_NAME_IF_EXISTS:conda_env>
 )
 
@@ -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.cpp)
   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.cpp b/cpp/benchmarks/fixture/nvbench_main.cpp
new file mode 100644
index 00000000000..f58eae62372
--- /dev/null
+++ b/cpp/benchmarks/fixture/nvbench_main.cpp
@@ -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);                                                             \