opencl optimization

src/orisys/HHD.py

@@ -1594,7 +1594,8 @@ class PoissonSolver(object):
         return self._opencl_crop_batch_kernel
 
     def _get_opencl_batch_mul_kernel(self, queue):
-        """Batched complex multiply: multiplies each batch plane by the same Green's spectrum."""
+        """Batched complex multiply: multiplies each batch plane by the same Green's spectrum.
+        Uses multiple elements per work item and explicit local size for better throughput."""
         if self._opencl_batch_mul_kernel is not None:
             return self._opencl_batch_mul_kernel
 
@@ -1604,14 +1605,14 @@ class PoissonSolver(object):
             __global const float2* g,
             const int n
         ) {
-            int i = get_global_id(0);
+            const int b = get_global_id(0);
-            int b = get_global_id(1);
+            const int i0 = get_global_id(1) * 4;
-            if (i < n) {
+            for (int i = i0; i < i0 + 4 && i < n; i++) {
                 float2 gv = g[i];
                 int idx = b * n + i;
                 float2 a = fft_batch[idx];
                 fft_batch[idx] = (float2)(a.x * gv.x - a.y * gv.y,
-                                           a.x * gv.y + a.y * gv.x);
+                                          a.x * gv.y + a.y * gv.x);
             }
         }
         """
@@ -1619,6 +1620,16 @@ class PoissonSolver(object):
         self._opencl_batch_mul_kernel = program.batch_complex_mul
         return self._opencl_batch_mul_kernel
 
+    def _get_opencl_alloc_host_ptr(self):
+        """Return an allocator that creates buffers with ALLOC_HOST_PTR for pinned host access (reduces H2D/D2H transfer cost)."""
+        ctx = self._opencl_context
+        mf = cl.mem_flags
+
+        def allocator(nbytes):
+            return cl.Buffer(ctx, mf.READ_WRITE | mf.ALLOC_HOST_PTR, size=nbytes)
+
+        return allocator
+
     @staticmethod
     def _center_crop_same(full_conv, target_shape):
         """Crop full convolution result to scipy 'same' output shape."""
@@ -1664,14 +1675,13 @@ class PoissonSolver(object):
         timing['queue_ms'] = (timeit.default_timer() - t0) * 1000.0
 
         t0 = timeit.default_timer()
-        f1_host = numpy.ascontiguousarray(f1, dtype=numpy.float32)
+        fshape = f1.shape
-        f2_host = numpy.ascontiguousarray(f2, dtype=numpy.float32)
+        full_shape = tuple(2 * s - 1 for s in fshape)
-        full_shape = tuple(2 * s - 1 for s in f1_host.shape)
         timing['host_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
 
         cache_valid = (
             self._opencl_cache is not None and
-            self._opencl_cache['fshape'] == f1_host.shape and
+            self._opencl_cache['fshape'] == fshape and
             self._opencl_cache['full_shape'] == full_shape and
             'f_batch_dev' in self._opencl_cache
         )
@@ -1689,16 +1699,19 @@ class PoissonSolver(object):
 
             # For circular conv at size (2s-1) with G of shape (2s-1) and f of shape s,
             # alias-free positions are [s-1, 2s-2], matching scipy 'same' extraction.
-            same_starts = [f1_host.shape[i] - 1 for i in range(self.dim)]
+            same_starts = [fshape[i] - 1 for i in range(self.dim)]
             fft_shape = tuple(list(full_shape[:-1]) + [full_shape[-1] // 2 + 1])
             fft_n_elements = 1
             for d in fft_shape:
                 fft_n_elements *= d
-            src_plane = f1_host.shape[0] * f1_host.shape[1]
+            src_plane = fshape[0] * fshape[1]
             dst_plane = full_shape[0] * full_shape[1]
 
+            # Pinned allocator for staging (H2D) and output (D2H) buffers
+            alloc_pinned = self._get_opencl_alloc_host_ptr()
+
             self._opencl_cache = {
-                'fshape': f1_host.shape,
+                'fshape': fshape,
                 'full_shape': full_shape,
                 'fft_shape': fft_shape,
                 'fft_n_elements': fft_n_elements,
@@ -1707,14 +1720,11 @@ class PoissonSolver(object):
                 'gdx_prod': numpy.prod(self.gdx),
                 'src_plane': src_plane,
                 'dst_plane': dst_plane,
-                # Stacked (batch=2) buffers for true batched FFT
                 'f_batch_dev': cla.zeros(queue, (2,) + full_shape, dtype=numpy.float32),
-                'f_stage_batch_dev': cla.empty(queue, (2,) + f1_host.shape, dtype=numpy.float32),
+                'f_stage_batch_dev': cla.empty(queue, (2,) + fshape, dtype=numpy.float32, allocator=alloc_pinned),
                 'f_fft_batch_dev': cla.empty(queue, (2,) + fft_shape, dtype=numpy.complex64),
                 'conv_batch_dev': cla.empty(queue, (2,) + full_shape, dtype=numpy.float32),
-                'p_batch_dev': cla.empty(queue, (2,) + f1_host.shape, dtype=numpy.float32),
+                'p_batch_dev': cla.empty(queue, (2,) + fshape, dtype=numpy.float32, allocator=alloc_pinned),
-                'p_batch_host': numpy.empty((2,) + f1_host.shape, dtype=numpy.float32),
-                'f_stage_batch_host': numpy.empty((2,) + f1_host.shape, dtype=numpy.float32),
             }
             timing['cache_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
 
@@ -1722,12 +1732,12 @@ class PoissonSolver(object):
         fshape_c = c['fshape']
         full_shape_c = c['full_shape']
 
-        # H2D: single transfer for both inputs, then batch scatter into padded buffer
+        # H2D: map pinned buffer, write f1/f2 directly (no f1_host/f2_host), unmap, then scatter
         t0 = timeit.default_timer()
-        f_stage_host = c['f_stage_batch_host']
+        mapped_stage = c['f_stage_batch_dev'].map_to_host(queue=queue, flags=cl.map_flags.WRITE, is_blocking=True)
-        f_stage_host[0] = f1_host
+        mapped_stage[0] = numpy.ascontiguousarray(f1, dtype=numpy.float32)
-        f_stage_host[1] = f2_host
+        mapped_stage[1] = numpy.ascontiguousarray(f2, dtype=numpy.float32)
-        c['f_stage_batch_dev'].set(f_stage_host, queue=queue)
+        unmap_evt = mapped_stage.base.release(queue=queue)
         scatter_batch = self._get_opencl_scatter_batch_kernel(queue)
         scatter_batch(
             queue, (fshape_c[1], fshape_c[0], 2), None,
@@ -1737,7 +1747,8 @@ class PoissonSolver(object):
             numpy.int32(fshape_c[0]),
             numpy.int32(full_shape_c[1]),
             numpy.int32(c['src_plane']),
-            numpy.int32(c['dst_plane'])
+            numpy.int32(c['dst_plane']),
+            wait_for=[unmap_evt]
         )
         if profile:
             queue.finish()
@@ -1750,14 +1761,18 @@ class PoissonSolver(object):
             queue.finish()
         timing['fft_f_ms'] = (timeit.default_timer() - t0) * 1000.0
 
-        # Batched multiply: each batch plane * same Green's spectrum
+        # Batched multiply: each batch plane * same Green's spectrum (4 elements/work item, explicit local size)
         t0 = timeit.default_timer()
+        n_fft = c['fft_n_elements']
+        global_n = max(1, (n_fft + 3) // 4)
+        local_n = min(256, global_n)
+        global_n = ((global_n + local_n - 1) // local_n) * local_n
         batch_mul = self._get_opencl_batch_mul_kernel(queue)
         batch_mul(
-            queue, (c['fft_n_elements'], 2), None,
+            queue, (2, global_n), (1, local_n),
             fft_batch.data,
             c['g_fft_dev'].data,
-            numpy.int32(c['fft_n_elements'])
+            numpy.int32(n_fft)
         )
         if profile:
             queue.finish()
@@ -1790,11 +1805,12 @@ class PoissonSolver(object):
             queue.finish()
         timing['crop_scale_ms'] = (timeit.default_timer() - t0) * 1000.0
 
-        # D2H: single transfer for both results
+        # D2H: map pinned buffer (ALLOC_HOST_PTR) for read, copy out p1/p2, unmap
         t0 = timeit.default_timer()
-        cl.enqueue_copy(queue, c['p_batch_host'], c['p_batch_dev'].data, is_blocking=True)
+        mapped_p = c['p_batch_dev'].map_to_host(queue=queue, flags=cl.map_flags.READ, is_blocking=True)
-        p1 = c['p_batch_host'][0]
+        p1 = numpy.asarray(mapped_p[0], dtype=numpy.float32).copy()
-        p2 = c['p_batch_host'][1]
+        p2 = numpy.asarray(mapped_p[1], dtype=numpy.float32).copy()
+        mapped_p.base.release(queue=queue)
         timing['d2h_get_ms'] = (timeit.default_timer() - t0) * 1000.0
 
         timing['total_ms'] = (timeit.default_timer() - t_total) * 1000.0