From c14514fbc5e8ace79395f952e6f26fa3cfd4b52d Mon Sep 17 00:00:00 2001
From: Joseph Redmon <pjreddie@gmail.com>
Date: Sat, 05 Nov 2016 21:10:27 +0000
Subject: [PATCH] need the actual chars

---
 src/cuda.c |   72 +++++++++++++++++++++++++----------
 1 files changed, 51 insertions(+), 21 deletions(-)

diff --git a/src/cuda.c b/src/cuda.c
index d773d0b..617e7b3 100644
--- a/src/cuda.c
+++ b/src/cuda.c
@@ -9,6 +9,20 @@
 #include <stdlib.h>
 #include <time.h>
 
+void cuda_set_device(int n)
+{
+    gpu_index = n;
+    cudaError_t status = cudaSetDevice(n);
+    check_error(status);
+}
+
+int cuda_get_device()
+{
+    int n = 0;
+    cudaError_t status = cudaGetDevice(&n);
+    check_error(status);
+    return n;
+}
 
 void check_error(cudaError_t status)
 {
@@ -38,26 +52,41 @@
     size_t x = k;
     size_t y = 1;
     if(x > 65535){
-         x = ceil(sqrt(k));
-         y = (n-1)/(x*BLOCK) + 1;
+        x = ceil(sqrt(k));
+        y = (n-1)/(x*BLOCK) + 1;
     }
     dim3 d = {x, y, 1};
     //printf("%ld %ld %ld %ld\n", n, x, y, x*y*BLOCK);
     return d;
 }
 
+#ifdef CUDNN
+cudnnHandle_t cudnn_handle()
+{
+    static int init[16] = {0};
+    static cudnnHandle_t handle[16];
+    int i = cuda_get_device();
+    if(!init[i]) {
+        cudnnCreate(&handle[i]);
+        init[i] = 1;
+    }
+    return handle[i];
+}
+#endif
+
 cublasHandle_t blas_handle()
 {
-    static int init = 0;
-    static cublasHandle_t handle;
-    if(!init) {
-        cublasCreate(&handle);
-        init = 1;
+    static int init[16] = {0};
+    static cublasHandle_t handle[16];
+    int i = cuda_get_device();
+    if(!init[i]) {
+        cublasCreate(&handle[i]);
+        init[i] = 1;
     }
-    return handle;
+    return handle[i];
 }
 
-float *cuda_make_array(float *x, int n)
+float *cuda_make_array(float *x, size_t n)
 {
     float *x_gpu;
     size_t size = sizeof(float)*n;
@@ -71,20 +100,21 @@
     return x_gpu;
 }
 
-void cuda_random(float *x_gpu, int n)
+void cuda_random(float *x_gpu, size_t n)
 {
-    static curandGenerator_t gen;
-    static int init = 0;
-    if(!init){
-        curandCreateGenerator(&gen, CURAND_RNG_PSEUDO_DEFAULT);
-        curandSetPseudoRandomGeneratorSeed(gen, time(0));
-        init = 1;
+    static curandGenerator_t gen[16];
+    static int init[16] = {0};
+    int i = cuda_get_device();
+    if(!init[i]){
+        curandCreateGenerator(&gen[i], CURAND_RNG_PSEUDO_DEFAULT);
+        curandSetPseudoRandomGeneratorSeed(gen[i], time(0));
+        init[i] = 1;
     }
-    curandGenerateUniform(gen, x_gpu, n);
+    curandGenerateUniform(gen[i], x_gpu, n);
     check_error(cudaPeekAtLastError());
 }
 
-float cuda_compare(float *x_gpu, float *x, int n, char *s)
+float cuda_compare(float *x_gpu, float *x, size_t n, char *s)
 {
     float *tmp = calloc(n, sizeof(float));
     cuda_pull_array(x_gpu, tmp, n);
@@ -97,7 +127,7 @@
     return err;
 }
 
-int *cuda_make_int_array(int n)
+int *cuda_make_int_array(size_t n)
 {
     int *x_gpu;
     size_t size = sizeof(int)*n;
@@ -112,14 +142,14 @@
     check_error(status);
 }
 
-void cuda_push_array(float *x_gpu, float *x, int n)
+void cuda_push_array(float *x_gpu, float *x, size_t n)
 {
     size_t size = sizeof(float)*n;
     cudaError_t status = cudaMemcpy(x_gpu, x, size, cudaMemcpyHostToDevice);
     check_error(status);
 }
 
-void cuda_pull_array(float *x_gpu, float *x, int n)
+void cuda_pull_array(float *x_gpu, float *x, size_t n)
 {
     size_t size = sizeof(float)*n;
     cudaError_t status = cudaMemcpy(x, x_gpu, size, cudaMemcpyDeviceToHost);

--
Gitblit v1.10.0