orisys-tech/orisys-sdk
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

optimized pyopencl HHD on radxa

Browse Source
This commit is contained in:
ykcheungab-dragonq6a
2026-02-17 19:38:16 +08:00
parent bbfbb913dd
commit b84ca53bbc
1 changed files with 195 additions and 90 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/orisys/HHD.py
+195 -90
View File
@@ -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