diff options
Diffstat (limited to 'gputhread.py')
| -rw-r--r-- | gputhread.py | 7 |
1 files changed, 4 insertions, 3 deletions
diff --git a/gputhread.py b/gputhread.py index 26520c4..12ae145 100644 --- a/gputhread.py +++ b/gputhread.py @@ -10,13 +10,14 @@ class Job(object): self.origins, self.directions = origins, directions class GPUThread(threading.Thread): - def __init__(self, device_id, geometry, jobs, output): + def __init__(self, device_id, geometry, jobs, output, nblocks=64): threading.Thread.__init__(self) self.device_id = device_id self.geometry = geometry self.jobs = jobs self.output = output + self.nblocks = nblocks self._stop = threading.Event() def stop(self): @@ -35,7 +36,7 @@ class GPUThread(threading.Thread): while not self.stopped(): try: - job = self.jobs.get(timeout=2) + job = self.jobs.get() except Queue.Empty: continue @@ -45,7 +46,7 @@ class GPUThread(threading.Thread): dest = np.empty(job.origins.size, dtype=np.int32) dest_gpu = cuda.to_device(dest) - propagate(np.int32(job.origins.size), origins_gpu, directions_gpu, np.int32(self.geometry.node_map.size-1), np.int32(self.geometry.first_node), dest_gpu, block=(64,1,1), grid=(job.origins.size//64+1,1), texrefs=texrefs) + propagate(np.int32(job.origins.size), origins_gpu, directions_gpu, np.int32(self.geometry.node_map.size-1), np.int32(self.geometry.first_node), dest_gpu, block=(self.nblocks,1,1), grid=(job.origins.size//self.nblocks+1,1), texrefs=texrefs) cuda.Context.synchronize() cuda.memcpy_dtoh(dest, dest_gpu) |