fix the gpu initialization for multiple cards

GPy/kern/_src/psi_comp/rbf_psi_gpucomp.py

@@ -242,7 +242,6 @@ gpu_code = """
 class PSICOMP_RBF_GPU(PSICOMP_RBF):
 
     def __init__(self, threadnum=128, blocknum=15, GPU_direct=False):
-        assert gpu_init.initSuccess, "GPU initialization failed!"
         self.GPU_direct = GPU_direct
         self.cublas_handle = gpu_init.cublas_handle
         self.gpuCache = None

GPy/kern/_src/psi_comp/ssrbf_psi_gpucomp.py

@@ -292,7 +292,6 @@ gpu_code = """
 class PSICOMP_SSRBF_GPU(PSICOMP_RBF):
 
     def __init__(self, threadnum=128, blocknum=15, GPU_direct=False):
-        assert gpu_init.initSuccess, "GPU initialization failed!"
         self.GPU_direct = GPU_direct
         self.cublas_handle = gpu_init.cublas_handle
         self.gpuCache = None

GPy/kern/_src/rbf.py

@@ -6,7 +6,6 @@ import numpy as np
 from stationary import Stationary
 from psi_comp import PSICOMP_RBF
 from psi_comp.rbf_psi_gpucomp import PSICOMP_RBF_GPU
-from ...util.gpu_init import initGPU
 from ...util.config import *
 
 class RBF(Stationary):
@@ -25,7 +24,6 @@ class RBF(Stationary):
         self.group_spike_prob = False
         self.psicomp = PSICOMP_RBF()
         if self.useGPU:
-            initGPU()
             self.psicomp = PSICOMP_RBF_GPU()
         else:
             self.psicomp = PSICOMP_RBF()

GPy/util/gpu_init.py

@@ -5,37 +5,43 @@ Global variables: initSuccess
 providing CUBLAS handle: cublas_handle
 """
 
+gpu_initialized = False
+gpu_device = None
+gpu_context = None
+MPI_enabled = False
+
+try:
+    from mpi4py import MPI
+    MPI_enabled = True
+except:
+    pass
+
+try:
+    if MPI_enabled and MPI.COMM_WORLD.size>1:
+        from .parallel import get_id_within_node
+        gpuid = get_id_within_node()
+        import pycuda.driver
+        pycuda.driver.init()
+        if gpuid>=pycuda.driver.Device.count():
+            print '['+MPI.Get_processor_name()+'] more processes than the GPU numbers!'
+            MPI.COMM_WORLD.Abort()
+            raise
+        gpu_device = pycuda.driver.Device(gpuid)
+        gpu_context = gpu_device.make_context()
+        gpu_initialized = True
+    else:
+        import pycuda.autoinit
+        gpu_initialized = True
+except:
+    pass
+
 try:
     from scikits.cuda import cublas
     import scikits.cuda.linalg as culinalg
     culinalg.init()
     cublas_handle = cublas.cublasCreate()
 except:
-
-gpu_initialized = False
-gpu_device = None
-gpu_context = None
-    
-def initGPU(gpuid=None):
-    if gpu_initialized:
-        return
-    if gpuid==None:
-        try:
-            import pycuda.autoinit
-            gpu_initialized = True
-        except:
-            pass
-    else:
-        try:
-            import pycuda.driver
-            pycuda.driver.init()
-            if gpuid>=pycuda.driver.Device.count():
-                return
-            gpu_device = pycuda.driver.Device(gpuid)
-            gpu_context = gpu_device.make_context()
-            gpu_initialized = True
-        except:
-            pass
+    pass
 
 def closeGPU():
     if gpu_context is not None:

GPy/util/parallel.py

@@ -0,0 +1,14 @@
+"""
+The module of tools for parallelization (MPI)
+"""
+
+try:
+    from mpi4py import MPI
+except:
+    pass
+
+def get_id_within_node(comm=MPI.COMM_WORLD):
+    rank = comm.rank
+    nodename =  MPI.Get_processor_name()
+    nodelist = comm.allgather(nodename)
+    return len([i for i in nodelist[:rank] if i==nodename])