From 607cab911c2fd33f9153d77afec9e998ad2a21d5 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:00:05 +0200 Subject: [PATCH 01/33] 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 02/33] 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 03/33] 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 04/33] 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 05/33] 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 c45019c23c48916b7d85a0448da6210270e2e496 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:26:52 +0200 Subject: [PATCH 06/33] 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 133be44595ac33a3f987b6de8603e0e1932a2ed2 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:04:40 +0200 Subject: [PATCH 07/33] 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 c508cce8bf3202436abdac637029b614b9f248a9 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:20:50 +0200 Subject: [PATCH 08/33] 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 6bb532fd1b63bd29166b84ffeb7d390d21fa10c9 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:00:05 +0200 Subject: [PATCH 09/33] 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 10/33] 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 11/33] 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 12/33] 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 13/33] 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 14/33] 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 15/33] 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 16/33] 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 17/33] 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 18/33] 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 f1cbb2d1200ce79c864c81dd4ef04280c85c198b Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:00:05 +0200 Subject: [PATCH 19/33] 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 20/33] 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 21/33] 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 22/33] 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 23/33] 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 f63a56031f63bcc4d3a3065cc4c42205b12d30c5 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 10:00:05 +0200 Subject: [PATCH 24/33] 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 25/33] 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 26/33] 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 27/33] 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 28/33] 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 29/33] 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 30/33] 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 31/33] 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 32/33] 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 33/33] 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):