Add NVIDIA stress loader selection and DCGM 4 support

iso/builder/bee-gpu-stress.c

@@ -29,6 +29,7 @@ typedef void *CUfunction;
 typedef void *CUstream;
 
 #define CU_SUCCESS 0
+#define CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT 16
 #define CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR 75
 #define CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR 76
 
@@ -97,6 +98,9 @@ typedef CUresult (*cuLaunchKernel_fn)(CUfunction,
                                       CUstream,
                                       void **,
                                       void **);
+typedef CUresult (*cuMemGetInfo_fn)(size_t *, size_t *);
+typedef CUresult (*cuStreamCreate_fn)(CUstream *, unsigned int);
+typedef CUresult (*cuStreamDestroy_fn)(CUstream);
 typedef CUresult (*cuGetErrorName_fn)(CUresult, const char **);
 typedef CUresult (*cuGetErrorString_fn)(CUresult, const char **);
 
@@ -118,6 +122,9 @@ struct cuda_api {
     cuModuleLoadDataEx_fn cuModuleLoadDataEx;
     cuModuleGetFunction_fn cuModuleGetFunction;
     cuLaunchKernel_fn cuLaunchKernel;
+    cuMemGetInfo_fn cuMemGetInfo;
+    cuStreamCreate_fn cuStreamCreate;
+    cuStreamDestroy_fn cuStreamDestroy;
     cuGetErrorName_fn cuGetErrorName;
     cuGetErrorString_fn cuGetErrorString;
 };
@@ -128,9 +135,10 @@ struct stress_report {
     int cc_major;
     int cc_minor;
     int buffer_mb;
+    int stream_count;
     unsigned long iterations;
     uint64_t checksum;
-    char details[1024];
+    char details[16384];
 };
 
 static int load_symbol(void *lib, const char *name, void **out) {
@@ -144,7 +152,7 @@ static int load_cuda(struct cuda_api *api) {
     if (!api->lib) {
         return 0;
     }
-    return
+    if (!(
         load_symbol(api->lib, "cuInit", (void **)&api->cuInit) &&
         load_symbol(api->lib, "cuDeviceGetCount", (void **)&api->cuDeviceGetCount) &&
         load_symbol(api->lib, "cuDeviceGet", (void **)&api->cuDeviceGet) &&
@@ -160,7 +168,17 @@ static int load_cuda(struct cuda_api *api) {
         load_symbol(api->lib, "cuMemcpyDtoH_v2", (void **)&api->cuMemcpyDtoH) &&
         load_symbol(api->lib, "cuModuleLoadDataEx", (void **)&api->cuModuleLoadDataEx) &&
         load_symbol(api->lib, "cuModuleGetFunction", (void **)&api->cuModuleGetFunction) &&
-        load_symbol(api->lib, "cuLaunchKernel", (void **)&api->cuLaunchKernel);
+        load_symbol(api->lib, "cuLaunchKernel", (void **)&api->cuLaunchKernel))) {
+        dlclose(api->lib);
+        memset(api, 0, sizeof(*api));
+        return 0;
+    }
+    load_symbol(api->lib, "cuMemGetInfo_v2", (void **)&api->cuMemGetInfo);
+    load_symbol(api->lib, "cuStreamCreate", (void **)&api->cuStreamCreate);
+    if (!load_symbol(api->lib, "cuStreamDestroy_v2", (void **)&api->cuStreamDestroy)) {
+        load_symbol(api->lib, "cuStreamDestroy", (void **)&api->cuStreamDestroy);
+    }
+    return 1;
 }
 
 static const char *cu_error_name(struct cuda_api *api, CUresult rc) {
@@ -220,6 +238,39 @@ static int query_compute_capability(struct cuda_api *api, CUdevice dev, int *maj
     return 1;
 }
 
+static int query_multiprocessor_count(struct cuda_api *api, CUdevice dev, int *count) {
+    int mp_count = 0;
+    if (!check_rc(api,
+                  "cuDeviceGetAttribute(multiprocessors)",
+                  api->cuDeviceGetAttribute(&mp_count, CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT, dev))) {
+        return 0;
+    }
+    *count = mp_count;
+    return 1;
+}
+
+static size_t clamp_budget_to_free_memory(struct cuda_api *api, size_t requested_bytes) {
+    size_t free_bytes = 0;
+    size_t total_bytes = 0;
+    size_t max_bytes = requested_bytes;
+
+    if (!api->cuMemGetInfo) {
+        return requested_bytes;
+    }
+    if (api->cuMemGetInfo(&free_bytes, &total_bytes) != CU_SUCCESS || free_bytes == 0) {
+        return requested_bytes;
+    }
+
+    max_bytes = (free_bytes * 9u) / 10u;
+    if (max_bytes < (size_t)4u * 1024u * 1024u) {
+        max_bytes = (size_t)4u * 1024u * 1024u;
+    }
+    if (requested_bytes > max_bytes) {
+        return max_bytes;
+    }
+    return requested_bytes;
+}
+
 #if HAVE_CUBLASLT_HEADERS
 static void append_detail(char *buf, size_t cap, const char *fmt, ...) {
     size_t len = strlen(buf);
@@ -1095,13 +1146,16 @@ static int run_cublaslt_stress(struct cuda_api *cuda,
 int main(int argc, char **argv) {
     int seconds = 5;
     int size_mb = 64;
+    int device_index = 0;
     for (int i = 1; i < argc; i++) {
         if ((strcmp(argv[i], "--seconds") == 0 || strcmp(argv[i], "-t") == 0) && i + 1 < argc) {
             seconds = atoi(argv[++i]);
         } else if ((strcmp(argv[i], "--size-mb") == 0 || strcmp(argv[i], "-m") == 0) && i + 1 < argc) {
             size_mb = atoi(argv[++i]);
+        } else if ((strcmp(argv[i], "--device") == 0 || strcmp(argv[i], "-d") == 0) && i + 1 < argc) {
+            device_index = atoi(argv[++i]);
         } else {
-            fprintf(stderr, "usage: %s [--seconds N] [--size-mb N]\n", argv[0]);
+            fprintf(stderr, "usage: %s [--seconds N] [--size-mb N] [--device N]\n", argv[0]);
             return 2;
         }
     }
@@ -1111,6 +1165,9 @@ int main(int argc, char **argv) {
     if (size_mb <= 0) {
         size_mb = 64;
     }
+    if (device_index < 0) {
+        device_index = 0;
+    }
 
     struct cuda_api cuda;
     if (!load_cuda(&cuda)) {
@@ -1133,8 +1190,13 @@ int main(int argc, char **argv) {
         return 1;
     }
 
+    if (device_index >= count) {
+        fprintf(stderr, "device index %d out of range (found %d CUDA device(s))\n", device_index, count);
+        return 1;
+    }
+
     CUdevice dev = 0;
-    if (!check_rc(&cuda, "cuDeviceGet", cuda.cuDeviceGet(&dev, 0))) {
+    if (!check_rc(&cuda, "cuDeviceGet", cuda.cuDeviceGet(&dev, device_index))) {
         return 1;
     }
 
@@ -1162,6 +1224,7 @@ int main(int argc, char **argv) {
     }
 
     printf("device=%s\n", report.device);
+    printf("device_index=%d\n", device_index);
     printf("compute_capability=%d.%d\n", report.cc_major, report.cc_minor);
     printf("backend=%s\n", report.backend);
     printf("duration_s=%d\n", seconds);