Extend api

cmake/cudawrappers-targets.cmake

@@ -4,9 +4,10 @@ get_filename_component(
 )
 
 # Define all the individual components that cudawrappers provides
-set(CUDAWRAPPERS_COMPONENTS cu cufft nvrtc nvtx)
+set(CUDAWRAPPERS_COMPONENTS cu cufft nvml nvrtc nvtx)
 set(LINK_cu CUDA::cuda_driver)
 set(LINK_cufft CUDA::cuda_driver CUDA::cufft)
+set(LINK_nvml CUDA::cuda_driver CUDA::nvml)
 set(LINK_nvrtc CUDA::cuda_driver CUDA::nvrtc)
 set(LINK_nvtx CUDA::nvToolsExt)
 

include/cudawrappers/cu.hpp

@@ -4,6 +4,7 @@
 #include <array>
 #include <cstddef>
 #include <exception>
+#include <iomanip>
 #include <map>
 #include <memory>
 #include <stdexcept>
@@ -118,6 +119,29 @@ class Device : public Wrapper<CUdevice> {
     return {name.data()};
   }
 
+  std::string getUuid() const {
+    CUuuid uuid;
+    checkCudaCall(cuDeviceGetUuid(&uuid, _obj));
+
+    // Convert a CUuuid to CUDA's string representation.
+    // The CUuuid contains an array of 16 bytes, the UUID has
+    // the form 'GPU-XXXXXXXX-XXXX-XXXX-XXXX-XXXXXXXX', with every
+    // X being an alphanumeric character.
+    std::stringstream result;
+    result << "GPU";
+
+    for (int i = 0; i < 16; ++i) {
+      if (i == 0 || i == 4 || i == 6 || i == 8 || i == 10) {
+        result << "-";
+      }
+      result << std::hex << std::setfill('0') << std::setw(2)
+             << static_cast<unsigned>(
+                    static_cast<unsigned char>(uuid.bytes[i]));
+    }
+
+    return result.str();
+  }
+
   size_t totalMem() const {
     size_t size{};
     checkCudaCall(cuDeviceTotalMem(&size, _obj));

include/cudawrappers/nvml.hpp

@@ -0,0 +1,51 @@
+#if !defined NVML_H
+#define NVML_H
+
+#include <nvml.h>
+
+#include <exception>
+
+#include <cudawrappers/cu.hpp>
+
+namespace nvml {
+class Error : public std::exception {
+ public:
+  explicit Error(nvmlReturn_t result) : _result(result) {}
+
+  const char* what() const noexcept { return nvmlErrorString(_result); }
+
+  operator nvmlReturn_t() const { return _result; }
+
+ private:
+  nvmlReturn_t _result;
+};
+
+inline void checkNvmlCall(nvmlReturn_t result) {
+  if (result != NVML_SUCCESS) throw Error(result);
+}
+
+class Device {
+ public:
+  Device(int index) {
+    checkNvmlCall(nvmlInit());
+    checkNvmlCall(nvmlDeviceGetHandleByIndex(index, &device_));
+  }
+
+  Device(cu::Device& device) {
+    checkNvmlCall(nvmlInit());
+    const std::string uuid = device.getUuid();
+    nvmlDeviceGetHandleByUUID(uuid.c_str(), &device_);
+  }
+
+  ~Device() { checkNvmlCall(nvmlShutdown()); }
+
+  void getFieldValues(int valuesCount, nvmlFieldValue_t* values) {
+    checkNvmlCall(nvmlDeviceGetFieldValues(device_, valuesCount, values));
+  }
+
+ private:
+  nvmlDevice_t device_;
+};
+}  // namespace nvml
+
+#endif