fix gpu initialziation

This commit is contained in:
Zhenwen Dai 2016-02-09 17:22:38 +00:00
parent 0dd52981d0
commit a77a675549
3 changed files with 29 additions and 28 deletions

View file

@ -5,7 +5,6 @@ The module for psi-statistics for RBF kernel
import numpy as np import numpy as np
from paramz.caching import Cache_this from paramz.caching import Cache_this
from . import PSICOMP_RBF from . import PSICOMP_RBF
from ....util import gpu_init
gpu_code = """ gpu_code = """
// define THREADNUM // define THREADNUM
@ -238,8 +237,6 @@ class PSICOMP_RBF_GPU(PSICOMP_RBF):
self.fall_back = PSICOMP_RBF() self.fall_back = PSICOMP_RBF()
from pycuda.compiler import SourceModule from pycuda.compiler import SourceModule
from ....util.gpu_init import initGPU
initGPU()
self.GPU_direct = GPU_direct self.GPU_direct = GPU_direct
self.gpuCache = None self.gpuCache = None

View file

@ -287,8 +287,6 @@ class PSICOMP_SSRBF_GPU(PSICOMP_RBF):
def __init__(self, threadnum=128, blocknum=15, GPU_direct=False): def __init__(self, threadnum=128, blocknum=15, GPU_direct=False):
from pycuda.compiler import SourceModule from pycuda.compiler import SourceModule
from ....util.gpu_init import initGPU
initGPU()
self.GPU_direct = GPU_direct self.GPU_direct = GPU_direct
self.gpuCache = None self.gpuCache = None

View file

@ -10,29 +10,35 @@ gpu_device = None
gpu_context = None gpu_context = None
MPI_enabled = False MPI_enabled = False
def initGPU(): try:
try: import pycuda.autoinit
from mpi4py import MPI gpu_initialized = True
MPI_enabled = True except:
except: pass
pass
try: # def initGPU():
if MPI_enabled and MPI.COMM_WORLD.size>1: # try:
from .parallel import get_id_within_node # from mpi4py import MPI
gpuid = get_id_within_node() # MPI_enabled = True
import pycuda.driver # except:
pycuda.driver.init() # pass
if gpuid>=pycuda.driver.Device.count(): # try:
print('['+MPI.Get_processor_name()+'] more processes than the GPU numbers!') # if MPI_enabled and MPI.COMM_WORLD.size>1:
raise # from .parallel import get_id_within_node
gpu_device = pycuda.driver.Device(gpuid) # gpuid = get_id_within_node()
gpu_context = gpu_device.make_context() # import pycuda.driver
gpu_initialized = True # pycuda.driver.init()
else: # if gpuid>=pycuda.driver.Device.count():
import pycuda.autoinit # print('['+MPI.Get_processor_name()+'] more processes than the GPU numbers!')
gpu_initialized = True # raise
except: # gpu_device = pycuda.driver.Device(gpuid)
pass # gpu_context = gpu_device.make_context()
# gpu_initialized = True
# else:
# import pycuda.autoinit
# gpu_initialized = True
# except:
# pass
def closeGPU(): def closeGPU():
if gpu_context is not None: if gpu_context is not None: