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

refractor for HHD and depth backend

Browse Source
This commit is contained in:
bluepulsar
2026-02-18 10:56:30 +08:00
parent ef3a25a158
commit 54baf1ecde
4 changed files with 494 additions and 609 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/orisys/HHD.py
+322 -382
View File
@@ -23,67 +23,42 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import numpy
import timeit
import sys
from scipy import signal, spatial
try:
import cupy as cp
import cupyx.scipy.signal as cp_signal
CUPY_AVAILABLE = True
# Try to import CuPy streams for parallel execution
if CUPY_AVAILABLE:
try:
from cupy.cuda import Stream
STREAMS_AVAILABLE = True
except ImportError:
Stream = None
STREAMS_AVAILABLE = False
# Create fused elementwise kernel for torque computation
# Computes cross products and combines them in a single kernel pass
try:
_torque_kernel = cp.ElementwiseKernel(
'float32 r_x, float32 r_y, float32 grid_0, float32 grid_1, '
'float32 max_0, float32 max_1, float32 min_0, float32 min_1',
'float32 torque',
'''
float cross_p = (grid_0 - max_0) * r_y - (grid_1 - max_1) * r_x;
float cross_n = (grid_0 - min_0) * r_y - (grid_1 - min_1) * r_x;
torque = cross_p + cross_n;
''',
'torque_kernel'
)
KERNEL_AVAILABLE = True
except Exception:
_torque_kernel = None
KERNEL_AVAILABLE = False
except ImportError:
cp = None
cp_signal = None
CUPY_AVAILABLE = False
STREAMS_AVAILABLE = False
KERNEL_AVAILABLE = False
Stream = None
_torque_kernel = None
from .backend_availability import (
CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE,
cp, cp_signal, _torque_kernel,
cl, cla, cl_rfftn, cl_irfftn
)
# ---------------------------------------------------------------------------
class nHHD(object):
# constructor
def __init__(self, cuda=True, **kwargs):
self.cuda = cuda
if self.cuda and CUPY_AVAILABLE:
# Select Poisson backend: 'auto' (default), 'cupy', 'opencl', or 'cpu'
backend = kwargs.get('poisson_backend', 'auto').lower()
if backend not in ('auto', 'cupy', 'opencl', 'cpu'):
raise ValueError("poisson_backend must be one of: 'auto', 'cupy', 'opencl', 'cpu'")
# Resolve 'auto' to a specific backend
if backend == 'auto':
backend = 'cupy' if CUPY_AVAILABLE else ('opencl' if OPENCL_AVAILABLE else 'cpu')
# Validate backend availability and fallback to 'cpu' if needed
if backend == 'cupy' and not CUPY_AVAILABLE:
backend = 'cpu'
elif backend == 'opencl' and not OPENCL_AVAILABLE:
backend = 'cpu'
self.poisson_backend = backend
# Initialize CUDA streams for cupy backend
if self.poisson_backend == 'cupy':
self.stream1 = cp.cuda.Stream()
self.stream2 = cp.cuda.Stream()
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'")
self.cuda = cuda
args = kwargs.keys()
@@ -95,7 +70,7 @@ class nHHD(object):
if 'grid' in args:
self.mesh = StructuredGrid(grid=kwargs['grid'], spacings=kwargs['spacings'])
self.psolver = PoissonSolver(grid=self.mesh.dims, spacings=self.mesh.dx)
self.psolver = PoissonSolver(grid=self.mesh.dims, spacings=self.mesh.dx, poisson_backend=self.poisson_backend)
self.dim = self.mesh.dim
if (self.dim != 2) and (self.dim != 3):
@@ -104,6 +79,24 @@ class nHHD(object):
self._force_cpu_cache = None
self.psolver.prepare(False)
def get_force(self, vfield, verbose=False):
if self.poisson_backend == 'cupy':
return self.get_force_cuda(vfield, verbose)
else:
return self.get_force_cpu(vfield, verbose)
def get_shear(self):
if self.poisson_backend == 'cpu' or self.poisson_backend == 'opencl':
return self.h
elif self.poisson_backend == 'cupy':
return cp.asnumpy(self._h_gpu)
def get_normal(self):
if self.poisson_backend == 'cpu' or self.poisson_backend == 'opencl':
return self.d
elif self.poisson_backend == 'cupy':
return cp.asnumpy(self._d_gpu)
def _solve_poisson_backend(self, field, verbose=False):
"""
Dispatch Poisson solve using the selected backend.
@@ -113,18 +106,13 @@ class nHHD(object):
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:
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)
if self.poisson_backend == 'cpu':
return self.psolver.solve(field, verbose)
# create the 3-component decomposition
def decompose(self, vfield, verbose=0):
# decompose for cpu and opencl backend
if vfield.shape[-1] != self.dim:
raise ValueError("nHHD.decompose requires a valid-dimensional vector field")
@@ -143,13 +131,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)
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)
self.nD, self.nRu = self.psolver.solve_opencl_batch2(self.div, self.curlw, verbose > 0, profile=False)
else:
self.nD = self._solve_poisson_backend(self.div, verbose > 0)
self.nRu = self._solve_poisson_backend(self.curlw, verbose > 0)
@@ -217,7 +199,7 @@ class nHHD(object):
if verbose > 0:
print( 'h =', self.h.shape, self.h.min(), self.h.max())
def decompose_gpu(self, vfield, verbose=0):
def decompose_cuda(self, vfield, verbose=0):
"""
GPU-optimized decomposition that keeps intermediate results on GPU.
Returns GPU arrays to avoid CPU-GPU transfers.
@@ -239,7 +221,7 @@ class nHHD(object):
raise ValueError("nHHD.decompose requires a valid-dimensional vector field")
if verbose > 0:
print('\n ==> nhhd.decompose_gpu', vfield.shape)
print('\n ==> nhhd.decompose_cuda', vfield.shape)
# Convert input to GPU if needed
if isinstance(vfield, numpy.ndarray):
@@ -261,42 +243,12 @@ class nHHD(object):
print('div =', div_gpu.shape, float(div_gpu.min()), float(div_gpu.max()))
print('curlw =', curlw_gpu.shape, float(curlw_gpu.min()), float(curlw_gpu.max()))
# 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=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:
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()
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()))
@@ -346,7 +298,7 @@ class nHHD(object):
self.decompose(vfield, verbose)
return None
def get_force(self, vfield, verbose=False):
def get_force_cpu(self, vfield, verbose=False):
shape = (vfield.shape[0], vfield.shape[1])
if (not hasattr(self, '_force_cpu_cache') or
self._force_cpu_cache is None or
@@ -354,7 +306,6 @@ class nHHD(object):
self._init_force_cpu_cache(shape)
grid = self._force_cpu_cache['grid']
self.decompose(vfield, verbose)
self.Sn = numpy.sum(numpy.absolute(self.d))
@@ -423,7 +374,7 @@ class nHHD(object):
self._force_gpu_cache['mgrid_1'], shape
).copy()
def get_force_gpu_v2(self, vfield, verbose=False):
def get_force_cuda(self, vfield, verbose=False):
"""
Further optimized GPU force computation using CuPy ElementwiseKernel.
Fuses all operations into a single kernel for maximum performance.
@@ -443,7 +394,7 @@ class nHHD(object):
return self.get_force(vfield, verbose)
# Decompose the vector field
self.decompose_gpu(vfield, verbose)
self.decompose_cuda(vfield, verbose)
# Keep data on GPU - avoid transfers by computing directly
# d_gpu = cp.asarray(self.d, dtype=cp.float32)
@@ -509,15 +460,7 @@ class nHHD(object):
if CUPY_AVAILABLE:
cp.get_default_memory_pool().free_all_blocks()
def get_shear(self):
if not self.cuda or not CUPY_AVAILABLE:
return self.h
return cp.asnumpy(self._h_gpu)
def get_normal(self):
if not self.cuda or not CUPY_AVAILABLE:
return self.d
return cp.asnumpy(self._d_gpu)
# ---------------------------------------------------------------------------
@@ -558,12 +501,12 @@ class Timer(object):
# ---------------------------------------------------------------------------
# Try to import CuPy for GPU acceleration
try:
import cupy as cp
CUPY_AVAILABLE = True
except ImportError:
cp = None
CUPY_AVAILABLE = False
# try:
# import cupy as cp
# CUPY_AVAILABLE = True
# except ImportError:
# cp = None
# CUPY_AVAILABLE = False
# ------------------------------------------------------------------------------
class StructuredGrid(object):
@@ -955,18 +898,7 @@ class StructuredGrid(object):
# 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
'''
@@ -1059,22 +991,28 @@ class PoissonSolver(object):
raise SyntaxError('Frequncy solver strongly suggested for structured grids')
self.stype = kwargs['solver']
if 'poisson_backend' in args:
self.poisson_backend = kwargs['poisson_backend']
if self.poisson_backend not in ('cpu', 'opencl', 'cupy'):
raise ValueError("poisson_backend must be one of: 'cpu', 'opencl', 'cupy'")
if (self.dim > 3):
raise ValueError("Poisson Solver works only for 1, 2, or 3 dimensions")
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_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.poisson_backend == 'cupy':
self._gpu_cache = None # Cache for GPU arrays
if self.poisson_backend == 'opencl':
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_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)
# elif (self.ptype == 'P'):
@@ -1360,6 +1298,48 @@ class PoissonSolver(object):
elif self.stype == 'S':
raise SyntaxError('Spatial solver not suggested for structured grids')
def _solve_gpu_internal(self, f_gpu):
"""Internal GPU solve that accepts and returns GPU arrays."""
"""
GPU-accelerated Poisson solver using CuPy.
Falls back to CPU if CuPy is not available.
Args:
f_gpu: Input field (numpy array)
Returns:
p_gpu: Solution (CuPy array)
"""
# ----------------------------------------------------------------------
# regular grid
if self.ptype == 'G':
# convolution in frequency domain using GPU
if self.stype == 'F':
# Initialize GPU cache if needed
if self._gpu_cache is None:
self._gpu_cache = {
'G_gpu': cp.asarray(self.G, dtype=cp.float32),
'gdx_prod': numpy.prod(self.gdx)
}
# Perform FFT convolution on GPU
# CuPy's fftconvolve is GPU-accelerated
p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache['G_gpu'], mode='same')
# Scale by grid spacing (on GPU)
p_gpu *= self._gpu_cache['gdx_prod']
return p_gpu # Return CuPy array
def clear_gpu_cache(self):
"""Clear GPU cache to free GPU memory"""
if self._gpu_cache is not None:
self._gpu_cache = None
if CUPY_AVAILABLE:
cp.get_default_memory_pool().free_all_blocks()
# ----------------------------------------------------------------------
# unstructured points - not optimized for GPU
# elif self.ptype == 'P':
@@ -1574,200 +1554,6 @@ 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):
"""
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 and
'f_dev' 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()
# 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 = [f_host.shape[i] - 1 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()
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):
"""
@@ -1781,8 +1567,8 @@ class PoissonSolver(object):
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 self.dim != 2:
# return self.solve_opencl(f1, verbose, profile), self.solve_opencl(f2, verbose, profile)
if not self.ready:
self.prepare(verbose)
@@ -1960,45 +1746,199 @@ class PoissonSolver(object):
)
return p1, p2
def _solve_gpu_internal(self, f_gpu):
"""Internal GPU solve that accepts and returns GPU arrays."""
"""
GPU-accelerated Poisson solver using CuPy.
Falls back to CPU if CuPy is not available.
Args:
f_gpu: Input field (numpy array)
Returns:
p_gpu: Solution (CuPy array)
"""
# ----------------------------------------------------------------------
# regular grid
if self.ptype == 'G':
# convolution in frequency domain using GPU
if self.stype == 'F':
# Initialize GPU cache if needed
if self._gpu_cache is None:
self._gpu_cache = {
'G_gpu': cp.asarray(self.G, dtype=cp.float32),
'gdx_prod': numpy.prod(self.gdx)
}
# 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 and
# 'f_dev' 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()
# # 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 = [f_host.shape[i] - 1 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']
# Perform FFT convolution on GPU
# CuPy's fftconvolve is GPU-accelerated
p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache['G_gpu'], mode='same')
# Scale by grid spacing (on GPU)
p_gpu *= self._gpu_cache['gdx_prod']
# 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
return p_gpu # Return CuPy array
# t0 = timeit.default_timer()
# 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 clear_gpu_cache(self):
"""Clear GPU cache to free GPU memory"""
if self._gpu_cache is not None:
self._gpu_cache = None
if CUPY_AVAILABLE:
cp.get_default_memory_pool().free_all_blocks()