fix fft correctness and changed resolution for vkfft speed

This commit is contained in:
ykcheungab-dragonq6a
2026-02-17 21:05:49 +08:00
parent 2479b1fa9a
commit 6a86007da1
4 changed files with 22 additions and 133 deletions
+2 -2
View File
@@ -6,8 +6,8 @@
"fps": 30,
"units": 15,
"units_used": 9,
"resolution_stitch": [384, 384],
"stitch_res_list": [384,256,128,64]
"resolution_stitch": [365, 365],
"stitch_res_list": [365,263,123,63]
},
"stitcher": {
"min_area_frac": 0.01,
+2 -2
View File
@@ -6,8 +6,8 @@
"fps": 30,
"units": 15,
"units_used": 9,
"resolution_stitch": [128, 128],
"stitch_res_list": [384,256,128,64]
"resolution_stitch": [123, 123],
"stitch_res_list": [365,263,123,63]
},
"stitcher": {
"min_area_frac": 0.01,
+2 -2
View File
@@ -6,8 +6,8 @@
"fps": 30,
"units": 15,
"units_used": 9,
"resolution_stitch": [256, 256],
"stitch_res_list": [384,256,128,64]
"resolution_stitch": [263, 263],
"stitch_res_list": [365,263,123,63]
},
"stitcher": {
"min_area_frac": 0.01,
+16 -127
View File
@@ -22,14 +22,8 @@ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
'''
import numpy
# from .nPoisson import PoissonSolver
# from .StructuredGrid import StructuredGrid
import timeit
import sys
from scipy import signal, spatial
from scipy.fft import next_fast_len
# from .UnstructuredGrid import UnstructuredGrid
# import matplotlib.pyplot as plt
try:
import cupy as cp
@@ -81,7 +75,7 @@ class nHHD(object):
self.cuda = cuda
if self.cuda and CUPY_AVAILABLE:
self.stream1 = cp.cuda.Stream()
self.stream2 = cp.cuda.Stream()
self.stream2 = cp.cuda.Stream()
else:
self.stream1 = None
self.stream2 = None
@@ -100,19 +94,9 @@ class nHHD(object):
raise SyntaxError("nHHD needs spacings for regular grid")
if 'grid' in args:
self.mesh = StructuredGrid(grid=kwargs['grid'], spacings=kwargs['spacings'])
self.psolver = PoissonSolver(grid=self.mesh.dims, spacings=self.mesh.dx)
# elif 'points' in args:
# if 'simplices' in args:
# self.mesh = UnstructuredGrid(vertices=kwargs['points'], simplices=kwargs['simplices'])
# else:
# self.mesh = UnstructuredGrid(vertices=kwargs['points'])
# self.psolver = PoissonSolver(points=self.mesh.vertices, pvolumes=self.mesh.pvolumes)
self.dim = self.mesh.dim
if (self.dim != 2) and (self.dim != 3):
raise ValueError("nHHD Solver works only for 2 and 3 dimensions")
@@ -125,8 +109,7 @@ class nHHD(object):
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=True)
return self.psolver.solve_opencl(field, verbose, profile=False)
if self.poisson_backend == 'cupy':
return self.psolver.solve_gpu(field, verbose)
@@ -134,8 +117,7 @@ class nHHD(object):
# 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=True)
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")
@@ -159,18 +141,10 @@ class nHHD(object):
print( 'div =', self.div.shape, self.div.min(), self.div.max())
print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max())
# # compute potentials
# self.nD = self.psolver.solve_gpu(self.div, verbose > 0) #1)
# if verbose > 0:
# 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:
# print( 'nR =', self.nRu.shape, self.nRu.min(), self.nRu.max())
# compute potentials
if self.poisson_backend == 'opencl' and ('OPENCL_AVAILABLE' in globals() and OPENCL_AVAILABLE):
try:
self.nD, self.nRu = self.psolver.solve_opencl_batch2(self.div, self.curlw, verbose > 0, profile=True)
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")
@@ -181,7 +155,6 @@ class nHHD(object):
self.nRu = self._solve_poisson_backend(self.curlw, verbose > 0)
if verbose > 0:
print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max())
if verbose > 0:
print( 'nR =', self.nRu.shape, self.nRu.min(), self.nRu.max())
# compute fields as gradients of potentials
@@ -1601,7 +1574,7 @@ class PoissonSolver(object):
slices = tuple(slice(starts[i], starts[i] + target_shape[i]) for i in range(len(target_shape)))
return full_conv[slices]
def solve_opencl(self, f, verbose=False, profile=True):
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.
@@ -1635,7 +1608,7 @@ class PoissonSolver(object):
t0 = timeit.default_timer()
f_host = numpy.ascontiguousarray(f, dtype=numpy.float32)
full_shape = tuple(next_fast_len(2 * s - 1) for s in f_host.shape)
full_shape = tuple(2 * s - 1 for s in f_host.shape)
timing['host_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
cache_valid = (
@@ -1655,8 +1628,9 @@ class PoissonSolver(object):
g_fft_dev = cl_rfftn(g_dev)
if profile:
queue.finish()
linear_shape = tuple(2 * s - 1 for s in f_host.shape)
same_starts = [(linear_shape[i] - f_host.shape[i]) // 2 for i in range(self.dim)]
# 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)
@@ -1795,11 +1769,12 @@ class PoissonSolver(object):
print('OpenCL solver not optimized for unstructured points, using CPU')
return self.solve(f, verbose)
def solve_opencl_batch2(self, f1, f2, verbose=False, profile=True):
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).
When profile=True, records and optionally prints per-stage timings.
"""
if not OPENCL_AVAILABLE:
if verbose:
@@ -1833,7 +1808,7 @@ class PoissonSolver(object):
t0 = timeit.default_timer()
f1_host = numpy.ascontiguousarray(f1, dtype=numpy.float32)
f2_host = numpy.ascontiguousarray(f2, dtype=numpy.float32)
full_shape = tuple(next_fast_len(2 * s - 1) for s in f1_host.shape)
full_shape = tuple(2 * s - 1 for s in f1_host.shape)
timing['host_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
cache_valid = (
@@ -1854,8 +1829,9 @@ class PoissonSolver(object):
if profile:
queue.finish()
linear_shape = tuple(2 * s - 1 for s in f1_host.shape)
same_starts = [(linear_shape[i] - f1_host.shape[i]) // 2 for i in range(self.dim)]
# 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 = [f1_host.shape[i] - 1 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:
@@ -2018,98 +1994,11 @@ class PoissonSolver(object):
p_gpu *= self._gpu_cache['gdx_prod']
return p_gpu # Return CuPy array
# def solve_gpu_batch(self, f_list, verbose=False):
# """
# Batch GPU-accelerated Poisson solver for multiple inputs.
# Solves multiple Poisson equations in a single GPU call.
# Args:
# f_list: List of input fields (numpy arrays)
# verbose: Print timing information
# Returns:
# List of solutions (numpy arrays on CPU)
# """
# if not CUPY_AVAILABLE:
# if verbose:
# print("CuPy not available, falling back to CPU solver")
# return [self.solve(f, verbose) for f in f_list]
# if not self.ready:
# self.prepare(verbose)
# gtimer = Timer()
# if verbose:
# print(f'\nSolving {len(f_list)} Poisson Eq. on GPU (batched)')
# # Verify all inputs have same shape
# fshape = f_list[0].shape
# for f in f_list:
# if f.shape != fshape:
# raise ValueError("All input fields must have the same shape for batch processing")
# # ----------------------------------------------------------------------
# # regular grid
# if self.ptype == 'G':
# if self.gdims != fshape:
# raise ValueError("Shape of function should match shape of grid")
# # convolution in frequency domain using GPU
# if self.stype == 'F':
# # Initialize GPU cache if needed
# if self._gpu_cache is None:
# if verbose:
# print(' - Transferring Green\'s function to GPU...')
# self._gpu_cache = {
# 'G_gpu': cp.asarray(self.G, dtype=cp.float32),
# 'gdx_prod': numpy.prod(self.gdx)
# }
# results = []
# # Process all inputs in batch
# for i, f in enumerate(f_list):
# # Transfer input to GPU
# f_gpu = cp.asarray(f, dtype=cp.float32)
# # Perform FFT convolution on GPU
# 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']
# # Transfer result back to CPU
# results.append(cp.asnumpy(p_gpu))
# if verbose:
# print(f' - Solved {i+1}/{len(f_list)}')
# if verbose:
# print('Poisson solver (GPU batch) finished')
# gtimer.end()
# return results
# elif self.stype == 'S':
# raise SyntaxError('Spatial solver not suggested for structured grids')
# # ----------------------------------------------------------------------
# # unstructured points - not optimized for GPU
# elif self.ptype == 'P':
# if verbose:
# print(' - GPU solver not optimized for unstructured points, using CPU')
# return [self.solve(f, verbose) for f in f_list]
# # If we reach here, it's an unsupported case
# raise ValueError(f"Unsupported point type '{self.ptype}' for GPU batch solver")
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()