opencl optimization
This commit is contained in:
radxa-d6a
+46
-30
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user