From c31541a4d881da79514adad0d829c0afbb3e620a Mon Sep 17 00:00:00 2001
From: Andrea Bocci <andrea.bocci@cern.ch>
Date: Thu, 2 Sep 2021 10:42:17 +0200
Subject: [PATCH] Print the NVIDIA and CUDA driver and library versions

---
 HeterogeneousCore/CUDAServices/BuildFile.xml  |  1 +
 .../CUDAServices/src/CUDAService.cc           | 41 ++++++++++++++++---
 2 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/HeterogeneousCore/CUDAServices/BuildFile.xml b/HeterogeneousCore/CUDAServices/BuildFile.xml
index 5fcaf5e5527b0..a48e1c639eaf3 100644
--- a/HeterogeneousCore/CUDAServices/BuildFile.xml
+++ b/HeterogeneousCore/CUDAServices/BuildFile.xml
@@ -5,6 +5,7 @@
   <use name="FWCore/Utilities"/>
   <use name="HeterogeneousCore/CUDAUtilities"/>
   <use name="cuda"/>
+  <use name="cuda-nvml"/>
   <export>
     <lib name="1"/>
   </export>
diff --git a/HeterogeneousCore/CUDAServices/src/CUDAService.cc b/HeterogeneousCore/CUDAServices/src/CUDAService.cc
index 1c0e1fc6d2516..a9156e1b1be9e 100644
--- a/HeterogeneousCore/CUDAServices/src/CUDAService.cc
+++ b/HeterogeneousCore/CUDAServices/src/CUDAService.cc
@@ -3,6 +3,8 @@
 #include <limits>
 
 #include <cuda.h>
+#include <cuda_runtime.h>
+#include <nvml.h>
 
 #include "FWCore/MessageLogger/interface/MessageLogger.h"
 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
@@ -10,13 +12,14 @@
 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
 #include "FWCore/Utilities/interface/ReusableObjectHolder.h"
 #include "HeterogeneousCore/CUDAServices/interface/CUDAService.h"
-#include "HeterogeneousCore/CUDAUtilities/interface/cudaCheck.h"
 #include "HeterogeneousCore/CUDAUtilities/interface/EventCache.h"
 #include "HeterogeneousCore/CUDAUtilities/interface/StreamCache.h"
+#include "HeterogeneousCore/CUDAUtilities/interface/cachingAllocators.h"
+#include "HeterogeneousCore/CUDAUtilities/interface/cudaCheck.h"
+#include "HeterogeneousCore/CUDAUtilities/interface/currentDevice.h"
 #include "HeterogeneousCore/CUDAUtilities/interface/device_unique_ptr.h"
 #include "HeterogeneousCore/CUDAUtilities/interface/host_unique_ptr.h"
-#include "HeterogeneousCore/CUDAUtilities/interface/currentDevice.h"
-#include "HeterogeneousCore/CUDAUtilities/interface/cachingAllocators.h"
+#include "HeterogeneousCore/CUDAUtilities/interface/nvmlCheck.h"
 
 void setCudaLimit(cudaLimit limit, const char* name, size_t request) {
   // read the current device
@@ -83,6 +86,10 @@ constexpr unsigned int getCudaCoresPerSM(unsigned int major, unsigned int minor)
   }
 }
 
+std::string decodeVersion(int version) {
+  return std::to_string(version / 1000) + '.' + std::to_string(version % 1000 / 10);
+}
+
 namespace {
   template <template <typename> typename UniquePtr, typename Allocate>
   void preallocate(Allocate allocate, const std::vector<unsigned int>& bufferSizes) {
@@ -132,10 +139,34 @@ CUDAService::CUDAService(edm::ParameterSet const& config) : verbose_(config.getU
     return;
   }
   computeCapabilities_.reserve(numberOfDevices_);
+
+  // NVIDIA system driver version, e.g. 470.57.02
+  char systemDriverVersion[NVML_SYSTEM_DRIVER_VERSION_BUFFER_SIZE];
+  nvmlCheck(nvmlInitWithFlags(NVML_INIT_FLAG_NO_GPUS | NVML_INIT_FLAG_NO_ATTACH));
+  nvmlCheck(nvmlSystemGetDriverVersion(systemDriverVersion, sizeof(systemDriverVersion)));
+  nvmlCheck(nvmlShutdown());
+
+  // CUDA driver version, e.g. 11.4
+  // the full version, like 11.4.1 or 11.4.100, is not reported
+  int driverVersion = 0;
+  cudaCheck(cudaDriverGetVersion(&driverVersion));
+
+  // CUDA runtime version, e.g. 11.4
+  // the full version, like 11.4.1 or 11.4.108, is not reported
+  int runtimeVersion = 0;
+  cudaCheck(cudaRuntimeGetVersion(&runtimeVersion));
+
   edm::LogInfo log("CUDAService");
-  log << "CUDA runtime successfully initialised, found " << numberOfDevices_ << " compute devices.\n";
   if (verbose_) {
-    log << '\n';
+    log << "NVIDIA driver:    " << systemDriverVersion << '\n';
+    log << "CUDA driver API:  " << decodeVersion(driverVersion) << " (compiled with " << decodeVersion(CUDA_VERSION)
+        << ")\n";
+    log << "CUDA runtime API: " << decodeVersion(runtimeVersion) << " (compiled with " << decodeVersion(CUDART_VERSION)
+        << ")\n";
+    log << "CUDA runtime successfully initialised, found " << numberOfDevices_ << " compute devices.\n\n";
+  } else {
+    log << "CUDA runtime version " << decodeVersion(runtimeVersion) << ", driver version "
+        << decodeVersion(driverVersion) << ", NVIDIA driver version " << systemDriverVersion << '\n';
   }
 
   auto const& limits = config.getUntrackedParameter<edm::ParameterSet>("limits");