Device ghost-exchange plans, fused block exchanges, single-reduction PCG (CUDA-aware MPI)

.gitignore

@@ -12,6 +12,9 @@ build
 *egg*
 *dist*
 *cache*
+# source files with 'cache' in their name (the rule above would ignore them)
+!feectools/utilities/cached_launch.py
+!feectools/linalg/tests/test_cached_launch.py
 prof
 
 *.swp

feectools/ddm/blocking_data_exchanger.py

@@ -42,6 +42,7 @@ def __init__( self, cart, dtype, *, coeff_shape=(),  assembly=False, axis=None,
         self._cart = cart
         self._comm = cart.comm_cart
         self._axis = axis
+        self._shape = shape
 
         if assembly:
             self._assembly_send_types, self._assembly_recv_types = self._create_assembly_buffer_types(
@@ -82,6 +83,12 @@ def start_update_ghost_regions( self, array, requests ):
 
         assert isinstance( array, xp.ndarray )
 
+        if xp.is_gpu( array ):
+            # contiguous device buffers instead of subarray datatypes (see CartDataExchanger)
+            for direction in range( self._cart.ndim ):
+                self._device_update_ghost_regions_along( array, direction )
+            return
+
         # MPI reads/writes `array` directly; on a device backend the
         # kernels that produced it must have finished first.
         synchronize_for_mpi( array )
@@ -127,6 +134,11 @@ def start_exchange_assembly_data( self, array ):
 
         assert isinstance( array, xp.ndarray )
 
+        if xp.is_gpu( array ):
+            # contiguous device buffers instead of subarray datatypes (see CartDataExchanger)
+            self._device_exchange_assembly_data( array )
+            return
+
         # MPI reads/writes `array` directly; on a device backend the
         # kernels that produced it must have finished first.
         synchronize_for_mpi( array )

feectools/ddm/nonblocking_data_exchanger.py

@@ -43,6 +43,7 @@ def __init__( self, cart, dtype, *, coeff_shape=(),  assembly=False, axis=None,
         self._cart = cart
         self._comm = cart.comm_cart
         self._axis = axis
+        self._shape = shape
 
         if assembly:
             self._assembly_send_types, self._assembly_recv_types = self._create_assembly_buffer_types(
@@ -74,6 +75,11 @@ def get_assembly_recv_type( self, *args ):
     # ...
     def prepare_communications(self, u):
 
+        # Device arrays need no persistent requests: their ghost regions are exchanged
+        # through contiguous buffers (`_device_ghost_plan`)
+        if xp.is_gpu(u):
+            return ()
+
         # Requests' handles
         requests = []
         cart = self._cart
@@ -99,19 +105,52 @@ def prepare_communications(self, u):
         return tuple(requests)
 
     def start_update_ghost_regions(self, array, requests ):
+        if xp.is_gpu(array):
+            # started here, finished in end_update_ghost_regions: several vectors (e.g. the
+            # components of a block vector) can exchange at the same time
+            plan = self._device_ghost_plan(array)
+            pending = self.__dict__.setdefault('_device_pending', {})
+            pending[id(array)] = (plan, plan.start(array))
+            return
         # The persistent requests read/write `array` directly; on a device
         # backend the kernels that produced it must have finished first.
         synchronize_for_mpi( array )
         MPI.Prequest.Startall( requests )
 
     def end_update_ghost_regions(self, array, requests):
+        if xp.is_gpu(array):
+            plan, device_requests = self._device_pending.pop(id(array))
+            plan.finish(array, device_requests)
+            return
         MPI.Prequest.Waitall  ( requests )
 
+    # ...
+    def _device_ghost_plan(self, array):
+        """The exchange of all ghost regions (every shift, as the persistent requests) of a device array."""
+        cart = self._cart
+
+        def transfers():
+            out = []
+            for shift in product( [-1,0,1], repeat=cart._ndims ):
+                if all(s==0 for s in shift):
+                    continue
+                info = cart.get_shift_info_non_blocking( shift )
+                out.append((self._region(info['send_starts'], info['buf_shape'], array.shape), info['rank_dest'],
+                            self._region(info['recv_starts'], info['buf_shape'], array.shape), info['rank_source']))
+            return out
+
+        return self._device_plan('ghost', array, cart._comm_cart, transfers)
+
     # ...
     def start_exchange_assembly_data( self, array ):
 
         assert isinstance( array, xp.ndarray )
 
+        if xp.is_gpu( array ):
+            # contiguous device buffers instead of subarray datatypes (see CartDataExchanger)
+            self._device_exchange_assembly_data( array )
+            return
+
         # MPI reads/writes `array` directly; on a device backend the
         # kernels that produced it must have finished first.
         synchronize_for_mpi( array )

feectools/linalg/basic.py

@@ -820,8 +820,9 @@ def dot(self, v, out=None):
         if out is not None:
             assert isinstance(out, Vector)
             assert out.space == self.codomain
-            out *= 0
-            out += v
+            if out is not v:
+                # one copy (instead of out *= 0; out += v, which also kept NaNs of the old out)
+                v.copy(out=out)
             return out
         else:
             return v.copy()

feectools/linalg/block.py

@@ -3,12 +3,13 @@
 # Copyright 2018 Jalal Lakhlili, Yaman Güçlü
 
 import cunumpy as xp
+import numpy as np
 
 from types import MappingProxyType
 from scipy.sparse import bmat, lil_matrix
 
 from feectools.linalg.basic    import VectorSpace, Vector, LinearOperator
-from feectools.linalg.stencil  import StencilMatrix
+from feectools.linalg.stencil  import StencilMatrix, _fake_cupy_active
 from feectools.ddm.cart        import InterfaceCartDecomposition
 from feectools.ddm.utilities   import get_data_exchanger
 
@@ -126,6 +127,19 @@ def inner(self, x, y):
         assert y.space is self
         return sum(Vi.inner(xi, yi) for Vi, xi, yi in zip(self.spaces, x.blocks, y.blocks))
 
+    def inner_local(self, x, y):
+        """The contribution of this process to :meth:`inner` (sum of the blocks', not reduced)."""
+        assert isinstance(x, BlockVector)
+        assert isinstance(y, BlockVector)
+        assert x.space is self
+        assert y.space is self
+        return sum(Vi.inner_local(xi, yi) for Vi, xi, yi in zip(self.spaces, x.blocks, y.blocks))
+
+    @property
+    def inner_comm(self):
+        """The communicator :meth:`inner` reduces over: that of the blocks (None if not parallel)."""
+        return self.spaces[0].inner_comm
+
     #...
     def axpy(self, a, x, y):
         """
@@ -399,14 +413,70 @@ def update_ghost_regions(self):
 
         req = self.start_update_interface_ghost_regions()
 
-        for vi in self.blocks:
-            vi.update_ghost_regions()
+        fused = self._fused_device_exchange()
+        if fused is not None:
+            # device data on one communicator: one message per neighbour for all blocks (see _fused_device_exchange)
+            arrays = [vi._data for vi in self.blocks]
+            fused.finish(arrays, fused.start(arrays))
+            for vi in self.blocks:
+                vi._sync = True
+        elif not _fake_cupy_active() and all(
+            hasattr(vi, '_start_update_ghost_regions') and xp.is_gpu(vi._data) for vi in self.blocks
+        ):
+            # device data: the exchanges of all blocks are in flight together, with one wait each at the
+            # end, instead of one complete exchange (and its fixed costs) after the other
+            for vi in self.blocks:
+                vi._start_update_ghost_regions()
+            for vi in self.blocks:
+                vi._end_update_ghost_regions()
+        else:
+            for vi in self.blocks:
+                vi.update_ghost_regions()
 
         self.end_update_interface_ghost_regions(req)
 
         # Flag ghost regions as up-to-date
         self._sync = True
 
+    def _fused_device_exchange(self):
+        """The ghost update of all blocks as one exchange, or None if it does not apply.
+
+        For blocks that are parallel StencilVectors with C-contiguous float64 device data on the same communicator and
+        without interfaces, the ghost regions of all blocks that go to the same rank travel as one message, with one
+        gather and one scatter kernel for all blocks (:class:`~feectools.ddm.basic._FusedDeviceExchange`), instead
+        of one message per block and neighbour. Built on first use and cached on the space (by the shapes).
+        """
+        from feectools.linalg.stencil import StencilVector
+
+        blocks = self.blocks
+        if not blocks or not all(isinstance(vi, StencilVector) for vi in blocks):
+            return None
+        V0 = blocks[0].space
+        from feectools.linalg.stencil import _fake_cupy_active
+        if _fake_cupy_active():
+            # cunumpy's fake CuPy (tests without a GPU) cannot run the gather and scatter kernels
+            return None
+        if not all(
+            vi.space.parallel
+            and not vi.space.cart.is_comm_null
+            and not vi.space.interfaces
+            and xp.is_gpu(vi._data)
+            and vi._data.dtype == np.float64
+            and vi._data.flags.c_contiguous
+            and vi.space.cart.comm_cart is V0.cart.comm_cart
+            for vi in blocks
+        ):
+            return None
+
+        cache = self.space.__dict__.setdefault('_fused_device_exchanges', {})
+        key = tuple(vi._data.shape for vi in blocks)
+        if key not in cache:
+            from feectools.ddm.basic import _FusedDeviceExchange
+
+            plans = [vi.space._synchronizer._device_ghost_plan(vi._data) for vi in blocks]
+            cache[key] = _FusedDeviceExchange(plans) if all(p is not None for p in plans) else None
+        return cache[key]
+
     def start_update_interface_ghost_regions(self):
         self._collect_interface_buf()
         req = {}
@@ -697,7 +767,6 @@ def dot(self, v, out=None):
     #...
     @staticmethod
     def _dot(blocks, v, out, n_rows, n_cols, inc):
-
         if n_rows == 1:
             for (_, j), L0j in blocks.items():
                 out += L0j.dot(v[j], out=inc)

feectools/linalg/kernels/stencil_dot_3d/__init__.py

@@ -1,6 +1,33 @@
 """The kernel of this folder: pyccel in ``stencil_dot_3d_kernels.py``, CUDA in ``stencil_dot_3d_cuda.cu``."""
 
+from functools import cache
+from pathlib import Path
+
 from cunumpy.kernels import Kernel
 
-# one thread per entry of `out` (the third argument)
-stencil_dot_3d = Kernel.from_folder(__name__, n_threads_from=lambda args: args[2].size)
+# STENCIL_DOT_LANES = 8 threads per entry of `out` (the third argument), see stencil_dot_3d_cuda.cu
+stencil_dot_3d = Kernel.from_folder(__name__, n_threads_from=lambda args: 8 * args[2].size)
+
+
+@cache
+def stencil_dot_3d_rows():
+    """The CUDA kernel with one thread per entry of `out` and no warp shuffles (``stencil_dot_3d_rows.cu``).
+
+    For launches on cunumpy's fake CuPy (tests without a GPU), whose CPU emulation runs one thread at a time and
+    rejects warp shuffles; on a GPU :data:`stencil_dot_3d` is used.
+    """
+    from cunumpy.kernels import CudaKernel
+
+    kernel = CudaKernel.from_file(Path(__file__).with_name("stencil_dot_3d_rows.cu"), "stencil_dot_3d_rows")
+    kernel.n_threads_from = lambda args: args[2].size
+    return kernel
+
+
+def stencil_dot_3d_fake_cupy(*args):
+    """:data:`stencil_dot_3d` on cunumpy's fake CuPy: :func:`stencil_dot_3d_rows` for device arguments (emulated
+    launches), the usual dispatch (the pyccel kernel) for host arguments."""
+    import cunumpy as xp
+
+    if any(xp.is_gpu(a) for a in args[:3]):
+        return stencil_dot_3d_rows()(*args)
+    return stencil_dot_3d(*args)

feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_cuda.cu

@@ -2,14 +2,23 @@
 #include "cunumpy/array_view.cuh"
 #include "cunumpy/index.cuh"
 
+// Lanes that share one row of `out` (a divisor of 32): n_threads = STENCIL_DOT_LANES * out.size, see __init__.py.
+// 8 was the fastest of 8, 16 and 32 for 3 x 5 x 7 diagonals at 128^3 on an H100 (1.52 ms vs 1.69 ms with one thread
+// per row).
+#define STENCIL_DOT_LANES 8
+
 /**
  * Stencil matrix-vector product out = mat @ x on the owned rows, as the pyccel kernel stencil_dot_3d.
  *
- * 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 the matrix has n_k = mat.shape[3 + 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, d2, d3 innermost).
+ * STENCIL_DOT_LANES consecutive threads per entry of `out` (the last axis of `out` varying fastest): the lanes split
+ * the diagonals of the row among themselves (lane k takes the flat diagonal indices k, k + LANES, ...), so that
+ * neighbouring threads read neighbouring diagonals of the row of `mat` (its last three, contiguous axes) instead of
+ * addresses n_diagonals * 8 bytes apart; the partial sums are then added with warp shuffles. The summation order
+ * differs from pyccel (round-off only).
+ *
+ * Rows outside the owned rows are not written, so the padding of `out` is left as it is, as in pyccel. Along each
+ * direction k the matrix has n_k = mat.shape[3 + 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.
  *
  * @param mat   matrix data (a 6D view), shape (rows of `out`..., diagonals...)
  * @param x     data of the domain vector, ghost regions included
@@ -21,32 +30,57 @@ extern "C" __global__ void stencil_dot_3d(Array6D<double> mat, Array3D<double> x
                                           const long long* s_out, const long long* e_out,
                                           const long long* p_out)
 {
-    CUNUMPY_THREAD_1D(thread, out.size());
-
-    const long long i1_loc = thread / (out.shape[1] * out.shape[2]) - p_out[0];  // local row indices
-    const long long i2_loc = (thread / out.shape[2]) % out.shape[1] - p_out[1];
-    const long long i3_loc = thread % out.shape[2] - p_out[2];
-    if (i1_loc < 0 || i1_loc > e_out[0] - s_out[0]) return;
-    if (i2_loc < 0 || i2_loc > e_out[1] - s_out[1]) return;
-    if (i3_loc < 0 || i3_loc > e_out[2] - s_out[2]) return;
-    const long long i1 = s_out[0] + i1_loc;  // global row indices
-    const long long i2 = s_out[1] + i2_loc;
-    const long long i3 = s_out[2] + i3_loc;
-
-    const long long nd1 = (i1 == e_out[0]) ? mat.shape[3] - 1 + add[0] : mat.shape[3];
-    const long long nd2 = (i2 == e_out[1]) ? mat.shape[4] - 1 + add[1] : mat.shape[4];
-    const long long nd3 = (i3 == e_out[2]) ? mat.shape[5] - 1 + add[2] : mat.shape[5];
-    // x index of diagonal 0 minus the global row index, per direction
-    const long long off1 = p_in[0] - (mat.shape[3] - 1) / 2 - s_in[0];
-    const long long off2 = p_in[1] - (mat.shape[4] - 1) / 2 - s_in[1];
-    const long long off3 = p_in[2] - (mat.shape[5] - 1) / 2 - s_in[2];
+    // CUDA-only: the row (entry of `out`) of this group of lanes and the lane within it. No thread returns before the
+    // shuffles at the end (they need every lane of the warp): threads past the end and rows outside the owned ones
+    // just have nothing to add.
+    const long long thread = (long long)blockDim.x * blockIdx.x + threadIdx.x;
+    const long long row = thread / STENCIL_DOT_LANES;
+    const int lane = (int)(thread % STENCIL_DOT_LANES);
+
+    const long long i1_loc = row / (out.shape[1] * out.shape[2]) - p_out[0];  // local row indices
+    const long long i2_loc = (row / out.shape[2]) % out.shape[1] - p_out[1];
+    const long long i3_loc = row % out.shape[2] - p_out[2];
+    const bool owned = row < out.size() && i1_loc >= 0 && i1_loc <= e_out[0] - s_out[0] && i2_loc >= 0 &&
+                       i2_loc <= e_out[1] - s_out[1] && i3_loc >= 0 && i3_loc <= e_out[2] - s_out[2];
+
+    const long long r1 = p_out[0] + i1_loc;  // row of `mat` and `out`
+    const long long r2 = p_out[1] + i2_loc;
+    const long long r3 = p_out[2] + i3_loc;
 
     double val = 0.;
-    for (long long d1 = 0; d1 < nd1; ++d1)
-        for (long long d2 = 0; d2 < nd2; ++d2)
-            for (long long d3 = 0; d3 < nd3; ++d3)
-                val += mat(p_out[0] + i1_loc, p_out[1] + i2_loc, p_out[2] + i3_loc, d1, d2, d3) *
-                       x(i1 + d1 + off1, i2 + d2 + off2, i3 + d3 + off3);
+    if (owned) {
+        const long long i1 = s_out[0] + i1_loc;  // global row indices
+        const long long i2 = s_out[1] + i2_loc;
+        const long long i3 = s_out[2] + i3_loc;
+
+        const long long nd1 = (i1 == e_out[0]) ? mat.shape[3] - 1 + add[0] : mat.shape[3];
+        const long long nd2 = (i2 == e_out[1]) ? mat.shape[4] - 1 + add[1] : mat.shape[4];
+        const long long nd3 = (i3 == e_out[2]) ? mat.shape[5] - 1 + add[2] : mat.shape[5];
+        // x index of diagonal 0 minus the global row index, per direction
+        const long long off1 = p_in[0] - (mat.shape[3] - 1) / 2 - s_in[0];
+        const long long off2 = p_in[1] - (mat.shape[4] - 1) / 2 - s_in[1];
+        const long long off3 = p_in[2] - (mat.shape[5] - 1) / 2 - s_in[2];
+
+        // 32-bit index arithmetic in the loop (64-bit division is slow on the GPU); the diagonals of the row and the
+        // matching entries of x through pointers and strides
+        const int n3 = (int)nd3, n23 = (int)(nd2 * nd3), nd = (int)(nd1 * nd2 * nd3);
+        const double* mat_row = &mat(r1, r2, r3, 0, 0, 0);
+        const double* x_row = &x(i1 + off1, i2 + off2, i3 + off3);
+        const long long m3 = mat.strides[3], m4 = mat.strides[4], m5 = mat.strides[5];
+        const long long x1 = x.strides[0], x2 = x.strides[1], x3 = x.strides[2];
+
+        for (int d = lane; d < nd; d += STENCIL_DOT_LANES) {
+            const int d1 = d / n23;
+            const int d23 = d - d1 * n23;
+            const int d2 = d23 / n3;
+            const int d3 = d23 - d2 * n3;
+            val += mat_row[d1 * m3 + d2 * m4 + d3 * m5] * x_row[d1 * x1 + d2 * x2 + d3 * x3];
+        }
+    }
+
+    // sum of the partial sums of the row's lanes (the first lane ends up with the total)
+    for (int offset = STENCIL_DOT_LANES / 2; offset > 0; offset /= 2)
+        val += __shfl_down_sync(0xffffffffu, val, offset, STENCIL_DOT_LANES);
 
-    out(p_out[0] + i1_loc, p_out[1] + i2_loc, p_out[2] + i3_loc) = val;
+    if (owned && lane == 0) out(r1, r2, r3) = val;
 }