fix fft correctness and changed resolution for vkfft speed
This commit is contained in:
ykcheungab-dragonq6a
+16
-127
@@ -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()
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user