adapt gpu initialization multiple gpu cards

This commit is contained in:
Zhenwen Dai 2014-06-26 12:17:45 +01:00
parent cf33808673
commit bdf478956d
2 changed files with 27 additions and 5 deletions

View file

@ -6,6 +6,7 @@ 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):
@ -24,6 +25,7 @@ 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()

View file

@ -6,17 +6,37 @@ providing CUBLAS handle: cublas_handle
"""
try:
import pycuda.autoinit
from scikits.cuda import cublas
import scikits.cuda.linalg as culinalg
culinalg.init()
cublas_handle = cublas.cublasCreate()
initSuccess = True
except:
initSuccess = False
gpu_initialized = False
gpu_device = None
gpu_context = None
def initGPU(gpuid=None):
if gpu_initialized:
return
if gpuid==None:
return pycuda.tools.make_default_context()
try:
import pycuda.autoinit
gpu_initialized = True
except:
pass
else:
return pycuda.driver.Device(gpuid).make_context()
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
def closeGPU():
if gpu_context is not None:
gpu_context.detach()