From 25f2976f729a7ef008e41b07c9189a12d8d3f642 Mon Sep 17 00:00:00 2001
From: Robert Maynard <rmaynard@nvidia.com>
Date: Tue, 14 Nov 2023 14:36:32 -0500
Subject: [PATCH] Ensure nvbench initializes nvml context when built statically

---
 .../nvbench/nvml_with_static_builds.diff      | 39 +++++++++++++++++++
 rapids-cmake/cpm/versions.json                |  9 ++++-
 2 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 rapids-cmake/cpm/patches/nvbench/nvml_with_static_builds.diff

diff --git a/rapids-cmake/cpm/patches/nvbench/nvml_with_static_builds.diff b/rapids-cmake/cpm/patches/nvbench/nvml_with_static_builds.diff
new file mode 100644
index 00000000..b66f6339
--- /dev/null
+++ b/rapids-cmake/cpm/patches/nvbench/nvml_with_static_builds.diff
@@ -0,0 +1,39 @@
+diff --git a/nvbench/device_info.cu b/nvbench/device_info.cu
+index 8edad22..a09def3 100644
+--- a/nvbench/device_info.cu
++++ b/nvbench/device_info.cu
+@@ -25,6 +25,10 @@
+
+ #include <cuda_runtime_api.h>
+
++#ifdef NVBENCH_HAS_NVML
++#include <nvbench/internal/nvml.cxx>
++#endif
++
+ #define UNUSED(x) (void)(x)
+
+ namespace nvbench
+@@ -47,6 +51,10 @@ device_info::device_info(int id)
+   NVBENCH_CUDA_CALL(cudaGetDeviceProperties(&m_prop, m_id));
+
+ #ifdef NVBENCH_HAS_NVML
++  // NVML's lifetime should extend for the entirety of the process, so store in a
++  // global.
++  [[maybe_unused]] static auto nvml_lifetime = NVMLLifetimeManager();
++
+   // Retrieve the current device's pci_id as a null-terminated string.
+   // Docs say 13 chars should always be sufficient.
+   constexpr int pci_id_len = 13;
+diff --git a/nvbench/internal/nvml.cxx b/nvbench/internal/nvml.cxx
+index 4f750bc..ce7d874 100644
+--- a/nvbench/internal/nvml.cxx
++++ b/nvbench/internal/nvml.cxx
+@@ -64,8 +64,4 @@ private:
+   bool m_inited{false};
+ };
+
+-// NVML's lifetime should extend for the entirety of the process, so store in a
+-// global.
+-auto nvml_lifetime = NVMLLifetimeManager{};
+-
+ } // namespace
diff --git a/rapids-cmake/cpm/versions.json b/rapids-cmake/cpm/versions.json
index a824698d..c8efbd98 100644
--- a/rapids-cmake/cpm/versions.json
+++ b/rapids-cmake/cpm/versions.json
@@ -59,7 +59,14 @@
       "version" : "0.0",
       "git_shallow" : false,
       "git_url" : "https://github.com/NVIDIA/nvbench.git",
-      "git_tag" : "b8f0f6b5832dd632006492fd2434c81fdef0f2de"
+      "git_tag" : "b8f0f6b5832dd632006492fd2434c81fdef0f2de",
+      "patches" : [
+        {
+          "file" : "nvbench/nvml_with_static_builds.diff",
+          "issue" : "Add support for nvml with static nvbench [https://github.com/NVIDIA/nvbench/pull/148]",
+          "fixed_in" : ""
+        }
+      ]
     },
     "nvcomp" : {
       "version" : "3.0.4",