apunkt/GPy
1
0
Fork 0
mirror of https://github.com/SheffieldML/GPy.git synced 2026-06-29 15:59:41 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix the gpu initialization for multiple cards

Browse source
This commit is contained in:
Zhenwen Dai 2014-06-26 13:46:47 +01:00
parent bdf478956d
commit c3482e7a94
5 changed files with 45 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

56
GPy/util/gpu_init.py
View file

@ -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: