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/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) 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.