optimized pyopencl HHD on radxa

src/orisys/HHD.py

@@ -27,6 +27,7 @@ import numpy
 import timeit
 import sys
 from scipy import signal, spatial
+from scipy.fft import next_fast_len
 # from .UnstructuredGrid import UnstructuredGrid
 # import matplotlib.pyplot as plt
 
@@ -125,7 +126,7 @@ class nHHD(object):
         """
         if self.poisson_backend == 'opencl':
             # Always enable profiling for OpenCL backend to see performance breakdown
-            return self.psolver.solve_opencl(field, verbose, profile=False)
+            return self.psolver.solve_opencl(field, verbose, profile=True)
 
         if self.poisson_backend == 'cupy':
             return self.psolver.solve_gpu(field, verbose)
@@ -134,7 +135,7 @@ class nHHD(object):
         if 'OPENCL_AVAILABLE' in globals() and OPENCL_AVAILABLE:
             try:
                 # Always enable profiling for OpenCL backend to see performance breakdown
-                return self.psolver.solve_opencl(field, verbose, profile=False)
+                return self.psolver.solve_opencl(field, verbose, profile=True)
             except Exception as exc:
                 if verbose:
                     print(f"OpenCL solve failed ({exc}), falling back to CuPy/CPU")
@@ -169,7 +170,7 @@ class nHHD(object):
             # compute potentials
             if self.poisson_backend == 'opencl' and ('OPENCL_AVAILABLE' in globals() and OPENCL_AVAILABLE):
                 try:
-                    self.nD, self.nRu = self.psolver.solve_opencl_batch2(self.div, self.curlw, verbose > 0, profile=False)
+                    self.nD, self.nRu = self.psolver.solve_opencl_batch2(self.div, self.curlw, verbose > 0, profile=True)
                 except Exception as exc:
                     if verbose > 0:
                         print(f"OpenCL batch solve failed ({exc}), falling back to single solves")
@@ -299,8 +300,8 @@ class nHHD(object):
                 try:
                     div_cpu = cp.asnumpy(div_gpu)
                     curlw_cpu = cp.asnumpy(curlw_gpu)
-                    nD_cpu = self.psolver.solve_opencl(div_cpu, verbose > 0, profile=False)
+                    nD_cpu = self.psolver.solve_opencl(div_cpu, verbose > 0, profile=True)
-                    nRu_cpu = self.psolver.solve_opencl(curlw_cpu, verbose > 0, profile=False)
+                    nRu_cpu = self.psolver.solve_opencl(curlw_cpu, verbose > 0, profile=True)
                     nD_gpu = cp.asarray(nD_cpu, dtype=cp.float32)
                     nRu_gpu = cp.asarray(nRu_cpu, dtype=cp.float32)
                 except Exception as exc:
@@ -1097,6 +1098,9 @@ class PoissonSolver(object):
         self._opencl_crop_kernel = None  # Cached OpenCL crop kernel
         self._opencl_scatter_kernel = None  # Cached scatter (reverse-crop) kernel
         self._opencl_fused_mul_kernel = None  # Cached fused complex multiply kernel
+        self._opencl_scatter_batch_kernel = None  # Cached batch scatter kernel
+        self._opencl_crop_batch_kernel = None  # Cached batch crop kernel
+        self._opencl_batch_mul_kernel = None  # Cached batch complex multiply kernel
         self._opencl_last_timing = None  # Last per-stage OpenCL timing in ms
         # if (self.ptype == 'G'):
         #     print( 'PoissonSolver:', self.dim, 'D grid =', self.gdims, 'with spacings', self.gdx)
@@ -1503,6 +1507,93 @@ class PoissonSolver(object):
         self._opencl_fused_mul_kernel = program.fused_complex_mul
         return self._opencl_fused_mul_kernel
 
+    def _get_opencl_scatter_batch_kernel(self, queue):
+        """Batched scatter: copies two unpadded inputs into a stacked zero-padded device buffer."""
+        if self._opencl_scatter_batch_kernel is not None:
+            return self._opencl_scatter_batch_kernel
+
+        kernel_code = """
+        __kernel void scatter_batch(
+            __global const float* src,
+            __global float* dst,
+            const int src_width,
+            const int src_height,
+            const int dst_width,
+            const int src_plane,
+            const int dst_plane
+        ) {
+            int gid_x = get_global_id(0);
+            int gid_y = get_global_id(1);
+            int gid_z = get_global_id(2);
+            if (gid_x < src_width && gid_y < src_height) {
+                int src_idx = gid_z * src_plane + gid_y * src_width + gid_x;
+                int dst_idx = gid_z * dst_plane + gid_y * dst_width + gid_x;
+                dst[dst_idx] = src[src_idx];
+            }
+        }
+        """
+        program = cl.Program(self._opencl_context, kernel_code).build()
+        self._opencl_scatter_batch_kernel = program.scatter_batch
+        return self._opencl_scatter_batch_kernel
+
+    def _get_opencl_crop_batch_kernel(self, queue):
+        """Batched crop+scale: extracts 'same' region from two stacked full-conv results."""
+        if self._opencl_crop_batch_kernel is not None:
+            return self._opencl_crop_batch_kernel
+
+        kernel_code = """
+        __kernel void crop_batch(
+            __global const float* src,
+            __global float* dst,
+            const int src_width,
+            const int dst_width,
+            const int dst_height,
+            const int start_x,
+            const int start_y,
+            const float scale,
+            const int src_plane,
+            const int dst_plane
+        ) {
+            int gid_x = get_global_id(0);
+            int gid_y = get_global_id(1);
+            int gid_z = get_global_id(2);
+            if (gid_x < dst_width && gid_y < dst_height) {
+                int src_idx = gid_z * src_plane + (start_y + gid_y) * src_width + (start_x + gid_x);
+                int dst_idx = gid_z * dst_plane + gid_y * dst_width + gid_x;
+                dst[dst_idx] = src[src_idx] * scale;
+            }
+        }
+        """
+        program = cl.Program(self._opencl_context, kernel_code).build()
+        self._opencl_crop_batch_kernel = program.crop_batch
+        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."""
+        if self._opencl_batch_mul_kernel is not None:
+            return self._opencl_batch_mul_kernel
+
+        kernel_code = """
+        __kernel void batch_complex_mul(
+            __global float2* fft_batch,
+            __global const float2* g,
+            const int n
+        ) {
+            int i = get_global_id(0);
+            int b = get_global_id(1);
+            if (i < n) {
+                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);
+            }
+        }
+        """
+        program = cl.Program(self._opencl_context, kernel_code).build()
+        self._opencl_batch_mul_kernel = program.batch_complex_mul
+        return self._opencl_batch_mul_kernel
+
     @staticmethod
     def _center_crop_same(full_conv, target_shape):
         """Crop full convolution result to scipy 'same' output shape."""
