added opencl backend

This commit is contained in:
ykcheungab-dragonq6a
2026-02-17 17:38:11 +08:00
parent 810060bb16
commit c8fc9b6a74
3 changed files with 651 additions and 22 deletions
+9 -1
View File
@@ -1,18 +1,26 @@
import os import os
import sys
os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0" os.environ["OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS"] = "0"
import cv2 import cv2
# Ensure benchmark uses local workspace code, not site-packages installation.
_local_src = os.path.join(os.path.dirname(__file__), "src")
if _local_src not in sys.path:
sys.path.insert(0, _local_src)
import orisys import orisys
import numpy as np import numpy as np
import time import time
def main(): def main():
print(f"Using orisys module: {orisys.__file__}")
# 步骤1:创建传感器对象 (必要步骤), # 步骤1:创建传感器对象 (必要步骤),
# config_name为配置文件名(内置配置文件名)或路径(自定义配置文件路径) # config_name为配置文件名(内置配置文件名)或路径(自定义配置文件路径)
# cal_path为校准文件路径, 每一个设备需要创建/使用一个校准文件 # cal_path为校准文件路径, 每一个设备需要创建/使用一个校准文件
# 方式1: 从视频文件读取 # 方式1: 从视频文件读取
sensor = orisys.Sensor("s1.mp4", config_name="./config/prototype_of3.json", cal_path= "./config/dd01a_0.npy", isstitch=True, cuda=True, verbose=True) sensor = orisys.Sensor("s1.mp4", config_name="./config/prototype_of3.json", cal_path= "./config/dd01a_0.npy", isstitch=True, cuda=False, verbose=True, poisson_backend='opencl')
# 方式2: 从摄像头读取(取消注释使用) # 方式2: 从摄像头读取(取消注释使用)
# sensor = orisys.Sensor(0, config_name="./config/prototype_of3.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True) # sensor = orisys.Sensor(0, config_name="./config/prototype_of3.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True)
+634 -18
View File
@@ -84,6 +84,11 @@ class nHHD(object):
self.stream1 = None self.stream1 = None
self.stream2 = 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() args = kwargs.keys()
if ('grid' in args) == ('points' in args): if ('grid' in args) == ('points' in args):
@@ -110,8 +115,30 @@ class nHHD(object):
if (self.dim != 2) and (self.dim != 3): if (self.dim != 2) and (self.dim != 3):
raise ValueError("nHHD Solver works only for 2 and 3 dimensions") raise ValueError("nHHD Solver works only for 2 and 3 dimensions")
self._force_cpu_cache = None
self.psolver.prepare(False) 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 # create the 3-component decomposition
def decompose(self, vfield, verbose=0): def decompose(self, vfield, verbose=0):
@@ -131,11 +158,19 @@ class nHHD(object):
print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max()) print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())
# compute potentials # compute potentials
self.nD = self.psolver.solve_gpu(self.div, verbose > 0) #1) 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: if verbose > 0:
print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max()) print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max())
self.nRu = self.psolver.solve_gpu(self.curlw, verbose > 0) #1)
if verbose > 0: if verbose > 0:
print( 'nR =', self.nRu.shape, self.nRu.min(), self.nRu.max()) print( 'nR =', self.nRu.shape, self.nRu.min(), self.nRu.max())
@@ -159,21 +194,21 @@ class nHHD(object):
print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max()) print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())
# compute potentials # 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: if verbose > 0:
print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max()) 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 self.nRu *= -1.0
if verbose > 0: if verbose > 0:
print( 'nRu =', self.nRu.shape, self.nRu.min(), self.nRu.max()) 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 self.nRv *= -1.0
if verbose > 0: if verbose > 0:
print( 'nRv =', self.nRv.shape, self.nRv.min(), self.nRv.max()) 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 self.nRw *= -1.0
if verbose > 0: if verbose > 0:
@@ -244,13 +279,41 @@ class nHHD(object):
print('curlw =', curlw_gpu.shape, float(curlw_gpu.min()), float(curlw_gpu.max())) 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) # In GPU decomposition, keep data on CuPy by default.
with self.stream1: # OpenCL here requires CuPy<->CPU copies and is only enabled
nD_gpu = self.psolver._solve_gpu_internal(div_gpu) # when explicitly requested.
with self.stream2: use_opencl = (
nRu_gpu = self.psolver._solve_gpu_internal(curlw_gpu) self.poisson_backend == 'opencl' and
self.stream1.synchronize() ('OPENCL_AVAILABLE' in globals() and OPENCL_AVAILABLE)
self.stream2.synchronize() )
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: if verbose > 0:
print('nD =', nD_gpu.shape, float(nD_gpu.min()), float(nD_gpu.max())) print('nD =', nD_gpu.shape, float(nD_gpu.min()), float(nD_gpu.max()))
@@ -301,10 +364,12 @@ class nHHD(object):
return None return None
def get_force(self, vfield, verbose=False): def get_force(self, vfield, verbose=False):
shape = (vfield.shape[0], vfield.shape[1])
mgrid = numpy.mgrid[0: vfield.shape[0], 0: vfield.shape[1]] if (not hasattr(self, '_force_cpu_cache') or
grid = numpy.stack((mgrid[0], mgrid[1]), axis =2) self._force_cpu_cache is None or
# print(grid.shape) self._force_cpu_cache['shape'] != shape):
self._init_force_cpu_cache(shape)
grid = self._force_cpu_cache['grid']
self.decompose(vfield, verbose) self.decompose(vfield, verbose)
@@ -339,6 +404,17 @@ class nHHD(object):
return flag,output return flag,output
# return self.d, self.h, numpy.cross(r_p, self.r) , numpy.cross(r_n, self.r) # 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): def _init_force_gpu_cache(self, shape):
""" """
Initialize cached GPU arrays for force computation. Initialize cached GPU arrays for force computation.
@@ -895,6 +971,20 @@ except ImportError:
cp_signal = None cp_signal = None
CUPY_AVAILABLE = False 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
''' '''
# ------------------------------------------------------------------------------ # ------------------------------------------------------------------------------
@@ -992,6 +1082,13 @@ class PoissonSolver(object):
self.ready = False self.ready = False
self._gpu_cache = None # Cache for GPU arrays 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'): # if (self.ptype == 'G'):
# print( 'PoissonSolver:', self.dim, 'D grid =', self.gdims, 'with spacings', self.gdx) # print( 'PoissonSolver:', self.dim, 'D grid =', self.gdims, 'with spacings', self.gdx)
# elif (self.ptype == 'P'): # elif (self.ptype == 'P'):
@@ -1288,6 +1385,525 @@ class PoissonSolver(object):
# return p # 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): def _solve_gpu_internal(self, f_gpu):
"""Internal GPU solve that accepts and returns GPU arrays.""" """Internal GPU solve that accepts and returns GPU arrays."""
""" """
+8 -3
View File
@@ -53,7 +53,7 @@ class Sensor:
# ========================================================================= # =========================================================================
def __init__(self, vid_src, cuda=True, isstitch=True, verbose=True, def __init__(self, vid_src, cuda=True, isstitch=True, verbose=True,
config_name="prototype", cal_path=None, override=False, box=None, mirror=True): config_name="prototype", cal_path=None, override=False, box=None, mirror=True, poisson_backend='auto'):
""" """
Initialize Sensor with video source and configuration. Initialize Sensor with video source and configuration.
@@ -141,6 +141,7 @@ class Sensor:
self.__tracker = Tracker(adaptive=True, cuda=cuda, verbose=self.verbose, dis_preset=self._config_dis_preset, config_tracker=self.__config_tracker) self.__tracker = Tracker(adaptive=True, cuda=cuda, verbose=self.verbose, dis_preset=self._config_dis_preset, config_tracker=self.__config_tracker)
self.__contact3d = Contact3D(config_contact=self.__config_contact, verbose=self.verbose) self.__contact3d = Contact3D(config_contact=self.__config_contact, verbose=self.verbose)
self.__hhd = None self.__hhd = None
self.poisson_backend = poisson_backend
# Setup video capture # Setup video capture
self._open_video_source(override, box) self._open_video_source(override, box)
@@ -243,10 +244,14 @@ class Sensor:
# Initialize HHD solver # Initialize HHD solver
if self.__hhd is None: if self.__hhd is None:
self.__hhd = nHHD(cuda=self.cuda, grid=flow.shape[:2], spacings=(18, 18)) self.__hhd = nHHD(cuda=self.cuda, grid=flow.shape[:2], spacings=(18, 18), poisson_backend=self.poisson_backend)
# Force computation # Force computation
if self.cuda and self.CUPY_AVAILABLE: # When OpenCL backend is selected, use CPU 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) _, forces = self.__hhd.get_force_gpu_v2(flow)
else: else:
_, forces = self.__hhd.get_force(flow) _, forces = self.__hhd.get_force(flow)