Merge pull request #3397 from vesis84/master

theano/sandbox/cuda/__init__.py

@@ -404,7 +404,11 @@ def use(device,
             pycuda_init_dev = theano.misc.pycuda_init.pycuda_available
 
         try:
-            if (device != 'gpu') and not pycuda_init_dev:
+            if pycuda_init_dev:
+                use.device_number = active_device_number()
+                # This is needed to initialize the cublas handle.
+                gpu_init(use.device_number, config.lib.cnmem)
+            elif(device != 'gpu'):
                 assert isinstance(device, int)
                 gpu_init(device, config.lib.cnmem)
                 use.device_number = device

theano/sandbox/cuda/cuda_ndarray.cu

@@ -2948,21 +2948,31 @@ PyObject *
 CudaNdarray_select_a_gpu(PyObject* _unused, PyObject* dummy)
 {
     void * rval = NULL;
+    cudaError_t err;
+    int num_gpus = 0;
 
-    cudaError_t err = cudaMalloc(&rval, 4);
+    err = cudaGetDeviceCount(&num_gpus);
     if (cudaSuccess != err){
         printf("ERR!\\n");
             PyErr_Format(PyExc_RuntimeError,
-                         "Not able to do basic stuff on the GPU (alloc of 4 bytes) (%s).",
+                         "Not able to get number of GPUs (%s).",
                          cudaGetErrorString(err));
             return NULL;
     }
-    err = cudaFree(rval);
+
+    for (int device = 0; device < num_gpus; device++) {
+        cudaSetDevice(device);
+        err = cudaDeviceSynchronize(); // << CUDA context gets created here.
+        cudaGetLastError(); // reset the error state     
+        if (cudaSuccess == err)
+            break;
+    }
+        
     if (cudaSuccess != err){
             printf("ERR!\\n");
                 PyErr_Format(PyExc_RuntimeError,
-                         "Not able to do basic stuff on the GPU (cudaFree failed) (%s).",
-                         cudaGetErrorString(err));
+                             "Not able to select available GPU from %d cards (%s).",
+                             num_gpus, cudaGetErrorString(err));
                 return NULL;
     }