diff --git a/src/orisys/HHD.py b/src/orisys/HHD.py index fbc1bab..10b2ab4 100644 --- a/src/orisys/HHD.py +++ b/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) - nRu_cpu = self.psolver.solve_opencl(curlw_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=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': numpy.empty(f_host.shape, dtype=numpy.float32), + 'p_host': p_host, } + # 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,8 +1762,11 @@ 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) - p = p_dev.get(ary=self._opencl_cache['p_host']) + 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: # Fallback for 3D: transfer full, crop on CPU @@ -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 - # (out-of-place rfftn does not modify the source) - 'f_dev': cla.zeros(queue, full_shape, dtype=numpy.float32), - 'f_dev2': cla.zeros(queue, full_shape, dtype=numpy.float32), - # Small staging buffers for input-sized H2D (~4x smaller) - 'f_stage_dev': cla.empty(queue, f1_host.shape, dtype=numpy.float32), - 'f_stage_dev2': cla.empty(queue, f1_host.shape, dtype=numpy.float32), - 'f_fft_dev': cla.empty(queue, fft_shape, dtype=numpy.complex64), - 'f_fft_dev2': cla.empty(queue, fft_shape, dtype=numpy.complex64), - '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), + '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_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_host': p_batch_host, + 'f_stage_batch_host': numpy.empty((2,) + 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) - self._opencl_cache['f_stage_dev2'].set(f2_host, queue=queue, async_=True) - scatter_kern = self._get_opencl_scatter_kernel(queue) - fshape_c = self._opencl_cache['fshape'] - full_shape_c = self._opencl_cache['full_shape'] - scatter_kern( - queue, (fshape_c[1], fshape_c[0]), None, - self._opencl_cache['f_stage_dev'].data, - self._opencl_cache['f_dev'].data, + f_stage_host = c['f_stage_batch_host'] + f_stage_host[0] = f1_host + f_stage_host[1] = f2_host + c['f_stage_batch_dev'].set(f_stage_host, queue=queue) + scatter_batch = self._get_opencl_scatter_batch_kernel(queue) + scatter_batch( + queue, (fshape_c[1], fshape_c[0], 2), None, + c['f_stage_batch_dev'].data, + c['f_batch_dev'].data, numpy.int32(fshape_c[1]), numpy.int32(fshape_c[0]), - numpy.int32(full_shape_c[1]) - ) - scatter_kern( - 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]) + numpy.int32(full_shape_c[1]), + numpy.int32(c['src_plane']), + numpy.int32(c['dst_plane']) ) 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']) - fft2 = cl_rfftn(self._opencl_cache['f_dev2'], self._opencl_cache['f_fft_dev2']) + fft_batch = cl_rfftn(c['f_batch_dev'], c['f_fft_batch_dev'], ndim=2) 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) - fused_mul( - queue, (self._opencl_cache['fft_n_elements'],), None, - fft1.data, fft2.data, - self._opencl_cache['g_fft_dev'].data, - numpy.int32(self._opencl_cache['fft_n_elements']) + batch_mul = self._get_opencl_batch_mul_kernel(queue) + batch_mul( + queue, (c['fft_n_elements'], 2), None, + fft_batch.data, + c['g_fft_dev'].data, + 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']) - full2 = cl_irfftn(fft2, self._opencl_cache['conv_full_dev2']) + conv_batch = cl_irfftn(fft_batch, c['conv_batch_dev'], ndim=2) 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) - same_starts = self._opencl_cache['same_starts'] - - crop_kernel( - queue, (fshape_c[1], fshape_c[0]), None, - full1.data, self._opencl_cache['p_dev'].data, + crop_batch = self._get_opencl_crop_batch_kernel(queue) + same_starts = c['same_starts'] + crop_batch( + queue, (fshape_c[1], fshape_c[0], 2), None, + conv_batch.data, c['p_batch_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']) - ) - crop_kernel( - 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']) + numpy.float32(c['gdx_prod']), + numpy.int32(c['dst_plane']), + numpy.int32(c['src_plane']) ) 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'], - 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'] - p2 = self._opencl_cache['p_host2'] + if c.get('_use_host_ptr'): + queue.finish() + else: + 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