Stencil 3D device kernels: 6D matrix views instead of the 2p+1 assumption

CUDA_STRATEGY.md

@@ -48,7 +48,7 @@ come from
 - **No silent CPU fallback** for folder kernels: on the CuPy backend a folder kernel without CUDA version raises.
   The other pyccel kernels (B-splines, field evaluation, DOF kernels, used at setup) are wrapped in
   `cunumpy.kernels.PyccelKernel` and copy their arrays to the host and back; see [Open questions](#open-questions).
-- **Same cunumpy as struphy.** `cunumpy >= 0.5.0, < 0.6`; names are imported from the submodules
+- **Same cunumpy as struphy.** `cunumpy >= 0.6.1` (the 3D stencil kernels need its 6D array views); names are imported from the submodules
   (`cunumpy.kernels`, `cunumpy.arguments`, `cunumpy.cuda`, `cunumpy.mpi`, `cunumpy.kernel_testing`); cunumpy 0.6 removed
   the old top-level names, and `CudaKernel`/`CudaKernelVariants` are in `cunumpy.kernels` since 0.6.
 - **Small steps.** Every PR keeps the NumPy path working and tested.
@@ -117,16 +117,11 @@ CuPy. There are no backend branches and no host staging at these call sites any
   of generated source is gone.
 - **Signature changes.** The inner kernels add their result to an argument `res` (`res[0] += ...`, the caller
   zeroes it; `inner` passes the MPI send buffer) instead of returning it, because a CUDA kernel cannot return a
-  value; on the GPU each block reduces in shared memory and adds its sum with one `atomicAdd`. The transpose
-  kernels take `e_in` (end of the row range of `mat`) after `s_in`; pyccel does not need it, the 3D CUDA kernel
-  needs it for the shape of `mat`.
-- **6D matrix data.** cunumpy's array views stop at `Array4D`, so the 3D kernels take the six-axis matrix data as
-  raw pointers (C-contiguity checked by `CudaKernel`) and derive their shape from the other arguments: rows from
-  `out` (dot) or from `s/e/p` (transpose), and `2 * p + 1` diagonals. This is exactly what the precompiled pyccel
-  kernels read, and they are wrong for other data too, so `StencilMatrix.set_backend` records whether the matrix
-  has no shifts and `2 * p + 1` diagonals (`_kernel_shapes_ok`), and `dot`, `vdot` and `transpose` raise
-  `NotImplementedError` otherwise, on both backends (no test or call site in the test suite builds such a matrix
-  for these kernels). Array views up to 6D in cunumpy would remove this restriction (see [Open questions](#open-questions)).
+  value; on the GPU each block reduces in shared memory and adds its sum with one `atomicAdd`. (The transpose
+  kernels had an argument `e_in` for the shape of the 3D matrix data; it is gone, see the 6D views below.)
+- **6D matrix data:** see [6D matrix views](#6d-matrix-views-stencil-6d-views). (In #88 the 3D kernels took the
+  six-axis data as raw pointers and assumed `2 * p + 1` diagonals; `dot`, `vdot` and `transpose` raised for other
+  matrices.)
 - **Launch sizes** are declared in each folder's `__init__.py` (`n_threads_from`): one thread per entry of `out`,
   `matT`, `v1` or `x`, the last axis varying fastest; threads outside the owned rows or diagonals return without
   writing, as the pyccel loops do. The inner kernels use blocks of 256 threads.
@@ -168,6 +163,37 @@ every CG iteration preconditioned by struphy's `MassMatrixPreconditioner` (strup
   device against host solves under `assert_no_transfers`, serial and with MPI. Before this change a 3D solve on
   the fake CuPy made 3 `to_host` copies (one per direction) and 3 uncounted copies back.
 
+## 6D matrix views (`stencil-6d-views`)
+
+With cunumpy 0.6.1 (`Array5D`/`Array6D`, `CArray5D`/`CArray6D` in `cunumpy/array_view.cuh`) the 3D kernels take
+the matrix data as views, like the 1D/2D kernels (`Array2D`/`Array4D`):
+
+- `stencil_dot_3d(Array6D<double> mat, Array3D<double> x, Array3D<double> out, s_in, p_in, add, s_out, e_out,
+  p_out)` and `stencil_transpose_3d(Array6D<double> mat, Array6D<double> matT, s_in, p_in, add, s_out, e_out,
+  p_out)`; the pyccel versions take `float[:, :, :, :, :, :]` with the same arguments in the same order. Strided
+  views (`Array6D`, not `CArray6D`), as in 1D/2D, so a non-contiguous matrix is not refused.
+- **No `e_in`.** The transpose kernels (all dimensions, the argument list is the same for all) lost `e_in`: it was
+  only needed for the row extents of the raw pointer in 3D.
+- **The number of diagonals comes from the data.** All six kernels (1D, 2D, 3D, pyccel and CUDA) read the number
+  of diagonals `n_k` of each direction from the shape of the matrix data instead of assuming `2 * p_in + 1`: the
+  pads of the matrix are `q_k = (n_k - 1) // 2` (at most the pads of the spaces, `StencilMatrix(V, W, pads=...)`),
+  diagonal `d` of row `i` is the column `i - q + d` (in `x` at `i - q + d - s_in + p_in`), and the last owned row
+  uses `n_k - 1 + add[k]` diagonals. With `q = p` this is the old loop, in the same order and with bitwise the same
+  results (checked against the kernels of #88 run as Python); the pyccel kernels are now one loop nest with the
+  number of diagonals picked per row instead of the eight (3D) spelled-out combinations.
+- **`StencilMatrix`.** `dot`, `vdot` and `transpose` no longer raise for matrices with fewer diagonals (blocks
+  between spaces of different degree, derivative-type stencils); `transpose(out=...)` asserts that `out` has the
+  pads of the matrix. The one restriction left is spaces with **shifts > 1**, which still raise
+  `NotImplementedError` (`_check_kernel_shifts`): the stencil kernels have never handled them, and there is no
+  reference to match, since the general psydac kernels (`matvec_<n>d`, `transpose_<n>d`) disagree with `toarray()`
+  and with each other (the transpose is not the adjoint of the product) for shifts > 1, and psydac never tested
+  them ("TODO: verify for s>1").
+- **Tests.** `MATRIX_CASES` (parity and CPU emulation) has matrices with fewer diagonals and non-periodic
+  rectangular blocks between spaces of different size per direction, including pads 0 (one diagonal);
+  `test_device_matvec.py` compares `dot`, `vdot` and `transpose` of such matrices with dense `toarray()`
+  references, and `test_mpi_device.py` checks a matrix with fewer diagonals and the adjoint identity of its
+  transpose against the global field on any number of ranks.
+
 ## Kernel folders
 
 ```
@@ -214,8 +240,9 @@ Conventions, as in struphy:
   `PyccelKernel` (host copies). They run at setup, not in the time loop; they move into kernel folders with CUDA
   versions when a profile shows they matter.
 - **GPU CI.** No GPU runner yet; GPU tests are run by hand.
-- **6D array views in cunumpy.** With `Array5D`/`Array6D` (also needed by struphy's matrix accumulations), the 3D
-  kernels could take the matrix data as views, drop the shape assumptions and the `e_in` argument.
+- **Shifts > 1.** `StencilMatrix.dot`/`transpose` raise for spaces with shifts > 1 (see
+  [6D matrix views](#6d-matrix-views-stencil-6d-views)); supporting them needs a verified definition of the data
+  layout first.
 - **Complex data on the device.** Not needed by struphy so far; would need a second CUDA kernel per folder (or
   dtype dispatch in `cunumpy.kernels.Kernel`).
 - **`inner` reduction.** One `atomicAdd` per block of 256 threads, then a copy of the 8-byte result to the host in

feectools/linalg/kernels/stencil_dot_1d/stencil_dot_1d_cuda.cu

@@ -7,15 +7,16 @@
  *
  * One thread per entry of `out` (n_threads = out.size, see __init__.py). A thread outside the owned rows
  * (local row index i1_loc outside [0, e_out - s_out]) returns without writing, so the padding of `out` is
- * left as it is, as in pyccel. Interior rows use 2 * p_in + 1 diagonals, the last owned row (i1 == e_out)
- * uses 2 * p_in + add, which is how a rectangular matrix is handled.
+ * left as it is, as in pyccel. The matrix has n = mat.shape[1] diagonals (its pads are q = (n - 1) / 2) and
+ * diagonal d1 of row i1 is the column i1 - q + d1. Interior rows use all n diagonals, the last owned row
+ * (i1 == e_out) uses n - 1 + add, which is how a rectangular matrix is handled.
  *
- * @param mat   matrix data, shape (rows of `out`, 2 * p_in + 1)
+ * @param mat   matrix data, shape (rows of `out`, diagonals)
  * @param x     data of the domain vector, ghost regions included
  * @param out   data of the codomain vector; the owned rows are written
  * @param s_in  global start of the domain of this process
  * @param p_in  padding of the domain
- * @param add   1 if the last row uses all 2 * p_in + 1 diagonals, else 0
+ * @param add   1 if the last row uses all diagonals, else 0
  * @param s_out global start of the codomain of this process
  * @param e_out global end (inclusive) of the codomain of this process
  * @param p_out padding of the codomain: the owned rows start at index p_out of `mat` and `out`
@@ -30,11 +31,13 @@ extern "C" __global__ void stencil_dot_1d(Array2D<double> mat, Array1D<double> x
     if (i1_loc < 0 || i1_loc > e_out - s_out) return;
     const long long i1 = s_out + i1_loc;  // global row index
 
-    const long long n_diags1 = (i1 == e_out) ? 2 * p_in + add : 2 * p_in + 1;
+    const long long n_diags1 = mat.shape[1];
+    const long long nd1 = (i1 == e_out) ? n_diags1 - 1 + add : n_diags1;
+    const long long off1 = p_in - (n_diags1 - 1) / 2 - s_in;  // x index of diagonal 0 minus i1
 
     double val = 0.;
-    for (long long d1 = 0; d1 < n_diags1; ++d1)
-        val += mat(p_out + i1_loc, d1) * x(i1 + d1 - s_in);
+    for (long long d1 = 0; d1 < nd1; ++d1)
+        val += mat(p_out + i1_loc, d1) * x(i1 + d1 + off1);
 
     out(p_out + i1_loc) = val;
 }

feectools/linalg/kernels/stencil_dot_1d/stencil_dot_1d_kernels.py

@@ -4,6 +4,10 @@
 
 Moved from ``feectools.linalg.stencil_dot_kernels.matvec_1d_kernel``. The CUDA version in ``stencil_dot_1d_cuda.cu``
 takes the same arguments in the same order.
+
+The number of diagonals is read from ``mat``: ``n = mat.shape[1]`` diagonals, the pads of the matrix are
+``q = (n - 1) // 2`` (at most the pads ``p_in`` of the domain), and diagonal ``d`` of row ``i`` is the column
+``i - q + d``. Interior rows use all ``n`` diagonals, the last owned row ``n - 1 + add``.
 """
 
 
@@ -17,19 +21,18 @@ def stencil_dot_1d(mat: 'float[:, :]',
                    e_out: int,
                    p_out: int):
 
-    for i1 in range(s_out, e_out):  # global row index
+    n_diags1 = mat.shape[1]
+    # x index of diagonal 0 minus the global row index: column i1 - q + d1 is at x[i1 - q + d1 - s_in + p_in]
+    off1 = p_in - (n_diags1 - 1) // 2 - s_in
+
+    for i1 in range(s_out, e_out + 1):  # global row index
         i1_loc = i1 - s_out  # local row index
+        nd1 = n_diags1
+        if i1 == e_out:
+            nd1 = n_diags1 - 1 + add
+
         val = 0.
-        for d1 in range(2*p_in + 1):
-            val += mat[p_out + i1_loc, d1] * x[i1 + d1 - s_in]
+        for d1 in range(nd1):
+            val += mat[p_out + i1_loc, d1] * x[i1 + d1 + off1]
 
         out[p_out + i1_loc] = val
-
-    # last row treated separately
-    i1 = e_out
-    i1_loc = i1 - s_out  # local row index
-    val = 0.
-    for d1 in range(2*p_in + add):
-        val += mat[p_out + i1_loc, d1] * x[i1 + d1 - s_in]
-
-    out[p_out + i1_loc] = val

feectools/linalg/kernels/stencil_dot_2d/stencil_dot_2d_cuda.cu

@@ -7,11 +7,11 @@
  *
  * One thread per entry of `out` (n_threads = out.size, see __init__.py), the last axis varying fastest. A
  * thread outside the owned rows returns without writing, so the padding of `out` is left as it is, as in
- * pyccel. Along each direction k, interior rows use 2 * p_in[k] + 1 diagonals and the last owned row
- * (i_k == e_out[k]) uses 2 * p_in[k] + add[k]; the pyccel kernel spells out the four combinations, this
- * kernel picks its own per thread. The diagonals are summed in the same order (d1 outer, d2 inner).
+ * pyccel. Along each direction k the matrix has n_k = mat.shape[2 + k] diagonals; interior rows use all of
+ * them, the last owned row (i_k == e_out[k]) uses n_k - 1 + add[k], as in the 1D kernel. The diagonals are
+ * summed in the same order as in pyccel (d1 outer, d2 inner).
  *
- * @param mat   matrix data, shape (rows of `out`..., 2 * p_in + 1...)
+ * @param mat   matrix data, shape (rows of `out`..., diagonals...)
  * @param x     data of the domain vector, ghost regions included
  * @param out   data of the codomain vector; the owned rows are written
  * @param s_in, p_in, add, s_out, e_out, p_out  per direction (length 2), as in the 1D kernel
@@ -30,13 +30,16 @@ extern "C" __global__ void stencil_dot_2d(Array4D<double> mat, Array2D<double> x
     const long long i1 = s_out[0] + i1_loc;  // global row indices
     const long long i2 = s_out[1] + i2_loc;
 
-    const long long n_diags1 = (i1 == e_out[0]) ? 2 * p_in[0] + add[0] : 2 * p_in[0] + 1;
-    const long long n_diags2 = (i2 == e_out[1]) ? 2 * p_in[1] + add[1] : 2 * p_in[1] + 1;
+    const long long nd1 = (i1 == e_out[0]) ? mat.shape[2] - 1 + add[0] : mat.shape[2];
+    const long long nd2 = (i2 == e_out[1]) ? mat.shape[3] - 1 + add[1] : mat.shape[3];
+    // x index of diagonal 0 minus the global row index, per direction
+    const long long off1 = p_in[0] - (mat.shape[2] - 1) / 2 - s_in[0];
+    const long long off2 = p_in[1] - (mat.shape[3] - 1) / 2 - s_in[1];
 
     double val = 0.;
-    for (long long d1 = 0; d1 < n_diags1; ++d1)
-        for (long long d2 = 0; d2 < n_diags2; ++d2)
-            val += mat(p_out[0] + i1_loc, p_out[1] + i2_loc, d1, d2) * x(i1 + d1 - s_in[0], i2 + d2 - s_in[1]);
+    for (long long d1 = 0; d1 < nd1; ++d1)
+        for (long long d2 = 0; d2 < nd2; ++d2)
+            val += mat(p_out[0] + i1_loc, p_out[1] + i2_loc, d1, d2) * x(i1 + d1 + off1, i2 + d2 + off2);
 
     out(p_out[0] + i1_loc, p_out[1] + i2_loc) = val;
 }

feectools/linalg/kernels/stencil_dot_2d/stencil_dot_2d_kernels.py

@@ -4,6 +4,10 @@
 
 Moved from ``feectools.linalg.stencil_dot_kernels.matvec_2d_kernel``. The CUDA version in ``stencil_dot_2d_cuda.cu``
 takes the same arguments in the same order.
+
+The number of diagonals is read from ``mat``, per direction as in ``stencil_dot_1d``: ``n_k = mat.shape[2 + k]``,
+the pads of the matrix are ``q_k = (n_k - 1) // 2`` and diagonal ``d_k`` of row ``i_k`` is the column
+``i_k - q_k + d_k``. Interior rows use all ``n_k`` diagonals, the last owned row ``n_k - 1 + add[k]``.
 """
 
 
@@ -17,83 +21,27 @@ def stencil_dot_2d(mat: 'float[:, :, :, :]',
                    e_out: 'int[:]',
                    p_out: 'int[:]'):
 
-    #####################################
-    #####################################
-    # without last row in 1st direction #
-    #####################################
-    #####################################
-    for i1 in range(s_out[0], e_out[0]):
-        i1_loc = i1 - s_out[0]
+    n_diags1 = mat.shape[2]
+    n_diags2 = mat.shape[3]
+    # x index of diagonal 0 minus the global row index, per direction
+    off1 = p_in[0] - (n_diags1 - 1) // 2 - s_in[0]
+    off2 = p_in[1] - (n_diags2 - 1) // 2 - s_in[1]
+
+    for i1 in range(s_out[0], e_out[0] + 1):  # global row indices
+        i1_loc = i1 - s_out[0]  # local row indices
+        nd1 = n_diags1
+        if i1 == e_out[0]:
+            nd1 = n_diags1 - 1 + add[0]
 
-        #####################################
-        # without last row in 2nd direction #
-        #####################################
-        for i2 in range(s_out[1], e_out[1]):
+        for i2 in range(s_out[1], e_out[1] + 1):
             i2_loc = i2 - s_out[1]
+            nd2 = n_diags2
+            if i2 == e_out[1]:
+                nd2 = n_diags2 - 1 + add[1]
 
             val = 0.
-            for d1 in range(2 * p_in[0] + 1):
-                for d2 in range(2 * p_in[1] + 1):
-                    val += mat[p_out[0] + i1_loc,
-                               p_out[1] + i2_loc,
-                               d1, d2] * x[i1 + d1 - s_in[0],
-                                           i2 + d2 - s_in[1]]
-            out[p_out[0] + i1_loc,
-                p_out[1] + i2_loc] = val
-
-        ##############################################
-        # treat last row in 2nd direction separately #
-        ##############################################
-        i2 = e_out[1]
-        i2_loc = i2 - s_out[1]
-
-        val = 0.
-        for d1 in range(2 * p_in[0] + 1):
-            for d2 in range(2 * p_in[1] + add[1]):
-                val += mat[p_out[0] + i1_loc,
-                           p_out[1] + i2_loc,
-                           d1, d2] * x[i1 + d1 - s_in[0],
-                                       i2 + d2 - s_in[1]]
-        out[p_out[0] + i1_loc,
-            p_out[1] + i2_loc] = val
-
-    ##############################################
-    ##############################################
-    # treat last row in 1st direction separately #
-    ##############################################
-    ##############################################
-    i1 = e_out[0]
-    i1_loc = i1 - s_out[0]
-
-    #####################################
-    # without last row in 2nd direction #
-    #####################################
-    for i2 in range(s_out[1], e_out[1]):
-        i2_loc = i2 - s_out[1]
-
-        val = 0.
-        for d1 in range(2 * p_in[0] + add[0]):
-            for d2 in range(2 * p_in[1] + 1):
-                val += mat[p_out[0] + i1_loc,
-                           p_out[1] + i2_loc,
-                           d1, d2] * x[i1 + d1 - s_in[0],
-                                       i2 + d2 - s_in[1]]
-        out[p_out[0] + i1_loc,
-            p_out[1] + i2_loc] = val
-
-    ##############################################
-    # treat last row in 2nd direction separately #
-    ##############################################
-    i2 = e_out[1]
-    i2_loc = i2 - s_out[1]
-
-    val = 0.
-    for d1 in range(2 * p_in[0] + add[0]):
-        for d2 in range(2 * p_in[1] + add[1]):
-            val += mat[p_out[0] + i1_loc,
-                       p_out[1] + i2_loc,
-                       d1, d2] * x[i1 + d1 - s_in[0],
-                                    i2 + d2 - s_in[1]]
+            for d1 in range(nd1):
+                for d2 in range(nd2):
+                    val += mat[p_out[0] + i1_loc, p_out[1] + i2_loc, d1, d2] * x[i1 + d1 + off1, i2 + d2 + off2]
 
-    out[p_out[0] + i1_loc,
-        p_out[1] + i2_loc] = val
+            out[p_out[0] + i1_loc, p_out[1] + i2_loc] = val