GPU stack 3/8: CUDA mass-matrix assembly kernel kernel_3d_mat

src/struphy/feec/kernels/kernel_3d_mat/__init__.py

@@ -1,8 +1,4 @@
-"""The kernel of this folder: pyccel in ``<name>_kernels.py``, CUDA in ``<name>_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
 )

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<long long> spans1, Array1D<long long> spans2,
+                                         Array1D<long long> 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<double> w1, Array2D<double> w2, Array2D<double> w3,
+                                         Array4D<double> bi1, Array4D<double> bi2, Array4D<double> bi3,
+                                         Array4D<double> bj1, Array4D<double> bj2, Array4D<double> bj3,
+                                         Array3D<double> mat_fun, Array6D<double> 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);
+}

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:

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)

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.