refractor for HHD and depth backend
This commit is contained in:
bluepulsar
+322
-382
@@ -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()
|
||||
|
||||
|
||||
@@ -0,0 +1,152 @@
|
||||
"""
|
||||
Centralized backend availability checking for GPU acceleration libraries.
|
||||
|
||||
This module provides a single source of truth for checking whether CuPy and OpenCL
|
||||
are available, avoiding duplicate import checks across the codebase.
|
||||
"""
|
||||
|
||||
# CuPy availability and related objects
|
||||
_cupy_available = False
|
||||
_cupy_cp = None
|
||||
_cupy_cp_signal = None
|
||||
_cupy_streams_available = False
|
||||
_cupy_stream = None
|
||||
_cupy_torque_kernel = None
|
||||
_cupy_kernel_available = False
|
||||
|
||||
# OpenCL availability and related objects
|
||||
_opencl_available = False
|
||||
_opencl_cl = None
|
||||
_opencl_cla = None
|
||||
_opencl_rfftn = None
|
||||
_opencl_irfftn = None
|
||||
|
||||
# Try to import CuPy
|
||||
try:
|
||||
import cupy as cp
|
||||
import cupyx.scipy.signal as cp_signal
|
||||
_cupy_cp = cp
|
||||
_cupy_cp_signal = cp_signal
|
||||
_cupy_available = True
|
||||
|
||||
# Try to import CuPy streams for parallel execution
|
||||
try:
|
||||
from cupy.cuda import Stream
|
||||
_cupy_stream = Stream
|
||||
_cupy_streams_available = True
|
||||
except ImportError:
|
||||
_cupy_stream = None
|
||||
_cupy_streams_available = False
|
||||
|
||||
# Create fused elementwise kernel for torque computation
|
||||
try:
|
||||
_cupy_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'
|
||||
)
|
||||
_cupy_kernel_available = True
|
||||
except Exception:
|
||||
_cupy_torque_kernel = None
|
||||
_cupy_kernel_available = False
|
||||
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# Try to import OpenCL
|
||||
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_cl = cl
|
||||
_opencl_cla = cla
|
||||
_opencl_rfftn = cl_rfftn
|
||||
_opencl_irfftn = cl_irfftn
|
||||
_opencl_available = True
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
|
||||
def check_cupy_available():
|
||||
"""Check if CuPy is available."""
|
||||
return _cupy_available
|
||||
|
||||
|
||||
def check_opencl_available():
|
||||
"""Check if OpenCL/pyvkfft is available."""
|
||||
return _opencl_available
|
||||
|
||||
|
||||
def get_cupy():
|
||||
"""Get CuPy module if available, None otherwise."""
|
||||
return _cupy_cp
|
||||
|
||||
|
||||
def get_cupy_signal():
|
||||
"""Get cupyx.scipy.signal module if available, None otherwise."""
|
||||
return _cupy_cp_signal
|
||||
|
||||
|
||||
def get_cupy_stream():
|
||||
"""Get CuPy Stream class if available, None otherwise."""
|
||||
return _cupy_stream
|
||||
|
||||
|
||||
def get_cupy_torque_kernel():
|
||||
"""Get CuPy torque kernel if available, None otherwise."""
|
||||
return _cupy_torque_kernel
|
||||
|
||||
|
||||
def is_cupy_streams_available():
|
||||
"""Check if CuPy streams are available."""
|
||||
return _cupy_streams_available
|
||||
|
||||
|
||||
def is_cupy_kernel_available():
|
||||
"""Check if CuPy torque kernel is available."""
|
||||
return _cupy_kernel_available
|
||||
|
||||
|
||||
def get_opencl():
|
||||
"""Get pyopencl module if available, None otherwise."""
|
||||
return _opencl_cl
|
||||
|
||||
|
||||
def get_opencl_array():
|
||||
"""Get pyopencl.array module if available, None otherwise."""
|
||||
return _opencl_cla
|
||||
|
||||
|
||||
def get_opencl_rfftn():
|
||||
"""Get pyvkfft.fft.rfftn function if available, None otherwise."""
|
||||
return _opencl_rfftn
|
||||
|
||||
|
||||
def get_opencl_irfftn():
|
||||
"""Get pyvkfft.fft.irfftn function if available, None otherwise."""
|
||||
return _opencl_irfftn
|
||||
|
||||
|
||||
# For backward compatibility, export as module-level constants
|
||||
CUPY_AVAILABLE = _cupy_available
|
||||
OPENCL_AVAILABLE = _opencl_available
|
||||
STREAMS_AVAILABLE = _cupy_streams_available
|
||||
KERNEL_AVAILABLE = _cupy_kernel_available
|
||||
|
||||
# Export commonly used objects
|
||||
cp = _cupy_cp
|
||||
cp_signal = _cupy_cp_signal
|
||||
Stream = _cupy_stream
|
||||
_torque_kernel = _cupy_torque_kernel
|
||||
cl = _opencl_cl
|
||||
cla = _opencl_cla
|
||||
cl_rfftn = _opencl_rfftn
|
||||
cl_irfftn = _opencl_irfftn
|
||||
|
||||
@@ -1,14 +1,8 @@
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
# Try to import CuPy for GPU acceleration
|
||||
try:
|
||||
import cupy as cp
|
||||
CUPY_AVAILABLE = True
|
||||
except ImportError:
|
||||
cp = None
|
||||
CUPY_AVAILABLE = False
|
||||
|
||||
from .backend_availability import CUPY_AVAILABLE, get_cupy
|
||||
cp = get_cupy()
|
||||
# 缓存形态学操作的核，避免重复创建
|
||||
_morphology_kernels = {}
|
||||
|
||||
|
||||
+18
-219
@@ -20,6 +20,7 @@ from .tracker import Tracker
|
||||
from .HHD import nHHD
|
||||
from . import configs
|
||||
from .get_depth import Contact3D
|
||||
from .backend_availability import (check_cupy_available, check_opencl_available)
|
||||
|
||||
class Sensor:
|
||||
"""
|
||||
@@ -67,12 +68,9 @@ class Sensor:
|
||||
override: Skip auto-detection, use provided box (default: False)
|
||||
box: Manual bounding boxes when override=True (default: None)
|
||||
"""
|
||||
# Check CUDA availability
|
||||
try:
|
||||
import cupy as cp
|
||||
self.CUPY_AVAILABLE = True
|
||||
except ImportError:
|
||||
self.CUPY_AVAILABLE = False
|
||||
|
||||
self.CUPY_AVAILABLE = check_cupy_available()
|
||||
self.OPENCL_AVAILABLE = check_opencl_available()
|
||||
|
||||
if vid_src is None:
|
||||
raise Exception("User must specify a video source.")
|
||||
@@ -124,18 +122,6 @@ class Sensor:
|
||||
self._last_capture_time = 0.0
|
||||
self._last_processing_time = 0.0
|
||||
|
||||
# Async pipeline state
|
||||
# self._async_enabled = False
|
||||
# self._capture_thread = None
|
||||
# self._flow_thread = None
|
||||
# self._compute_thread = None
|
||||
# self._capture_queue = None
|
||||
# self._flow_queue = None
|
||||
# self._result_queue = None
|
||||
# self._stop_event = None
|
||||
|
||||
|
||||
|
||||
# Initialize components
|
||||
self.stitcher = Stitcher(isstitch=isstitch, config_stitcher=self.__config_stitcher, verbose=self.verbose)
|
||||
self.__tracker = Tracker(adaptive=True, cuda=cuda, verbose=self.verbose, dis_preset=self._config_dis_preset, config_tracker=self.__config_tracker)
|
||||
@@ -149,9 +135,13 @@ class Sensor:
|
||||
if self.verbose:
|
||||
print(f" Video source: {vid_src}")
|
||||
print(f" Stitch mode: {self.stitch_valid}")
|
||||
print(f" CUDA acceleration: {self.cuda}")
|
||||
print(f" CuPy available: {self.CUPY_AVAILABLE}")
|
||||
print(f" Optical flow algorithm: {self.__config_tracker['algorithm']}")
|
||||
print(f" Poisson backend: {self.poisson_backend}")
|
||||
# print(f" CUDA acceleration: {self.cuda}")
|
||||
print(f" CuPy available: {self.CUPY_AVAILABLE}")
|
||||
print(f" OpenCL available: {self.OPENCL_AVAILABLE}")
|
||||
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Public API - Core Methods (Most Frequently Used)
|
||||
@@ -247,14 +237,15 @@ class Sensor:
|
||||
self.__hhd = nHHD(cuda=self.cuda, grid=flow.shape[:2], spacings=(18, 18), poisson_backend=self.poisson_backend)
|
||||
|
||||
# Force computation
|
||||
# When OpenCL backend is selected, use CPU force path to ensure
|
||||
# When OpenCL backend is selected, use host force path to ensure
|
||||
# Poisson solve goes through PoissonSolver.solve_opencl().
|
||||
if self.poisson_backend == 'opencl':
|
||||
_, forces = self.__hhd.get_force(flow)
|
||||
elif self.cuda and self.CUPY_AVAILABLE:
|
||||
_, forces = self.__hhd.get_force_gpu_v2(flow)
|
||||
else:
|
||||
_, forces = self.__hhd.get_force(flow)
|
||||
_, forces = self.__hhd.get_force(flow)
|
||||
# if self.poisson_backend == 'opencl':
|
||||
# _, forces = self.__hhd.get_force_host(flow)
|
||||
# elif self.cuda and self.CUPY_AVAILABLE:
|
||||
# _, forces = self.__hhd.get_force_gpu_v2(flow)
|
||||
# else:
|
||||
# _, forces = self.__hhd.get_force_host(flow)
|
||||
self.fnormal, self.fshearx, self.fsheary, self.ftor = forces
|
||||
|
||||
# Get components
|
||||
@@ -359,8 +350,6 @@ class Sensor:
|
||||
|
||||
def disconnect(self):
|
||||
"""Release resources and stop async threads."""
|
||||
# if self._async_enabled:
|
||||
# self.disable_async()
|
||||
if self.cap.isOpened():
|
||||
self.cap.release()
|
||||
if self.verbose:
|
||||
@@ -521,193 +510,3 @@ class Sensor:
|
||||
print(f" ─────────────────────────")
|
||||
print(f" TOTAL: {timings['total']:7.2f} ms ({1000/timings['total']:.1f} FPS)")
|
||||
|
||||
|
||||
|
||||
|
||||
# =========================================================================
|
||||
# Public API - Async Pipeline
|
||||
# =========================================================================
|
||||
|
||||
# def enable_async(self, queue_size=2, compute_contact=False):
|
||||
# """
|
||||
# Enable 3-stage async pipeline: Capture → Flow → HHD.
|
||||
|
||||
# Each stage runs in parallel for maximum throughput.
|
||||
# """
|
||||
# if self._async_enabled:
|
||||
# return
|
||||
# self._compute_contact_enabled = compute_contact
|
||||
|
||||
# self._capture_queue = Queue(maxsize=queue_size)
|
||||
# self._flow_queue = Queue(maxsize=queue_size)
|
||||
# self._result_queue = Queue(maxsize=queue_size)
|
||||
# if compute_contact:
|
||||
# self._contact_queue = Queue(maxsize=queue_size)
|
||||
# self._stop_event = Event()
|
||||
# self._async_enabled = True
|
||||
|
||||
# self._capture_thread = Thread(target=self._capture_worker, daemon=True)
|
||||
# self._flow_thread = Thread(target=self._flow_worker, daemon=True)
|
||||
# self._compute_thread = Thread(target=self._hhd_worker, daemon=True)
|
||||
# if compute_contact:
|
||||
# self._contact_thread = Thread(target=self._contact_worker, daemon=True)
|
||||
|
||||
# self._capture_thread.start()
|
||||
# self._flow_thread.start()
|
||||
# self._compute_thread.start()
|
||||
# if compute_contact:
|
||||
# self._contact_thread.start()
|
||||
|
||||
# if self.verbose:
|
||||
# if compute_contact:
|
||||
# print("Async pipeline enabled: Capture → Flow → HHD → Contact")
|
||||
# else:
|
||||
# print("Async pipeline enabled: Capture → Flow → HHD")
|
||||
|
||||
# def get_async(self, timeout=0.05):
|
||||
# """
|
||||
# Get latest result from async pipeline (non-blocking).
|
||||
|
||||
# Returns:
|
||||
# dict or None: Sensor data if available
|
||||
# """
|
||||
# compute_contact = self._compute_contact_enabled
|
||||
# if not self._async_enabled:
|
||||
# return None
|
||||
# try:
|
||||
# result = self._result_queue.get(timeout=timeout)
|
||||
# t = time.perf_counter()
|
||||
# result['fps'] = 1.0 / (t - self.last_log_time) if self.last_log_time else 0
|
||||
# self.last_log_time = t
|
||||
|
||||
# # Sync instance variables
|
||||
# # for k in ['img', 'flow', 'normal', 'shear', 'depth', 'fnormal', 'fshearx', 'fsheary', 'ftor', 'fps']:
|
||||
# if compute_contact:
|
||||
# for k in ['img', 'flow', 'normal', 'shear', 'depth', 'centroid', 'centroid_found', 'contour', 'fps', 'fnormal', 'fshearx', 'fsheary', 'ftor']:
|
||||
# setattr(self, k, result.get(k))
|
||||
# else:
|
||||
# for k in ['img', 'flow', 'normal', 'shear', 'fnormal', 'fshearx', 'fsheary', 'ftor', 'fps']:
|
||||
# setattr(self, k, result.get(k))
|
||||
# self.vupdate = True
|
||||
# return result
|
||||
# except Empty:
|
||||
# return None
|
||||
|
||||
# def disable_async(self):
|
||||
# """Stop async pipeline."""
|
||||
# if not self._async_enabled:
|
||||
# return
|
||||
# self._stop_event.set()
|
||||
# for t in [self._capture_thread, self._flow_thread, self._compute_thread]:
|
||||
# if t: t.join(timeout=2.0)
|
||||
# self._async_enabled = False
|
||||
# self._capture_thread = self._flow_thread = self._compute_thread = None
|
||||
# self._capture_queue = self._flow_queue = self._result_queue = self._stop_event = None
|
||||
# if self.verbose:
|
||||
# print("Async pipeline disabled")
|
||||
|
||||
# =========================================================================
|
||||
# Private - Async Workers
|
||||
# =========================================================================
|
||||
|
||||
# def _queue_put(self, q, item):
|
||||
# """Put item in queue, dropping old if full."""
|
||||
# if q.full():
|
||||
# try: q.get_nowait()
|
||||
# except Empty: pass
|
||||
# q.put(item)
|
||||
|
||||
# def _capture_worker(self):
|
||||
# """Stage 1: Capture and preprocess frames by reusing get_img()."""
|
||||
# while not self._stop_event.is_set():
|
||||
# try:
|
||||
# # Reuse get_img() to avoid code duplication
|
||||
# if self.get_img() is None:
|
||||
# continue
|
||||
|
||||
# # Put the processed image in the queue
|
||||
# # get_img() updates self.img, so we use that
|
||||
# if self.img is not None:
|
||||
# self._queue_put(self._capture_queue, self.img.copy())
|
||||
# except Exception as e:
|
||||
# if self.verbose: print(f"Capture error: {e}")
|
||||
|
||||
# def _flow_worker(self):
|
||||
# """Stage 2: Compute optical flow by reusing __compute_flow()."""
|
||||
# while not self._stop_event.is_set():
|
||||
# try:
|
||||
# img = self._capture_queue.get(timeout=0.1)
|
||||
# # Set __img to reuse __compute_flow()
|
||||
# self.img = img
|
||||
# self.__compute_flow()
|
||||
# flow = self.flow.copy() if self.flow is not None else None
|
||||
# self._queue_put(self._flow_queue, (img, flow))
|
||||
# except Empty:
|
||||
# continue
|
||||
# except Exception as e:
|
||||
# if self.verbose: print(f"Flow error: {e}")
|
||||
|
||||
# def _hhd_worker(self):
|
||||
# """Stage 3: HHD decomposition by reusing __compute_hhd()."""
|
||||
# while not self._stop_event.is_set():
|
||||
# try:
|
||||
# img, flow = self._flow_queue.get(timeout=0.1)
|
||||
|
||||
# # Set __flow to reuse __compute_hhd()
|
||||
# self.flow = flow
|
||||
# self.__compute_hhd()
|
||||
# # Get results from internal variables set by __compute_hhd()
|
||||
# result = {
|
||||
# 'img': img, 'flow': flow,
|
||||
# 'normal': self.vnormal.copy() if self.vnormal is not None else None,
|
||||
# 'shear': self.vshear.copy() if self.vshear is not None else None,
|
||||
# # 'depth': getFilteredDensity(flow, use_cuda=self.cuda),
|
||||
# 'fnormal': self.fnormal, 'fshearx': self.fshearx,
|
||||
# 'fsheary': self.fsheary, 'ftor': self.ftor,
|
||||
# }
|
||||
|
||||
# # Always include contact fields to ensure consistent result structure
|
||||
# # When compute_contact is True, contact worker will update these values
|
||||
# # When compute_contact is False, they remain None
|
||||
# result.update({
|
||||
# 'centroid': None,
|
||||
# 'centroid_found': False,
|
||||
# 'contour': None,
|
||||
# 'depth_map': None,
|
||||
# 'depth': None, # Alias for depth_map for compatibility
|
||||
# })
|
||||
|
||||
# self._queue_put(self._result_queue, result)
|
||||
# except Empty:
|
||||
# continue
|
||||
# except Exception as e:
|
||||
# if self.verbose: print(f"HHD error: {e}")
|
||||
|
||||
# def _contact_worker(self):
|
||||
# """Stage 4: Contact computation using vnormal from HHD results."""
|
||||
# while not self._stop_event.is_set():
|
||||
# try:
|
||||
# # Get results from HHD worker which includes vnormal
|
||||
# hhd_result = self._result_queue.get(timeout=0.1)
|
||||
# vnormal = hhd_result.get('normal') # This is vnormal from HHD
|
||||
|
||||
# if vnormal is not None:
|
||||
# # Compute contact using vnormal directly
|
||||
# contact_result = self.compute_contact(vnormal=vnormal)
|
||||
# # Add contact results to the existing HHD result
|
||||
# depth_map = contact_result.get('divergence_map')
|
||||
# hhd_result.update({
|
||||
# 'centroid': contact_result.get('centroid'),
|
||||
# 'centroid_found': contact_result.get('centroid_found'),
|
||||
# 'contour': contact_result.get('contour'),
|
||||
# 'depth_map': depth_map,
|
||||
# 'depth': depth_map, # Alias for depth_map for compatibility
|
||||
# })
|
||||
|
||||
# # Put updated result back to queue
|
||||
# self._queue_put(self._result_queue, hhd_result)
|
||||
# except Empty:
|
||||
# continue
|
||||
# except Exception as e:
|
||||
# if self.verbose: print(f"Contact error: {e}")
|
||||
|
||||
|
||||
Reference in New Issue
Block a user