@@ -1510,7 +1601,7 @@ class PoissonSolver(object):
         slices = tuple(slice(starts[i], starts[i] + target_shape[i]) for i in range(len(target_shape)))
         return full_conv[slices]
 
-    def solve_opencl(self, f, verbose=False, profile=False):
+    def solve_opencl(self, f, verbose=False, profile=True):
         """
         OpenCL Poisson solver using pyvkfft FFT convolution.
         Replaces scipy.signal.fftconvolve(..., mode='same') for structured grids.
@@ -1544,13 +1635,14 @@ class PoissonSolver(object):
 
                 t0 = timeit.default_timer()
                 f_host = numpy.ascontiguousarray(f, dtype=numpy.float32)
-                full_shape = tuple(2 * s - 1 for s in f_host.shape)
+                full_shape = tuple(next_fast_len(2 * s - 1) for s in f_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'] == f_host.shape and
-                    self._opencl_cache['full_shape'] == full_shape
+                    self._opencl_cache['full_shape'] == full_shape and
+                    'f_dev' in self._opencl_cache
                 )
 
                 if not cache_valid:
@@ -1569,6 +1661,7 @@ class PoissonSolver(object):
                         for i in range(self.dim)
                     )
                     fft_shape = tuple(list(full_shape[:-1]) + [full_shape[-1] // 2 + 1])
+                    p_host = numpy.empty(f_host.shape, dtype=numpy.float32)
                     self._opencl_cache = {
                         'fshape': f_host.shape,
                         'full_shape': full_shape,
@@ -1583,9 +1676,21 @@ class PoissonSolver(object):
                         'f_stage_dev': cla.empty(queue, f_host.shape, dtype=numpy.float32),
                         'f_fft_dev': cla.empty(queue, fft_shape, dtype=numpy.complex64),
                         'conv_full_dev': cla.empty(queue, full_shape, dtype=numpy.float32),
-                        'p_dev': cla.empty(queue, f_host.shape, dtype=numpy.float32),
+                        'p_host': p_host,
-                        'p_host': numpy.empty(f_host.shape, dtype=numpy.float32),
                     }
+                    # USE_HOST_PTR: crop kernel writes directly into host numpy array
+                    try:
+                        p_buf = cl.Buffer(queue.context,
+                                          cl.mem_flags.READ_WRITE | cl.mem_flags.USE_HOST_PTR,
+                                          hostbuf=p_host)
+                        self._opencl_cache['p_dev'] = cla.Array(
+                            queue, f_host.shape, dtype=numpy.float32, data=p_buf)
+                        self._opencl_cache['_p_buf'] = p_buf
+                        self._opencl_cache['_use_host_ptr'] = True
+                    except Exception:
+                        self._opencl_cache['p_dev'] = cla.empty(
+                            queue, f_host.shape, dtype=numpy.float32)
+                        self._opencl_cache['_use_host_ptr'] = False
                     timing['cache_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
 
                 # H2D: upload only input-sized data, scatter into padded device buffer
@@ -1657,7 +1762,10 @@ class PoissonSolver(object):
                     timing['crop_scale_ms'] = (timeit.default_timer() - t0) * 1000.0
 
                     t0 = timeit.default_timer()
-                    # Transfer only the cropped result (much smaller than full convolution)
+                    if self._opencl_cache.get('_use_host_ptr'):
+                        queue.finish()
+                        p = self._opencl_cache['p_host']
+                    else:
                         p = p_dev.get(ary=self._opencl_cache['p_host'])
                     timing['d2h_get_ms'] = (timeit.default_timer() - t0) * 1000.0
                 else:
@@ -1703,7 +1811,7 @@ class PoissonSolver(object):
                 print('OpenCL solver not optimized for unstructured points, using CPU')
             return self.solve(f, verbose)
 
-    def solve_opencl_batch2(self, f1, f2, verbose=False, profile=False):
+    def solve_opencl_batch2(self, f1, f2, verbose=False, profile=True):
         """
         Batched OpenCL Poisson solver for two 2D fields.
         Optimized: scatter H2D (~4x smaller transfer), fused multiply
