diff options
Diffstat (limited to 'gpu.py')
| -rw-r--r-- | gpu.py | 32 |
1 files changed, 18 insertions, 14 deletions
@@ -7,6 +7,7 @@ import pycuda.driver as cuda from pycuda import gpuarray from copy import copy from itertools import izip +from tools import timeit import src from geometry import standard_wavelengths @@ -42,18 +43,16 @@ class GPU(object): '''Initialize a GPU context on the specified device. If device_id==None, the default device is used.''' - if device_id == None: + if device_id is None: self.context = make_default_context() else: - device = pycuda.driver.Device(device_id) + print 'device_id = %s' % device_id + device = cuda.Device(device_id) self.context = device.make_context() print 'device %s' % self.context.get_device().name() self.module = SourceModule(src.kernel, options=['-I' + src.dir], no_extern_c=True, cache_dir=False) - self.geo_funcs = CUDAFuncs(self.module, ['set_wavelength_range', 'set_material', - 'set_surface', - 'set_global_mesh_variables']) - - + self.geo_funcs = CUDAFuncs(self.module, ['set_wavelength_range', 'set_material', 'set_surface', 'set_global_mesh_variables', 'color_solids']) + self.prop_funcs = CUDAFuncs(self.module, ['init_rng', 'propagate']) self.nblocks = 64 self.max_nthreads = 200000 @@ -154,16 +153,21 @@ class GPU(object): self.node_map_tex.set_format(cuda.array_format.UNSIGNED_INT32, 1) self.node_length_tex.set_format(cuda.array_format.UNSIGNED_INT32, 1) + self.geometry = geometry + self.print_device_usage() - def color_solids(self, solid_ids, colors): - solid_id_map = self.solid_id_map_gpu.get() - triangle_colors = self.colors_gpu.get() + @timeit + def reset_colors(self): + self.colors_gpu.set(self.geometry.colors.astype(np.uint32)) - for i, color in izip(solid_ids, colors): - triangle_colors[solid_id_map == i] = color + @timeit + def color_solids(self, solid_ids, colors): + solid_ids_gpu = gpuarray.to_gpu(np.array(solid_ids, dtype=np.int32)) + solid_colors_gpu = gpuarray.to_gpu(np.array(colors, dtype=np.uint32)) - self.colors_gpu.set(triangle_colors) + self.geo_funcs.color_solids(np.int32(solid_ids_gpu.size), np.uint32(self.triangles_gpu.size), self.solid_id_map_gpu, solid_ids_gpu, solid_colors_gpu, block=(self.nblocks,1,1), grid=(solid_ids_gpu.size//self.nblocks+1,1)) + self.context.synchronize() def setup_propagate(self, seed=1): self.rng_states_gpu = cuda.mem_alloc(self.max_nthreads*sizeof('curandStateXORWOW', '#include <curand_kernel.h>')) @@ -277,5 +281,5 @@ class GPU(object): def get_hits(self): return self.earliest_time_gpu.get() - def shutdown(self): + def __del__(self): self.context.pop() |