Merge branch 'radxa' of yahoo.com.hk@192.168.192.190:D:\repo\orisyssdk_git into radxa
This commit is contained in:
bluepulsar
+640
-16
@@ -84,6 +84,11 @@ class nHHD(object):
|
||||
else:
|
||||
self.stream1 = None
|
||||
self.stream2 = None
|
||||
|
||||
# Select Poisson backend: 'auto' (default), 'cupy', or 'opencl'
|
||||
self.poisson_backend = kwargs.get('poisson_backend', 'auto').lower()
|
||||
if self.poisson_backend not in ('auto', 'cupy', 'opencl'):
|
||||
raise ValueError("poisson_backend must be one of: 'auto', 'cupy', 'opencl'")
|
||||
|
||||
args = kwargs.keys()
|
||||
|
||||
@@ -111,7 +116,29 @@ class nHHD(object):
|
||||
if (self.dim != 2) and (self.dim != 3):
|
||||
raise ValueError("nHHD Solver works only for 2 and 3 dimensions")
|
||||
|
||||
self._force_cpu_cache = None
|
||||
self.psolver.prepare(False)
|
||||
|
||||
def _solve_poisson_backend(self, field, verbose=False):
|
||||
"""
|
||||
Dispatch Poisson solve using the selected backend.
|
||||
"""
|
||||
if self.poisson_backend == 'opencl':
|
||||
# Always enable profiling for OpenCL backend to see performance breakdown
|
||||
return self.psolver.solve_opencl(field, verbose, profile=False)
|
||||
|
||||
if self.poisson_backend == 'cupy':
|
||||
return self.psolver.solve_gpu(field, verbose)
|
||||
|
||||
# auto: prefer OpenCL, then fallback
|
||||
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)
|
||||
except Exception as exc:
|
||||
if verbose:
|
||||
print(f"OpenCL solve failed ({exc}), falling back to CuPy/CPU")
|
||||
return self.psolver.solve_gpu(field, verbose)
|
||||
|
||||
# create the 3-component decomposition
|
||||
def decompose(self, vfield, verbose=0):
|
||||
@@ -139,6 +166,22 @@ class nHHD(object):
|
||||
# self.nRu = self.psolver.solve_gpu(self.curlw, verbose > 0) #1)
|
||||
# if verbose > 0:
|
||||
# print( 'nR =', self.nRu.shape, self.nRu.min(), self.nRu.max())
|
||||
# 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)
|
||||
except Exception as exc:
|
||||
if verbose > 0:
|
||||
print(f"OpenCL batch solve failed ({exc}), falling back to single solves")
|
||||
self.nD = self._solve_poisson_backend(self.div, verbose > 0)
|
||||
self.nRu = self._solve_poisson_backend(self.curlw, verbose > 0)
|
||||
else:
|
||||
self.nD = self._solve_poisson_backend(self.div, verbose > 0)
|
||||
self.nRu = self._solve_poisson_backend(self.curlw, verbose > 0)
|
||||
if verbose > 0:
|
||||
print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max())
|
||||
if verbose > 0:
|
||||
print( 'nR =', self.nRu.shape, self.nRu.min(), self.nRu.max())
|
||||
|
||||
# compute fields as gradients of potentials
|
||||
self.d = self.mesh.gradient(self.nD)
|
||||
@@ -160,21 +203,21 @@ class nHHD(object):
|
||||
print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())
|
||||
|
||||
# compute potentials
|
||||
self.nD = self.psolver.solve_gpu(self.div, verbose > 0) #1)
|
||||
self.nD = self._solve_poisson_backend(self.div, verbose > 0)
|
||||
if verbose > 0:
|
||||
print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max())
|
||||
|
||||
self.nRu = self.psolver.solve_gpu(self.curlu, verbose > 0)
|
||||
self.nRu = self._solve_poisson_backend(self.curlu, verbose > 0)
|
||||
self.nRu *= -1.0
|
||||
if verbose > 0:
|
||||
print( 'nRu =', self.nRu.shape, self.nRu.min(), self.nRu.max())
|
||||
|
||||
self.nRv = self.psolver.solve_gpu(self.curlv, verbose > 0)
|
||||
self.nRv = self._solve_poisson_backend(self.curlv, verbose > 0)
|
||||
self.nRv *= -1.0
|
||||
if verbose > 0:
|
||||
print( 'nRv =', self.nRv.shape, self.nRv.min(), self.nRv.max())
|
||||
|
||||
self.nRw = self.psolver.solve_gpu(self.curlw, verbose > 0)
|
||||
self.nRw = self._solve_poisson_backend(self.curlw, verbose > 0)
|
||||
self.nRw *= -1.0
|
||||
|
||||
if verbose > 0:
|
||||
@@ -245,13 +288,41 @@ class nHHD(object):
|
||||
print('curlw =', curlw_gpu.shape, float(curlw_gpu.min()), float(curlw_gpu.max()))
|
||||
|
||||
|
||||
# Compute potentials on GPU (solve_gpu returns numpy, so we convert back)
|
||||
with self.stream1:
|
||||
nD_gpu = self.psolver._solve_gpu_internal(div_gpu)
|
||||
with self.stream2:
|
||||
nRu_gpu = self.psolver._solve_gpu_internal(curlw_gpu)
|
||||
self.stream1.synchronize()
|
||||
self.stream2.synchronize()
|
||||
# In GPU decomposition, keep data on CuPy by default.
|
||||
# OpenCL here requires CuPy<->CPU copies and is only enabled
|
||||
# when explicitly requested.
|
||||
use_opencl = (
|
||||
self.poisson_backend == 'opencl' and
|
||||
('OPENCL_AVAILABLE' in globals() and OPENCL_AVAILABLE)
|
||||
)
|
||||
if use_opencl:
|
||||
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_gpu = cp.asarray(nD_cpu, dtype=cp.float32)
|
||||
nRu_gpu = cp.asarray(nRu_cpu, dtype=cp.float32)
|
||||
except Exception as exc:
|
||||
if self.poisson_backend == 'opencl':
|
||||
raise
|
||||
if verbose > 0:
|
||||
print(f"OpenCL path failed ({exc}), using CuPy internal solver")
|
||||
with self.stream1:
|
||||
nD_gpu = self.psolver._solve_gpu_internal(div_gpu)
|
||||
with self.stream2:
|
||||
nRu_gpu = self.psolver._solve_gpu_internal(curlw_gpu)
|
||||
self.stream1.synchronize()
|
||||
self.stream2.synchronize()
|
||||
else:
|
||||
if self.poisson_backend == 'opencl':
|
||||
raise RuntimeError("poisson_backend='opencl' requested but OpenCL is not available")
|
||||
with self.stream1:
|
||||
nD_gpu = self.psolver._solve_gpu_internal(div_gpu)
|
||||
with self.stream2:
|
||||
nRu_gpu = self.psolver._solve_gpu_internal(curlw_gpu)
|
||||
self.stream1.synchronize()
|
||||
self.stream2.synchronize()
|
||||
|
||||
if verbose > 0:
|
||||
print('nD =', nD_gpu.shape, float(nD_gpu.min()), float(nD_gpu.max()))
|
||||
@@ -302,10 +373,12 @@ class nHHD(object):
|
||||
return None
|
||||
|
||||
def get_force(self, vfield, verbose=False):
|
||||
|
||||
mgrid = numpy.mgrid[0: vfield.shape[0], 0: vfield.shape[1]]
|
||||
grid = numpy.stack((mgrid[0], mgrid[1]), axis =2)
|
||||
# print(grid.shape)
|
||||
shape = (vfield.shape[0], vfield.shape[1])
|
||||
if (not hasattr(self, '_force_cpu_cache') or
|
||||
self._force_cpu_cache is None or
|
||||
self._force_cpu_cache['shape'] != shape):
|
||||
self._init_force_cpu_cache(shape)
|
||||
grid = self._force_cpu_cache['grid']
|
||||
|
||||
|
||||
self.decompose(vfield, verbose)
|
||||
@@ -339,6 +412,17 @@ class nHHD(object):
|
||||
# print(flag)
|
||||
return flag,output
|
||||
# return self.d, self.h, numpy.cross(r_p, self.r) , numpy.cross(r_n, self.r)
|
||||
|
||||
def _init_force_cpu_cache(self, shape):
|
||||
"""
|
||||
Initialize cached CPU grid arrays for force computation.
|
||||
Avoids per-frame mgrid/stack allocations and reduces jitter.
|
||||
"""
|
||||
mgrid = numpy.mgrid[0:shape[0], 0:shape[1]]
|
||||
self._force_cpu_cache = {
|
||||
'shape': shape,
|
||||
'grid': numpy.stack((mgrid[0], mgrid[1]), axis=2),
|
||||
}
|
||||
|
||||
def _init_force_gpu_cache(self, shape):
|
||||
"""
|
||||
@@ -896,6 +980,20 @@ class StructuredGrid(object):
|
||||
# cp_signal = None
|
||||
# CUPY_AVAILABLE = False
|
||||
|
||||
# Try to import OpenCL + pyvkfft for FFT-based convolution
|
||||
try:
|
||||
import pyopencl as cl
|
||||
import pyopencl.array as cla
|
||||
from pyvkfft.fft import rfftn as cl_rfftn
|
||||
from pyvkfft.fft import irfftn as cl_irfftn
|
||||
OPENCL_AVAILABLE = True
|
||||
except ImportError:
|
||||
cl = None
|
||||
cla = None
|
||||
cl_rfftn = None
|
||||
cl_irfftn = None
|
||||
OPENCL_AVAILABLE = False
|
||||
|
||||
|
||||
'''
|
||||
# ------------------------------------------------------------------------------
|
||||
@@ -993,6 +1091,13 @@ class PoissonSolver(object):
|
||||
|
||||
self.ready = False
|
||||
self._gpu_cache = None # Cache for GPU arrays
|
||||
self._opencl_cache = None # Cache for OpenCL FFT buffers
|
||||
self._opencl_queue = None # OpenCL command queue
|
||||
self._opencl_context = None # OpenCL context (set when queue is created)
|
||||
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_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)
|
||||
# elif (self.ptype == 'P'):
|
||||
@@ -1288,6 +1393,525 @@ class PoissonSolver(object):
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
# return p
|
||||
|
||||
def _get_opencl_queue(self):
|
||||
"""Create and cache an OpenCL queue on first available GPU."""
|
||||
if self._opencl_queue is not None:
|
||||
return self._opencl_queue
|
||||
|
||||
if not OPENCL_AVAILABLE:
|
||||
raise RuntimeError("OpenCL/pyvkfft is not available")
|
||||
|
||||
ctx = None
|
||||
for platform in cl.get_platforms():
|
||||
for device in platform.get_devices():
|
||||
if device.type & cl.device_type.GPU == 0:
|
||||
continue
|
||||
ctx = cl.Context(devices=(device,))
|
||||
break
|
||||
if ctx is not None:
|
||||
break
|
||||
|
||||
if ctx is None:
|
||||
raise RuntimeError("No OpenCL GPU device found")
|
||||
|
||||
self._opencl_queue = cl.CommandQueue(ctx)
|
||||
self._opencl_context = ctx
|
||||
return self._opencl_queue
|
||||
|
||||
def _get_opencl_crop_kernel(self, queue):
|
||||
"""Create and cache OpenCL kernel for cropping convolution result."""
|
||||
if hasattr(self, '_opencl_crop_kernel') and self._opencl_crop_kernel is not None:
|
||||
return self._opencl_crop_kernel
|
||||
|
||||
# Simple 2D copy kernel
|
||||
kernel_code = """
|
||||
__kernel void crop_and_scale(
|
||||
__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
|
||||
) {
|
||||
int gid_x = get_global_id(0);
|
||||
int gid_y = get_global_id(1);
|
||||
|
||||
if (gid_x < dst_width && gid_y < dst_height) {
|
||||
int src_idx = (start_y + gid_y) * src_width + (start_x + gid_x);
|
||||
int dst_idx = 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_kernel = program.crop_and_scale
|
||||
return self._opencl_crop_kernel
|
||||
|
||||
def _get_opencl_scatter_kernel(self, queue):
|
||||
"""Kernel that copies unpadded input into a zero-padded device buffer (reverse of crop)."""
|
||||
if self._opencl_scatter_kernel is not None:
|
||||
return self._opencl_scatter_kernel
|
||||
|
||||
kernel_code = """
|
||||
__kernel void scatter_to_padded(
|
||||
__global const float* src,
|
||||
__global float* dst,
|
||||
const int src_width,
|
||||
const int src_height,
|
||||
const int dst_width
|
||||
) {
|
||||
int gid_x = get_global_id(0);
|
||||
int gid_y = get_global_id(1);
|
||||
if (gid_x < src_width && gid_y < src_height) {
|
||||
dst[gid_y * dst_width + gid_x] = src[gid_y * src_width + gid_x];
|
||||
}
|
||||
}
|
||||
"""
|
||||
program = cl.Program(self._opencl_context, kernel_code).build()
|
||||
self._opencl_scatter_kernel = program.scatter_to_padded
|
||||
return self._opencl_scatter_kernel
|
||||
|
||||
def _get_opencl_fused_mul_kernel(self, queue):
|
||||
"""Kernel that multiplies two complex spectra by one Green's spectrum in a single pass."""
|
||||
if self._opencl_fused_mul_kernel is not None:
|
||||
return self._opencl_fused_mul_kernel
|
||||
|
||||
kernel_code = """
|
||||
__kernel void fused_complex_mul(
|
||||
__global float2* fft1,
|
||||
__global float2* fft2,
|
||||
__global const float2* g,
|
||||
const int n
|
||||
) {
|
||||
int i = get_global_id(0);
|
||||
if (i < n) {
|
||||
float2 gv = g[i];
|
||||
float2 a = fft1[i];
|
||||
float2 b = fft2[i];
|
||||
fft1[i] = (float2)(a.x * gv.x - a.y * gv.y,
|
||||
a.x * gv.y + a.y * gv.x);
|
||||
fft2[i] = (float2)(b.x * gv.x - b.y * gv.y,
|
||||
b.x * gv.y + b.y * gv.x);
|
||||
}
|
||||
}
|
||||
"""
|
||||
program = cl.Program(self._opencl_context, kernel_code).build()
|
||||
self._opencl_fused_mul_kernel = program.fused_complex_mul
|
||||
return self._opencl_fused_mul_kernel
|
||||
|
||||
@staticmethod
|
||||
def _center_crop_same(full_conv, target_shape):
|
||||
"""Crop full convolution result to scipy 'same' output shape."""
|
||||
starts = [(full_conv.shape[i] - target_shape[i]) // 2 for i in range(len(target_shape))]
|
||||
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):
|
||||
"""
|
||||
OpenCL Poisson solver using pyvkfft FFT convolution.
|
||||
Replaces scipy.signal.fftconvolve(..., mode='same') for structured grids.
|
||||
When profile=True, records and optionally prints per-stage timings.
|
||||
"""
|
||||
if not OPENCL_AVAILABLE:
|
||||
if verbose:
|
||||
print("OpenCL/pyvkfft not available, falling back to CPU solver")
|
||||
return self.solve(f, verbose)
|
||||
|
||||
if not self.ready:
|
||||
self.prepare(verbose)
|
||||
|
||||
fshape = f.shape
|
||||
gtimer = Timer()
|
||||
if verbose:
|
||||
print('\nSolving Poisson Eq. with OpenCL FFT')
|
||||
|
||||
if self.ptype == 'G':
|
||||
if self.gdims != fshape:
|
||||
print(self.gdims, fshape)
|
||||
raise ValueError("Shape of function should match shape of grid")
|
||||
|
||||
if self.stype == 'F':
|
||||
timing = {'cache_prepare_ms': 0.0}
|
||||
t_total = timeit.default_timer()
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
queue = self._get_opencl_queue()
|
||||
timing['queue_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
f_host = numpy.ascontiguousarray(f, dtype=numpy.float32)
|
||||
full_shape = tuple(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
|
||||
)
|
||||
|
||||
if not cache_valid:
|
||||
t0 = timeit.default_timer()
|
||||
g_host = numpy.asarray(self.G, dtype=numpy.float32)
|
||||
g_pad = numpy.zeros(full_shape, dtype=numpy.float32)
|
||||
g_slices = tuple(slice(0, s) for s in g_host.shape)
|
||||
g_pad[g_slices] = g_host
|
||||
g_dev = cla.to_device(queue, g_pad)
|
||||
g_fft_dev = cl_rfftn(g_dev)
|
||||
if profile:
|
||||
queue.finish()
|
||||
same_starts = [(full_shape[i] - f_host.shape[i]) // 2 for i in range(self.dim)]
|
||||
same_slices = tuple(
|
||||
slice(same_starts[i], same_starts[i] + f_host.shape[i])
|
||||
for i in range(self.dim)
|
||||
)
|
||||
fft_shape = tuple(list(full_shape[:-1]) + [full_shape[-1] // 2 + 1])
|
||||
self._opencl_cache = {
|
||||
'fshape': f_host.shape,
|
||||
'full_shape': full_shape,
|
||||
'fft_shape': fft_shape,
|
||||
'g_fft_dev': g_fft_dev,
|
||||
'same_slices': same_slices,
|
||||
'same_starts': same_starts,
|
||||
'gdx_prod': numpy.prod(self.gdx),
|
||||
# Zero-init so padding stays zero across frames
|
||||
'f_dev': cla.zeros(queue, full_shape, dtype=numpy.float32),
|
||||
# Small staging buffer for input-sized H2D
|
||||
'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),
|
||||
}
|
||||
timing['cache_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# H2D: upload only input-sized data, scatter into padded device buffer
|
||||
t0 = timeit.default_timer()
|
||||
self._opencl_cache['f_stage_dev'].set(f_host, queue=queue)
|
||||
if self.dim == 2:
|
||||
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,
|
||||
numpy.int32(fshape_c[1]),
|
||||
numpy.int32(fshape_c[0]),
|
||||
numpy.int32(full_shape_c[1])
|
||||
)
|
||||
else:
|
||||
# 3D fallback: host-side pad + full upload
|
||||
f_pad = numpy.zeros(full_shape, dtype=numpy.float32)
|
||||
f_slices = tuple(slice(0, s) for s in f_host.shape)
|
||||
f_pad[f_slices] = f_host
|
||||
self._opencl_cache['f_dev'].set(f_pad, queue=queue)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['h2d_set_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
f_fft_dev = cl_rfftn(self._opencl_cache['f_dev'], self._opencl_cache['f_fft_dev'])
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['fft_f_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
f_fft_dev *= self._opencl_cache['g_fft_dev']
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['mul_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
conv_full_dev = cl_irfftn(f_fft_dev, self._opencl_cache['conv_full_dev'])
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['ifft_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# Optimized path: use OpenCL kernel to crop and scale on device, then transfer only cropped result.
|
||||
t0 = timeit.default_timer()
|
||||
p_dev = self._opencl_cache['p_dev']
|
||||
same_starts = self._opencl_cache['same_starts']
|
||||
fshape = self._opencl_cache['fshape']
|
||||
full_shape_cache = self._opencl_cache['full_shape']
|
||||
|
||||
if self.dim == 2:
|
||||
# Use custom kernel for 2D cropping and scaling
|
||||
crop_kernel = self._get_opencl_crop_kernel(queue)
|
||||
crop_kernel(
|
||||
queue, (fshape[1], fshape[0]), None,
|
||||
conv_full_dev.data,
|
||||
p_dev.data,
|
||||
numpy.int32(full_shape_cache[1]), # src_width
|
||||
numpy.int32(fshape[1]), # dst_width
|
||||
numpy.int32(fshape[0]), # dst_height
|
||||
numpy.int32(same_starts[1]), # start_x
|
||||
numpy.int32(same_starts[0]), # start_y
|
||||
numpy.float32(self._opencl_cache['gdx_prod']) # scale
|
||||
)
|
||||
if profile:
|
||||
queue.finish()
|
||||
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'])
|
||||
timing['d2h_get_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
else:
|
||||
# Fallback for 3D: transfer full, crop on CPU
|
||||
conv_full_dev *= self._opencl_cache['gdx_prod']
|
||||
queue.finish()
|
||||
timing['crop_scale_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
conv_full = conv_full_dev.get()
|
||||
timing['d2h_get_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
p = conv_full[self._opencl_cache['same_slices']].copy()
|
||||
timing['crop_scale_ms'] += (timeit.default_timer() - t0) * 1000.0
|
||||
timing['total_ms'] = (timeit.default_timer() - t_total) * 1000.0
|
||||
self._opencl_last_timing = timing
|
||||
|
||||
if verbose:
|
||||
print('Poisson solver (OpenCL FFT) finished')
|
||||
if profile:
|
||||
print(
|
||||
"OpenCL timing [ms]: "
|
||||
f"total={timing['total_ms']:.3f}, "
|
||||
f"queue={timing['queue_ms']:.3f}, "
|
||||
f"cache={timing['cache_prepare_ms']:.3f}, "
|
||||
f"host={timing['host_prepare_ms']:.3f}, "
|
||||
f"h2d={timing['h2d_set_ms']:.3f}, "
|
||||
f"fft={timing['fft_f_ms']:.3f}, "
|
||||
f"mul={timing['mul_ms']:.3f}, "
|
||||
f"ifft={timing['ifft_ms']:.3f}, "
|
||||
f"d2h={timing['d2h_get_ms']:.3f}, "
|
||||
f"crop={timing['crop_scale_ms']:.3f}"
|
||||
)
|
||||
gtimer.end()
|
||||
return p
|
||||
|
||||
elif self.stype == 'S':
|
||||
raise SyntaxError('Spatial solver not suggested for structured grids')
|
||||
|
||||
elif self.ptype == 'P':
|
||||
if verbose:
|
||||
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):
|
||||
"""
|
||||
Batched OpenCL Poisson solver for two 2D fields.
|
||||
Optimized: scatter H2D (~4x smaller transfer), fused multiply
|
||||
(reads g_fft once), async D2H (one sync instead of two).
|
||||
"""
|
||||
if not OPENCL_AVAILABLE:
|
||||
if verbose:
|
||||
print("OpenCL/pyvkfft not available, falling back to CPU solver")
|
||||
return self.solve(f1, verbose), self.solve(f2, verbose)
|
||||
|
||||
if self.dim != 2:
|
||||
return self.solve_opencl(f1, verbose, profile), self.solve_opencl(f2, verbose, profile)
|
||||
|
||||
if not self.ready:
|
||||
self.prepare(verbose)
|
||||
|
||||
fshape = f1.shape
|
||||
if fshape != f2.shape:
|
||||
raise ValueError("solve_opencl_batch2 requires inputs with identical shapes")
|
||||
|
||||
if self.ptype != 'G' or self.stype != 'F':
|
||||
return self.solve_opencl(f1, verbose, profile), self.solve_opencl(f2, verbose, profile)
|
||||
|
||||
if self.gdims != fshape:
|
||||
print(self.gdims, fshape)
|
||||
raise ValueError("Shape of function should match shape of grid")
|
||||
|
||||
timing = {'cache_prepare_ms': 0.0}
|
||||
t_total = timeit.default_timer()
|
||||
|
||||
t0 = timeit.default_timer()
|
||||
queue = self._get_opencl_queue()
|
||||
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)
|
||||
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
|
||||
)
|
||||
|
||||
if not cache_valid:
|
||||
t0 = timeit.default_timer()
|
||||
g_host = numpy.asarray(self.G, dtype=numpy.float32)
|
||||
g_pad = numpy.zeros(full_shape, dtype=numpy.float32)
|
||||
g_slices = tuple(slice(0, s) for s in g_host.shape)
|
||||
g_pad[g_slices] = g_host
|
||||
g_dev = cla.to_device(queue, g_pad)
|
||||
g_fft_dev = cl_rfftn(g_dev)
|
||||
if profile:
|
||||
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
|
||||
|
||||
self._opencl_cache = {
|
||||
'fshape': f1_host.shape,
|
||||
'full_shape': full_shape,
|
||||
'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),
|
||||
}
|
||||
timing['cache_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# H2D: async upload input-sized data, scatter into padded device 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,
|
||||
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])
|
||||
)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['h2d_set_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# FFT both fields
|
||||
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'])
|
||||
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
|
||||
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'])
|
||||
)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['mul_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# iFFT 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'])
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['ifft_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# Crop + scale on device (both fields, back-to-back enqueue)
|
||||
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,
|
||||
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'])
|
||||
)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['crop_scale_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# D2H: async enqueue both copies, single sync point
|
||||
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']
|
||||
timing['d2h_get_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
timing['total_ms'] = (timeit.default_timer() - t_total) * 1000.0
|
||||
self._opencl_last_timing = timing
|
||||
|
||||
if verbose:
|
||||
print('Poisson solver (OpenCL FFT batch2) finished')
|
||||
if profile:
|
||||
print(
|
||||
"OpenCL batch2 timing [ms]: "
|
||||
f"total={timing['total_ms']:.3f}, "
|
||||
f"queue={timing['queue_ms']:.3f}, "
|
||||
f"cache={timing['cache_prepare_ms']:.3f}, "
|
||||
f"host={timing['host_prepare_ms']:.3f}, "
|
||||
f"h2d={timing['h2d_set_ms']:.3f}, "
|
||||
f"fft={timing['fft_f_ms']:.3f}, "
|
||||
f"mul={timing['mul_ms']:.3f}, "
|
||||
f"ifft={timing['ifft_ms']:.3f}, "
|
||||
f"d2h={timing['d2h_get_ms']:.3f}, "
|
||||
f"crop={timing['crop_scale_ms']:.3f}"
|
||||
)
|
||||
return p1, p2
|
||||
|
||||
def _solve_gpu_internal(self, f_gpu):
|
||||
"""Internal GPU solve that accepts and returns GPU arrays."""
|
||||
@@ -1417,4 +2041,4 @@ class PoissonSolver(object):
|
||||
if CUPY_AVAILABLE:
|
||||
cp.get_default_memory_pool().free_all_blocks()
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user