@@ -1741,14 +1849,14 @@ class PoissonSolver(object):
         t0 = timeit.default_timer()
         f1_host = numpy.ascontiguousarray(f1, dtype=numpy.float32)
         f2_host = numpy.ascontiguousarray(f2, dtype=numpy.float32)
-        full_shape = tuple(2 * s - 1 for s in f1_host.shape)
+        full_shape = tuple(next_fast_len(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['full_shape'] == full_shape and
-            'f_dev2' in self._opencl_cache
+            'f_batch_dev' in self._opencl_cache
         )
 
         if not cache_valid:
@@ -1763,11 +1871,15 @@ class PoissonSolver(object):
                 queue.finish()
 
             same_starts = [(full_shape[i] - f1_host.shape[i]) // 2 for i in range(self.dim)]
-            same_slices = tuple(slice(same_starts[i], same_starts[i] + f1_host.shape[i]) 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]
+            dst_plane = full_shape[0] * full_shape[1]
+
+            p_shape = (2,) + f1_host.shape
+            p_batch_host = numpy.empty(p_shape, dtype=numpy.float32)
 
             self._opencl_cache = {
                 'fshape': f1_host.shape,
@@ -1775,121 +1887,114 @@ class PoissonSolver(object):
                 'fft_shape': fft_shape,
                 'fft_n_elements': fft_n_elements,
                 'g_fft_dev': g_fft_dev,
-                'same_slices': same_slices,
                 'same_starts': same_starts,
                 'gdx_prod': numpy.prod(self.gdx),
-                # Zero-init padded buffers so padding stays zero across frames
+                'src_plane': src_plane,
-                # (out-of-place rfftn does not modify the source)
+                'dst_plane': dst_plane,
-                'f_dev': cla.zeros(queue, full_shape, dtype=numpy.float32),
+                # Stacked (batch=2) buffers for true batched FFT
-                'f_dev2': cla.zeros(queue, full_shape, dtype=numpy.float32),
+                'f_batch_dev': cla.zeros(queue, (2,) + full_shape, dtype=numpy.float32),
-                # Small staging buffers for input-sized H2D (~4x smaller)
+                'f_stage_batch_dev': cla.empty(queue, (2,) + f1_host.shape, dtype=numpy.float32),
-                'f_stage_dev': cla.empty(queue, f1_host.shape, dtype=numpy.float32),
+                'f_fft_batch_dev': cla.empty(queue, (2,) + fft_shape, dtype=numpy.complex64),
-                'f_stage_dev2': cla.empty(queue, f1_host.shape, dtype=numpy.float32),
+                'conv_batch_dev': cla.empty(queue, (2,) + full_shape, dtype=numpy.float32),
-                'f_fft_dev': cla.empty(queue, fft_shape, dtype=numpy.complex64),
+                'p_batch_host': p_batch_host,
-                'f_fft_dev2': cla.empty(queue, fft_shape, dtype=numpy.complex64),
+                'f_stage_batch_host': numpy.empty((2,) + f1_host.shape, dtype=numpy.float32),
-                'conv_full_dev': cla.empty(queue, full_shape, dtype=numpy.float32),
-                'conv_full_dev2': cla.empty(queue, full_shape, dtype=numpy.float32),
-                'p_dev': cla.empty(queue, f1_host.shape, dtype=numpy.float32),
-                'p_dev2': cla.empty(queue, f1_host.shape, dtype=numpy.float32),
-                'p_host': numpy.empty(f1_host.shape, dtype=numpy.float32),
-                'p_host2': numpy.empty(f1_host.shape, dtype=numpy.float32),
             }
+            # USE_HOST_PTR: crop kernel writes directly into host numpy array
+            # On unified memory (Qualcomm), this avoids the explicit D2H memcpy
+            try:
+                p_buf = cl.Buffer(queue.context,
+                                  cl.mem_flags.READ_WRITE | cl.mem_flags.USE_HOST_PTR,
+                                  hostbuf=p_batch_host)
+                self._opencl_cache['p_batch_dev'] = cla.Array(
+                    queue, p_shape, dtype=numpy.float32, data=p_buf)
+                self._opencl_cache['_p_buf'] = p_buf
+                self._opencl_cache['_use_host_ptr'] = True
+            except Exception:
+                self._opencl_cache['p_batch_dev'] = cla.empty(
+                    queue, p_shape, dtype=numpy.float32)
+                self._opencl_cache['_use_host_ptr'] = False
             timing['cache_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
 
-        # H2D: async upload input-sized data, scatter into padded device buffer
+        c = self._opencl_cache
+        fshape_c = c['fshape']
+        full_shape_c = c['full_shape']
+
+        # H2D: single transfer for both inputs, then batch scatter into padded buffer
         t0 = timeit.default_timer()
-        self._opencl_cache['f_stage_dev'].set(f1_host, queue=queue, async_=True)
+        f_stage_host = c['f_stage_batch_host']
-        self._opencl_cache['f_stage_dev2'].set(f2_host, queue=queue, async_=True)
+        f_stage_host[0] = f1_host
-        scatter_kern = self._get_opencl_scatter_kernel(queue)
+        f_stage_host[1] = f2_host
-        fshape_c = self._opencl_cache['fshape']
+        c['f_stage_batch_dev'].set(f_stage_host, queue=queue)
-        full_shape_c = self._opencl_cache['full_shape']
+        scatter_batch = self._get_opencl_scatter_batch_kernel(queue)
-        scatter_kern(
+        scatter_batch(
-            queue, (fshape_c[1], fshape_c[0]), None,
+            queue, (fshape_c[1], fshape_c[0], 2), None,
-            self._opencl_cache['f_stage_dev'].data,
+            c['f_stage_batch_dev'].data,
-            self._opencl_cache['f_dev'].data,
+            c['f_batch_dev'].data,
             numpy.int32(fshape_c[1]),
             numpy.int32(fshape_c[0]),
-            numpy.int32(full_shape_c[1])
+            numpy.int32(full_shape_c[1]),
-        )
+            numpy.int32(c['src_plane']),
-        scatter_kern(
+            numpy.int32(c['dst_plane'])
-            queue, (fshape_c[1], fshape_c[0]), None,
-            self._opencl_cache['f_stage_dev2'].data,
-            self._opencl_cache['f_dev2'].data,
-            numpy.int32(fshape_c[1]),
-            numpy.int32(fshape_c[0]),
-            numpy.int32(full_shape_c[1])
         )
         if profile:
             queue.finish()
         timing['h2d_set_ms'] = (timeit.default_timer() - t0) * 1000.0
 
-        # FFT both fields
+        # Single batched 2D FFT for both fields (ndim=2 treats dim-0 as batch)
         t0 = timeit.default_timer()
-        fft1 = cl_rfftn(self._opencl_cache['f_dev'], self._opencl_cache['f_fft_dev'])
+        fft_batch = cl_rfftn(c['f_batch_dev'], c['f_fft_batch_dev'], ndim=2)
-        fft2 = cl_rfftn(self._opencl_cache['f_dev2'], self._opencl_cache['f_fft_dev2'])
         if profile:
             queue.finish()
         timing['fft_f_ms'] = (timeit.default_timer() - t0) * 1000.0
 
-        # Fused multiply: one kernel, reads g_fft once, multiplies both spectra
+        # Batched multiply: each batch plane * same Green's spectrum
         t0 = timeit.default_timer()
-        fused_mul = self._get_opencl_fused_mul_kernel(queue)
+        batch_mul = self._get_opencl_batch_mul_kernel(queue)
-        fused_mul(
+        batch_mul(
-            queue, (self._opencl_cache['fft_n_elements'],), None,
+            queue, (c['fft_n_elements'], 2), None,
-            fft1.data, fft2.data,
+            fft_batch.data,
-            self._opencl_cache['g_fft_dev'].data,
+            c['g_fft_dev'].data,
-            numpy.int32(self._opencl_cache['fft_n_elements'])
+            numpy.int32(c['fft_n_elements'])
         )
         if profile:
             queue.finish()
         timing['mul_ms'] = (timeit.default_timer() - t0) * 1000.0
 
-        # iFFT both fields
+        # Single batched 2D iFFT for both fields
         t0 = timeit.default_timer()
-        full1 = cl_irfftn(fft1, self._opencl_cache['conv_full_dev'])
+        conv_batch = cl_irfftn(fft_batch, c['conv_batch_dev'], ndim=2)
-        full2 = cl_irfftn(fft2, self._opencl_cache['conv_full_dev2'])
         if profile:
             queue.finish()
         timing['ifft_ms'] = (timeit.default_timer() - t0) * 1000.0
 
-        # Crop + scale on device (both fields, back-to-back enqueue)
+        # Batch crop + scale on device
         t0 = timeit.default_timer()
-        crop_kernel = self._get_opencl_crop_kernel(queue)
+        crop_batch = self._get_opencl_crop_batch_kernel(queue)
-        same_starts = self._opencl_cache['same_starts']
+        same_starts = c['same_starts']
-
+        crop_batch(
-        crop_kernel(
+            queue, (fshape_c[1], fshape_c[0], 2), None,
-            queue, (fshape_c[1], fshape_c[0]), None,
+            conv_batch.data, c['p_batch_dev'].data,
-            full1.data, self._opencl_cache['p_dev'].data,
             numpy.int32(full_shape_c[1]),
             numpy.int32(fshape_c[1]),
             numpy.int32(fshape_c[0]),
             numpy.int32(same_starts[1]),
             numpy.int32(same_starts[0]),
-            numpy.float32(self._opencl_cache['gdx_prod'])
+            numpy.float32(c['gdx_prod']),
-        )
+            numpy.int32(c['dst_plane']),
-        crop_kernel(
+            numpy.int32(c['src_plane'])
-            queue, (fshape_c[1], fshape_c[0]), None,
-            full2.data, self._opencl_cache['p_dev2'].data,
-            numpy.int32(full_shape_c[1]),
-            numpy.int32(fshape_c[1]),
-            numpy.int32(fshape_c[0]),
-            numpy.int32(same_starts[1]),
-            numpy.int32(same_starts[0]),
-            numpy.float32(self._opencl_cache['gdx_prod'])
         )
         if profile:
             queue.finish()
         timing['crop_scale_ms'] = (timeit.default_timer() - t0) * 1000.0
 
-        # D2H: async enqueue both copies, single sync point
+        # D2H: zero-copy if USE_HOST_PTR, else explicit transfer
         t0 = timeit.default_timer()
-        cl.enqueue_copy(queue, self._opencl_cache['p_host'],
+        if c.get('_use_host_ptr'):
-                        self._opencl_cache['p_dev'].data, is_blocking=False)
-        cl.enqueue_copy(queue, self._opencl_cache['p_host2'],
-                        self._opencl_cache['p_dev2'].data, is_blocking=False)
             queue.finish()
-        p1 = self._opencl_cache['p_host']
+        else:
-        p2 = self._opencl_cache['p_host2']
+            cl.enqueue_copy(queue, c['p_batch_host'], c['p_batch_dev'].data, is_blocking=True)
+        p1 = c['p_batch_host'][0]
+        p2 = c['p_batch_host'][1]
         timing['d2h_get_ms'] = (timeit.default_timer() - t0) * 1000.0
 
         timing['total_ms'] = (timeit.default_timer() - t_total) * 1000.0