From a3bf6fd4ef9f83218511671247b3c10b1405d174 Mon Sep 17 00:00:00 2001
From: jsteube <jens.steube@gmail.com>
Date: Tue, 11 Oct 2016 12:18:06 +0200
Subject: [PATCH] Make use of nvmlDeviceGetCount()

---
 include/ext_nvml.h |  2 ++
 src/hwmon.c        | 47 +++++++++++++++++++++++++++++++++-------------
 2 files changed, 36 insertions(+), 13 deletions(-)

diff --git a/include/ext_nvml.h b/include/ext_nvml.h
index b9ff6444f..9342b4e6e 100644
--- a/include/ext_nvml.h
+++ b/include/ext_nvml.h
@@ -168,6 +168,7 @@ typedef nvmlDevice_t HM_ADAPTER_NVML;
 typedef const char * (*NVML_API_CALL NVML_ERROR_STRING) (nvmlReturn_t);
 typedef int (*NVML_API_CALL NVML_INIT) (void);
 typedef int (*NVML_API_CALL NVML_SHUTDOWN) (void);
+typedef nvmlReturn_t (*NVML_API_CALL NVML_DEVICE_GET_COUNT) (unsigned int *);
 typedef nvmlReturn_t (*NVML_API_CALL NVML_DEVICE_GET_NAME) (nvmlDevice_t, char *, unsigned int);
 typedef nvmlReturn_t (*NVML_API_CALL NVML_DEVICE_GET_HANDLE_BY_INDEX) (unsigned int, nvmlDevice_t *);
 typedef nvmlReturn_t (*NVML_API_CALL NVML_DEVICE_GET_TEMPERATURE) (nvmlDevice_t, nvmlTemperatureSensors_t, unsigned int *);
@@ -199,6 +200,7 @@ typedef struct hm_nvml_lib
   NVML_ERROR_STRING nvmlErrorString;
   NVML_INIT nvmlInit;
   NVML_SHUTDOWN nvmlShutdown;
+  NVML_DEVICE_GET_COUNT nvmlDeviceGetCount;
   NVML_DEVICE_GET_NAME nvmlDeviceGetName;
   NVML_DEVICE_GET_HANDLE_BY_INDEX nvmlDeviceGetHandleByIndex;
   NVML_DEVICE_GET_TEMPERATURE nvmlDeviceGetTemperature;
diff --git a/src/hwmon.c b/src/hwmon.c
index 95ad3d41f..780500479 100644
--- a/src/hwmon.c
+++ b/src/hwmon.c
@@ -79,6 +79,7 @@ static int nvml_init (hashcat_ctx_t *hashcat_ctx)
   HC_LOAD_FUNC(nvml, nvmlErrorString, NVML_ERROR_STRING, NVML, 0)
   HC_LOAD_FUNC(nvml, nvmlInit, NVML_INIT, NVML, 0)
   HC_LOAD_FUNC(nvml, nvmlShutdown, NVML_SHUTDOWN, NVML, 0)
+  HC_LOAD_FUNC(nvml, nvmlDeviceGetCount, NVML_DEVICE_GET_COUNT, NVML, 0)
   HC_LOAD_FUNC(nvml, nvmlDeviceGetName, NVML_DEVICE_GET_NAME, NVML, 0)
   HC_LOAD_FUNC(nvml, nvmlDeviceGetHandleByIndex, NVML_DEVICE_GET_HANDLE_BY_INDEX, NVML, 0)
   HC_LOAD_FUNC(nvml, nvmlDeviceGetTemperature, NVML_DEVICE_GET_TEMPERATURE, NVML, 0)
@@ -160,6 +161,26 @@ static int hm_NVML_nvmlShutdown (hashcat_ctx_t *hashcat_ctx)
   return 0;
 }
 
+static int hm_NVML_nvmlDeviceGetCount (hashcat_ctx_t *hashcat_ctx, unsigned int *deviceCount)
+{
+  hwmon_ctx_t *hwmon_ctx = hashcat_ctx->hwmon_ctx;
+
+  NVML_PTR *nvml = hwmon_ctx->hm_nvml;
+
+  const nvmlReturn_t nvml_rc = nvml->nvmlDeviceGetCount (deviceCount);
+
+  if (nvml_rc != NVML_SUCCESS)
+  {
+    const char *string = hm_NVML_nvmlErrorString (nvml, nvml_rc);
+
+    event_log_error (hashcat_ctx, "nvmlDeviceGetCount(): %s", string);
+
+    return -1;
+  }
+
+  return 0;
+}
+
 static int hm_NVML_nvmlDeviceGetHandleByIndex (hashcat_ctx_t *hashcat_ctx, unsigned int index, nvmlDevice_t *device)
 {
   hwmon_ctx_t *hwmon_ctx = hashcat_ctx->hwmon_ctx;
@@ -1869,27 +1890,27 @@ static int hm_get_adapter_index_nvapi (hashcat_ctx_t *hashcat_ctx, HM_ADAPTER_NV
 
 static int hm_get_adapter_index_nvml (hashcat_ctx_t *hashcat_ctx, HM_ADAPTER_NVML *nvmlGPUHandle)
 {
-  int pGpuCount = 0;
+  unsigned int deviceCount = 0;
 
-  for (u32 i = 0; i < DEVICES_MAX; i++)
-  {
-    if (hm_NVML_nvmlDeviceGetHandleByIndex (hashcat_ctx, i, &nvmlGPUHandle[i]) == -1) break;
+  hm_NVML_nvmlDeviceGetCount (hashcat_ctx, &deviceCount);
 
-    // can be used to determine if the device by index matches the cuda device by index
-    // char name[100]; memset (name, 0, sizeof (name));
-    // hm_NVML_nvmlDeviceGetName (hashcat_ctx, nvGPUHandle[i], name, sizeof (name) - 1);
-
-    pGpuCount++;
-  }
-
-  if (pGpuCount == 0)
+  if (deviceCount == 0)
   {
     event_log_error (hashcat_ctx, "No NVML adapters found");
 
     return 0;
   }
 
-  return (pGpuCount);
+  for (u32 i = 0; i < deviceCount; i++)
+  {
+    if (hm_NVML_nvmlDeviceGetHandleByIndex (hashcat_ctx, i, &nvmlGPUHandle[i]) == -1) break;
+
+    // can be used to determine if the device by index matches the cuda device by index
+    // char name[100]; memset (name, 0, sizeof (name));
+    // hm_NVML_nvmlDeviceGetName (hashcat_ctx, nvGPUHandle[i], name, sizeof (name) - 1);
+  }
+
+  return (deviceCount);
 }
 
 static void hm_sort_adl_adapters_by_busid_devid (u32 *valid_adl_device_list, int num_adl_adapters, LPAdapterInfo lpAdapterInfo)