you can now specify block size for a gputhread

2 files changed, 6 insertions, 4 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)
diff --git a/threadtest.py b/threadtest.py
index 11bdaa4..fe90b4e 100644
--- a/threadtest.py
+++ b/threadtest.py
@@ -65,6 +65,7 @@ if __name__ == '__main__':
     parser = optparse.OptionParser('%prog')
     parser.add_option('-b', type='int', dest='nbits', default=8)
     parser.add_option('-j', type='int', dest='ndevices', default=1)
+    parser.add_option('-n', type='int', dest='nblocks', default=64)
     options, args = parser.parse_args()
 
     lbne = LBNE()
@@ -74,7 +75,7 @@ if __name__ == '__main__':
 
     gputhreads = []
     for i in range(options.ndevices):
-        gputhreads.append(GPUThread(i, lbne, jobs, output))
+        gputhreads.append(GPUThread(i, lbne, jobs, output, options.nblocks))
         gputhreads[-1].start()
 
     try: