improve gpu computation in HHD computation
This commit is contained in:
unknown
+214
-155
@@ -102,7 +102,11 @@ class nHHD(object):
|
||||
if (self.dim != 2) and (self.dim != 3):
|
||||
raise ValueError("nHHD Solver works only for 2 and 3 dimensions")
|
||||
|
||||
self.psolver.prepare(False)
|
||||
self.psolver.prepare(False)
|
||||
|
||||
# Example pattern for stream usage:
|
||||
self.stream1 = cp.cuda.Stream()
|
||||
self.stream2 = cp.cuda.Stream()
|
||||
|
||||
# create the 3-component decomposition
|
||||
def decompose(self, vfield, verbose=0):
|
||||
@@ -228,20 +232,21 @@ class nHHD(object):
|
||||
dudy, dudx = cp.gradient(vfield_gpu[:,:,0], dx0, dx1)
|
||||
dvdy, dvdx = cp.gradient(vfield_gpu[:,:,1], dx0, dx1)
|
||||
|
||||
div_gpu = dudx + dvdy
|
||||
curlw_gpu = dvdx - dudy
|
||||
div_gpu = cp.add(dudx, dvdy)
|
||||
curlw_gpu = cp.subtract(dvdx, dudy)
|
||||
|
||||
if verbose > 0:
|
||||
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()))
|
||||
|
||||
|
||||
# Compute potentials on GPU (solve_gpu returns numpy, so we convert back)
|
||||
nD_cpu = self.psolver.solve_gpu(cp.asnumpy(div_gpu), verbose > 0)
|
||||
nRu_cpu = self.psolver.solve_gpu(cp.asnumpy(curlw_gpu), verbose > 0)
|
||||
|
||||
# Keep potentials on GPU
|
||||
nD_gpu = cp.asarray(nD_cpu, dtype=cp.float32)
|
||||
nRu_gpu = cp.asarray(nRu_cpu, dtype=cp.float32)
|
||||
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()))
|
||||
@@ -256,15 +261,15 @@ class nHHD(object):
|
||||
|
||||
# Compute rotated gradient on GPU
|
||||
rdy, rdx = cp.gradient(nRu_gpu, dx0, dx1)
|
||||
rdy *= -1.0
|
||||
rdy = cp.multiply(rdy, -1.0)
|
||||
r_gpu = cp.stack((rdy, rdx), axis=-1)
|
||||
|
||||
if verbose > 0:
|
||||
print('r =', r_gpu.shape, float(r_gpu.min()), float(r_gpu.max()))
|
||||
|
||||
# Compute h on GPU
|
||||
h_gpu = d_gpu + r_gpu
|
||||
h_gpu = vfield_gpu - h_gpu
|
||||
h_gpu = cp.add(d_gpu, r_gpu)
|
||||
h_gpu = cp.subtract(vfield_gpu, h_gpu)
|
||||
|
||||
if verbose > 0:
|
||||
print('h =', h_gpu.shape, float(h_gpu.min()), float(h_gpu.max()))
|
||||
@@ -276,11 +281,11 @@ class nHHD(object):
|
||||
self._nRu_gpu = nRu_gpu
|
||||
|
||||
# Also store CPU copies for compatibility
|
||||
self.d = cp.asnumpy(d_gpu)
|
||||
self.h = cp.asnumpy(h_gpu)
|
||||
self.r = cp.asnumpy(r_gpu)
|
||||
self.nRu = cp.asnumpy(nRu_gpu)
|
||||
self.nD = nD_cpu
|
||||
# self.d = cp.asnumpy(d_gpu)
|
||||
# self.h = cp.asnumpy(h_gpu)
|
||||
# self.r = cp.asnumpy(r_gpu)
|
||||
# self.nRu = cp.asnumpy(nRu_gpu)
|
||||
# self.nD = cp.asnumpy(nD_gpu)
|
||||
|
||||
return (d_gpu, h_gpu, r_gpu, nRu_gpu)
|
||||
|
||||
@@ -375,13 +380,18 @@ class nHHD(object):
|
||||
return self.get_force(vfield, verbose)
|
||||
|
||||
# Decompose the vector field
|
||||
self.decompose(vfield, verbose)
|
||||
self.decompose_gpu(vfield, verbose)
|
||||
|
||||
# Keep data on GPU - avoid transfers by computing directly
|
||||
d_gpu = cp.asarray(self.d, dtype=cp.float32)
|
||||
h_gpu = cp.asarray(self.h, dtype=cp.float32)
|
||||
r_gpu = cp.asarray(self.r, dtype=cp.float32)
|
||||
nRu_gpu = cp.asarray(self.nRu, dtype=cp.float32)
|
||||
# d_gpu = cp.asarray(self.d, dtype=cp.float32)
|
||||
# h_gpu = cp.asarray(self.h, dtype=cp.float32)
|
||||
# r_gpu = cp.asarray(self.r, dtype=cp.float32)
|
||||
# nRu_gpu = cp.asarray(self.nRu, dtype=cp.float32)
|
||||
|
||||
d_gpu = self._d_gpu
|
||||
h_gpu = self._h_gpu
|
||||
r_gpu = self._r_gpu
|
||||
nRu_gpu = self._nRu_gpu
|
||||
|
||||
# Use reduction kernels for better performance
|
||||
# Compute Sn using L1 norm (sum of absolute values)
|
||||
@@ -410,10 +420,17 @@ class nHHD(object):
|
||||
r_y = r_gpu[:, :, 1]
|
||||
|
||||
# Compute torque with fused sum
|
||||
Stor = float(cp.sum(
|
||||
(mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x +
|
||||
(mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x
|
||||
))
|
||||
# Stor = float(cp.sum(
|
||||
# (mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x +
|
||||
# (mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x
|
||||
# ))
|
||||
# Replace lines 413-416:
|
||||
if KERNEL_AVAILABLE and _torque_kernel is not None:
|
||||
torque = _torque_kernel(
|
||||
r_x, r_y, mgrid_0, mgrid_1,
|
||||
float(max_0), float(max_1), float(min_0), float(min_1)
|
||||
)
|
||||
Stor = float(cp.sum(torque))
|
||||
|
||||
self.Sn, self.Stx, self.Sty, self.Stor = Sn, Stx, Sty, Stor
|
||||
output = numpy.array([Sn, Stx, Sty, Stor], dtype=numpy.float32)
|
||||
@@ -430,12 +447,14 @@ class nHHD(object):
|
||||
cp.get_default_memory_pool().free_all_blocks()
|
||||
|
||||
def get_shear(self):
|
||||
if not CUPY_AVAILABLE:
|
||||
return self.h
|
||||
return cp.asnumpy(self._h_gpu)
|
||||
|
||||
return self.h
|
||||
|
||||
def get_normal(self):
|
||||
|
||||
return self.d
|
||||
if not CUPY_AVAILABLE:
|
||||
return self.d
|
||||
return cp.asnumpy(self._d_gpu)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -574,72 +593,72 @@ class StructuredGrid(object):
|
||||
|
||||
return (dudx, dwdy, dudz, dvdx)
|
||||
|
||||
def divcurl_gpu(self, vfield, verbose=False):
|
||||
"""
|
||||
GPU-accelerated computation of divergence and curl of vector field.
|
||||
Uses CuPy for faster gradient computation on GPU.
|
||||
"""
|
||||
if not CUPY_AVAILABLE:
|
||||
if verbose:
|
||||
print('CuPy not available, falling back to CPU')
|
||||
return self.divcurl(vfield, verbose)
|
||||
# def divcurl_gpu(self, vfield, verbose=False):
|
||||
# """
|
||||
# GPU-accelerated computation of divergence and curl of vector field.
|
||||
# Uses CuPy for faster gradient computation on GPU.
|
||||
# """
|
||||
# if not CUPY_AVAILABLE:
|
||||
# if verbose:
|
||||
# print('CuPy not available, falling back to CPU')
|
||||
# return self.divcurl(vfield, verbose)
|
||||
|
||||
if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims):
|
||||
raise ValueError("Dimensions of vector field should match that of the grid")
|
||||
# if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims):
|
||||
# raise ValueError("Dimensions of vector field should match that of the grid")
|
||||
|
||||
if verbose:
|
||||
print(' Computing div and curl on GPU...', end ='')
|
||||
sys.stdout.flush()
|
||||
mtimer = Timer()
|
||||
# if verbose:
|
||||
# print(' Computing div and curl on GPU...', end ='')
|
||||
# sys.stdout.flush()
|
||||
# mtimer = Timer()
|
||||
|
||||
if self.dim == 2:
|
||||
# Transfer to GPU
|
||||
vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
|
||||
# if self.dim == 2:
|
||||
# # Transfer to GPU
|
||||
# vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
|
||||
|
||||
# Compute gradients on GPU
|
||||
dudy, dudx = cp.gradient(vfield_gpu[:,:,0], self.dx[0], self.dx[1])
|
||||
dvdy, dvdx = cp.gradient(vfield_gpu[:,:,1], self.dx[0], self.dx[1])
|
||||
# # Compute gradients on GPU
|
||||
# dudy, dudx = cp.gradient(vfield_gpu[:,:,0], self.dx[0], self.dx[1])
|
||||
# dvdy, dvdx = cp.gradient(vfield_gpu[:,:,1], self.dx[0], self.dx[1])
|
||||
|
||||
# Compute div and curl on GPU
|
||||
div_gpu = dudx + dvdy
|
||||
curl_gpu = dvdx - dudy
|
||||
# # Compute div and curl on GPU
|
||||
# div_gpu = dudx + dvdy
|
||||
# curl_gpu = dvdx - dudy
|
||||
|
||||
# Transfer back to CPU
|
||||
div = cp.asnumpy(div_gpu)
|
||||
curl = cp.asnumpy(curl_gpu)
|
||||
# # Transfer back to CPU
|
||||
# div = cp.asnumpy(div_gpu)
|
||||
# curl = cp.asnumpy(curl_gpu)
|
||||
|
||||
if verbose:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
# if verbose:
|
||||
# print(' Done!', end ='')
|
||||
# mtimer.end()
|
||||
|
||||
return (div, curl)
|
||||
# return (div, curl)
|
||||
|
||||
elif self.dim == 3:
|
||||
# Transfer to GPU
|
||||
vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
|
||||
# elif self.dim == 3:
|
||||
# # Transfer to GPU
|
||||
# vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
|
||||
|
||||
# Compute gradients on GPU
|
||||
dudz, dudy, dudx = cp.gradient(vfield_gpu[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
|
||||
dvdz, dvdy, dvdx = cp.gradient(vfield_gpu[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
|
||||
dwdz, dwdy, dwdx = cp.gradient(vfield_gpu[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])
|
||||
# # Compute gradients on GPU
|
||||
# dudz, dudy, dudx = cp.gradient(vfield_gpu[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
|
||||
# dvdz, dvdy, dvdx = cp.gradient(vfield_gpu[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
|
||||
# dwdz, dwdy, dwdx = cp.gradient(vfield_gpu[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])
|
||||
|
||||
# Compute div and curl components on GPU
|
||||
div_gpu = dudx + dvdy + dwdz
|
||||
curlu_gpu = dwdy - dvdz
|
||||
curlv_gpu = dudz - dwdx
|
||||
curlw_gpu = dvdx - dudy
|
||||
# # Compute div and curl components on GPU
|
||||
# div_gpu = dudx + dvdy + dwdz
|
||||
# curlu_gpu = dwdy - dvdz
|
||||
# curlv_gpu = dudz - dwdx
|
||||
# curlw_gpu = dvdx - dudy
|
||||
|
||||
# Transfer back to CPU
|
||||
div = cp.asnumpy(div_gpu)
|
||||
curlu = cp.asnumpy(curlu_gpu)
|
||||
curlv = cp.asnumpy(curlv_gpu)
|
||||
curlw = cp.asnumpy(curlw_gpu)
|
||||
# # Transfer back to CPU
|
||||
# div = cp.asnumpy(div_gpu)
|
||||
# curlu = cp.asnumpy(curlu_gpu)
|
||||
# curlv = cp.asnumpy(curlv_gpu)
|
||||
# curlw = cp.asnumpy(curlw_gpu)
|
||||
|
||||
if verbose:
|
||||
print(' Done!', end ='')
|
||||
mtimer.end()
|
||||
# if verbose:
|
||||
# print(' Done!', end ='')
|
||||
# mtimer.end()
|
||||
|
||||
return (div, curlu, curlv, curlw)
|
||||
# return (div, curlu, curlv, curlw)
|
||||
|
||||
def curl3D(self, vfield, verbose=False):
|
||||
|
||||
@@ -1187,7 +1206,7 @@ class PoissonSolver(object):
|
||||
# --------------------------------------------------------------------------
|
||||
# GPU-accelerated solve using CuPy
|
||||
# --------------------------------------------------------------------------
|
||||
def solve_gpu(self, f, verbose=False):
|
||||
def solve_gpu(self, f, verbose=False, return_gpu=False):
|
||||
"""
|
||||
GPU-accelerated Poisson solver using CuPy.
|
||||
Falls back to CPU if CuPy is not available.
|
||||
@@ -1240,111 +1259,151 @@ class PoissonSolver(object):
|
||||
|
||||
# Scale by grid spacing (on GPU)
|
||||
p_gpu *= self._gpu_cache['gdx_prod']
|
||||
|
||||
if verbose:
|
||||
print('Poisson solver (GPU) finished')
|
||||
gtimer.end()
|
||||
|
||||
# Transfer result back to CPU
|
||||
p = cp.asnumpy(p_gpu)
|
||||
if return_gpu:
|
||||
return p_gpu # Return CuPy array
|
||||
return cp.asnumpy(p_gpu)
|
||||
# p = cp.asnumpy(p_gpu)
|
||||
|
||||
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)
|
||||
# elif self.ptype == 'P':
|
||||
# if verbose:
|
||||
# print(' - GPU solver not optimized for unstructured points, using CPU')
|
||||
# return self.solve(f, verbose)
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
if verbose:
|
||||
print('Poisson solver (GPU) finished')
|
||||
gtimer.end()
|
||||
|
||||
return p
|
||||
|
||||
# return p
|
||||
|
||||
def solve_gpu_batch(self, f_list, verbose=False):
|
||||
def _solve_gpu_internal(self, f_gpu):
|
||||
"""Internal GPU solve that accepts and returns GPU arrays."""
|
||||
"""
|
||||
Batch GPU-accelerated Poisson solver for multiple inputs.
|
||||
Solves multiple Poisson equations in a single GPU call.
|
||||
GPU-accelerated Poisson solver using CuPy.
|
||||
Falls back to CPU if CuPy is not available.
|
||||
|
||||
Args:
|
||||
f_list: List of input fields (numpy arrays)
|
||||
verbose: Print timing information
|
||||
f_gpu: Input field (numpy array)
|
||||
|
||||
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")
|
||||
p_gpu: Solution (CuPy array)
|
||||
"""
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# 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 = []
|
||||
# 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')
|
||||
|
||||
# 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')
|
||||
# Scale by grid spacing (on GPU)
|
||||
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.
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# 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]
|
||||
# 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 we reach here, it's an unsupported case
|
||||
raise ValueError(f"Unsupported point type '{self.ptype}' for GPU batch solver")
|
||||
# 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"""
|
||||
|
||||
Reference in New Issue
Block a user