From 78dde20838c4508fa7301feb6a0958a83f96220c Mon Sep 17 00:00:00 2001 From: radxa-d6a Date: Sun, 22 Feb 2026 12:57:53 +0800 Subject: [PATCH] opencl optimization --- src/orisys/HHD.py | 76 ++++++++++++++++++++++++++++------------------- 1 file changed, 46 insertions(+), 30 deletions(-) diff --git a/src/orisys/HHD.py b/src/orisys/HHD.py index 9b0eb3a..d3e00ab 100644 --- a/src/orisys/HHD.py +++ b/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); - int b = get_global_id(1); - if (i < n) { + const int b = get_global_id(0); + const int i0 = get_global_id(1) * 4; + 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) - f2_host = numpy.ascontiguousarray(f2, dtype=numpy.float32) - full_shape = tuple(2 * s - 1 for s in f1_host.shape) + fshape = f1.shape + full_shape = tuple(2 * s - 1 for s in fshape) 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_host': numpy.empty((2,) + f1_host.shape, dtype=numpy.float32), - 'f_stage_batch_host': numpy.empty((2,) + f1_host.shape, dtype=numpy.float32), + 'p_batch_dev': cla.empty(queue, (2,) + fshape, dtype=numpy.float32, allocator=alloc_pinned), } 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'] - f_stage_host[0] = f1_host - f_stage_host[1] = f2_host - c['f_stage_batch_dev'].set(f_stage_host, queue=queue) + mapped_stage = c['f_stage_batch_dev'].map_to_host(queue=queue, flags=cl.map_flags.WRITE, is_blocking=True) + mapped_stage[0] = numpy.ascontiguousarray(f1, dtype=numpy.float32) + mapped_stage[1] = numpy.ascontiguousarray(f2, dtype=numpy.float32) + 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) - p1 = c['p_batch_host'][0] - p2 = c['p_batch_host'][1] + mapped_p = c['p_batch_dev'].map_to_host(queue=queue, flags=cl.map_flags.READ, is_blocking=True) + p1 = numpy.asarray(mapped_p[0], dtype=numpy.float32).copy() + 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