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,6 +3,7 @@
 # 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
@@ -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,68 @@ 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 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 = {}
@@ -682,7 +750,9 @@ def dot(self, v, out=None):
                 assert isinstance(out, BlockVector)
 
             assert out.space is self.codomain
-            out *= 0.0
+            # `_dot` writes every block row of `out` itself (see there); other kernels add to it
+            if self._func is not self._dot:
+                out *= 0.0
         else:
             out = self.codomain.zeros()
 
@@ -697,16 +767,33 @@ def dot(self, v, out=None):
     #...
     @staticmethod
     def _dot(blocks, v, out, n_rows, n_cols, inc):
+        # The first block of each block row writes its product straight into that row of `out`; further blocks of
+        # the row go through `inc` and are added. A row with a single block (e.g. the diagonal blocks of a mass
+        # matrix on a Cuboid, the gradient) thus costs no zeroing pass and no addition; rows without any block are
+        # set to zero.
 
         if n_rows == 1:
+            written = False
             for (_, j), L0j in blocks.items():
-                out += L0j.dot(v[j], out=inc)
-        elif n_cols == 1:
-            for (i, _), Li0 in blocks.items():
-                out[i] += Li0.dot(v, out=inc[i])
+                if written:
+                    out += L0j.dot(v[j], out=inc)
+                else:
+                    L0j.dot(v[j], out=out)
+                    written = True
+            if not written:
+                out *= 0.0
         else:
+            written = set()
             for (i, j), Lij in blocks.items():
-                out[i] += Lij.dot(v[j], out=inc[i])
+                vj = v if n_cols == 1 else v[j]
+                if i in written:
+                    out[i] += Lij.dot(vj, out=inc[i])
+                else:
+                    Lij.dot(vj, out=out[i])
+                    written.add(i)
+            for i in range(n_rows):
+                if i not in written:
+                    out[i] *= 0.0
 
     # ...
     def transpose(self, conjugate=False, out=None):

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)