Adapt work group size based on available lmem

inducer · Jul 15, 2015 · 2b98b07 · 2b98b07
1 parent 2181288
commit 2b98b07
Showing 1 changed file with 25 additions and 1 deletion.
diff --git a/pyopencl/bitonic_sort.py b/pyopencl/bitonic_sort.py
@@ -176,7 +176,31 @@ def sort_b_prepare_wl(self, argsort, key_dtype, idx_dtype, shape, axis):
         allowb8 = True
         allowb16 = True
 
-        wg = min(ds, self.context.devices[0].max_work_group_size)
+        dev = self.context.devices[0]
+
+        # {{{ find workgroup size
+
+        wg = min(ds, dev.max_work_group_size)
+
+        available_lmem = dev.local_mem_size
+        while True:
+            lmem_size = wg*4*key_dtype.itemsize
+            if argsort:
+                lmem_size += wg*4*idx_dtype.itemsize
+
+            if lmem_size + 512 > available_lmem:
+                wg //= 2
+
+                if not wg:
+                    raise RuntimeError(
+                        "too little local memory available on '%s'"
+                        % dev)
+
+            else:
+                break
+
+        # }}}
+
         length = wg >> 1
         prg = self.get_program(
                 'BLO', argsort, (1, 1, key_ctype, idx_ctype, ds, ns))