diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index fbfe99639..ba8fb9d3d 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -12,7 +12,7 @@ The work is split into small PRs that can be reviewed and merged one at a time. - [x] **PR 4: `Pusher` accepts `Kernel`** — the kernel for the active backend is chosen once, when the pusher is created; a plain `PyccelKernel` is wrapped, so the propagators do not change (no behaviour change on CPU). - [x] **PR 5: `Domain` on the GPU** — domain arguments are selected and stored at domain construction; CUDA arguments reference device arrays, and deepcopy/unpickling rebuilds the arguments from the copied or restored arrays. - [x] **PR 6: `Particles` on the GPU** — `Particles` can be created on the CuPy backend, and `Particles.args_markers` is selected as the CUDA or Pyccel argument bundle at construction. -- [ ] **PR 7: `Derham` on the GPU** — `Derham` can be created on the CuPy backend, plus `Derham.cuda_args_derham`. +- [x] **PR 7: `Derham` on the GPU** — `Derham` can be created on the CuPy backend, plus `Derham.cuda_args_derham`. - [ ] **PR 8: Shared CUDA headers for the argument classes** — one `.cuh` per argument class instead of long flat kernel signatures. - [ ] **PR 9: One folder per kernel, starting with `pic/pushing`** — pure refactor, no behaviour change. - [ ] **PR 10: Device versions of helper kernels** — B-spline evaluation, mapping evaluation (per domain), small linear algebra, as `__device__` functions in `.cuh` headers. @@ -41,12 +41,12 @@ CUDA kernels can be added one by one. If the code runs on the GPU and needs a ke - **No silent CPU fallback on the GPU.** A kernel without a CUDA version raises an error on the GPU backend. Falling back would mean copying data to the host and back at every call. - **Small steps.** Every PR keeps the CPU code path working and tested. -## Current state (PR 6) +## Current state (PR 7) | File | Content | |---|---| | `src/struphy/utils/kernel_backends.py` | `is_cuda_backend()`, `CudaKernel` (wraps a `cupy.RawKernel`, compiled lazily; expands `Argument.get_cuda_args()` and takes `n_threads`), `Kernel` and `KernelCatalog` for backend selection and discovery | -| `src/struphy/utils/cuda_arguments.py` | `Argument` contract plus `CudaMarkerArguments` and `CudaDomainArguments`; CUDA arrays are stored individually and returned in signature order by `get_cuda_args()` | +| `src/struphy/utils/cuda_arguments.py` | `Argument` contract plus `CudaMarkerArguments`, `CudaDerhamArguments` and `CudaDomainArguments`; CUDA arrays are stored individually and returned in signature order by `get_cuda_args()` | | `src/struphy/geometry/base.py` | `Domain.args_domain` is selected once at construction; CUDA domains use device arrays, while direct Pyccel geometry calls retain a host argument bundle | | `src/struphy/pic/base.py` | `Particles` arrays and `args_markers` use the backend selected at construction; direct Pyccel methods retain a private host bundle | | `src/struphy/pic/tests/test_kernel_backends.py` | the demo kernel pair `push_eta_linear` (pyccel function compiled with `epyccel` at test time, CUDA source string) and tests on both backends | @@ -126,7 +126,7 @@ kernel = catalog["push_eta_stage"] # Kernel: pyccel or CUDA depending on the ba - The CUDA argument objects hold references. If an owner reallocates an array (today the markers are allocated once), it must rebuild its CUDA arguments at the same place, exactly like for the pyccel arguments. - First these classes must be creatable on the CuPy backend at all: - `Particles`: wrap Python lists in `xp.array` before reductions, and use host buffers for scalar MPI gathers (`pic/base.py`). - - `Derham`: NumPy arrays from feectools reach `cupy.ascontiguousarray`. + - `Derham`: feectools and struphy moved host data (knots, grids, collocation matrices) to CuPy before calling pyccel kernels (see PR 7 below). - `Domain`: deepcopy and unpickling on CuPy failed (see PR 5 below). ### PR 5: `Domain` on the GPU (complete) @@ -142,7 +142,19 @@ kernel = catalog["push_eta_stage"] # Kernel: pyccel or CUDA depending on the ba - Particle arrays, validity masks, and boundary-condition codes are allocated through `cunumpy`, so they live on CuPy when the CuPy backend is active. - `args_markers` is built as `CudaMarkerArguments` from device arrays on CuPy, or as `MarkerArguments` on NumPy. A private host bundle remains for direct Pyccel calls. - Domain decomposition now wraps the Python `nprocs` list with `xp.array` before calling `xp.prod`. Scalar MPI gathers use small NumPy buffers and copy the results back to the active array backend, avoiding unsupported CuPy buffers in MPI calls. -- Full GPU particle pushing still depends on CUDA versions of the required kernels and on PR 7's `Derham` support. +- Full GPU particle pushing still depends on CUDA versions of the required kernels. + +### PR 7: `Derham` on the GPU (complete) + +- feectools and the struphy code that builds `Derham` followed `xp` everywhere. So on CuPy, the knots, quadrature grids and decomposition metadata became device arrays and then reached pyccel kernels, SciPy or MPI, which only take host arrays. + - feectools: struphy-hub/feectools#85, the first of the feectools CUDA PRs, makes feectools run on the CuPy backend. It must be merged, and the submodule or the feectools version bumped, before `Derham` can be created on CuPy. + - struphy: data that describes the spline spaces is host data on every backend. Only the coefficients (`StencilVector` data) and stencil matrices live on the device. The projection and quadrature grids of `Derham` (`get_pts_and_wts`, ...) and `spline_types_pyccel` are NumPy. `domain_array`, `index_array(_N/_D)` and `neighbours` are gathered with NumPy MPI buffers and then converted with `xp.asarray`, so they are device arrays on CuPy like `Particles.domain_array`. +- `Derham.args_derham` is built from the host knots, degrees and starts on both backends. `Derham.cuda_args_derham` lazily builds `CudaDerhamArguments` with one device copy of these small arrays. On NumPy it raises (host arrays are never copied to the device). The pyccel scratch arrays (`bn1`, ..., `bd3`) are not part of it; they become per-thread local arrays in CUDA (PR 10). +- Not supported on CuPy yet: + - Local projectors (`DerhamOptions.local_projectors=True`) raise `NotImplementedError` when the `Derham` is created. `CommutingProjectorLocal` builds its data with `xp` and calls pyccel kernels on it, like `Derham` did. + - Polar splines need a spline mapping, which cannot be created on CuPy yet (see PR 5). + - Field evaluation (`SplineFunction.__call__`, ...) still calls pyccel kernels with the coefficients, which are device arrays on CuPy. It needs CUDA evaluation kernels (PR 10+). +- Tests: `feec/tests/test_derham_gpu.py`. Without a GPU, a strict host stand-in for CuPy (rejects host/device mixing, cannot run kernels; not part of the repository) was used, on top of struphy-hub/feectools#85. With it, a `Derham` created on the "CuPy" backend matches the NumPy one on 1, 2 and 4 MPI processes. ### PR 8: Argument structs in shared headers diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index e73f4cdf2..e16af24c5 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -44,6 +44,8 @@ from struphy.polar.extraction_operators import PolarExtractionBlocksC1 from struphy.polar.linear_operators import PolarExtractionOperator, PolarLinearOperator from struphy.topology.grids import TensorProductGrid +from struphy.utils.cuda_arguments import CudaDerhamArguments +from struphy.utils.kernel_backends import is_cuda_backend NonTrivialBC = LiteralOptions.OptsNonTrivialBoundaryCondition space_to_form = { @@ -57,14 +59,6 @@ logger = logging.getLogger("struphy") -def _to_numpy_for_kernel(value): - """Convert CuPy arrays to NumPy for compiled kernel calls.""" - if hasattr(value, "get"): - # This is a CuPy array - return value.get() - return value - - class DiscreteDerham: """Discrete 3D de Rham sequence built from four FE spaces. @@ -423,7 +417,8 @@ def __init__( self._quad_grid_spans[-1] = tuple(self.quad_grid_spans[-1]) self._quad_grid_bases[-1] = tuple(self.quad_grid_bases[-1]) - self._spline_types_pyccel[-1] = xp.array( + # inputs of pyccel kernels, kept on the host + self._spline_types_pyccel[-1] = np.array( self._spline_types_pyccel[-1], ) @@ -485,7 +480,7 @@ def spline_types(self) -> tuple[tuple[str]]: return self._spline_types @property - def spline_types_pyccel(self) -> tuple[xp.ndarray]: + def spline_types_pyccel(self) -> tuple[np.ndarray]: """Tuple of spline types in each direction as integers (0 for 'B', 1 for 'M') for each component of the vector space.""" return self._spline_types_pyccel @@ -608,6 +603,10 @@ def __init__( polar_splines = options.polar_splines # local commuting projectors local_projectors = options.local_projectors + if local_projectors and is_cuda_backend(): + raise NotImplementedError( + "Local projectors (DerhamOptions.local_projectors=True) are not supported on the CuPy backend yet." + ) # number of elements and spline degrees in each direction assert len(num_elements) == 3 @@ -899,14 +898,20 @@ def __init__( self._neighbours = self._get_neighbours() - # collect arguments for kernels - self._args_derham = DerhamArguments( - _to_numpy_for_kernel(xp.array(self.degree)), - _to_numpy_for_kernel(self.V0fem.knots[0]), - _to_numpy_for_kernel(self.V0fem.knots[1]), - _to_numpy_for_kernel(self.V0fem.knots[2]), - _to_numpy_for_kernel(xp.array(self.V0.starts)), + # collect arguments for kernels (the knots of feectools are host arrays on every array backend) + self._pyccel_args_derham = DerhamArguments( + np.array(self.degree), + *self.V0fem.knots, + np.array(self.V0.starts), ) + if is_cuda_backend(): + self._args_derham = CudaDerhamArguments( + xp.asarray(self._pyccel_args_derham.pn), + *(xp.asarray(t) for t in self.V0fem.knots), + xp.asarray(self._pyccel_args_derham.starts), + ) + else: + self._args_derham = self._pyccel_args_derham logger.debug("\nDERHAM:") logger.debug(f"{'number of elements:'.ljust(25)} {num_elements}") @@ -1477,8 +1482,8 @@ def div_bcfree(self): return self._div_bcfree @property - def args_derham(self): - """Collection of mandatory arguments for pusher kernels.""" + def args_derham(self) -> DerhamArguments | CudaDerhamArguments: + """Mandatory pusher kernel arguments for the backend used at initialization.""" return self._args_derham # -------------------------- @@ -1778,7 +1783,7 @@ def _discretize_space( ) # Create uniform grid - grids = [xp.linspace(xmin, xmax, num=ne + 1) for xmin, xmax, ne in zip(min_coords, max_coords, ncells)] + grids = [np.linspace(xmin, xmax, num=ne + 1) for xmin, xmax, ne in zip(min_coords, max_coords, ncells)] # Create 1D finite element spaces and precompute quadrature data spaces_1d = [ @@ -1925,11 +1930,11 @@ def _get_domain_array(self): else: nproc = 1 - # send buffer - dom_arr_loc = xp.zeros(9, dtype=float) + # send buffer (MPI buffers are host arrays on every array backend) + dom_arr_loc = np.zeros(9, dtype=float) # main array (receive buffers) - dom_arr = xp.zeros(nproc * 9, dtype=float) + dom_arr = np.zeros(nproc * 9, dtype=float) # Get global starts and ends of domain decomposition gl_s = self.domain_decomposition.starts @@ -1947,7 +1952,7 @@ def _get_domain_array(self): else: dom_arr[:] = dom_arr_loc - return dom_arr.reshape(nproc, 9) + return xp.asarray(dom_arr.reshape(nproc, 9)) def _get_index_array(self, decomposition): """ @@ -1972,11 +1977,11 @@ def _get_index_array(self, decomposition): else: nproc = 1 - # send buffer - ind_arr_loc = xp.zeros(6, dtype=int) + # send buffer (MPI buffers are host arrays on every array backend) + ind_arr_loc = np.zeros(6, dtype=int) # main array (receive buffers) - ind_arr = xp.zeros(nproc * 6, dtype=int) + ind_arr = np.zeros(nproc * 6, dtype=int) # Get global starts and ends of cart OR domain decomposition gl_s = decomposition.starts @@ -1993,7 +1998,7 @@ def _get_index_array(self, decomposition): else: ind_arr[:] = ind_arr_loc - return ind_arr.reshape(nproc, 6) + return xp.asarray(ind_arr.reshape(nproc, 6)) def _get_neighbours(self): """ @@ -2025,7 +2030,7 @@ def _get_neighbours(self): neighbours along the edges only have one 1, neighbours along the edges have no 1 in the index. """ - neighs = xp.empty((3, 3, 3), dtype=int) + neighs = np.empty((3, 3, 3), dtype=int) for i in range(3): for j in range(3): @@ -2034,7 +2039,7 @@ def _get_neighbours(self): ind = tuple(comp) neighs[ind] = self._get_neighbour_one_component(comp) - return neighs + return xp.asarray(neighs) def _get_neighbour_one_component(self, comp): """ @@ -2070,8 +2075,9 @@ def _get_neighbour_one_component(self, comp): if comp == [1, 1, 1]: return neigh_id - comp = xp.array(comp) - kinds = xp.array(kinds) + # computed on the host: the start/end indices below are an object array (with None entries) + comp = np.array(comp) + kinds = np.array(kinds) # if only one process: check if comp is neighbour in non-peridic directions, if this is not the case then return the rank as neighbour id if size == 1: @@ -2084,12 +2090,13 @@ def _get_neighbour_one_component(self, comp): # elements with index 2n are the starts and 2n + 1 are the ends. neigh_inds = [None] * 6 + index_array = xp.to_numpy(self.index_array) # in each direction find start/end index for neighbour for k, co in enumerate(comp): if co == 1: - neigh_inds[2 * k + 0] = self.index_array[rank, 2 * k + 0] - neigh_inds[2 * k + 1] = self.index_array[rank, 2 * k + 1] + neigh_inds[2 * k + 0] = index_array[rank, 2 * k + 0] + neigh_inds[2 * k + 1] = index_array[rank, 2 * k + 1] elif co == 0: neigh_inds[2 * k + 1] = gl_s[k] - 1 @@ -2106,15 +2113,15 @@ def _get_neighbour_one_component(self, comp): "Wrong value for component; must be 0 or 1 or 2 !", ) - neigh_inds = xp.array(neigh_inds) + neigh_inds = np.array(neigh_inds) # only use indices where information is present to find the neighbours rank - inds = xp.where(xp.not_equal(neigh_inds, None)) + inds = np.where(np.not_equal(neigh_inds, None)) # find ranks (row index of domain_array) which agree in start/end indices - index_temp = xp.squeeze(self.index_array[:, inds]) - unique_ranks = xp.where( - xp.equal(index_temp, neigh_inds[inds]).all(1), + index_temp = np.squeeze(index_array[:, inds]) + unique_ranks = np.where( + np.equal(index_temp, neigh_inds[inds]).all(1), )[0] # if any row satisfies condition, return its index (=rank of neighbour) @@ -3493,18 +3500,16 @@ def get_pts_and_wts(space_1d, start, end, n_quad=None, polar_shift=False): histopol_loc = space_1d.histopolation_grid[start : end + 2].copy() # make sure that greville points used for interpolation are in [0, 1] - # Use numpy for comparison since greville points are NumPy arrays - greville_loc_np = greville_loc.get() if hasattr(greville_loc, "get") else greville_loc - assert np.all(np.logical_and(greville_loc_np >= 0.0, greville_loc_np <= 1.0)) + assert np.all(np.logical_and(greville_loc >= 0.0, greville_loc <= 1.0)) # interpolation if space_1d.basis == "B": x_grid = greville_loc pts = greville_loc[:, None] - wts = xp.ones(pts.shape, dtype=float) + wts = np.ones(pts.shape, dtype=float) # sub-interval index is always 0 for interpolation. - subs = xp.zeros(pts.shape[0], dtype=int) + subs = np.zeros(pts.shape[0], dtype=int) # !! shift away first interpolation point in eta_1 direction for polar domains !! if pts[0] == 0.0 and polar_shift: @@ -3518,32 +3523,32 @@ def get_pts_and_wts(space_1d, start, end, n_quad=None, polar_shift=False): union_breaks = space_1d.breaks[:-1] # Make union of Greville and break points - # tmp = set(xp.round(space_1d.histopolation_grid, decimals=14)).union( - # xp.round(union_breaks, decimals=14), + # tmp = set(np.round(space_1d.histopolation_grid, decimals=14)).union( + # np.round(union_breaks, decimals=14), # ) # tmp = list(tmp) # tmp.sort() - # tmp_a = xp.array(tmp) + # tmp_a = np.array(tmp) - tmp = set(xp.round(space_1d.histopolation_grid, decimals=14).tolist()).union( - xp.round(union_breaks, decimals=14).tolist() + tmp = set(np.round(space_1d.histopolation_grid, decimals=14).tolist()).union( + np.round(union_breaks, decimals=14).tolist() ) tmp = sorted(tmp) - tmp_a = xp.array(tmp) + tmp_a = np.array(tmp) x_grid = tmp_a[ - xp.logical_and( + np.logical_and( tmp_a - >= xp.min( + >= np.min( histopol_loc, ) - 1e-14, - tmp_a <= xp.max(histopol_loc) + 1e-14, + tmp_a <= np.max(histopol_loc) + 1e-14, ) ] # determine subinterval index (= 0 or 1): - subs = xp.zeros(x_grid[:-1].size, dtype=int) + subs = np.zeros(x_grid[:-1].size, dtype=int) for n, x_h in enumerate(x_grid[:-1]): add = 1 for x_g in histopol_loc: @@ -3558,12 +3563,6 @@ def get_pts_and_wts(space_1d, start, end, n_quad=None, polar_shift=False): pts_loc, wts_loc = np.polynomial.legendre.leggauss(n_quad) - if "cupy" in xp.__name__: - import cupy as cp - - pts_loc = cp.array(pts_loc) - wts_loc = cp.array(wts_loc) - x, wts = bsp.quadrature_grid(x_grid, pts_loc, wts_loc) pts = x % 1.0 @@ -3575,7 +3574,7 @@ def get_pts_and_wts_quasi( space_1d: SplineSpace, *, polar_shift: bool = False, -) -> tuple[xp.ndarray, xp.ndarray]: +) -> tuple[np.ndarray, np.ndarray]: r"""Obtain local projection point sets and weights in one grid direction for the quasi-interpolation method. The quasi-interpolation points are :math:`\nu - \mu +p` equidistant points :math:`\{ x^i_j \}_{0 \leq j < \nu - \mu +p}` in the sub-interval :math:`Q = [\eta_\mu , \eta_\nu]` given by: @@ -3623,12 +3622,12 @@ def get_pts_and_wts_quasi( # interpolation if space_1d.basis == "B": if degree == 1 and h != 1.0: - x_grid = xp.linspace(-(degree - 1) * h, 1.0 - h + (h / 2.0), (N + degree - 1) * 2) + x_grid = np.linspace(-(degree - 1) * h, 1.0 - h + (h / 2.0), (N + degree - 1) * 2) else: - x_grid = xp.linspace(-(degree - 1) * h, 1.0 - h, (N + degree - 1) * 2 - 1) + x_grid = np.linspace(-(degree - 1) * h, 1.0 - h, (N + degree - 1) * 2 - 1) pts = x_grid[:, None] % 1.0 - wts = xp.ones(pts.shape, dtype=float) + wts = np.ones(pts.shape, dtype=float) # !! shift away first interpolation point in eta_1 direction for polar domains !! if pts[0] == 0.0 and polar_shift: @@ -3639,16 +3638,16 @@ def get_pts_and_wts_quasi( # The computation of histopolation points breaks in case we have num_elements=1 and periodic boundary conditions since we end up with only one x_grid point. # We need to build the histopolation points by hand in this scenario. if degree == 0 and h == 1.0: - x_grid = xp.array([0.0, 0.5, 1.0]) + x_grid = np.array([0.0, 0.5, 1.0]) elif degree == 0 and h != 1.0: - x_grid = xp.linspace(-degree * h, 1.0 - h + (h / 2.0), (N + degree) * 2) + x_grid = np.linspace(-degree * h, 1.0 - h + (h / 2.0), (N + degree) * 2) else: - x_grid = xp.linspace(-degree * h, 1.0 - h, (N + degree) * 2 - 1) + x_grid = np.linspace(-degree * h, 1.0 - h, (N + degree) * 2 - 1) n_quad = degree + 1 # Gauss - Legendre quadrature points and weights # products of basis functions are integrated exactly - pts_loc, wts_loc = xp.polynomial.legendre.leggauss(n_quad) + pts_loc, wts_loc = np.polynomial.legendre.leggauss(n_quad) x, wts = bsp.quadrature_grid(x_grid, pts_loc, wts_loc) pts = x % 1.0 @@ -3662,26 +3661,26 @@ def get_pts_and_wts_quasi( N_b = N + degree # Filling the quasi-interpolation points for i=0 and i=1 (since they are equal) - x_grid = xp.linspace(0.0, knots[degree + 1], degree + 1) - x_aux = xp.linspace(0.0, knots[degree + 1], degree + 1) - x_grid = xp.append(x_grid, x_aux) + x_grid = np.linspace(0.0, knots[degree + 1], degree + 1) + x_aux = np.linspace(0.0, knots[degree + 1], degree + 1) + x_grid = np.append(x_grid, x_aux) # Now we append those for 1 tuple: ) +class CudaDerhamArguments(Argument): + """CUDA version of :class:`~struphy.kernel_arguments.pusher_args_kernels.DerhamArguments`. + + CUDA signature of :meth:`get_cuda_args`: ``long long* pn, double* tn1, double* tn2, double* tn3, long long* starts`` + + The scratch arrays of the pyccel class (``bn1``, ..., ``bd3``) are not part of it; CUDA kernels use + per-thread local arrays instead. + + Parameters + ---------- + pn : cupy.ndarray[int] + Spline degrees of :class:`~struphy.feec.psydac_derham.Derham` (int64). + + tn1, tn2, tn3 : cupy.ndarray[float] + Knot sequences of :class:`~struphy.feec.psydac_derham.Derham`. + + starts : cupy.ndarray[int] + Start indices (current MPI process) of :class:`~struphy.feec.psydac_derham.Derham` (int64). + """ + + def __init__(self, pn, tn1, tn2, tn3, starts): + self.pn = _cupy_array("pn", pn, np.int64) + self.tn1, self.tn2, self.tn3 = (_cupy_array("tn", t, np.float64) for t in (tn1, tn2, tn3)) + self.starts = _cupy_array("starts", starts, np.int64) + + def get_cuda_args(self) -> tuple: + return (self.pn, self.tn1, self.tn2, self.tn3, self.starts) + + class CudaDomainArguments(Argument): """CUDA version of :class:`~struphy.kernel_arguments.pusher_args_kernels.DomainArguments`.