Allow to specify wg_size

1 files changed, 7 insertions, 1 deletions

diff --git a/src/solvers/opencl.rs b/src/solvers/opencl.rs
index 8421d53..cdedd37 100644
--- a/src/solvers/opencl.rs
+++ b/src/solvers/opencl.rs
@@ -63,6 +63,7 @@ impl GpuSolver {
         n: u32,
         h: u32,
         w: u32,
+        mut wg_size: usize,
         src: &str,
     ) -> ocl::Result<Vec<Sender<Job>>> {
         let platform = ocl::Platform::default();
@@ -86,7 +87,12 @@ impl GpuSolver {
 
         let mut senders = Vec::with_capacity(h as usize);
         let mut receivers = Vec::with_capacity(h as usize);
-        let wg_size = device.max_wg_size()?;
+        let max_wg_size = device.max_wg_size()?;
+        if wg_size == 0 {
+            wg_size = max_wg_size;
+        } else if wg_size > max_wg_size {
+            return Err(ocl::Error::from("invalid workgroup size"));
+        }
         for _ in 0..h {
             let (sx, rx) = std::sync::mpsc::channel();
             senders.push(sx);