Merge pull request #476 from notoraptor/improve-cuda-detection-win

Loop over supported CUDA versions to find installed CUDA on Windows and Mac.

Author: lamblin

src/gpuarray_buffer_cuda.c

@@ -129,7 +129,6 @@ static int setup_done = 0;
 static int major = -1;
 static int minor = -1;
 static int setup_lib(void) {
-  const char *ver;
   CUresult err;
   int res, tmp;
 
@@ -140,20 +139,28 @@ static int setup_lib(void) {
     err = cuInit(0);
     if (err != CUDA_SUCCESS)
       return error_cuda(global_err, "cuInit", err);
-    ver = getenv("GPUARRAY_CUDA_VERSION");
-    if (ver == NULL || strlen(ver) != 2) {
-      err = cuDriverGetVersion(&tmp);
-      if (err != CUDA_SUCCESS)
-        return error_set(global_err, GA_IMPL_ERROR, "cuDriverGetVersion failed");
-      major = tmp / 1000;
-      minor = (tmp / 10) % 10;
-    } else {
-      major = ver[0] - '0';
-      minor = ver[1] - '0';
-    }
-    if (major > 9 || major < 0 || minor > 9 || minor < 0)
-      return error_fmt(global_err, GA_VALUE_ERROR, "Invalid cuda version: %d.%d", major, minor);
+    err = cuDriverGetVersion(&tmp);
+    if (err != CUDA_SUCCESS)
+      return error_set(global_err, GA_IMPL_ERROR, "cuDriverGetVersion failed");
+    major = tmp / 1000;
+    minor = (tmp / 10) % 10;
+    /* Let's try to load a nvrtc corresponding to detected CUDA version. */
     res = load_libnvrtc(major, minor, global_err);
+    if (res != GA_NO_ERROR) {
+      /* Else, let's try to find a nvrtc corresponding to supported CUDA versions. */
+      int versions[][2] = {{8, 0}, {7, 5}, {7, 0}};
+      int versions_length = sizeof(versions) / sizeof(versions[0]);
+      int i = 0;
+      /* Skip versions that are higher or equal to the driver version */
+      while (versions[i][0] > major ||
+             (versions[i][0] == major && versions[i][1] >= minor)) i++;
+      do {
+        major = versions[i][0];
+        minor = versions[i][1];
+        res = load_libnvrtc(major, minor, global_err);
+        i++;
+      } while (res != GA_NO_ERROR && i < versions_length);
+    }
     if (res != GA_NO_ERROR)
       return res;
     setup_done = 1;

src/loaders/libcublas.c

@@ -1,4 +1,5 @@
-#include <stdlib.h>
+/* To be able to use snprintf with any compiler including MSVC2008. */
+#include <private_config.h>
 
 #include "libcublas.h"
 #include "dyn_load.h"
@@ -41,21 +42,27 @@ int load_libcublas(int major, int minor, error *e) {
 
 #if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)
   {
-    static const char DIGITS[] = "0123456789";
-    char libname[] = "cublas64_??.dll";
-
-    libname[9] = DIGITS[major];
-    libname[10] = DIGITS[minor];
-
+    char libname[64];
+    int n;
+    #ifdef DEBUG
+    fprintf(stderr, "Loading cuBLAS %d.%d.\n", major, minor);
+    #endif
+    n = snprintf(libname, sizeof(libname), "cublas64_%d%d.dll", major, minor);
+    if (n < 0 || n >= sizeof(libname))
+      return error_set(e, GA_SYS_ERROR, "snprintf");
     lib = ga_load_library(libname, e);
   }
 #else /* Unix */
 #ifdef __APPLE__
   {
-    static const char DIGITS[] = "0123456789";
-    char libname[] = "/Developer/NVIDIA/CUDA-?.?/lib/libcublas.dylib";
-    libname[23] = DIGITS[major];
-    libname[25] = DIGITS[minor];
+    char libname[128];
+    int n;
+    #ifdef DEBUG
+    fprintf(stderr, "Loading cuBLAS %d.%d.\n", major, minor);
+    #endif
+    n = snprintf(libname, sizeof(libname), "/Developer/NVIDIA/CUDA-%d.%d/lib/libcublas.dylib", major, minor);
+    if (n < 0 || n >= sizeof(libname))
+      return error_set(e, GA_SYS_ERROR, "snprintf");
     lib = ga_load_library(libname, e);
   }
 #else

src/loaders/libnvrtc.c

@@ -1,4 +1,5 @@
-#include <stdlib.h>
+/* To be able to use snprintf with any compiler including MSVC2008. */
+#include <private_config.h>
 
 #include "libcuda.h"
 #include "libnvrtc.h"
@@ -27,22 +28,28 @@ int load_libnvrtc(int major, int minor, error *e) {
 
 #if defined(WIN32) || defined(_WIN32) || defined(WIN64) || defined(_WIN64)
   {
-    static const char DIGITS[] = "0123456789";
-    char libname[] = "nvrtc64_??.dll";
-
-    libname[8] = DIGITS[major];
-    libname[9] = DIGITS[minor];
+    char libname[64];
+    int n;
+    #ifdef DEBUG
+    fprintf(stderr, "Loading nvrtc %d.%d.\n", major, minor);
+    #endif
+    n = snprintf(libname, sizeof(libname), "nvrtc64_%d%d.dll", major, minor);
+    if (n < 0 || n >= sizeof(libname))
+      return error_set(e, GA_SYS_ERROR, "snprintf");
 
     lib = ga_load_library(libname, e);
   }
 #else /* Unix */
 #ifdef __APPLE__
   {
-    static const char DIGITS[] = "0123456789";
-    /* Try the usual fullpath first */
-    char libname[] = "/Developer/NVIDIA/CUDA-?.?/lib/libnvrtc.dylib";
-    libname[23] = DIGITS[major];
-    libname[25] = DIGITS[minor];
+    char libname[128];
+    int n;
+    #ifdef DEBUG
+    fprintf(stderr, "Loading nvrtc %d.%d.\n", major, minor);
+    #endif
+    n = snprintf(libname, sizeof(libname), "/Developer/NVIDIA/CUDA-%d.%d/lib/libnvrtc.dylib", major, minor);
+    if (n < 0 || n >= sizeof(libname))
+      return error_set(e, GA_SYS_ERROR, "snprintf");
     lib = ga_load_library(libname, e);
   }
 #else