''' Copyright (c) 2015, Harsh Bhatia (bhatia4@llnl.gov) All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. ''' import numpy import timeit import sys # import cv2 from scipy import signal, spatial from numba import njit from .backend_availability import ( CUPY_AVAILABLE, OPENCL_AVAILABLE, KERNEL_AVAILABLE, cp, cp_signal, _torque_kernel, cl, cla, cl_rfftn, cl_irfftn ) # -- Numba-accelerated finite-difference kernels for 2D HHD ------------------ # Matches numpy.gradient's second-order boundary stencil exactly. @njit(cache=True) def _divcurl_2d_nb(vfield, dy, dx, div_out, curl_out): """Fused divergence + curl of a 2D vector field via finite differences.""" H = vfield.shape[0] W = vfield.shape[1] idy2 = 1.0 / (2.0 * dy) idx2 = 1.0 / (2.0 * dx) idy_edge = 1.0 / (2.0 * dy) idx_edge = 1.0 / (2.0 * dx) for y in range(H): for x in range(W): # du/dx if x == 0: dudx = (-3.0*vfield[y,0,0] + 4.0*vfield[y,1,0] - vfield[y,2,0]) * idx_edge elif x == W - 1: dudx = (vfield[y,W-3,0] - 4.0*vfield[y,W-2,0] + 3.0*vfield[y,W-1,0]) * idx_edge else: dudx = (vfield[y,x+1,0] - vfield[y,x-1,0]) * idx2 # dv/dy if y == 0: dvdy = (-3.0*vfield[0,x,1] + 4.0*vfield[1,x,1] - vfield[2,x,1]) * idy_edge elif y == H - 1: dvdy = (vfield[H-3,x,1] - 4.0*vfield[H-2,x,1] + 3.0*vfield[H-1,x,1]) * idy_edge else: dvdy = (vfield[y+1,x,1] - vfield[y-1,x,1]) * idy2 # dv/dx if x == 0: dvdx = (-3.0*vfield[y,0,1] + 4.0*vfield[y,1,1] - vfield[y,2,1]) * idx_edge elif x == W - 1: dvdx = (vfield[y,W-3,1] - 4.0*vfield[y,W-2,1] + 3.0*vfield[y,W-1,1]) * idx_edge else: dvdx = (vfield[y,x+1,1] - vfield[y,x-1,1]) * idx2 # du/dy if y == 0: dudy = (-3.0*vfield[0,x,0] + 4.0*vfield[1,x,0] - vfield[2,x,0]) * idy_edge elif y == H - 1: dudy = (vfield[H-3,x,0] - 4.0*vfield[H-2,x,0] + 3.0*vfield[H-1,x,0]) * idy_edge else: dudy = (vfield[y+1,x,0] - vfield[y-1,x,0]) * idy2 div_out[y, x] = dudx + dvdy curl_out[y, x] = dvdx - dudy @njit(cache=True) def _grad_rotgrad_h_2d_nb(nD, nRu, vfield, dy, dx, d_out, r_out, h_out): """Fused gradient(nD) + rotated_gradient(nRu) + h = vfield - d - r.""" H = nD.shape[0] W = nD.shape[1] idy2 = 1.0 / (2.0 * dy) idx2 = 1.0 / (2.0 * dx) idy_edge = 1.0 / (2.0 * dy) idx_edge = 1.0 / (2.0 * dx) for y in range(H): for x in range(W): # gradient(nD) → d = (ddx_nD, ddy_nD) if x == 0: ddx_nD = (-3.0*nD[y,0] + 4.0*nD[y,1] - nD[y,2]) * idx_edge ddx_nRu = (-3.0*nRu[y,0] + 4.0*nRu[y,1] - nRu[y,2]) * idx_edge elif x == W - 1: ddx_nD = (nD[y,W-3] - 4.0*nD[y,W-2] + 3.0*nD[y,W-1]) * idx_edge ddx_nRu = (nRu[y,W-3] - 4.0*nRu[y,W-2] + 3.0*nRu[y,W-1]) * idx_edge else: ddx_nD = (nD[y,x+1] - nD[y,x-1]) * idx2 ddx_nRu = (nRu[y,x+1] - nRu[y,x-1]) * idx2 if y == 0: ddy_nD = (-3.0*nD[0,x] + 4.0*nD[1,x] - nD[2,x]) * idy_edge ddy_nRu = (-3.0*nRu[0,x] + 4.0*nRu[1,x] - nRu[2,x]) * idy_edge elif y == H - 1: ddy_nD = (nD[H-3,x] - 4.0*nD[H-2,x] + 3.0*nD[H-1,x]) * idy_edge ddy_nRu = (nRu[H-3,x] - 4.0*nRu[H-2,x] + 3.0*nRu[H-1,x]) * idy_edge else: ddy_nD = (nD[y+1,x] - nD[y-1,x]) * idy2 ddy_nRu = (nRu[y+1,x] - nRu[y-1,x]) * idy2 d0 = ddx_nD d1 = ddy_nD r0 = -ddy_nRu r1 = ddx_nRu d_out[y, x, 0] = d0 d_out[y, x, 1] = d1 r_out[y, x, 0] = r0 r_out[y, x, 1] = r1 h_out[y, x, 0] = vfield[y, x, 0] - d0 - r0 h_out[y, x, 1] = vfield[y, x, 1] - d1 - r1 # --------------------------------------------------------------------------- class nHHD(object): # constructor def __init__(self, cuda=True, **kwargs): # Select Poisson backend: 'auto' (default), 'cupy', 'opencl', or 'cpu' backend = kwargs.get('poisson_backend', 'auto').lower() if backend not in ('auto', 'cupy', 'opencl', 'cpu'): raise ValueError("poisson_backend must be one of: 'auto', 'cupy', 'opencl', 'cpu'") # Resolve 'auto' to a specific backend if backend == 'auto': backend = 'cupy' if CUPY_AVAILABLE else ('opencl' if OPENCL_AVAILABLE else 'cpu') # Validate backend availability and fallback to 'cpu' if needed if backend == 'cupy' and not CUPY_AVAILABLE: backend = 'cpu' elif backend == 'opencl' and not OPENCL_AVAILABLE: backend = 'cpu' self.debug = kwargs.get('debug', False) self.poisson_backend = backend # Initialize CUDA streams for cupy backend if self.poisson_backend == 'cupy': self.stream1 = cp.cuda.Stream() self.stream2 = cp.cuda.Stream() self.cuda = cuda args = kwargs.keys() if ('grid' in args) == ('points' in args): raise SyntaxError("nHHD 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("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, poisson_backend=self.poisson_backend) self.dim = self.mesh.dim if (self.dim != 2) and (self.dim != 3): raise ValueError("nHHD Solver works only for 2 and 3 dimensions") self._force_cpu_cache = None self._decompose_bufs = None self.psolver.prepare(False) def _ensure_decompose_bufs(self, shape2d): """Lazily allocate reusable buffers for the 2D decompose fast path.""" if self._decompose_bufs is not None and self._decompose_bufs['shape'] == shape2d: return self._decompose_bufs self._decompose_bufs = { 'shape': shape2d, 'div': numpy.empty(shape2d, dtype=numpy.float32), 'curl': numpy.empty(shape2d, dtype=numpy.float32), 'd': numpy.empty(shape2d + (2,), dtype=numpy.float32), 'r': numpy.empty(shape2d + (2,), dtype=numpy.float32), 'h': numpy.empty(shape2d + (2,), dtype=numpy.float32), } return self._decompose_bufs def get_force(self, vfield, verbose=False): if self.poisson_backend == 'cupy': return self.get_force_cuda(vfield, verbose) else: return self.get_force_cpu(vfield, verbose) def get_shear(self): if self.poisson_backend == 'cpu' or self.poisson_backend == 'opencl': return self.h elif self.poisson_backend == 'cupy': # Use CPU copy from get_force_cuda when set (stable snapshot; avoids jumps) if hasattr(self, 'h') and self.h is not None: return self.h return cp.asnumpy(self._h_gpu) def get_normal(self): if self.poisson_backend == 'cpu' or self.poisson_backend == 'opencl': return self.d elif self.poisson_backend == 'cupy': # Use CPU copy from get_force_cuda when set (stable snapshot; avoids jumps) if hasattr(self, 'd') and self.d is not None: return self.d return cp.asnumpy(self._d_gpu) def _solve_poisson_backend(self, field, verbose=False): """ Dispatch Poisson solve using the selected backend. """ if self.poisson_backend == 'opencl': return self.psolver.solve_opencl(field, verbose, profile=self.debug) if self.poisson_backend == 'cupy': return self.psolver.solve_gpu(field, verbose) if self.poisson_backend == 'cpu': return self.psolver.solve(field, verbose) # create the 3-component decomposition def decompose(self, vfield, verbose=0): # decompose for cpu and opencl backend if vfield.shape[-1] != self.dim: raise ValueError("nHHD.decompose requires a valid-dimensional vector field") if verbose > 0: print( '\n ==> nhhd.decompose', vfield.shape) if self.dim == 2: shape2d = (vfield.shape[0], vfield.shape[1]) bufs = self._ensure_decompose_bufs(shape2d) dy, dx = self.mesh.dx[0], self.mesh.dx[1] vf32 = vfield if vfield.dtype == numpy.float32 else vfield.astype(numpy.float32) _divcurl_2d_nb(vf32, dy, dx, bufs['div'], bufs['curl']) self.div = bufs['div'] self.curlw = bufs['curl'] if verbose > 0: print( 'div =', self.div.shape, self.div.min(), self.div.max()) print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max()) if self.poisson_backend == 'opencl' and ('OPENCL_AVAILABLE' in globals() and OPENCL_AVAILABLE): self.nD, self.nRu = self.psolver.solve_opencl_batch2(self.div, self.curlw, verbose > 0, profile=self.debug) else: self.nD = self._solve_poisson_backend(self.div, verbose > 0) self.nRu = self._solve_poisson_backend(self.curlw, verbose > 0) if verbose > 0: print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max()) print( 'nR =', self.nRu.shape, self.nRu.min(), self.nRu.max()) nD_f32 = self.nD if self.nD.dtype == numpy.float32 else self.nD.astype(numpy.float32) nRu_f32 = self.nRu if self.nRu.dtype == numpy.float32 else self.nRu.astype(numpy.float32) _grad_rotgrad_h_2d_nb(nD_f32, nRu_f32, vf32, dy, dx, bufs['d'], bufs['r'], bufs['h']) self.d = bufs['d'] self.r = bufs['r'] self.h = bufs['h'] if verbose > 0: print( 'd =', self.d.shape, self.d.min(), self.d.max()) print( 'r =', self.r.shape, self.r.min(), self.r.max()) print( 'h =', self.h.shape, self.h.min(), self.h.max()) return else: # compute div curl (self.div, self.curlu, self.curlv, self.curlw) = self.mesh.divcurl(vfield) if verbose > 0: print( 'div =', self.div.shape, self.div.min(), self.div.max()) print( 'curlu =', self.curlu.shape, self.curlu.min(), self.curlu.max()) print( 'curlv =', self.curlv.shape, self.curlv.min(), self.curlv.max()) print( 'curlw =', self.curlw.shape, self.curlw.min(), self.curlw.max()) # compute potentials self.nD = self._solve_poisson_backend(self.div, verbose > 0) if verbose > 0: print( 'nD =', self.nD.shape, self.nD.min(), self.nD.max()) self.nRu = self._solve_poisson_backend(self.curlu, verbose > 0) self.nRu *= -1.0 if verbose > 0: print( 'nRu =', self.nRu.shape, self.nRu.min(), self.nRu.max()) self.nRv = self._solve_poisson_backend(self.curlv, verbose > 0) self.nRv *= -1.0 if verbose > 0: print( 'nRv =', self.nRv.shape, self.nRv.min(), self.nRv.max()) self.nRw = self._solve_poisson_backend(self.curlw, verbose > 0) self.nRw *= -1.0 if verbose > 0: print( 'nRw =', self.nRw.shape, self.nRw.min(), self.nRw.max()) # compute divergent field as gradient of D potential self.d = self.mesh.gradient(self.nD) if verbose > 0: print( 'd =', self.d.shape, self.d.min(), self.d.max()) # temporary use of r self.r = numpy.stack((self.nRu, self.nRv, self.nRw), axis = -1) # compute rotational field as curl of R potential (cu, cv, cw) = self.mesh.curl3D(self.r) self.r = numpy.stack((cu, cv, cw), axis = -1) if verbose > 0: print( 'r =', self.r.shape, self.r.min(), self.r.max()) self.h = numpy.add(self.d, self.r) numpy.subtract(vfield, self.h, self.h) if verbose > 0: print( 'h =', self.h.shape, self.h.min(), self.h.max()) def decompose_cuda(self, vfield, verbose=0): """ GPU-optimized decomposition that keeps intermediate results on GPU. Returns GPU arrays to avoid CPU-GPU transfers. Args: vfield: Vector field (numpy array or CuPy array) verbose: Verbosity level (default: 0) Returns: Tuple of GPU arrays: (d_gpu, h_gpu, r_gpu, nRu_gpu) for 2D fields """ if not CUPY_AVAILABLE: if verbose > 0: print("CuPy not available, falling back to CPU decompose") self.decompose(vfield, verbose) return None if vfield.shape[-1] != self.dim: raise ValueError("nHHD.decompose requires a valid-dimensional vector field") if verbose > 0: print('\n ==> nhhd.decompose_cuda', vfield.shape) # Convert input to GPU if needed if isinstance(vfield, numpy.ndarray): vfield_gpu = cp.asarray(vfield, dtype=cp.float32) else: vfield_gpu = vfield if self.dim == 2: # Compute div curl on GPU dx0, dx1 = self.mesh.dx[0], self.mesh.dx[1] dudy, dudx = cp.gradient(vfield_gpu[:,:,0], dx0, dx1) dvdy, dvdx = cp.gradient(vfield_gpu[:,:,1], dx0, dx1) 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())) # Use a single stream for both solves. Running concurrent FFT-based # solves on separate streams can cause occasional long-run jitter. nD_gpu = self.psolver._solve_gpu_internal(div_gpu) nRu_gpu = self.psolver._solve_gpu_internal(curlw_gpu) cp.cuda.get_current_stream().synchronize() if verbose > 0: print('nD =', nD_gpu.shape, float(nD_gpu.min()), float(nD_gpu.max())) print('nR =', nRu_gpu.shape, float(nRu_gpu.min()), float(nRu_gpu.max())) # Compute gradients on GPU ddy, ddx = cp.gradient(nD_gpu, dx0, dx1) d_gpu = cp.stack((ddx, ddy), axis=-1) if verbose > 0: print('d =', d_gpu.shape, float(d_gpu.min()), float(d_gpu.max())) # Compute rotated gradient on GPU rdy, rdx = cp.gradient(nRu_gpu, dx0, dx1) 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 = 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())) # Ensure all default-stream work (gradient, etc.) is complete before use cp.cuda.get_current_stream().synchronize() # Store GPU arrays self._d_gpu = d_gpu self._h_gpu = h_gpu self._r_gpu = r_gpu 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 = cp.asnumpy(nD_gpu) return (d_gpu, h_gpu, r_gpu, nRu_gpu) else: # For 3D, fall back to CPU for now if verbose > 0: print("3D GPU decompose not yet implemented, using CPU") self.decompose(vfield, verbose) return None def get_force_cpu(self, vfield, verbose=False): shape = (vfield.shape[0], vfield.shape[1]) if (not hasattr(self, '_force_cpu_cache') or self._force_cpu_cache is None or self._force_cpu_cache['shape'] != shape): self._init_force_cpu_cache(shape) grid = self._force_cpu_cache['grid'] self.decompose(vfield, verbose) # Optimized: use abs instead of absolute (slightly faster) self.Sn = numpy.sum(numpy.abs(self.d)) # Optimized: direct indexing and sum (already efficient) self.Stx = numpy.sum(self.h[:, :, 0]) self.Sty = numpy.sum(self.h[:, :, 1]) # Optimized: find both max and min indices in single pass through flattened array nRu_flat = self.nRu.ravel() nR_max_idx = numpy.argmax(nRu_flat) nR_min_idx = numpy.argmin(nRu_flat) # Convert flat indices to 2D coordinates (faster than unravel_index) nR_max_loc = (nR_max_idx // shape[1], nR_max_idx % shape[1]) nR_min_loc = (nR_min_idx // shape[1], nR_min_idx % shape[1]) # Optimized: compute cross products directly without creating intermediate arrays # For 2D: cross(a, b) = a[0]*b[1] - a[1]*b[0] # Extract components once for reuse r_x = self.r[:, :, 0] r_y = self.r[:, :, 1] grid_x = grid[:, :, 0] grid_y = grid[:, :, 1] # Compute cross products directly (faster than numpy.cross for 2D) # cross(r_p, r) where r_p = grid - nR_max_loc max_y, max_x = nR_max_loc[0], nR_max_loc[1] cross_p = (grid_x - max_x) * r_y - (grid_y - max_y) * r_x # cross(r_n, r) where r_n = grid - nR_min_loc min_y, min_x = nR_min_loc[0], nR_min_loc[1] cross_n = (grid_x - min_x) * r_y - (grid_y - min_y) * r_x self.Stor = numpy.sum(cross_p) + numpy.sum(cross_n) # Optimized: pre-allocate output array with explicit dtype output = numpy.array([self.Sn, self.Stx, self.Sty, self.Stor], dtype=numpy.float32) # Optimized: combine threshold checks (short-circuit evaluation) flag = (self.Sn >= 300 and self.Stx >= 300 and self.Sty >= 300 and numpy.abs(self.Stor) >= 1000) return flag, output # 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): """ Initialize cached GPU arrays for force computation. Call once when shape is known to avoid repeated allocations. Args: shape: Tuple (height, width) of the vector field """ if not CUPY_AVAILABLE: return # Cache grid arrays (only need to create once per shape) self._force_gpu_cache = { 'shape': shape, 'mgrid_0': cp.arange(shape[0], dtype=cp.float32)[:, None], 'mgrid_1': cp.arange(shape[1], dtype=cp.float32)[None, :], } # Pre-broadcast for efficiency self._force_gpu_cache['mgrid_0_bc'] = cp.broadcast_to( self._force_gpu_cache['mgrid_0'], shape ).copy() # .copy() to make contiguous self._force_gpu_cache['mgrid_1_bc'] = cp.broadcast_to( self._force_gpu_cache['mgrid_1'], shape ).copy() def get_force_cuda(self, vfield, verbose=False): """ Further optimized GPU force computation using CuPy ElementwiseKernel. Fuses all operations into a single kernel for maximum performance. Args: vfield: Vector field (numpy array) verbose: Print debug information (default: False) Returns: tuple: (flag, output) where: - flag: Boolean indicating if force thresholds are met - output: Array [Sn, Stx, Sty, Stor] containing force components """ if not CUPY_AVAILABLE: if verbose: print("CuPy not available, falling back to CPU get_force") return self.get_force(vfield, verbose) # Decompose the vector field self.decompose_cuda(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 = 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 as L1 sum directly on GPU (faster than linalg.norm for this case) Sn = float(cp.sum(cp.abs(d_gpu))) # Compute Stx, Sty Stx = float(cp.sum(h_gpu[:, :, 0])) Sty = float(cp.sum(h_gpu[:, :, 1])) # Find extrema indices nR_max_idx = int(cp.argmax(nRu_gpu)) nR_min_idx = int(cp.argmin(nRu_gpu)) shape = nRu_gpu.shape max_0, max_1 = nR_max_idx // shape[1], nR_max_idx % shape[1] min_0, min_1 = nR_min_idx // shape[1], nR_min_idx % shape[1] # Initialize cache if needed if not hasattr(self, '_force_gpu_cache') or self._force_gpu_cache is None or self._force_gpu_cache['shape'] != shape: self._init_force_gpu_cache(shape) mgrid_0 = self._force_gpu_cache['mgrid_0_bc'] mgrid_1 = self._force_gpu_cache['mgrid_1_bc'] r_x = r_gpu[:, :, 0] 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 # )) # 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) # Keep stable CPU copies for get_normal()/get_shear() so callers get this frame's # result instead of reading GPU memory that can change (avoids CuPy-only jumps) self.d = cp.asnumpy(self._d_gpu) self.h = cp.asnumpy(self._h_gpu) flag = (Sn >= 300 and Stx >= 300 and Sty >= 300 and abs(Stor) >= 1000) return flag, output def clear_force_gpu_cache(self): """Clear the GPU cache for force computation to free memory.""" if hasattr(self, '_force_gpu_cache'): self._force_gpu_cache = None if CUPY_AVAILABLE: cp.get_default_memory_pool().free_all_blocks() # --------------------------------------------------------------------------- class Timer(object): def display(self): tseconds = self.etime - self.stime print(' [[ elapsed time:', end = "") if tseconds < numpy.power(10.,-6.): print("%.3f micro-sec." % (tseconds*numpy.power(10,6)),end = "") elif tseconds < numpy.power(10.,-3.): print("%.3f milli-sec." % (tseconds*numpy.power(10.,3.)),end = "") elif tseconds < 60.0: print("%.3f sec." % (tseconds),end = "") else: m = int(tseconds/ 60.0) print("%d min. %.3f sec." % (m, (tseconds - 60*m)),end = "") print(']]') def start(self): self.stime = timeit.default_timer() def end(self, print_time=True): self.etime = timeit.default_timer() if print_time: self.display() def __init__(self): self.start() # --------------------------------------------------------------------------- # Try to import CuPy for GPU acceleration # try: # import cupy as cp # CUPY_AVAILABLE = True # except ImportError: # cp = None # CUPY_AVAILABLE = False # ------------------------------------------------------------------------------ class StructuredGrid(object): # -------------------------------------------------------------------------- # -------------------------------------------------------------------------- def __init__(self, **kwargs): ''' kwargs: grid: ndarray of grid dimensions (Y,X) or (Z,Y,X) spacings: ndarray of grid spacings (dy, dx) or (dz, dy, dx) verbose: verbosity level ''' args = kwargs.keys() if ('grid' not in args) or ('spacings' not in args): raise SyntaxError("Dimensions and spacings of the grid are required") verbose = 0 if 'verbose' in args: verbose = kwargs['verbose'] self.dims = kwargs['grid'] self.dx = kwargs['spacings'] self.dim = len(self.dims) if self.dim != 2 and self.dim != 3: raise ValueError("StructuredGrid works for 2D and 3D only") if self.dim != len(self.dx): raise ValueError("Dimensions of spacings should match that of the grid") # Initialize GPU cache self._gpu_divcurl_cache = None if verbose > 0: print(' Initializing', self.dim, 'D structured grid...', end ='') sys.stdout.flush() mtimer = Timer() print(' Done!', end ='') mtimer.end() # -------------------------------------------------------------------------- # -------------------------------------------------------------------------- def divcurl(self, vfield, verbose=False): # compute div and curl of vfield (vector field) #if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] - self.dims).any(): 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...', end ='') sys.stdout.flush() mtimer = Timer() if self.dim == 2: # self.dx = (dy,dx) dudy, dudx = numpy.gradient(vfield[:,:,0], self.dx[0], self.dx[1]) dvdy, dvdx = numpy.gradient(vfield[:,:,1], self.dx[0], self.dx[1]) numpy.add(dudx, dvdy, dudx) numpy.subtract(dvdx, dudy, dvdx) if verbose: print(' Done!', end ='') mtimer.end() return (dudx, dvdx) elif self.dim == 3: # self.dx = (dz,dy,dx) dudz, dudy, dudx = numpy.gradient(vfield[:,:,:,0], self.dx[0], self.dx[1], self.dx[2]) dvdz, dvdy, dvdx = numpy.gradient(vfield[:,:,:,1], self.dx[0], self.dx[1], self.dx[2]) dwdz, dwdy, dwdx = numpy.gradient(vfield[:,:,:,2], self.dx[0], self.dx[1], self.dx[2]) numpy.add(dudx, dvdy, dudx) numpy.add(dudx, dvdz, dudx) numpy.subtract(dwdy, dvdz, dwdy) numpy.subtract(dudz, dwdx, dudz) numpy.subtract(dvdx, dudy, dvdx) if verbose: print(' Done!', end ='') mtimer.end() 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) # 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 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 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) # if verbose: # print(' Done!', end ='') # mtimer.end() # return (div, curl) # 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 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) # if verbose: # print(' Done!', end ='') # mtimer.end() # return (div, curlu, curlv, curlw) def curl3D(self, vfield, verbose=False): #if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] - self.dims).any(): 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 self.dim != 3: raise ValueError("curl3D works only for 2D") if verbose: print(' Computing curl...', end ='') sys.stdout.flush() mtimer = Timer() # self.dx = (dz,dy,dx) dudz, dudy, dudx = numpy.gradient(vfield[:,:,:,0], self.dx[0], self.dx[1], self.dx[2]) dvdz, dvdy, dvdx = numpy.gradient(vfield[:,:,:,1], self.dx[0], self.dx[1], self.dx[2]) dwdz, dwdy, dwdx = numpy.gradient(vfield[:,:,:,2], self.dx[0], self.dx[1], self.dx[2]) numpy.subtract(dwdy, dvdz, dwdy) numpy.subtract(dudz, dwdx, dudz) numpy.subtract(dvdx, dudy, dvdx) if verbose: print(' Done!', end ='') mtimer.end() return (dwdy, dudz, dvdx) def curl3D_gpu(self, vfield, verbose=False): """ GPU-accelerated computation of curl for 3D 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.curl3D(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 self.dim != 3: raise ValueError("curl3D works only for 3D") if verbose: print(' Computing curl on GPU...', end ='') sys.stdout.flush() mtimer = Timer() # 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 curl components on GPU curl_u = dwdy - dvdz curl_v = dudz - dwdx curl_w = dvdx - dudy # Transfer back to CPU curlu = cp.asnumpy(curl_u) curlv = cp.asnumpy(curl_v) curlw = cp.asnumpy(curl_w) if verbose: print(' Done!', end ='') mtimer.end() return (curlu, curlv, curlw) def rotated_gradient(self, sfield, verbose=False): if (sfield.shape != self.dims): #if (sfield.shape - self.dims).any(): raise ValueError("Dimensions of scalar field should match that of the grid") if self.dim != 2: raise ValueError("rotated_gradient works only for 2D") if verbose: print(' Computing rotated gradient...', end ='') sys.stdout.flush() mtimer = Timer() ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1]) ddy *= -1.0 grad = numpy.stack((ddy, ddx), axis=-1) if verbose: print(' Done!', end ='') mtimer.end() return grad def rotated_gradient_gpu(self, sfield, verbose=False): """ GPU-accelerated computation of rotated gradient for 2D scalar 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.rotated_gradient(sfield, verbose) if (sfield.shape != self.dims): raise ValueError("Dimensions of scalar field should match that of the grid") if self.dim != 2: raise ValueError("rotated_gradient works only for 2D") if verbose: print(' Computing rotated gradient on GPU...', end ='') sys.stdout.flush() mtimer = Timer() # Transfer to GPU sfield_gpu = cp.asarray(sfield, dtype=cp.float32) # Compute gradients on GPU ddy, ddx = cp.gradient(sfield_gpu, self.dx[0], self.dx[1]) ddy *= -1.0 # Stack and transfer back to CPU grad_gpu = cp.stack((ddy, ddx), axis=-1) grad = cp.asnumpy(grad_gpu) if verbose: print(' Done!', end ='') mtimer.end() return grad def gradient(self, sfield, verbose=False): if (sfield.shape != self.dims): #if (sfield.shape - self.dims).any(): raise ValueError("Dimensions of scalar field should match that of the grid") if verbose: print(' Computing gradient...', end ='') sys.stdout.flush() mtimer = Timer() if self.dim == 2: # self.dx = (dy,dx) ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1]) grad = numpy.stack((ddx, ddy), axis = -1) elif self.dim == 3: # self.dx = (dz,dy,dx) ddz, ddy, ddx = numpy.gradient(sfield, self.dx[0], self.dx[1], self.dx[2]) grad = numpy.stack((ddx, ddy, ddz), axis = -1) if verbose: print(' Done!', end ='') mtimer.end() return grad def gradient_gpu(self, sfield, verbose=False): """ GPU-accelerated computation of gradient for scalar 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.gradient(sfield, verbose) if (sfield.shape != self.dims): raise ValueError("Dimensions of scalar field should match that of the grid") if verbose: print(' Computing gradient on GPU...', end ='') sys.stdout.flush() mtimer = Timer() # Transfer to GPU sfield_gpu = cp.asarray(sfield, dtype=cp.float32) if self.dim == 2: # Compute gradients on GPU ddy, ddx = cp.gradient(sfield_gpu, self.dx[0], self.dx[1]) grad_gpu = cp.stack((ddx, ddy), axis=-1) elif self.dim == 3: # Compute gradients on GPU ddz, ddy, ddx = cp.gradient(sfield_gpu, self.dx[0], self.dx[1], self.dx[2]) grad_gpu = cp.stack((ddx, ddy, ddz), axis=-1) # Transfer back to CPU grad = cp.asnumpy(grad_gpu) if verbose: print(' Done!', end ='') mtimer.end() return grad def clear_gpu_cache(self): """Clear GPU cache to free GPU memory""" if self._gpu_divcurl_cache is not None: self._gpu_divcurl_cache = None if CUPY_AVAILABLE: cp.get_default_memory_pool().free_all_blocks() # -------------------------------------------------------------------------- # -------------------------------------------------------------------------- # import numpy # from scipy import signal, spatial # # from .timer import Timer # # Try to import CuPy for GPU acceleration # try: # import cupy as cp # import cupyx.scipy.signal as cp_signal # CUPY_AVAILABLE = True # except ImportError: # cp = None # cp_signal = None # CUPY_AVAILABLE = False # Try to import OpenCL + pyvkfft for FFT-based convolution ''' # ------------------------------------------------------------------------------ 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 # def solve_opencl(self, f, verbose=False, profile=True): # """ # 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 and # 'f_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 = [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) # ) # 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() # 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)