From afe80fc806673d55212f691b0f7dc86da40d32f5 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 09:59:59 +0200 Subject: [PATCH 01/55] Update feectools (device ghost exchange plans, single-reduction PCG); add the pcg_sr solver option feectools: struphy-hub/feectools#104 (fix-cuda-aware-mpi-issues into devel-tiny). pcg_sr (PConjugateGradientSingleReduction) is accepted wherever pcg is, including the preconditioned solvers of the Poisson, diffusion, curl-curl and variational propagators. Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- src/struphy/io/options.py | 2 +- src/struphy/propagators/curl_curl_solve.py | 6 +++--- src/struphy/propagators/gyrokinetic_poisson_solve.py | 2 +- src/struphy/propagators/implicit_diffusion.py | 6 +++--- src/struphy/propagators/poisson_solve.py | 2 +- src/struphy/propagators/variational_density_evolve.py | 4 ++-- src/struphy/propagators/variational_momentum_advection.py | 4 ++-- 8 files changed, 14 insertions(+), 14 deletions(-) diff --git a/feectools b/feectools index 0c548d17a..8ae781d4d 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 0c548d17a3202bdd0531b0f9d81abccb27d8887e +Subproject commit 8ae781d4db109824a773a7bc6f8f0f6bf114b3ac diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 4f8919ed6..dadcc4800 100644 --- a/src/struphy/io/options.py +++ b/src/struphy/io/options.py @@ -74,7 +74,7 @@ class LiteralOptions: GivenInBasis = Literal["0", "1", "2", "3", "v", "physical", "physical_at_eta", "norm", None] # solvers - OptsSymmSolver = Literal["pcg", "cg"] + OptsSymmSolver = Literal["pcg", "pcg_sr", "cg"] OptsGenSolver = Literal["pbicgstab", "bicgstab", "gmres"] OptsMassPrecond = Literal["MassMatrixPreconditioner", "PolarMassMatrixPreconditioner", None] OptsShearAlfvenPrecond = Literal["MassMatrixPreconditioner", "PolarMassMatrixPreconditioner", "MultiGrid", None] diff --git a/src/struphy/propagators/curl_curl_solve.py b/src/struphy/propagators/curl_curl_solve.py index 7628178ac..f6b8e5b2f 100644 --- a/src/struphy/propagators/curl_curl_solve.py +++ b/src/struphy/propagators/curl_curl_solve.py @@ -191,7 +191,7 @@ def __post_init__(self): check_option(self.solver, LiteralOptions.OptsSymmSolver) check_option(self.precond, LiteralOptions.OptsCurlCurlPrecond) if self.precond is not None: - assert self.solver == "pcg", f"precond={self.precond!r} requires solver='pcg'." + assert self.solver in ("pcg", "pcg_sr"), f"precond={self.precond!r} requires solver='pcg' or 'pcg_sr'." # defaults if self.precond_params is None: @@ -323,8 +323,8 @@ def verify_rhs(j) -> Vector | FEECVariable | AccumulatorVector: pc_class = getattr(preconditioner, precond) pc = pc_class(self.mass_ops.M1, **self.options.precond_params) - # (only "pcg" takes a preconditioner) - pc_kwargs = {"pc": pc} if self.options.solver == "pcg" else {} + # (only "pcg" and "pcg_sr" take a preconditioner) + pc_kwargs = {"pc": pc} if self.options.solver in ("pcg", "pcg_sr") else {} self._solver = inverse( lhs, self.options.solver, diff --git a/src/struphy/propagators/gyrokinetic_poisson_solve.py b/src/struphy/propagators/gyrokinetic_poisson_solve.py index b5bc5a89e..18ab5eb5c 100644 --- a/src/struphy/propagators/gyrokinetic_poisson_solve.py +++ b/src/struphy/propagators/gyrokinetic_poisson_solve.py @@ -190,7 +190,7 @@ def __post_init__(self): check_option(self.precond, LiteralOptions.OptsDiffusionPrecond) check_option(self.which_geometry, self.OptsGeometry) if self.precond is not None: - assert self.solver == "pcg", f"precond={self.precond!r} requires solver='pcg'." + assert self.solver in ("pcg", "pcg_sr"), f"precond={self.precond!r} requires solver='pcg' or 'pcg_sr'." # defaults if self.precond_params is None: diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index c6e01c065..46143adc6 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -280,7 +280,7 @@ def __post_init__(self): check_option(self.solver, LiteralOptions.OptsSymmSolver) check_option(self.precond, LiteralOptions.OptsDiffusionPrecond) if self.precond is not None: - assert self.solver == "pcg", f"precond={self.precond!r} requires solver='pcg'." + assert self.solver in ("pcg", "pcg_sr"), f"precond={self.precond!r} requires solver='pcg' or 'pcg_sr'." preconditioner._parse_directions(self.stab_average) # defaults @@ -436,8 +436,8 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: pc = pc_class(self.mass_ops.M0, **self.options.precond_params) # solver just with A_2, but will be set during call with dt - # (only "pcg" takes a preconditioner) - pc_kwargs = {"pc": pc} if self.options.solver == "pcg" else {} + # (only "pcg" and "pcg_sr" take a preconditioner) + pc_kwargs = {"pc": pc} if self.options.solver in ("pcg", "pcg_sr") else {} self._solver = inverse( self._diffusion_op, self.options.solver, diff --git a/src/struphy/propagators/poisson_solve.py b/src/struphy/propagators/poisson_solve.py index 79137400c..7a765b811 100644 --- a/src/struphy/propagators/poisson_solve.py +++ b/src/struphy/propagators/poisson_solve.py @@ -145,7 +145,7 @@ def __post_init__(self): check_option(self.solver, LiteralOptions.OptsSymmSolver) check_option(self.precond, LiteralOptions.OptsDiffusionPrecond) if self.precond is not None: - assert self.solver == "pcg", f"precond={self.precond!r} requires solver='pcg'." + assert self.solver in ("pcg", "pcg_sr"), f"precond={self.precond!r} requires solver='pcg' or 'pcg_sr'." _parse_directions(self.stab_average) # defaults diff --git a/src/struphy/propagators/variational_density_evolve.py b/src/struphy/propagators/variational_density_evolve.py index 2b54296f7..cd98c37f9 100644 --- a/src/struphy/propagators/variational_density_evolve.py +++ b/src/struphy/propagators/variational_density_evolve.py @@ -207,7 +207,7 @@ def allocate(self): self._Mrho, self.options.solver, # "cg" takes no preconditioner - **({"pc": pc} if self.options.solver == "pcg" else {}), + **({"pc": pc} if self.options.solver in ("pcg", "pcg_sr") else {}), tol=1e-16, maxiter=500, recycle=True, @@ -416,7 +416,7 @@ def _initialize_projectors_and_mass(self): self.mass_ops.Mv, self.options.solver, # "cg" takes no preconditioner - **({"pc": self.pc_Mv} if self.options.solver == "pcg" else {}), + **({"pc": self.pc_Mv} if self.options.solver in ("pcg", "pcg_sr") else {}), tol=1e-16, maxiter=1000, verbose=False, diff --git a/src/struphy/propagators/variational_momentum_advection.py b/src/struphy/propagators/variational_momentum_advection.py index 7fc6e64ce..4ba1a5b65 100644 --- a/src/struphy/propagators/variational_momentum_advection.py +++ b/src/struphy/propagators/variational_momentum_advection.py @@ -135,7 +135,7 @@ def allocate(self): self._Mrho, self.options.solver, # "cg" takes no preconditioner - **({"pc": pc} if self.options.solver == "pcg" else {}), + **({"pc": pc} if self.options.solver in ("pcg", "pcg_sr") else {}), tol=1e-16, maxiter=500, recycle=True, @@ -284,7 +284,7 @@ def _initialize_mass(self): self.mass_ops.Mv, self.options.solver, # "cg" takes no preconditioner - **({"pc": self.pc_Mv} if self.options.solver == "pcg" else {}), + **({"pc": self.pc_Mv} if self.options.solver in ("pcg", "pcg_sr") else {}), tol=1e-16, maxiter=1000, verbose=False, From 607cab911c2fd33f9153d77afec9e998ad2a21d5 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:00:05 +0200 Subject: [PATCH 02/55] CuPy backend: host/device conversions in post-processing; math.sqrt for host scalars in Species Post-processing evaluates fields on the device and writes them from the host: the evaluation grids, point data and coordinates are converted where they cross, and HDF5 data is moved to the device before it is written into the spline coefficients. Co-Authored-By: Claude Opus 5.5 --- src/struphy/models/species.py | 4 ++-- src/struphy/post_processing/output.py | 31 ++++++++++++++++++--------- 2 files changed, 23 insertions(+), 12 deletions(-) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 3c8b47657..4bc2dcac1 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,9 +1,9 @@ import logging +import math import warnings from abc import ABCMeta, abstractmethod from dataclasses import fields, is_dataclass -import cunumpy as xp from maybempi import MPI from struphy.models.variables import Variable @@ -203,7 +203,7 @@ def __init__( con = ConstantsOfNature() # relevant frequencies - om_p = xp.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) + om_p = math.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) om_c = Z * con.e * units.B / (A * con.mH) # compute equation parameters diff --git a/src/struphy/post_processing/output.py b/src/struphy/post_processing/output.py index 57e3c0e32..c3ae64218 100644 --- a/src/struphy/post_processing/output.py +++ b/src/struphy/post_processing/output.py @@ -85,6 +85,13 @@ def load_plasma_plots() -> bool: return True +def _to_host(value): + """`value` (an array, or a list/tuple of arrays for vector-valued fields) as NumPy array(s).""" + if isinstance(value, (list, tuple)): + return type(value)(_to_host(v) for v in value) + return xp.to_numpy(value) + + class ProductMapping(Mapping[str, xr.DataArray]): """A discoverable mapping whose products are loaded on first access.""" @@ -568,13 +575,14 @@ def _attach_physical_coords( self, array: xr.DataArray, etas: tuple[Any, Any, Any], dims: tuple[str, ...] ) -> xr.DataArray: """Map a raw FEEC evaluation grid and attach its Cartesian coordinates.""" - grids = [np.atleast_1d(np.asarray(eta, dtype=float)) for eta in etas] - mapped = self.domain(*grids) + # the logical grids may be device arrays (CuPy backend): map them where they live, keep the coordinates on the host + grids = [np.atleast_1d(np.asarray(xp.to_numpy(eta), dtype=float)) for eta in etas] + mapped = self.domain(*(xp.asarray(grid) for grid in grids)) shape = tuple(len(grid) for grid in grids) keep = tuple(slice(None) if dim in dims else 0 for dim in ("eta1", "eta2", "eta3")) coordinates = {} for name, values in zip(("X", "Y", "Z"), mapped): - values = np.asarray(values).reshape(shape)[keep] + values = np.asarray(xp.to_numpy(values)).reshape(shape)[keep] coordinates[name] = (dims, values) if dims else values.item() return array.assign_coords(coordinates) @@ -1295,7 +1303,9 @@ def _process_fields( # the physical grid is only needed for output, hence it is only built on rank 0 if self._pproc_rank == 0: - grids_phy = list(self.domain(*grids_log)) + # output coordinates, on the host (the logical grids stay on the device for the field evaluation) + grids_phy = [xp.to_numpy(grid) for grid in self.domain(*grids_log)] + grids_log_host = [xp.to_numpy(grid) for grid in grids_log] else: grids_phy = None @@ -1327,8 +1337,9 @@ def _process_fields( if self._pproc_rank == 0: for species, vars in vals.items(): for name, val in vars.items(): - point_data[species][name][t] = val - point_data_phy[species][name][t] = vals_phy[species][name] + # evaluated on the device on the CuPy backend; written to files from the host + point_data[species][name][t] = _to_host(val) + point_data_phy[species][name][t] = _to_host(vals_phy[species][name]) # directory for the vtk files path_fields = os.path.join(self.path_pproc, "fields_data") @@ -1338,10 +1349,10 @@ def _process_fields( for species, vars in point_data.items(): variables = {} for name, val in vars.items(): - variables[name] = wrap_field_data(val, grids_log, grids_phy=grids_phy, name=name) + variables[name] = wrap_field_data(val, grids_log_host, grids_phy=grids_phy, name=name) if physical: variables[name + "_xyz"] = wrap_field_data( - point_data_phy[species][name], grids_log, grids_phy=grids_phy, name=name + "_xyz" + point_data_phy[species][name], grids_log_host, grids_phy=grids_phy, name=name + "_xyz" ) store.write_group(store.store_path(self.path_pproc), f"/{species}", xr.Dataset(variables)) @@ -1517,7 +1528,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[n * step, p1:-p1, p2:-p2, p3:-p3]) # HDF5 data is on the host # vector-valued field else: @@ -1530,7 +1541,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3]) vector.update_ghost_regions() if is_polar: From 4de785df2a75c6cbb5c55f9e88f52fd54cf9a717 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:07:26 +0200 Subject: [PATCH 03/55] CUDA version of the mass-matrix assembly kernel kernel_3d_mat - kernel_3d_mat_cuda.cu: one thread per (element, codomain basis function, domain basis function); replaces the host fallback on the CuPy backend. - WeightedMassOperator.assemble reads the quadrature grids (spans, weights, basis values) as backend arrays, copied to the device once and cached. - Simulation.allocate compiles the assembly kernel before it is first called, so the compilation is not counted in the assembly's profiling regions. Co-Authored-By: Claude Opus 5.5 --- .../feec/kernels/kernel_3d_mat/__init__.py | 18 ++-- .../kernel_3d_mat/kernel_3d_mat_cuda.cu | 92 +++++++++++++++++++ src/struphy/feec/mass.py | 21 ++++- src/struphy/simulation/sim.py | 18 ++++ 4 files changed, 139 insertions(+), 10 deletions(-) create mode 100644 src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu diff --git a/src/struphy/feec/kernels/kernel_3d_mat/__init__.py b/src/struphy/feec/kernels/kernel_3d_mat/__init__.py index a881301b4..76da51673 100644 --- a/src/struphy/feec/kernels/kernel_3d_mat/__init__.py +++ b/src/struphy/feec/kernels/kernel_3d_mat/__init__.py @@ -1,8 +1,4 @@ -"""The kernel of this folder: pyccel in ``_kernels.py``, CUDA in ``_cuda.cu`` once ported. - -No CUDA version yet: on the CuPy backend the pyccel kernel runs on host copies of the arrays (a warning is -emitted once), as before the move into this folder. -""" +"""The kernel of this folder: pyccel in ``kernel_3d_mat_kernels.py``, CUDA in ``kernel_3d_mat_cuda.cu``.""" from cunumpy.kernels import Kernel @@ -11,6 +7,16 @@ # the arguments the kernel writes to, by position: data OUTPUTS = (25,) + +def _n_threads(args): + """One thread per (element, codomain basis function, domain basis function), see kernel_3d_mat_cuda.cu.""" + n_elements = args[0].size * args[1].size * args[2].size + n_pairs = 1 + for degree in args[3:9]: # pi1, pi2, pi3, pj1, pj2, pj3 + n_pairs *= int(degree) + 1 + return n_elements * n_pairs + + kernel_3d_mat = Kernel.from_folder( - __name__, structs=CUDA_STRUCTS, missing_cuda="fallback", host_options={"outputs": OUTPUTS} + __name__, structs=CUDA_STRUCTS, host_options={"outputs": OUTPUTS}, n_threads_from=_n_threads ) diff --git a/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu b/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu new file mode 100644 index 000000000..6afab9398 --- /dev/null +++ b/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu @@ -0,0 +1,92 @@ +// CUDA version of kernel_3d_mat (kernel_3d_mat_kernels.py), with the same arguments in the same order. +#include "cunumpy/array_view.cuh" +#include "cunumpy/atomic.cuh" +#include "cunumpy/index.cuh" + +/** + * Integral of Lambda_i * mat_fun * Lambda_j over the local elements, added to the stencil matrix data, as in + * kernel_3d_mat_kernels.kernel_3d_mat. + * + * One thread per (element, codomain basis function (il1, il2, il3), domain basis function (jl1, jl2, jl3)): + * n_threads = ne1 * ne2 * ne3 * (pi1 + 1) * (pi2 + 1) * (pi3 + 1) * (pj1 + 1) * (pj2 + 1) * (pj3 + 1) (see + * __init__.py), the element index varying slowest and jl3 fastest. Each thread computes the quadrature sum of its + * pair exactly as the pyccel kernel (same factors, same order) and adds it to its entry of `data` with an atomic + * addition: neighbouring elements contribute to the same entries, so only the summation order over elements + * differs from pyccel (round-off). + * + * @param spans1, spans2, spans3 Knot span index of each local element, per direction. + * @param pi1, pi2, pi3 Degrees of the codomain basis functions. + * @param pj1, pj2, pj3 Degrees of the domain basis functions. + * @param starts1, starts2, starts3 Global start index of the rows on this process. + * @param pads1, pads2, pads3 Paddings (ghost regions) of `data`. + * @param w1, w2, w3 Quadrature weights [element, quadrature point]. + * @param bi1, bi2, bi3 Codomain basis functions [element, local basis function, derivative, quadrature point]. + * @param bj1, bj2, bj3 Domain basis functions, same layout. + * @param mat_fun Weight function at the quadrature points [flattened point 1, flattened point 2, flattened point 3]. + * @param data StencilMatrix data (rows with ghost regions..., diagonals...); the results are added (not set). + */ +extern "C" __global__ void kernel_3d_mat(Array1D spans1, Array1D spans2, + Array1D spans3, int pi1, int pi2, int pi3, int pj1, int pj2, + int pj3, int starts1, int starts2, int starts3, int pads1, int pads2, + int pads3, Array2D w1, Array2D w2, Array2D w3, + Array4D bi1, Array4D bi2, Array4D bi3, + Array4D bj1, Array4D bj2, Array4D bj3, + Array3D mat_fun, Array6D data) +{ + const long long ne1 = spans1.size(); + const long long ne2 = spans2.size(); + const long long ne3 = spans3.size(); + + const long long ni1 = pi1 + 1, ni2 = pi2 + 1, ni3 = pi3 + 1; + const long long nj1 = pj1 + 1, nj2 = pj2 + 1, nj3 = pj3 + 1; + const long long n_pairs = ni1 * ni2 * ni3 * nj1 * nj2 * nj3; + + CUNUMPY_THREAD_1D(thread, ne1 * ne2 * ne3 * n_pairs); + + // CUDA-only: split the thread index into the pyccel loop variables (iel1, iel2, iel3, il1, il2, il3, jl1, jl2, + // jl3), jl3 fastest + long long rest = thread; + const long long jl3 = rest % nj3; rest /= nj3; + const long long jl2 = rest % nj2; rest /= nj2; + const long long jl1 = rest % nj1; rest /= nj1; + const long long il3 = rest % ni3; rest /= ni3; + const long long il2 = rest % ni2; rest /= ni2; + const long long il1 = rest % ni1; rest /= ni1; + const long long iel3 = rest % ne3; rest /= ne3; + const long long iel2 = rest % ne2; rest /= ne2; + const long long iel1 = rest; + + const long long nq1 = w1.shape[1]; + const long long nq2 = w2.shape[1]; + const long long nq3 = w3.shape[1]; + + double value = 0.0; + + for (long long q1 = 0; q1 < nq1; ++q1) { + for (long long q2 = 0; q2 < nq2; ++q2) { + for (long long q3 = 0; q3 < nq3; ++q3) { + const double wvol = w1(iel1, q1) * w2(iel2, q2) * w3(iel3, q3) * + mat_fun(iel1 * nq1 + q1, iel2 * nq2 + q2, iel3 * nq3 + q3); + + const double bi = bi1(iel1, il1, 0, q1) * bi2(iel2, il2, 0, q2) * bi3(iel3, il3, 0, q3); + const double bj = bj1(iel1, jl1, 0, q1) * bj2(iel2, jl2, 0, q2) * bj3(iel3, jl3, 0, q3); + + value += wvol * bi * bj; + } + } + } + + // global spline indices + const long long i_global1 = spans1(iel1) - pi1 + il1; + const long long i_global2 = spans2(iel2) - pi2 + il2; + const long long i_global3 = spans3(iel3) - pi3 + il3; + + // local spline indices (- starts --> can be negative, will therefore be written to ghost regions) + const long long i_local1 = i_global1 - starts1; + const long long i_local2 = i_global2 - starts2; + const long long i_local3 = i_global3 - starts3; + + cunumpy_atomic_add(&data(pads1 + i_local1, pads2 + i_local2, pads3 + i_local3, pads1 + jl1 - il1, + pads2 + jl2 - il2, pads3 + jl3 - il3), + value); +} diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 8ff1d6395..c99b0c68a 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -2167,6 +2167,19 @@ def transpose(self, conjugate=False): return M + def _quad_grid(self, space_name, attribute, component): + """Quadrature-grid arrays (spans, weights or basis values) of one component of a space, as backend arrays. + + The Derham keeps them as host arrays; on the CuPy backend the assembly kernel needs them on the device, + so they are copied there once and cached (they do not change between assemblies). + """ + cache = self.__dict__.setdefault("_quad_grid_cache", {}) + key = (space_name, attribute, component) + if key not in cache: + arrays = getattr(self.derham.spline_attributes[space_name], attribute)[component] + cache[key] = [xp.asarray(array) for array in arrays] + return cache[key] + def assemble(self, weights=None, clear=True): r""" Assembles the weighted mass matrix, i.e. computes the integrals @@ -2243,7 +2256,7 @@ def assemble(self, weights=None, clear=True): # loop over codomain spaces (rows) for a, codomain_space in enumerate(codomain_spaces): # knot span indices of elements of local domain - codomain_spans = spline_attr[W_name].quad_grid_spans[a] + codomain_spans = self._quad_grid(W_name, "quad_grid_spans", a) # global start spline index on process codomain_starts = [int(start) for start in codomain_space.coeff_space.starts] @@ -2255,10 +2268,10 @@ def assemble(self, weights=None, clear=True): pts = self._quad_pts(W_name, a) # global quadrature weights in format (local element, local weight) - wts = spline_attr[W_name].quad_grid_wts[a] + wts = self._quad_grid(W_name, "quad_grid_wts", a) # evaluated basis functions at quadrature points of codomain space - codomain_basis = spline_attr[W_name].quad_grid_bases[a] + codomain_basis = self._quad_grid(W_name, "quad_grid_bases", a) # loop over domain spaces (columns) for b, domain_space in enumerate(domain_spaces): @@ -2298,7 +2311,7 @@ def assemble(self, weights=None, clear=True): ) # evaluated basis functions at quadrature points of domain space - domain_basis = spline_attr[V_name].quad_grid_bases[b] + domain_basis = self._quad_grid(V_name, "quad_grid_bases", b) # assemble matrix (if mat_w is not zero) by calling the appropriate kernel (1d, 2d or 3d) if not_weight_zero or self._is_scalar: diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 2719095f5..3df498eea 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -243,6 +243,8 @@ def allocate(self): logger.debug("\nAllocating simulation data ...") + self._compile_setup_cuda_kernels() + with ProfileManager.profile_region("setup: allocate"): # feec with ProfileManager.profile_region("allocate feec", functions=[self._allocate_feec]): @@ -280,6 +282,22 @@ def kernels(self) -> tuple[Kernel, ...]: owners += self.model.prop_list return collect_kernels(*owners) + def _compile_setup_cuda_kernels(self) -> None: + """Compile the CUDA kernels that :meth:`allocate` itself calls (mass-matrix assembly) before it calls them. + + :meth:`compile_cuda_kernels` runs after :meth:`allocate`, so without this the first assembly would include + the compilation (~1 s per kernel on a node whose CuPy cache does not have it yet), in the assembly's + profiling regions. Here it is counted in ``"setup: compile cuda kernels (assembly)"``. Nothing happens on + the NumPy backend. + """ + if xp.get_backend() != "cupy": + return + + from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat + + with ProfileManager.profile_region("setup: compile cuda kernels (assembly)"): + kernel_3d_mat.compile() + def compile_cuda_kernels(self) -> tuple[Kernel, ...]: """Compile the CUDA versions of the kernels of :meth:`kernels` now, before the time stepping. From 773d7b120a5037a5f9d8fb74ae1b3e02c9116856 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:26:26 +0200 Subject: [PATCH 04/55] Parity cases for kernel_3d_mat Random quadrature data on 3 x 2 x 2 elements with mixed degrees, with the first row local or a ghost row. test_cuda_kernels_have_parity_cases requires a case for every kernel with a CUDA version; the emulation test compares it with the pyccel kernel without a GPU. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/cuda_parity_cases.py | 26 ++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index 7e4bb00e2..adb1b2bb3 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -313,6 +313,30 @@ def kernel_pullpush_args(case): ) +def kernel_3d_mat_threads(args): + """One thread per (element, codomain basis function, domain basis function), as the kernel's launch size.""" + from struphy.feec.kernels.kernel_3d_mat import _n_threads + + return _n_threads(args) + + +def kernel_3d_mat_args(case): + """Mass-matrix assembly on 3 x 2 x 2 elements with mixed degrees and random quadrature data; `case` is the start + index of the rows in direction 1 (1: the first row is a ghost row, local index -1).""" + rng = np.random.default_rng(20 + case) + ne, nq = (3, 2, 2), (3, 2, 2) + pi, pj, pads = (2, 1, 1), (1, 2, 1), (2, 2, 2) + starts = (case, 0, 0) + spans = [xp.asarray(np.arange(n, dtype=np.int64) + p) for n, p in zip(ne, pi)] + w = [xp.asarray(rng.random((n, q))) for n, q in zip(ne, nq)] + bi = [xp.asarray(rng.random((n, p + 1, 1, q))) for n, p, q in zip(ne, pi, nq)] + bj = [xp.asarray(rng.random((n, p + 1, 1, q))) for n, p, q in zip(ne, pj, nq)] + mat_fun = xp.asarray(rng.normal(size=tuple(n * q for n, q in zip(ne, nq)))) + rows = tuple(n + p + 2 * pad for n, p, pad in zip(ne, pi, pads)) + data = xp.asarray(rng.normal(size=rows + tuple(2 * pad + 1 for pad in pads))) # added to, not overwritten + return (*spans, *pi, *pj, *starts, *pads, *w, *bi, *bj, mat_fun, data) + + # ---------------------------------------------------------------- all kernels with a CUDA version PUSHER_TOLERANCES = {"rtol": 1e-13, "atol": 1e-14} @@ -331,6 +355,8 @@ def size_of(index, per_thread=1): PARITY_CASES = { + # atomic additions: the summation order over elements differs from pyccel + "kernel_3d_mat": ParityCases((0, 1), kernel_3d_mat_args, rtol=1e-12, atol=1e-13, n_threads=kernel_3d_mat_threads), "push_eta_stage": ParityCases( [ (bc, method, stage) From 1d0e0ab7a30b1aff6ca8adf3aa4a5e9df20c5817 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:08:01 +0200 Subject: [PATCH 05/55] Apply separable mass matrices as Kronecker products of 1D banded matrices WeightedMassOperator.dot applies the blocks of its matrix that are Kronecker products of 1D banded matrices (constant weights, e.g. on a Cuboid) as three banded passes (struphy.feec.kronecker_mass, CUDA kernel kron_band_pass on the GPU) instead of the 3D stencil product. The replacement is built on first use and invalidated by assembly, copies and in-place arithmetic; a matrix changed by arithmetic that is not a Kronecker product is not checked again. STRUPHY_KRONECKER_MASS=0 switches it off. The boundary and extraction operators around the matrix are skipped when they are identities instead of being applied as copies. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/kron_band_pass_cuda.cu | 35 ++++ src/struphy/feec/kronecker_mass.py | 247 ++++++++++++++++++++++++ src/struphy/feec/mass.py | 85 ++++++-- 3 files changed, 355 insertions(+), 12 deletions(-) create mode 100644 src/struphy/feec/kron_band_pass_cuda.cu create mode 100644 src/struphy/feec/kronecker_mass.py diff --git a/src/struphy/feec/kron_band_pass_cuda.cu b/src/struphy/feec/kron_band_pass_cuda.cu new file mode 100644 index 000000000..745e782f4 --- /dev/null +++ b/src/struphy/feec/kron_band_pass_cuda.cu @@ -0,0 +1,35 @@ +/** + * One pass of a Kronecker product of 1D banded matrices (struphy.feec.kronecker_mass._band_pass): + * + * out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...] along `axis`, + * + * with the other two axes unchanged (out and inp have the same extent along them). One thread per entry of `out`, + * the last axis fastest, so that neighbouring threads read neighbouring entries of `inp` along every axis. `out` may + * be a strided view (the owned rows of a stencil vector's data). + * + * @param A Banded 1D matrix, shape (out.shape[axis], n_diagonals). + * @param inp Input array (3D). + * @param out Output array (3D, any strides); written, not added to. + * @param axis Axis of the pass (0, 1 or 2). + * @param offset Index in `inp` of diagonal 0 of row 0 along `axis`. + */ +#include "cunumpy/array_view.cuh" + +extern "C" __global__ void kron_band_pass(Array2D A, Array3D inp, Array3D out, int axis, + long long offset) { + const long long n1 = out.shape[1], n2 = out.shape[2]; + const long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; + if (t >= out.shape[0] * n1 * n2) return; + const long long i2 = t % n2; + const long long i1 = (t / n2) % n1; + const long long i0 = t / (n1 * n2); + const long long i = axis == 0 ? i0 : (axis == 1 ? i1 : i2); + + double v = 0.0; + for (long long d = 0; d < A.shape[1]; ++d) { + const long long j = i + offset + d; + const double x = axis == 0 ? inp(j, i1, i2) : (axis == 1 ? inp(i0, j, i2) : inp(i0, i1, j)); + v += A(i, d) * x; + } + out(i0, i1, i2) = v; +} diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py new file mode 100644 index 000000000..2050edb82 --- /dev/null +++ b/src/struphy/feec/kronecker_mass.py @@ -0,0 +1,247 @@ +"""Fast application of assembled mass matrices that are Kronecker products of 1D banded matrices. + +A mass matrix block ``M[i, j] = int w Lambda_i Lambda_j`` of tensor-product splines factorizes as +``M = c * M_1 (x) M_2 (x) M_3`` when its weight ``w`` is constant (e.g. on a Cuboid). Applied as a +:class:`~feectools.linalg.stencil.StencilMatrix`, every row reads all ``n_1 * n_2 * n_3`` diagonals (105 for +degrees (1, 2, 3)) from memory; as a Kronecker product it is three 1D banded passes over the vector, ``n_1 + n_2 + +n_3`` operations per entry and no 3D matrix data. :func:`kronecker_apply_operator` checks each assembled block and +returns an operator for the matrix-vector products in which the blocks that are Kronecker products (to round-off) +are replaced by :class:`KroneckerBandMatrix`. + +The 1D factors are read from the assembled data itself (rows owned by this process), so they hold exactly what the +assembly computed, on any domain decomposition, and the check is the assembled data, not an assumption about the +mapping or the weights. +""" + +import logging +import os +from functools import cache +from pathlib import Path + +import cunumpy as xp +from feectools.linalg.basic import LinearOperator +from feectools.linalg.block import BlockLinearOperator +from feectools.linalg.stencil import StencilMatrix, StencilVector + +logger = logging.getLogger("struphy") + +# relative tolerance of the Kronecker check: the 3D assembly and the product of the factors differ by round-off only +KRONECKER_RTOL = 1e-12 + + +def kronecker_mass_enabled() -> bool: + """Whether :func:`kronecker_apply_operator` may replace blocks (``STRUPHY_KRONECKER_MASS=0`` switches it off).""" + return os.environ.get("STRUPHY_KRONECKER_MASS", "1").lower() not in ("0", "false", "no") + + +@cache +def _band_pass_cuda(): + """The CUDA kernel of one banded pass (compiled on first call).""" + from cunumpy.kernels import CudaKernel + + return CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass") + + +def _band_pass(A, inp, out, axis, offset): + """``out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...]`` along ``axis`` (other axes unchanged).""" + if xp.is_gpu(out): + _band_pass_cuda()(A, inp, out, axis, offset, n_threads=out.size) + return + m = out.shape[axis] + shape = [1, 1, 1] + shape[axis] = m + out[...] = 0.0 + for d in range(A.shape[1]): + idx = [slice(None)] * 3 + idx[axis] = slice(offset + d, offset + d + m) + out += A[:, d].reshape(shape) * inp[tuple(idx)] + + +class KroneckerBandMatrix(LinearOperator): + """A square 3D :class:`~feectools.linalg.stencil.StencilMatrix` applied as a Kronecker product of 1D banded + matrices, ``(A_1 (x) A_2 (x) A_3) x`` on the rows owned by this process. + + ``A_k[i, d]`` is the entry of owned row ``i`` (along direction ``k``) and diagonal ``d``, column + ``i - q_k + d`` with ``q_k = (n_k - 1) // 2``, as in the stencil matrix. The product is computed in three passes, + each over the owned rows of its direction and the full (padded) extent of the directions after it, so the passes + need no communication beyond the ghost regions of the input. Use :meth:`from_stencil`. + + Parameters + ---------- + stencil : StencilMatrix + The assembled matrix (kept for ``toarray``, ``tosparse`` and ``transpose``). + factors : tuple of arrays + ``A_1, A_2, A_3``, shapes ``(m_k, n_k)``, C-contiguous backend arrays. + """ + + def __init__(self, stencil: StencilMatrix, factors): + self._stencil = stencil + self._factors = factors + + V = stencil.domain + self._m = tuple(int(e - s + 1) for s, e in zip(V.starts, V.ends)) + self._pads = tuple(int(p) for p in V.pads) + self._q = tuple((int(n) - 1) // 2 for n in stencil._data.shape[3:]) + + # temporaries of the first two passes: owned rows along the passes done, padded extent along the others + x_shape = V.zeros()._data.shape + self._t1 = xp.empty((self._m[0], x_shape[1], x_shape[2]), dtype=stencil._data.dtype) + self._t2 = xp.empty((self._m[0], self._m[1], x_shape[2]), dtype=stencil._data.dtype) + + @classmethod + def from_stencil(cls, S: StencilMatrix, rtol: float = KRONECKER_RTOL): + """The Kronecker form of ``S`` if it is one (to ``rtol``, relative to its largest entry), else None. + + Only square 3D matrices on one space (domain is codomain) with shifts 1 qualify, e.g. the diagonal blocks of + the mass matrices. The factors are read from the owned rows of ``S``: with reference rows ``r_k`` (the first + owned row along each direction) and central diagonals ``q_k``, + ``A_1[i, :] = S[i, r_2, r_3, :, q_2, q_3]``, ``A_2[i, :] = S[r_1, i, r_3, q_1, :, q_3] / S[r, q]``, + ``A_3[i, :] = S[r_1, r_2, i, q_1, q_2, :] / S[r, q]``; the product is then compared with all owned rows. + """ + if not isinstance(S, StencilMatrix) or S.domain is not S.codomain or S.domain.ndim != 3: + return None + V = S.domain + if any(int(s) != 1 for s in V.shifts) or S._data.ndim != 6: + return None + + p = [int(p) for p in V.pads] + m = [int(e - s + 1) for s, e in zip(V.starts, V.ends)] + q = [(int(n) - 1) // 2 for n in S._data.shape[3:]] + if min(m) < 1: + return None + data = S._data + + den = data[p[0], p[1], p[2], q[0], q[1], q[2]] + if float(den) == 0.0: + return None + A1 = xp.ascontiguousarray(data[p[0] : p[0] + m[0], p[1], p[2], :, q[1], q[2]]) + A2 = xp.ascontiguousarray(data[p[0], p[1] : p[1] + m[1], p[2], q[0], :, q[2]] / den) + A3 = xp.ascontiguousarray(data[p[0], p[1], p[2] : p[2] + m[2], q[0], q[1], :] / den) + + # compare with every owned row: planes (rows along direction 1) in chunks, the maxima kept on the device + # (one read at the end); a few planes first, which rejects most matrices that are not Kronecker products + # (e.g. ones assembled from particle data, checked again after every change) at once + owned = data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + A23 = A2[:, None, None, :, None] * A3[None, :, None, None, :] # (m2, m3, 1, n2, n3) + + def max_diff_and_abs(rows): + block = owned[rows] + recon = A1[rows][:, None, None, :, None, None] * A23[None] + return xp.max(xp.abs(block - recon)), xp.max(xp.abs(block)) + + def fails(max_diff, max_abs): + max_diff, max_abs = float(max_diff), float(max_abs) + return max_abs == 0.0 or max_diff > rtol * max_abs + + quick = sorted({0, m[0] // 2, m[0] - 1}) + if fails(*max_diff_and_abs(xp.asarray(quick))): + return None + + # chunks of at most ~64 MB of temporaries + plane = max(1, owned[0].size) + chunk = max(1, min(m[0], (8 << 20) // plane)) + max_diff = xp.zeros(()) + max_abs = xp.zeros(()) + for start in range(0, m[0], chunk): + d, a = max_diff_and_abs(slice(start, min(start + chunk, m[0]))) + max_diff = xp.maximum(max_diff, d) + max_abs = xp.maximum(max_abs, a) + if fails(max_diff, max_abs): + return None + return cls(S, (A1, A2, A3)) + + # ------------------------------------------------------------------ + # LinearOperator interface + # ------------------------------------------------------------------ + @property + def domain(self): + return self._stencil.domain + + @property + def codomain(self): + return self._stencil.codomain + + @property + def dtype(self): + return self._stencil.dtype + + @property + def factors(self): + """The 1D factors ``A_1, A_2, A_3`` (rows owned by this process).""" + return self._factors + + def tosparse(self): + return self._stencil.tosparse() + + def toarray(self): + return self._stencil.toarray() + + def transpose(self, conjugate=False): + return self._stencil.transpose(conjugate=conjugate) + + def dot(self, v, out=None): + assert isinstance(v, StencilVector) + assert v.space is self.domain + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.codomain + + if not v.ghost_regions_in_sync: + v.update_ghost_regions() + + A1, A2, A3 = self._factors + p, q, m = self._pads, self._q, self._m + owned = out._data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + + _band_pass(A1, v._data, self._t1, 0, p[0] - q[0]) + _band_pass(A2, self._t1, self._t2, 1, p[1] - q[1]) + _band_pass(A3, self._t2, owned, 2, p[2] - q[2]) + + out.ghost_regions_in_sync = False + return out + + +def kronecker_apply_operator(mat, name: str = ""): + """The operator to apply ``mat`` with: ``mat`` itself, or a copy whose Kronecker blocks are replaced. + + Parameters + ---------- + mat : StencilMatrix or BlockLinearOperator + An assembled (block) mass matrix. + name : str + Name of the operator, for the log. + + Returns + ------- + LinearOperator + ``mat`` if no block is a Kronecker product (or :func:`kronecker_mass_enabled` is False); otherwise a + :class:`KroneckerBandMatrix` (for a StencilMatrix) or a BlockLinearOperator with the same blocks except the + Kronecker ones. It is valid only as long as ``mat`` is not modified. + """ + if not kronecker_mass_enabled(): + return mat + + if isinstance(mat, StencilMatrix): + kron = KroneckerBandMatrix.from_stencil(mat) + if kron is not None: + logger.info(f"Mass matrix {name}: applied as a Kronecker product of 1D banded matrices") + return mat if kron is None else kron + + if isinstance(mat, BlockLinearOperator): + blocks = {} + replaced = [] + for key, block in mat._blocks.items(): + kron = KroneckerBandMatrix.from_stencil(block) + if kron is None: + blocks[key] = block + else: + blocks[key] = kron + replaced.append(key) + if not replaced: + return mat + logger.info(f"Mass matrix {name}: blocks {replaced} applied as Kronecker products of 1D banded matrices") + return BlockLinearOperator(mat.domain, mat.codomain, blocks=blocks) + + return mat diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index c99b0c68a..ad0e5d05a 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -25,6 +25,7 @@ from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat from struphy.feec.kernels.kernel_3d_matrixfree import kernel_3d_matrixfree from struphy.feec.kernels.kernel_3d_vec import kernel_3d_vec +from struphy.feec.kronecker_mass import kronecker_apply_operator from struphy.feec.linear_operators import BoundaryOperator from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.feec.utilities import LocalProjectionMatrix, LocalRotationMatrix, get_quad_grids @@ -1628,6 +1629,8 @@ def __init__( blocks, self._weights = self._init_blocks_from_weights(weights_info) self._mat = self._wrap_blocks(blocks) + # the operator dot() applies self._mat with (Kronecker blocks replaced), built on first use, see _apply_mat + self._apply_mat_cache = None # the weights are stored in the block order of the (possibly transposed) operator if transposed: @@ -2046,6 +2049,32 @@ def M0(self): def matrix(self): return self._mat + @property + def _apply_mat(self): + """The operator :meth:`dot` applies: :attr:`matrix`, with its blocks that are Kronecker products of 1D banded + matrices (constant weights, e.g. on a Cuboid) applied as such (:mod:`struphy.feec.kronecker_mass`). + + Built on first use after every change of the matrix (assembly, copy into, in-place arithmetic). + """ + if self._apply_mat_cache is None: + if getattr(self, "_kronecker_rejected", False): + self._apply_mat_cache = self._mat + else: + self._apply_mat_cache = kronecker_apply_operator(self._mat, name=self.name) + # a matrix changed by in-place arithmetic that is not a Kronecker product (e.g. the Schur complement + # M1 - dt^2 BC with particle data, rebuilt every step) is not checked again after further arithmetic: + # the check reads the whole matrix. Assembly and copies check again. + if getattr(self, "_changed_by_arithmetic", False) and self._apply_mat_cache is self._mat: + self._kronecker_rejected = True + return self._apply_mat_cache + + def _matrix_changed(self, by_arithmetic: bool): + """Invalidate :attr:`_apply_mat` after a change of the matrix (in-place arithmetic, or assembly/copy).""" + self._apply_mat_cache = None + self._changed_by_arithmetic = by_arithmetic + if not by_arithmetic: + self._kronecker_rejected = False + @property def nquads(self): if self._nquads is None: @@ -2093,29 +2122,54 @@ def dot(self, v, out=None, apply_bc=True): if apply_bc: if self._transposed: - self._W_boundary_op_T.dot(v, out=self._temp_WB) - self._W_extraction_op_T.dot(self._temp_WB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) - self._V_extraction_op.dot(self._temp_VE, out=self._temp_VB) - out = self._V_boundary_op.dot(self._temp_VB, out=out) + ops_in = (self._W_boundary_op_T, self._W_extraction_op_T) + ops_out = (self._V_extraction_op, self._V_boundary_op) + tmps_in = (self._temp_WB, self._temp_mat) + tmps_out = (self._temp_VE, self._temp_VB) else: - self._V_boundary_op_T.dot(v, out=self._temp_VB) - self._V_extraction_op_T.dot(self._temp_VB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) - self._W_extraction_op.dot(self._temp_WE, out=self._temp_WB) - out = self._W_boundary_op.dot(self._temp_WB, out=out) + ops_in = (self._V_boundary_op_T, self._V_extraction_op_T) + ops_out = (self._W_extraction_op, self._W_boundary_op) + tmps_in = (self._temp_VB, self._temp_mat) + tmps_out = (self._temp_WE, self._temp_WB) + out = self._apply_chain(v, ops_in, tmps_in, ops_out, tmps_out, out) else: if self._transposed: self._W_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) + self._apply_mat.dot(self._temp_mat, out=self._temp_VE) out = self._V_extraction_op.dot(self._temp_VE, out=out) else: self._V_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) + self._apply_mat.dot(self._temp_mat, out=self._temp_WE) out = self._W_extraction_op.dot(self._temp_WE, out=out) return out + def _apply_chain(self, v, ops_in, tmps_in, ops_out, tmps_out, out): + """``out = B_out E_out M E_in B_in v`` (boundary and extraction operators around the matrix), see :meth:`dot`. + + Identity operators (e.g. the extraction operators without polar splines, or boundary operators without + essential boundary conditions) are skipped instead of applied as copies: each vector is passed on as it is, + and the matrix writes into the next temporary that is actually used (or into ``out``). + """ + if out is v: + # in place: the chain must not read v after writing the output, so it reads a copy + v = v.copy() + + x = v + for op, tmp in zip(ops_in, tmps_in): + if not isinstance(op, IdentityOperator): + x = op.dot(x, out=tmp) + + # the first output operator that is not an identity reads the matrix product from its temporary + remaining = [(op, tmp) for op, tmp in zip(ops_out, tmps_out) if not isinstance(op, IdentityOperator)] + if not remaining: + return self._apply_mat.dot(x, out=out) + y = self._apply_mat.dot(x, out=remaining[0][1]) + for k, (op, _) in enumerate(remaining): + target = out if k == len(remaining) - 1 else remaining[k + 1][1] + y = op.dot(y, out=target) + return y + def transpose(self, conjugate=False): """ Returns the transposed operator. @@ -2202,6 +2256,7 @@ def assemble(self, weights=None, clear=True): Whether to first set all data to zero before assembly. If False, the new contributions are added to existing ones. """ + self._matrix_changed(by_arithmetic=False) # the matrix changes assert not self._dry_run, ( "A dry-run operator has no matrix data and cannot be assembled (memory estimation only)." ) @@ -2477,6 +2532,7 @@ def copy(self, out=None): out._weights = [list(row) for row in self._weights] self._mat.copy(out=out._mat) + out._matrix_changed(by_arithmetic=False) if self._creation_info is None: out._creation_info = None @@ -2487,6 +2543,7 @@ def copy(self, out=None): def __imul__(self, a): self._mat *= a + self._matrix_changed(by_arithmetic=True) self._creation_info = None return self @@ -2497,11 +2554,13 @@ def __iadd__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat += M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat += M + self._matrix_changed(by_arithmetic=True) return self else: @@ -2514,11 +2573,13 @@ def __isub__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat -= M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat -= M + self._matrix_changed(by_arithmetic=True) return self else: From e4152816bcef07831902ee514cfb0425f51ea04d Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:08:29 +0200 Subject: [PATCH 06/55] Poisson/diffusion solve: fewer passes over the vectors per iteration - BoundaryOperator.dot on device Stencil/Block vectors: out = v * mask, one kernel per block (one for all blocks when they are contiguous float64), with masks found once from apply_essential_bc_to_array, instead of a copy and one fill per boundary and component. - WeightedMassOperator.without_boundary_ops / MassWithoutBoundaryOps: the mass operator without its boundary operators. In grad.T @ M1 @ grad, grad already applies the same 0/1 masks, so this is the same operator with two fewer passes per product (not with the multigrid preconditioner). - ImplicitDiffusion: the automatic 1e-14 stabilization is left out of the solver operator when a Dirichlet boundary already makes grad.T M1 grad positive definite (diffusion_mat M1); no stab_mat product for the right-hand side when sigma_2 = 0 (Poisson). Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/linear_operators.py | 105 +++++++++++++++++- src/struphy/feec/mass.py | 52 +++++++++ src/struphy/propagators/implicit_diffusion.py | 40 ++++++- 3 files changed, 192 insertions(+), 5 deletions(-) diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index b36318a38..ac2d680b6 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -1,10 +1,75 @@ +import cunumpy as xp from feectools.linalg.basic import LinearOperator, Vector, VectorSpace -from feectools.linalg.stencil import StencilVectorSpace +from feectools.linalg.block import BlockVector +from feectools.linalg.stencil import StencilVector, StencilVectorSpace from struphy.feec.utilities import apply_essential_bc_to_array from struphy.polar.basic import PolarDerhamSpace +def _blocks(v): + """The StencilVector blocks of a StencilVector or BlockVector (of StencilVectors), else None.""" + if isinstance(v, StencilVector): + return [v] + if isinstance(v, BlockVector) and all(isinstance(b, StencilVector) for b in v.blocks): + return list(v.blocks) + return None + + +_MASKED_COPY_SRC = r""" +extern "C" __global__ void masked_copy(const unsigned long long* __restrict__ ptrs, const long long* __restrict__ offsets, + int n_comp, long long total) +{ + long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= total) return; + int c = 0; + while (c < n_comp - 1 && i >= offsets[c + 1]) ++c; + long long j = i - offsets[c]; + const double* v = (const double*)ptrs[3 * c]; + double* out = (double*)ptrs[3 * c + 1]; + const double* mask = (const double*)ptrs[3 * c + 2]; + out[j] = v[j] * mask[j]; +} +""" + + +def _masked_copy_fused(op, v_blocks, out_blocks, masks) -> bool: + """``out_blocks[c] = v_blocks[c] * masks[c]`` for all blocks with one kernel (pointer tables cached on ``op``). + + False (nothing done) unless every block is C-contiguous float64 of the shape of its mask.""" + arrays = [(vb._data, ob._data, m) for vb, ob, (m, _) in zip(v_blocks, out_blocks, masks)] + for a, b, m in arrays: + if not ( + a.dtype == b.dtype == m.dtype == xp.float64 + and a.shape == b.shape == m.shape + and a.flags.c_contiguous + and b.flags.c_contiguous + and m.flags.c_contiguous + ): + return False + import cupy as cp + import numpy as np + + cache = op.__dict__.setdefault("_masked_copy_cache", {}) + key = tuple(int(x.data.ptr) for triple in arrays for x in triple) + entry = cache.get(key) + if entry is None: + if "kernel" not in cache: + cache["kernel"] = cp.RawKernel(_MASKED_COPY_SRC, "masked_copy") + sizes = [a.size for a, _, _ in arrays] + offsets = np.concatenate([[0], np.cumsum(sizes)]).astype(np.int64) + entry = (cp.asarray(np.array(key, dtype=np.uint64)), cp.asarray(offsets), int(offsets[-1])) + if len(cache) > 64: + kernel = cache["kernel"] + cache.clear() + cache["kernel"] = kernel + cache[key] = entry + ptrs, offsets, total = entry + if total: + cache["kernel"](((total + 255) // 256,), (256,), (ptrs, offsets, np.int32(len(arrays)), np.int64(total))) + return True + + class BoundaryOperator(LinearOperator): r""" Applies homogeneous Dirichlet boundary conditions to a vector. @@ -196,6 +261,22 @@ def dot(self, v, out=None): else: assert isinstance(out, Vector) assert out.space == self._codomain + masks = self._device_masks(v) + if masks is not None: + # copy and boundary conditions in one kernel per block (see _device_masks) + in_sync = True + v_blocks, out_blocks = _blocks(v), _blocks(out) + if len(v_blocks) > 1 and _masked_copy_fused(self, v_blocks, out_blocks, masks): + pass + else: + for v_blk, out_blk, (mask, _) in zip(v_blocks, out_blocks, masks): + xp.multiply(v_blk._data, mask, out=out_blk._data) + for v_blk, (_, stale) in zip(v_blocks, masks): + in_sync = in_sync and v_blk.ghost_regions_in_sync and not stale + # one flag for all blocks: a BlockVector flagged in sync with stale blocks would have them updated + # one at a time by the operators that check each block, instead of together + out.ghost_regions_in_sync = in_sync + return out v.copy(out=out) # apply boundary conditions to output vector @@ -203,6 +284,28 @@ def dot(self, v, out=None): return out + def _device_masks(self, v): + """Per block of ``v``: (0/1 mask of the entries the boundary conditions keep, whether they stale the ghosts). + + For device (CuPy) data of Stencil/Block vectors, ``dot`` computes ``out = v * mask`` (one kernel per block) + instead of a copy and one fill per boundary and component. The masks are found once, by applying + :func:`apply_essential_bc_to_array` to a vector of ones, so they zero exactly the entries it zeroes (ghost + entries included) and flag the same blocks as having stale ghost regions. None on the host backend and for + other vectors (e.g. polar vectors), which use the copy and :func:`apply_essential_bc_to_array`. + """ + if not hasattr(self, "_masks_cache"): + self._masks_cache = None + blocks = _blocks(v) + if blocks is not None and all(xp.is_gpu(b._data) for b in blocks): + probe = v.space.zeros() + probe_blocks = _blocks(probe) + for b in probe_blocks: + b._data[...] = 1.0 + b.ghost_regions_in_sync = True + apply_essential_bc_to_array(self._space_id, probe, self.bc) + self._masks_cache = [(b._data.copy(), not b.ghost_regions_in_sync) for b in probe_blocks] + return self._masks_cache + def transpose(self, conjugate=False): """ Returns the transposed operator. diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index ad0e5d05a..dd115a459 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -2170,6 +2170,15 @@ def _apply_chain(self, v, ops_in, tmps_in, ops_out, tmps_out, out): y = op.dot(y, out=target) return y + def without_boundary_ops(self) -> "MassWithoutBoundaryOps": + """This operator applied without its boundary operators, see :class:`MassWithoutBoundaryOps`.""" + return MassWithoutBoundaryOps(self) + + @property + def boundary_ops(self): + """The boundary operators ``(B_V, B_W)`` applied on the domain and codomain side.""" + return self._V_boundary_op, self._W_boundary_op + def transpose(self, conjugate=False): """ Returns the transposed operator. @@ -2714,6 +2723,49 @@ def info(self, use_rst=False): return info(self, use_rst=use_rst) +class MassWithoutBoundaryOps(LinearOperator): + """A :class:`WeightedMassOperator` applied without its boundary operators: ``E_W M E_V^T`` instead of + ``B_W E_W M E_V^T B_V^T`` (:meth:`WeightedMassOperator.dot` with ``apply_bc=False``). + + Between operators that apply the same boundary operators on the outside, e.g. in ``grad.T @ M1 @ grad`` with + ``grad = B1 G B0^T``, this is the same operator: boundary operators are 0/1 diagonal masks, hence symmetric and + idempotent (``B1^T B1 = B1``), and the masked products are bit for bit the same, with two fewer passes over the + vectors per product. Get it with :meth:`WeightedMassOperator.without_boundary_ops`. + """ + + def __init__(self, mass: "WeightedMassOperator"): + self._mass = mass + + @property + def mass(self): + """The wrapped mass operator.""" + return self._mass + + @property + def domain(self): + return self._mass.domain + + @property + def codomain(self): + return self._mass.codomain + + @property + def dtype(self): + return self._mass.dtype + + def tosparse(self): + raise NotImplementedError() + + def toarray(self): + raise NotImplementedError() + + def transpose(self, conjugate=False): + return MassWithoutBoundaryOps(self._mass.transpose(conjugate=conjugate)) + + def dot(self, v, out=None): + return self._mass.dot(v, out=out, apply_bc=False) + + class StencilMatrixFreeMassOperator(LinearOperator): r"""Class implementing matrix-free weighted mass operators between StencilVectorSpaces. diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index 46143adc6..0b16de4f7 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -306,10 +306,24 @@ def options(self, new): @profile def allocate(self): # always stabilize + auto_stabilized = False if xp.abs(self.options.sigma_1) < 1e-14: self.options.sigma_1 = 1e-14 + auto_stabilized = True logger.warning(f"Stabilizing Poisson solve with {self.options.sigma_1 =}") + # The automatic stabilization only makes singular problems (e.g. periodic in all directions) solvable. With a + # Dirichlet boundary in some direction, grad.T @ M1 @ grad is already positive definite and 1e-14 * stab_mat + # changes the solution only at round-off, but applying it costs a whole stab_mat product per solver iteration; + # it is then left out of the solver's operator. (M1perp can be degenerate, so it keeps the stabilization.) + self._stabilize_lhs = not ( + auto_stabilized + and self.options.diffusion_mat == "M1" + and any(any(sides) for sides in self.derham.dirichlet_bc) + ) + if not self._stabilize_lhs: + logger.info("Dirichlet boundary present: the automatic stabilization is left out of the solver operator") + # model parameters self._sigma_1 = self.options.sigma_1 self._sigma_2 = self.options.sigma_2 @@ -408,7 +422,17 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: # Set lhs matrices (without dt) self._stab_mat = stab_mat - self._diffusion_op = self.derham.grad.T @ diffusion_mat @ self.derham.grad + # grad = B1 G B0^T already applies the boundary operator B1 of the 1-forms on both sides of diffusion_mat; if + # diffusion_mat applies the same one, it is left out there (the same operator, see MassWithoutBoundaryOps), + # except for the multigrid preconditioner, which coarsens the operator by its structure + mid = diffusion_mat + grad, grad_t = self.derham.grad, self.derham.grad.T + B1 = self.derham.boundary_ops["1"] + if self.options.precond != "MultiGrid" and all( + op is B1 for op in getattr(diffusion_mat, "boundary_ops", (None,)) + ): + mid = diffusion_mat.without_boundary_ops() + self._diffusion_op = grad_t @ mid @ grad # preconditioner and solver for Ax=b # the multigrid and stiffness preconditioners depend on sigma_1, they are updated in __call__ @@ -537,8 +561,13 @@ def __call__(self, dt): # compute rhs phin = self.variables.phi.spline.vector - rhs = self._stab_mat.dot(phin, out=self._rhs) - rhs *= sig_2 + if sig_2 == 0.0: + # no stab_mat product only to multiply it by zero (e.g. Poisson: sigma_2 = 0) + rhs = self._rhs + rhs *= 0.0 + else: + rhs = self._stab_mat.dot(phin, out=self._rhs) + rhs *= sig_2 self._rhs2 *= 0.0 for src, coeff in zip(self.sources, self.coeffs): @@ -582,7 +611,10 @@ def __call__(self, dt): self.diagnostic.spline.vector = proj.solve(rhs) # compute lhs - self._solver.linop = sig_1 * self._stab_mat + self._diffusion_op + if self._stabilize_lhs: + self._solver.linop = sig_1 * self._stab_mat + self._diffusion_op + else: + self._solver.linop = self._diffusion_op if sig_1 != self._pc_sig_1: if self._mg is not None: self._mg.update(self._solver.linop) From 33ccffc2d38e78888b13989cf0cef19aecdcd9fd Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:08:47 +0200 Subject: [PATCH 07/55] GPU marker buffer: column-major storage; marker rows prepared inside the push kernels - Particles allocates the marker buffer column-major on the CuPy backend: kernels and the Python column operations access a few columns of many markers, so with a row-major buffer every access touched whole rows and the threads of a warp read addresses a row apart. The CUDA kernels address the buffer through its strides (Array2D in MarkerArgs); NumPy keeps C order for the pyccel kernels. - push_eta_stage and push_vxb_analytic prepare each marker row themselves at stage 0 (MarkerArgs.prepare_rows), so Pusher._push skips the separate prepare_push pass when no init or evaluation kernel runs before them. prepare_push is launched on a 2D grid. Co-Authored-By: Claude Opus 5.5 --- src/struphy/kernel_arguments/pusher_args.cuh | 1 + .../kernel_arguments/pusher_args_cuda.py | 18 +++++++-- src/struphy/pic/base.py | 7 +++- .../push_eta_stage/push_eta_stage_cuda.cu | 3 ++ .../push_vxb_analytic_cuda.cu | 4 ++ src/struphy/pic/pushing/prepare_push_cuda.cu | 26 ++++++------- src/struphy/pic/pushing/pusher.py | 37 +++++++++++++------ .../pic/pushing/pusher_utilities_kernels.cuh | 14 +++++++ src/struphy/pic/tests/test_kernel_backends.py | 2 +- 9 files changed, 80 insertions(+), 32 deletions(-) diff --git a/src/struphy/kernel_arguments/pusher_args.cuh b/src/struphy/kernel_arguments/pusher_args.cuh index e3a547e1e..0ace87644 100644 --- a/src/struphy/kernel_arguments/pusher_args.cuh +++ b/src/struphy/kernel_arguments/pusher_args.cuh @@ -19,6 +19,7 @@ struct MarkerArgs { int mu_idx; long long* bc_type; bool pole; + int prepare_rows; }; struct DerhamArgs { diff --git a/src/struphy/kernel_arguments/pusher_args_cuda.py b/src/struphy/kernel_arguments/pusher_args_cuda.py index feac42098..9b2a2cabc 100644 --- a/src/struphy/kernel_arguments/pusher_args_cuda.py +++ b/src/struphy/kernel_arguments/pusher_args_cuda.py @@ -12,15 +12,21 @@ from cunumpy.arguments import CudaStructArguments -def _device_array(name, arr, dtype, ndim=None): - """Check that ``arr`` is a C-contiguous CuPy array of ``dtype`` (and ``ndim``) on the current device (never copies).""" +def _device_array(name, arr, dtype, ndim=None, strided=False): + """Check that ``arr`` is a CuPy array of ``dtype`` (and ``ndim``) on the current device (never copies). + + It must be C-contiguous, unless ``strided``: arrays passed as strided views (``Array2D``, e.g. the + column-major marker buffer) may have any strides. + """ import cupy as cp if not isinstance(arr, cp.ndarray): raise TypeError(f"{name} must be a CuPy array, got {type(arr)}") if arr.device.id != cp.cuda.runtime.getDevice(): raise ValueError(f"{name} must be on the current CUDA device") - if arr.dtype != dtype or not arr.flags.c_contiguous: + if arr.dtype != dtype: + raise TypeError(f"{name} must be an array of dtype {np.dtype(dtype)}") + if not strided and not arr.flags.c_contiguous: raise TypeError(f"{name} must be a C-contiguous array of dtype {np.dtype(dtype)}") if ndim is not None and arr.ndim != ndim: raise TypeError(f"{name} must be a {ndim}-dimensional array, got {arr.ndim} dimensions") @@ -46,6 +52,8 @@ class CudaMarkerArguments(CudaStructArguments): ("mu_idx", "int"), ("bc_type", "long long*"), ("pole", "bool"), + # CUDA-only: 1 while a pusher has its kernel prepare each marker row at stage 0 (Pusher._push), else 0 + ("prepare_rows", "int"), ) def __init__( @@ -64,7 +72,8 @@ def __init__( bc_type, pole: bool, ): - self.markers = _device_array("markers", markers, np.float64) + # a strided view in MarkerArgs: Particles allocates the buffer column-major on the GPU + self.markers = _device_array("markers", markers, np.float64, strided=True) if markers.ndim != 2: raise TypeError("markers must be a two-dimensional array.") self.valid_mks = _device_array("valid_mks", valid_mks, np.bool_) @@ -80,6 +89,7 @@ def __init__( self.mu_idx = mu_idx self.bc_type = _device_array("bc_type", bc_type, np.int64) self.pole = pole + self.prepare_rows = 0 self.pack() diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index e3529b9c4..8f238267b 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2634,7 +2634,12 @@ def _allocate_marker_array(self, dry_run: bool = False): if dry_run: return - self._markers = xp.zeros((self.n_rows, self.n_cols), dtype=float) + # column-major on the GPU: the kernels (and the column operations on the Python side) access a few columns of + # many markers, so with a row-major buffer every access touches the whole row (n_cols * 8 bytes) and the + # threads of a warp read addresses a row apart. The CUDA kernels address the buffer through its strides + # (Array2D in MarkerArgs). The pyccel kernels expect C order, so NumPy keeps it. + order = "F" if xp.get_backend() == "cupy" else "C" + self._markers = xp.zeros((self.n_rows, self.n_cols), dtype=float, order=order) # allocate auxiliary arrays self._holes = xp.zeros(self.n_rows, dtype=bool) diff --git a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu index 9a6c24d52..eeb1ac9ca 100644 --- a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu @@ -26,6 +26,9 @@ extern "C" __global__ void push_eta_stage(double dt, int stage, MarkerArgs args_ int n_markers = args_markers.n_markers; if (ip >= n_markers) return; + // CUDA-only: prepare this row here instead of in a separate pass of prepare_push (Pusher._push) + if (stage == 0 && args_markers.prepare_rows) pusher_utilities_kernels::prepare_marker_row(ip, args_markers); + int first_pusher_idx = args_markers.first_pusher_idx; int first_free_idx = args_markers.first_free_idx; double last = stage == n_stages - 1 ? 1. : 0.; diff --git a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu index 7d2a270ec..d4822afcf 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu @@ -2,6 +2,7 @@ #include "struphy/bsplines/evaluation_kernels_3d.cuh" #include "struphy/geometry/evaluation_kernels.cuh" #include "struphy/linear_algebra/linalg_kernels.cuh" +#include "struphy/pic/pushing/pusher_utilities_kernels.cuh" using namespace struphy_cuda; @@ -30,6 +31,9 @@ extern "C" __global__ void push_vxb_analytic(double dt, int stage, MarkerArgs ar int n_markers = args_markers.n_markers; if (ip >= n_markers) return; + // CUDA-only: prepare this row here instead of in a separate pass of prepare_push (Pusher._push) + if (stage == 0 && args_markers.prepare_rows) pusher_utilities_kernels::prepare_marker_row(ip, args_markers); + int first_pusher_idx = args_markers.first_pusher_idx; // check if marker is a hole diff --git a/src/struphy/pic/pushing/prepare_push_cuda.cu b/src/struphy/pic/pushing/prepare_push_cuda.cu index 71c02556d..d22518b55 100644 --- a/src/struphy/pic/pushing/prepare_push_cuda.cu +++ b/src/struphy/pic/pushing/prepare_push_cuda.cu @@ -5,24 +5,20 @@ * markers[:, first_shift_idx:residual_idx] = 0.0 * markers[:, residual_idx:-2] = 0.0 * - * One thread per entry of the contiguous column range [first_pusher_idx, n_cols - 2) of every row, so that - * neighbouring threads write neighbouring addresses of the row-major buffer and the whole range is written - * in a single pass (instead of one strided pass over the buffer per slice). + * A 2D launch, n_threads = (n_rows, n_cols - 2 - first_pusher_idx): x is the row, y the column within the range + * [first_pusher_idx, n_cols - 2), so the whole range is written in a single pass. The buffer is addressed through + * a strided view: on the CuPy backend it is column-major (Particles allocates it in Fortran order), so neighbouring + * threads (rows) of one column write neighbouring addresses. * - * @param markers Marker buffer (n_rows x n_cols, row-major), every row including holes. - * @param n_rows Number of rows of the buffer. - * @param n_cols Number of columns of the buffer. + * @param markers Marker buffer (n_rows x n_cols, any strides), every row including holes. * @param first_pusher_idx First column of the saved initial phase space coordinates. * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_pusher_idx + n_init). */ -extern "C" __global__ void prepare_push(double* markers, long long n_rows, int n_cols, int first_pusher_idx, - int n_init) { - long long width = n_cols - 2 - first_pusher_idx; - long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; - if (t >= n_rows * width) return; +#include "cunumpy/array_view.cuh" - long long ip = t / width; - int col = first_pusher_idx + (int)(t - ip * width); - double* row = markers + ip * n_cols; - row[col] = col < first_pusher_idx + n_init ? row[col - first_pusher_idx] : 0.0; +extern "C" __global__ void prepare_push(Array2D markers, int first_pusher_idx, int n_init) { + const long long ip = (long long)blockDim.x * blockIdx.x + threadIdx.x; + const long long k = (long long)blockDim.y * blockIdx.y + threadIdx.y; + if (ip >= markers.shape[0] || k >= markers.shape[1] - 2 - first_pusher_idx) return; + markers(ip, first_pusher_idx + k) = k < n_init ? markers(ip, k) : 0.0; } diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index c9271883f..3617791a5 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -18,6 +18,11 @@ logger = logging.getLogger("struphy") +# CUDA kernels that prepare each marker row themselves at stage 0 when MarkerArgs.prepare_rows is set (see +# pusher_utilities_kernels.prepare_marker_row), so that Pusher._push can skip the separate pass of prepare_push +_KERNELS_PREPARING_ROWS = frozenset({"push_eta_stage", "push_vxb_analytic"}) + + @cache def _prepare_push_cuda() -> CudaKernel: """The CUDA kernel that fuses the column-slice assignments at the start of :meth:`Pusher._push` @@ -217,17 +222,22 @@ def _push(self, dt: float): logger.debug(f"{residual_idx =}") logger.debug(f"{self.particles.n_cols =}") - if self._cuda: + # On the GPU the kernels of _KERNELS_PREPARING_ROWS prepare each row themselves at stage 0, saving a pass over + # the marker buffer; only if no other kernel runs before them (init and evaluation kernels may read the + # prepared columns). + prepare_in_kernel = ( + self._cuda + and self.kernel.name in _KERNELS_PREPARING_ROWS + and not self.init_kernels + and not self.eval_kernels + ) + + if prepare_in_kernel: + pass + elif self._cuda: # the three assignments below in one pass over the buffer (each slice is a strided pass on its own) n_rows, n_cols = markers.shape - _prepare_push_cuda()( - markers, - n_rows, - n_cols, - first_pusher_idx, - 3 + vdim, - n_threads=n_rows * (n_cols - 2 - first_pusher_idx), - ) + _prepare_push_cuda()(markers, first_pusher_idx, 3 + vdim, n_threads=(n_rows, n_cols - 2 - first_pusher_idx)) else: init_slice = slice(first_pusher_idx, first_shift_idx) shift_slice = slice(first_shift_idx, residual_idx) @@ -279,15 +289,20 @@ def _push(self, dt: float): # sort according to alpha-weighted average self._sort_for_alpha(self._alpha_in_kernel) - # push markers + # push markers (the first call prepares the rows if prepare_in_kernel) + args_markers = self.particles.args_markers + if prepare_in_kernel and stage == 0 and k == 1: + args_markers.prepare_rows = 1 with ProfileManager.profile_region("kernel: " + self.kernel.name): self.kernel( dt, stage, - self.particles.args_markers, + args_markers, self._args_domain, *self._args_kernel, ) + if prepare_in_kernel and args_markers.prepare_rows: + args_markers.prepare_rows = 0 # markers have moved if self.pushes_eta: diff --git a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh index 259091753..3e0f5e8ea 100644 --- a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh +++ b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh @@ -43,6 +43,20 @@ __device__ inline void reflect_velocity(long long ip, Array2D markers, * the host. The marker view carries shape and strides; j is a CUDA-only * column index replacing Pyccel array slices. */ +/** + * Prepare marker row ip for a push, as the kernel prepare_push (prepare_push_cuda.cu) does for every row: copy the + * phase space coordinates (columns 0 to 3 + vdim) to the columns from first_pusher_idx and zero the columns after + * them up to n_cols - 2. Called at stage 0 by the push kernels that take over the preparation from prepare_push + * (args_markers.prepare_rows, see Pusher._push), for every row including holes, so that it is exactly the same. + */ +__device__ inline void prepare_marker_row(int ip, const MarkerArgs& args_markers) { + const int first = args_markers.first_pusher_idx; + const int n_init = 3 + args_markers.vdim; + const int end = (int)args_markers.markers.shape[1] - 2; + for (int col = first; col < end; ++col) + args_markers.markers(ip, col) = col < first + n_init ? args_markers.markers(ip, col - first) : 0.0; +} + __device__ inline void apply_kinetic_bc_marker(int ip, const MarkerArgs& args_markers, const DomainArgs& args_domain, bool newton) { const long long* bc_type = args_markers.bc_type; diff --git a/src/struphy/pic/tests/test_kernel_backends.py b/src/struphy/pic/tests/test_kernel_backends.py index 75e3178d3..add0614d3 100644 --- a/src/struphy/pic/tests/test_kernel_backends.py +++ b/src/struphy/pic/tests/test_kernel_backends.py @@ -37,7 +37,7 @@ # (CUDA class, pyccel source, pyccel class, struct members that only the CUDA class has, its header) ARGUMENT_PAIRS = ( - (CudaMarkerArguments, "pusher_args_kernels.py", "MarkerArguments", {"n_markers"}, HEADER), + (CudaMarkerArguments, "pusher_args_kernels.py", "MarkerArguments", {"n_markers", "prepare_rows"}, HEADER), (CudaDerhamArguments, "pusher_args_kernels.py", "DerhamArguments", {"nt1", "nt2", "nt3"}, HEADER), (CudaDomainArguments, "pusher_args_kernels.py", "DomainArguments", set(), HEADER), ( From c45019c23c48916b7d85a0448da6210270e2e496 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:26:52 +0200 Subject: [PATCH 08/55] BoundaryOperator: no fused mask kernel on the fake CuPy cunumpy's fake CuPy (tests without a GPU) cannot launch raw kernels; the per-block multiplication is used there (test_mass_preconditioners_fake_cupy failed). Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/linear_operators.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/struphy/feec/linear_operators.py b/src/struphy/feec/linear_operators.py index ac2d680b6..056817730 100644 --- a/src/struphy/feec/linear_operators.py +++ b/src/struphy/feec/linear_operators.py @@ -36,7 +36,12 @@ def _blocks(v): def _masked_copy_fused(op, v_blocks, out_blocks, masks) -> bool: """``out_blocks[c] = v_blocks[c] * masks[c]`` for all blocks with one kernel (pointer tables cached on ``op``). - False (nothing done) unless every block is C-contiguous float64 of the shape of its mask.""" + False (nothing done) unless every block is C-contiguous float64 of the shape of its mask, and on cunumpy's fake + CuPy (tests without a GPU), which cannot launch raw kernels.""" + from cunumpy.kernel_testing import fake_cupy_active + + if fake_cupy_active(): + return False arrays = [(vb._data, ob._data, m) for vb, ob, (m, _) in zip(v_blocks, out_blocks, masks)] for a, b, m in arrays: if not ( From 074151ef16ff7956915225d4b6e3d1f26b518f37 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:28:24 +0200 Subject: [PATCH 09/55] Column-major marker buffer opt-in (STRUPHY_MARKERS_COLUMN_MAJOR=1); prepare_rows in the CUDA emulation Kernels without a CUDA version run on host copies of the device arrays, which keep the F order (cunumpy.to_numpy uses get(order='A')), and the pyccel kernels reject it: the sorting-box kernels failed on the fake CuPy (test_sorting_boxes_on_cupy_fake_cupy), and the same holds on a GPU for every marker kernel without a CUDA version. The buffer is therefore C-ordered unless the variable is set. The CUDA emulation of the tests reads the struct fields from the pyccel argument objects, which do not have the CUDA-only prepare_rows flag (every emulated marker kernel failed). Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/base.py | 21 ++++++++++++++++----- src/struphy/pic/tests/cuda_emulation.py | 5 ++++- 2 files changed, 20 insertions(+), 6 deletions(-) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 8f238267b..75947a4ea 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -75,6 +75,16 @@ """Orbit quantities of the position columns, which post-processing maps to physical coordinates.""" +def markers_column_major() -> bool: + """Whether the marker buffer is allocated column-major on the CuPy backend (``STRUPHY_MARKERS_COLUMN_MAJOR=1``). + + Off by default: kernels without a CUDA version run on host copies of the device arrays, which keep the F order + (``cunumpy.to_numpy``), and the pyccel kernels reject it (e.g. the sorting-box kernels). Switch it on for runs + whose marker kernels all have CUDA versions (e.g. the Vlasov push profiling case). + """ + return os.environ.get("STRUPHY_MARKERS_COLUMN_MAJOR", "0").lower() in ("1", "true", "yes") + + class Particles(metaclass=ABCMeta): r""" Base class for particle species. @@ -2634,11 +2644,12 @@ def _allocate_marker_array(self, dry_run: bool = False): if dry_run: return - # column-major on the GPU: the kernels (and the column operations on the Python side) access a few columns of - # many markers, so with a row-major buffer every access touches the whole row (n_cols * 8 bytes) and the - # threads of a warp read addresses a row apart. The CUDA kernels address the buffer through its strides - # (Array2D in MarkerArgs). The pyccel kernels expect C order, so NumPy keeps it. - order = "F" if xp.get_backend() == "cupy" else "C" + # column-major on the GPU (opt-in, see markers_column_major): the kernels (and the column operations on the + # Python side) access a few columns of many markers, so with a row-major buffer every access touches the whole + # row (n_cols * 8 bytes) and the threads of a warp read addresses a row apart. The CUDA kernels address the + # buffer through its strides (Array2D in MarkerArgs). The pyccel kernels expect C order, so NumPy + # keeps it. + order = "F" if xp.get_backend() == "cupy" and markers_column_major() else "C" self._markers = xp.zeros((self.n_rows, self.n_cols), dtype=float, order=order) # allocate auxiliary arrays diff --git a/src/struphy/pic/tests/cuda_emulation.py b/src/struphy/pic/tests/cuda_emulation.py index ca3c76d01..31071a699 100644 --- a/src/struphy/pic/tests/cuda_emulation.py +++ b/src/struphy/pic/tests/cuda_emulation.py @@ -50,10 +50,13 @@ def _struct_field(value, name): """Field `name` of a struct, read from the pyccel argument object `value`. The CUDA structs have the attributes of their pyccel classes plus the knot lengths ``nt1``, ``nt2``, - ``nt3`` of ``DerhamArgs``, which pyccel gets from ``len(tn1)``. + ``nt3`` of ``DerhamArgs``, which pyccel gets from ``len(tn1)``, and the flag ``prepare_rows`` of ``MarkerArgs``. """ if name in ("nt1", "nt2", "nt3"): return len(getattr(value, "tn" + name[-1])) + if name == "prepare_rows": + # CUDA-only flag of MarkerArgs (set by Pusher._push around a launch); 0: the rows are prepared beforehand + return getattr(value, name, 0) return getattr(value, name) From 289da51e568b3b27e14e297f4defbadcff45623b Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:09:16 +0200 Subject: [PATCH 10/55] Particles: chunked loading, row-based weight update, leaner MPI sorting; PushVxB with a static field - Pseudo-random markers are drawn in chunks of 2**20 (was 10000): 16 instead of 1600 iterations for 16M markers, each a few small kernels and a device sync on the GPU. - update_weights finds the rows of the valid markers once and reads the columns it needs, instead of selecting every property with the boolean mask and gathering whole rows. - mpi_sort_markers: the ghost bookkeeping is skipped without sorting boxes (no ghosts can exist); no barriers around the device exchange, which waits for every send and receive; the sorting positions for alpha = 1 are one add instead of five kernels. - PushVxB without a magnetic field variable computes the total field once (no ghost exchange, an MPI synchronization point, before every push). Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/base.py | 63 ++++++++++++++++---- src/struphy/pic/tests/test_sorting_device.py | 4 +- src/struphy/propagators/push_vxb.py | 16 +++-- 3 files changed, 64 insertions(+), 19 deletions(-) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 75947a4ea..071d28a3a 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1240,7 +1240,10 @@ def draw_markers( # counting integers num_loaded_particles_loc = 0 # number of particles alreday loaded (local) num_loaded_particles_glob = 0 # number of particles already loaded (each clone) - chunk_size = 10000 # TODO: number of particle chunk + # markers drawn per chunk: large chunks keep the number of iterations (each a few small kernels and + # a device sync on the GPU) low, e.g. 16 instead of 1600 for 16M markers; a chunk of temporaries + # takes ~(3 + vdim + 4) * 8 bytes per marker (~50 MB). The NumPy draws do not depend on it. + chunk_size = 2**20 # Total number of markers to draw (sum over all clones) while num_loaded_particles_glob < int(self.Np): @@ -1563,8 +1566,16 @@ def update_weights(self): """ from struphy.pic.particles import ParticlesSPH - if xp.size(self.markers_wo_holes_and_ghost) == 0: + # the rows of the valid markers, found once: each property (positions, weights0, ...) would select them + # again with the boolean mask, and gathering markers_wo_holes_and_ghost just to test for markers copies + # every column of every valid row + rows = xp.flatnonzero(self.valid_mks) + if rows.size == 0: return + markers = self._markers + + def column(c): + return xp.take(markers[:, c], rows) if isinstance(self, ParticlesSPH): f0 = self.f0.n0(self.positions) @@ -1572,7 +1583,9 @@ def update_weights(self): # in case of CanonicalMaxwellian, evaluate constants_of_motion # if isinstance(self.f0, CanonicalMaxwellian): # self.save_constants_of_motion() - f0 = self.f0(*self.f_coords.T) + idx = self.f_coords_index + cols = range(markers.shape[1])[idx] if isinstance(idx, slice) else np.atleast_1d(idx).tolist() + f0 = self.f0(*[column(c) for c in cols]) # if f_init is vol-form, transform to 0-form if self.is_volume_form[0]: @@ -1581,7 +1594,7 @@ def update_weights(self): if self.is_volume_form[1]: f0 /= self.f0.velocity_jacobian_det(*self.f_jacobian_coords.T) - self.weights = self.weights0 - f0 / self.sampling_density_values / self.Np + markers[rows, self.index["weights"]] = column(self.index["w0"]) - f0 / column(self.index["s0"]) / self.Np @profile def binning( @@ -1823,10 +1836,25 @@ def mpi_sort_markers( (staged through host memory otherwise), see :meth:`_sendrecv_markers_device`. The per-rank counts are small host arrays on every backend (:meth:`_sendrecv_all_to_all`). """ + # Ghost particles (last column -2) only exist with sorting boxes (SPH ghost boxes, + # _prepare_ghost_particles); without them the ghost mask stays all False and its + # bookkeeping (full passes over the buffer, a nonzero) can be skipped. + ghosts_possible = self._initialized_sorting + + # On the GPU the marker rows go through cunumpy.mpi.exchange, which waits for every send and + # receive (_sendrecv_markers_device): no barriers are needed around the exchange. The host path + # does not wait for its sends, so it keeps them. + barriers = xp.get_array_backend(self._markers) != "cupy" + if remove_ghost: - self._remove_ghost_particles() + if ghosts_possible: + self._remove_ghost_particles() + else: + # what _remove_ghost_particles leaves when there are no ghosts: refreshed holes + self.update_holes() - self._Barrier() + if barriers: + self._Barrier() # before sorting, apply kinetic bc if apply_bc: @@ -1856,7 +1884,8 @@ def mpi_sort_markers( # refresh ghost mask: received markers may land in rows that previously held # ghost particles. update_holes alone recomputes valid_mks from a stale # _ghost_particles mask, which would wrongly exclude these incoming real markers. - self._update_ghost_particles() + if ghosts_possible: + self._update_ghost_particles() # check if all markers are on the right process after sorting if do_test: @@ -1870,7 +1899,8 @@ def mpi_sort_markers( assert all_on_right_proc # assert self.phasespace_coords.size > 0, f'No particles on process {self.mpi_rank}, please rebalance, aborting ...' - self._Barrier() + if barriers: + self._Barrier() @profile @ProfileManager.profile("apply_kinetic_bc") @@ -4485,10 +4515,19 @@ def _sendrecv_determine_mtbs( assert alpha.size == 3 assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) bi = self.first_pusher_idx - self._sorting_etas[:] = ( - alpha * (self.markers[:, :3] + self.markers[:, bi + 3 + self.vdim : bi + 3 + self.vdim + 3]) - + (1.0 - alpha) * self.markers[:, bi : bi + 3] - ) + if isinstance(alpha, float) and alpha == 1.0: + # the common case: alpha * (pos + shift) + 0 * init is exactly pos + shift for finite values, + # computed with one kernel instead of five + xp.add( + self.markers[:, :3], + self.markers[:, bi + 3 + self.vdim : bi + 3 + self.vdim + 3], + out=self._sorting_etas, + ) + else: + self._sorting_etas[:] = ( + alpha * (self.markers[:, :3] + self.markers[:, bi + 3 + self.vdim : bi + 3 + self.vdim + 3]) + + (1.0 - alpha) * self.markers[:, bi : bi + 3] + ) # for non-periodic axes, eta = 1 is the right end of the domain and must not be wrapped to eta = 0 non_periodic_axes = [axis for axis in range(3) if axis not in self._periodic_axes] diff --git a/src/struphy/pic/tests/test_sorting_device.py b/src/struphy/pic/tests/test_sorting_device.py index ac5c7aa88..cd8929176 100644 --- a/src/struphy/pic/tests/test_sorting_device.py +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -11,6 +11,7 @@ default. The staged path is forced with :func:`cunumpy.mpi.set_mpi_cuda_aware`. """ +import contextlib import sys from contextlib import contextmanager @@ -138,7 +139,8 @@ def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", with on_device(): dev.markers[:] = xp.to_cupy(new) dev.update_holes() - with xp.profiling.count_transfers() as transfers, _count_host_reads() as reads: + reads_ctx = _count_host_reads() if count_reads else contextlib.nullcontext([]) + with xp.profiling.count_transfers() as transfers, reads_ctx as reads: dev.mpi_sort_markers() if cuda_aware: diff --git a/src/struphy/propagators/push_vxb.py b/src/struphy/propagators/push_vxb.py index cfe13d330..59adef3f1 100644 --- a/src/struphy/propagators/push_vxb.py +++ b/src/struphy/propagators/push_vxb.py @@ -152,12 +152,16 @@ def allocate(self): @profile def __call__(self, dt): - # sum up total magnetic field in the tensor-product space (the variable is extracted first) - b_full: BlockVector = self.b2_0.copy(out=self._b_full) - if self.b2_var is not None: - b_full += self._E2T.dot(self.b2_var.spline.vector, out=self._tmp) - b_full.update_ghost_regions() - b_full /= self._epsilon + # sum up total magnetic field in the tensor-product space (the variable is extracted first); without a + # magnetic field variable it is the equilibrium field alone, which does not change: computed once, so that the + # step has no ghost exchange (an MPI synchronization point) before the kernel + if self.b2_var is not None or not getattr(self, "_b_full_static_done", False): + b_full: BlockVector = self.b2_0.copy(out=self._b_full) + if self.b2_var is not None: + b_full += self._E2T.dot(self.b2_var.spline.vector, out=self._tmp) + b_full.update_ghost_regions() + b_full /= self._epsilon + self._b_full_static_done = self.b2_var is None # call pusher kernel self._pusher(dt) From aaf14a7f7733c80fb3399f9d73bccaa7cec9efa7 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 13:18:19 +0200 Subject: [PATCH 11/55] Update feectools: stencil_dot_3d without warp shuffles on the fake CuPy feectools b70cb6c (struphy-hub/feectools#104): the emulated GPU tests (fake CuPy) failed at the first stencil product. Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index 8ae781d4d..b70cb6cf6 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 8ae781d4db109824a773a7bc6f8f0f6bf114b3ac +Subproject commit b70cb6cf6077dd8ce8dde52ce4dd61d2035d58b0 From 133be44595ac33a3f987b6de8603e0e1932a2ed2 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:04:40 +0200 Subject: [PATCH 12/55] Fused gradient with boundary operators (struphy.feec.fused_derivatives) fuse_boundary_gradient(derham.grad) and fuse_boundary_gradient(derham.grad.T) apply B_out G B_in with one CUDA kernel on the GPU (fused_derivatives_cuda.cu) instead of the composed operators (a masked copy, three derivatives, a masked copy). The kernel computes the same operations in the same order, so the products are bit for bit the same; the input is read with its ghost regions exchanged, times the input mask with exchanged ghost regions, which equals the exchanged masked input in serial and in parallel. Host data and cunumpy's fake CuPy use the composed operators. ImplicitDiffusion uses it for grad and grad.T (not with the multigrid preconditioner), as on cuda-aware-mpi-fixes, where the module was imported but never committed. test_fused_derivatives: recognition of the gradient and its transpose, and the kernels against the composed operators on the fake CuPy with emulated launches (Dirichlet at both ends and at one end). Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/fused_derivatives.py | 217 ++++++++++++++++++ src/struphy/feec/fused_derivatives_cuda.cu | 97 ++++++++ .../feec/tests/test_fused_derivatives.py | 145 ++++++++++++ src/struphy/propagators/implicit_diffusion.py | 4 + 4 files changed, 463 insertions(+) create mode 100644 src/struphy/feec/fused_derivatives.py create mode 100644 src/struphy/feec/fused_derivatives_cuda.cu create mode 100644 src/struphy/feec/tests/test_fused_derivatives.py diff --git a/src/struphy/feec/fused_derivatives.py b/src/struphy/feec/fused_derivatives.py new file mode 100644 index 000000000..dd3ef5919 --- /dev/null +++ b/src/struphy/feec/fused_derivatives.py @@ -0,0 +1,217 @@ +"""The discrete gradient with its boundary operators as one CUDA kernel. + +``Derham.grad`` is ``B1 @ G @ B0.T``: the boundary operator of the 0-forms (a 0/1 mask), the derivatives of the +three directions (feectools ``DirectionalDerivativeOperator``, a difference of neighbouring coefficients) and the +boundary operator of the 1-forms. Applied as composed operators, each product makes three passes over the vectors +(and one more over the input for the masked copy), and its transpose four. :func:`fuse_boundary_gradient` returns an +operator that applies the whole chain with one kernel on the GPU (``fused_derivatives_cuda.cu``), with the same +operations in the same order, so the result is bit for bit the same, and the composed operator otherwise. +""" + +from functools import cache +from pathlib import Path + +import cunumpy as xp +import numpy as np +from feectools.feec.derivatives import DirectionalDerivativeOperator +from feectools.linalg.basic import ComposedLinearOperator, LinearOperator +from feectools.linalg.block import BlockLinearOperator, BlockVector +from feectools.linalg.stencil import StencilVector + +from struphy.feec.linear_operators import BoundaryOperator + +# entries of the kernel metadata per component, see fused_derivatives_cuda.cu +_META = 8 + + +@cache +def _kernels(): + """The forward and transposed CUDA kernels (compiled on first launch).""" + from cunumpy.kernels import CudaKernel + + path = Path(__file__).with_name("fused_derivatives_cuda.cu") + return ( + CudaKernel.from_file(path, "fused_boundary_gradient"), + CudaKernel.from_file(path, "fused_boundary_gradient_t"), + ) + + +def _derivative_blocks(block: BlockLinearOperator): + """The derivatives of a gradient block (one block column) or its transpose (one block row), else None.""" + if not isinstance(block, BlockLinearOperator): + return None + n_rows, n_cols = block.n_block_rows, block.n_block_cols + if min(n_rows, n_cols) != 1 or not 1 <= max(n_rows, n_cols) <= 3: + return None + ops = [block[k, 0] if n_cols == 1 else block[0, k] for k in range(max(n_rows, n_cols))] + if not all(isinstance(op, DirectionalDerivativeOperator) and op.domain.ndim == 3 for op in ops): + return None + transposed = n_cols > 1 or (n_rows == 1 and ops[0].transposed) + if any(op.transposed != transposed for op in ops): + return None + return ops, transposed + + +def fuse_boundary_gradient(op: LinearOperator) -> LinearOperator: + """``op`` as a :class:`FusedBoundaryGradient` if it is a gradient with boundary operators (``B_out @ G @ B_in``, + ``G`` a block of directional derivatives) or the transpose of one, else ``op`` itself. + + Parameters + ---------- + op : LinearOperator + E.g. ``derham.grad`` or ``derham.grad.T``. + + Returns + ------- + LinearOperator + The same operator; on the GPU applied with one kernel. + """ + if isinstance(op, FusedBoundaryGradient) or not isinstance(op, ComposedLinearOperator): + return op + factors = op.multiplicands + if len(factors) != 3: + return op + b_out, block, b_in = factors + if not (isinstance(b_out, BoundaryOperator) and isinstance(b_in, BoundaryOperator)): + return op + found = _derivative_blocks(block) + if found is None: + return op + derivatives, transposed = found + return FusedBoundaryGradient(op, b_out, derivatives, b_in, transposed) + + +class FusedBoundaryGradient(LinearOperator): + """``B_out @ G @ B_in`` (a gradient with its boundary operators, or its transpose) with one kernel on the GPU. + + The kernel computes the masked differences on the owned entries of the output from the input with its ghost + regions, as the composed operators do (see ``fused_derivatives_cuda.cu``). It is used for device data on a GPU; the + composed operator is applied in every other case (host data, cunumpy's fake CuPy, polar or other vectors). + + Parameters + ---------- + op : ComposedLinearOperator + The composed operator, applied when the kernel does not apply. + b_out, b_in : BoundaryOperator + The boundary operators on the output and input side. + derivatives : list[DirectionalDerivativeOperator] + The derivatives per component (of the output, or of the input if transposed). + transposed : bool + Whether ``op`` is the transpose of a gradient (1-forms to 0-forms). + """ + + def __init__(self, op, b_out, derivatives, b_in, transposed: bool): + self._op = op + self._b_out = b_out + self._b_in = b_in + self._derivatives = derivatives + self._transposed = transposed + self._meta = None + self._input_masks = None + + @property + def operator(self) -> LinearOperator: + """The composed operator.""" + return self._op + + @property + def domain(self): + return self._op.domain + + @property + def codomain(self): + return self._op.codomain + + @property + def dtype(self): + return self._op.dtype + + def tosparse(self): + return self._op.tosparse() + + def toarray(self): + return self._op.toarray() + + def transpose(self, conjugate=False): + return fuse_boundary_gradient(self._op.transpose(conjugate=conjugate)) + + def dot(self, v, out=None): + if self._use_kernel(v): + return self._dot_fused(v, out) + return self._op.dot(v, out=out) + + def _use_kernel(self, v) -> bool: + """Whether :meth:`dot` uses the kernel: device data on a real GPU (not cunumpy's fake CuPy).""" + from cunumpy.kernel_testing import fake_cupy_active + + blocks = v.blocks if isinstance(v, BlockVector) else [v] + if not all(isinstance(b, StencilVector) and xp.is_gpu(b._data) for b in blocks): + return False + return not fake_cupy_active() + + def _metadata(self): + """Owned region, direction and sign of each derivative, as the kernels take them (device array, cached).""" + if self._meta is None: + rows = [] + for d in self._derivatives: + # the owned entries of the derivative's codomain, the same padded indices in its domain + idslice = d._idslice + rows.append( + [s.start for s in idslice] + [s.stop - s.start for s in idslice] + [d.diffdir, int(d.negative)] + ) + assert all(len(r) == _META for r in rows) + self._meta = xp.asarray(np.array(rows, dtype=np.int64).ravel()) + return self._meta + + def _input_mask_arrays(self): + """The 0/1 mask of the input boundary operator per block, with its ghost regions exchanged (cached). + + The composed operators mask the input and then exchange the ghost regions of the masked vector (the boundary + operator leaves copies of boundary coefficients in the ghost regions unmasked). The kernels read the input with + its ghost regions exchanged instead, times this mask, whose ghost entries are then those of the owners: the + products are the exchanged values of the masked input, in serial and in parallel. + """ + if self._input_masks is None: + ones = self.domain.zeros() + for b in getattr(ones, "blocks", [ones]): + b._data[...] = 1.0 + ones.ghost_regions_in_sync = True + mask = self._b_in.dot(ones) + mask.update_ghost_regions() + self._input_masks = [b._data for b in getattr(mask, "blocks", [mask])] + return self._input_masks + + def _dot_fused(self, v, out=None): + """The product with the kernels (device data).""" + if out is None: + out = self.codomain.zeros() + masks_out = self._b_out._device_masks(out) + if masks_out is None: + return self._op.dot(v, out=out) + masks_in = self._input_mask_arrays() + + # the kernels read the ghost regions of the input + if not v.ghost_regions_in_sync: + v.update_ghost_regions() + + forward, transposed = _kernels() + meta = self._metadata() + n_comp = len(self._derivatives) + + def three(arrays): + return [arrays[min(k, len(arrays) - 1)] for k in range(3)] + + if not self._transposed: + outs = three([b._data for b in out.blocks]) + m_out = three([m for m, _ in masks_out]) + n_threads = sum(int(np.prod([s.stop - s.start for s in d._idslice])) for d in self._derivatives) + forward(v._data, masks_in[0], *outs, *m_out, meta, n_comp, n_threads=n_threads) + else: + ys = three([b._data for b in v.blocks]) + m_in = three(masks_in) + n_threads = int(np.prod([s.stop - s.start for s in self._derivatives[0]._idslice])) + transposed(*ys, *m_in, out._data, masks_out[0][0], meta, n_comp, n_threads=n_threads) + + # as the composed operators: the ghost regions of the output are not written + out.ghost_regions_in_sync = False + return out diff --git a/src/struphy/feec/fused_derivatives_cuda.cu b/src/struphy/feec/fused_derivatives_cuda.cu new file mode 100644 index 000000000..5fb66fbc3 --- /dev/null +++ b/src/struphy/feec/fused_derivatives_cuda.cu @@ -0,0 +1,97 @@ +// The discrete gradient with its boundary operators as one kernel (struphy.feec.fused_derivatives). +#include "cunumpy/array_view.cuh" +#include "cunumpy/index.cuh" + +// Per component c, 8 entries of `meta`: first owned index (3, padded coordinates), owned sizes (3), derivative +// direction, negative (0 or 1). +#define FUSED_META 8 + +__device__ inline Array3D pick(int c, Array3D a0, Array3D a1, Array3D a2) { + return c == 0 ? a0 : (c == 1 ? a1 : a2); +} + +/** + * out_c = B_out,c D_c B_in x for the components c of a gradient, B the 0/1 masks of the boundary operators and D_c the + * derivative in direction dir_c (feectools DirectionalDerivativeOperator): on the owned entries i of out_c (padded + * coordinates, shared by x and out_c as in DirectionalDerivativeOperator), + * + * out_c[i] = ((x[i + e] * m_in[i + e]) - (x[i] * m_in[i])) * m_out_c[i] (operands swapped if negative), + * + * the operations of the composed operators in the same order (bit for bit the same). One thread per owned entry of + * all components (n_threads = sum of the owned sizes), component 0 first; the ghost regions of out are not written. + * + * @param x, m_in Input vector data (ghost regions in sync) and the mask of the input boundary operator. + * @param out0, out1, out2, m0, m1, m2 Data of the output components and the masks of the output boundary operator. + * @param meta FUSED_META entries per component, see above. + * @param n_comp Number of components (1 to 3). + */ +extern "C" __global__ void fused_boundary_gradient(Array3D x, Array3D m_in, Array3D out0, + Array3D out1, Array3D out2, Array3D m0, + Array3D m1, Array3D m2, const long long* meta, + int n_comp) { + long long n_total = 0; + for (int c = 0; c < n_comp; ++c) + n_total += meta[FUSED_META * c + 3] * meta[FUSED_META * c + 4] * meta[FUSED_META * c + 5]; + CUNUMPY_THREAD_1D(t, n_total); + + // CUDA-only: the component and owned entry of this thread + int c = 0; + long long local = t; + while (c < n_comp - 1) { + const long long size = meta[FUSED_META * c + 3] * meta[FUSED_META * c + 4] * meta[FUSED_META * c + 5]; + if (local < size) break; + local -= size; + ++c; + } + const long long* m = meta + FUSED_META * c; + const long long i3 = m[2] + local % m[5]; + const long long i2 = m[1] + (local / m[5]) % m[4]; + const long long i1 = m[0] + local / (m[4] * m[5]); + const int dir = (int)m[6]; + const long long j1 = i1 + (dir == 0), j2 = i2 + (dir == 1), j3 = i3 + (dir == 2); + + const double a = x(j1, j2, j3) * m_in(j1, j2, j3); + const double b = x(i1, i2, i3) * m_in(i1, i2, i3); + const double d = m[7] ? b - a : a - b; + Array3D out = pick(c, out0, out1, out2); + Array3D mask = pick(c, m0, m1, m2); + out(i1, i2, i3) = d * mask(i1, i2, i3); +} + +/** + * The transpose: out = B_out (sum_c D_c^T B_in,c y_c), with D_c^T the transposed derivative; on the owned entries i of + * out, + * + * t_c = (y_c[i - e_c] * m_c[i - e_c]) - (y_c[i] * m_c[i]) (operands swapped if negative), + * out[i] = ((t_0 + t_1) + t_2) * m_out[i], + * + * summed in the order of BlockLinearOperator.dot. One thread per owned entry of out (n_threads = its owned size). + * + * @param y0, y1, y2, m0, m1, m2 Input components (ghost regions in sync) and the masks of the input boundary operator. + * @param out, m_out Output data and the mask of the output boundary operator. + * @param meta FUSED_META entries per component (the owned region of out, the same for every component). + * @param n_comp Number of components (1 to 3). + */ +extern "C" __global__ void fused_boundary_gradient_t(Array3D y0, Array3D y1, Array3D y2, + Array3D m0, Array3D m1, Array3D m2, + Array3D out, Array3D m_out, + const long long* meta, int n_comp) { + CUNUMPY_THREAD_1D(t, meta[3] * meta[4] * meta[5]); + const long long i3 = meta[2] + t % meta[5]; + const long long i2 = meta[1] + (t / meta[5]) % meta[4]; + const long long i1 = meta[0] + t / (meta[4] * meta[5]); + + double acc = 0.0; + for (int c = 0; c < n_comp; ++c) { + const long long* m = meta + FUSED_META * c; + const int dir = (int)m[6]; + const long long j1 = i1 - (dir == 0), j2 = i2 - (dir == 1), j3 = i3 - (dir == 2); + Array3D y = pick(c, y0, y1, y2); + Array3D mask = pick(c, m0, m1, m2); + const double a = y(j1, j2, j3) * mask(j1, j2, j3); + const double b = y(i1, i2, i3) * mask(i1, i2, i3); + const double d = m[7] ? b - a : a - b; + acc = c == 0 ? d : acc + d; + } + out(i1, i2, i3) = acc * m_out(i1, i2, i3); +} diff --git a/src/struphy/feec/tests/test_fused_derivatives.py b/src/struphy/feec/tests/test_fused_derivatives.py new file mode 100644 index 000000000..b6e3707d5 --- /dev/null +++ b/src/struphy/feec/tests/test_fused_derivatives.py @@ -0,0 +1,145 @@ +"""The fused gradient with boundary operators (struphy.feec.fused_derivatives) against the composed operators. + +On the host the fused operator applies the composed one. Without a GPU, the CUDA kernels run on cunumpy's fake CuPy with +the launches emulated on the CPU (in a subprocess) and must give bit for bit the products of the composed operators. +""" + +import numpy as np +import pytest + +from struphy.geometry.tests.test_domain import _cupy_installed, run_fake_cupy_child + +# (degree, boundary conditions): Dirichlet at both ends, at one end, none (periodic) +CASES = { + "dirichlet": ((2, 3, 1), (("dirichlet", "dirichlet"), None, None)), + "mixed": ((1, 2, 3), (None, ("dirichlet", "free"), None)), + "periodic": ((3, 1, 2), (None, None, None)), +} + + +def _derham(case): + from maybempi import MPI + + from struphy import domains + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + degree, bcs = CASES[case] + return Derham( + TensorProductGrid(num_elements=(6, 5, 4)), + DerhamOptions(degree=degree, bcs=bcs), + comm=MPI.COMM_WORLD, + domain=domains.Cuboid(), + ) + + +def _random(space, seed): + import cunumpy as xp + + rng = np.random.default_rng(seed) + v = space.zeros() + for b in getattr(v, "blocks", [v]): + b._data[...] = xp.asarray(rng.standard_normal(b._data.shape)) + v.update_ghost_regions() + return v + + +def _owned(v): + import cunumpy as xp + + return xp.to_numpy(v.toarray()) + + +@pytest.mark.parametrize("case", list(CASES)) +def test_fuse_recognizes_gradient(case): + """``derham.grad`` and its transpose are fused; on the host the fused operators give the composed products. + + Without essential boundary conditions the gradient has no boundary operators (only the derivatives), and it is + returned as it is.""" + from struphy.feec.fused_derivatives import FusedBoundaryGradient, fuse_boundary_gradient + + derham = _derham(case) + for op in (derham.grad, derham.grad.T): + fused = fuse_boundary_gradient(op) + if case == "periodic": + assert fused is op + continue + assert isinstance(fused, FusedBoundaryGradient) + assert fused.domain is op.domain and fused.codomain is op.codomain + v = _random(op.domain, 1) + np.testing.assert_array_equal(_owned(fused.dot(v)), _owned(op.dot(v))) + assert isinstance(fused.T, FusedBoundaryGradient) + # other operators are returned as they are + assert fuse_boundary_gradient(derham.curl) is derham.curl + + +def check_fused_gradient_on_device(): + """On the CuPy backend: the kernels give bit for bit the products of the composed operators.""" + import cunumpy as xp + + import struphy.feec.fused_derivatives as fd + from struphy.feec.fused_derivatives import fuse_boundary_gradient + + # count the launches of the two kernels: the checks below must not pass through the composed operators + kernels, launches = fd._kernels, [0, 0] + + def counted(k, i): + def launch(*args, **kwargs): + launches[i] += 1 + return k(*args, **kwargs) + + return launch + + fd._kernels = lambda: tuple(counted(k, i) for i, k in enumerate(kernels())) + try: + _check_fused_gradient(xp, fuse_boundary_gradient) + finally: + fd._kernels = kernels + assert launches[0] > 0 and launches[1] > 0, launches + + +def _check_fused_gradient(xp, fuse_boundary_gradient): + with xp.use_backend("cupy"): + for case in CASES: + derham = _derham(case) + for op in (derham.grad, derham.grad.T): + fused = fuse_boundary_gradient(op) + if fused is op: + continue + v = _random(op.domain, 2) + expected = _owned(op.dot(v)) + out = op.codomain.zeros() + result = fused._dot_fused(v, out=out) + assert result is out + assert not out.ghost_regions_in_sync + assert np.abs(expected).max() > 0, case + np.testing.assert_array_equal(_owned(result), expected, err_msg=case) + # an input whose ghost regions are not in sync + v.ghost_regions_in_sync = False + np.testing.assert_array_equal(_owned(fused._dot_fused(v)), expected, err_msg=case) + + +def check_fused_gradient_on_fake_cupy(): + """:func:`check_fused_gradient_on_device` with the CUDA kernel launches emulated on the CPU.""" + from cunumpy.kernel_testing import emulated_launches, fake_cupy_active + + from struphy.pic.tests.cuda_emulation import EMULATION_OPTIONS + + assert fake_cupy_active() + with emulated_launches(options=EMULATION_OPTIONS): + check_fused_gradient_on_device() + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +def test_fused_gradient_fake_cupy(): + """Without a GPU: the kernels against the composed operators on cunumpy's fake CuPy (serial subprocess).""" + run_fake_cupy_child( + "from struphy.feec.tests.test_fused_derivatives import check_fused_gradient_on_fake_cupy; " + "check_fused_gradient_on_fake_cupy()" + ) + + +@pytest.mark.skipif(not _cupy_installed(), reason="needs CuPy and a GPU") +def test_fused_gradient_on_gpu(): + check_fused_gradient_on_device() diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index 2993419bb..315ebf6e4 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -10,6 +10,7 @@ from scope_profiler import ProfileManager from struphy.feec import preconditioner +from struphy.feec.fused_derivatives import fuse_boundary_gradient from struphy.feec.mass import L2Projector, WeightedAverageProjection, WeightedMassOperator from struphy.io.options import LiteralOptions, OptionsBase from struphy.linear_algebra.multigrid.preconditioner import MultiGridOptions, MultiGridPreconditioner @@ -432,6 +433,9 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector: op is B1 for op in getattr(diffusion_mat, "boundary_ops", (None,)) ): mid = diffusion_mat.without_boundary_ops() + if self.options.precond != "MultiGrid": + # on the GPU, the gradient and its transpose with their boundary operators as one kernel each + grad, grad_t = fuse_boundary_gradient(grad), fuse_boundary_gradient(grad_t) self._diffusion_op = grad_t @ mid @ grad # preconditioner and solver for Ax=b From e7c4f09a8fda5792547f309ef89983c285082dc5 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:12:58 +0200 Subject: [PATCH 13/55] Device MPI sort: destinations of all markers in one kernel (struphy.pic.mpi_sort_cuda) On the GPU, mpi_sort_markers finds the destination of every marker with one kernel (sort_destinations, mpi_sort_cuda.cu) and counts the markers per destination, instead of a pass over the buffer per direction and per rank; the leaving markers are listed once, stably sorted by destination and gathered once per destination rank. The result is the one of the host path (the same markers in the same columns, the lost markers become holes); the counts are one 8-byte read per rank. STRUPHY_DEVICE_SORT=1/0 forces it on (also on the fake CuPy) or off. On cuda-aware-mpi-fixes Particles._mpi_sort_device imported this module, which was never committed; it is written for the (n_cols, n_markers) buffer of devel (#774, #775). test_mpi_sort_markers_device_kernel_fake_cupy: the kernel (launches emulated) against the NumPy sort, CUDA-aware and staged, periodic and removing boundaries (checked on 2 and 4 ranks). Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/base.py | 172 +++++++++++++++++-- src/struphy/pic/mpi_sort_cuda.cu | 79 +++++++++ src/struphy/pic/mpi_sort_cuda.py | 99 +++++++++++ src/struphy/pic/tests/test_sorting_device.py | 42 ++++- 4 files changed, 376 insertions(+), 16 deletions(-) create mode 100644 src/struphy/pic/mpi_sort_cuda.cu create mode 100644 src/struphy/pic/mpi_sort_cuda.py diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 29cc95d4a..1839c06aa 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1863,24 +1863,29 @@ def mpi_sort_markers( if isinstance(alpha, int) or isinstance(alpha, float): alpha = (alpha, alpha, alpha) - # create new markers_to_be_sent array and make corresponding holes in markers array - hole_inds_after_send, send_inds = self._sendrecv_determine_mtbs(alpha=alpha) + if self._use_device_sort(): + # destinations in one kernel pass, one gather per destination rank (see _mpi_sort_device) + n_emptied, recv_info = self._mpi_sort_device(alpha) + else: + # create new markers_to_be_sent array and make corresponding holes in markers array + hole_inds_after_send, send_inds = self._sendrecv_determine_mtbs(alpha=alpha) - # determine where to send markers_to_be_sent - send_info = self._sendrecv_get_destinations(send_inds) + # determine where to send markers_to_be_sent + send_info = self._sendrecv_get_destinations(send_inds) - # set new holes in markers array to -1 - self._markers[:, send_inds] = -1.0 + # set new holes in markers array to -1 + self._markers[:, send_inds] = -1.0 - # transpose send_info - recv_info = self._sendrecv_all_to_all(send_info) + # transpose send_info + recv_info = self._sendrecv_all_to_all(send_info) - # send and receive markers - self._sendrecv_markers(recv_info, hole_inds_after_send) + # send and receive markers + self._sendrecv_markers(recv_info, hole_inds_after_send) + n_emptied = int(send_inds.size) # the sent markers are holes now, the received ones filled the lowest holes (host counts, no device read) n_received = int(np.sum(recv_info)) - int(recv_info[self.mpi_rank]) - self._track_markers(n_emptied=int(send_inds.size), n_filled=n_received) + self._track_markers(n_emptied=n_emptied, n_filled=n_received) # new holes and new number of holes and markers on process self.update_holes() @@ -4649,6 +4654,151 @@ def _sendrecv_determine_mtbs( return hole_inds_after_send, send_inds + def _use_device_sort(self) -> bool: + """Whether :meth:`mpi_sort_markers` uses :meth:`_mpi_sort_device`: markers on a GPU (not on cunumpy's fake CuPy, + which cannot launch kernels). ``STRUPHY_DEVICE_SORT=1`` or ``0`` forces it on (also on the fake CuPy, e.g. in + tests with emulated launches) or off.""" + forced = os.environ.get("STRUPHY_DEVICE_SORT", "").strip().lower() + if forced in ("0", "false", "no"): + return False + if xp.get_array_backend(self._markers) != "cupy": + return False + if forced in ("1", "true", "yes"): + return True + from cunumpy.kernel_testing import fake_cupy_active + + return not fake_cupy_active() + + def _mpi_sort_device(self, alpha) -> tuple[int, np.ndarray]: + """The exchange of :meth:`mpi_sort_markers` for markers on the GPU. + + The same result as :meth:`_sendrecv_determine_mtbs`, :meth:`_sendrecv_get_destinations` and + :meth:`_sendrecv_markers` (the same markers in the same columns), with fewer passes over the marker buffer: + + - one kernel finds the destination of every marker (:func:`~struphy.pic.mpi_sort_cuda.sort_destinations`) + and counts the markers per destination, instead of a pass per direction and per rank; + - the leaving markers are listed once, stably sorted by destination (so the markers to each rank keep their + order, as before) and gathered once per destination rank. + + Markers that leave towards no rank's domain (e.g. a NaN position) become holes and are not sent, as before. + + Parameters + ---------- + alpha : tuple + Weights of the sorting position, see :meth:`mpi_sort_markers`. + + Returns + ------- + n_emptied : int + Number of markers that left (became holes), for :meth:`_track_markers`. + recv_info : numpy.ndarray + Number of markers received from each rank. + """ + from struphy.pic import mpi_sort_cuda + + alpha_host = np.array(alpha, dtype=float).ravel() + assert alpha_host.size == 3 + assert np.all(alpha_host >= 0.0) and np.all(alpha_host <= 1.0) + + if not hasattr(self, "_dev_sort_bufs"): + ranks = np.array(self._neighbor_ranks + self._non_neighbor_ranks, dtype=np.int32) + periodic = np.array([1 if axis in self._periodic_axes else 0 for axis in range(3)], dtype=np.int32) + bounds = mpi_sort_cuda.domain_bounds(self._domain_array_host) + self._dev_sort_bufs = { + "ranks": xp.asarray(ranks), + "periodic": xp.asarray(periodic), + "bounds": xp.asarray(bounds), + "dest": xp.empty(self.n_markers_max, dtype=np.int32), + "counts": xp.zeros(self.mpi_size + 2, dtype=np.uint64), + "alpha": xp.zeros(3, dtype=float), + } + bufs = self._dev_sort_bufs + use_alpha = not np.all(alpha_host == 1.0) + if use_alpha: + bufs["alpha"][:] = xp.asarray(alpha_host) + + n = self.n_markers_used + bi = self.first_pusher_idx + dest = bufs["dest"][:n] + counts_dev = bufs["counts"] + counts_dev.fill(0) + mpi_sort_cuda.sort_destinations( + self._markers, + self.holes, + self.ghost_particles, + bi + 3 + self.vdim, + bi, + bufs["alpha"], + use_alpha, + bufs["periodic"], + bufs["bounds"], + bufs["ranks"], + self.mpi_rank, + n, + bufs["dest"], + counts_dev, + ) + # markers per destination (lost first): a read of 8 bytes per rank to size the messages, like the one the + # nonzero() of the host path makes to size its result (marker data stays on the device) + counts = xp.to_numpy(counts_dev).astype(int) + n_lost = int(counts[0]) + send_info = counts[2:].copy() + + # the leaving markers, sorted by destination (stable: ascending columns per destination); lost markers first + leaving = dest != -1 + leave_inds = xp.flatnonzero(leaving) + n_leaving = n_lost + int(send_info.sum()) + if n_leaving: + leave_sorted = leave_inds[xp.argsort(dest[leave_inds], kind="stable")] + else: + leave_sorted = leave_inds + + # the columns that are holes after the send: the old holes (all columns from n_markers_used on are holes) + # and the leaving markers; the received markers fill the lowest of them + holes_after = self.holes.copy() + xp.logical_or(holes_after[:n], leaving, out=holes_after[:n]) + hole_inds_after_send = xp.flatnonzero(holes_after) + + # one contiguous buffer (n_cols, count) per destination rank, in the order of the host path + first_send = np.cumsum(send_info) - send_info + sends = [] + for rank in range(self.mpi_size): + if rank == self.mpi_rank or send_info[rank] == 0: + self._send_list[rank] = self._markers[:, :0] + continue + start = n_lost + int(first_send[rank]) + cols = leave_sorted[start : start + int(send_info[rank])] + self._send_list[rank] = xp.ascontiguousarray(self._markers[:, cols]) + sends.append((self._send_list[rank], rank, self.mpi_rank)) + + # the leaving markers are holes now (also the lost ones) + if n_leaving: + self._markers[:, leave_inds] = -1.0 + + recv_info = self._sendrecv_all_to_all(send_info) + n_recv = int(recv_info.sum()) + if hole_inds_after_send.size < n_recv: + warnings.warn( + f"Strong load imbalance detected: number of holes ({hole_inds_after_send.size}) on rank " + f"{self.mpi_rank} is smaller than number of incoming particles ({n_recv}). " + 'Increase "bufsize" in the markers parameters for the next run.', + ) + self.mpi_comm.Abort() + + first_recv = np.cumsum(recv_info) - recv_info + receives = [] + for rank in range(self.mpi_size): + if rank != self.mpi_rank and recv_info[rank] > 0: + receives.append((xp.empty((self.n_cols, int(recv_info[rank])), dtype=self._markers.dtype), rank, rank)) + + xp.mpi.exchange(self.mpi_comm, sends=sends, receives=receives, cuda_aware=self._mpi_cuda_aware()) + + for recvbuf, rank, _ in receives: + start = int(first_recv[rank]) + self._markers[:, hole_inds_after_send[start : start + int(recv_info[rank])]] = recvbuf + + return n_leaving, recv_info + def _compute_neighbor_ranks(self) -> tuple[list[int], list[int]]: """Split every other rank into geometric neighbours of this rank's sub-domain box and everyone else, for :meth:`_sendrecv_get_destinations`. diff --git a/src/struphy/pic/mpi_sort_cuda.cu b/src/struphy/pic/mpi_sort_cuda.cu new file mode 100644 index 000000000..ff9c3d1e5 --- /dev/null +++ b/src/struphy/pic/mpi_sort_cuda.cu @@ -0,0 +1,79 @@ +// Destinations of the markers in Particles.mpi_sort_markers on the GPU (struphy.pic.mpi_sort_cuda). +#include "cunumpy/array_view.cuh" +#include "cunumpy/atomic.cuh" +#include "cunumpy/index.cuh" + +// The domain test of Particles._is_on_domain_of_rank for one rank: [left, right) in every direction. +__device__ inline bool on_domain(const double* eta, const double* bounds, int r) { + for (int a = 0; a < 3; ++a) { + const double left = bounds[(r * 3 + a) * 2], right = bounds[(r * 3 + a) * 2 + 1]; + if (!(eta[a] >= left && eta[a] < right)) return false; + } + return true; +} + +/** + * The rank each marker has to go to, as Particles._sendrecv_determine_mtbs and _sendrecv_get_destinations decide it. + * + * Marker ip (a column of the buffer, ip < n_used) is sorted by eta = pos + shift (use_alpha == 0) or + * alpha * (pos + shift) + (1 - alpha) * init per direction (rows 0-2, shift_row.., init_row..), reduced modulo 1 as + * numpy.mod does; on a non-periodic axis eta = 1 stays 1. It stays (dest -1) if it is a hole or a ghost or lies in + * the domain of this rank, else it goes to the rank whose domain contains eta (bounds: [left, right) per rank and + * direction, the right end 1 shifted to the next float, see domain_bounds), or it is lost (dest -2) if no domain + * contains it (e.g. a NaN position). The domains do not overlap, so the rank is unique. + * + * One thread per marker (n_threads = n_used). counts[dest + 2] counts the leaving markers (dest -2 or a rank). + * + * @param markers Marker buffer (n_cols x n_markers_max, one column per marker). + * @param holes, ghosts Hole and ghost masks (n_markers_max). + * @param shift_row, init_row First rows of the boundary shifts and of the initial positions. + * @param alpha, use_alpha The three weights, used if use_alpha (else alpha = 1, eta = pos + shift). + * @param periodic Per direction, 1 if periodic. + * @param bounds (n_ranks, 3, 2): left and right end of each rank's domain per direction. + * @param ranks The other ranks in the order they are checked (neighbours first), n_others entries. + * @param n_others Number of entries of ranks. + * @param rank This rank. + * @param n_used Number of markers (columns) that can hold markers. + * @param dest Output (n_markers_max): -1 stay, -2 lost, else the destination rank (written for ip < n_used). + * @param counts Output (n_ranks + 2), added to: leaving markers per destination (dest + 2). + */ +extern "C" __global__ void sort_destinations(Array2D markers, const bool* holes, const bool* ghosts, + int shift_row, int init_row, const double* alpha, int use_alpha, + const int* periodic, + const double* bounds, const int* ranks, int n_others, int rank, + long long n_used, int* dest, unsigned long long* counts) { + CUNUMPY_THREAD_1D(ip, n_used); + if (holes[ip] || ghosts[ip]) { + dest[ip] = -1; + return; + } + + double eta[3]; + for (int a = 0; a < 3; ++a) { + const double moved = markers(a, ip) + markers(shift_row + a, ip); + double e = use_alpha ? alpha[a] * moved + (1.0 - alpha[a]) * markers(init_row + a, ip) : moved; + const bool right_end = !periodic[a] && e == 1.0; + // numpy.mod(e, 1.0): the remainder with the sign of the divisor + double r = fmod(e, 1.0); + if (r != 0.0) { + if (r < 0.0) r += 1.0; + } else { + r = 0.0; + } + eta[a] = right_end ? 1.0 : r; + } + + if (on_domain(eta, bounds, rank)) { + dest[ip] = -1; + return; + } + int d = -2; + for (int k = 0; k < n_others; ++k) { + if (on_domain(eta, bounds, ranks[k])) { + d = ranks[k]; + break; + } + } + dest[ip] = d; + cunumpy_atomic_add(&counts[d + 2], 1ULL); +} diff --git a/src/struphy/pic/mpi_sort_cuda.py b/src/struphy/pic/mpi_sort_cuda.py new file mode 100644 index 000000000..a862d8831 --- /dev/null +++ b/src/struphy/pic/mpi_sort_cuda.py @@ -0,0 +1,99 @@ +"""The destinations of :meth:`~struphy.pic.base.Particles.mpi_sort_markers` on the GPU (``mpi_sort_cuda.cu``). + +:meth:`~struphy.pic.base.Particles._mpi_sort_device` finds the rank of every leaving marker with one kernel +(:func:`sort_destinations`) instead of a pass over the buffer per rank and per direction. +""" + +from functools import cache +from pathlib import Path + +import numpy as np + + +@cache +def _kernel(): + from cunumpy.kernels import CudaKernel + + return CudaKernel.from_file(Path(__file__).with_name("mpi_sort_cuda.cu"), "sort_destinations") + + +def domain_bounds(domain_array) -> np.ndarray: + """The ends of each rank's domain per direction, as :meth:`~struphy.pic.base.Particles._is_on_domain_of_rank` tests + them: an array (n_ranks, 3, 2) of [left, right), with the right end 1 shifted to the next float (eta = 1 belongs to + the last rank). + + Parameters + ---------- + domain_array : array + (n_ranks, 9) host array of :attr:`~struphy.pic.base.Particles.domain_array` (left, right, cells per direction). + """ + domain_array = np.asarray(domain_array, dtype=float) + left = domain_array[:, 0::3] + right = domain_array[:, 1::3] + right = np.where(right == 1.0, np.nextafter(right, 2.0), right) + return np.ascontiguousarray(np.stack([left, right], axis=-1)) + + +def sort_destinations( + markers, + holes, + ghosts, + shift_row: int, + init_row: int, + alpha, + use_alpha: bool, + periodic, + bounds, + ranks, + rank: int, + n_used: int, + dest, + counts, +): + """Launch the kernel ``sort_destinations`` (one thread per marker, see ``mpi_sort_cuda.cu``). + + Parameters + ---------- + markers : cupy.ndarray + Marker buffer (n_cols, n_markers_max). + holes, ghosts : cupy.ndarray + Hole and ghost masks (bool, n_markers_max). + shift_row, init_row : int + First rows of the boundary shifts and of the initial positions. + alpha : cupy.ndarray + The three sorting weights (float64, 3). + use_alpha : bool + Whether to use them; else the sorting position is position + shift (alpha = 1, computed exactly). + periodic : cupy.ndarray + Per direction, 1 if periodic (int32, 3). + bounds : cupy.ndarray + :func:`domain_bounds` on the device (float64). + ranks : cupy.ndarray + The other ranks, neighbours first (int32). + rank, n_used : int + This rank and the number of markers (columns) that can hold markers. + dest : cupy.ndarray + Output, int32 (n_markers_max): -1 stay, -2 lost, else the destination rank. + counts : cupy.ndarray + Output, uint64 (n_ranks + 2), added to: leaving markers per destination + 2. + """ + if n_used == 0: + return + _kernel()( + markers, + holes, + ghosts, + int(shift_row), + int(init_row), + alpha, + int(bool(use_alpha)), + periodic, + bounds, + ranks, + int(ranks.size), + int(rank), + int(n_used), + dest, + counts, + n_threads=int(n_used), + ) diff --git a/src/struphy/pic/tests/test_sorting_device.py b/src/struphy/pic/tests/test_sorting_device.py index aa7c4e5c0..e66af628c 100644 --- a/src/struphy/pic/tests/test_sorting_device.py +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -12,6 +12,7 @@ """ import contextlib +import os import sys from contextlib import contextmanager @@ -90,12 +91,31 @@ def _shift_positions(markers: np.ndarray, valid: np.ndarray, step: int) -> np.nd return new -def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", count_reads: bool = True): +def check_mpi_sort_markers_device( + cuda_aware: bool | None, bc: str = "periodic", count_reads: bool = True, device_sort: bool = False +): """Sort the same markers on NumPy and CuPy across all ranks and compare the marker arrays. ``cuda_aware=None`` uses the probe (:func:`cunumpy.mpi.mpi_is_cuda_aware`); ``True``/``False`` - force the device-buffer or the host-staged exchange. + force the device-buffer or the host-staged exchange. With ``device_sort`` the CuPy markers are sorted with the + destination kernel (:meth:`Particles._mpi_sort_device`, launches emulated on the fake CuPy), which reads the + counts per destination (one copy to the host per sort). """ + if device_sort: + import os + + from cunumpy.kernel_testing import emulated_launches + + from struphy.pic.tests.cuda_emulation import EMULATION_OPTIONS + + os.environ["STRUPHY_DEVICE_SORT"] = "1" + try: + with emulated_launches(options=EMULATION_OPTIONS): + check_mpi_sort_markers_device(cuda_aware, bc, count_reads=False) + finally: + del os.environ["STRUPHY_DEVICE_SORT"] + return + comm = MPI.COMM_WORLD rank, size = comm.Get_rank(), comm.Get_size() @@ -143,16 +163,19 @@ def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", with xp.profiling.count_transfers() as transfers, reads_ctx as reads: dev.mpi_sort_markers() + # the destination kernel reads its counts per destination (8 bytes per rank) once per sort + with on_device(): + counts_read = 1 if dev._use_device_sort() else 0 if cuda_aware: - assert transfers.total == 0, transfers.report() + assert transfers.total == counts_read, transfers.report() else: # exchange stages each non-empty send and receive buffer separately n_dest = sum(1 for i in range(size) if i != rank and dev._send_list[i].shape[1] > 0) n_sources = sum(1 for i in range(size) if i != rank and ref._recvbufs[i].shape[1] > 0) - assert transfers.to_host == n_dest, transfers.report() + assert transfers.to_host == n_dest + counts_read, transfers.report() assert transfers.to_device == n_sources, transfers.report() sent_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_host") - assert sent_bytes == n_send * ref.n_cols * 8 + assert sent_bytes == n_send * ref.n_cols * 8 + counts_read * (size + 2) * 8 received_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_device") assert received_bytes == n_recv * ref.n_cols * 8 @@ -183,6 +206,15 @@ def test_mpi_sort_markers_fake_cupy(cuda_aware, bc): check_mpi_sort_markers_device(cuda_aware, bc=bc) +@requires_fake_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("cuda_aware", [True, False]) +@pytest.mark.parametrize("bc", ["periodic", "remove"]) +def test_mpi_sort_markers_device_kernel_fake_cupy(cuda_aware, bc): + """The destination kernel (Particles._mpi_sort_device) against the NumPy sort, launches emulated.""" + check_mpi_sort_markers_device(cuda_aware, bc=bc, device_sort=True) + + @requires_cupy @pytest.mark.mpi @pytest.mark.parametrize("bc", ["periodic", "remove"]) From 8c7d0c62af43d690aecaf72394b864c8c9409a87 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:15:37 +0200 Subject: [PATCH 14/55] Update feectools: CachedLaunch for the stencil axpy and inner kernels feectools 1d8d47c (struphy-hub/feectools#104). Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index b70cb6cf6..1d8d47c4a 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit b70cb6cf6077dd8ce8dde52ce4dd61d2035d58b0 +Subproject commit 1d8d47c4a3bc8e0b594992cc22944b1ac4230c0f From c508cce8bf3202436abdac637029b614b9f248a9 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:20:50 +0200 Subject: [PATCH 15/55] Kronecker mass: launch the band-pass kernel through CachedLaunch The banded passes run with the same buffers in every solver iteration (as on cuda-aware-mpi-fixes, now that feectools has CachedLaunch). Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/kronecker_mass.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py index 2050edb82..5456b2eeb 100644 --- a/src/struphy/feec/kronecker_mass.py +++ b/src/struphy/feec/kronecker_mass.py @@ -38,8 +38,10 @@ def kronecker_mass_enabled() -> bool: def _band_pass_cuda(): """The CUDA kernel of one banded pass (compiled on first call).""" from cunumpy.kernels import CudaKernel + from feectools.utilities.cached_launch import CachedLaunch - return CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass") + # launched with the same buffers in every solver iteration: cache the launch preparation + return CachedLaunch(CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass")) def _band_pass(A, inp, out, axis, offset): From 0166e833806bbb71af1c7339a1b2f3c2aa4c1803 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:41:05 +0200 Subject: [PATCH 16/55] pcg_sr also for the multigrid and Hiptmair-Xu preconditioners of the shear Alfven solves ShearAlfven and ShearAlfvenCurlCurl (#778, #780) accepted these preconditioners only with solver='pcg'; pcg_sr takes a preconditioner as well, as in the other propagators. Co-Authored-By: Claude Opus 5.5 --- src/struphy/propagators/shear_alfven_curl_curl_propagator.py | 2 +- src/struphy/propagators/shear_alfven_propagator.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/struphy/propagators/shear_alfven_curl_curl_propagator.py b/src/struphy/propagators/shear_alfven_curl_curl_propagator.py index 80d21555e..729ae0f95 100644 --- a/src/struphy/propagators/shear_alfven_curl_curl_propagator.py +++ b/src/struphy/propagators/shear_alfven_curl_curl_propagator.py @@ -185,7 +185,7 @@ def __post_init__(self): if self.algo == "explicit": self.precond = "MassMatrixPreconditioner" else: - assert self.solver == "pcg", 'precond="HiptmairXu" requires solver="pcg".' + assert self.solver in ("pcg", "pcg_sr"), 'precond="HiptmairXu" requires solver="pcg" or "pcg_sr".' # defaults if self.solver_params is None: diff --git a/src/struphy/propagators/shear_alfven_propagator.py b/src/struphy/propagators/shear_alfven_propagator.py index 70d4cd293..cc6004ec4 100644 --- a/src/struphy/propagators/shear_alfven_propagator.py +++ b/src/struphy/propagators/shear_alfven_propagator.py @@ -148,7 +148,7 @@ def __post_init__(self): check_option(self.precond, LiteralOptions.OptsShearAlfvenPrecond) if self.precond in ("MultiGrid", "StiffnessPreconditioner"): assert self.algo == "implicit", f'precond="{self.precond}" requires algo="implicit".' - assert self.solver == "pcg", f'precond="{self.precond}" requires solver="pcg".' + assert self.solver in ("pcg", "pcg_sr"), f'precond="{self.precond}" requires solver="pcg" or "pcg_sr".' if self.precond == "StiffnessPreconditioner": assert self.u_space == "Hdiv", 'precond="StiffnessPreconditioner" requires u_space="Hdiv".' if self.precond == "MultiGrid" and self.multigrid is None: From 6bb532fd1b63bd29166b84ffeb7d390d21fa10c9 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:00:05 +0200 Subject: [PATCH 17/55] CuPy backend: host/device conversions in post-processing; math.sqrt for host scalars in Species Post-processing evaluates fields on the device and writes them from the host: the evaluation grids, point data and coordinates are converted where they cross, and HDF5 data is moved to the device before it is written into the spline coefficients. Co-Authored-By: Claude Opus 5.5 --- src/struphy/models/species.py | 4 ++-- src/struphy/post_processing/output.py | 31 ++++++++++++++++++--------- 2 files changed, 23 insertions(+), 12 deletions(-) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 3c8b47657..4bc2dcac1 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,9 +1,9 @@ import logging +import math import warnings from abc import ABCMeta, abstractmethod from dataclasses import fields, is_dataclass -import cunumpy as xp from maybempi import MPI from struphy.models.variables import Variable @@ -203,7 +203,7 @@ def __init__( con = ConstantsOfNature() # relevant frequencies - om_p = xp.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) + om_p = math.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) om_c = Z * con.e * units.B / (A * con.mH) # compute equation parameters diff --git a/src/struphy/post_processing/output.py b/src/struphy/post_processing/output.py index bfcf5fd58..cc65d56f8 100644 --- a/src/struphy/post_processing/output.py +++ b/src/struphy/post_processing/output.py @@ -85,6 +85,13 @@ def load_plasma_plots() -> bool: return True +def _to_host(value): + """`value` (an array, or a list/tuple of arrays for vector-valued fields) as NumPy array(s).""" + if isinstance(value, (list, tuple)): + return type(value)(_to_host(v) for v in value) + return xp.to_numpy(value) + + class ProductMapping(Mapping[str, xr.DataArray]): """A discoverable mapping whose products are loaded on first access.""" @@ -568,13 +575,14 @@ def _attach_physical_coords( self, array: xr.DataArray, etas: tuple[Any, Any, Any], dims: tuple[str, ...] ) -> xr.DataArray: """Map a raw FEEC evaluation grid and attach its Cartesian coordinates.""" - grids = [np.atleast_1d(np.asarray(eta, dtype=float)) for eta in etas] - mapped = self.domain(*grids) + # the logical grids may be device arrays (CuPy backend): map them where they live, keep the coordinates on the host + grids = [np.atleast_1d(np.asarray(xp.to_numpy(eta), dtype=float)) for eta in etas] + mapped = self.domain(*(xp.asarray(grid) for grid in grids)) shape = tuple(len(grid) for grid in grids) keep = tuple(slice(None) if dim in dims else 0 for dim in ("eta1", "eta2", "eta3")) coordinates = {} for name, values in zip(("X", "Y", "Z"), mapped): - values = np.asarray(values).reshape(shape)[keep] + values = np.asarray(xp.to_numpy(values)).reshape(shape)[keep] coordinates[name] = (dims, values) if dims else values.item() return array.assign_coords(coordinates) @@ -1295,7 +1303,9 @@ def _process_fields( # the physical grid is only needed for output, hence it is only built on rank 0 if self._pproc_rank == 0: - grids_phy = list(self.domain(*grids_log)) + # output coordinates, on the host (the logical grids stay on the device for the field evaluation) + grids_phy = [xp.to_numpy(grid) for grid in self.domain(*grids_log)] + grids_log_host = [xp.to_numpy(grid) for grid in grids_log] else: grids_phy = None @@ -1327,8 +1337,9 @@ def _process_fields( if self._pproc_rank == 0: for species, vars in vals.items(): for name, val in vars.items(): - point_data[species][name][t] = val - point_data_phy[species][name][t] = vals_phy[species][name] + # evaluated on the device on the CuPy backend; written to files from the host + point_data[species][name][t] = _to_host(val) + point_data_phy[species][name][t] = _to_host(vals_phy[species][name]) # directory for the vtk files path_fields = os.path.join(self.path_pproc, "fields_data") @@ -1338,10 +1349,10 @@ def _process_fields( for species, vars in point_data.items(): variables = {} for name, val in vars.items(): - variables[name] = wrap_field_data(val, grids_log, grids_phy=grids_phy, name=name) + variables[name] = wrap_field_data(val, grids_log_host, grids_phy=grids_phy, name=name) if physical: variables[name + "_xyz"] = wrap_field_data( - point_data_phy[species][name], grids_log, grids_phy=grids_phy, name=name + "_xyz" + point_data_phy[species][name], grids_log_host, grids_phy=grids_phy, name=name + "_xyz" ) store.write_group(store.store_path(self.path_pproc), f"/{species}", xr.Dataset(variables)) @@ -1517,7 +1528,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[n * step, p1:-p1, p2:-p2, p3:-p3]) # HDF5 data is on the host # vector-valued field else: @@ -1530,7 +1541,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3]) vector.update_ghost_regions() if is_polar: From 5e2dced2cb8186408654a37a6f0484b4f920b676 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:07:26 +0200 Subject: [PATCH 18/55] CUDA version of the mass-matrix assembly kernel kernel_3d_mat - kernel_3d_mat_cuda.cu: one thread per (element, codomain basis function, domain basis function); replaces the host fallback on the CuPy backend. - WeightedMassOperator.assemble reads the quadrature grids (spans, weights, basis values) as backend arrays, copied to the device once and cached. - Simulation.allocate compiles the assembly kernel before it is first called, so the compilation is not counted in the assembly's profiling regions. Co-Authored-By: Claude Opus 5.5 --- .../feec/kernels/kernel_3d_mat/__init__.py | 18 ++-- .../kernel_3d_mat/kernel_3d_mat_cuda.cu | 92 +++++++++++++++++++ src/struphy/feec/mass.py | 21 ++++- src/struphy/simulation/sim.py | 18 ++++ 4 files changed, 139 insertions(+), 10 deletions(-) create mode 100644 src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu diff --git a/src/struphy/feec/kernels/kernel_3d_mat/__init__.py b/src/struphy/feec/kernels/kernel_3d_mat/__init__.py index a881301b4..76da51673 100644 --- a/src/struphy/feec/kernels/kernel_3d_mat/__init__.py +++ b/src/struphy/feec/kernels/kernel_3d_mat/__init__.py @@ -1,8 +1,4 @@ -"""The kernel of this folder: pyccel in ``_kernels.py``, CUDA in ``_cuda.cu`` once ported. - -No CUDA version yet: on the CuPy backend the pyccel kernel runs on host copies of the arrays (a warning is -emitted once), as before the move into this folder. -""" +"""The kernel of this folder: pyccel in ``kernel_3d_mat_kernels.py``, CUDA in ``kernel_3d_mat_cuda.cu``.""" from cunumpy.kernels import Kernel @@ -11,6 +7,16 @@ # the arguments the kernel writes to, by position: data OUTPUTS = (25,) + +def _n_threads(args): + """One thread per (element, codomain basis function, domain basis function), see kernel_3d_mat_cuda.cu.""" + n_elements = args[0].size * args[1].size * args[2].size + n_pairs = 1 + for degree in args[3:9]: # pi1, pi2, pi3, pj1, pj2, pj3 + n_pairs *= int(degree) + 1 + return n_elements * n_pairs + + kernel_3d_mat = Kernel.from_folder( - __name__, structs=CUDA_STRUCTS, missing_cuda="fallback", host_options={"outputs": OUTPUTS} + __name__, structs=CUDA_STRUCTS, host_options={"outputs": OUTPUTS}, n_threads_from=_n_threads ) diff --git a/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu b/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu new file mode 100644 index 000000000..6afab9398 --- /dev/null +++ b/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu @@ -0,0 +1,92 @@ +// CUDA version of kernel_3d_mat (kernel_3d_mat_kernels.py), with the same arguments in the same order. +#include "cunumpy/array_view.cuh" +#include "cunumpy/atomic.cuh" +#include "cunumpy/index.cuh" + +/** + * Integral of Lambda_i * mat_fun * Lambda_j over the local elements, added to the stencil matrix data, as in + * kernel_3d_mat_kernels.kernel_3d_mat. + * + * One thread per (element, codomain basis function (il1, il2, il3), domain basis function (jl1, jl2, jl3)): + * n_threads = ne1 * ne2 * ne3 * (pi1 + 1) * (pi2 + 1) * (pi3 + 1) * (pj1 + 1) * (pj2 + 1) * (pj3 + 1) (see + * __init__.py), the element index varying slowest and jl3 fastest. Each thread computes the quadrature sum of its + * pair exactly as the pyccel kernel (same factors, same order) and adds it to its entry of `data` with an atomic + * addition: neighbouring elements contribute to the same entries, so only the summation order over elements + * differs from pyccel (round-off). + * + * @param spans1, spans2, spans3 Knot span index of each local element, per direction. + * @param pi1, pi2, pi3 Degrees of the codomain basis functions. + * @param pj1, pj2, pj3 Degrees of the domain basis functions. + * @param starts1, starts2, starts3 Global start index of the rows on this process. + * @param pads1, pads2, pads3 Paddings (ghost regions) of `data`. + * @param w1, w2, w3 Quadrature weights [element, quadrature point]. + * @param bi1, bi2, bi3 Codomain basis functions [element, local basis function, derivative, quadrature point]. + * @param bj1, bj2, bj3 Domain basis functions, same layout. + * @param mat_fun Weight function at the quadrature points [flattened point 1, flattened point 2, flattened point 3]. + * @param data StencilMatrix data (rows with ghost regions..., diagonals...); the results are added (not set). + */ +extern "C" __global__ void kernel_3d_mat(Array1D spans1, Array1D spans2, + Array1D spans3, int pi1, int pi2, int pi3, int pj1, int pj2, + int pj3, int starts1, int starts2, int starts3, int pads1, int pads2, + int pads3, Array2D w1, Array2D w2, Array2D w3, + Array4D bi1, Array4D bi2, Array4D bi3, + Array4D bj1, Array4D bj2, Array4D bj3, + Array3D mat_fun, Array6D data) +{ + const long long ne1 = spans1.size(); + const long long ne2 = spans2.size(); + const long long ne3 = spans3.size(); + + const long long ni1 = pi1 + 1, ni2 = pi2 + 1, ni3 = pi3 + 1; + const long long nj1 = pj1 + 1, nj2 = pj2 + 1, nj3 = pj3 + 1; + const long long n_pairs = ni1 * ni2 * ni3 * nj1 * nj2 * nj3; + + CUNUMPY_THREAD_1D(thread, ne1 * ne2 * ne3 * n_pairs); + + // CUDA-only: split the thread index into the pyccel loop variables (iel1, iel2, iel3, il1, il2, il3, jl1, jl2, + // jl3), jl3 fastest + long long rest = thread; + const long long jl3 = rest % nj3; rest /= nj3; + const long long jl2 = rest % nj2; rest /= nj2; + const long long jl1 = rest % nj1; rest /= nj1; + const long long il3 = rest % ni3; rest /= ni3; + const long long il2 = rest % ni2; rest /= ni2; + const long long il1 = rest % ni1; rest /= ni1; + const long long iel3 = rest % ne3; rest /= ne3; + const long long iel2 = rest % ne2; rest /= ne2; + const long long iel1 = rest; + + const long long nq1 = w1.shape[1]; + const long long nq2 = w2.shape[1]; + const long long nq3 = w3.shape[1]; + + double value = 0.0; + + for (long long q1 = 0; q1 < nq1; ++q1) { + for (long long q2 = 0; q2 < nq2; ++q2) { + for (long long q3 = 0; q3 < nq3; ++q3) { + const double wvol = w1(iel1, q1) * w2(iel2, q2) * w3(iel3, q3) * + mat_fun(iel1 * nq1 + q1, iel2 * nq2 + q2, iel3 * nq3 + q3); + + const double bi = bi1(iel1, il1, 0, q1) * bi2(iel2, il2, 0, q2) * bi3(iel3, il3, 0, q3); + const double bj = bj1(iel1, jl1, 0, q1) * bj2(iel2, jl2, 0, q2) * bj3(iel3, jl3, 0, q3); + + value += wvol * bi * bj; + } + } + } + + // global spline indices + const long long i_global1 = spans1(iel1) - pi1 + il1; + const long long i_global2 = spans2(iel2) - pi2 + il2; + const long long i_global3 = spans3(iel3) - pi3 + il3; + + // local spline indices (- starts --> can be negative, will therefore be written to ghost regions) + const long long i_local1 = i_global1 - starts1; + const long long i_local2 = i_global2 - starts2; + const long long i_local3 = i_global3 - starts3; + + cunumpy_atomic_add(&data(pads1 + i_local1, pads2 + i_local2, pads3 + i_local3, pads1 + jl1 - il1, + pads2 + jl2 - il2, pads3 + jl3 - il3), + value); +} diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index f75c6de17..e035e4fe4 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -2242,6 +2242,19 @@ def transpose(self, conjugate=False): return M + def _quad_grid(self, space_name, attribute, component): + """Quadrature-grid arrays (spans, weights or basis values) of one component of a space, as backend arrays. + + The Derham keeps them as host arrays; on the CuPy backend the assembly kernel needs them on the device, + so they are copied there once and cached (they do not change between assemblies). + """ + cache = self.__dict__.setdefault("_quad_grid_cache", {}) + key = (space_name, attribute, component) + if key not in cache: + arrays = getattr(self.derham.spline_attributes[space_name], attribute)[component] + cache[key] = [xp.asarray(array) for array in arrays] + return cache[key] + def assemble(self, weights=None, clear=True): r""" Assembles the weighted mass matrix, i.e. computes the integrals @@ -2320,7 +2333,7 @@ def assemble(self, weights=None, clear=True): # loop over codomain spaces (rows) for a, codomain_space in enumerate(codomain_spaces): # knot span indices of elements of local domain - codomain_spans = spline_attr[W_name].quad_grid_spans[a] + codomain_spans = self._quad_grid(W_name, "quad_grid_spans", a) # global start spline index on process codomain_starts = [int(start) for start in codomain_space.coeff_space.starts] @@ -2332,10 +2345,10 @@ def assemble(self, weights=None, clear=True): pts = self._quad_pts(W_name, a) # global quadrature weights in format (local element, local weight) - wts = spline_attr[W_name].quad_grid_wts[a] + wts = self._quad_grid(W_name, "quad_grid_wts", a) # evaluated basis functions at quadrature points of codomain space - codomain_basis = spline_attr[W_name].quad_grid_bases[a] + codomain_basis = self._quad_grid(W_name, "quad_grid_bases", a) # loop over domain spaces (columns) for b, domain_space in enumerate(domain_spaces): @@ -2375,7 +2388,7 @@ def assemble(self, weights=None, clear=True): ) # evaluated basis functions at quadrature points of domain space - domain_basis = spline_attr[V_name].quad_grid_bases[b] + domain_basis = self._quad_grid(V_name, "quad_grid_bases", b) # assemble matrix (if mat_w is not zero) by calling the appropriate kernel (1d, 2d or 3d) if not_weight_zero or self._is_scalar: diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 9452e81c5..e6f358832 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -243,6 +243,8 @@ def allocate(self): logger.debug("\nAllocating simulation data ...") + self._compile_setup_cuda_kernels() + with ProfileManager.profile_region("setup: allocate"): # feec with ProfileManager.profile_region("allocate feec", functions=[self._allocate_feec]): @@ -280,6 +282,22 @@ def kernels(self) -> tuple[Kernel, ...]: owners += self.model.prop_list return collect_kernels(*owners) + def _compile_setup_cuda_kernels(self) -> None: + """Compile the CUDA kernels that :meth:`allocate` itself calls (mass-matrix assembly) before it calls them. + + :meth:`compile_cuda_kernels` runs after :meth:`allocate`, so without this the first assembly would include + the compilation (~1 s per kernel on a node whose CuPy cache does not have it yet), in the assembly's + profiling regions. Here it is counted in ``"setup: compile cuda kernels (assembly)"``. Nothing happens on + the NumPy backend. + """ + if xp.get_backend() != "cupy": + return + + from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat + + with ProfileManager.profile_region("setup: compile cuda kernels (assembly)"): + kernel_3d_mat.compile() + def compile_cuda_kernels(self) -> tuple[Kernel, ...]: """Compile the CUDA versions of the kernels of :meth:`kernels` now, before the time stepping. From b586cde364b0cd021de76dd25a83f0896a25946f Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:26:26 +0200 Subject: [PATCH 19/55] Parity cases for kernel_3d_mat Random quadrature data on 3 x 2 x 2 elements with mixed degrees, with the first row local or a ghost row. test_cuda_kernels_have_parity_cases requires a case for every kernel with a CUDA version; the emulation test compares it with the pyccel kernel without a GPU. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/cuda_parity_cases.py | 26 ++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index a2396ee5b..6e10707be 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -315,6 +315,30 @@ def kernel_pullpush_args(case): ) +def kernel_3d_mat_threads(args): + """One thread per (element, codomain basis function, domain basis function), as the kernel's launch size.""" + from struphy.feec.kernels.kernel_3d_mat import _n_threads + + return _n_threads(args) + + +def kernel_3d_mat_args(case): + """Mass-matrix assembly on 3 x 2 x 2 elements with mixed degrees and random quadrature data; `case` is the start + index of the rows in direction 1 (1: the first row is a ghost row, local index -1).""" + rng = np.random.default_rng(20 + case) + ne, nq = (3, 2, 2), (3, 2, 2) + pi, pj, pads = (2, 1, 1), (1, 2, 1), (2, 2, 2) + starts = (case, 0, 0) + spans = [xp.asarray(np.arange(n, dtype=np.int64) + p) for n, p in zip(ne, pi)] + w = [xp.asarray(rng.random((n, q))) for n, q in zip(ne, nq)] + bi = [xp.asarray(rng.random((n, p + 1, 1, q))) for n, p, q in zip(ne, pi, nq)] + bj = [xp.asarray(rng.random((n, p + 1, 1, q))) for n, p, q in zip(ne, pj, nq)] + mat_fun = xp.asarray(rng.normal(size=tuple(n * q for n, q in zip(ne, nq)))) + rows = tuple(n + p + 2 * pad for n, p, pad in zip(ne, pi, pads)) + data = xp.asarray(rng.normal(size=rows + tuple(2 * pad + 1 for pad in pads))) # added to, not overwritten + return (*spans, *pi, *pj, *starts, *pads, *w, *bi, *bj, mat_fun, data) + + # ---------------------------------------------------------------- all kernels with a CUDA version PUSHER_TOLERANCES = {"rtol": 1e-13, "atol": 1e-14} @@ -338,6 +362,8 @@ def size_of(index, per_thread=1): PARITY_CASES = { + # atomic additions: the summation order over elements differs from pyccel + "kernel_3d_mat": ParityCases((0, 1), kernel_3d_mat_args, rtol=1e-12, atol=1e-13, n_threads=kernel_3d_mat_threads), "push_eta_stage": ParityCases( [ (bc, method, stage) From 4147d1a156ee0a98022751342499f14c30040425 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:08:01 +0200 Subject: [PATCH 20/55] Apply separable mass matrices as Kronecker products of 1D banded matrices WeightedMassOperator.dot applies the blocks of its matrix that are Kronecker products of 1D banded matrices (constant weights, e.g. on a Cuboid) as three banded passes (struphy.feec.kronecker_mass, CUDA kernel kron_band_pass on the GPU) instead of the 3D stencil product. The replacement is built on first use and invalidated by assembly, copies and in-place arithmetic; a matrix changed by arithmetic that is not a Kronecker product is not checked again. STRUPHY_KRONECKER_MASS=0 switches it off. The boundary and extraction operators around the matrix are skipped when they are identities instead of being applied as copies. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/kron_band_pass_cuda.cu | 35 ++++ src/struphy/feec/kronecker_mass.py | 247 ++++++++++++++++++++++++ src/struphy/feec/mass.py | 85 ++++++-- 3 files changed, 355 insertions(+), 12 deletions(-) create mode 100644 src/struphy/feec/kron_band_pass_cuda.cu create mode 100644 src/struphy/feec/kronecker_mass.py diff --git a/src/struphy/feec/kron_band_pass_cuda.cu b/src/struphy/feec/kron_band_pass_cuda.cu new file mode 100644 index 000000000..745e782f4 --- /dev/null +++ b/src/struphy/feec/kron_band_pass_cuda.cu @@ -0,0 +1,35 @@ +/** + * One pass of a Kronecker product of 1D banded matrices (struphy.feec.kronecker_mass._band_pass): + * + * out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...] along `axis`, + * + * with the other two axes unchanged (out and inp have the same extent along them). One thread per entry of `out`, + * the last axis fastest, so that neighbouring threads read neighbouring entries of `inp` along every axis. `out` may + * be a strided view (the owned rows of a stencil vector's data). + * + * @param A Banded 1D matrix, shape (out.shape[axis], n_diagonals). + * @param inp Input array (3D). + * @param out Output array (3D, any strides); written, not added to. + * @param axis Axis of the pass (0, 1 or 2). + * @param offset Index in `inp` of diagonal 0 of row 0 along `axis`. + */ +#include "cunumpy/array_view.cuh" + +extern "C" __global__ void kron_band_pass(Array2D A, Array3D inp, Array3D out, int axis, + long long offset) { + const long long n1 = out.shape[1], n2 = out.shape[2]; + const long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; + if (t >= out.shape[0] * n1 * n2) return; + const long long i2 = t % n2; + const long long i1 = (t / n2) % n1; + const long long i0 = t / (n1 * n2); + const long long i = axis == 0 ? i0 : (axis == 1 ? i1 : i2); + + double v = 0.0; + for (long long d = 0; d < A.shape[1]; ++d) { + const long long j = i + offset + d; + const double x = axis == 0 ? inp(j, i1, i2) : (axis == 1 ? inp(i0, j, i2) : inp(i0, i1, j)); + v += A(i, d) * x; + } + out(i0, i1, i2) = v; +} diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py new file mode 100644 index 000000000..2050edb82 --- /dev/null +++ b/src/struphy/feec/kronecker_mass.py @@ -0,0 +1,247 @@ +"""Fast application of assembled mass matrices that are Kronecker products of 1D banded matrices. + +A mass matrix block ``M[i, j] = int w Lambda_i Lambda_j`` of tensor-product splines factorizes as +``M = c * M_1 (x) M_2 (x) M_3`` when its weight ``w`` is constant (e.g. on a Cuboid). Applied as a +:class:`~feectools.linalg.stencil.StencilMatrix`, every row reads all ``n_1 * n_2 * n_3`` diagonals (105 for +degrees (1, 2, 3)) from memory; as a Kronecker product it is three 1D banded passes over the vector, ``n_1 + n_2 + +n_3`` operations per entry and no 3D matrix data. :func:`kronecker_apply_operator` checks each assembled block and +returns an operator for the matrix-vector products in which the blocks that are Kronecker products (to round-off) +are replaced by :class:`KroneckerBandMatrix`. + +The 1D factors are read from the assembled data itself (rows owned by this process), so they hold exactly what the +assembly computed, on any domain decomposition, and the check is the assembled data, not an assumption about the +mapping or the weights. +""" + +import logging +import os +from functools import cache +from pathlib import Path + +import cunumpy as xp +from feectools.linalg.basic import LinearOperator +from feectools.linalg.block import BlockLinearOperator +from feectools.linalg.stencil import StencilMatrix, StencilVector + +logger = logging.getLogger("struphy") + +# relative tolerance of the Kronecker check: the 3D assembly and the product of the factors differ by round-off only +KRONECKER_RTOL = 1e-12 + + +def kronecker_mass_enabled() -> bool: + """Whether :func:`kronecker_apply_operator` may replace blocks (``STRUPHY_KRONECKER_MASS=0`` switches it off).""" + return os.environ.get("STRUPHY_KRONECKER_MASS", "1").lower() not in ("0", "false", "no") + + +@cache +def _band_pass_cuda(): + """The CUDA kernel of one banded pass (compiled on first call).""" + from cunumpy.kernels import CudaKernel + + return CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass") + + +def _band_pass(A, inp, out, axis, offset): + """``out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...]`` along ``axis`` (other axes unchanged).""" + if xp.is_gpu(out): + _band_pass_cuda()(A, inp, out, axis, offset, n_threads=out.size) + return + m = out.shape[axis] + shape = [1, 1, 1] + shape[axis] = m + out[...] = 0.0 + for d in range(A.shape[1]): + idx = [slice(None)] * 3 + idx[axis] = slice(offset + d, offset + d + m) + out += A[:, d].reshape(shape) * inp[tuple(idx)] + + +class KroneckerBandMatrix(LinearOperator): + """A square 3D :class:`~feectools.linalg.stencil.StencilMatrix` applied as a Kronecker product of 1D banded + matrices, ``(A_1 (x) A_2 (x) A_3) x`` on the rows owned by this process. + + ``A_k[i, d]`` is the entry of owned row ``i`` (along direction ``k``) and diagonal ``d``, column + ``i - q_k + d`` with ``q_k = (n_k - 1) // 2``, as in the stencil matrix. The product is computed in three passes, + each over the owned rows of its direction and the full (padded) extent of the directions after it, so the passes + need no communication beyond the ghost regions of the input. Use :meth:`from_stencil`. + + Parameters + ---------- + stencil : StencilMatrix + The assembled matrix (kept for ``toarray``, ``tosparse`` and ``transpose``). + factors : tuple of arrays + ``A_1, A_2, A_3``, shapes ``(m_k, n_k)``, C-contiguous backend arrays. + """ + + def __init__(self, stencil: StencilMatrix, factors): + self._stencil = stencil + self._factors = factors + + V = stencil.domain + self._m = tuple(int(e - s + 1) for s, e in zip(V.starts, V.ends)) + self._pads = tuple(int(p) for p in V.pads) + self._q = tuple((int(n) - 1) // 2 for n in stencil._data.shape[3:]) + + # temporaries of the first two passes: owned rows along the passes done, padded extent along the others + x_shape = V.zeros()._data.shape + self._t1 = xp.empty((self._m[0], x_shape[1], x_shape[2]), dtype=stencil._data.dtype) + self._t2 = xp.empty((self._m[0], self._m[1], x_shape[2]), dtype=stencil._data.dtype) + + @classmethod + def from_stencil(cls, S: StencilMatrix, rtol: float = KRONECKER_RTOL): + """The Kronecker form of ``S`` if it is one (to ``rtol``, relative to its largest entry), else None. + + Only square 3D matrices on one space (domain is codomain) with shifts 1 qualify, e.g. the diagonal blocks of + the mass matrices. The factors are read from the owned rows of ``S``: with reference rows ``r_k`` (the first + owned row along each direction) and central diagonals ``q_k``, + ``A_1[i, :] = S[i, r_2, r_3, :, q_2, q_3]``, ``A_2[i, :] = S[r_1, i, r_3, q_1, :, q_3] / S[r, q]``, + ``A_3[i, :] = S[r_1, r_2, i, q_1, q_2, :] / S[r, q]``; the product is then compared with all owned rows. + """ + if not isinstance(S, StencilMatrix) or S.domain is not S.codomain or S.domain.ndim != 3: + return None + V = S.domain + if any(int(s) != 1 for s in V.shifts) or S._data.ndim != 6: + return None + + p = [int(p) for p in V.pads] + m = [int(e - s + 1) for s, e in zip(V.starts, V.ends)] + q = [(int(n) - 1) // 2 for n in S._data.shape[3:]] + if min(m) < 1: + return None + data = S._data + + den = data[p[0], p[1], p[2], q[0], q[1], q[2]] + if float(den) == 0.0: + return None + A1 = xp.ascontiguousarray(data[p[0] : p[0] + m[0], p[1], p[2], :, q[1], q[2]]) + A2 = xp.ascontiguousarray(data[p[0], p[1] : p[1] + m[1], p[2], q[0], :, q[2]] / den) + A3 = xp.ascontiguousarray(data[p[0], p[1], p[2] : p[2] + m[2], q[0], q[1], :] / den) + + # compare with every owned row: planes (rows along direction 1) in chunks, the maxima kept on the device + # (one read at the end); a few planes first, which rejects most matrices that are not Kronecker products + # (e.g. ones assembled from particle data, checked again after every change) at once + owned = data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + A23 = A2[:, None, None, :, None] * A3[None, :, None, None, :] # (m2, m3, 1, n2, n3) + + def max_diff_and_abs(rows): + block = owned[rows] + recon = A1[rows][:, None, None, :, None, None] * A23[None] + return xp.max(xp.abs(block - recon)), xp.max(xp.abs(block)) + + def fails(max_diff, max_abs): + max_diff, max_abs = float(max_diff), float(max_abs) + return max_abs == 0.0 or max_diff > rtol * max_abs + + quick = sorted({0, m[0] // 2, m[0] - 1}) + if fails(*max_diff_and_abs(xp.asarray(quick))): + return None + + # chunks of at most ~64 MB of temporaries + plane = max(1, owned[0].size) + chunk = max(1, min(m[0], (8 << 20) // plane)) + max_diff = xp.zeros(()) + max_abs = xp.zeros(()) + for start in range(0, m[0], chunk): + d, a = max_diff_and_abs(slice(start, min(start + chunk, m[0]))) + max_diff = xp.maximum(max_diff, d) + max_abs = xp.maximum(max_abs, a) + if fails(max_diff, max_abs): + return None + return cls(S, (A1, A2, A3)) + + # ------------------------------------------------------------------ + # LinearOperator interface + # ------------------------------------------------------------------ + @property + def domain(self): + return self._stencil.domain + + @property + def codomain(self): + return self._stencil.codomain + + @property + def dtype(self): + return self._stencil.dtype + + @property + def factors(self): + """The 1D factors ``A_1, A_2, A_3`` (rows owned by this process).""" + return self._factors + + def tosparse(self): + return self._stencil.tosparse() + + def toarray(self): + return self._stencil.toarray() + + def transpose(self, conjugate=False): + return self._stencil.transpose(conjugate=conjugate) + + def dot(self, v, out=None): + assert isinstance(v, StencilVector) + assert v.space is self.domain + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.codomain + + if not v.ghost_regions_in_sync: + v.update_ghost_regions() + + A1, A2, A3 = self._factors + p, q, m = self._pads, self._q, self._m + owned = out._data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + + _band_pass(A1, v._data, self._t1, 0, p[0] - q[0]) + _band_pass(A2, self._t1, self._t2, 1, p[1] - q[1]) + _band_pass(A3, self._t2, owned, 2, p[2] - q[2]) + + out.ghost_regions_in_sync = False + return out + + +def kronecker_apply_operator(mat, name: str = ""): + """The operator to apply ``mat`` with: ``mat`` itself, or a copy whose Kronecker blocks are replaced. + + Parameters + ---------- + mat : StencilMatrix or BlockLinearOperator + An assembled (block) mass matrix. + name : str + Name of the operator, for the log. + + Returns + ------- + LinearOperator + ``mat`` if no block is a Kronecker product (or :func:`kronecker_mass_enabled` is False); otherwise a + :class:`KroneckerBandMatrix` (for a StencilMatrix) or a BlockLinearOperator with the same blocks except the + Kronecker ones. It is valid only as long as ``mat`` is not modified. + """ + if not kronecker_mass_enabled(): + return mat + + if isinstance(mat, StencilMatrix): + kron = KroneckerBandMatrix.from_stencil(mat) + if kron is not None: + logger.info(f"Mass matrix {name}: applied as a Kronecker product of 1D banded matrices") + return mat if kron is None else kron + + if isinstance(mat, BlockLinearOperator): + blocks = {} + replaced = [] + for key, block in mat._blocks.items(): + kron = KroneckerBandMatrix.from_stencil(block) + if kron is None: + blocks[key] = block + else: + blocks[key] = kron + replaced.append(key) + if not replaced: + return mat + logger.info(f"Mass matrix {name}: blocks {replaced} applied as Kronecker products of 1D banded matrices") + return BlockLinearOperator(mat.domain, mat.codomain, blocks=blocks) + + return mat diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index e035e4fe4..1020f0ef6 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -25,6 +25,7 @@ from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat from struphy.feec.kernels.kernel_3d_matrixfree import kernel_3d_matrixfree from struphy.feec.kernels.kernel_3d_vec import kernel_3d_vec +from struphy.feec.kronecker_mass import kronecker_apply_operator from struphy.feec.linear_operators import BoundaryOperator from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.feec.utilities import LocalProjectionMatrix, LocalRotationMatrix, get_quad_grids @@ -1703,6 +1704,8 @@ def __init__( blocks, self._weights = self._init_blocks_from_weights(weights_info) self._mat = self._wrap_blocks(blocks) + # the operator dot() applies self._mat with (Kronecker blocks replaced), built on first use, see _apply_mat + self._apply_mat_cache = None # the weights are stored in the block order of the (possibly transposed) operator if transposed: @@ -2121,6 +2124,32 @@ def M0(self): def matrix(self): return self._mat + @property + def _apply_mat(self): + """The operator :meth:`dot` applies: :attr:`matrix`, with its blocks that are Kronecker products of 1D banded + matrices (constant weights, e.g. on a Cuboid) applied as such (:mod:`struphy.feec.kronecker_mass`). + + Built on first use after every change of the matrix (assembly, copy into, in-place arithmetic). + """ + if self._apply_mat_cache is None: + if getattr(self, "_kronecker_rejected", False): + self._apply_mat_cache = self._mat + else: + self._apply_mat_cache = kronecker_apply_operator(self._mat, name=self.name) + # a matrix changed by in-place arithmetic that is not a Kronecker product (e.g. the Schur complement + # M1 - dt^2 BC with particle data, rebuilt every step) is not checked again after further arithmetic: + # the check reads the whole matrix. Assembly and copies check again. + if getattr(self, "_changed_by_arithmetic", False) and self._apply_mat_cache is self._mat: + self._kronecker_rejected = True + return self._apply_mat_cache + + def _matrix_changed(self, by_arithmetic: bool): + """Invalidate :attr:`_apply_mat` after a change of the matrix (in-place arithmetic, or assembly/copy).""" + self._apply_mat_cache = None + self._changed_by_arithmetic = by_arithmetic + if not by_arithmetic: + self._kronecker_rejected = False + @property def nquads(self): if self._nquads is None: @@ -2168,29 +2197,54 @@ def dot(self, v, out=None, apply_bc=True): if apply_bc: if self._transposed: - self._W_boundary_op_T.dot(v, out=self._temp_WB) - self._W_extraction_op_T.dot(self._temp_WB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) - self._V_extraction_op.dot(self._temp_VE, out=self._temp_VB) - out = self._V_boundary_op.dot(self._temp_VB, out=out) + ops_in = (self._W_boundary_op_T, self._W_extraction_op_T) + ops_out = (self._V_extraction_op, self._V_boundary_op) + tmps_in = (self._temp_WB, self._temp_mat) + tmps_out = (self._temp_VE, self._temp_VB) else: - self._V_boundary_op_T.dot(v, out=self._temp_VB) - self._V_extraction_op_T.dot(self._temp_VB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) - self._W_extraction_op.dot(self._temp_WE, out=self._temp_WB) - out = self._W_boundary_op.dot(self._temp_WB, out=out) + ops_in = (self._V_boundary_op_T, self._V_extraction_op_T) + ops_out = (self._W_extraction_op, self._W_boundary_op) + tmps_in = (self._temp_VB, self._temp_mat) + tmps_out = (self._temp_WE, self._temp_WB) + out = self._apply_chain(v, ops_in, tmps_in, ops_out, tmps_out, out) else: if self._transposed: self._W_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) + self._apply_mat.dot(self._temp_mat, out=self._temp_VE) out = self._V_extraction_op.dot(self._temp_VE, out=out) else: self._V_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) + self._apply_mat.dot(self._temp_mat, out=self._temp_WE) out = self._W_extraction_op.dot(self._temp_WE, out=out) return out + def _apply_chain(self, v, ops_in, tmps_in, ops_out, tmps_out, out): + """``out = B_out E_out M E_in B_in v`` (boundary and extraction operators around the matrix), see :meth:`dot`. + + Identity operators (e.g. the extraction operators without polar splines, or boundary operators without + essential boundary conditions) are skipped instead of applied as copies: each vector is passed on as it is, + and the matrix writes into the next temporary that is actually used (or into ``out``). + """ + if out is v: + # in place: the chain must not read v after writing the output, so it reads a copy + v = v.copy() + + x = v + for op, tmp in zip(ops_in, tmps_in): + if not isinstance(op, IdentityOperator): + x = op.dot(x, out=tmp) + + # the first output operator that is not an identity reads the matrix product from its temporary + remaining = [(op, tmp) for op, tmp in zip(ops_out, tmps_out) if not isinstance(op, IdentityOperator)] + if not remaining: + return self._apply_mat.dot(x, out=out) + y = self._apply_mat.dot(x, out=remaining[0][1]) + for k, (op, _) in enumerate(remaining): + target = out if k == len(remaining) - 1 else remaining[k + 1][1] + y = op.dot(y, out=target) + return y + def transpose(self, conjugate=False): """ Returns the transposed operator. @@ -2277,6 +2331,7 @@ def assemble(self, weights=None, clear=True): Whether to first set all data to zero before assembly. If False, the new contributions are added to existing ones. """ + self._matrix_changed(by_arithmetic=False) # the matrix changes assert not self._dry_run, ( "A dry-run operator has no matrix data and cannot be assembled (memory estimation only)." ) @@ -2559,6 +2614,7 @@ def copy(self, out=None): out._weights = [list(row) for row in self._weights] self._mat.copy(out=out._mat) + out._matrix_changed(by_arithmetic=False) if self._creation_info is None: out._creation_info = None @@ -2569,6 +2625,7 @@ def copy(self, out=None): def __imul__(self, a): self._mat *= a + self._matrix_changed(by_arithmetic=True) self._creation_info = None return self @@ -2579,11 +2636,13 @@ def __iadd__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat += M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat += M + self._matrix_changed(by_arithmetic=True) return self else: @@ -2596,11 +2655,13 @@ def __isub__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat -= M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat -= M + self._matrix_changed(by_arithmetic=True) return self else: From f76683086b2812177a352584a61626c763df141d Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:20:50 +0200 Subject: [PATCH 21/55] Kronecker mass: launch the band-pass kernel through CachedLaunch The banded passes run with the same buffers in every solver iteration (as on cuda-aware-mpi-fixes, now that feectools has CachedLaunch). Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/kronecker_mass.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py index 2050edb82..5456b2eeb 100644 --- a/src/struphy/feec/kronecker_mass.py +++ b/src/struphy/feec/kronecker_mass.py @@ -38,8 +38,10 @@ def kronecker_mass_enabled() -> bool: def _band_pass_cuda(): """The CUDA kernel of one banded pass (compiled on first call).""" from cunumpy.kernels import CudaKernel + from feectools.utilities.cached_launch import CachedLaunch - return CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass") + # launched with the same buffers in every solver iteration: cache the launch preparation + return CachedLaunch(CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass")) def _band_pass(A, inp, out, axis, offset): From d0a8f31c75bc0c5b4e3e385ed9bb6db9c55c9a5f Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:00:05 +0200 Subject: [PATCH 22/55] CuPy backend: host/device conversions in post-processing; math.sqrt for host scalars in Species Post-processing evaluates fields on the device and writes them from the host: the evaluation grids, point data and coordinates are converted where they cross, and HDF5 data is moved to the device before it is written into the spline coefficients. Co-Authored-By: Claude Opus 5.5 --- src/struphy/models/species.py | 4 ++-- src/struphy/post_processing/output.py | 31 ++++++++++++++++++--------- 2 files changed, 23 insertions(+), 12 deletions(-) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 3c8b47657..4bc2dcac1 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,9 +1,9 @@ import logging +import math import warnings from abc import ABCMeta, abstractmethod from dataclasses import fields, is_dataclass -import cunumpy as xp from maybempi import MPI from struphy.models.variables import Variable @@ -203,7 +203,7 @@ def __init__( con = ConstantsOfNature() # relevant frequencies - om_p = xp.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) + om_p = math.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) om_c = Z * con.e * units.B / (A * con.mH) # compute equation parameters diff --git a/src/struphy/post_processing/output.py b/src/struphy/post_processing/output.py index bfcf5fd58..cc65d56f8 100644 --- a/src/struphy/post_processing/output.py +++ b/src/struphy/post_processing/output.py @@ -85,6 +85,13 @@ def load_plasma_plots() -> bool: return True +def _to_host(value): + """`value` (an array, or a list/tuple of arrays for vector-valued fields) as NumPy array(s).""" + if isinstance(value, (list, tuple)): + return type(value)(_to_host(v) for v in value) + return xp.to_numpy(value) + + class ProductMapping(Mapping[str, xr.DataArray]): """A discoverable mapping whose products are loaded on first access.""" @@ -568,13 +575,14 @@ def _attach_physical_coords( self, array: xr.DataArray, etas: tuple[Any, Any, Any], dims: tuple[str, ...] ) -> xr.DataArray: """Map a raw FEEC evaluation grid and attach its Cartesian coordinates.""" - grids = [np.atleast_1d(np.asarray(eta, dtype=float)) for eta in etas] - mapped = self.domain(*grids) + # the logical grids may be device arrays (CuPy backend): map them where they live, keep the coordinates on the host + grids = [np.atleast_1d(np.asarray(xp.to_numpy(eta), dtype=float)) for eta in etas] + mapped = self.domain(*(xp.asarray(grid) for grid in grids)) shape = tuple(len(grid) for grid in grids) keep = tuple(slice(None) if dim in dims else 0 for dim in ("eta1", "eta2", "eta3")) coordinates = {} for name, values in zip(("X", "Y", "Z"), mapped): - values = np.asarray(values).reshape(shape)[keep] + values = np.asarray(xp.to_numpy(values)).reshape(shape)[keep] coordinates[name] = (dims, values) if dims else values.item() return array.assign_coords(coordinates) @@ -1295,7 +1303,9 @@ def _process_fields( # the physical grid is only needed for output, hence it is only built on rank 0 if self._pproc_rank == 0: - grids_phy = list(self.domain(*grids_log)) + # output coordinates, on the host (the logical grids stay on the device for the field evaluation) + grids_phy = [xp.to_numpy(grid) for grid in self.domain(*grids_log)] + grids_log_host = [xp.to_numpy(grid) for grid in grids_log] else: grids_phy = None @@ -1327,8 +1337,9 @@ def _process_fields( if self._pproc_rank == 0: for species, vars in vals.items(): for name, val in vars.items(): - point_data[species][name][t] = val - point_data_phy[species][name][t] = vals_phy[species][name] + # evaluated on the device on the CuPy backend; written to files from the host + point_data[species][name][t] = _to_host(val) + point_data_phy[species][name][t] = _to_host(vals_phy[species][name]) # directory for the vtk files path_fields = os.path.join(self.path_pproc, "fields_data") @@ -1338,10 +1349,10 @@ def _process_fields( for species, vars in point_data.items(): variables = {} for name, val in vars.items(): - variables[name] = wrap_field_data(val, grids_log, grids_phy=grids_phy, name=name) + variables[name] = wrap_field_data(val, grids_log_host, grids_phy=grids_phy, name=name) if physical: variables[name + "_xyz"] = wrap_field_data( - point_data_phy[species][name], grids_log, grids_phy=grids_phy, name=name + "_xyz" + point_data_phy[species][name], grids_log_host, grids_phy=grids_phy, name=name + "_xyz" ) store.write_group(store.store_path(self.path_pproc), f"/{species}", xr.Dataset(variables)) @@ -1517,7 +1528,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[n * step, p1:-p1, p2:-p2, p3:-p3]) # HDF5 data is on the host # vector-valued field else: @@ -1530,7 +1541,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3]) vector.update_ghost_regions() if is_polar: From 5c099f3a49c908373763cbd7aadb041289ff47c7 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:07:26 +0200 Subject: [PATCH 23/55] CUDA version of the mass-matrix assembly kernel kernel_3d_mat - kernel_3d_mat_cuda.cu: one thread per (element, codomain basis function, domain basis function); replaces the host fallback on the CuPy backend. - WeightedMassOperator.assemble reads the quadrature grids (spans, weights, basis values) as backend arrays, copied to the device once and cached. - Simulation.allocate compiles the assembly kernel before it is first called, so the compilation is not counted in the assembly's profiling regions. Co-Authored-By: Claude Opus 5.5 --- .../feec/kernels/kernel_3d_mat/__init__.py | 18 ++-- .../kernel_3d_mat/kernel_3d_mat_cuda.cu | 92 +++++++++++++++++++ src/struphy/feec/mass.py | 21 ++++- src/struphy/simulation/sim.py | 18 ++++ 4 files changed, 139 insertions(+), 10 deletions(-) create mode 100644 src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu diff --git a/src/struphy/feec/kernels/kernel_3d_mat/__init__.py b/src/struphy/feec/kernels/kernel_3d_mat/__init__.py index a881301b4..76da51673 100644 --- a/src/struphy/feec/kernels/kernel_3d_mat/__init__.py +++ b/src/struphy/feec/kernels/kernel_3d_mat/__init__.py @@ -1,8 +1,4 @@ -"""The kernel of this folder: pyccel in ``_kernels.py``, CUDA in ``_cuda.cu`` once ported. - -No CUDA version yet: on the CuPy backend the pyccel kernel runs on host copies of the arrays (a warning is -emitted once), as before the move into this folder. -""" +"""The kernel of this folder: pyccel in ``kernel_3d_mat_kernels.py``, CUDA in ``kernel_3d_mat_cuda.cu``.""" from cunumpy.kernels import Kernel @@ -11,6 +7,16 @@ # the arguments the kernel writes to, by position: data OUTPUTS = (25,) + +def _n_threads(args): + """One thread per (element, codomain basis function, domain basis function), see kernel_3d_mat_cuda.cu.""" + n_elements = args[0].size * args[1].size * args[2].size + n_pairs = 1 + for degree in args[3:9]: # pi1, pi2, pi3, pj1, pj2, pj3 + n_pairs *= int(degree) + 1 + return n_elements * n_pairs + + kernel_3d_mat = Kernel.from_folder( - __name__, structs=CUDA_STRUCTS, missing_cuda="fallback", host_options={"outputs": OUTPUTS} + __name__, structs=CUDA_STRUCTS, host_options={"outputs": OUTPUTS}, n_threads_from=_n_threads ) diff --git a/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu b/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu new file mode 100644 index 000000000..6afab9398 --- /dev/null +++ b/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu @@ -0,0 +1,92 @@ +// CUDA version of kernel_3d_mat (kernel_3d_mat_kernels.py), with the same arguments in the same order. +#include "cunumpy/array_view.cuh" +#include "cunumpy/atomic.cuh" +#include "cunumpy/index.cuh" + +/** + * Integral of Lambda_i * mat_fun * Lambda_j over the local elements, added to the stencil matrix data, as in + * kernel_3d_mat_kernels.kernel_3d_mat. + * + * One thread per (element, codomain basis function (il1, il2, il3), domain basis function (jl1, jl2, jl3)): + * n_threads = ne1 * ne2 * ne3 * (pi1 + 1) * (pi2 + 1) * (pi3 + 1) * (pj1 + 1) * (pj2 + 1) * (pj3 + 1) (see + * __init__.py), the element index varying slowest and jl3 fastest. Each thread computes the quadrature sum of its + * pair exactly as the pyccel kernel (same factors, same order) and adds it to its entry of `data` with an atomic + * addition: neighbouring elements contribute to the same entries, so only the summation order over elements + * differs from pyccel (round-off). + * + * @param spans1, spans2, spans3 Knot span index of each local element, per direction. + * @param pi1, pi2, pi3 Degrees of the codomain basis functions. + * @param pj1, pj2, pj3 Degrees of the domain basis functions. + * @param starts1, starts2, starts3 Global start index of the rows on this process. + * @param pads1, pads2, pads3 Paddings (ghost regions) of `data`. + * @param w1, w2, w3 Quadrature weights [element, quadrature point]. + * @param bi1, bi2, bi3 Codomain basis functions [element, local basis function, derivative, quadrature point]. + * @param bj1, bj2, bj3 Domain basis functions, same layout. + * @param mat_fun Weight function at the quadrature points [flattened point 1, flattened point 2, flattened point 3]. + * @param data StencilMatrix data (rows with ghost regions..., diagonals...); the results are added (not set). + */ +extern "C" __global__ void kernel_3d_mat(Array1D spans1, Array1D spans2, + Array1D spans3, int pi1, int pi2, int pi3, int pj1, int pj2, + int pj3, int starts1, int starts2, int starts3, int pads1, int pads2, + int pads3, Array2D w1, Array2D w2, Array2D w3, + Array4D bi1, Array4D bi2, Array4D bi3, + Array4D bj1, Array4D bj2, Array4D bj3, + Array3D mat_fun, Array6D data) +{ + const long long ne1 = spans1.size(); + const long long ne2 = spans2.size(); + const long long ne3 = spans3.size(); + + const long long ni1 = pi1 + 1, ni2 = pi2 + 1, ni3 = pi3 + 1; + const long long nj1 = pj1 + 1, nj2 = pj2 + 1, nj3 = pj3 + 1; + const long long n_pairs = ni1 * ni2 * ni3 * nj1 * nj2 * nj3; + + CUNUMPY_THREAD_1D(thread, ne1 * ne2 * ne3 * n_pairs); + + // CUDA-only: split the thread index into the pyccel loop variables (iel1, iel2, iel3, il1, il2, il3, jl1, jl2, + // jl3), jl3 fastest + long long rest = thread; + const long long jl3 = rest % nj3; rest /= nj3; + const long long jl2 = rest % nj2; rest /= nj2; + const long long jl1 = rest % nj1; rest /= nj1; + const long long il3 = rest % ni3; rest /= ni3; + const long long il2 = rest % ni2; rest /= ni2; + const long long il1 = rest % ni1; rest /= ni1; + const long long iel3 = rest % ne3; rest /= ne3; + const long long iel2 = rest % ne2; rest /= ne2; + const long long iel1 = rest; + + const long long nq1 = w1.shape[1]; + const long long nq2 = w2.shape[1]; + const long long nq3 = w3.shape[1]; + + double value = 0.0; + + for (long long q1 = 0; q1 < nq1; ++q1) { + for (long long q2 = 0; q2 < nq2; ++q2) { + for (long long q3 = 0; q3 < nq3; ++q3) { + const double wvol = w1(iel1, q1) * w2(iel2, q2) * w3(iel3, q3) * + mat_fun(iel1 * nq1 + q1, iel2 * nq2 + q2, iel3 * nq3 + q3); + + const double bi = bi1(iel1, il1, 0, q1) * bi2(iel2, il2, 0, q2) * bi3(iel3, il3, 0, q3); + const double bj = bj1(iel1, jl1, 0, q1) * bj2(iel2, jl2, 0, q2) * bj3(iel3, jl3, 0, q3); + + value += wvol * bi * bj; + } + } + } + + // global spline indices + const long long i_global1 = spans1(iel1) - pi1 + il1; + const long long i_global2 = spans2(iel2) - pi2 + il2; + const long long i_global3 = spans3(iel3) - pi3 + il3; + + // local spline indices (- starts --> can be negative, will therefore be written to ghost regions) + const long long i_local1 = i_global1 - starts1; + const long long i_local2 = i_global2 - starts2; + const long long i_local3 = i_global3 - starts3; + + cunumpy_atomic_add(&data(pads1 + i_local1, pads2 + i_local2, pads3 + i_local3, pads1 + jl1 - il1, + pads2 + jl2 - il2, pads3 + jl3 - il3), + value); +} diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index f75c6de17..e035e4fe4 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -2242,6 +2242,19 @@ def transpose(self, conjugate=False): return M + def _quad_grid(self, space_name, attribute, component): + """Quadrature-grid arrays (spans, weights or basis values) of one component of a space, as backend arrays. + + The Derham keeps them as host arrays; on the CuPy backend the assembly kernel needs them on the device, + so they are copied there once and cached (they do not change between assemblies). + """ + cache = self.__dict__.setdefault("_quad_grid_cache", {}) + key = (space_name, attribute, component) + if key not in cache: + arrays = getattr(self.derham.spline_attributes[space_name], attribute)[component] + cache[key] = [xp.asarray(array) for array in arrays] + return cache[key] + def assemble(self, weights=None, clear=True): r""" Assembles the weighted mass matrix, i.e. computes the integrals @@ -2320,7 +2333,7 @@ def assemble(self, weights=None, clear=True): # loop over codomain spaces (rows) for a, codomain_space in enumerate(codomain_spaces): # knot span indices of elements of local domain - codomain_spans = spline_attr[W_name].quad_grid_spans[a] + codomain_spans = self._quad_grid(W_name, "quad_grid_spans", a) # global start spline index on process codomain_starts = [int(start) for start in codomain_space.coeff_space.starts] @@ -2332,10 +2345,10 @@ def assemble(self, weights=None, clear=True): pts = self._quad_pts(W_name, a) # global quadrature weights in format (local element, local weight) - wts = spline_attr[W_name].quad_grid_wts[a] + wts = self._quad_grid(W_name, "quad_grid_wts", a) # evaluated basis functions at quadrature points of codomain space - codomain_basis = spline_attr[W_name].quad_grid_bases[a] + codomain_basis = self._quad_grid(W_name, "quad_grid_bases", a) # loop over domain spaces (columns) for b, domain_space in enumerate(domain_spaces): @@ -2375,7 +2388,7 @@ def assemble(self, weights=None, clear=True): ) # evaluated basis functions at quadrature points of domain space - domain_basis = spline_attr[V_name].quad_grid_bases[b] + domain_basis = self._quad_grid(V_name, "quad_grid_bases", b) # assemble matrix (if mat_w is not zero) by calling the appropriate kernel (1d, 2d or 3d) if not_weight_zero or self._is_scalar: diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 9452e81c5..e6f358832 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -243,6 +243,8 @@ def allocate(self): logger.debug("\nAllocating simulation data ...") + self._compile_setup_cuda_kernels() + with ProfileManager.profile_region("setup: allocate"): # feec with ProfileManager.profile_region("allocate feec", functions=[self._allocate_feec]): @@ -280,6 +282,22 @@ def kernels(self) -> tuple[Kernel, ...]: owners += self.model.prop_list return collect_kernels(*owners) + def _compile_setup_cuda_kernels(self) -> None: + """Compile the CUDA kernels that :meth:`allocate` itself calls (mass-matrix assembly) before it calls them. + + :meth:`compile_cuda_kernels` runs after :meth:`allocate`, so without this the first assembly would include + the compilation (~1 s per kernel on a node whose CuPy cache does not have it yet), in the assembly's + profiling regions. Here it is counted in ``"setup: compile cuda kernels (assembly)"``. Nothing happens on + the NumPy backend. + """ + if xp.get_backend() != "cupy": + return + + from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat + + with ProfileManager.profile_region("setup: compile cuda kernels (assembly)"): + kernel_3d_mat.compile() + def compile_cuda_kernels(self) -> tuple[Kernel, ...]: """Compile the CUDA versions of the kernels of :meth:`kernels` now, before the time stepping. From bcbb46a2474fbe07d469ae5d457bd8db96025c0d Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:26:26 +0200 Subject: [PATCH 24/55] Parity cases for kernel_3d_mat Random quadrature data on 3 x 2 x 2 elements with mixed degrees, with the first row local or a ghost row. test_cuda_kernels_have_parity_cases requires a case for every kernel with a CUDA version; the emulation test compares it with the pyccel kernel without a GPU. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/cuda_parity_cases.py | 26 ++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index a2396ee5b..6e10707be 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -315,6 +315,30 @@ def kernel_pullpush_args(case): ) +def kernel_3d_mat_threads(args): + """One thread per (element, codomain basis function, domain basis function), as the kernel's launch size.""" + from struphy.feec.kernels.kernel_3d_mat import _n_threads + + return _n_threads(args) + + +def kernel_3d_mat_args(case): + """Mass-matrix assembly on 3 x 2 x 2 elements with mixed degrees and random quadrature data; `case` is the start + index of the rows in direction 1 (1: the first row is a ghost row, local index -1).""" + rng = np.random.default_rng(20 + case) + ne, nq = (3, 2, 2), (3, 2, 2) + pi, pj, pads = (2, 1, 1), (1, 2, 1), (2, 2, 2) + starts = (case, 0, 0) + spans = [xp.asarray(np.arange(n, dtype=np.int64) + p) for n, p in zip(ne, pi)] + w = [xp.asarray(rng.random((n, q))) for n, q in zip(ne, nq)] + bi = [xp.asarray(rng.random((n, p + 1, 1, q))) for n, p, q in zip(ne, pi, nq)] + bj = [xp.asarray(rng.random((n, p + 1, 1, q))) for n, p, q in zip(ne, pj, nq)] + mat_fun = xp.asarray(rng.normal(size=tuple(n * q for n, q in zip(ne, nq)))) + rows = tuple(n + p + 2 * pad for n, p, pad in zip(ne, pi, pads)) + data = xp.asarray(rng.normal(size=rows + tuple(2 * pad + 1 for pad in pads))) # added to, not overwritten + return (*spans, *pi, *pj, *starts, *pads, *w, *bi, *bj, mat_fun, data) + + # ---------------------------------------------------------------- all kernels with a CUDA version PUSHER_TOLERANCES = {"rtol": 1e-13, "atol": 1e-14} @@ -338,6 +362,8 @@ def size_of(index, per_thread=1): PARITY_CASES = { + # atomic additions: the summation order over elements differs from pyccel + "kernel_3d_mat": ParityCases((0, 1), kernel_3d_mat_args, rtol=1e-12, atol=1e-13, n_threads=kernel_3d_mat_threads), "push_eta_stage": ParityCases( [ (bc, method, stage) From bbe66edeeea9cd1ea9283264832165398607c65e Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:08:01 +0200 Subject: [PATCH 25/55] Apply separable mass matrices as Kronecker products of 1D banded matrices WeightedMassOperator.dot applies the blocks of its matrix that are Kronecker products of 1D banded matrices (constant weights, e.g. on a Cuboid) as three banded passes (struphy.feec.kronecker_mass, CUDA kernel kron_band_pass on the GPU) instead of the 3D stencil product. The replacement is built on first use and invalidated by assembly, copies and in-place arithmetic; a matrix changed by arithmetic that is not a Kronecker product is not checked again. STRUPHY_KRONECKER_MASS=0 switches it off. The boundary and extraction operators around the matrix are skipped when they are identities instead of being applied as copies. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/kron_band_pass_cuda.cu | 35 ++++ src/struphy/feec/kronecker_mass.py | 247 ++++++++++++++++++++++++ src/struphy/feec/mass.py | 85 ++++++-- 3 files changed, 355 insertions(+), 12 deletions(-) create mode 100644 src/struphy/feec/kron_band_pass_cuda.cu create mode 100644 src/struphy/feec/kronecker_mass.py diff --git a/src/struphy/feec/kron_band_pass_cuda.cu b/src/struphy/feec/kron_band_pass_cuda.cu new file mode 100644 index 000000000..745e782f4 --- /dev/null +++ b/src/struphy/feec/kron_band_pass_cuda.cu @@ -0,0 +1,35 @@ +/** + * One pass of a Kronecker product of 1D banded matrices (struphy.feec.kronecker_mass._band_pass): + * + * out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...] along `axis`, + * + * with the other two axes unchanged (out and inp have the same extent along them). One thread per entry of `out`, + * the last axis fastest, so that neighbouring threads read neighbouring entries of `inp` along every axis. `out` may + * be a strided view (the owned rows of a stencil vector's data). + * + * @param A Banded 1D matrix, shape (out.shape[axis], n_diagonals). + * @param inp Input array (3D). + * @param out Output array (3D, any strides); written, not added to. + * @param axis Axis of the pass (0, 1 or 2). + * @param offset Index in `inp` of diagonal 0 of row 0 along `axis`. + */ +#include "cunumpy/array_view.cuh" + +extern "C" __global__ void kron_band_pass(Array2D A, Array3D inp, Array3D out, int axis, + long long offset) { + const long long n1 = out.shape[1], n2 = out.shape[2]; + const long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; + if (t >= out.shape[0] * n1 * n2) return; + const long long i2 = t % n2; + const long long i1 = (t / n2) % n1; + const long long i0 = t / (n1 * n2); + const long long i = axis == 0 ? i0 : (axis == 1 ? i1 : i2); + + double v = 0.0; + for (long long d = 0; d < A.shape[1]; ++d) { + const long long j = i + offset + d; + const double x = axis == 0 ? inp(j, i1, i2) : (axis == 1 ? inp(i0, j, i2) : inp(i0, i1, j)); + v += A(i, d) * x; + } + out(i0, i1, i2) = v; +} diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py new file mode 100644 index 000000000..2050edb82 --- /dev/null +++ b/src/struphy/feec/kronecker_mass.py @@ -0,0 +1,247 @@ +"""Fast application of assembled mass matrices that are Kronecker products of 1D banded matrices. + +A mass matrix block ``M[i, j] = int w Lambda_i Lambda_j`` of tensor-product splines factorizes as +``M = c * M_1 (x) M_2 (x) M_3`` when its weight ``w`` is constant (e.g. on a Cuboid). Applied as a +:class:`~feectools.linalg.stencil.StencilMatrix`, every row reads all ``n_1 * n_2 * n_3`` diagonals (105 for +degrees (1, 2, 3)) from memory; as a Kronecker product it is three 1D banded passes over the vector, ``n_1 + n_2 + +n_3`` operations per entry and no 3D matrix data. :func:`kronecker_apply_operator` checks each assembled block and +returns an operator for the matrix-vector products in which the blocks that are Kronecker products (to round-off) +are replaced by :class:`KroneckerBandMatrix`. + +The 1D factors are read from the assembled data itself (rows owned by this process), so they hold exactly what the +assembly computed, on any domain decomposition, and the check is the assembled data, not an assumption about the +mapping or the weights. +""" + +import logging +import os +from functools import cache +from pathlib import Path + +import cunumpy as xp +from feectools.linalg.basic import LinearOperator +from feectools.linalg.block import BlockLinearOperator +from feectools.linalg.stencil import StencilMatrix, StencilVector + +logger = logging.getLogger("struphy") + +# relative tolerance of the Kronecker check: the 3D assembly and the product of the factors differ by round-off only +KRONECKER_RTOL = 1e-12 + + +def kronecker_mass_enabled() -> bool: + """Whether :func:`kronecker_apply_operator` may replace blocks (``STRUPHY_KRONECKER_MASS=0`` switches it off).""" + return os.environ.get("STRUPHY_KRONECKER_MASS", "1").lower() not in ("0", "false", "no") + + +@cache +def _band_pass_cuda(): + """The CUDA kernel of one banded pass (compiled on first call).""" + from cunumpy.kernels import CudaKernel + + return CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass") + + +def _band_pass(A, inp, out, axis, offset): + """``out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...]`` along ``axis`` (other axes unchanged).""" + if xp.is_gpu(out): + _band_pass_cuda()(A, inp, out, axis, offset, n_threads=out.size) + return + m = out.shape[axis] + shape = [1, 1, 1] + shape[axis] = m + out[...] = 0.0 + for d in range(A.shape[1]): + idx = [slice(None)] * 3 + idx[axis] = slice(offset + d, offset + d + m) + out += A[:, d].reshape(shape) * inp[tuple(idx)] + + +class KroneckerBandMatrix(LinearOperator): + """A square 3D :class:`~feectools.linalg.stencil.StencilMatrix` applied as a Kronecker product of 1D banded + matrices, ``(A_1 (x) A_2 (x) A_3) x`` on the rows owned by this process. + + ``A_k[i, d]`` is the entry of owned row ``i`` (along direction ``k``) and diagonal ``d``, column + ``i - q_k + d`` with ``q_k = (n_k - 1) // 2``, as in the stencil matrix. The product is computed in three passes, + each over the owned rows of its direction and the full (padded) extent of the directions after it, so the passes + need no communication beyond the ghost regions of the input. Use :meth:`from_stencil`. + + Parameters + ---------- + stencil : StencilMatrix + The assembled matrix (kept for ``toarray``, ``tosparse`` and ``transpose``). + factors : tuple of arrays + ``A_1, A_2, A_3``, shapes ``(m_k, n_k)``, C-contiguous backend arrays. + """ + + def __init__(self, stencil: StencilMatrix, factors): + self._stencil = stencil + self._factors = factors + + V = stencil.domain + self._m = tuple(int(e - s + 1) for s, e in zip(V.starts, V.ends)) + self._pads = tuple(int(p) for p in V.pads) + self._q = tuple((int(n) - 1) // 2 for n in stencil._data.shape[3:]) + + # temporaries of the first two passes: owned rows along the passes done, padded extent along the others + x_shape = V.zeros()._data.shape + self._t1 = xp.empty((self._m[0], x_shape[1], x_shape[2]), dtype=stencil._data.dtype) + self._t2 = xp.empty((self._m[0], self._m[1], x_shape[2]), dtype=stencil._data.dtype) + + @classmethod + def from_stencil(cls, S: StencilMatrix, rtol: float = KRONECKER_RTOL): + """The Kronecker form of ``S`` if it is one (to ``rtol``, relative to its largest entry), else None. + + Only square 3D matrices on one space (domain is codomain) with shifts 1 qualify, e.g. the diagonal blocks of + the mass matrices. The factors are read from the owned rows of ``S``: with reference rows ``r_k`` (the first + owned row along each direction) and central diagonals ``q_k``, + ``A_1[i, :] = S[i, r_2, r_3, :, q_2, q_3]``, ``A_2[i, :] = S[r_1, i, r_3, q_1, :, q_3] / S[r, q]``, + ``A_3[i, :] = S[r_1, r_2, i, q_1, q_2, :] / S[r, q]``; the product is then compared with all owned rows. + """ + if not isinstance(S, StencilMatrix) or S.domain is not S.codomain or S.domain.ndim != 3: + return None + V = S.domain + if any(int(s) != 1 for s in V.shifts) or S._data.ndim != 6: + return None + + p = [int(p) for p in V.pads] + m = [int(e - s + 1) for s, e in zip(V.starts, V.ends)] + q = [(int(n) - 1) // 2 for n in S._data.shape[3:]] + if min(m) < 1: + return None + data = S._data + + den = data[p[0], p[1], p[2], q[0], q[1], q[2]] + if float(den) == 0.0: + return None + A1 = xp.ascontiguousarray(data[p[0] : p[0] + m[0], p[1], p[2], :, q[1], q[2]]) + A2 = xp.ascontiguousarray(data[p[0], p[1] : p[1] + m[1], p[2], q[0], :, q[2]] / den) + A3 = xp.ascontiguousarray(data[p[0], p[1], p[2] : p[2] + m[2], q[0], q[1], :] / den) + + # compare with every owned row: planes (rows along direction 1) in chunks, the maxima kept on the device + # (one read at the end); a few planes first, which rejects most matrices that are not Kronecker products + # (e.g. ones assembled from particle data, checked again after every change) at once + owned = data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + A23 = A2[:, None, None, :, None] * A3[None, :, None, None, :] # (m2, m3, 1, n2, n3) + + def max_diff_and_abs(rows): + block = owned[rows] + recon = A1[rows][:, None, None, :, None, None] * A23[None] + return xp.max(xp.abs(block - recon)), xp.max(xp.abs(block)) + + def fails(max_diff, max_abs): + max_diff, max_abs = float(max_diff), float(max_abs) + return max_abs == 0.0 or max_diff > rtol * max_abs + + quick = sorted({0, m[0] // 2, m[0] - 1}) + if fails(*max_diff_and_abs(xp.asarray(quick))): + return None + + # chunks of at most ~64 MB of temporaries + plane = max(1, owned[0].size) + chunk = max(1, min(m[0], (8 << 20) // plane)) + max_diff = xp.zeros(()) + max_abs = xp.zeros(()) + for start in range(0, m[0], chunk): + d, a = max_diff_and_abs(slice(start, min(start + chunk, m[0]))) + max_diff = xp.maximum(max_diff, d) + max_abs = xp.maximum(max_abs, a) + if fails(max_diff, max_abs): + return None + return cls(S, (A1, A2, A3)) + + # ------------------------------------------------------------------ + # LinearOperator interface + # ------------------------------------------------------------------ + @property + def domain(self): + return self._stencil.domain + + @property + def codomain(self): + return self._stencil.codomain + + @property + def dtype(self): + return self._stencil.dtype + + @property + def factors(self): + """The 1D factors ``A_1, A_2, A_3`` (rows owned by this process).""" + return self._factors + + def tosparse(self): + return self._stencil.tosparse() + + def toarray(self): + return self._stencil.toarray() + + def transpose(self, conjugate=False): + return self._stencil.transpose(conjugate=conjugate) + + def dot(self, v, out=None): + assert isinstance(v, StencilVector) + assert v.space is self.domain + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.codomain + + if not v.ghost_regions_in_sync: + v.update_ghost_regions() + + A1, A2, A3 = self._factors + p, q, m = self._pads, self._q, self._m + owned = out._data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + + _band_pass(A1, v._data, self._t1, 0, p[0] - q[0]) + _band_pass(A2, self._t1, self._t2, 1, p[1] - q[1]) + _band_pass(A3, self._t2, owned, 2, p[2] - q[2]) + + out.ghost_regions_in_sync = False + return out + + +def kronecker_apply_operator(mat, name: str = ""): + """The operator to apply ``mat`` with: ``mat`` itself, or a copy whose Kronecker blocks are replaced. + + Parameters + ---------- + mat : StencilMatrix or BlockLinearOperator + An assembled (block) mass matrix. + name : str + Name of the operator, for the log. + + Returns + ------- + LinearOperator + ``mat`` if no block is a Kronecker product (or :func:`kronecker_mass_enabled` is False); otherwise a + :class:`KroneckerBandMatrix` (for a StencilMatrix) or a BlockLinearOperator with the same blocks except the + Kronecker ones. It is valid only as long as ``mat`` is not modified. + """ + if not kronecker_mass_enabled(): + return mat + + if isinstance(mat, StencilMatrix): + kron = KroneckerBandMatrix.from_stencil(mat) + if kron is not None: + logger.info(f"Mass matrix {name}: applied as a Kronecker product of 1D banded matrices") + return mat if kron is None else kron + + if isinstance(mat, BlockLinearOperator): + blocks = {} + replaced = [] + for key, block in mat._blocks.items(): + kron = KroneckerBandMatrix.from_stencil(block) + if kron is None: + blocks[key] = block + else: + blocks[key] = kron + replaced.append(key) + if not replaced: + return mat + logger.info(f"Mass matrix {name}: blocks {replaced} applied as Kronecker products of 1D banded matrices") + return BlockLinearOperator(mat.domain, mat.codomain, blocks=blocks) + + return mat diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index e035e4fe4..1020f0ef6 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -25,6 +25,7 @@ from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat from struphy.feec.kernels.kernel_3d_matrixfree import kernel_3d_matrixfree from struphy.feec.kernels.kernel_3d_vec import kernel_3d_vec +from struphy.feec.kronecker_mass import kronecker_apply_operator from struphy.feec.linear_operators import BoundaryOperator from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.feec.utilities import LocalProjectionMatrix, LocalRotationMatrix, get_quad_grids @@ -1703,6 +1704,8 @@ def __init__( blocks, self._weights = self._init_blocks_from_weights(weights_info) self._mat = self._wrap_blocks(blocks) + # the operator dot() applies self._mat with (Kronecker blocks replaced), built on first use, see _apply_mat + self._apply_mat_cache = None # the weights are stored in the block order of the (possibly transposed) operator if transposed: @@ -2121,6 +2124,32 @@ def M0(self): def matrix(self): return self._mat + @property + def _apply_mat(self): + """The operator :meth:`dot` applies: :attr:`matrix`, with its blocks that are Kronecker products of 1D banded + matrices (constant weights, e.g. on a Cuboid) applied as such (:mod:`struphy.feec.kronecker_mass`). + + Built on first use after every change of the matrix (assembly, copy into, in-place arithmetic). + """ + if self._apply_mat_cache is None: + if getattr(self, "_kronecker_rejected", False): + self._apply_mat_cache = self._mat + else: + self._apply_mat_cache = kronecker_apply_operator(self._mat, name=self.name) + # a matrix changed by in-place arithmetic that is not a Kronecker product (e.g. the Schur complement + # M1 - dt^2 BC with particle data, rebuilt every step) is not checked again after further arithmetic: + # the check reads the whole matrix. Assembly and copies check again. + if getattr(self, "_changed_by_arithmetic", False) and self._apply_mat_cache is self._mat: + self._kronecker_rejected = True + return self._apply_mat_cache + + def _matrix_changed(self, by_arithmetic: bool): + """Invalidate :attr:`_apply_mat` after a change of the matrix (in-place arithmetic, or assembly/copy).""" + self._apply_mat_cache = None + self._changed_by_arithmetic = by_arithmetic + if not by_arithmetic: + self._kronecker_rejected = False + @property def nquads(self): if self._nquads is None: @@ -2168,29 +2197,54 @@ def dot(self, v, out=None, apply_bc=True): if apply_bc: if self._transposed: - self._W_boundary_op_T.dot(v, out=self._temp_WB) - self._W_extraction_op_T.dot(self._temp_WB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) - self._V_extraction_op.dot(self._temp_VE, out=self._temp_VB) - out = self._V_boundary_op.dot(self._temp_VB, out=out) + ops_in = (self._W_boundary_op_T, self._W_extraction_op_T) + ops_out = (self._V_extraction_op, self._V_boundary_op) + tmps_in = (self._temp_WB, self._temp_mat) + tmps_out = (self._temp_VE, self._temp_VB) else: - self._V_boundary_op_T.dot(v, out=self._temp_VB) - self._V_extraction_op_T.dot(self._temp_VB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) - self._W_extraction_op.dot(self._temp_WE, out=self._temp_WB) - out = self._W_boundary_op.dot(self._temp_WB, out=out) + ops_in = (self._V_boundary_op_T, self._V_extraction_op_T) + ops_out = (self._W_extraction_op, self._W_boundary_op) + tmps_in = (self._temp_VB, self._temp_mat) + tmps_out = (self._temp_WE, self._temp_WB) + out = self._apply_chain(v, ops_in, tmps_in, ops_out, tmps_out, out) else: if self._transposed: self._W_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) + self._apply_mat.dot(self._temp_mat, out=self._temp_VE) out = self._V_extraction_op.dot(self._temp_VE, out=out) else: self._V_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) + self._apply_mat.dot(self._temp_mat, out=self._temp_WE) out = self._W_extraction_op.dot(self._temp_WE, out=out) return out + def _apply_chain(self, v, ops_in, tmps_in, ops_out, tmps_out, out): + """``out = B_out E_out M E_in B_in v`` (boundary and extraction operators around the matrix), see :meth:`dot`. + + Identity operators (e.g. the extraction operators without polar splines, or boundary operators without + essential boundary conditions) are skipped instead of applied as copies: each vector is passed on as it is, + and the matrix writes into the next temporary that is actually used (or into ``out``). + """ + if out is v: + # in place: the chain must not read v after writing the output, so it reads a copy + v = v.copy() + + x = v + for op, tmp in zip(ops_in, tmps_in): + if not isinstance(op, IdentityOperator): + x = op.dot(x, out=tmp) + + # the first output operator that is not an identity reads the matrix product from its temporary + remaining = [(op, tmp) for op, tmp in zip(ops_out, tmps_out) if not isinstance(op, IdentityOperator)] + if not remaining: + return self._apply_mat.dot(x, out=out) + y = self._apply_mat.dot(x, out=remaining[0][1]) + for k, (op, _) in enumerate(remaining): + target = out if k == len(remaining) - 1 else remaining[k + 1][1] + y = op.dot(y, out=target) + return y + def transpose(self, conjugate=False): """ Returns the transposed operator. @@ -2277,6 +2331,7 @@ def assemble(self, weights=None, clear=True): Whether to first set all data to zero before assembly. If False, the new contributions are added to existing ones. """ + self._matrix_changed(by_arithmetic=False) # the matrix changes assert not self._dry_run, ( "A dry-run operator has no matrix data and cannot be assembled (memory estimation only)." ) @@ -2559,6 +2614,7 @@ def copy(self, out=None): out._weights = [list(row) for row in self._weights] self._mat.copy(out=out._mat) + out._matrix_changed(by_arithmetic=False) if self._creation_info is None: out._creation_info = None @@ -2569,6 +2625,7 @@ def copy(self, out=None): def __imul__(self, a): self._mat *= a + self._matrix_changed(by_arithmetic=True) self._creation_info = None return self @@ -2579,11 +2636,13 @@ def __iadd__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat += M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat += M + self._matrix_changed(by_arithmetic=True) return self else: @@ -2596,11 +2655,13 @@ def __isub__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat -= M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat -= M + self._matrix_changed(by_arithmetic=True) return self else: From 07fd091596812c93fdd1afa1a8b68b4199c49fdc Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:20:50 +0200 Subject: [PATCH 26/55] Kronecker mass: launch the band-pass kernel through CachedLaunch The banded passes run with the same buffers in every solver iteration (as on cuda-aware-mpi-fixes, now that feectools has CachedLaunch). Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/kronecker_mass.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py index 2050edb82..5456b2eeb 100644 --- a/src/struphy/feec/kronecker_mass.py +++ b/src/struphy/feec/kronecker_mass.py @@ -38,8 +38,10 @@ def kronecker_mass_enabled() -> bool: def _band_pass_cuda(): """The CUDA kernel of one banded pass (compiled on first call).""" from cunumpy.kernels import CudaKernel + from feectools.utilities.cached_launch import CachedLaunch - return CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass") + # launched with the same buffers in every solver iteration: cache the launch preparation + return CachedLaunch(CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass")) def _band_pass(A, inp, out, axis, offset): From 2f091cdcf4a887b96849c84c907e1bfcd37d14bd Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 09:17:35 +0200 Subject: [PATCH 27/55] Bump feectools version number --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index 1d8d47c4a..e75b4c766 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 1d8d47c4a3bc8e0b594992cc22944b1ac4230c0f +Subproject commit e75b4c76616ef93f0091d17e752e4155f430c5d2 From f1cbb2d1200ce79c864c81dd4ef04280c85c198b Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:00:05 +0200 Subject: [PATCH 28/55] CuPy backend: host/device conversions in post-processing; math.sqrt for host scalars in Species Post-processing evaluates fields on the device and writes them from the host: the evaluation grids, point data and coordinates are converted where they cross, and HDF5 data is moved to the device before it is written into the spline coefficients. Co-Authored-By: Claude Opus 5.5 --- src/struphy/models/species.py | 4 ++-- src/struphy/post_processing/output.py | 31 ++++++++++++++++++--------- 2 files changed, 23 insertions(+), 12 deletions(-) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index 3c8b47657..4bc2dcac1 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,9 +1,9 @@ import logging +import math import warnings from abc import ABCMeta, abstractmethod from dataclasses import fields, is_dataclass -import cunumpy as xp from maybempi import MPI from struphy.models.variables import Variable @@ -203,7 +203,7 @@ def __init__( con = ConstantsOfNature() # relevant frequencies - om_p = xp.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) + om_p = math.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) om_c = Z * con.e * units.B / (A * con.mH) # compute equation parameters diff --git a/src/struphy/post_processing/output.py b/src/struphy/post_processing/output.py index bfcf5fd58..cc65d56f8 100644 --- a/src/struphy/post_processing/output.py +++ b/src/struphy/post_processing/output.py @@ -85,6 +85,13 @@ def load_plasma_plots() -> bool: return True +def _to_host(value): + """`value` (an array, or a list/tuple of arrays for vector-valued fields) as NumPy array(s).""" + if isinstance(value, (list, tuple)): + return type(value)(_to_host(v) for v in value) + return xp.to_numpy(value) + + class ProductMapping(Mapping[str, xr.DataArray]): """A discoverable mapping whose products are loaded on first access.""" @@ -568,13 +575,14 @@ def _attach_physical_coords( self, array: xr.DataArray, etas: tuple[Any, Any, Any], dims: tuple[str, ...] ) -> xr.DataArray: """Map a raw FEEC evaluation grid and attach its Cartesian coordinates.""" - grids = [np.atleast_1d(np.asarray(eta, dtype=float)) for eta in etas] - mapped = self.domain(*grids) + # the logical grids may be device arrays (CuPy backend): map them where they live, keep the coordinates on the host + grids = [np.atleast_1d(np.asarray(xp.to_numpy(eta), dtype=float)) for eta in etas] + mapped = self.domain(*(xp.asarray(grid) for grid in grids)) shape = tuple(len(grid) for grid in grids) keep = tuple(slice(None) if dim in dims else 0 for dim in ("eta1", "eta2", "eta3")) coordinates = {} for name, values in zip(("X", "Y", "Z"), mapped): - values = np.asarray(values).reshape(shape)[keep] + values = np.asarray(xp.to_numpy(values)).reshape(shape)[keep] coordinates[name] = (dims, values) if dims else values.item() return array.assign_coords(coordinates) @@ -1295,7 +1303,9 @@ def _process_fields( # the physical grid is only needed for output, hence it is only built on rank 0 if self._pproc_rank == 0: - grids_phy = list(self.domain(*grids_log)) + # output coordinates, on the host (the logical grids stay on the device for the field evaluation) + grids_phy = [xp.to_numpy(grid) for grid in self.domain(*grids_log)] + grids_log_host = [xp.to_numpy(grid) for grid in grids_log] else: grids_phy = None @@ -1327,8 +1337,9 @@ def _process_fields( if self._pproc_rank == 0: for species, vars in vals.items(): for name, val in vars.items(): - point_data[species][name][t] = val - point_data_phy[species][name][t] = vals_phy[species][name] + # evaluated on the device on the CuPy backend; written to files from the host + point_data[species][name][t] = _to_host(val) + point_data_phy[species][name][t] = _to_host(vals_phy[species][name]) # directory for the vtk files path_fields = os.path.join(self.path_pproc, "fields_data") @@ -1338,10 +1349,10 @@ def _process_fields( for species, vars in point_data.items(): variables = {} for name, val in vars.items(): - variables[name] = wrap_field_data(val, grids_log, grids_phy=grids_phy, name=name) + variables[name] = wrap_field_data(val, grids_log_host, grids_phy=grids_phy, name=name) if physical: variables[name + "_xyz"] = wrap_field_data( - point_data_phy[species][name], grids_log, grids_phy=grids_phy, name=name + "_xyz" + point_data_phy[species][name], grids_log_host, grids_phy=grids_phy, name=name + "_xyz" ) store.write_group(store.store_path(self.path_pproc), f"/{species}", xr.Dataset(variables)) @@ -1517,7 +1528,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[n * step, p1:-p1, p2:-p2, p3:-p3]) # HDF5 data is on the host # vector-valued field else: @@ -1530,7 +1541,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3]) vector.update_ghost_regions() if is_polar: From ef927d1bd4fc2cef3c17a566dd1e0fc30d0043e6 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:07:26 +0200 Subject: [PATCH 29/55] CUDA version of the mass-matrix assembly kernel kernel_3d_mat - kernel_3d_mat_cuda.cu: one thread per (element, codomain basis function, domain basis function); replaces the host fallback on the CuPy backend. - WeightedMassOperator.assemble reads the quadrature grids (spans, weights, basis values) as backend arrays, copied to the device once and cached. - Simulation.allocate compiles the assembly kernel before it is first called, so the compilation is not counted in the assembly's profiling regions. Co-Authored-By: Claude Opus 5.5 --- .../feec/kernels/kernel_3d_mat/__init__.py | 18 ++-- .../kernel_3d_mat/kernel_3d_mat_cuda.cu | 92 +++++++++++++++++++ src/struphy/feec/mass.py | 21 ++++- src/struphy/simulation/sim.py | 18 ++++ 4 files changed, 139 insertions(+), 10 deletions(-) create mode 100644 src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu diff --git a/src/struphy/feec/kernels/kernel_3d_mat/__init__.py b/src/struphy/feec/kernels/kernel_3d_mat/__init__.py index a881301b4..76da51673 100644 --- a/src/struphy/feec/kernels/kernel_3d_mat/__init__.py +++ b/src/struphy/feec/kernels/kernel_3d_mat/__init__.py @@ -1,8 +1,4 @@ -"""The kernel of this folder: pyccel in ``_kernels.py``, CUDA in ``_cuda.cu`` once ported. - -No CUDA version yet: on the CuPy backend the pyccel kernel runs on host copies of the arrays (a warning is -emitted once), as before the move into this folder. -""" +"""The kernel of this folder: pyccel in ``kernel_3d_mat_kernels.py``, CUDA in ``kernel_3d_mat_cuda.cu``.""" from cunumpy.kernels import Kernel @@ -11,6 +7,16 @@ # the arguments the kernel writes to, by position: data OUTPUTS = (25,) + +def _n_threads(args): + """One thread per (element, codomain basis function, domain basis function), see kernel_3d_mat_cuda.cu.""" + n_elements = args[0].size * args[1].size * args[2].size + n_pairs = 1 + for degree in args[3:9]: # pi1, pi2, pi3, pj1, pj2, pj3 + n_pairs *= int(degree) + 1 + return n_elements * n_pairs + + kernel_3d_mat = Kernel.from_folder( - __name__, structs=CUDA_STRUCTS, missing_cuda="fallback", host_options={"outputs": OUTPUTS} + __name__, structs=CUDA_STRUCTS, host_options={"outputs": OUTPUTS}, n_threads_from=_n_threads ) diff --git a/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu b/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu new file mode 100644 index 000000000..6afab9398 --- /dev/null +++ b/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu @@ -0,0 +1,92 @@ +// CUDA version of kernel_3d_mat (kernel_3d_mat_kernels.py), with the same arguments in the same order. +#include "cunumpy/array_view.cuh" +#include "cunumpy/atomic.cuh" +#include "cunumpy/index.cuh" + +/** + * Integral of Lambda_i * mat_fun * Lambda_j over the local elements, added to the stencil matrix data, as in + * kernel_3d_mat_kernels.kernel_3d_mat. + * + * One thread per (element, codomain basis function (il1, il2, il3), domain basis function (jl1, jl2, jl3)): + * n_threads = ne1 * ne2 * ne3 * (pi1 + 1) * (pi2 + 1) * (pi3 + 1) * (pj1 + 1) * (pj2 + 1) * (pj3 + 1) (see + * __init__.py), the element index varying slowest and jl3 fastest. Each thread computes the quadrature sum of its + * pair exactly as the pyccel kernel (same factors, same order) and adds it to its entry of `data` with an atomic + * addition: neighbouring elements contribute to the same entries, so only the summation order over elements + * differs from pyccel (round-off). + * + * @param spans1, spans2, spans3 Knot span index of each local element, per direction. + * @param pi1, pi2, pi3 Degrees of the codomain basis functions. + * @param pj1, pj2, pj3 Degrees of the domain basis functions. + * @param starts1, starts2, starts3 Global start index of the rows on this process. + * @param pads1, pads2, pads3 Paddings (ghost regions) of `data`. + * @param w1, w2, w3 Quadrature weights [element, quadrature point]. + * @param bi1, bi2, bi3 Codomain basis functions [element, local basis function, derivative, quadrature point]. + * @param bj1, bj2, bj3 Domain basis functions, same layout. + * @param mat_fun Weight function at the quadrature points [flattened point 1, flattened point 2, flattened point 3]. + * @param data StencilMatrix data (rows with ghost regions..., diagonals...); the results are added (not set). + */ +extern "C" __global__ void kernel_3d_mat(Array1D spans1, Array1D spans2, + Array1D spans3, int pi1, int pi2, int pi3, int pj1, int pj2, + int pj3, int starts1, int starts2, int starts3, int pads1, int pads2, + int pads3, Array2D w1, Array2D w2, Array2D w3, + Array4D bi1, Array4D bi2, Array4D bi3, + Array4D bj1, Array4D bj2, Array4D bj3, + Array3D mat_fun, Array6D data) +{ + const long long ne1 = spans1.size(); + const long long ne2 = spans2.size(); + const long long ne3 = spans3.size(); + + const long long ni1 = pi1 + 1, ni2 = pi2 + 1, ni3 = pi3 + 1; + const long long nj1 = pj1 + 1, nj2 = pj2 + 1, nj3 = pj3 + 1; + const long long n_pairs = ni1 * ni2 * ni3 * nj1 * nj2 * nj3; + + CUNUMPY_THREAD_1D(thread, ne1 * ne2 * ne3 * n_pairs); + + // CUDA-only: split the thread index into the pyccel loop variables (iel1, iel2, iel3, il1, il2, il3, jl1, jl2, + // jl3), jl3 fastest + long long rest = thread; + const long long jl3 = rest % nj3; rest /= nj3; + const long long jl2 = rest % nj2; rest /= nj2; + const long long jl1 = rest % nj1; rest /= nj1; + const long long il3 = rest % ni3; rest /= ni3; + const long long il2 = rest % ni2; rest /= ni2; + const long long il1 = rest % ni1; rest /= ni1; + const long long iel3 = rest % ne3; rest /= ne3; + const long long iel2 = rest % ne2; rest /= ne2; + const long long iel1 = rest; + + const long long nq1 = w1.shape[1]; + const long long nq2 = w2.shape[1]; + const long long nq3 = w3.shape[1]; + + double value = 0.0; + + for (long long q1 = 0; q1 < nq1; ++q1) { + for (long long q2 = 0; q2 < nq2; ++q2) { + for (long long q3 = 0; q3 < nq3; ++q3) { + const double wvol = w1(iel1, q1) * w2(iel2, q2) * w3(iel3, q3) * + mat_fun(iel1 * nq1 + q1, iel2 * nq2 + q2, iel3 * nq3 + q3); + + const double bi = bi1(iel1, il1, 0, q1) * bi2(iel2, il2, 0, q2) * bi3(iel3, il3, 0, q3); + const double bj = bj1(iel1, jl1, 0, q1) * bj2(iel2, jl2, 0, q2) * bj3(iel3, jl3, 0, q3); + + value += wvol * bi * bj; + } + } + } + + // global spline indices + const long long i_global1 = spans1(iel1) - pi1 + il1; + const long long i_global2 = spans2(iel2) - pi2 + il2; + const long long i_global3 = spans3(iel3) - pi3 + il3; + + // local spline indices (- starts --> can be negative, will therefore be written to ghost regions) + const long long i_local1 = i_global1 - starts1; + const long long i_local2 = i_global2 - starts2; + const long long i_local3 = i_global3 - starts3; + + cunumpy_atomic_add(&data(pads1 + i_local1, pads2 + i_local2, pads3 + i_local3, pads1 + jl1 - il1, + pads2 + jl2 - il2, pads3 + jl3 - il3), + value); +} diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index f75c6de17..e035e4fe4 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -2242,6 +2242,19 @@ def transpose(self, conjugate=False): return M + def _quad_grid(self, space_name, attribute, component): + """Quadrature-grid arrays (spans, weights or basis values) of one component of a space, as backend arrays. + + The Derham keeps them as host arrays; on the CuPy backend the assembly kernel needs them on the device, + so they are copied there once and cached (they do not change between assemblies). + """ + cache = self.__dict__.setdefault("_quad_grid_cache", {}) + key = (space_name, attribute, component) + if key not in cache: + arrays = getattr(self.derham.spline_attributes[space_name], attribute)[component] + cache[key] = [xp.asarray(array) for array in arrays] + return cache[key] + def assemble(self, weights=None, clear=True): r""" Assembles the weighted mass matrix, i.e. computes the integrals @@ -2320,7 +2333,7 @@ def assemble(self, weights=None, clear=True): # loop over codomain spaces (rows) for a, codomain_space in enumerate(codomain_spaces): # knot span indices of elements of local domain - codomain_spans = spline_attr[W_name].quad_grid_spans[a] + codomain_spans = self._quad_grid(W_name, "quad_grid_spans", a) # global start spline index on process codomain_starts = [int(start) for start in codomain_space.coeff_space.starts] @@ -2332,10 +2345,10 @@ def assemble(self, weights=None, clear=True): pts = self._quad_pts(W_name, a) # global quadrature weights in format (local element, local weight) - wts = spline_attr[W_name].quad_grid_wts[a] + wts = self._quad_grid(W_name, "quad_grid_wts", a) # evaluated basis functions at quadrature points of codomain space - codomain_basis = spline_attr[W_name].quad_grid_bases[a] + codomain_basis = self._quad_grid(W_name, "quad_grid_bases", a) # loop over domain spaces (columns) for b, domain_space in enumerate(domain_spaces): @@ -2375,7 +2388,7 @@ def assemble(self, weights=None, clear=True): ) # evaluated basis functions at quadrature points of domain space - domain_basis = spline_attr[V_name].quad_grid_bases[b] + domain_basis = self._quad_grid(V_name, "quad_grid_bases", b) # assemble matrix (if mat_w is not zero) by calling the appropriate kernel (1d, 2d or 3d) if not_weight_zero or self._is_scalar: diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 9452e81c5..e6f358832 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -243,6 +243,8 @@ def allocate(self): logger.debug("\nAllocating simulation data ...") + self._compile_setup_cuda_kernels() + with ProfileManager.profile_region("setup: allocate"): # feec with ProfileManager.profile_region("allocate feec", functions=[self._allocate_feec]): @@ -280,6 +282,22 @@ def kernels(self) -> tuple[Kernel, ...]: owners += self.model.prop_list return collect_kernels(*owners) + def _compile_setup_cuda_kernels(self) -> None: + """Compile the CUDA kernels that :meth:`allocate` itself calls (mass-matrix assembly) before it calls them. + + :meth:`compile_cuda_kernels` runs after :meth:`allocate`, so without this the first assembly would include + the compilation (~1 s per kernel on a node whose CuPy cache does not have it yet), in the assembly's + profiling regions. Here it is counted in ``"setup: compile cuda kernels (assembly)"``. Nothing happens on + the NumPy backend. + """ + if xp.get_backend() != "cupy": + return + + from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat + + with ProfileManager.profile_region("setup: compile cuda kernels (assembly)"): + kernel_3d_mat.compile() + def compile_cuda_kernels(self) -> tuple[Kernel, ...]: """Compile the CUDA versions of the kernels of :meth:`kernels` now, before the time stepping. From c1f8875b99a8f3e4ef889c2ec8063af5b76092e1 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:26:26 +0200 Subject: [PATCH 30/55] Parity cases for kernel_3d_mat Random quadrature data on 3 x 2 x 2 elements with mixed degrees, with the first row local or a ghost row. test_cuda_kernels_have_parity_cases requires a case for every kernel with a CUDA version; the emulation test compares it with the pyccel kernel without a GPU. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/cuda_parity_cases.py | 26 ++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index a2396ee5b..6e10707be 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -315,6 +315,30 @@ def kernel_pullpush_args(case): ) +def kernel_3d_mat_threads(args): + """One thread per (element, codomain basis function, domain basis function), as the kernel's launch size.""" + from struphy.feec.kernels.kernel_3d_mat import _n_threads + + return _n_threads(args) + + +def kernel_3d_mat_args(case): + """Mass-matrix assembly on 3 x 2 x 2 elements with mixed degrees and random quadrature data; `case` is the start + index of the rows in direction 1 (1: the first row is a ghost row, local index -1).""" + rng = np.random.default_rng(20 + case) + ne, nq = (3, 2, 2), (3, 2, 2) + pi, pj, pads = (2, 1, 1), (1, 2, 1), (2, 2, 2) + starts = (case, 0, 0) + spans = [xp.asarray(np.arange(n, dtype=np.int64) + p) for n, p in zip(ne, pi)] + w = [xp.asarray(rng.random((n, q))) for n, q in zip(ne, nq)] + bi = [xp.asarray(rng.random((n, p + 1, 1, q))) for n, p, q in zip(ne, pi, nq)] + bj = [xp.asarray(rng.random((n, p + 1, 1, q))) for n, p, q in zip(ne, pj, nq)] + mat_fun = xp.asarray(rng.normal(size=tuple(n * q for n, q in zip(ne, nq)))) + rows = tuple(n + p + 2 * pad for n, p, pad in zip(ne, pi, pads)) + data = xp.asarray(rng.normal(size=rows + tuple(2 * pad + 1 for pad in pads))) # added to, not overwritten + return (*spans, *pi, *pj, *starts, *pads, *w, *bi, *bj, mat_fun, data) + + # ---------------------------------------------------------------- all kernels with a CUDA version PUSHER_TOLERANCES = {"rtol": 1e-13, "atol": 1e-14} @@ -338,6 +362,8 @@ def size_of(index, per_thread=1): PARITY_CASES = { + # atomic additions: the summation order over elements differs from pyccel + "kernel_3d_mat": ParityCases((0, 1), kernel_3d_mat_args, rtol=1e-12, atol=1e-13, n_threads=kernel_3d_mat_threads), "push_eta_stage": ParityCases( [ (bc, method, stage) From 890b6d93cd3a5e22fb05141b5ec460375e472e23 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:08:01 +0200 Subject: [PATCH 31/55] Apply separable mass matrices as Kronecker products of 1D banded matrices WeightedMassOperator.dot applies the blocks of its matrix that are Kronecker products of 1D banded matrices (constant weights, e.g. on a Cuboid) as three banded passes (struphy.feec.kronecker_mass, CUDA kernel kron_band_pass on the GPU) instead of the 3D stencil product. The replacement is built on first use and invalidated by assembly, copies and in-place arithmetic; a matrix changed by arithmetic that is not a Kronecker product is not checked again. STRUPHY_KRONECKER_MASS=0 switches it off. The boundary and extraction operators around the matrix are skipped when they are identities instead of being applied as copies. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/kron_band_pass_cuda.cu | 35 ++++ src/struphy/feec/kronecker_mass.py | 247 ++++++++++++++++++++++++ src/struphy/feec/mass.py | 85 ++++++-- 3 files changed, 355 insertions(+), 12 deletions(-) create mode 100644 src/struphy/feec/kron_band_pass_cuda.cu create mode 100644 src/struphy/feec/kronecker_mass.py diff --git a/src/struphy/feec/kron_band_pass_cuda.cu b/src/struphy/feec/kron_band_pass_cuda.cu new file mode 100644 index 000000000..745e782f4 --- /dev/null +++ b/src/struphy/feec/kron_band_pass_cuda.cu @@ -0,0 +1,35 @@ +/** + * One pass of a Kronecker product of 1D banded matrices (struphy.feec.kronecker_mass._band_pass): + * + * out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...] along `axis`, + * + * with the other two axes unchanged (out and inp have the same extent along them). One thread per entry of `out`, + * the last axis fastest, so that neighbouring threads read neighbouring entries of `inp` along every axis. `out` may + * be a strided view (the owned rows of a stencil vector's data). + * + * @param A Banded 1D matrix, shape (out.shape[axis], n_diagonals). + * @param inp Input array (3D). + * @param out Output array (3D, any strides); written, not added to. + * @param axis Axis of the pass (0, 1 or 2). + * @param offset Index in `inp` of diagonal 0 of row 0 along `axis`. + */ +#include "cunumpy/array_view.cuh" + +extern "C" __global__ void kron_band_pass(Array2D A, Array3D inp, Array3D out, int axis, + long long offset) { + const long long n1 = out.shape[1], n2 = out.shape[2]; + const long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; + if (t >= out.shape[0] * n1 * n2) return; + const long long i2 = t % n2; + const long long i1 = (t / n2) % n1; + const long long i0 = t / (n1 * n2); + const long long i = axis == 0 ? i0 : (axis == 1 ? i1 : i2); + + double v = 0.0; + for (long long d = 0; d < A.shape[1]; ++d) { + const long long j = i + offset + d; + const double x = axis == 0 ? inp(j, i1, i2) : (axis == 1 ? inp(i0, j, i2) : inp(i0, i1, j)); + v += A(i, d) * x; + } + out(i0, i1, i2) = v; +} diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py new file mode 100644 index 000000000..2050edb82 --- /dev/null +++ b/src/struphy/feec/kronecker_mass.py @@ -0,0 +1,247 @@ +"""Fast application of assembled mass matrices that are Kronecker products of 1D banded matrices. + +A mass matrix block ``M[i, j] = int w Lambda_i Lambda_j`` of tensor-product splines factorizes as +``M = c * M_1 (x) M_2 (x) M_3`` when its weight ``w`` is constant (e.g. on a Cuboid). Applied as a +:class:`~feectools.linalg.stencil.StencilMatrix`, every row reads all ``n_1 * n_2 * n_3`` diagonals (105 for +degrees (1, 2, 3)) from memory; as a Kronecker product it is three 1D banded passes over the vector, ``n_1 + n_2 + +n_3`` operations per entry and no 3D matrix data. :func:`kronecker_apply_operator` checks each assembled block and +returns an operator for the matrix-vector products in which the blocks that are Kronecker products (to round-off) +are replaced by :class:`KroneckerBandMatrix`. + +The 1D factors are read from the assembled data itself (rows owned by this process), so they hold exactly what the +assembly computed, on any domain decomposition, and the check is the assembled data, not an assumption about the +mapping or the weights. +""" + +import logging +import os +from functools import cache +from pathlib import Path + +import cunumpy as xp +from feectools.linalg.basic import LinearOperator +from feectools.linalg.block import BlockLinearOperator +from feectools.linalg.stencil import StencilMatrix, StencilVector + +logger = logging.getLogger("struphy") + +# relative tolerance of the Kronecker check: the 3D assembly and the product of the factors differ by round-off only +KRONECKER_RTOL = 1e-12 + + +def kronecker_mass_enabled() -> bool: + """Whether :func:`kronecker_apply_operator` may replace blocks (``STRUPHY_KRONECKER_MASS=0`` switches it off).""" + return os.environ.get("STRUPHY_KRONECKER_MASS", "1").lower() not in ("0", "false", "no") + + +@cache +def _band_pass_cuda(): + """The CUDA kernel of one banded pass (compiled on first call).""" + from cunumpy.kernels import CudaKernel + + return CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass") + + +def _band_pass(A, inp, out, axis, offset): + """``out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...]`` along ``axis`` (other axes unchanged).""" + if xp.is_gpu(out): + _band_pass_cuda()(A, inp, out, axis, offset, n_threads=out.size) + return + m = out.shape[axis] + shape = [1, 1, 1] + shape[axis] = m + out[...] = 0.0 + for d in range(A.shape[1]): + idx = [slice(None)] * 3 + idx[axis] = slice(offset + d, offset + d + m) + out += A[:, d].reshape(shape) * inp[tuple(idx)] + + +class KroneckerBandMatrix(LinearOperator): + """A square 3D :class:`~feectools.linalg.stencil.StencilMatrix` applied as a Kronecker product of 1D banded + matrices, ``(A_1 (x) A_2 (x) A_3) x`` on the rows owned by this process. + + ``A_k[i, d]`` is the entry of owned row ``i`` (along direction ``k``) and diagonal ``d``, column + ``i - q_k + d`` with ``q_k = (n_k - 1) // 2``, as in the stencil matrix. The product is computed in three passes, + each over the owned rows of its direction and the full (padded) extent of the directions after it, so the passes + need no communication beyond the ghost regions of the input. Use :meth:`from_stencil`. + + Parameters + ---------- + stencil : StencilMatrix + The assembled matrix (kept for ``toarray``, ``tosparse`` and ``transpose``). + factors : tuple of arrays + ``A_1, A_2, A_3``, shapes ``(m_k, n_k)``, C-contiguous backend arrays. + """ + + def __init__(self, stencil: StencilMatrix, factors): + self._stencil = stencil + self._factors = factors + + V = stencil.domain + self._m = tuple(int(e - s + 1) for s, e in zip(V.starts, V.ends)) + self._pads = tuple(int(p) for p in V.pads) + self._q = tuple((int(n) - 1) // 2 for n in stencil._data.shape[3:]) + + # temporaries of the first two passes: owned rows along the passes done, padded extent along the others + x_shape = V.zeros()._data.shape + self._t1 = xp.empty((self._m[0], x_shape[1], x_shape[2]), dtype=stencil._data.dtype) + self._t2 = xp.empty((self._m[0], self._m[1], x_shape[2]), dtype=stencil._data.dtype) + + @classmethod + def from_stencil(cls, S: StencilMatrix, rtol: float = KRONECKER_RTOL): + """The Kronecker form of ``S`` if it is one (to ``rtol``, relative to its largest entry), else None. + + Only square 3D matrices on one space (domain is codomain) with shifts 1 qualify, e.g. the diagonal blocks of + the mass matrices. The factors are read from the owned rows of ``S``: with reference rows ``r_k`` (the first + owned row along each direction) and central diagonals ``q_k``, + ``A_1[i, :] = S[i, r_2, r_3, :, q_2, q_3]``, ``A_2[i, :] = S[r_1, i, r_3, q_1, :, q_3] / S[r, q]``, + ``A_3[i, :] = S[r_1, r_2, i, q_1, q_2, :] / S[r, q]``; the product is then compared with all owned rows. + """ + if not isinstance(S, StencilMatrix) or S.domain is not S.codomain or S.domain.ndim != 3: + return None + V = S.domain + if any(int(s) != 1 for s in V.shifts) or S._data.ndim != 6: + return None + + p = [int(p) for p in V.pads] + m = [int(e - s + 1) for s, e in zip(V.starts, V.ends)] + q = [(int(n) - 1) // 2 for n in S._data.shape[3:]] + if min(m) < 1: + return None + data = S._data + + den = data[p[0], p[1], p[2], q[0], q[1], q[2]] + if float(den) == 0.0: + return None + A1 = xp.ascontiguousarray(data[p[0] : p[0] + m[0], p[1], p[2], :, q[1], q[2]]) + A2 = xp.ascontiguousarray(data[p[0], p[1] : p[1] + m[1], p[2], q[0], :, q[2]] / den) + A3 = xp.ascontiguousarray(data[p[0], p[1], p[2] : p[2] + m[2], q[0], q[1], :] / den) + + # compare with every owned row: planes (rows along direction 1) in chunks, the maxima kept on the device + # (one read at the end); a few planes first, which rejects most matrices that are not Kronecker products + # (e.g. ones assembled from particle data, checked again after every change) at once + owned = data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + A23 = A2[:, None, None, :, None] * A3[None, :, None, None, :] # (m2, m3, 1, n2, n3) + + def max_diff_and_abs(rows): + block = owned[rows] + recon = A1[rows][:, None, None, :, None, None] * A23[None] + return xp.max(xp.abs(block - recon)), xp.max(xp.abs(block)) + + def fails(max_diff, max_abs): + max_diff, max_abs = float(max_diff), float(max_abs) + return max_abs == 0.0 or max_diff > rtol * max_abs + + quick = sorted({0, m[0] // 2, m[0] - 1}) + if fails(*max_diff_and_abs(xp.asarray(quick))): + return None + + # chunks of at most ~64 MB of temporaries + plane = max(1, owned[0].size) + chunk = max(1, min(m[0], (8 << 20) // plane)) + max_diff = xp.zeros(()) + max_abs = xp.zeros(()) + for start in range(0, m[0], chunk): + d, a = max_diff_and_abs(slice(start, min(start + chunk, m[0]))) + max_diff = xp.maximum(max_diff, d) + max_abs = xp.maximum(max_abs, a) + if fails(max_diff, max_abs): + return None + return cls(S, (A1, A2, A3)) + + # ------------------------------------------------------------------ + # LinearOperator interface + # ------------------------------------------------------------------ + @property + def domain(self): + return self._stencil.domain + + @property + def codomain(self): + return self._stencil.codomain + + @property + def dtype(self): + return self._stencil.dtype + + @property + def factors(self): + """The 1D factors ``A_1, A_2, A_3`` (rows owned by this process).""" + return self._factors + + def tosparse(self): + return self._stencil.tosparse() + + def toarray(self): + return self._stencil.toarray() + + def transpose(self, conjugate=False): + return self._stencil.transpose(conjugate=conjugate) + + def dot(self, v, out=None): + assert isinstance(v, StencilVector) + assert v.space is self.domain + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.codomain + + if not v.ghost_regions_in_sync: + v.update_ghost_regions() + + A1, A2, A3 = self._factors + p, q, m = self._pads, self._q, self._m + owned = out._data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + + _band_pass(A1, v._data, self._t1, 0, p[0] - q[0]) + _band_pass(A2, self._t1, self._t2, 1, p[1] - q[1]) + _band_pass(A3, self._t2, owned, 2, p[2] - q[2]) + + out.ghost_regions_in_sync = False + return out + + +def kronecker_apply_operator(mat, name: str = ""): + """The operator to apply ``mat`` with: ``mat`` itself, or a copy whose Kronecker blocks are replaced. + + Parameters + ---------- + mat : StencilMatrix or BlockLinearOperator + An assembled (block) mass matrix. + name : str + Name of the operator, for the log. + + Returns + ------- + LinearOperator + ``mat`` if no block is a Kronecker product (or :func:`kronecker_mass_enabled` is False); otherwise a + :class:`KroneckerBandMatrix` (for a StencilMatrix) or a BlockLinearOperator with the same blocks except the + Kronecker ones. It is valid only as long as ``mat`` is not modified. + """ + if not kronecker_mass_enabled(): + return mat + + if isinstance(mat, StencilMatrix): + kron = KroneckerBandMatrix.from_stencil(mat) + if kron is not None: + logger.info(f"Mass matrix {name}: applied as a Kronecker product of 1D banded matrices") + return mat if kron is None else kron + + if isinstance(mat, BlockLinearOperator): + blocks = {} + replaced = [] + for key, block in mat._blocks.items(): + kron = KroneckerBandMatrix.from_stencil(block) + if kron is None: + blocks[key] = block + else: + blocks[key] = kron + replaced.append(key) + if not replaced: + return mat + logger.info(f"Mass matrix {name}: blocks {replaced} applied as Kronecker products of 1D banded matrices") + return BlockLinearOperator(mat.domain, mat.codomain, blocks=blocks) + + return mat diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index e035e4fe4..1020f0ef6 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -25,6 +25,7 @@ from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat from struphy.feec.kernels.kernel_3d_matrixfree import kernel_3d_matrixfree from struphy.feec.kernels.kernel_3d_vec import kernel_3d_vec +from struphy.feec.kronecker_mass import kronecker_apply_operator from struphy.feec.linear_operators import BoundaryOperator from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.feec.utilities import LocalProjectionMatrix, LocalRotationMatrix, get_quad_grids @@ -1703,6 +1704,8 @@ def __init__( blocks, self._weights = self._init_blocks_from_weights(weights_info) self._mat = self._wrap_blocks(blocks) + # the operator dot() applies self._mat with (Kronecker blocks replaced), built on first use, see _apply_mat + self._apply_mat_cache = None # the weights are stored in the block order of the (possibly transposed) operator if transposed: @@ -2121,6 +2124,32 @@ def M0(self): def matrix(self): return self._mat + @property + def _apply_mat(self): + """The operator :meth:`dot` applies: :attr:`matrix`, with its blocks that are Kronecker products of 1D banded + matrices (constant weights, e.g. on a Cuboid) applied as such (:mod:`struphy.feec.kronecker_mass`). + + Built on first use after every change of the matrix (assembly, copy into, in-place arithmetic). + """ + if self._apply_mat_cache is None: + if getattr(self, "_kronecker_rejected", False): + self._apply_mat_cache = self._mat + else: + self._apply_mat_cache = kronecker_apply_operator(self._mat, name=self.name) + # a matrix changed by in-place arithmetic that is not a Kronecker product (e.g. the Schur complement + # M1 - dt^2 BC with particle data, rebuilt every step) is not checked again after further arithmetic: + # the check reads the whole matrix. Assembly and copies check again. + if getattr(self, "_changed_by_arithmetic", False) and self._apply_mat_cache is self._mat: + self._kronecker_rejected = True + return self._apply_mat_cache + + def _matrix_changed(self, by_arithmetic: bool): + """Invalidate :attr:`_apply_mat` after a change of the matrix (in-place arithmetic, or assembly/copy).""" + self._apply_mat_cache = None + self._changed_by_arithmetic = by_arithmetic + if not by_arithmetic: + self._kronecker_rejected = False + @property def nquads(self): if self._nquads is None: @@ -2168,29 +2197,54 @@ def dot(self, v, out=None, apply_bc=True): if apply_bc: if self._transposed: - self._W_boundary_op_T.dot(v, out=self._temp_WB) - self._W_extraction_op_T.dot(self._temp_WB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) - self._V_extraction_op.dot(self._temp_VE, out=self._temp_VB) - out = self._V_boundary_op.dot(self._temp_VB, out=out) + ops_in = (self._W_boundary_op_T, self._W_extraction_op_T) + ops_out = (self._V_extraction_op, self._V_boundary_op) + tmps_in = (self._temp_WB, self._temp_mat) + tmps_out = (self._temp_VE, self._temp_VB) else: - self._V_boundary_op_T.dot(v, out=self._temp_VB) - self._V_extraction_op_T.dot(self._temp_VB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) - self._W_extraction_op.dot(self._temp_WE, out=self._temp_WB) - out = self._W_boundary_op.dot(self._temp_WB, out=out) + ops_in = (self._V_boundary_op_T, self._V_extraction_op_T) + ops_out = (self._W_extraction_op, self._W_boundary_op) + tmps_in = (self._temp_VB, self._temp_mat) + tmps_out = (self._temp_WE, self._temp_WB) + out = self._apply_chain(v, ops_in, tmps_in, ops_out, tmps_out, out) else: if self._transposed: self._W_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) + self._apply_mat.dot(self._temp_mat, out=self._temp_VE) out = self._V_extraction_op.dot(self._temp_VE, out=out) else: self._V_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) + self._apply_mat.dot(self._temp_mat, out=self._temp_WE) out = self._W_extraction_op.dot(self._temp_WE, out=out) return out + def _apply_chain(self, v, ops_in, tmps_in, ops_out, tmps_out, out): + """``out = B_out E_out M E_in B_in v`` (boundary and extraction operators around the matrix), see :meth:`dot`. + + Identity operators (e.g. the extraction operators without polar splines, or boundary operators without + essential boundary conditions) are skipped instead of applied as copies: each vector is passed on as it is, + and the matrix writes into the next temporary that is actually used (or into ``out``). + """ + if out is v: + # in place: the chain must not read v after writing the output, so it reads a copy + v = v.copy() + + x = v + for op, tmp in zip(ops_in, tmps_in): + if not isinstance(op, IdentityOperator): + x = op.dot(x, out=tmp) + + # the first output operator that is not an identity reads the matrix product from its temporary + remaining = [(op, tmp) for op, tmp in zip(ops_out, tmps_out) if not isinstance(op, IdentityOperator)] + if not remaining: + return self._apply_mat.dot(x, out=out) + y = self._apply_mat.dot(x, out=remaining[0][1]) + for k, (op, _) in enumerate(remaining): + target = out if k == len(remaining) - 1 else remaining[k + 1][1] + y = op.dot(y, out=target) + return y + def transpose(self, conjugate=False): """ Returns the transposed operator. @@ -2277,6 +2331,7 @@ def assemble(self, weights=None, clear=True): Whether to first set all data to zero before assembly. If False, the new contributions are added to existing ones. """ + self._matrix_changed(by_arithmetic=False) # the matrix changes assert not self._dry_run, ( "A dry-run operator has no matrix data and cannot be assembled (memory estimation only)." ) @@ -2559,6 +2614,7 @@ def copy(self, out=None): out._weights = [list(row) for row in self._weights] self._mat.copy(out=out._mat) + out._matrix_changed(by_arithmetic=False) if self._creation_info is None: out._creation_info = None @@ -2569,6 +2625,7 @@ def copy(self, out=None): def __imul__(self, a): self._mat *= a + self._matrix_changed(by_arithmetic=True) self._creation_info = None return self @@ -2579,11 +2636,13 @@ def __iadd__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat += M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat += M + self._matrix_changed(by_arithmetic=True) return self else: @@ -2596,11 +2655,13 @@ def __isub__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat -= M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat -= M + self._matrix_changed(by_arithmetic=True) return self else: From df53602f02f6adb33b3489b5d8aeb13458cd5fe7 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:20:50 +0200 Subject: [PATCH 32/55] Kronecker mass: launch the band-pass kernel through CachedLaunch The banded passes run with the same buffers in every solver iteration (as on cuda-aware-mpi-fixes, now that feectools has CachedLaunch). Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/kronecker_mass.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py index 2050edb82..5456b2eeb 100644 --- a/src/struphy/feec/kronecker_mass.py +++ b/src/struphy/feec/kronecker_mass.py @@ -38,8 +38,10 @@ def kronecker_mass_enabled() -> bool: def _band_pass_cuda(): """The CUDA kernel of one banded pass (compiled on first call).""" from cunumpy.kernels import CudaKernel + from feectools.utilities.cached_launch import CachedLaunch - return CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass") + # launched with the same buffers in every solver iteration: cache the launch preparation + return CachedLaunch(CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass")) def _band_pass(A, inp, out, axis, offset): From 20b77a56c939d11274fd490ec8cea0beafe03a82 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 13:59:36 +0200 Subject: [PATCH 33/55] Allow fused solver inner products in the GPU transfer budget --- src/struphy/models/tests/gpu_e2e.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/src/struphy/models/tests/gpu_e2e.py b/src/struphy/models/tests/gpu_e2e.py index ecfa0c9a2..46dbe2765 100644 --- a/src/struphy/models/tests/gpu_e2e.py +++ b/src/struphy/models/tests/gpu_e2e.py @@ -66,13 +66,18 @@ def device_context(): yield +# most inner products a feectools solver iteration copies to the host at once (pipelined PCG: 3) +MAX_FUSED_REDUCTIONS = 3 + + @dataclass class TransferBudget: """Transfers recorded per phase of the time loop (the setup before it is not counted). - ``integrate`` (the propagators) may only copy scalars to the host: the convergence test of the Krylov solvers - reads the 8-byte residual norm once per iteration (feectools keeps everything else of the solve on the device). - No arrays, no host-to-device copies, no host kernels. ``diagnostics`` (scalars, saved markers, binned + ``integrate`` (the propagators) may only copy scalars to the host: the Krylov solvers copy the inner products + of one iteration (residual norm and step coefficients, at most ``MAX_FUSED_REDUCTIONS`` of them) to the host + together, once per iteration (feectools keeps everything else of the solve on the device). No other arrays, no + host-to-device copies, no host kernels. ``diagnostics`` (scalars, saved markers, binned distribution functions) may copy small results to the host: at most one scalar per tracked scalar and call, never a host-to-device copy or a host kernel. ``output`` (``DataContainer.save_data``) may copy each saved dataset to the host once, nothing more. @@ -99,10 +104,10 @@ def check(self, n_scalars: int): report = self.report() assert self.n_integrate_calls == self.n_steps, report - # 1. the propagators: only scalars to the host (solver convergence tests), nothing else + # 1. the propagators: only (fused) scalars to the host (solver inner products), nothing else for e in self.integrate.events: if e.kind == "to_host": - assert e.nbytes is not None and e.nbytes <= 8, report + assert e.nbytes is not None and e.nbytes <= 8 * MAX_FUSED_REDUCTIONS, report else: assert e.kind not in HOST_DEVICE_KINDS, report From 9e6eeec20284c765afe6852284d359f00ac8115e Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 14:56:25 +0200 Subject: [PATCH 34/55] polar MPI decomposition parity requires more than one rank --- src/struphy/polar/tests/test_polar_cupy.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/struphy/polar/tests/test_polar_cupy.py b/src/struphy/polar/tests/test_polar_cupy.py index 2e7dc2c5f..6f23408a6 100644 --- a/src/struphy/polar/tests/test_polar_cupy.py +++ b/src/struphy/polar/tests/test_polar_cupy.py @@ -342,9 +342,13 @@ def check_polar_mpi(cuda_aware, mpi_dims_mask, fake=True): """Backend parity with both Kronecker redistributions and ranks away from the pole.""" from contextlib import nullcontext + from maybempi import MPI from struphy.pic.tests.cuda_emulation import emulated_launches from struphy.polar.linear_operators import set_device_sparse_module + if MPI.COMM_WORLD.Get_size() == 1: + pytest.skip("polar MPI decomposition parity requires more than one rank") + previous_mpi = cunumpy.mpi.get_mpi_cuda_aware() previous_sparse = set_device_sparse_module(DenseDeviceSparse) if fake else None try: From 3f090ceb8b21bee6c1a1300e8052247089a31115 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 15:30:29 +0200 Subject: [PATCH 35/55] Fix fake CuPy MPI and solver paths Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- feectools | 2 +- src/struphy/polar/tests/test_polar_cupy.py | 4 ---- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/feectools b/feectools index e75b4c766..717edc98b 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit e75b4c76616ef93f0091d17e752e4155f430c5d2 +Subproject commit 717edc98b09464358b9865dfbe91cea0d00ca5a5 diff --git a/src/struphy/polar/tests/test_polar_cupy.py b/src/struphy/polar/tests/test_polar_cupy.py index 6f23408a6..2e7dc2c5f 100644 --- a/src/struphy/polar/tests/test_polar_cupy.py +++ b/src/struphy/polar/tests/test_polar_cupy.py @@ -342,13 +342,9 @@ def check_polar_mpi(cuda_aware, mpi_dims_mask, fake=True): """Backend parity with both Kronecker redistributions and ranks away from the pole.""" from contextlib import nullcontext - from maybempi import MPI from struphy.pic.tests.cuda_emulation import emulated_launches from struphy.polar.linear_operators import set_device_sparse_module - if MPI.COMM_WORLD.Get_size() == 1: - pytest.skip("polar MPI decomposition parity requires more than one rank") - previous_mpi = cunumpy.mpi.get_mpi_cuda_aware() previous_sparse = set_device_sparse_module(DenseDeviceSparse) if fake else None try: From 29e918008adb0095bbe6dda82eb25490825010f2 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 15:40:11 +0200 Subject: [PATCH 36/55] Account for MPI scalar diagnostic transfers Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/struphy/models/tests/gpu_e2e.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/struphy/models/tests/gpu_e2e.py b/src/struphy/models/tests/gpu_e2e.py index 46dbe2765..d58e20af4 100644 --- a/src/struphy/models/tests/gpu_e2e.py +++ b/src/struphy/models/tests/gpu_e2e.py @@ -114,8 +114,11 @@ def check(self, n_scalars: int): # 2. diagnostics: no host-to-device copies, no host kernels, only scalar-sized downloads for kind in ("to_device", "kernel_conversion", "fallback"): assert self.diagnostics.count(kind) == 0, report - assert self.diagnostics.to_host <= n_scalars * self.n_diagnostics_calls, report - assert self.diagnostics.bytes_to_host <= 8 * n_scalars * self.n_diagnostics_calls, report + # MPI scalar reductions copy one value per scalar, followed by the consolidated scalar buffer. + max_diagnostic_copies = (n_scalars + 1) * self.n_diagnostics_calls + assert self.diagnostics.to_host <= max_diagnostic_copies, report + max_diagnostic_bytes = 16 * n_scalars * self.n_diagnostics_calls + assert self.diagnostics.bytes_to_host <= max_diagnostic_bytes, report # 3. output: each saved dataset once to the host for kind in ("to_device", "kernel_conversion", "fallback"): From 080cead991ddb4d6d1cd225f564feedf1e1439b2 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 16:20:02 +0200 Subject: [PATCH 37/55] Fix fake CuPy block operations and serial inner products --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index 717edc98b..b3de2933b 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 717edc98b09464358b9865dfbe91cea0d00ca5a5 +Subproject commit b3de2933b5071a6762aa0432838c39f201cdc73a From f63a56031f63bcc4d3a3065cc4c42205b12d30c5 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:00:05 +0200 Subject: [PATCH 38/55] CuPy backend: host/device conversions in post-processing; math.sqrt for host scalars in Species Post-processing evaluates fields on the device and writes them from the host: the evaluation grids, point data and coordinates are converted where they cross, and HDF5 data is moved to the device before it is written into the spline coefficients. Co-Authored-By: Claude Opus 5.5 --- src/struphy/models/species.py | 4 ++-- src/struphy/post_processing/output.py | 31 ++++++++++++++++++--------- 2 files changed, 23 insertions(+), 12 deletions(-) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index b43464713..0ad27067d 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,9 +1,9 @@ import logging +import math import warnings from abc import ABCMeta, abstractmethod from dataclasses import fields, is_dataclass -import cunumpy as xp from maybempi import MPI from struphy.models.variables import Variable @@ -220,7 +220,7 @@ def __init__( con = ConstantsOfNature() # relevant frequencies - om_p = xp.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) + om_p = math.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) om_c = Z * con.e * units.B / (A * con.mH) # compute equation parameters diff --git a/src/struphy/post_processing/output.py b/src/struphy/post_processing/output.py index bfcf5fd58..cc65d56f8 100644 --- a/src/struphy/post_processing/output.py +++ b/src/struphy/post_processing/output.py @@ -85,6 +85,13 @@ def load_plasma_plots() -> bool: return True +def _to_host(value): + """`value` (an array, or a list/tuple of arrays for vector-valued fields) as NumPy array(s).""" + if isinstance(value, (list, tuple)): + return type(value)(_to_host(v) for v in value) + return xp.to_numpy(value) + + class ProductMapping(Mapping[str, xr.DataArray]): """A discoverable mapping whose products are loaded on first access.""" @@ -568,13 +575,14 @@ def _attach_physical_coords( self, array: xr.DataArray, etas: tuple[Any, Any, Any], dims: tuple[str, ...] ) -> xr.DataArray: """Map a raw FEEC evaluation grid and attach its Cartesian coordinates.""" - grids = [np.atleast_1d(np.asarray(eta, dtype=float)) for eta in etas] - mapped = self.domain(*grids) + # the logical grids may be device arrays (CuPy backend): map them where they live, keep the coordinates on the host + grids = [np.atleast_1d(np.asarray(xp.to_numpy(eta), dtype=float)) for eta in etas] + mapped = self.domain(*(xp.asarray(grid) for grid in grids)) shape = tuple(len(grid) for grid in grids) keep = tuple(slice(None) if dim in dims else 0 for dim in ("eta1", "eta2", "eta3")) coordinates = {} for name, values in zip(("X", "Y", "Z"), mapped): - values = np.asarray(values).reshape(shape)[keep] + values = np.asarray(xp.to_numpy(values)).reshape(shape)[keep] coordinates[name] = (dims, values) if dims else values.item() return array.assign_coords(coordinates) @@ -1295,7 +1303,9 @@ def _process_fields( # the physical grid is only needed for output, hence it is only built on rank 0 if self._pproc_rank == 0: - grids_phy = list(self.domain(*grids_log)) + # output coordinates, on the host (the logical grids stay on the device for the field evaluation) + grids_phy = [xp.to_numpy(grid) for grid in self.domain(*grids_log)] + grids_log_host = [xp.to_numpy(grid) for grid in grids_log] else: grids_phy = None @@ -1327,8 +1337,9 @@ def _process_fields( if self._pproc_rank == 0: for species, vars in vals.items(): for name, val in vars.items(): - point_data[species][name][t] = val - point_data_phy[species][name][t] = vals_phy[species][name] + # evaluated on the device on the CuPy backend; written to files from the host + point_data[species][name][t] = _to_host(val) + point_data_phy[species][name][t] = _to_host(vals_phy[species][name]) # directory for the vtk files path_fields = os.path.join(self.path_pproc, "fields_data") @@ -1338,10 +1349,10 @@ def _process_fields( for species, vars in point_data.items(): variables = {} for name, val in vars.items(): - variables[name] = wrap_field_data(val, grids_log, grids_phy=grids_phy, name=name) + variables[name] = wrap_field_data(val, grids_log_host, grids_phy=grids_phy, name=name) if physical: variables[name + "_xyz"] = wrap_field_data( - point_data_phy[species][name], grids_log, grids_phy=grids_phy, name=name + "_xyz" + point_data_phy[species][name], grids_log_host, grids_phy=grids_phy, name=name + "_xyz" ) store.write_group(store.store_path(self.path_pproc), f"/{species}", xr.Dataset(variables)) @@ -1517,7 +1528,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[n * step, p1:-p1, p2:-p2, p3:-p3]) # HDF5 data is on the host # vector-valued field else: @@ -1530,7 +1541,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3]) vector.update_ghost_regions() if is_polar: From 1a3939ef1d19f549b1a8a1e43e5e1390f1bb6695 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:07:26 +0200 Subject: [PATCH 39/55] CUDA version of the mass-matrix assembly kernel kernel_3d_mat - kernel_3d_mat_cuda.cu: one thread per (element, codomain basis function, domain basis function); replaces the host fallback on the CuPy backend. - WeightedMassOperator.assemble reads the quadrature grids (spans, weights, basis values) as backend arrays, copied to the device once and cached. - Simulation.allocate compiles the assembly kernel before it is first called, so the compilation is not counted in the assembly's profiling regions. Co-Authored-By: Claude Opus 5.5 --- .../feec/kernels/kernel_3d_mat/__init__.py | 18 ++-- .../kernel_3d_mat/kernel_3d_mat_cuda.cu | 92 +++++++++++++++++++ src/struphy/feec/mass.py | 21 ++++- src/struphy/simulation/sim.py | 18 ++++ 4 files changed, 139 insertions(+), 10 deletions(-) create mode 100644 src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu diff --git a/src/struphy/feec/kernels/kernel_3d_mat/__init__.py b/src/struphy/feec/kernels/kernel_3d_mat/__init__.py index a881301b4..76da51673 100644 --- a/src/struphy/feec/kernels/kernel_3d_mat/__init__.py +++ b/src/struphy/feec/kernels/kernel_3d_mat/__init__.py @@ -1,8 +1,4 @@ -"""The kernel of this folder: pyccel in ``_kernels.py``, CUDA in ``_cuda.cu`` once ported. - -No CUDA version yet: on the CuPy backend the pyccel kernel runs on host copies of the arrays (a warning is -emitted once), as before the move into this folder. -""" +"""The kernel of this folder: pyccel in ``kernel_3d_mat_kernels.py``, CUDA in ``kernel_3d_mat_cuda.cu``.""" from cunumpy.kernels import Kernel @@ -11,6 +7,16 @@ # the arguments the kernel writes to, by position: data OUTPUTS = (25,) + +def _n_threads(args): + """One thread per (element, codomain basis function, domain basis function), see kernel_3d_mat_cuda.cu.""" + n_elements = args[0].size * args[1].size * args[2].size + n_pairs = 1 + for degree in args[3:9]: # pi1, pi2, pi3, pj1, pj2, pj3 + n_pairs *= int(degree) + 1 + return n_elements * n_pairs + + kernel_3d_mat = Kernel.from_folder( - __name__, structs=CUDA_STRUCTS, missing_cuda="fallback", host_options={"outputs": OUTPUTS} + __name__, structs=CUDA_STRUCTS, host_options={"outputs": OUTPUTS}, n_threads_from=_n_threads ) diff --git a/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu b/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu new file mode 100644 index 000000000..6afab9398 --- /dev/null +++ b/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu @@ -0,0 +1,92 @@ +// CUDA version of kernel_3d_mat (kernel_3d_mat_kernels.py), with the same arguments in the same order. +#include "cunumpy/array_view.cuh" +#include "cunumpy/atomic.cuh" +#include "cunumpy/index.cuh" + +/** + * Integral of Lambda_i * mat_fun * Lambda_j over the local elements, added to the stencil matrix data, as in + * kernel_3d_mat_kernels.kernel_3d_mat. + * + * One thread per (element, codomain basis function (il1, il2, il3), domain basis function (jl1, jl2, jl3)): + * n_threads = ne1 * ne2 * ne3 * (pi1 + 1) * (pi2 + 1) * (pi3 + 1) * (pj1 + 1) * (pj2 + 1) * (pj3 + 1) (see + * __init__.py), the element index varying slowest and jl3 fastest. Each thread computes the quadrature sum of its + * pair exactly as the pyccel kernel (same factors, same order) and adds it to its entry of `data` with an atomic + * addition: neighbouring elements contribute to the same entries, so only the summation order over elements + * differs from pyccel (round-off). + * + * @param spans1, spans2, spans3 Knot span index of each local element, per direction. + * @param pi1, pi2, pi3 Degrees of the codomain basis functions. + * @param pj1, pj2, pj3 Degrees of the domain basis functions. + * @param starts1, starts2, starts3 Global start index of the rows on this process. + * @param pads1, pads2, pads3 Paddings (ghost regions) of `data`. + * @param w1, w2, w3 Quadrature weights [element, quadrature point]. + * @param bi1, bi2, bi3 Codomain basis functions [element, local basis function, derivative, quadrature point]. + * @param bj1, bj2, bj3 Domain basis functions, same layout. + * @param mat_fun Weight function at the quadrature points [flattened point 1, flattened point 2, flattened point 3]. + * @param data StencilMatrix data (rows with ghost regions..., diagonals...); the results are added (not set). + */ +extern "C" __global__ void kernel_3d_mat(Array1D spans1, Array1D spans2, + Array1D spans3, int pi1, int pi2, int pi3, int pj1, int pj2, + int pj3, int starts1, int starts2, int starts3, int pads1, int pads2, + int pads3, Array2D w1, Array2D w2, Array2D w3, + Array4D bi1, Array4D bi2, Array4D bi3, + Array4D bj1, Array4D bj2, Array4D bj3, + Array3D mat_fun, Array6D data) +{ + const long long ne1 = spans1.size(); + const long long ne2 = spans2.size(); + const long long ne3 = spans3.size(); + + const long long ni1 = pi1 + 1, ni2 = pi2 + 1, ni3 = pi3 + 1; + const long long nj1 = pj1 + 1, nj2 = pj2 + 1, nj3 = pj3 + 1; + const long long n_pairs = ni1 * ni2 * ni3 * nj1 * nj2 * nj3; + + CUNUMPY_THREAD_1D(thread, ne1 * ne2 * ne3 * n_pairs); + + // CUDA-only: split the thread index into the pyccel loop variables (iel1, iel2, iel3, il1, il2, il3, jl1, jl2, + // jl3), jl3 fastest + long long rest = thread; + const long long jl3 = rest % nj3; rest /= nj3; + const long long jl2 = rest % nj2; rest /= nj2; + const long long jl1 = rest % nj1; rest /= nj1; + const long long il3 = rest % ni3; rest /= ni3; + const long long il2 = rest % ni2; rest /= ni2; + const long long il1 = rest % ni1; rest /= ni1; + const long long iel3 = rest % ne3; rest /= ne3; + const long long iel2 = rest % ne2; rest /= ne2; + const long long iel1 = rest; + + const long long nq1 = w1.shape[1]; + const long long nq2 = w2.shape[1]; + const long long nq3 = w3.shape[1]; + + double value = 0.0; + + for (long long q1 = 0; q1 < nq1; ++q1) { + for (long long q2 = 0; q2 < nq2; ++q2) { + for (long long q3 = 0; q3 < nq3; ++q3) { + const double wvol = w1(iel1, q1) * w2(iel2, q2) * w3(iel3, q3) * + mat_fun(iel1 * nq1 + q1, iel2 * nq2 + q2, iel3 * nq3 + q3); + + const double bi = bi1(iel1, il1, 0, q1) * bi2(iel2, il2, 0, q2) * bi3(iel3, il3, 0, q3); + const double bj = bj1(iel1, jl1, 0, q1) * bj2(iel2, jl2, 0, q2) * bj3(iel3, jl3, 0, q3); + + value += wvol * bi * bj; + } + } + } + + // global spline indices + const long long i_global1 = spans1(iel1) - pi1 + il1; + const long long i_global2 = spans2(iel2) - pi2 + il2; + const long long i_global3 = spans3(iel3) - pi3 + il3; + + // local spline indices (- starts --> can be negative, will therefore be written to ghost regions) + const long long i_local1 = i_global1 - starts1; + const long long i_local2 = i_global2 - starts2; + const long long i_local3 = i_global3 - starts3; + + cunumpy_atomic_add(&data(pads1 + i_local1, pads2 + i_local2, pads3 + i_local3, pads1 + jl1 - il1, + pads2 + jl2 - il2, pads3 + jl3 - il3), + value); +} diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index f75c6de17..e035e4fe4 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -2242,6 +2242,19 @@ def transpose(self, conjugate=False): return M + def _quad_grid(self, space_name, attribute, component): + """Quadrature-grid arrays (spans, weights or basis values) of one component of a space, as backend arrays. + + The Derham keeps them as host arrays; on the CuPy backend the assembly kernel needs them on the device, + so they are copied there once and cached (they do not change between assemblies). + """ + cache = self.__dict__.setdefault("_quad_grid_cache", {}) + key = (space_name, attribute, component) + if key not in cache: + arrays = getattr(self.derham.spline_attributes[space_name], attribute)[component] + cache[key] = [xp.asarray(array) for array in arrays] + return cache[key] + def assemble(self, weights=None, clear=True): r""" Assembles the weighted mass matrix, i.e. computes the integrals @@ -2320,7 +2333,7 @@ def assemble(self, weights=None, clear=True): # loop over codomain spaces (rows) for a, codomain_space in enumerate(codomain_spaces): # knot span indices of elements of local domain - codomain_spans = spline_attr[W_name].quad_grid_spans[a] + codomain_spans = self._quad_grid(W_name, "quad_grid_spans", a) # global start spline index on process codomain_starts = [int(start) for start in codomain_space.coeff_space.starts] @@ -2332,10 +2345,10 @@ def assemble(self, weights=None, clear=True): pts = self._quad_pts(W_name, a) # global quadrature weights in format (local element, local weight) - wts = spline_attr[W_name].quad_grid_wts[a] + wts = self._quad_grid(W_name, "quad_grid_wts", a) # evaluated basis functions at quadrature points of codomain space - codomain_basis = spline_attr[W_name].quad_grid_bases[a] + codomain_basis = self._quad_grid(W_name, "quad_grid_bases", a) # loop over domain spaces (columns) for b, domain_space in enumerate(domain_spaces): @@ -2375,7 +2388,7 @@ def assemble(self, weights=None, clear=True): ) # evaluated basis functions at quadrature points of domain space - domain_basis = spline_attr[V_name].quad_grid_bases[b] + domain_basis = self._quad_grid(V_name, "quad_grid_bases", b) # assemble matrix (if mat_w is not zero) by calling the appropriate kernel (1d, 2d or 3d) if not_weight_zero or self._is_scalar: diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 9452e81c5..e6f358832 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -243,6 +243,8 @@ def allocate(self): logger.debug("\nAllocating simulation data ...") + self._compile_setup_cuda_kernels() + with ProfileManager.profile_region("setup: allocate"): # feec with ProfileManager.profile_region("allocate feec", functions=[self._allocate_feec]): @@ -280,6 +282,22 @@ def kernels(self) -> tuple[Kernel, ...]: owners += self.model.prop_list return collect_kernels(*owners) + def _compile_setup_cuda_kernels(self) -> None: + """Compile the CUDA kernels that :meth:`allocate` itself calls (mass-matrix assembly) before it calls them. + + :meth:`compile_cuda_kernels` runs after :meth:`allocate`, so without this the first assembly would include + the compilation (~1 s per kernel on a node whose CuPy cache does not have it yet), in the assembly's + profiling regions. Here it is counted in ``"setup: compile cuda kernels (assembly)"``. Nothing happens on + the NumPy backend. + """ + if xp.get_backend() != "cupy": + return + + from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat + + with ProfileManager.profile_region("setup: compile cuda kernels (assembly)"): + kernel_3d_mat.compile() + def compile_cuda_kernels(self) -> tuple[Kernel, ...]: """Compile the CUDA versions of the kernels of :meth:`kernels` now, before the time stepping. From 361b541175a0eee72f5a6f8145c18992d804fc05 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:26:26 +0200 Subject: [PATCH 40/55] Parity cases for kernel_3d_mat Random quadrature data on 3 x 2 x 2 elements with mixed degrees, with the first row local or a ghost row. test_cuda_kernels_have_parity_cases requires a case for every kernel with a CUDA version; the emulation test compares it with the pyccel kernel without a GPU. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/cuda_parity_cases.py | 26 ++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index a2396ee5b..6e10707be 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -315,6 +315,30 @@ def kernel_pullpush_args(case): ) +def kernel_3d_mat_threads(args): + """One thread per (element, codomain basis function, domain basis function), as the kernel's launch size.""" + from struphy.feec.kernels.kernel_3d_mat import _n_threads + + return _n_threads(args) + + +def kernel_3d_mat_args(case): + """Mass-matrix assembly on 3 x 2 x 2 elements with mixed degrees and random quadrature data; `case` is the start + index of the rows in direction 1 (1: the first row is a ghost row, local index -1).""" + rng = np.random.default_rng(20 + case) + ne, nq = (3, 2, 2), (3, 2, 2) + pi, pj, pads = (2, 1, 1), (1, 2, 1), (2, 2, 2) + starts = (case, 0, 0) + spans = [xp.asarray(np.arange(n, dtype=np.int64) + p) for n, p in zip(ne, pi)] + w = [xp.asarray(rng.random((n, q))) for n, q in zip(ne, nq)] + bi = [xp.asarray(rng.random((n, p + 1, 1, q))) for n, p, q in zip(ne, pi, nq)] + bj = [xp.asarray(rng.random((n, p + 1, 1, q))) for n, p, q in zip(ne, pj, nq)] + mat_fun = xp.asarray(rng.normal(size=tuple(n * q for n, q in zip(ne, nq)))) + rows = tuple(n + p + 2 * pad for n, p, pad in zip(ne, pi, pads)) + data = xp.asarray(rng.normal(size=rows + tuple(2 * pad + 1 for pad in pads))) # added to, not overwritten + return (*spans, *pi, *pj, *starts, *pads, *w, *bi, *bj, mat_fun, data) + + # ---------------------------------------------------------------- all kernels with a CUDA version PUSHER_TOLERANCES = {"rtol": 1e-13, "atol": 1e-14} @@ -338,6 +362,8 @@ def size_of(index, per_thread=1): PARITY_CASES = { + # atomic additions: the summation order over elements differs from pyccel + "kernel_3d_mat": ParityCases((0, 1), kernel_3d_mat_args, rtol=1e-12, atol=1e-13, n_threads=kernel_3d_mat_threads), "push_eta_stage": ParityCases( [ (bc, method, stage) From 62afbad5fe32a41827fd129dfd562f755d828871 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:08:01 +0200 Subject: [PATCH 41/55] Apply separable mass matrices as Kronecker products of 1D banded matrices WeightedMassOperator.dot applies the blocks of its matrix that are Kronecker products of 1D banded matrices (constant weights, e.g. on a Cuboid) as three banded passes (struphy.feec.kronecker_mass, CUDA kernel kron_band_pass on the GPU) instead of the 3D stencil product. The replacement is built on first use and invalidated by assembly, copies and in-place arithmetic; a matrix changed by arithmetic that is not a Kronecker product is not checked again. STRUPHY_KRONECKER_MASS=0 switches it off. The boundary and extraction operators around the matrix are skipped when they are identities instead of being applied as copies. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/kron_band_pass_cuda.cu | 35 ++++ src/struphy/feec/kronecker_mass.py | 247 ++++++++++++++++++++++++ src/struphy/feec/mass.py | 85 ++++++-- 3 files changed, 355 insertions(+), 12 deletions(-) create mode 100644 src/struphy/feec/kron_band_pass_cuda.cu create mode 100644 src/struphy/feec/kronecker_mass.py diff --git a/src/struphy/feec/kron_band_pass_cuda.cu b/src/struphy/feec/kron_band_pass_cuda.cu new file mode 100644 index 000000000..745e782f4 --- /dev/null +++ b/src/struphy/feec/kron_band_pass_cuda.cu @@ -0,0 +1,35 @@ +/** + * One pass of a Kronecker product of 1D banded matrices (struphy.feec.kronecker_mass._band_pass): + * + * out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...] along `axis`, + * + * with the other two axes unchanged (out and inp have the same extent along them). One thread per entry of `out`, + * the last axis fastest, so that neighbouring threads read neighbouring entries of `inp` along every axis. `out` may + * be a strided view (the owned rows of a stencil vector's data). + * + * @param A Banded 1D matrix, shape (out.shape[axis], n_diagonals). + * @param inp Input array (3D). + * @param out Output array (3D, any strides); written, not added to. + * @param axis Axis of the pass (0, 1 or 2). + * @param offset Index in `inp` of diagonal 0 of row 0 along `axis`. + */ +#include "cunumpy/array_view.cuh" + +extern "C" __global__ void kron_band_pass(Array2D A, Array3D inp, Array3D out, int axis, + long long offset) { + const long long n1 = out.shape[1], n2 = out.shape[2]; + const long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; + if (t >= out.shape[0] * n1 * n2) return; + const long long i2 = t % n2; + const long long i1 = (t / n2) % n1; + const long long i0 = t / (n1 * n2); + const long long i = axis == 0 ? i0 : (axis == 1 ? i1 : i2); + + double v = 0.0; + for (long long d = 0; d < A.shape[1]; ++d) { + const long long j = i + offset + d; + const double x = axis == 0 ? inp(j, i1, i2) : (axis == 1 ? inp(i0, j, i2) : inp(i0, i1, j)); + v += A(i, d) * x; + } + out(i0, i1, i2) = v; +} diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py new file mode 100644 index 000000000..2050edb82 --- /dev/null +++ b/src/struphy/feec/kronecker_mass.py @@ -0,0 +1,247 @@ +"""Fast application of assembled mass matrices that are Kronecker products of 1D banded matrices. + +A mass matrix block ``M[i, j] = int w Lambda_i Lambda_j`` of tensor-product splines factorizes as +``M = c * M_1 (x) M_2 (x) M_3`` when its weight ``w`` is constant (e.g. on a Cuboid). Applied as a +:class:`~feectools.linalg.stencil.StencilMatrix`, every row reads all ``n_1 * n_2 * n_3`` diagonals (105 for +degrees (1, 2, 3)) from memory; as a Kronecker product it is three 1D banded passes over the vector, ``n_1 + n_2 + +n_3`` operations per entry and no 3D matrix data. :func:`kronecker_apply_operator` checks each assembled block and +returns an operator for the matrix-vector products in which the blocks that are Kronecker products (to round-off) +are replaced by :class:`KroneckerBandMatrix`. + +The 1D factors are read from the assembled data itself (rows owned by this process), so they hold exactly what the +assembly computed, on any domain decomposition, and the check is the assembled data, not an assumption about the +mapping or the weights. +""" + +import logging +import os +from functools import cache +from pathlib import Path + +import cunumpy as xp +from feectools.linalg.basic import LinearOperator +from feectools.linalg.block import BlockLinearOperator +from feectools.linalg.stencil import StencilMatrix, StencilVector + +logger = logging.getLogger("struphy") + +# relative tolerance of the Kronecker check: the 3D assembly and the product of the factors differ by round-off only +KRONECKER_RTOL = 1e-12 + + +def kronecker_mass_enabled() -> bool: + """Whether :func:`kronecker_apply_operator` may replace blocks (``STRUPHY_KRONECKER_MASS=0`` switches it off).""" + return os.environ.get("STRUPHY_KRONECKER_MASS", "1").lower() not in ("0", "false", "no") + + +@cache +def _band_pass_cuda(): + """The CUDA kernel of one banded pass (compiled on first call).""" + from cunumpy.kernels import CudaKernel + + return CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass") + + +def _band_pass(A, inp, out, axis, offset): + """``out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...]`` along ``axis`` (other axes unchanged).""" + if xp.is_gpu(out): + _band_pass_cuda()(A, inp, out, axis, offset, n_threads=out.size) + return + m = out.shape[axis] + shape = [1, 1, 1] + shape[axis] = m + out[...] = 0.0 + for d in range(A.shape[1]): + idx = [slice(None)] * 3 + idx[axis] = slice(offset + d, offset + d + m) + out += A[:, d].reshape(shape) * inp[tuple(idx)] + + +class KroneckerBandMatrix(LinearOperator): + """A square 3D :class:`~feectools.linalg.stencil.StencilMatrix` applied as a Kronecker product of 1D banded + matrices, ``(A_1 (x) A_2 (x) A_3) x`` on the rows owned by this process. + + ``A_k[i, d]`` is the entry of owned row ``i`` (along direction ``k``) and diagonal ``d``, column + ``i - q_k + d`` with ``q_k = (n_k - 1) // 2``, as in the stencil matrix. The product is computed in three passes, + each over the owned rows of its direction and the full (padded) extent of the directions after it, so the passes + need no communication beyond the ghost regions of the input. Use :meth:`from_stencil`. + + Parameters + ---------- + stencil : StencilMatrix + The assembled matrix (kept for ``toarray``, ``tosparse`` and ``transpose``). + factors : tuple of arrays + ``A_1, A_2, A_3``, shapes ``(m_k, n_k)``, C-contiguous backend arrays. + """ + + def __init__(self, stencil: StencilMatrix, factors): + self._stencil = stencil + self._factors = factors + + V = stencil.domain + self._m = tuple(int(e - s + 1) for s, e in zip(V.starts, V.ends)) + self._pads = tuple(int(p) for p in V.pads) + self._q = tuple((int(n) - 1) // 2 for n in stencil._data.shape[3:]) + + # temporaries of the first two passes: owned rows along the passes done, padded extent along the others + x_shape = V.zeros()._data.shape + self._t1 = xp.empty((self._m[0], x_shape[1], x_shape[2]), dtype=stencil._data.dtype) + self._t2 = xp.empty((self._m[0], self._m[1], x_shape[2]), dtype=stencil._data.dtype) + + @classmethod + def from_stencil(cls, S: StencilMatrix, rtol: float = KRONECKER_RTOL): + """The Kronecker form of ``S`` if it is one (to ``rtol``, relative to its largest entry), else None. + + Only square 3D matrices on one space (domain is codomain) with shifts 1 qualify, e.g. the diagonal blocks of + the mass matrices. The factors are read from the owned rows of ``S``: with reference rows ``r_k`` (the first + owned row along each direction) and central diagonals ``q_k``, + ``A_1[i, :] = S[i, r_2, r_3, :, q_2, q_3]``, ``A_2[i, :] = S[r_1, i, r_3, q_1, :, q_3] / S[r, q]``, + ``A_3[i, :] = S[r_1, r_2, i, q_1, q_2, :] / S[r, q]``; the product is then compared with all owned rows. + """ + if not isinstance(S, StencilMatrix) or S.domain is not S.codomain or S.domain.ndim != 3: + return None + V = S.domain + if any(int(s) != 1 for s in V.shifts) or S._data.ndim != 6: + return None + + p = [int(p) for p in V.pads] + m = [int(e - s + 1) for s, e in zip(V.starts, V.ends)] + q = [(int(n) - 1) // 2 for n in S._data.shape[3:]] + if min(m) < 1: + return None + data = S._data + + den = data[p[0], p[1], p[2], q[0], q[1], q[2]] + if float(den) == 0.0: + return None + A1 = xp.ascontiguousarray(data[p[0] : p[0] + m[0], p[1], p[2], :, q[1], q[2]]) + A2 = xp.ascontiguousarray(data[p[0], p[1] : p[1] + m[1], p[2], q[0], :, q[2]] / den) + A3 = xp.ascontiguousarray(data[p[0], p[1], p[2] : p[2] + m[2], q[0], q[1], :] / den) + + # compare with every owned row: planes (rows along direction 1) in chunks, the maxima kept on the device + # (one read at the end); a few planes first, which rejects most matrices that are not Kronecker products + # (e.g. ones assembled from particle data, checked again after every change) at once + owned = data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + A23 = A2[:, None, None, :, None] * A3[None, :, None, None, :] # (m2, m3, 1, n2, n3) + + def max_diff_and_abs(rows): + block = owned[rows] + recon = A1[rows][:, None, None, :, None, None] * A23[None] + return xp.max(xp.abs(block - recon)), xp.max(xp.abs(block)) + + def fails(max_diff, max_abs): + max_diff, max_abs = float(max_diff), float(max_abs) + return max_abs == 0.0 or max_diff > rtol * max_abs + + quick = sorted({0, m[0] // 2, m[0] - 1}) + if fails(*max_diff_and_abs(xp.asarray(quick))): + return None + + # chunks of at most ~64 MB of temporaries + plane = max(1, owned[0].size) + chunk = max(1, min(m[0], (8 << 20) // plane)) + max_diff = xp.zeros(()) + max_abs = xp.zeros(()) + for start in range(0, m[0], chunk): + d, a = max_diff_and_abs(slice(start, min(start + chunk, m[0]))) + max_diff = xp.maximum(max_diff, d) + max_abs = xp.maximum(max_abs, a) + if fails(max_diff, max_abs): + return None + return cls(S, (A1, A2, A3)) + + # ------------------------------------------------------------------ + # LinearOperator interface + # ------------------------------------------------------------------ + @property + def domain(self): + return self._stencil.domain + + @property + def codomain(self): + return self._stencil.codomain + + @property + def dtype(self): + return self._stencil.dtype + + @property + def factors(self): + """The 1D factors ``A_1, A_2, A_3`` (rows owned by this process).""" + return self._factors + + def tosparse(self): + return self._stencil.tosparse() + + def toarray(self): + return self._stencil.toarray() + + def transpose(self, conjugate=False): + return self._stencil.transpose(conjugate=conjugate) + + def dot(self, v, out=None): + assert isinstance(v, StencilVector) + assert v.space is self.domain + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.codomain + + if not v.ghost_regions_in_sync: + v.update_ghost_regions() + + A1, A2, A3 = self._factors + p, q, m = self._pads, self._q, self._m + owned = out._data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + + _band_pass(A1, v._data, self._t1, 0, p[0] - q[0]) + _band_pass(A2, self._t1, self._t2, 1, p[1] - q[1]) + _band_pass(A3, self._t2, owned, 2, p[2] - q[2]) + + out.ghost_regions_in_sync = False + return out + + +def kronecker_apply_operator(mat, name: str = ""): + """The operator to apply ``mat`` with: ``mat`` itself, or a copy whose Kronecker blocks are replaced. + + Parameters + ---------- + mat : StencilMatrix or BlockLinearOperator + An assembled (block) mass matrix. + name : str + Name of the operator, for the log. + + Returns + ------- + LinearOperator + ``mat`` if no block is a Kronecker product (or :func:`kronecker_mass_enabled` is False); otherwise a + :class:`KroneckerBandMatrix` (for a StencilMatrix) or a BlockLinearOperator with the same blocks except the + Kronecker ones. It is valid only as long as ``mat`` is not modified. + """ + if not kronecker_mass_enabled(): + return mat + + if isinstance(mat, StencilMatrix): + kron = KroneckerBandMatrix.from_stencil(mat) + if kron is not None: + logger.info(f"Mass matrix {name}: applied as a Kronecker product of 1D banded matrices") + return mat if kron is None else kron + + if isinstance(mat, BlockLinearOperator): + blocks = {} + replaced = [] + for key, block in mat._blocks.items(): + kron = KroneckerBandMatrix.from_stencil(block) + if kron is None: + blocks[key] = block + else: + blocks[key] = kron + replaced.append(key) + if not replaced: + return mat + logger.info(f"Mass matrix {name}: blocks {replaced} applied as Kronecker products of 1D banded matrices") + return BlockLinearOperator(mat.domain, mat.codomain, blocks=blocks) + + return mat diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index e035e4fe4..1020f0ef6 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -25,6 +25,7 @@ from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat from struphy.feec.kernels.kernel_3d_matrixfree import kernel_3d_matrixfree from struphy.feec.kernels.kernel_3d_vec import kernel_3d_vec +from struphy.feec.kronecker_mass import kronecker_apply_operator from struphy.feec.linear_operators import BoundaryOperator from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.feec.utilities import LocalProjectionMatrix, LocalRotationMatrix, get_quad_grids @@ -1703,6 +1704,8 @@ def __init__( blocks, self._weights = self._init_blocks_from_weights(weights_info) self._mat = self._wrap_blocks(blocks) + # the operator dot() applies self._mat with (Kronecker blocks replaced), built on first use, see _apply_mat + self._apply_mat_cache = None # the weights are stored in the block order of the (possibly transposed) operator if transposed: @@ -2121,6 +2124,32 @@ def M0(self): def matrix(self): return self._mat + @property + def _apply_mat(self): + """The operator :meth:`dot` applies: :attr:`matrix`, with its blocks that are Kronecker products of 1D banded + matrices (constant weights, e.g. on a Cuboid) applied as such (:mod:`struphy.feec.kronecker_mass`). + + Built on first use after every change of the matrix (assembly, copy into, in-place arithmetic). + """ + if self._apply_mat_cache is None: + if getattr(self, "_kronecker_rejected", False): + self._apply_mat_cache = self._mat + else: + self._apply_mat_cache = kronecker_apply_operator(self._mat, name=self.name) + # a matrix changed by in-place arithmetic that is not a Kronecker product (e.g. the Schur complement + # M1 - dt^2 BC with particle data, rebuilt every step) is not checked again after further arithmetic: + # the check reads the whole matrix. Assembly and copies check again. + if getattr(self, "_changed_by_arithmetic", False) and self._apply_mat_cache is self._mat: + self._kronecker_rejected = True + return self._apply_mat_cache + + def _matrix_changed(self, by_arithmetic: bool): + """Invalidate :attr:`_apply_mat` after a change of the matrix (in-place arithmetic, or assembly/copy).""" + self._apply_mat_cache = None + self._changed_by_arithmetic = by_arithmetic + if not by_arithmetic: + self._kronecker_rejected = False + @property def nquads(self): if self._nquads is None: @@ -2168,29 +2197,54 @@ def dot(self, v, out=None, apply_bc=True): if apply_bc: if self._transposed: - self._W_boundary_op_T.dot(v, out=self._temp_WB) - self._W_extraction_op_T.dot(self._temp_WB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) - self._V_extraction_op.dot(self._temp_VE, out=self._temp_VB) - out = self._V_boundary_op.dot(self._temp_VB, out=out) + ops_in = (self._W_boundary_op_T, self._W_extraction_op_T) + ops_out = (self._V_extraction_op, self._V_boundary_op) + tmps_in = (self._temp_WB, self._temp_mat) + tmps_out = (self._temp_VE, self._temp_VB) else: - self._V_boundary_op_T.dot(v, out=self._temp_VB) - self._V_extraction_op_T.dot(self._temp_VB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) - self._W_extraction_op.dot(self._temp_WE, out=self._temp_WB) - out = self._W_boundary_op.dot(self._temp_WB, out=out) + ops_in = (self._V_boundary_op_T, self._V_extraction_op_T) + ops_out = (self._W_extraction_op, self._W_boundary_op) + tmps_in = (self._temp_VB, self._temp_mat) + tmps_out = (self._temp_WE, self._temp_WB) + out = self._apply_chain(v, ops_in, tmps_in, ops_out, tmps_out, out) else: if self._transposed: self._W_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) + self._apply_mat.dot(self._temp_mat, out=self._temp_VE) out = self._V_extraction_op.dot(self._temp_VE, out=out) else: self._V_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) + self._apply_mat.dot(self._temp_mat, out=self._temp_WE) out = self._W_extraction_op.dot(self._temp_WE, out=out) return out + def _apply_chain(self, v, ops_in, tmps_in, ops_out, tmps_out, out): + """``out = B_out E_out M E_in B_in v`` (boundary and extraction operators around the matrix), see :meth:`dot`. + + Identity operators (e.g. the extraction operators without polar splines, or boundary operators without + essential boundary conditions) are skipped instead of applied as copies: each vector is passed on as it is, + and the matrix writes into the next temporary that is actually used (or into ``out``). + """ + if out is v: + # in place: the chain must not read v after writing the output, so it reads a copy + v = v.copy() + + x = v + for op, tmp in zip(ops_in, tmps_in): + if not isinstance(op, IdentityOperator): + x = op.dot(x, out=tmp) + + # the first output operator that is not an identity reads the matrix product from its temporary + remaining = [(op, tmp) for op, tmp in zip(ops_out, tmps_out) if not isinstance(op, IdentityOperator)] + if not remaining: + return self._apply_mat.dot(x, out=out) + y = self._apply_mat.dot(x, out=remaining[0][1]) + for k, (op, _) in enumerate(remaining): + target = out if k == len(remaining) - 1 else remaining[k + 1][1] + y = op.dot(y, out=target) + return y + def transpose(self, conjugate=False): """ Returns the transposed operator. @@ -2277,6 +2331,7 @@ def assemble(self, weights=None, clear=True): Whether to first set all data to zero before assembly. If False, the new contributions are added to existing ones. """ + self._matrix_changed(by_arithmetic=False) # the matrix changes assert not self._dry_run, ( "A dry-run operator has no matrix data and cannot be assembled (memory estimation only)." ) @@ -2559,6 +2614,7 @@ def copy(self, out=None): out._weights = [list(row) for row in self._weights] self._mat.copy(out=out._mat) + out._matrix_changed(by_arithmetic=False) if self._creation_info is None: out._creation_info = None @@ -2569,6 +2625,7 @@ def copy(self, out=None): def __imul__(self, a): self._mat *= a + self._matrix_changed(by_arithmetic=True) self._creation_info = None return self @@ -2579,11 +2636,13 @@ def __iadd__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat += M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat += M + self._matrix_changed(by_arithmetic=True) return self else: @@ -2596,11 +2655,13 @@ def __isub__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat -= M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat -= M + self._matrix_changed(by_arithmetic=True) return self else: From 1a0f47773f606479705135ba450156478b9e8dfc Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:20:50 +0200 Subject: [PATCH 42/55] Kronecker mass: launch the band-pass kernel through CachedLaunch The banded passes run with the same buffers in every solver iteration (as on cuda-aware-mpi-fixes, now that feectools has CachedLaunch). Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/kronecker_mass.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py index 2050edb82..5456b2eeb 100644 --- a/src/struphy/feec/kronecker_mass.py +++ b/src/struphy/feec/kronecker_mass.py @@ -38,8 +38,10 @@ def kronecker_mass_enabled() -> bool: def _band_pass_cuda(): """The CUDA kernel of one banded pass (compiled on first call).""" from cunumpy.kernels import CudaKernel + from feectools.utilities.cached_launch import CachedLaunch - return CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass") + # launched with the same buffers in every solver iteration: cache the launch preparation + return CachedLaunch(CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass")) def _band_pass(A, inp, out, axis, offset): From 035b8a9521d517a3f4f740fb3a3f9018e6d9a41e Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:00:05 +0200 Subject: [PATCH 43/55] CuPy backend: host/device conversions in post-processing; math.sqrt for host scalars in Species Post-processing evaluates fields on the device and writes them from the host: the evaluation grids, point data and coordinates are converted where they cross, and HDF5 data is moved to the device before it is written into the spline coefficients. Co-Authored-By: Claude Opus 5.5 --- src/struphy/models/species.py | 4 ++-- src/struphy/post_processing/output.py | 31 ++++++++++++++++++--------- 2 files changed, 23 insertions(+), 12 deletions(-) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index b43464713..0ad27067d 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,9 +1,9 @@ import logging +import math import warnings from abc import ABCMeta, abstractmethod from dataclasses import fields, is_dataclass -import cunumpy as xp from maybempi import MPI from struphy.models.variables import Variable @@ -220,7 +220,7 @@ def __init__( con = ConstantsOfNature() # relevant frequencies - om_p = xp.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) + om_p = math.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) om_c = Z * con.e * units.B / (A * con.mH) # compute equation parameters diff --git a/src/struphy/post_processing/output.py b/src/struphy/post_processing/output.py index bfcf5fd58..cc65d56f8 100644 --- a/src/struphy/post_processing/output.py +++ b/src/struphy/post_processing/output.py @@ -85,6 +85,13 @@ def load_plasma_plots() -> bool: return True +def _to_host(value): + """`value` (an array, or a list/tuple of arrays for vector-valued fields) as NumPy array(s).""" + if isinstance(value, (list, tuple)): + return type(value)(_to_host(v) for v in value) + return xp.to_numpy(value) + + class ProductMapping(Mapping[str, xr.DataArray]): """A discoverable mapping whose products are loaded on first access.""" @@ -568,13 +575,14 @@ def _attach_physical_coords( self, array: xr.DataArray, etas: tuple[Any, Any, Any], dims: tuple[str, ...] ) -> xr.DataArray: """Map a raw FEEC evaluation grid and attach its Cartesian coordinates.""" - grids = [np.atleast_1d(np.asarray(eta, dtype=float)) for eta in etas] - mapped = self.domain(*grids) + # the logical grids may be device arrays (CuPy backend): map them where they live, keep the coordinates on the host + grids = [np.atleast_1d(np.asarray(xp.to_numpy(eta), dtype=float)) for eta in etas] + mapped = self.domain(*(xp.asarray(grid) for grid in grids)) shape = tuple(len(grid) for grid in grids) keep = tuple(slice(None) if dim in dims else 0 for dim in ("eta1", "eta2", "eta3")) coordinates = {} for name, values in zip(("X", "Y", "Z"), mapped): - values = np.asarray(values).reshape(shape)[keep] + values = np.asarray(xp.to_numpy(values)).reshape(shape)[keep] coordinates[name] = (dims, values) if dims else values.item() return array.assign_coords(coordinates) @@ -1295,7 +1303,9 @@ def _process_fields( # the physical grid is only needed for output, hence it is only built on rank 0 if self._pproc_rank == 0: - grids_phy = list(self.domain(*grids_log)) + # output coordinates, on the host (the logical grids stay on the device for the field evaluation) + grids_phy = [xp.to_numpy(grid) for grid in self.domain(*grids_log)] + grids_log_host = [xp.to_numpy(grid) for grid in grids_log] else: grids_phy = None @@ -1327,8 +1337,9 @@ def _process_fields( if self._pproc_rank == 0: for species, vars in vals.items(): for name, val in vars.items(): - point_data[species][name][t] = val - point_data_phy[species][name][t] = vals_phy[species][name] + # evaluated on the device on the CuPy backend; written to files from the host + point_data[species][name][t] = _to_host(val) + point_data_phy[species][name][t] = _to_host(vals_phy[species][name]) # directory for the vtk files path_fields = os.path.join(self.path_pproc, "fields_data") @@ -1338,10 +1349,10 @@ def _process_fields( for species, vars in point_data.items(): variables = {} for name, val in vars.items(): - variables[name] = wrap_field_data(val, grids_log, grids_phy=grids_phy, name=name) + variables[name] = wrap_field_data(val, grids_log_host, grids_phy=grids_phy, name=name) if physical: variables[name + "_xyz"] = wrap_field_data( - point_data_phy[species][name], grids_log, grids_phy=grids_phy, name=name + "_xyz" + point_data_phy[species][name], grids_log_host, grids_phy=grids_phy, name=name + "_xyz" ) store.write_group(store.store_path(self.path_pproc), f"/{species}", xr.Dataset(variables)) @@ -1517,7 +1528,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[n * step, p1:-p1, p2:-p2, p3:-p3]) # HDF5 data is on the host # vector-valued field else: @@ -1530,7 +1541,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3]) vector.update_ghost_regions() if is_polar: From 565eea9cac2657ad2000d6a7725a3c129bc6ff05 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:07:26 +0200 Subject: [PATCH 44/55] CUDA version of the mass-matrix assembly kernel kernel_3d_mat - kernel_3d_mat_cuda.cu: one thread per (element, codomain basis function, domain basis function); replaces the host fallback on the CuPy backend. - WeightedMassOperator.assemble reads the quadrature grids (spans, weights, basis values) as backend arrays, copied to the device once and cached. - Simulation.allocate compiles the assembly kernel before it is first called, so the compilation is not counted in the assembly's profiling regions. Co-Authored-By: Claude Opus 5.5 --- .../feec/kernels/kernel_3d_mat/__init__.py | 18 ++-- .../kernel_3d_mat/kernel_3d_mat_cuda.cu | 92 +++++++++++++++++++ src/struphy/feec/mass.py | 21 ++++- src/struphy/simulation/sim.py | 18 ++++ 4 files changed, 139 insertions(+), 10 deletions(-) create mode 100644 src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu diff --git a/src/struphy/feec/kernels/kernel_3d_mat/__init__.py b/src/struphy/feec/kernels/kernel_3d_mat/__init__.py index a881301b4..76da51673 100644 --- a/src/struphy/feec/kernels/kernel_3d_mat/__init__.py +++ b/src/struphy/feec/kernels/kernel_3d_mat/__init__.py @@ -1,8 +1,4 @@ -"""The kernel of this folder: pyccel in ``_kernels.py``, CUDA in ``_cuda.cu`` once ported. - -No CUDA version yet: on the CuPy backend the pyccel kernel runs on host copies of the arrays (a warning is -emitted once), as before the move into this folder. -""" +"""The kernel of this folder: pyccel in ``kernel_3d_mat_kernels.py``, CUDA in ``kernel_3d_mat_cuda.cu``.""" from cunumpy.kernels import Kernel @@ -11,6 +7,16 @@ # the arguments the kernel writes to, by position: data OUTPUTS = (25,) + +def _n_threads(args): + """One thread per (element, codomain basis function, domain basis function), see kernel_3d_mat_cuda.cu.""" + n_elements = args[0].size * args[1].size * args[2].size + n_pairs = 1 + for degree in args[3:9]: # pi1, pi2, pi3, pj1, pj2, pj3 + n_pairs *= int(degree) + 1 + return n_elements * n_pairs + + kernel_3d_mat = Kernel.from_folder( - __name__, structs=CUDA_STRUCTS, missing_cuda="fallback", host_options={"outputs": OUTPUTS} + __name__, structs=CUDA_STRUCTS, host_options={"outputs": OUTPUTS}, n_threads_from=_n_threads ) diff --git a/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu b/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu new file mode 100644 index 000000000..6afab9398 --- /dev/null +++ b/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu @@ -0,0 +1,92 @@ +// CUDA version of kernel_3d_mat (kernel_3d_mat_kernels.py), with the same arguments in the same order. +#include "cunumpy/array_view.cuh" +#include "cunumpy/atomic.cuh" +#include "cunumpy/index.cuh" + +/** + * Integral of Lambda_i * mat_fun * Lambda_j over the local elements, added to the stencil matrix data, as in + * kernel_3d_mat_kernels.kernel_3d_mat. + * + * One thread per (element, codomain basis function (il1, il2, il3), domain basis function (jl1, jl2, jl3)): + * n_threads = ne1 * ne2 * ne3 * (pi1 + 1) * (pi2 + 1) * (pi3 + 1) * (pj1 + 1) * (pj2 + 1) * (pj3 + 1) (see + * __init__.py), the element index varying slowest and jl3 fastest. Each thread computes the quadrature sum of its + * pair exactly as the pyccel kernel (same factors, same order) and adds it to its entry of `data` with an atomic + * addition: neighbouring elements contribute to the same entries, so only the summation order over elements + * differs from pyccel (round-off). + * + * @param spans1, spans2, spans3 Knot span index of each local element, per direction. + * @param pi1, pi2, pi3 Degrees of the codomain basis functions. + * @param pj1, pj2, pj3 Degrees of the domain basis functions. + * @param starts1, starts2, starts3 Global start index of the rows on this process. + * @param pads1, pads2, pads3 Paddings (ghost regions) of `data`. + * @param w1, w2, w3 Quadrature weights [element, quadrature point]. + * @param bi1, bi2, bi3 Codomain basis functions [element, local basis function, derivative, quadrature point]. + * @param bj1, bj2, bj3 Domain basis functions, same layout. + * @param mat_fun Weight function at the quadrature points [flattened point 1, flattened point 2, flattened point 3]. + * @param data StencilMatrix data (rows with ghost regions..., diagonals...); the results are added (not set). + */ +extern "C" __global__ void kernel_3d_mat(Array1D spans1, Array1D spans2, + Array1D spans3, int pi1, int pi2, int pi3, int pj1, int pj2, + int pj3, int starts1, int starts2, int starts3, int pads1, int pads2, + int pads3, Array2D w1, Array2D w2, Array2D w3, + Array4D bi1, Array4D bi2, Array4D bi3, + Array4D bj1, Array4D bj2, Array4D bj3, + Array3D mat_fun, Array6D data) +{ + const long long ne1 = spans1.size(); + const long long ne2 = spans2.size(); + const long long ne3 = spans3.size(); + + const long long ni1 = pi1 + 1, ni2 = pi2 + 1, ni3 = pi3 + 1; + const long long nj1 = pj1 + 1, nj2 = pj2 + 1, nj3 = pj3 + 1; + const long long n_pairs = ni1 * ni2 * ni3 * nj1 * nj2 * nj3; + + CUNUMPY_THREAD_1D(thread, ne1 * ne2 * ne3 * n_pairs); + + // CUDA-only: split the thread index into the pyccel loop variables (iel1, iel2, iel3, il1, il2, il3, jl1, jl2, + // jl3), jl3 fastest + long long rest = thread; + const long long jl3 = rest % nj3; rest /= nj3; + const long long jl2 = rest % nj2; rest /= nj2; + const long long jl1 = rest % nj1; rest /= nj1; + const long long il3 = rest % ni3; rest /= ni3; + const long long il2 = rest % ni2; rest /= ni2; + const long long il1 = rest % ni1; rest /= ni1; + const long long iel3 = rest % ne3; rest /= ne3; + const long long iel2 = rest % ne2; rest /= ne2; + const long long iel1 = rest; + + const long long nq1 = w1.shape[1]; + const long long nq2 = w2.shape[1]; + const long long nq3 = w3.shape[1]; + + double value = 0.0; + + for (long long q1 = 0; q1 < nq1; ++q1) { + for (long long q2 = 0; q2 < nq2; ++q2) { + for (long long q3 = 0; q3 < nq3; ++q3) { + const double wvol = w1(iel1, q1) * w2(iel2, q2) * w3(iel3, q3) * + mat_fun(iel1 * nq1 + q1, iel2 * nq2 + q2, iel3 * nq3 + q3); + + const double bi = bi1(iel1, il1, 0, q1) * bi2(iel2, il2, 0, q2) * bi3(iel3, il3, 0, q3); + const double bj = bj1(iel1, jl1, 0, q1) * bj2(iel2, jl2, 0, q2) * bj3(iel3, jl3, 0, q3); + + value += wvol * bi * bj; + } + } + } + + // global spline indices + const long long i_global1 = spans1(iel1) - pi1 + il1; + const long long i_global2 = spans2(iel2) - pi2 + il2; + const long long i_global3 = spans3(iel3) - pi3 + il3; + + // local spline indices (- starts --> can be negative, will therefore be written to ghost regions) + const long long i_local1 = i_global1 - starts1; + const long long i_local2 = i_global2 - starts2; + const long long i_local3 = i_global3 - starts3; + + cunumpy_atomic_add(&data(pads1 + i_local1, pads2 + i_local2, pads3 + i_local3, pads1 + jl1 - il1, + pads2 + jl2 - il2, pads3 + jl3 - il3), + value); +} diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index f75c6de17..e035e4fe4 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -2242,6 +2242,19 @@ def transpose(self, conjugate=False): return M + def _quad_grid(self, space_name, attribute, component): + """Quadrature-grid arrays (spans, weights or basis values) of one component of a space, as backend arrays. + + The Derham keeps them as host arrays; on the CuPy backend the assembly kernel needs them on the device, + so they are copied there once and cached (they do not change between assemblies). + """ + cache = self.__dict__.setdefault("_quad_grid_cache", {}) + key = (space_name, attribute, component) + if key not in cache: + arrays = getattr(self.derham.spline_attributes[space_name], attribute)[component] + cache[key] = [xp.asarray(array) for array in arrays] + return cache[key] + def assemble(self, weights=None, clear=True): r""" Assembles the weighted mass matrix, i.e. computes the integrals @@ -2320,7 +2333,7 @@ def assemble(self, weights=None, clear=True): # loop over codomain spaces (rows) for a, codomain_space in enumerate(codomain_spaces): # knot span indices of elements of local domain - codomain_spans = spline_attr[W_name].quad_grid_spans[a] + codomain_spans = self._quad_grid(W_name, "quad_grid_spans", a) # global start spline index on process codomain_starts = [int(start) for start in codomain_space.coeff_space.starts] @@ -2332,10 +2345,10 @@ def assemble(self, weights=None, clear=True): pts = self._quad_pts(W_name, a) # global quadrature weights in format (local element, local weight) - wts = spline_attr[W_name].quad_grid_wts[a] + wts = self._quad_grid(W_name, "quad_grid_wts", a) # evaluated basis functions at quadrature points of codomain space - codomain_basis = spline_attr[W_name].quad_grid_bases[a] + codomain_basis = self._quad_grid(W_name, "quad_grid_bases", a) # loop over domain spaces (columns) for b, domain_space in enumerate(domain_spaces): @@ -2375,7 +2388,7 @@ def assemble(self, weights=None, clear=True): ) # evaluated basis functions at quadrature points of domain space - domain_basis = spline_attr[V_name].quad_grid_bases[b] + domain_basis = self._quad_grid(V_name, "quad_grid_bases", b) # assemble matrix (if mat_w is not zero) by calling the appropriate kernel (1d, 2d or 3d) if not_weight_zero or self._is_scalar: diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 9452e81c5..e6f358832 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -243,6 +243,8 @@ def allocate(self): logger.debug("\nAllocating simulation data ...") + self._compile_setup_cuda_kernels() + with ProfileManager.profile_region("setup: allocate"): # feec with ProfileManager.profile_region("allocate feec", functions=[self._allocate_feec]): @@ -280,6 +282,22 @@ def kernels(self) -> tuple[Kernel, ...]: owners += self.model.prop_list return collect_kernels(*owners) + def _compile_setup_cuda_kernels(self) -> None: + """Compile the CUDA kernels that :meth:`allocate` itself calls (mass-matrix assembly) before it calls them. + + :meth:`compile_cuda_kernels` runs after :meth:`allocate`, so without this the first assembly would include + the compilation (~1 s per kernel on a node whose CuPy cache does not have it yet), in the assembly's + profiling regions. Here it is counted in ``"setup: compile cuda kernels (assembly)"``. Nothing happens on + the NumPy backend. + """ + if xp.get_backend() != "cupy": + return + + from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat + + with ProfileManager.profile_region("setup: compile cuda kernels (assembly)"): + kernel_3d_mat.compile() + def compile_cuda_kernels(self) -> tuple[Kernel, ...]: """Compile the CUDA versions of the kernels of :meth:`kernels` now, before the time stepping. From 08ea857ff80d2408ad8a7997f9a957e23a901ac7 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:26:26 +0200 Subject: [PATCH 45/55] Parity cases for kernel_3d_mat Random quadrature data on 3 x 2 x 2 elements with mixed degrees, with the first row local or a ghost row. test_cuda_kernels_have_parity_cases requires a case for every kernel with a CUDA version; the emulation test compares it with the pyccel kernel without a GPU. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/cuda_parity_cases.py | 26 ++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index a2396ee5b..6e10707be 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -315,6 +315,30 @@ def kernel_pullpush_args(case): ) +def kernel_3d_mat_threads(args): + """One thread per (element, codomain basis function, domain basis function), as the kernel's launch size.""" + from struphy.feec.kernels.kernel_3d_mat import _n_threads + + return _n_threads(args) + + +def kernel_3d_mat_args(case): + """Mass-matrix assembly on 3 x 2 x 2 elements with mixed degrees and random quadrature data; `case` is the start + index of the rows in direction 1 (1: the first row is a ghost row, local index -1).""" + rng = np.random.default_rng(20 + case) + ne, nq = (3, 2, 2), (3, 2, 2) + pi, pj, pads = (2, 1, 1), (1, 2, 1), (2, 2, 2) + starts = (case, 0, 0) + spans = [xp.asarray(np.arange(n, dtype=np.int64) + p) for n, p in zip(ne, pi)] + w = [xp.asarray(rng.random((n, q))) for n, q in zip(ne, nq)] + bi = [xp.asarray(rng.random((n, p + 1, 1, q))) for n, p, q in zip(ne, pi, nq)] + bj = [xp.asarray(rng.random((n, p + 1, 1, q))) for n, p, q in zip(ne, pj, nq)] + mat_fun = xp.asarray(rng.normal(size=tuple(n * q for n, q in zip(ne, nq)))) + rows = tuple(n + p + 2 * pad for n, p, pad in zip(ne, pi, pads)) + data = xp.asarray(rng.normal(size=rows + tuple(2 * pad + 1 for pad in pads))) # added to, not overwritten + return (*spans, *pi, *pj, *starts, *pads, *w, *bi, *bj, mat_fun, data) + + # ---------------------------------------------------------------- all kernels with a CUDA version PUSHER_TOLERANCES = {"rtol": 1e-13, "atol": 1e-14} @@ -338,6 +362,8 @@ def size_of(index, per_thread=1): PARITY_CASES = { + # atomic additions: the summation order over elements differs from pyccel + "kernel_3d_mat": ParityCases((0, 1), kernel_3d_mat_args, rtol=1e-12, atol=1e-13, n_threads=kernel_3d_mat_threads), "push_eta_stage": ParityCases( [ (bc, method, stage) From 715cc5c8c49394e53357038c86827b1348d10336 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:08:01 +0200 Subject: [PATCH 46/55] Apply separable mass matrices as Kronecker products of 1D banded matrices WeightedMassOperator.dot applies the blocks of its matrix that are Kronecker products of 1D banded matrices (constant weights, e.g. on a Cuboid) as three banded passes (struphy.feec.kronecker_mass, CUDA kernel kron_band_pass on the GPU) instead of the 3D stencil product. The replacement is built on first use and invalidated by assembly, copies and in-place arithmetic; a matrix changed by arithmetic that is not a Kronecker product is not checked again. STRUPHY_KRONECKER_MASS=0 switches it off. The boundary and extraction operators around the matrix are skipped when they are identities instead of being applied as copies. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/kron_band_pass_cuda.cu | 35 ++++ src/struphy/feec/kronecker_mass.py | 247 ++++++++++++++++++++++++ src/struphy/feec/mass.py | 85 ++++++-- 3 files changed, 355 insertions(+), 12 deletions(-) create mode 100644 src/struphy/feec/kron_band_pass_cuda.cu create mode 100644 src/struphy/feec/kronecker_mass.py diff --git a/src/struphy/feec/kron_band_pass_cuda.cu b/src/struphy/feec/kron_band_pass_cuda.cu new file mode 100644 index 000000000..745e782f4 --- /dev/null +++ b/src/struphy/feec/kron_band_pass_cuda.cu @@ -0,0 +1,35 @@ +/** + * One pass of a Kronecker product of 1D banded matrices (struphy.feec.kronecker_mass._band_pass): + * + * out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...] along `axis`, + * + * with the other two axes unchanged (out and inp have the same extent along them). One thread per entry of `out`, + * the last axis fastest, so that neighbouring threads read neighbouring entries of `inp` along every axis. `out` may + * be a strided view (the owned rows of a stencil vector's data). + * + * @param A Banded 1D matrix, shape (out.shape[axis], n_diagonals). + * @param inp Input array (3D). + * @param out Output array (3D, any strides); written, not added to. + * @param axis Axis of the pass (0, 1 or 2). + * @param offset Index in `inp` of diagonal 0 of row 0 along `axis`. + */ +#include "cunumpy/array_view.cuh" + +extern "C" __global__ void kron_band_pass(Array2D A, Array3D inp, Array3D out, int axis, + long long offset) { + const long long n1 = out.shape[1], n2 = out.shape[2]; + const long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; + if (t >= out.shape[0] * n1 * n2) return; + const long long i2 = t % n2; + const long long i1 = (t / n2) % n1; + const long long i0 = t / (n1 * n2); + const long long i = axis == 0 ? i0 : (axis == 1 ? i1 : i2); + + double v = 0.0; + for (long long d = 0; d < A.shape[1]; ++d) { + const long long j = i + offset + d; + const double x = axis == 0 ? inp(j, i1, i2) : (axis == 1 ? inp(i0, j, i2) : inp(i0, i1, j)); + v += A(i, d) * x; + } + out(i0, i1, i2) = v; +} diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py new file mode 100644 index 000000000..2050edb82 --- /dev/null +++ b/src/struphy/feec/kronecker_mass.py @@ -0,0 +1,247 @@ +"""Fast application of assembled mass matrices that are Kronecker products of 1D banded matrices. + +A mass matrix block ``M[i, j] = int w Lambda_i Lambda_j`` of tensor-product splines factorizes as +``M = c * M_1 (x) M_2 (x) M_3`` when its weight ``w`` is constant (e.g. on a Cuboid). Applied as a +:class:`~feectools.linalg.stencil.StencilMatrix`, every row reads all ``n_1 * n_2 * n_3`` diagonals (105 for +degrees (1, 2, 3)) from memory; as a Kronecker product it is three 1D banded passes over the vector, ``n_1 + n_2 + +n_3`` operations per entry and no 3D matrix data. :func:`kronecker_apply_operator` checks each assembled block and +returns an operator for the matrix-vector products in which the blocks that are Kronecker products (to round-off) +are replaced by :class:`KroneckerBandMatrix`. + +The 1D factors are read from the assembled data itself (rows owned by this process), so they hold exactly what the +assembly computed, on any domain decomposition, and the check is the assembled data, not an assumption about the +mapping or the weights. +""" + +import logging +import os +from functools import cache +from pathlib import Path + +import cunumpy as xp +from feectools.linalg.basic import LinearOperator +from feectools.linalg.block import BlockLinearOperator +from feectools.linalg.stencil import StencilMatrix, StencilVector + +logger = logging.getLogger("struphy") + +# relative tolerance of the Kronecker check: the 3D assembly and the product of the factors differ by round-off only +KRONECKER_RTOL = 1e-12 + + +def kronecker_mass_enabled() -> bool: + """Whether :func:`kronecker_apply_operator` may replace blocks (``STRUPHY_KRONECKER_MASS=0`` switches it off).""" + return os.environ.get("STRUPHY_KRONECKER_MASS", "1").lower() not in ("0", "false", "no") + + +@cache +def _band_pass_cuda(): + """The CUDA kernel of one banded pass (compiled on first call).""" + from cunumpy.kernels import CudaKernel + + return CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass") + + +def _band_pass(A, inp, out, axis, offset): + """``out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...]`` along ``axis`` (other axes unchanged).""" + if xp.is_gpu(out): + _band_pass_cuda()(A, inp, out, axis, offset, n_threads=out.size) + return + m = out.shape[axis] + shape = [1, 1, 1] + shape[axis] = m + out[...] = 0.0 + for d in range(A.shape[1]): + idx = [slice(None)] * 3 + idx[axis] = slice(offset + d, offset + d + m) + out += A[:, d].reshape(shape) * inp[tuple(idx)] + + +class KroneckerBandMatrix(LinearOperator): + """A square 3D :class:`~feectools.linalg.stencil.StencilMatrix` applied as a Kronecker product of 1D banded + matrices, ``(A_1 (x) A_2 (x) A_3) x`` on the rows owned by this process. + + ``A_k[i, d]`` is the entry of owned row ``i`` (along direction ``k``) and diagonal ``d``, column + ``i - q_k + d`` with ``q_k = (n_k - 1) // 2``, as in the stencil matrix. The product is computed in three passes, + each over the owned rows of its direction and the full (padded) extent of the directions after it, so the passes + need no communication beyond the ghost regions of the input. Use :meth:`from_stencil`. + + Parameters + ---------- + stencil : StencilMatrix + The assembled matrix (kept for ``toarray``, ``tosparse`` and ``transpose``). + factors : tuple of arrays + ``A_1, A_2, A_3``, shapes ``(m_k, n_k)``, C-contiguous backend arrays. + """ + + def __init__(self, stencil: StencilMatrix, factors): + self._stencil = stencil + self._factors = factors + + V = stencil.domain + self._m = tuple(int(e - s + 1) for s, e in zip(V.starts, V.ends)) + self._pads = tuple(int(p) for p in V.pads) + self._q = tuple((int(n) - 1) // 2 for n in stencil._data.shape[3:]) + + # temporaries of the first two passes: owned rows along the passes done, padded extent along the others + x_shape = V.zeros()._data.shape + self._t1 = xp.empty((self._m[0], x_shape[1], x_shape[2]), dtype=stencil._data.dtype) + self._t2 = xp.empty((self._m[0], self._m[1], x_shape[2]), dtype=stencil._data.dtype) + + @classmethod + def from_stencil(cls, S: StencilMatrix, rtol: float = KRONECKER_RTOL): + """The Kronecker form of ``S`` if it is one (to ``rtol``, relative to its largest entry), else None. + + Only square 3D matrices on one space (domain is codomain) with shifts 1 qualify, e.g. the diagonal blocks of + the mass matrices. The factors are read from the owned rows of ``S``: with reference rows ``r_k`` (the first + owned row along each direction) and central diagonals ``q_k``, + ``A_1[i, :] = S[i, r_2, r_3, :, q_2, q_3]``, ``A_2[i, :] = S[r_1, i, r_3, q_1, :, q_3] / S[r, q]``, + ``A_3[i, :] = S[r_1, r_2, i, q_1, q_2, :] / S[r, q]``; the product is then compared with all owned rows. + """ + if not isinstance(S, StencilMatrix) or S.domain is not S.codomain or S.domain.ndim != 3: + return None + V = S.domain + if any(int(s) != 1 for s in V.shifts) or S._data.ndim != 6: + return None + + p = [int(p) for p in V.pads] + m = [int(e - s + 1) for s, e in zip(V.starts, V.ends)] + q = [(int(n) - 1) // 2 for n in S._data.shape[3:]] + if min(m) < 1: + return None + data = S._data + + den = data[p[0], p[1], p[2], q[0], q[1], q[2]] + if float(den) == 0.0: + return None + A1 = xp.ascontiguousarray(data[p[0] : p[0] + m[0], p[1], p[2], :, q[1], q[2]]) + A2 = xp.ascontiguousarray(data[p[0], p[1] : p[1] + m[1], p[2], q[0], :, q[2]] / den) + A3 = xp.ascontiguousarray(data[p[0], p[1], p[2] : p[2] + m[2], q[0], q[1], :] / den) + + # compare with every owned row: planes (rows along direction 1) in chunks, the maxima kept on the device + # (one read at the end); a few planes first, which rejects most matrices that are not Kronecker products + # (e.g. ones assembled from particle data, checked again after every change) at once + owned = data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + A23 = A2[:, None, None, :, None] * A3[None, :, None, None, :] # (m2, m3, 1, n2, n3) + + def max_diff_and_abs(rows): + block = owned[rows] + recon = A1[rows][:, None, None, :, None, None] * A23[None] + return xp.max(xp.abs(block - recon)), xp.max(xp.abs(block)) + + def fails(max_diff, max_abs): + max_diff, max_abs = float(max_diff), float(max_abs) + return max_abs == 0.0 or max_diff > rtol * max_abs + + quick = sorted({0, m[0] // 2, m[0] - 1}) + if fails(*max_diff_and_abs(xp.asarray(quick))): + return None + + # chunks of at most ~64 MB of temporaries + plane = max(1, owned[0].size) + chunk = max(1, min(m[0], (8 << 20) // plane)) + max_diff = xp.zeros(()) + max_abs = xp.zeros(()) + for start in range(0, m[0], chunk): + d, a = max_diff_and_abs(slice(start, min(start + chunk, m[0]))) + max_diff = xp.maximum(max_diff, d) + max_abs = xp.maximum(max_abs, a) + if fails(max_diff, max_abs): + return None + return cls(S, (A1, A2, A3)) + + # ------------------------------------------------------------------ + # LinearOperator interface + # ------------------------------------------------------------------ + @property + def domain(self): + return self._stencil.domain + + @property + def codomain(self): + return self._stencil.codomain + + @property + def dtype(self): + return self._stencil.dtype + + @property + def factors(self): + """The 1D factors ``A_1, A_2, A_3`` (rows owned by this process).""" + return self._factors + + def tosparse(self): + return self._stencil.tosparse() + + def toarray(self): + return self._stencil.toarray() + + def transpose(self, conjugate=False): + return self._stencil.transpose(conjugate=conjugate) + + def dot(self, v, out=None): + assert isinstance(v, StencilVector) + assert v.space is self.domain + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.codomain + + if not v.ghost_regions_in_sync: + v.update_ghost_regions() + + A1, A2, A3 = self._factors + p, q, m = self._pads, self._q, self._m + owned = out._data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + + _band_pass(A1, v._data, self._t1, 0, p[0] - q[0]) + _band_pass(A2, self._t1, self._t2, 1, p[1] - q[1]) + _band_pass(A3, self._t2, owned, 2, p[2] - q[2]) + + out.ghost_regions_in_sync = False + return out + + +def kronecker_apply_operator(mat, name: str = ""): + """The operator to apply ``mat`` with: ``mat`` itself, or a copy whose Kronecker blocks are replaced. + + Parameters + ---------- + mat : StencilMatrix or BlockLinearOperator + An assembled (block) mass matrix. + name : str + Name of the operator, for the log. + + Returns + ------- + LinearOperator + ``mat`` if no block is a Kronecker product (or :func:`kronecker_mass_enabled` is False); otherwise a + :class:`KroneckerBandMatrix` (for a StencilMatrix) or a BlockLinearOperator with the same blocks except the + Kronecker ones. It is valid only as long as ``mat`` is not modified. + """ + if not kronecker_mass_enabled(): + return mat + + if isinstance(mat, StencilMatrix): + kron = KroneckerBandMatrix.from_stencil(mat) + if kron is not None: + logger.info(f"Mass matrix {name}: applied as a Kronecker product of 1D banded matrices") + return mat if kron is None else kron + + if isinstance(mat, BlockLinearOperator): + blocks = {} + replaced = [] + for key, block in mat._blocks.items(): + kron = KroneckerBandMatrix.from_stencil(block) + if kron is None: + blocks[key] = block + else: + blocks[key] = kron + replaced.append(key) + if not replaced: + return mat + logger.info(f"Mass matrix {name}: blocks {replaced} applied as Kronecker products of 1D banded matrices") + return BlockLinearOperator(mat.domain, mat.codomain, blocks=blocks) + + return mat diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index e035e4fe4..1020f0ef6 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -25,6 +25,7 @@ from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat from struphy.feec.kernels.kernel_3d_matrixfree import kernel_3d_matrixfree from struphy.feec.kernels.kernel_3d_vec import kernel_3d_vec +from struphy.feec.kronecker_mass import kronecker_apply_operator from struphy.feec.linear_operators import BoundaryOperator from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.feec.utilities import LocalProjectionMatrix, LocalRotationMatrix, get_quad_grids @@ -1703,6 +1704,8 @@ def __init__( blocks, self._weights = self._init_blocks_from_weights(weights_info) self._mat = self._wrap_blocks(blocks) + # the operator dot() applies self._mat with (Kronecker blocks replaced), built on first use, see _apply_mat + self._apply_mat_cache = None # the weights are stored in the block order of the (possibly transposed) operator if transposed: @@ -2121,6 +2124,32 @@ def M0(self): def matrix(self): return self._mat + @property + def _apply_mat(self): + """The operator :meth:`dot` applies: :attr:`matrix`, with its blocks that are Kronecker products of 1D banded + matrices (constant weights, e.g. on a Cuboid) applied as such (:mod:`struphy.feec.kronecker_mass`). + + Built on first use after every change of the matrix (assembly, copy into, in-place arithmetic). + """ + if self._apply_mat_cache is None: + if getattr(self, "_kronecker_rejected", False): + self._apply_mat_cache = self._mat + else: + self._apply_mat_cache = kronecker_apply_operator(self._mat, name=self.name) + # a matrix changed by in-place arithmetic that is not a Kronecker product (e.g. the Schur complement + # M1 - dt^2 BC with particle data, rebuilt every step) is not checked again after further arithmetic: + # the check reads the whole matrix. Assembly and copies check again. + if getattr(self, "_changed_by_arithmetic", False) and self._apply_mat_cache is self._mat: + self._kronecker_rejected = True + return self._apply_mat_cache + + def _matrix_changed(self, by_arithmetic: bool): + """Invalidate :attr:`_apply_mat` after a change of the matrix (in-place arithmetic, or assembly/copy).""" + self._apply_mat_cache = None + self._changed_by_arithmetic = by_arithmetic + if not by_arithmetic: + self._kronecker_rejected = False + @property def nquads(self): if self._nquads is None: @@ -2168,29 +2197,54 @@ def dot(self, v, out=None, apply_bc=True): if apply_bc: if self._transposed: - self._W_boundary_op_T.dot(v, out=self._temp_WB) - self._W_extraction_op_T.dot(self._temp_WB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) - self._V_extraction_op.dot(self._temp_VE, out=self._temp_VB) - out = self._V_boundary_op.dot(self._temp_VB, out=out) + ops_in = (self._W_boundary_op_T, self._W_extraction_op_T) + ops_out = (self._V_extraction_op, self._V_boundary_op) + tmps_in = (self._temp_WB, self._temp_mat) + tmps_out = (self._temp_VE, self._temp_VB) else: - self._V_boundary_op_T.dot(v, out=self._temp_VB) - self._V_extraction_op_T.dot(self._temp_VB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) - self._W_extraction_op.dot(self._temp_WE, out=self._temp_WB) - out = self._W_boundary_op.dot(self._temp_WB, out=out) + ops_in = (self._V_boundary_op_T, self._V_extraction_op_T) + ops_out = (self._W_extraction_op, self._W_boundary_op) + tmps_in = (self._temp_VB, self._temp_mat) + tmps_out = (self._temp_WE, self._temp_WB) + out = self._apply_chain(v, ops_in, tmps_in, ops_out, tmps_out, out) else: if self._transposed: self._W_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) + self._apply_mat.dot(self._temp_mat, out=self._temp_VE) out = self._V_extraction_op.dot(self._temp_VE, out=out) else: self._V_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) + self._apply_mat.dot(self._temp_mat, out=self._temp_WE) out = self._W_extraction_op.dot(self._temp_WE, out=out) return out + def _apply_chain(self, v, ops_in, tmps_in, ops_out, tmps_out, out): + """``out = B_out E_out M E_in B_in v`` (boundary and extraction operators around the matrix), see :meth:`dot`. + + Identity operators (e.g. the extraction operators without polar splines, or boundary operators without + essential boundary conditions) are skipped instead of applied as copies: each vector is passed on as it is, + and the matrix writes into the next temporary that is actually used (or into ``out``). + """ + if out is v: + # in place: the chain must not read v after writing the output, so it reads a copy + v = v.copy() + + x = v + for op, tmp in zip(ops_in, tmps_in): + if not isinstance(op, IdentityOperator): + x = op.dot(x, out=tmp) + + # the first output operator that is not an identity reads the matrix product from its temporary + remaining = [(op, tmp) for op, tmp in zip(ops_out, tmps_out) if not isinstance(op, IdentityOperator)] + if not remaining: + return self._apply_mat.dot(x, out=out) + y = self._apply_mat.dot(x, out=remaining[0][1]) + for k, (op, _) in enumerate(remaining): + target = out if k == len(remaining) - 1 else remaining[k + 1][1] + y = op.dot(y, out=target) + return y + def transpose(self, conjugate=False): """ Returns the transposed operator. @@ -2277,6 +2331,7 @@ def assemble(self, weights=None, clear=True): Whether to first set all data to zero before assembly. If False, the new contributions are added to existing ones. """ + self._matrix_changed(by_arithmetic=False) # the matrix changes assert not self._dry_run, ( "A dry-run operator has no matrix data and cannot be assembled (memory estimation only)." ) @@ -2559,6 +2614,7 @@ def copy(self, out=None): out._weights = [list(row) for row in self._weights] self._mat.copy(out=out._mat) + out._matrix_changed(by_arithmetic=False) if self._creation_info is None: out._creation_info = None @@ -2569,6 +2625,7 @@ def copy(self, out=None): def __imul__(self, a): self._mat *= a + self._matrix_changed(by_arithmetic=True) self._creation_info = None return self @@ -2579,11 +2636,13 @@ def __iadd__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat += M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat += M + self._matrix_changed(by_arithmetic=True) return self else: @@ -2596,11 +2655,13 @@ def __isub__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat -= M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat -= M + self._matrix_changed(by_arithmetic=True) return self else: From c78a6f71d6ceeae10df3589076ae15b8d16a7444 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:20:50 +0200 Subject: [PATCH 47/55] Kronecker mass: launch the band-pass kernel through CachedLaunch The banded passes run with the same buffers in every solver iteration (as on cuda-aware-mpi-fixes, now that feectools has CachedLaunch). Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/kronecker_mass.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py index 2050edb82..5456b2eeb 100644 --- a/src/struphy/feec/kronecker_mass.py +++ b/src/struphy/feec/kronecker_mass.py @@ -38,8 +38,10 @@ def kronecker_mass_enabled() -> bool: def _band_pass_cuda(): """The CUDA kernel of one banded pass (compiled on first call).""" from cunumpy.kernels import CudaKernel + from feectools.utilities.cached_launch import CachedLaunch - return CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass") + # launched with the same buffers in every solver iteration: cache the launch preparation + return CachedLaunch(CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass")) def _band_pass(A, inp, out, axis, offset): From 352a4bb73513a98f3e2d0905b06837299a704956 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 24 Oct 2026 13:39:58 +0200 Subject: [PATCH 48/55] Set feectools to devel-tiny --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index b3de2933b..c3c7d941e 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit b3de2933b5071a6762aa0432838c39f201cdc73a +Subproject commit c3c7d941e97f5adb40dbe6e6704d261fb48c3a11 From b3d5f2ecfe73b497a1812599039a207a157de1b3 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 24 Oct 2026 13:40:24 +0200 Subject: [PATCH 49/55] Update feectools dependency to 0.10.0 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index c106ba648..e00ff8276 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ dependencies = [ "cunumpy>=0.6.3, <=0.7.0", "pyccel>=2.2.0, <=2.2.3", "maybempi>=0.1.2", - "feectools>=0.9.0, <=0.9.0", + "feectools>=0.10.0, <=0.10.0", "scipy<=1.18.0", "h5py<=3.16.0", "h5netcdf<=1.8.1", From 8b7d737fd46cdd27a9cdaa707da6e686e1d4cf0c Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:00:05 +0200 Subject: [PATCH 50/55] CuPy backend: host/device conversions in post-processing; math.sqrt for host scalars in Species Post-processing evaluates fields on the device and writes them from the host: the evaluation grids, point data and coordinates are converted where they cross, and HDF5 data is moved to the device before it is written into the spline coefficients. Co-Authored-By: Claude Opus 5.5 --- src/struphy/models/species.py | 4 ++-- src/struphy/post_processing/output.py | 31 ++++++++++++++++++--------- 2 files changed, 23 insertions(+), 12 deletions(-) diff --git a/src/struphy/models/species.py b/src/struphy/models/species.py index b43464713..0ad27067d 100644 --- a/src/struphy/models/species.py +++ b/src/struphy/models/species.py @@ -1,9 +1,9 @@ import logging +import math import warnings from abc import ABCMeta, abstractmethod from dataclasses import fields, is_dataclass -import cunumpy as xp from maybempi import MPI from struphy.models.variables import Variable @@ -220,7 +220,7 @@ def __init__( con = ConstantsOfNature() # relevant frequencies - om_p = xp.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) + om_p = math.sqrt(units.n * (Z * con.e) ** 2 / (con.eps0 * A * con.mH)) om_c = Z * con.e * units.B / (A * con.mH) # compute equation parameters diff --git a/src/struphy/post_processing/output.py b/src/struphy/post_processing/output.py index bfcf5fd58..cc65d56f8 100644 --- a/src/struphy/post_processing/output.py +++ b/src/struphy/post_processing/output.py @@ -85,6 +85,13 @@ def load_plasma_plots() -> bool: return True +def _to_host(value): + """`value` (an array, or a list/tuple of arrays for vector-valued fields) as NumPy array(s).""" + if isinstance(value, (list, tuple)): + return type(value)(_to_host(v) for v in value) + return xp.to_numpy(value) + + class ProductMapping(Mapping[str, xr.DataArray]): """A discoverable mapping whose products are loaded on first access.""" @@ -568,13 +575,14 @@ def _attach_physical_coords( self, array: xr.DataArray, etas: tuple[Any, Any, Any], dims: tuple[str, ...] ) -> xr.DataArray: """Map a raw FEEC evaluation grid and attach its Cartesian coordinates.""" - grids = [np.atleast_1d(np.asarray(eta, dtype=float)) for eta in etas] - mapped = self.domain(*grids) + # the logical grids may be device arrays (CuPy backend): map them where they live, keep the coordinates on the host + grids = [np.atleast_1d(np.asarray(xp.to_numpy(eta), dtype=float)) for eta in etas] + mapped = self.domain(*(xp.asarray(grid) for grid in grids)) shape = tuple(len(grid) for grid in grids) keep = tuple(slice(None) if dim in dims else 0 for dim in ("eta1", "eta2", "eta3")) coordinates = {} for name, values in zip(("X", "Y", "Z"), mapped): - values = np.asarray(values).reshape(shape)[keep] + values = np.asarray(xp.to_numpy(values)).reshape(shape)[keep] coordinates[name] = (dims, values) if dims else values.item() return array.assign_coords(coordinates) @@ -1295,7 +1303,9 @@ def _process_fields( # the physical grid is only needed for output, hence it is only built on rank 0 if self._pproc_rank == 0: - grids_phy = list(self.domain(*grids_log)) + # output coordinates, on the host (the logical grids stay on the device for the field evaluation) + grids_phy = [xp.to_numpy(grid) for grid in self.domain(*grids_log)] + grids_log_host = [xp.to_numpy(grid) for grid in grids_log] else: grids_phy = None @@ -1327,8 +1337,9 @@ def _process_fields( if self._pproc_rank == 0: for species, vars in vals.items(): for name, val in vars.items(): - point_data[species][name][t] = val - point_data_phy[species][name][t] = vals_phy[species][name] + # evaluated on the device on the CuPy backend; written to files from the host + point_data[species][name][t] = _to_host(val) + point_data_phy[species][name][t] = _to_host(vals_phy[species][name]) # directory for the vtk files path_fields = os.path.join(self.path_pproc, "fields_data") @@ -1338,10 +1349,10 @@ def _process_fields( for species, vars in point_data.items(): variables = {} for name, val in vars.items(): - variables[name] = wrap_field_data(val, grids_log, grids_phy=grids_phy, name=name) + variables[name] = wrap_field_data(val, grids_log_host, grids_phy=grids_phy, name=name) if physical: variables[name + "_xyz"] = wrap_field_data( - point_data_phy[species][name], grids_log, grids_phy=grids_phy, name=name + "_xyz" + point_data_phy[species][name], grids_log_host, grids_phy=grids_phy, name=name + "_xyz" ) store.write_group(store.store_path(self.path_pproc), f"/{species}", xr.Dataset(variables)) @@ -1517,7 +1528,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[n * step, p1:-p1, p2:-p2, p3:-p3]) # HDF5 data is on the host # vector-valued field else: @@ -1530,7 +1541,7 @@ def _load_femfields(self, fields: dict, files: list, n: int, step: int = 1): s1 : e1 + 1, s2 : e2 + 1, s3 : e3 + 1, - ] = ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3] + ] = xp.asarray(ddset[str(comp + 1)][n * step, p1:-p1, p2:-p2, p3:-p3]) vector.update_ghost_regions() if is_polar: From 0191d399593ca7b9d20b77cb522006debeb7a620 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:07:26 +0200 Subject: [PATCH 51/55] CUDA version of the mass-matrix assembly kernel kernel_3d_mat - kernel_3d_mat_cuda.cu: one thread per (element, codomain basis function, domain basis function); replaces the host fallback on the CuPy backend. - WeightedMassOperator.assemble reads the quadrature grids (spans, weights, basis values) as backend arrays, copied to the device once and cached. - Simulation.allocate compiles the assembly kernel before it is first called, so the compilation is not counted in the assembly's profiling regions. Co-Authored-By: Claude Opus 5.5 --- .../feec/kernels/kernel_3d_mat/__init__.py | 18 ++-- .../kernel_3d_mat/kernel_3d_mat_cuda.cu | 92 +++++++++++++++++++ src/struphy/feec/mass.py | 21 ++++- src/struphy/simulation/sim.py | 18 ++++ 4 files changed, 139 insertions(+), 10 deletions(-) create mode 100644 src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu diff --git a/src/struphy/feec/kernels/kernel_3d_mat/__init__.py b/src/struphy/feec/kernels/kernel_3d_mat/__init__.py index a881301b4..76da51673 100644 --- a/src/struphy/feec/kernels/kernel_3d_mat/__init__.py +++ b/src/struphy/feec/kernels/kernel_3d_mat/__init__.py @@ -1,8 +1,4 @@ -"""The kernel of this folder: pyccel in ``_kernels.py``, CUDA in ``_cuda.cu`` once ported. - -No CUDA version yet: on the CuPy backend the pyccel kernel runs on host copies of the arrays (a warning is -emitted once), as before the move into this folder. -""" +"""The kernel of this folder: pyccel in ``kernel_3d_mat_kernels.py``, CUDA in ``kernel_3d_mat_cuda.cu``.""" from cunumpy.kernels import Kernel @@ -11,6 +7,16 @@ # the arguments the kernel writes to, by position: data OUTPUTS = (25,) + +def _n_threads(args): + """One thread per (element, codomain basis function, domain basis function), see kernel_3d_mat_cuda.cu.""" + n_elements = args[0].size * args[1].size * args[2].size + n_pairs = 1 + for degree in args[3:9]: # pi1, pi2, pi3, pj1, pj2, pj3 + n_pairs *= int(degree) + 1 + return n_elements * n_pairs + + kernel_3d_mat = Kernel.from_folder( - __name__, structs=CUDA_STRUCTS, missing_cuda="fallback", host_options={"outputs": OUTPUTS} + __name__, structs=CUDA_STRUCTS, host_options={"outputs": OUTPUTS}, n_threads_from=_n_threads ) diff --git a/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu b/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu new file mode 100644 index 000000000..6afab9398 --- /dev/null +++ b/src/struphy/feec/kernels/kernel_3d_mat/kernel_3d_mat_cuda.cu @@ -0,0 +1,92 @@ +// CUDA version of kernel_3d_mat (kernel_3d_mat_kernels.py), with the same arguments in the same order. +#include "cunumpy/array_view.cuh" +#include "cunumpy/atomic.cuh" +#include "cunumpy/index.cuh" + +/** + * Integral of Lambda_i * mat_fun * Lambda_j over the local elements, added to the stencil matrix data, as in + * kernel_3d_mat_kernels.kernel_3d_mat. + * + * One thread per (element, codomain basis function (il1, il2, il3), domain basis function (jl1, jl2, jl3)): + * n_threads = ne1 * ne2 * ne3 * (pi1 + 1) * (pi2 + 1) * (pi3 + 1) * (pj1 + 1) * (pj2 + 1) * (pj3 + 1) (see + * __init__.py), the element index varying slowest and jl3 fastest. Each thread computes the quadrature sum of its + * pair exactly as the pyccel kernel (same factors, same order) and adds it to its entry of `data` with an atomic + * addition: neighbouring elements contribute to the same entries, so only the summation order over elements + * differs from pyccel (round-off). + * + * @param spans1, spans2, spans3 Knot span index of each local element, per direction. + * @param pi1, pi2, pi3 Degrees of the codomain basis functions. + * @param pj1, pj2, pj3 Degrees of the domain basis functions. + * @param starts1, starts2, starts3 Global start index of the rows on this process. + * @param pads1, pads2, pads3 Paddings (ghost regions) of `data`. + * @param w1, w2, w3 Quadrature weights [element, quadrature point]. + * @param bi1, bi2, bi3 Codomain basis functions [element, local basis function, derivative, quadrature point]. + * @param bj1, bj2, bj3 Domain basis functions, same layout. + * @param mat_fun Weight function at the quadrature points [flattened point 1, flattened point 2, flattened point 3]. + * @param data StencilMatrix data (rows with ghost regions..., diagonals...); the results are added (not set). + */ +extern "C" __global__ void kernel_3d_mat(Array1D spans1, Array1D spans2, + Array1D spans3, int pi1, int pi2, int pi3, int pj1, int pj2, + int pj3, int starts1, int starts2, int starts3, int pads1, int pads2, + int pads3, Array2D w1, Array2D w2, Array2D w3, + Array4D bi1, Array4D bi2, Array4D bi3, + Array4D bj1, Array4D bj2, Array4D bj3, + Array3D mat_fun, Array6D data) +{ + const long long ne1 = spans1.size(); + const long long ne2 = spans2.size(); + const long long ne3 = spans3.size(); + + const long long ni1 = pi1 + 1, ni2 = pi2 + 1, ni3 = pi3 + 1; + const long long nj1 = pj1 + 1, nj2 = pj2 + 1, nj3 = pj3 + 1; + const long long n_pairs = ni1 * ni2 * ni3 * nj1 * nj2 * nj3; + + CUNUMPY_THREAD_1D(thread, ne1 * ne2 * ne3 * n_pairs); + + // CUDA-only: split the thread index into the pyccel loop variables (iel1, iel2, iel3, il1, il2, il3, jl1, jl2, + // jl3), jl3 fastest + long long rest = thread; + const long long jl3 = rest % nj3; rest /= nj3; + const long long jl2 = rest % nj2; rest /= nj2; + const long long jl1 = rest % nj1; rest /= nj1; + const long long il3 = rest % ni3; rest /= ni3; + const long long il2 = rest % ni2; rest /= ni2; + const long long il1 = rest % ni1; rest /= ni1; + const long long iel3 = rest % ne3; rest /= ne3; + const long long iel2 = rest % ne2; rest /= ne2; + const long long iel1 = rest; + + const long long nq1 = w1.shape[1]; + const long long nq2 = w2.shape[1]; + const long long nq3 = w3.shape[1]; + + double value = 0.0; + + for (long long q1 = 0; q1 < nq1; ++q1) { + for (long long q2 = 0; q2 < nq2; ++q2) { + for (long long q3 = 0; q3 < nq3; ++q3) { + const double wvol = w1(iel1, q1) * w2(iel2, q2) * w3(iel3, q3) * + mat_fun(iel1 * nq1 + q1, iel2 * nq2 + q2, iel3 * nq3 + q3); + + const double bi = bi1(iel1, il1, 0, q1) * bi2(iel2, il2, 0, q2) * bi3(iel3, il3, 0, q3); + const double bj = bj1(iel1, jl1, 0, q1) * bj2(iel2, jl2, 0, q2) * bj3(iel3, jl3, 0, q3); + + value += wvol * bi * bj; + } + } + } + + // global spline indices + const long long i_global1 = spans1(iel1) - pi1 + il1; + const long long i_global2 = spans2(iel2) - pi2 + il2; + const long long i_global3 = spans3(iel3) - pi3 + il3; + + // local spline indices (- starts --> can be negative, will therefore be written to ghost regions) + const long long i_local1 = i_global1 - starts1; + const long long i_local2 = i_global2 - starts2; + const long long i_local3 = i_global3 - starts3; + + cunumpy_atomic_add(&data(pads1 + i_local1, pads2 + i_local2, pads3 + i_local3, pads1 + jl1 - il1, + pads2 + jl2 - il2, pads3 + jl3 - il3), + value); +} diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index f75c6de17..e035e4fe4 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -2242,6 +2242,19 @@ def transpose(self, conjugate=False): return M + def _quad_grid(self, space_name, attribute, component): + """Quadrature-grid arrays (spans, weights or basis values) of one component of a space, as backend arrays. + + The Derham keeps them as host arrays; on the CuPy backend the assembly kernel needs them on the device, + so they are copied there once and cached (they do not change between assemblies). + """ + cache = self.__dict__.setdefault("_quad_grid_cache", {}) + key = (space_name, attribute, component) + if key not in cache: + arrays = getattr(self.derham.spline_attributes[space_name], attribute)[component] + cache[key] = [xp.asarray(array) for array in arrays] + return cache[key] + def assemble(self, weights=None, clear=True): r""" Assembles the weighted mass matrix, i.e. computes the integrals @@ -2320,7 +2333,7 @@ def assemble(self, weights=None, clear=True): # loop over codomain spaces (rows) for a, codomain_space in enumerate(codomain_spaces): # knot span indices of elements of local domain - codomain_spans = spline_attr[W_name].quad_grid_spans[a] + codomain_spans = self._quad_grid(W_name, "quad_grid_spans", a) # global start spline index on process codomain_starts = [int(start) for start in codomain_space.coeff_space.starts] @@ -2332,10 +2345,10 @@ def assemble(self, weights=None, clear=True): pts = self._quad_pts(W_name, a) # global quadrature weights in format (local element, local weight) - wts = spline_attr[W_name].quad_grid_wts[a] + wts = self._quad_grid(W_name, "quad_grid_wts", a) # evaluated basis functions at quadrature points of codomain space - codomain_basis = spline_attr[W_name].quad_grid_bases[a] + codomain_basis = self._quad_grid(W_name, "quad_grid_bases", a) # loop over domain spaces (columns) for b, domain_space in enumerate(domain_spaces): @@ -2375,7 +2388,7 @@ def assemble(self, weights=None, clear=True): ) # evaluated basis functions at quadrature points of domain space - domain_basis = spline_attr[V_name].quad_grid_bases[b] + domain_basis = self._quad_grid(V_name, "quad_grid_bases", b) # assemble matrix (if mat_w is not zero) by calling the appropriate kernel (1d, 2d or 3d) if not_weight_zero or self._is_scalar: diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 9452e81c5..e6f358832 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -243,6 +243,8 @@ def allocate(self): logger.debug("\nAllocating simulation data ...") + self._compile_setup_cuda_kernels() + with ProfileManager.profile_region("setup: allocate"): # feec with ProfileManager.profile_region("allocate feec", functions=[self._allocate_feec]): @@ -280,6 +282,22 @@ def kernels(self) -> tuple[Kernel, ...]: owners += self.model.prop_list return collect_kernels(*owners) + def _compile_setup_cuda_kernels(self) -> None: + """Compile the CUDA kernels that :meth:`allocate` itself calls (mass-matrix assembly) before it calls them. + + :meth:`compile_cuda_kernels` runs after :meth:`allocate`, so without this the first assembly would include + the compilation (~1 s per kernel on a node whose CuPy cache does not have it yet), in the assembly's + profiling regions. Here it is counted in ``"setup: compile cuda kernels (assembly)"``. Nothing happens on + the NumPy backend. + """ + if xp.get_backend() != "cupy": + return + + from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat + + with ProfileManager.profile_region("setup: compile cuda kernels (assembly)"): + kernel_3d_mat.compile() + def compile_cuda_kernels(self) -> tuple[Kernel, ...]: """Compile the CUDA versions of the kernels of :meth:`kernels` now, before the time stepping. From 540cb59ded4200d05be2d5ff9871c34f3516ec4e Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:26:26 +0200 Subject: [PATCH 52/55] Parity cases for kernel_3d_mat Random quadrature data on 3 x 2 x 2 elements with mixed degrees, with the first row local or a ghost row. test_cuda_kernels_have_parity_cases requires a case for every kernel with a CUDA version; the emulation test compares it with the pyccel kernel without a GPU. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/cuda_parity_cases.py | 26 ++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index a2396ee5b..6e10707be 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -315,6 +315,30 @@ def kernel_pullpush_args(case): ) +def kernel_3d_mat_threads(args): + """One thread per (element, codomain basis function, domain basis function), as the kernel's launch size.""" + from struphy.feec.kernels.kernel_3d_mat import _n_threads + + return _n_threads(args) + + +def kernel_3d_mat_args(case): + """Mass-matrix assembly on 3 x 2 x 2 elements with mixed degrees and random quadrature data; `case` is the start + index of the rows in direction 1 (1: the first row is a ghost row, local index -1).""" + rng = np.random.default_rng(20 + case) + ne, nq = (3, 2, 2), (3, 2, 2) + pi, pj, pads = (2, 1, 1), (1, 2, 1), (2, 2, 2) + starts = (case, 0, 0) + spans = [xp.asarray(np.arange(n, dtype=np.int64) + p) for n, p in zip(ne, pi)] + w = [xp.asarray(rng.random((n, q))) for n, q in zip(ne, nq)] + bi = [xp.asarray(rng.random((n, p + 1, 1, q))) for n, p, q in zip(ne, pi, nq)] + bj = [xp.asarray(rng.random((n, p + 1, 1, q))) for n, p, q in zip(ne, pj, nq)] + mat_fun = xp.asarray(rng.normal(size=tuple(n * q for n, q in zip(ne, nq)))) + rows = tuple(n + p + 2 * pad for n, p, pad in zip(ne, pi, pads)) + data = xp.asarray(rng.normal(size=rows + tuple(2 * pad + 1 for pad in pads))) # added to, not overwritten + return (*spans, *pi, *pj, *starts, *pads, *w, *bi, *bj, mat_fun, data) + + # ---------------------------------------------------------------- all kernels with a CUDA version PUSHER_TOLERANCES = {"rtol": 1e-13, "atol": 1e-14} @@ -338,6 +362,8 @@ def size_of(index, per_thread=1): PARITY_CASES = { + # atomic additions: the summation order over elements differs from pyccel + "kernel_3d_mat": ParityCases((0, 1), kernel_3d_mat_args, rtol=1e-12, atol=1e-13, n_threads=kernel_3d_mat_threads), "push_eta_stage": ParityCases( [ (bc, method, stage) From b0e73eb19601a1d8628e3ad8fff12035b5ae9672 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:08:01 +0200 Subject: [PATCH 53/55] Apply separable mass matrices as Kronecker products of 1D banded matrices WeightedMassOperator.dot applies the blocks of its matrix that are Kronecker products of 1D banded matrices (constant weights, e.g. on a Cuboid) as three banded passes (struphy.feec.kronecker_mass, CUDA kernel kron_band_pass on the GPU) instead of the 3D stencil product. The replacement is built on first use and invalidated by assembly, copies and in-place arithmetic; a matrix changed by arithmetic that is not a Kronecker product is not checked again. STRUPHY_KRONECKER_MASS=0 switches it off. The boundary and extraction operators around the matrix are skipped when they are identities instead of being applied as copies. Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/kron_band_pass_cuda.cu | 35 ++++ src/struphy/feec/kronecker_mass.py | 247 ++++++++++++++++++++++++ src/struphy/feec/mass.py | 85 ++++++-- 3 files changed, 355 insertions(+), 12 deletions(-) create mode 100644 src/struphy/feec/kron_band_pass_cuda.cu create mode 100644 src/struphy/feec/kronecker_mass.py diff --git a/src/struphy/feec/kron_band_pass_cuda.cu b/src/struphy/feec/kron_band_pass_cuda.cu new file mode 100644 index 000000000..745e782f4 --- /dev/null +++ b/src/struphy/feec/kron_band_pass_cuda.cu @@ -0,0 +1,35 @@ +/** + * One pass of a Kronecker product of 1D banded matrices (struphy.feec.kronecker_mass._band_pass): + * + * out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...] along `axis`, + * + * with the other two axes unchanged (out and inp have the same extent along them). One thread per entry of `out`, + * the last axis fastest, so that neighbouring threads read neighbouring entries of `inp` along every axis. `out` may + * be a strided view (the owned rows of a stencil vector's data). + * + * @param A Banded 1D matrix, shape (out.shape[axis], n_diagonals). + * @param inp Input array (3D). + * @param out Output array (3D, any strides); written, not added to. + * @param axis Axis of the pass (0, 1 or 2). + * @param offset Index in `inp` of diagonal 0 of row 0 along `axis`. + */ +#include "cunumpy/array_view.cuh" + +extern "C" __global__ void kron_band_pass(Array2D A, Array3D inp, Array3D out, int axis, + long long offset) { + const long long n1 = out.shape[1], n2 = out.shape[2]; + const long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; + if (t >= out.shape[0] * n1 * n2) return; + const long long i2 = t % n2; + const long long i1 = (t / n2) % n1; + const long long i0 = t / (n1 * n2); + const long long i = axis == 0 ? i0 : (axis == 1 ? i1 : i2); + + double v = 0.0; + for (long long d = 0; d < A.shape[1]; ++d) { + const long long j = i + offset + d; + const double x = axis == 0 ? inp(j, i1, i2) : (axis == 1 ? inp(i0, j, i2) : inp(i0, i1, j)); + v += A(i, d) * x; + } + out(i0, i1, i2) = v; +} diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py new file mode 100644 index 000000000..2050edb82 --- /dev/null +++ b/src/struphy/feec/kronecker_mass.py @@ -0,0 +1,247 @@ +"""Fast application of assembled mass matrices that are Kronecker products of 1D banded matrices. + +A mass matrix block ``M[i, j] = int w Lambda_i Lambda_j`` of tensor-product splines factorizes as +``M = c * M_1 (x) M_2 (x) M_3`` when its weight ``w`` is constant (e.g. on a Cuboid). Applied as a +:class:`~feectools.linalg.stencil.StencilMatrix`, every row reads all ``n_1 * n_2 * n_3`` diagonals (105 for +degrees (1, 2, 3)) from memory; as a Kronecker product it is three 1D banded passes over the vector, ``n_1 + n_2 + +n_3`` operations per entry and no 3D matrix data. :func:`kronecker_apply_operator` checks each assembled block and +returns an operator for the matrix-vector products in which the blocks that are Kronecker products (to round-off) +are replaced by :class:`KroneckerBandMatrix`. + +The 1D factors are read from the assembled data itself (rows owned by this process), so they hold exactly what the +assembly computed, on any domain decomposition, and the check is the assembled data, not an assumption about the +mapping or the weights. +""" + +import logging +import os +from functools import cache +from pathlib import Path + +import cunumpy as xp +from feectools.linalg.basic import LinearOperator +from feectools.linalg.block import BlockLinearOperator +from feectools.linalg.stencil import StencilMatrix, StencilVector + +logger = logging.getLogger("struphy") + +# relative tolerance of the Kronecker check: the 3D assembly and the product of the factors differ by round-off only +KRONECKER_RTOL = 1e-12 + + +def kronecker_mass_enabled() -> bool: + """Whether :func:`kronecker_apply_operator` may replace blocks (``STRUPHY_KRONECKER_MASS=0`` switches it off).""" + return os.environ.get("STRUPHY_KRONECKER_MASS", "1").lower() not in ("0", "false", "no") + + +@cache +def _band_pass_cuda(): + """The CUDA kernel of one banded pass (compiled on first call).""" + from cunumpy.kernels import CudaKernel + + return CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass") + + +def _band_pass(A, inp, out, axis, offset): + """``out[..., i, ...] = sum_d A[i, d] * inp[..., i + offset + d, ...]`` along ``axis`` (other axes unchanged).""" + if xp.is_gpu(out): + _band_pass_cuda()(A, inp, out, axis, offset, n_threads=out.size) + return + m = out.shape[axis] + shape = [1, 1, 1] + shape[axis] = m + out[...] = 0.0 + for d in range(A.shape[1]): + idx = [slice(None)] * 3 + idx[axis] = slice(offset + d, offset + d + m) + out += A[:, d].reshape(shape) * inp[tuple(idx)] + + +class KroneckerBandMatrix(LinearOperator): + """A square 3D :class:`~feectools.linalg.stencil.StencilMatrix` applied as a Kronecker product of 1D banded + matrices, ``(A_1 (x) A_2 (x) A_3) x`` on the rows owned by this process. + + ``A_k[i, d]`` is the entry of owned row ``i`` (along direction ``k``) and diagonal ``d``, column + ``i - q_k + d`` with ``q_k = (n_k - 1) // 2``, as in the stencil matrix. The product is computed in three passes, + each over the owned rows of its direction and the full (padded) extent of the directions after it, so the passes + need no communication beyond the ghost regions of the input. Use :meth:`from_stencil`. + + Parameters + ---------- + stencil : StencilMatrix + The assembled matrix (kept for ``toarray``, ``tosparse`` and ``transpose``). + factors : tuple of arrays + ``A_1, A_2, A_3``, shapes ``(m_k, n_k)``, C-contiguous backend arrays. + """ + + def __init__(self, stencil: StencilMatrix, factors): + self._stencil = stencil + self._factors = factors + + V = stencil.domain + self._m = tuple(int(e - s + 1) for s, e in zip(V.starts, V.ends)) + self._pads = tuple(int(p) for p in V.pads) + self._q = tuple((int(n) - 1) // 2 for n in stencil._data.shape[3:]) + + # temporaries of the first two passes: owned rows along the passes done, padded extent along the others + x_shape = V.zeros()._data.shape + self._t1 = xp.empty((self._m[0], x_shape[1], x_shape[2]), dtype=stencil._data.dtype) + self._t2 = xp.empty((self._m[0], self._m[1], x_shape[2]), dtype=stencil._data.dtype) + + @classmethod + def from_stencil(cls, S: StencilMatrix, rtol: float = KRONECKER_RTOL): + """The Kronecker form of ``S`` if it is one (to ``rtol``, relative to its largest entry), else None. + + Only square 3D matrices on one space (domain is codomain) with shifts 1 qualify, e.g. the diagonal blocks of + the mass matrices. The factors are read from the owned rows of ``S``: with reference rows ``r_k`` (the first + owned row along each direction) and central diagonals ``q_k``, + ``A_1[i, :] = S[i, r_2, r_3, :, q_2, q_3]``, ``A_2[i, :] = S[r_1, i, r_3, q_1, :, q_3] / S[r, q]``, + ``A_3[i, :] = S[r_1, r_2, i, q_1, q_2, :] / S[r, q]``; the product is then compared with all owned rows. + """ + if not isinstance(S, StencilMatrix) or S.domain is not S.codomain or S.domain.ndim != 3: + return None + V = S.domain + if any(int(s) != 1 for s in V.shifts) or S._data.ndim != 6: + return None + + p = [int(p) for p in V.pads] + m = [int(e - s + 1) for s, e in zip(V.starts, V.ends)] + q = [(int(n) - 1) // 2 for n in S._data.shape[3:]] + if min(m) < 1: + return None + data = S._data + + den = data[p[0], p[1], p[2], q[0], q[1], q[2]] + if float(den) == 0.0: + return None + A1 = xp.ascontiguousarray(data[p[0] : p[0] + m[0], p[1], p[2], :, q[1], q[2]]) + A2 = xp.ascontiguousarray(data[p[0], p[1] : p[1] + m[1], p[2], q[0], :, q[2]] / den) + A3 = xp.ascontiguousarray(data[p[0], p[1], p[2] : p[2] + m[2], q[0], q[1], :] / den) + + # compare with every owned row: planes (rows along direction 1) in chunks, the maxima kept on the device + # (one read at the end); a few planes first, which rejects most matrices that are not Kronecker products + # (e.g. ones assembled from particle data, checked again after every change) at once + owned = data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + A23 = A2[:, None, None, :, None] * A3[None, :, None, None, :] # (m2, m3, 1, n2, n3) + + def max_diff_and_abs(rows): + block = owned[rows] + recon = A1[rows][:, None, None, :, None, None] * A23[None] + return xp.max(xp.abs(block - recon)), xp.max(xp.abs(block)) + + def fails(max_diff, max_abs): + max_diff, max_abs = float(max_diff), float(max_abs) + return max_abs == 0.0 or max_diff > rtol * max_abs + + quick = sorted({0, m[0] // 2, m[0] - 1}) + if fails(*max_diff_and_abs(xp.asarray(quick))): + return None + + # chunks of at most ~64 MB of temporaries + plane = max(1, owned[0].size) + chunk = max(1, min(m[0], (8 << 20) // plane)) + max_diff = xp.zeros(()) + max_abs = xp.zeros(()) + for start in range(0, m[0], chunk): + d, a = max_diff_and_abs(slice(start, min(start + chunk, m[0]))) + max_diff = xp.maximum(max_diff, d) + max_abs = xp.maximum(max_abs, a) + if fails(max_diff, max_abs): + return None + return cls(S, (A1, A2, A3)) + + # ------------------------------------------------------------------ + # LinearOperator interface + # ------------------------------------------------------------------ + @property + def domain(self): + return self._stencil.domain + + @property + def codomain(self): + return self._stencil.codomain + + @property + def dtype(self): + return self._stencil.dtype + + @property + def factors(self): + """The 1D factors ``A_1, A_2, A_3`` (rows owned by this process).""" + return self._factors + + def tosparse(self): + return self._stencil.tosparse() + + def toarray(self): + return self._stencil.toarray() + + def transpose(self, conjugate=False): + return self._stencil.transpose(conjugate=conjugate) + + def dot(self, v, out=None): + assert isinstance(v, StencilVector) + assert v.space is self.domain + if out is None: + out = self.codomain.zeros() + else: + assert isinstance(out, StencilVector) + assert out.space is self.codomain + + if not v.ghost_regions_in_sync: + v.update_ghost_regions() + + A1, A2, A3 = self._factors + p, q, m = self._pads, self._q, self._m + owned = out._data[p[0] : p[0] + m[0], p[1] : p[1] + m[1], p[2] : p[2] + m[2]] + + _band_pass(A1, v._data, self._t1, 0, p[0] - q[0]) + _band_pass(A2, self._t1, self._t2, 1, p[1] - q[1]) + _band_pass(A3, self._t2, owned, 2, p[2] - q[2]) + + out.ghost_regions_in_sync = False + return out + + +def kronecker_apply_operator(mat, name: str = ""): + """The operator to apply ``mat`` with: ``mat`` itself, or a copy whose Kronecker blocks are replaced. + + Parameters + ---------- + mat : StencilMatrix or BlockLinearOperator + An assembled (block) mass matrix. + name : str + Name of the operator, for the log. + + Returns + ------- + LinearOperator + ``mat`` if no block is a Kronecker product (or :func:`kronecker_mass_enabled` is False); otherwise a + :class:`KroneckerBandMatrix` (for a StencilMatrix) or a BlockLinearOperator with the same blocks except the + Kronecker ones. It is valid only as long as ``mat`` is not modified. + """ + if not kronecker_mass_enabled(): + return mat + + if isinstance(mat, StencilMatrix): + kron = KroneckerBandMatrix.from_stencil(mat) + if kron is not None: + logger.info(f"Mass matrix {name}: applied as a Kronecker product of 1D banded matrices") + return mat if kron is None else kron + + if isinstance(mat, BlockLinearOperator): + blocks = {} + replaced = [] + for key, block in mat._blocks.items(): + kron = KroneckerBandMatrix.from_stencil(block) + if kron is None: + blocks[key] = block + else: + blocks[key] = kron + replaced.append(key) + if not replaced: + return mat + logger.info(f"Mass matrix {name}: blocks {replaced} applied as Kronecker products of 1D banded matrices") + return BlockLinearOperator(mat.domain, mat.codomain, blocks=blocks) + + return mat diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index e035e4fe4..1020f0ef6 100644 --- a/src/struphy/feec/mass.py +++ b/src/struphy/feec/mass.py @@ -25,6 +25,7 @@ from struphy.feec.kernels.kernel_3d_mat import kernel_3d_mat from struphy.feec.kernels.kernel_3d_matrixfree import kernel_3d_matrixfree from struphy.feec.kernels.kernel_3d_vec import kernel_3d_vec +from struphy.feec.kronecker_mass import kronecker_apply_operator from struphy.feec.linear_operators import BoundaryOperator from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.feec.utilities import LocalProjectionMatrix, LocalRotationMatrix, get_quad_grids @@ -1703,6 +1704,8 @@ def __init__( blocks, self._weights = self._init_blocks_from_weights(weights_info) self._mat = self._wrap_blocks(blocks) + # the operator dot() applies self._mat with (Kronecker blocks replaced), built on first use, see _apply_mat + self._apply_mat_cache = None # the weights are stored in the block order of the (possibly transposed) operator if transposed: @@ -2121,6 +2124,32 @@ def M0(self): def matrix(self): return self._mat + @property + def _apply_mat(self): + """The operator :meth:`dot` applies: :attr:`matrix`, with its blocks that are Kronecker products of 1D banded + matrices (constant weights, e.g. on a Cuboid) applied as such (:mod:`struphy.feec.kronecker_mass`). + + Built on first use after every change of the matrix (assembly, copy into, in-place arithmetic). + """ + if self._apply_mat_cache is None: + if getattr(self, "_kronecker_rejected", False): + self._apply_mat_cache = self._mat + else: + self._apply_mat_cache = kronecker_apply_operator(self._mat, name=self.name) + # a matrix changed by in-place arithmetic that is not a Kronecker product (e.g. the Schur complement + # M1 - dt^2 BC with particle data, rebuilt every step) is not checked again after further arithmetic: + # the check reads the whole matrix. Assembly and copies check again. + if getattr(self, "_changed_by_arithmetic", False) and self._apply_mat_cache is self._mat: + self._kronecker_rejected = True + return self._apply_mat_cache + + def _matrix_changed(self, by_arithmetic: bool): + """Invalidate :attr:`_apply_mat` after a change of the matrix (in-place arithmetic, or assembly/copy).""" + self._apply_mat_cache = None + self._changed_by_arithmetic = by_arithmetic + if not by_arithmetic: + self._kronecker_rejected = False + @property def nquads(self): if self._nquads is None: @@ -2168,29 +2197,54 @@ def dot(self, v, out=None, apply_bc=True): if apply_bc: if self._transposed: - self._W_boundary_op_T.dot(v, out=self._temp_WB) - self._W_extraction_op_T.dot(self._temp_WB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) - self._V_extraction_op.dot(self._temp_VE, out=self._temp_VB) - out = self._V_boundary_op.dot(self._temp_VB, out=out) + ops_in = (self._W_boundary_op_T, self._W_extraction_op_T) + ops_out = (self._V_extraction_op, self._V_boundary_op) + tmps_in = (self._temp_WB, self._temp_mat) + tmps_out = (self._temp_VE, self._temp_VB) else: - self._V_boundary_op_T.dot(v, out=self._temp_VB) - self._V_extraction_op_T.dot(self._temp_VB, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) - self._W_extraction_op.dot(self._temp_WE, out=self._temp_WB) - out = self._W_boundary_op.dot(self._temp_WB, out=out) + ops_in = (self._V_boundary_op_T, self._V_extraction_op_T) + ops_out = (self._W_extraction_op, self._W_boundary_op) + tmps_in = (self._temp_VB, self._temp_mat) + tmps_out = (self._temp_WE, self._temp_WB) + out = self._apply_chain(v, ops_in, tmps_in, ops_out, tmps_out, out) else: if self._transposed: self._W_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_VE) + self._apply_mat.dot(self._temp_mat, out=self._temp_VE) out = self._V_extraction_op.dot(self._temp_VE, out=out) else: self._V_extraction_op_T.dot(v, out=self._temp_mat) - self._mat.dot(self._temp_mat, out=self._temp_WE) + self._apply_mat.dot(self._temp_mat, out=self._temp_WE) out = self._W_extraction_op.dot(self._temp_WE, out=out) return out + def _apply_chain(self, v, ops_in, tmps_in, ops_out, tmps_out, out): + """``out = B_out E_out M E_in B_in v`` (boundary and extraction operators around the matrix), see :meth:`dot`. + + Identity operators (e.g. the extraction operators without polar splines, or boundary operators without + essential boundary conditions) are skipped instead of applied as copies: each vector is passed on as it is, + and the matrix writes into the next temporary that is actually used (or into ``out``). + """ + if out is v: + # in place: the chain must not read v after writing the output, so it reads a copy + v = v.copy() + + x = v + for op, tmp in zip(ops_in, tmps_in): + if not isinstance(op, IdentityOperator): + x = op.dot(x, out=tmp) + + # the first output operator that is not an identity reads the matrix product from its temporary + remaining = [(op, tmp) for op, tmp in zip(ops_out, tmps_out) if not isinstance(op, IdentityOperator)] + if not remaining: + return self._apply_mat.dot(x, out=out) + y = self._apply_mat.dot(x, out=remaining[0][1]) + for k, (op, _) in enumerate(remaining): + target = out if k == len(remaining) - 1 else remaining[k + 1][1] + y = op.dot(y, out=target) + return y + def transpose(self, conjugate=False): """ Returns the transposed operator. @@ -2277,6 +2331,7 @@ def assemble(self, weights=None, clear=True): Whether to first set all data to zero before assembly. If False, the new contributions are added to existing ones. """ + self._matrix_changed(by_arithmetic=False) # the matrix changes assert not self._dry_run, ( "A dry-run operator has no matrix data and cannot be assembled (memory estimation only)." ) @@ -2559,6 +2614,7 @@ def copy(self, out=None): out._weights = [list(row) for row in self._weights] self._mat.copy(out=out._mat) + out._matrix_changed(by_arithmetic=False) if self._creation_info is None: out._creation_info = None @@ -2569,6 +2625,7 @@ def copy(self, out=None): def __imul__(self, a): self._mat *= a + self._matrix_changed(by_arithmetic=True) self._creation_info = None return self @@ -2579,11 +2636,13 @@ def __iadd__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat += M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat += M + self._matrix_changed(by_arithmetic=True) return self else: @@ -2596,11 +2655,13 @@ def __isub__(self, M): if isinstance(M, WeightedMassOperator): self._creation_info = None self._mat -= M._mat + self._matrix_changed(by_arithmetic=True) return self elif isinstance(M, LinearOperator): self._creation_info = None self._mat -= M + self._matrix_changed(by_arithmetic=True) return self else: From a5a56e315eda78c56a104e91ef946e319cf67529 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:20:50 +0200 Subject: [PATCH 54/55] Kronecker mass: launch the band-pass kernel through CachedLaunch The banded passes run with the same buffers in every solver iteration (as on cuda-aware-mpi-fixes, now that feectools has CachedLaunch). Co-Authored-By: Claude Opus 5.5 --- src/struphy/feec/kronecker_mass.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/feec/kronecker_mass.py b/src/struphy/feec/kronecker_mass.py index 2050edb82..5456b2eeb 100644 --- a/src/struphy/feec/kronecker_mass.py +++ b/src/struphy/feec/kronecker_mass.py @@ -38,8 +38,10 @@ def kronecker_mass_enabled() -> bool: def _band_pass_cuda(): """The CUDA kernel of one banded pass (compiled on first call).""" from cunumpy.kernels import CudaKernel + from feectools.utilities.cached_launch import CachedLaunch - return CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass") + # launched with the same buffers in every solver iteration: cache the launch preparation + return CachedLaunch(CudaKernel.from_file(Path(__file__).with_name("kron_band_pass_cuda.cu"), "kron_band_pass")) def _band_pass(A, inp, out, axis, offset): From 37ad300769f55d29849a565039144e9bc29d1f42 Mon Sep 17 00:00:00 2001 From: Max Date: Tue, 27 Oct 2026 15:41:47 +0100 Subject: [PATCH 55/55] Evaluate the control variate at f_coords in update_weights --- src/struphy/pic/base.py | 8 +++++++- src/struphy/pic/particles.py | 4 ++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 1839c06aa..0e30c2d5e 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1555,6 +1555,12 @@ def initialize_weights( else: self.weights = self.weights0 + def _f_coords_of_valid(self, quantity, rows): + """Coordinates of the distribution function at the valid markers, as a list of arrays. + ``quantity(j)`` gives row ``j`` of the marker buffer; subclasses whose :attr:`f_coords` differs + from the stored rows override this.""" + return [quantity(j) for j in rows] + @profile def update_weights(self): """ @@ -1583,7 +1589,7 @@ def quantity(j): # self.save_constants_of_motion() coords = self.f_coords_index rows = range(markers.shape[0])[coords] if isinstance(coords, slice) else np.atleast_1d(coords).tolist() - f0 = self.f0(*[quantity(j) for j in rows]) + f0 = self.f0(*self._f_coords_of_valid(quantity, rows)) # if f_init is vol-form, transform to 0-form if self.is_volume_form[0]: diff --git a/src/struphy/pic/particles.py b/src/struphy/pic/particles.py index 00854437c..4691e6547 100644 --- a/src/struphy/pic/particles.py +++ b/src/struphy/pic/particles.py @@ -395,6 +395,10 @@ def f_coords(self): def f_coords(self, new): Particles6D.f_coords.fset(self, new) + def _f_coords_of_valid(self, quantity, rows): + # the control variate is evaluated at the gyrotropic velocity of f_coords, not at the stored v_par b0 + return list(self.f_coords) + def _maxwellians_in(f): """The Maxwellians of a kinetic background (summands of a :class:`~struphy.kinetic_background.base.SumKineticBackground`)."""