refractor HHD
This commit is contained in:
ykcheungab
@@ -0,0 +1,954 @@
|
||||
import numpy
|
||||
import timeit
|
||||
import sys
|
||||
# import cv2
|
||||
from scipy import signal, spatial
|
||||
|
||||
from ..backend_availability import (
|
||||
CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE,
|
||||
cp, cp_signal, _torque_kernel,
|
||||
cl, cla, cl_rfftn, cl_irfftn
|
||||
)
|
||||
from .timing import Timer
|
||||
|
||||
'''
|
||||
# ------------------------------------------------------------------------------
|
||||
Poisson Solver for 1D, 2D, and 3D scalar fields.
|
||||
Data may be defined on a rectlinear grid or unstructured points.
|
||||
|
||||
1) create an object
|
||||
PoissonSolver(args)
|
||||
'solver' : type of solver
|
||||
'S' for spatial convolution
|
||||
only option for unstructured points
|
||||
'F' for convolution in frequency space
|
||||
default option for regular grids
|
||||
|
||||
--> for regular grid
|
||||
'grid' : shape of the rectlinear grid
|
||||
tuple of dimensions (Z,Y,X)
|
||||
'spacings' : spacings of the rectlinear grid
|
||||
tuple of spacings (dz,dy,dx)
|
||||
|
||||
--> for unstrucutred points
|
||||
'points' : unstructured points
|
||||
numpy array of shape (N, dim)
|
||||
'pvolumes' : point volumes for each point
|
||||
numpy array of shape (N, 1)
|
||||
|
||||
2) call prepare()
|
||||
this initializes internal datastructures depending upon the mesh.
|
||||
|
||||
Once initialized, solve() can be called any number of times to
|
||||
solve the Poisson equation on the same domain. If the domain
|
||||
changes (e.g., a Lagrangian mesh), create a new object.
|
||||
|
||||
3) call solve(f), where, f is a numpy array defining a scalar field
|
||||
|
||||
for rectlinear grid, f.shape = grid.shape
|
||||
for unstructured points, f.shape = (points.shape[0], 1)
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
'''
|
||||
|
||||
class PoissonSolver(object):
|
||||
|
||||
# constructor
|
||||
def __init__(self, **kwargs):
|
||||
|
||||
args = kwargs.keys()
|
||||
|
||||
if ('grid' in args) == ('points' in args):
|
||||
raise SyntaxError("Poisson Solver needs either shape of a regular grid or the points in an rectilinear/unstructured grid")
|
||||
|
||||
if ('grid' in args) != ('spacings' in args):
|
||||
raise SyntaxError("Poisson Solver needs spacings for regular grid")
|
||||
|
||||
if ('points' in args) != ('pvolumes' in args):
|
||||
raise SyntaxError("Poisson Solver needs point volumes for unstructured points")
|
||||
|
||||
# (default) data type
|
||||
self.dtype = numpy.float32
|
||||
if 'dtype' in args:
|
||||
self.dtype = kwargs['dtype']
|
||||
|
||||
if 'grid' in args:
|
||||
self.ptype = 'G'
|
||||
self.stype = 'F' # default
|
||||
self.gdims = kwargs['grid']
|
||||
self.gdx = kwargs['spacings']
|
||||
self.dim = len(self.gdims)
|
||||
if self.dim != len(self.gdx):
|
||||
raise ValueError("Dimensions of spacings should match that of the grid")
|
||||
|
||||
elif 'points' in args:
|
||||
self.ptype = 'P'
|
||||
self.stype = 'S' # only option
|
||||
self.points = kwargs['points']
|
||||
self.pvolumes = kwargs['pvolumes']
|
||||
self.dim = len(self.points.shape)
|
||||
if self.pvolumes.shape[0] != self.points.shape[0]:
|
||||
raise ValueError("Number of pvolumes should match that of the points")
|
||||
|
||||
if 'solver' in args:
|
||||
if (kwargs['solver'] != 'S' and kwargs['solver'] != 'F'):
|
||||
raise SyntaxError('Solver can be only of type S(patial) or F(requency)')
|
||||
|
||||
if (kwargs['solver'] == 'F') and ('points' in args):
|
||||
raise SyntaxError('Only Spatial solver supported for unstructured points')
|
||||
|
||||
if (kwargs['solver'] == 'S') and ('grid' in args):
|
||||
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
|
||||
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
|
||||
# Cache for OpenCV FFT (CPU backend)
|
||||
self._opencv_cache = None # Cache for OpenCV FFT buffers
|
||||
# if (self.ptype == 'G'):
|
||||
# print( 'PoissonSolver:', self.dim, 'D grid =', self.gdims, 'with spacings', self.gdx)
|
||||
# elif (self.ptype == 'P'):
|
||||
# print( 'PoissonSolver', self.dim, 'D points =', self.points.shape)
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# create the Green's function in appropriate representation
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
# generate a grid -(N-1) to (N-1) to define distance function
|
||||
# where N is the size of the original grid
|
||||
def _create_radialGrid(self, verbose):
|
||||
|
||||
mtimer = Timer()
|
||||
if verbose:
|
||||
print( ' - creating distance kernel:', end ='')
|
||||
|
||||
rdims = tuple(2*d - 1 for d in self.gdims)
|
||||
|
||||
R = numpy.indices(rdims, dtype=self.dtype)
|
||||
for d in range(self.dim):
|
||||
R[d] += 1-self.gdims[d]
|
||||
R[d] *= self.gdx[d]
|
||||
|
||||
if self.dim == 1:
|
||||
numpy.absolute(R[0], R[0])
|
||||
R = R[0]
|
||||
|
||||
elif self.dim == 2:
|
||||
numpy.hypot(R[1], R[0], R[0])
|
||||
R = R[0]
|
||||
|
||||
# half of the average grid spacing
|
||||
zval = 0.5 * (self.gdx[0]+self.gdx[1]) / 2.0
|
||||
R[self.gdims[0]-1, self.gdims[1]-1] = zval
|
||||
|
||||
elif self.dim == 3:
|
||||
numpy.hypot(R[2], R[1], R[1])
|
||||
numpy.hypot(R[1], R[0], R[0])
|
||||
R = R[0]
|
||||
|
||||
# half of the average grid spacing
|
||||
zval = 0.5 * (self.gdx[0]+self.gdx[1]+self.gdx[2]) / 3.0
|
||||
R[self.gdims[0]-1, self.gdims[1]-1, self.gdims[2]-1] = zval
|
||||
|
||||
if verbose:
|
||||
print( R.shape, R.min(), R.max())
|
||||
mtimer.end()
|
||||
|
||||
return R
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# scale a function by point volumes for correct integration
|
||||
def _scale(self, func):
|
||||
|
||||
if self.ptype == 'P':
|
||||
return numpy.multiply(func, self.pvolumes)
|
||||
|
||||
elif self.ptype == 'G':
|
||||
sfunc = func * numpy.prod(self.gdx)
|
||||
|
||||
'''
|
||||
# volume of the boundary elements is only half the voxel
|
||||
if self.dim == 1:
|
||||
sfunc[ 0] *= 0.5
|
||||
sfunc[-1] *= 0.5
|
||||
elif self.dim == 2:
|
||||
sfunc[ 0, :] *= 0.5
|
||||
sfunc[-1, :] *= 0.5
|
||||
sfunc[ :, 0] *= 0.5
|
||||
sfunc[ :,-1] *= 0.5
|
||||
elif self.dim == 3:
|
||||
sfunc[ 0, :, :] *= 0.5
|
||||
sfunc[-1, :, :] *= 0.5
|
||||
sfunc[ :, 0, :] *= 0.5
|
||||
sfunc[ :,-1, :] *= 0.5
|
||||
sfunc[ :, :, 0] *= 0.5
|
||||
sfunc[ :, :,-1] *= 0.5
|
||||
'''
|
||||
return sfunc
|
||||
|
||||
def prepare(self, verbose=False):
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
gtimer = Timer()
|
||||
if verbose:
|
||||
print( '\nInitializing Poisson solver, type =', self.stype)
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# generate a grid to store pairwise distances to store Green's function
|
||||
self.G = None # definition depends upon solver type
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# for Fourier solution,
|
||||
if self.stype == 'F':
|
||||
|
||||
if self.ptype == 'G':
|
||||
self.G = self._create_radialGrid(verbose)
|
||||
|
||||
elif self.ptype == 'P':
|
||||
raise SyntaxError('Frequncy solver not supported for unstructured points')
|
||||
#self.G = create_distArray()
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# for spatial solution,
|
||||
# use the pairwise distance matrix, N x N
|
||||
# where N is the number of points
|
||||
elif self.stype == 'S':
|
||||
|
||||
if self.ptype == 'G':
|
||||
raise SyntaxError('Spatial solver not suggested for structured grids')
|
||||
#self.G = self.create_points(verbose)
|
||||
|
||||
elif self.ptype == 'P':
|
||||
# for 1D points, add an extra axis
|
||||
if (self.dim == 1) and (len(self.points.shape) == 1):
|
||||
self.points = self.points[:,numpy.newaxis]
|
||||
|
||||
# create pairwise distance matrix now
|
||||
ltimer = Timer()
|
||||
self.G = spatial.distance.cdist(self.points, self.points)
|
||||
|
||||
# TODO: use a good approximation here
|
||||
zval = numpy.power(10.0, -10.0)
|
||||
self.G += zval * numpy.identity(self.G.shape[0])
|
||||
|
||||
if verbose:
|
||||
print(' - created pairwise distance matrix:', self.G.shape, self.G.min(), self.G.max(), end = "")
|
||||
ltimer.end()
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# compute the Green's function
|
||||
ltimer = Timer()
|
||||
if verbose:
|
||||
print( ' - computing the Green\'s function:')
|
||||
|
||||
if self.dim == 1:
|
||||
numpy.multiply(self.G, 0.5, self.G)
|
||||
|
||||
elif self.dim == 2:
|
||||
numpy.log(self.G, self.G)
|
||||
numpy.multiply(self.G, (0.5 / numpy.pi), self.G)
|
||||
|
||||
elif self.dim == 3:
|
||||
numpy.reciprocal(self.G, self.G)
|
||||
numpy.multiply(self.G, (-0.25 / numpy.pi), self.G)
|
||||
|
||||
if verbose:
|
||||
print( self.G.shape, self.G.min(), self.G.max(), end = "")
|
||||
ltimer.end()
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
self.ready = True
|
||||
|
||||
if verbose:
|
||||
print( 'Poisson solver initialized', end ='')
|
||||
gtimer.end()
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# compute the integral solution
|
||||
# --------------------------------------------------------------------------
|
||||
def solve(self, f, verbose=False):
|
||||
|
||||
if not self.ready:
|
||||
self.prepare(verbose)
|
||||
|
||||
fshape = f.shape
|
||||
|
||||
gtimer = Timer()
|
||||
if verbose:
|
||||
print( '\nSolving Poisson Eq.')
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# regular grid
|
||||
if self.ptype == 'G':
|
||||
|
||||
#if (self.gdims - fshape).any():
|
||||
if self.gdims != fshape:
|
||||
print( self.gdims, fshape)
|
||||
raise ValueError("Shape of function should match shape of grid")
|
||||
|
||||
# convolution in frequency domain using OpenCV FFT (2D) or numpy FFT (3D)
|
||||
if self.stype == 'F':
|
||||
if self.dim == 2:
|
||||
# # Use OpenCV DFT for 2D (optimized with caching)
|
||||
# # Convert to float32 for OpenCV (use ascontiguousarray for better performance)
|
||||
# f_32 = numpy.ascontiguousarray(f, dtype=numpy.float32)
|
||||
|
||||
# # For 'same' mode convolution, pad to (2*N-1) size
|
||||
# full_shape = tuple(2 * s - 1 for s in f_32.shape)
|
||||
|
||||
# # Initialize or validate cache
|
||||
# cache_valid = (
|
||||
# self._opencv_cache is not None and
|
||||
# self._opencv_cache['fshape'] == fshape and
|
||||
# self._opencv_cache['full_shape'] == full_shape
|
||||
# )
|
||||
|
||||
# if not cache_valid:
|
||||
# # Cache G_fft (G doesn't change, so compute once)
|
||||
# G_32 = numpy.ascontiguousarray(self.G, dtype=numpy.float32)
|
||||
# G_fft = cv2.dft(G_32, flags=cv2.DFT_COMPLEX_OUTPUT)
|
||||
|
||||
# # Pre-compute crop slices for 'same' mode extraction
|
||||
# # For circular convolution at size (2*H-1, 2*W-1), the center corresponds to
|
||||
# # position [H-1, W-1], which matches scipy's fftconvolve mode='same' extraction
|
||||
# starts = [fshape[i] - 1 for i in range(len(fshape))]
|
||||
# slices = tuple(slice(starts[i], starts[i] + fshape[i]) for i in range(len(fshape)))
|
||||
|
||||
# # Pre-allocate reusable buffers
|
||||
# self._opencv_cache = {
|
||||
# 'fshape': fshape,
|
||||
# 'full_shape': full_shape,
|
||||
# 'G_fft': G_fft, # Cached: G doesn't change, so compute FFT once
|
||||
# 'slices': slices, # Cached: crop slices pre-computed
|
||||
# 'gdx_prod': numpy.prod(self.gdx), # Cached: grid spacing product
|
||||
# # Reusable buffers (avoid reallocation on each call)
|
||||
# 'f_pad': numpy.zeros(full_shape, dtype=numpy.float32), # Reused for padding
|
||||
# 'conv_fft': numpy.empty((full_shape[0], full_shape[1], 2), dtype=numpy.float32), # Output for mulSpectrums
|
||||
# }
|
||||
|
||||
# cache = self._opencv_cache
|
||||
|
||||
# # Zero-pad f into pre-allocated buffer (faster than creating new array)
|
||||
# f_pad = cache['f_pad']
|
||||
# f_pad.fill(0.0) # Clear previous data efficiently
|
||||
# f_slices = tuple(slice(0, s) for s in f_32.shape)
|
||||
# f_pad[f_slices] = f_32
|
||||
|
||||
# # Compute FFT of f (cv2.dft returns new array, but we can reuse for mulSpectrums)
|
||||
# f_fft = cv2.dft(f_pad, flags=cv2.DFT_COMPLEX_OUTPUT)
|
||||
|
||||
# # Multiply in frequency domain using OpenCV's optimized mulSpectrums
|
||||
# # This is faster than manual complex multiplication and supports output buffer
|
||||
# # Signature: mulSpectrums(a, b, flags, c=None, conjB=False)
|
||||
# cv2.mulSpectrums(f_fft, cache['G_fft'], cv2.DFT_COMPLEX_OUTPUT,
|
||||
# cache['conv_fft'], False)
|
||||
|
||||
# # Compute inverse FFT (cv2.idft returns new array)
|
||||
# conv_full = cv2.idft(cache['conv_fft'], flags=cv2.DFT_SCALE | cv2.DFT_REAL_OUTPUT)
|
||||
|
||||
# # Extract 'same' mode result (use cached slices, scale in-place)
|
||||
# p = conv_full[cache['slices']].copy()
|
||||
# numpy.multiply(p, cache['gdx_prod'], p)
|
||||
# from scipy import signal
|
||||
p = signal.fftconvolve(f, self.G, mode='same')
|
||||
# Scale by grid spacing
|
||||
numpy.multiply(p, numpy.prod(self.gdx), p)
|
||||
else:
|
||||
# For 3D, use numpy FFT (OpenCV DFT only supports 2D)
|
||||
# from scipy import signal
|
||||
p = signal.fftconvolve(f, self.G, mode='same')
|
||||
# Scale by grid spacing
|
||||
numpy.multiply(p, numpy.prod(self.gdx), p)
|
||||
|
||||
# convolution in spatial domain
|
||||
elif self.stype == 'S':
|
||||
raise SyntaxError('Spatial solver not suggested for structured grids')
|
||||
#p = (self.G).dot(f.reshape(-1)).reshape(fshape)
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# unstructured points
|
||||
elif self.ptype == 'P':
|
||||
|
||||
if (self.points.shape[0] != fshape[0]) or (1 != len(fshape)):
|
||||
raise ValueError("Shape of function should be (N,1) with N as the number of points")
|
||||
|
||||
# convolution in spatial domain
|
||||
# Dec 13, 2017 -- pre scaling works for unstructured grids!
|
||||
if self.stype == 'S':
|
||||
p = numpy.multiply(f, self.pvolumes)
|
||||
numpy.dot(self.G, p, p)
|
||||
#numpy.multiply(p, self.pvolumes, p)
|
||||
|
||||
elif self.stype == 'F':
|
||||
raise SyntaxError('Frequncy solver not supported for unstructured points')
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
if verbose:
|
||||
print( 'Poisson solver finished')
|
||||
gtimer.end()
|
||||
|
||||
return p
|
||||
|
||||
# --------------------------------------------------------------------------
|
||||
# GPU-accelerated solve using CuPy
|
||||
# --------------------------------------------------------------------------
|
||||
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.
|
||||
|
||||
Args:
|
||||
f: Input field (numpy array)
|
||||
verbose: Print timing information
|
||||
|
||||
Returns:
|
||||
p: Solution (numpy array on CPU)
|
||||
"""
|
||||
if not CUPY_AVAILABLE:
|
||||
if verbose:
|
||||
print("CuPy 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. on GPU')
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# regular grid
|
||||
if self.ptype == 'G':
|
||||
if self.gdims != fshape:
|
||||
print(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)
|
||||
}
|
||||
|
||||
# Transfer input to GPU
|
||||
f_gpu = cp.asarray(f, dtype=cp.float32)
|
||||
|
||||
# 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 verbose:
|
||||
print('Poisson solver (GPU) finished')
|
||||
gtimer.end()
|
||||
|
||||
# Transfer result back to CPU
|
||||
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')
|
||||
|
||||
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':
|
||||
# if verbose:
|
||||
# print(' - GPU solver not optimized for unstructured points, using CPU')
|
||||
# return self.solve(f, verbose)
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
|
||||
# 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
|
||||
|
||||
def _get_opencl_scatter_batch_kernel(self, queue):
|
||||
"""Batched scatter: copies two unpadded inputs into a stacked zero-padded device buffer."""
|
||||
if self._opencl_scatter_batch_kernel is not None:
|
||||
return self._opencl_scatter_batch_kernel
|
||||
|
||||
kernel_code = """
|
||||
__kernel void scatter_batch(
|
||||
__global const float* src,
|
||||
__global float* dst,
|
||||
const int src_width,
|
||||
const int src_height,
|
||||
const int dst_width,
|
||||
const int src_plane,
|
||||
const int dst_plane
|
||||
) {
|
||||
int gid_x = get_global_id(0);
|
||||
int gid_y = get_global_id(1);
|
||||
int gid_z = get_global_id(2);
|
||||
if (gid_x < src_width && gid_y < src_height) {
|
||||
int src_idx = gid_z * src_plane + gid_y * src_width + gid_x;
|
||||
int dst_idx = gid_z * dst_plane + gid_y * dst_width + gid_x;
|
||||
dst[dst_idx] = src[src_idx];
|
||||
}
|
||||
}
|
||||
"""
|
||||
program = cl.Program(self._opencl_context, kernel_code).build()
|
||||
self._opencl_scatter_batch_kernel = program.scatter_batch
|
||||
return self._opencl_scatter_batch_kernel
|
||||
|
||||
def _get_opencl_crop_batch_kernel(self, queue):
|
||||
"""Batched crop+scale: extracts 'same' region from two stacked full-conv results."""
|
||||
if self._opencl_crop_batch_kernel is not None:
|
||||
return self._opencl_crop_batch_kernel
|
||||
|
||||
kernel_code = """
|
||||
__kernel void crop_batch(
|
||||
__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,
|
||||
const int src_plane,
|
||||
const int dst_plane
|
||||
) {
|
||||
int gid_x = get_global_id(0);
|
||||
int gid_y = get_global_id(1);
|
||||
int gid_z = get_global_id(2);
|
||||
if (gid_x < dst_width && gid_y < dst_height) {
|
||||
int src_idx = gid_z * src_plane + (start_y + gid_y) * src_width + (start_x + gid_x);
|
||||
int dst_idx = gid_z * dst_plane + 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_batch_kernel = program.crop_batch
|
||||
return self._opencl_crop_batch_kernel
|
||||
|
||||
def _get_opencl_batch_mul_kernel(self, queue):
|
||||
"""Batched complex multiply: multiplies each batch plane by the same Green's spectrum.
|
||||
Uses multiple elements per work item, __constant for broadcast-read Green's spectrum,
|
||||
and fma for better throughput/precision."""
|
||||
if self._opencl_batch_mul_kernel is not None:
|
||||
return self._opencl_batch_mul_kernel
|
||||
|
||||
kernel_code = """
|
||||
__kernel void batch_complex_mul(
|
||||
__global float2* fft_batch,
|
||||
__global const float2* g,
|
||||
const int n
|
||||
) {
|
||||
const int b = get_global_id(0);
|
||||
const int base = b * n;
|
||||
const int i0 = get_global_id(1) * 4;
|
||||
const int i_end = min(i0 + 4, n);
|
||||
for (int i = i0; i < i_end; i++) {
|
||||
float2 gv = g[i];
|
||||
float2 a = fft_batch[base + i];
|
||||
fft_batch[base + i] = (float2)(fma(a.x, gv.x, -(a.y * gv.y)),
|
||||
fma(a.x, gv.y, a.y * gv.x));
|
||||
}
|
||||
}
|
||||
"""
|
||||
program = cl.Program(self._opencl_context, kernel_code).build()
|
||||
self._opencl_batch_mul_kernel = program.batch_complex_mul
|
||||
return self._opencl_batch_mul_kernel
|
||||
|
||||
def _get_opencl_alloc_host_ptr(self):
|
||||
"""Return an allocator that creates buffers with ALLOC_HOST_PTR for pinned host access (reduces H2D/D2H transfer cost)."""
|
||||
ctx = self._opencl_context
|
||||
mf = cl.mem_flags
|
||||
|
||||
def allocator(nbytes):
|
||||
return cl.Buffer(ctx, mf.READ_WRITE | mf.ALLOC_HOST_PTR, size=nbytes)
|
||||
|
||||
return allocator
|
||||
|
||||
@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_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:
|
||||
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()
|
||||
fshape = f1.shape
|
||||
full_shape = tuple(2 * s - 1 for s in fshape)
|
||||
timing['host_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
cache_valid = (
|
||||
self._opencl_cache is not None and
|
||||
self._opencl_cache['fshape'] == fshape and
|
||||
self._opencl_cache['full_shape'] == full_shape and
|
||||
'f_batch_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 = [fshape[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:
|
||||
fft_n_elements *= d
|
||||
src_plane = fshape[0] * fshape[1]
|
||||
dst_plane = full_shape[0] * full_shape[1]
|
||||
|
||||
# Pinned allocator for staging (H2D) and output (D2H) buffers
|
||||
alloc_pinned = self._get_opencl_alloc_host_ptr()
|
||||
|
||||
self._opencl_cache = {
|
||||
'fshape': fshape,
|
||||
'full_shape': full_shape,
|
||||
'fft_shape': fft_shape,
|
||||
'fft_n_elements': fft_n_elements,
|
||||
'g_fft_dev': g_fft_dev,
|
||||
'same_starts': same_starts,
|
||||
'gdx_prod': numpy.prod(self.gdx),
|
||||
'src_plane': src_plane,
|
||||
'dst_plane': dst_plane,
|
||||
'f_batch_dev': cla.zeros(queue, (2,) + full_shape, dtype=numpy.float32),
|
||||
'host_stage': numpy.empty((2,) + fshape, dtype=numpy.float32),
|
||||
'f_stage_dev': cla.empty(queue, (2,) + fshape, dtype=numpy.float32),
|
||||
'f_fft_batch_dev': cla.empty(queue, (2,) + fft_shape, dtype=numpy.complex64),
|
||||
'conv_batch_dev': cla.empty(queue, (2,) + full_shape, dtype=numpy.float32),
|
||||
'p_batch_dev': cla.empty(queue, (2,) + fshape, dtype=numpy.float32, allocator=alloc_pinned),
|
||||
'p1_host': numpy.empty(fshape, dtype=numpy.float32),
|
||||
'p2_host': numpy.empty(fshape, dtype=numpy.float32),
|
||||
}
|
||||
timing['cache_prepare_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
c = self._opencl_cache
|
||||
fshape_c = c['fshape']
|
||||
full_shape_c = c['full_shape']
|
||||
|
||||
# H2D: compact transfer of unpadded data, then scatter to padded buffer on device
|
||||
t0 = timeit.default_timer()
|
||||
H, W = fshape_c
|
||||
f1_f32 = f1 if (f1.flags['C_CONTIGUOUS'] and f1.dtype == numpy.float32) else numpy.ascontiguousarray(f1, dtype=numpy.float32)
|
||||
f2_f32 = f2 if (f2.flags['C_CONTIGUOUS'] and f2.dtype == numpy.float32) else numpy.ascontiguousarray(f2, dtype=numpy.float32)
|
||||
host_stage = c['host_stage']
|
||||
numpy.copyto(host_stage[0], f1_f32)
|
||||
numpy.copyto(host_stage[1], f2_f32)
|
||||
f_stage_dev = c['f_stage_dev']
|
||||
cl.enqueue_copy(queue, f_stage_dev.data, host_stage, is_blocking=False)
|
||||
c['f_batch_dev'].fill(0, queue=queue)
|
||||
scatter_batch = self._get_opencl_scatter_batch_kernel(queue)
|
||||
scatter_batch(
|
||||
queue, (W, H, 2), None,
|
||||
f_stage_dev.data, c['f_batch_dev'].data,
|
||||
numpy.int32(W), numpy.int32(H),
|
||||
numpy.int32(full_shape_c[1]),
|
||||
numpy.int32(H * W),
|
||||
numpy.int32(c['dst_plane']),
|
||||
)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['h2d_set_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# Single batched 2D FFT for both fields (ndim=2 treats dim-0 as batch)
|
||||
t0 = timeit.default_timer()
|
||||
fft_batch = cl_rfftn(c['f_batch_dev'], c['f_fft_batch_dev'], ndim=2)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['fft_f_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# Batched multiply: each batch plane * same Green's spectrum (4 elements/work item, explicit local size)
|
||||
t0 = timeit.default_timer()
|
||||
n_fft = c['fft_n_elements']
|
||||
global_n = max(1, (n_fft + 3) // 4)
|
||||
local_n = min(256, global_n)
|
||||
global_n = ((global_n + local_n - 1) // local_n) * local_n
|
||||
batch_mul = self._get_opencl_batch_mul_kernel(queue)
|
||||
batch_mul(
|
||||
queue, (2, global_n), (1, local_n),
|
||||
fft_batch.data,
|
||||
c['g_fft_dev'].data,
|
||||
numpy.int32(n_fft)
|
||||
)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['mul_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# Single batched 2D iFFT for both fields
|
||||
t0 = timeit.default_timer()
|
||||
conv_batch = cl_irfftn(fft_batch, c['conv_batch_dev'], ndim=2)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['ifft_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# Batch crop + scale on device
|
||||
t0 = timeit.default_timer()
|
||||
crop_batch = self._get_opencl_crop_batch_kernel(queue)
|
||||
same_starts = c['same_starts']
|
||||
crop_batch(
|
||||
queue, (fshape_c[1], fshape_c[0], 2), None,
|
||||
conv_batch.data, c['p_batch_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(c['gdx_prod']),
|
||||
numpy.int32(c['dst_plane']),
|
||||
numpy.int32(c['src_plane'])
|
||||
)
|
||||
if profile:
|
||||
queue.finish()
|
||||
timing['crop_scale_ms'] = (timeit.default_timer() - t0) * 1000.0
|
||||
|
||||
# D2H: map pinned buffer, copy into pre-allocated host arrays (zero allocation)
|
||||
t0 = timeit.default_timer()
|
||||
p1 = c['p1_host']
|
||||
p2 = c['p2_host']
|
||||
mapped_p = c['p_batch_dev'].map_to_host(queue=queue, flags=cl.map_flags.READ, is_blocking=True)
|
||||
numpy.copyto(p1, mapped_p[0])
|
||||
numpy.copyto(p2, mapped_p[1])
|
||||
mapped_p.base.release(queue=queue)
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user