From a8d5338c0d660ccc5826ef42f5b9399be368e675 Mon Sep 17 00:00:00 2001
From: Max Lindqvist <max.lindqvist@ipp.mpg.de>
Date: Tue, 13 Oct 2026 17:14:54 +0200
Subject: [PATCH 01/10] fixes

---
 feectools/ddm/basic.py                      | 177 ++++++++++++++++++++
 feectools/ddm/blocking_data_exchanger.py    |  12 ++
 feectools/ddm/nonblocking_data_exchanger.py |  33 ++++
 feectools/linalg/stencil.py                 |  18 +-
 4 files changed, 235 insertions(+), 5 deletions(-)

diff --git a/feectools/ddm/basic.py b/feectools/ddm/basic.py
index 8d1d5c23a..37b8cf5f6 100644
--- a/feectools/ddm/basic.py
+++ b/feectools/ddm/basic.py
@@ -5,8 +5,109 @@
 #---------------------------------------------------------------------------#
 from abc import ABC, abstractmethod
 
+import cunumpy as xp
+import numpy as np
+from cunumpy.mpi import exchange, get_mpi_cuda_aware
+from maybempi import MPI
+
 
 __all__ = ('CartDataExchanger',)
+
+#===============================================================================
+class _DeviceRegionExchange:
+    """
+    Copy regions of a device (CuPy) array between ranks without MPI derived datatypes.
+
+    Open MPI's CUDA-aware path moves a subarray datatype on device memory one contiguous
+    piece at a time (~15 us per piece on H100 nodes), so a strided ghost region of a few
+    hundred bytes can take milliseconds. Here all regions are gathered on the device into
+    one contiguous buffer with a single kernel (flat indices computed once), the regions
+    going to the same rank travel as one message, a region that a rank sends to itself
+    (periodic direction that is not decomposed) is copied without MPI, and one more kernel
+    scatters everything received into the array.
+
+    All regions are read (gathered) before any is written, as in a nonblocking MPI exchange.
+
+    Parameters
+    ----------
+    comm : mpi4py.MPI.Comm
+        Communicator of the ranks in `transfers`.
+
+    transfers : list of (send_idx, dest, recv_idx, source)
+        Send the region `send_idx` of the array to rank `dest` and receive the region
+        `recv_idx` from rank `source` (MPI.PROC_NULL: no message). Every rank lists the
+        transfers in the same order: the regions exchanged between two ranks are matched
+        by that order.
+
+    shape, dtype :
+        Shape and dtype of the (C-contiguous) arrays the exchange is used for.
+    """
+    TAG = 4242
+
+    def __init__(self, comm, transfers, shape, dtype):
+
+        me = comm.Get_rank()
+        self._comm = comm
+        flat = np.arange(int(np.prod(shape)), dtype=np.int64).reshape(shape)
+
+        local_send, local_recv, sends, recvs = [], [], {}, {}
+        for send_idx, dest, recv_idx, source in transfers:
+            if dest == me and source == me:
+                local_send.append(flat[send_idx].ravel())
+                local_recv.append(flat[recv_idx].ravel())
+                continue
+            if dest != MPI.PROC_NULL:
+                sends.setdefault(dest, []).append(flat[send_idx].ravel())
+            if source != MPI.PROC_NULL:
+                recvs.setdefault(source, []).append(flat[recv_idx].ravel())
+
+        def layout(groups, local):
+            """Flat indices of [message to/from each rank..., local regions], and each message's slice."""
+            parts, slices, offset = [], [], 0
+            for rank, regions in groups.items():
+                n = sum(r.size for r in regions)
+                parts += regions
+                slices.append((rank, slice(offset, offset + n)))
+                offset += n
+            local_slice = slice(offset, offset + sum(r.size for r in local))
+            parts += local
+            index = np.concatenate(parts) if parts else np.zeros(0, dtype=np.int64)
+            return xp.asarray(index), slices, local_slice
+
+        self._send_index, send_slices, self._local_send = layout(sends, local_send)
+        self._recv_index, recv_slices, self._local_recv = layout(recvs, local_recv)
+        self._send_buf = xp.empty(self._send_index.size, dtype=dtype)
+        self._recv_buf = xp.empty(self._recv_index.size, dtype=dtype)
+        self._sends = [(self._send_buf[sl], rank) for rank, sl in send_slices]
+        self._recvs = [(self._recv_buf[sl], rank) for rank, sl in recv_slices]
+
+    def __call__(self, array):
+
+        assert array.flags.c_contiguous, "device ghost exchange needs a C-contiguous array"
+        values = array.reshape(-1)
+
+        # gather every region to send (or copy locally) with one kernel
+        if self._send_index.size:
+            xp.take(values, self._send_index, out=self._send_buf)
+
+        if self._sends or self._recvs:
+            # device buffers straight to MPI if it is CUDA-aware (assumed when unknown, as for the
+            # datatype exchange this replaces), staged through the host otherwise
+            cuda_aware = get_mpi_cuda_aware()
+            exchange(
+                self._comm,
+                sends=[(buf, rank, self.TAG) for buf, rank in self._sends],
+                receives=[(buf, rank, self.TAG) for buf, rank in self._recvs],
+                cuda_aware=True if cuda_aware is None else cuda_aware,
+            )
+
+        # scatter everything received (and the local copies) with one kernel
+        if self._recv_index.size:
+            if self._local_recv.stop > self._local_recv.start:
+                self._recv_buf[self._local_recv] = self._send_buf[self._local_send]
+            values[self._recv_index] = self._recv_buf
+
+
 #===============================================================================
 class CartDataExchanger(ABC):
     """
@@ -80,3 +181,79 @@ def start_exchange_assembly_data( self, array ):
     def end_exchange_assembly_data( self, array ):
         pass
 
+    #---------------------------------------------------------------------------
+    # Device (CuPy) arrays: contiguous buffers instead of MPI derived datatypes
+    #---------------------------------------------------------------------------
+    @staticmethod
+    def _region(starts, buf_shape, array_shape):
+        """Index of the region `buf_shape` at `starts` of the decomposed axes, whole along the coefficient axes."""
+        idx = [slice(int(s), int(s) + int(b)) for s, b in zip(starts, buf_shape)]
+        idx += [slice(0, int(n)) for n in array_shape[len(idx):]]
+        return tuple(idx)
+
+    def _device_exchange(self, key, array, comm, transfers):
+        """Run the `_DeviceRegionExchange` of `transfers` (built by the callable on first use) on `array`."""
+        plans = self.__dict__.setdefault('_device_plans', {})
+        plan_key = (key, tuple(array.shape), np.dtype(array.dtype))
+        plan = plans.get(plan_key)
+        if plan is None:
+            plan = plans[plan_key] = _DeviceRegionExchange(comm, transfers(), array.shape, array.dtype)
+        plan(array)
+
+    def _device_update_ghost_regions_along(self, array, direction):
+        """Ghost regions of `direction` (both sides) from the neighbours, as the blocking exchange."""
+        cart = self._cart
+
+        def transfers():
+            out = []
+            for disp in [-1, 1]:
+                info = cart.get_shift_info(direction, disp)
+                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
+
+        self._device_exchange(('ghost', direction), array, self._comm, transfers)
+
+    def _device_exchange_assembly_data(self, array):
+        """`start_exchange_assembly_data` for a device array: same transfers and sums, contiguous buffers."""
+        cart  = self._cart
+        gcomm = self._comm
+        ndim  = cart.ndim
+        axis  = self._axis
+
+        for direction in range(ndim):
+            if direction == axis:
+                continue
+            comm = cart.subcomm[direction] if axis is not None else gcomm
+            disp = 1
+            info = cart.get_shift_info(direction, disp)
+
+            def transfers(info=info, comm=comm):
+                rank_source, rank_dest = info['rank_source'], info['rank_dest']
+                if axis is not None:
+                    rank_source = gcomm.group.Translate_ranks([int(rank_source)], comm.group)[0]
+                    rank_dest   = gcomm.group.Translate_ranks([int(rank_dest)], comm.group)[0]
+                buf_shape   = list(info['buf_shape'])
+                send_starts = list(info['send_assembly_starts'])
+                recv_starts = list(info['recv_assembly_starts'])
+                if axis is not None:
+                    buf_shape[axis]   = self._shape[axis]
+                    send_starts[axis] = 0
+                    recv_starts[axis] = 0
+                return [(self._region(send_starts, buf_shape, array.shape), rank_dest,
+                         self._region(recv_starts, buf_shape, array.shape), rank_source)]
+
+            self._device_exchange(('assembly', direction), array, comm, transfers)
+            self._add_assembly_ghosts(array, direction)
+
+    def _add_assembly_ghosts(self, array, direction):
+        """Add the received assembly data of `direction` (disp = 1) to the owned entries."""
+        cart = self._cart
+        ndim = cart.ndim
+        info = cart.get_shift_info(direction, 1)
+        pads = [0]*ndim
+        pads[direction] = cart._pads[direction]*cart._shifts[direction]
+        idx_from = tuple(slice(s,s+b) for s,b in zip(info['recv_starts'],info['buf_shape']))
+        idx_to   = tuple(slice(s+p,s+b+p) for s,b,p in zip(info['recv_starts'],info['buf_shape'],pads))
+        array[idx_to] += array[idx_from]
+
diff --git a/feectools/ddm/blocking_data_exchanger.py b/feectools/ddm/blocking_data_exchanger.py
index 68776f764..02f2fbc86 100644
--- a/feectools/ddm/blocking_data_exchanger.py
+++ b/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 )
diff --git a/feectools/ddm/nonblocking_data_exchanger.py b/feectools/ddm/nonblocking_data_exchanger.py
index 3f70b933c..29167253a 100644
--- a/feectools/ddm/nonblocking_data_exchanger.py
+++ b/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_update_ghost_regions`)
+        if xp.is_gpu(u):
+            return ()
+
         # Requests' handles
         requests = []
         cart = self._cart
@@ -99,19 +105,46 @@ def prepare_communications(self, u):
         return tuple(requests)
 
     def start_update_ghost_regions(self, array, requests ):
+        if xp.is_gpu(array):
+            self._device_update_ghost_regions(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):
+            return
         MPI.Prequest.Waitall  ( requests )
 
+    # ...
+    def _device_update_ghost_regions(self, array):
+        """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
+
+        self._device_exchange('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 )
diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py
index cc8299189..3e446b5c8 100644
--- a/feectools/linalg/stencil.py
+++ b/feectools/linalg/stencil.py
@@ -322,11 +322,19 @@ def inner(self, x, y):
             self._inner_func(x._data, y._data, *self._inner_consts, res)
 
         if self.parallel:
-            # The send buffer may have just been written on the device; MPI reads it directly.
-            synchronize_for_mpi(x._dot_send_data, x._dot_recv_data)
-            self.cart.global_comm.Allreduce((x._dot_send_data, self.mpi_type),
-                                            (x._dot_recv_data, self.mpi_type),
-                                             op=MPI.SUM )
+            if xp.is_gpu(x._dot_send_data):
+                # One number: an Allreduce of a host copy (~10 us) is much cheaper than the
+                # CUDA-aware Allreduce of the device buffer (~40 us on H100 nodes). The result
+                # goes back into the device buffer, so the inner product is still a device array.
+                send = xp.to_numpy(x._dot_send_data)
+                recv = send.copy()
+                self.cart.global_comm.Allreduce((send, self.mpi_type), (recv, self.mpi_type), op=MPI.SUM)
+                x._dot_recv_data[...] = xp.asarray(recv)
+            else:
+                synchronize_for_mpi(x._dot_send_data, x._dot_recv_data)
+                self.cart.global_comm.Allreduce((x._dot_send_data, self.mpi_type),
+                                                (x._dot_recv_data, self.mpi_type),
+                                                 op=MPI.SUM )
             res = x._dot_recv_data
 
         if xp.is_gpu(res):

From 0ab90db424fe98b169e01ca6a68bc8f6a70e4058 Mon Sep 17 00:00:00 2001
From: Max Lindqvist <max.lindqvist@ipp.mpg.de>
Date: Wed, 14 Oct 2026 17:21:21 +0200
Subject: [PATCH 02/10] Build device exchange plans from per-region indices

---
 feectools/ddm/basic.py | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/feectools/ddm/basic.py b/feectools/ddm/basic.py
index 37b8cf5f6..a75138717 100644
--- a/feectools/ddm/basic.py
+++ b/feectools/ddm/basic.py
@@ -4,6 +4,7 @@
 # for full license details.                                                 #
 #---------------------------------------------------------------------------#
 from abc import ABC, abstractmethod
+from functools import reduce
 
 import cunumpy as xp
 import numpy as np
@@ -48,18 +49,27 @@ def __init__(self, comm, transfers, shape, dtype):
 
         me = comm.Get_rank()
         self._comm = comm
-        flat = np.arange(int(np.prod(shape)), dtype=np.int64).reshape(shape)
+        # C-order element strides: the flat indices of a region are built from per-axis ranges, so the work and
+        # memory scale with the region, not with the whole array (a stencil matrix block has ~10^8 entries)
+        strides = [1] * len(shape)
+        for axis in range(len(shape) - 2, -1, -1):
+            strides[axis] = strides[axis + 1] * int(shape[axis + 1])
+
+        def flat(idx):
+            """Flat (C-order) indices of the region `idx` of the array, in C order of the region."""
+            axes = [np.arange(*sl.indices(int(n)), dtype=np.int64) * st for sl, n, st in zip(idx, shape, strides)]
+            return reduce(np.add, np.ix_(*axes)).ravel()
 
         local_send, local_recv, sends, recvs = [], [], {}, {}
         for send_idx, dest, recv_idx, source in transfers:
             if dest == me and source == me:
-                local_send.append(flat[send_idx].ravel())
-                local_recv.append(flat[recv_idx].ravel())
+                local_send.append(flat(send_idx))
+                local_recv.append(flat(recv_idx))
                 continue
             if dest != MPI.PROC_NULL:
-                sends.setdefault(dest, []).append(flat[send_idx].ravel())
+                sends.setdefault(dest, []).append(flat(send_idx))
             if source != MPI.PROC_NULL:
-                recvs.setdefault(source, []).append(flat[recv_idx].ravel())
+                recvs.setdefault(source, []).append(flat(recv_idx))
 
         def layout(groups, local):
             """Flat indices of [message to/from each rank..., local regions], and each message's slice."""

From 7664ff8d92531aa7a85d685d9590297a72d67158 Mon Sep 17 00:00:00 2001
From: Max Lindqvist <max.lindqvist@ipp.mpg.de>
Date: Wed, 14 Oct 2026 17:21:21 +0200
Subject: [PATCH 03/10] BlockLinearOperator.dot: write the first block of each
 row directly

---
 feectools/linalg/block.py | 31 +++++++++++++++++++++++++------
 1 file changed, 25 insertions(+), 6 deletions(-)

diff --git a/feectools/linalg/block.py b/feectools/linalg/block.py
index 541f1c4a0..3e461202d 100644
--- a/feectools/linalg/block.py
+++ b/feectools/linalg/block.py
@@ -682,7 +682,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 +699,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):

From c10b33bfa1c2bc3ccb2962db683ee58c1b7c8e09 Mon Sep 17 00:00:00 2001
From: Max Lindqvist <max.lindqvist@ipp.mpg.de>
Date: Wed, 14 Oct 2026 17:21:21 +0200
Subject: [PATCH 04/10] stencil_dot_3d (CUDA): 8 threads per row, coalesced
 reads

---
 .../linalg/kernels/stencil_dot_3d/__init__.py |  4 +-
 .../stencil_dot_3d/stencil_dot_3d_cuda.cu     | 94 +++++++++++++------
 2 files changed, 66 insertions(+), 32 deletions(-)

diff --git a/feectools/linalg/kernels/stencil_dot_3d/__init__.py b/feectools/linalg/kernels/stencil_dot_3d/__init__.py
index 4c486c986..3477cdd4d 100644
--- a/feectools/linalg/kernels/stencil_dot_3d/__init__.py
+++ b/feectools/linalg/kernels/stencil_dot_3d/__init__.py
@@ -2,5 +2,5 @@
 
 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)
diff --git a/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_cuda.cu b/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_cuda.cu
index 066fa86d3..80c3e3b48 100644
--- a/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_cuda.cu
+++ b/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;
 }

From 6571ebd75e9f6e52bed7f186d5bca59e2cede335 Mon Sep 17 00:00:00 2001
From: Max Lindqvist <max.lindqvist@ipp.mpg.de>
Date: Wed, 14 Oct 2026 23:06:19 +0200
Subject: [PATCH 05/10] MPI improvements

---
 feectools/ddm/basic.py                      |  49 +++--
 feectools/ddm/nonblocking_data_exchanger.py |  16 +-
 feectools/linalg/block.py                   |  25 ++-
 feectools/linalg/solvers.py                 | 187 +++++++++++++++++++-
 feectools/linalg/stencil.py                 | 127 ++++++++++++-
 feectools/linalg/tests/test_solvers.py      |  14 +-
 6 files changed, 381 insertions(+), 37 deletions(-)

diff --git a/feectools/ddm/basic.py b/feectools/ddm/basic.py
index a75138717..346a1d46b 100644
--- a/feectools/ddm/basic.py
+++ b/feectools/ddm/basic.py
@@ -8,7 +8,7 @@
 
 import cunumpy as xp
 import numpy as np
-from cunumpy.mpi import exchange, get_mpi_cuda_aware
+from cunumpy.mpi import exchange, get_mpi_cuda_aware, synchronize_for_mpi
 from maybempi import MPI
 
 
@@ -91,8 +91,15 @@ def layout(groups, local):
         self._sends = [(self._send_buf[sl], rank) for rank, sl in send_slices]
         self._recvs = [(self._recv_buf[sl], rank) for rank, sl in recv_slices]
 
-    def __call__(self, array):
+    def start(self, array):
+        """Gather the regions to send and post the messages; returns the requests for :meth:`finish`.
 
+        With CUDA-aware MPI the device buffers go to MPI directly after one stream synchronization
+        (which also orders the receives after the previous scatter from the receive buffer), and
+        the transfers proceed until :meth:`finish`: several exchanges (e.g. the components of a
+        block vector) can be in flight at once. Otherwise the buffers are staged through the host
+        and the exchange completes here.
+        """
         assert array.flags.c_contiguous, "device ghost exchange needs a C-contiguous array"
         values = array.reshape(-1)
 
@@ -100,23 +107,41 @@ def __call__(self, array):
         if self._send_index.size:
             xp.take(values, self._send_index, out=self._send_buf)
 
-        if self._sends or self._recvs:
-            # device buffers straight to MPI if it is CUDA-aware (assumed when unknown, as for the
-            # datatype exchange this replaces), staged through the host otherwise
-            cuda_aware = get_mpi_cuda_aware()
+        if not (self._sends or self._recvs):
+            return []
+
+        # device buffers straight to MPI if it is CUDA-aware (assumed when unknown, as for the
+        # datatype exchange this replaces), staged through the host otherwise
+        cuda_aware = get_mpi_cuda_aware()
+        if cuda_aware is False:
             exchange(
                 self._comm,
                 sends=[(buf, rank, self.TAG) for buf, rank in self._sends],
                 receives=[(buf, rank, self.TAG) for buf, rank in self._recvs],
-                cuda_aware=True if cuda_aware is None else cuda_aware,
+                cuda_aware=False,
             )
+            return []
+
+        synchronize_for_mpi(self._send_buf)
+        requests = [self._comm.Irecv(buf, source=rank, tag=self.TAG) for buf, rank in self._recvs]
+        requests += [self._comm.Isend(buf, dest=rank, tag=self.TAG) for buf, rank in self._sends]
+        return requests
 
+    def finish(self, array, requests):
+        """Wait for the messages of :meth:`start` and scatter everything received (and the local copies)."""
+        if requests:
+            MPI.Request.Waitall(requests)
+
+        values = array.reshape(-1)
         # scatter everything received (and the local copies) with one kernel
         if self._recv_index.size:
             if self._local_recv.stop > self._local_recv.start:
                 self._recv_buf[self._local_recv] = self._send_buf[self._local_send]
             values[self._recv_index] = self._recv_buf
 
+    def __call__(self, array):
+        self.finish(array, self.start(array))
+
 
 #===============================================================================
 class CartDataExchanger(ABC):
@@ -201,14 +226,18 @@ def _region(starts, buf_shape, array_shape):
         idx += [slice(0, int(n)) for n in array_shape[len(idx):]]
         return tuple(idx)
 
-    def _device_exchange(self, key, array, comm, transfers):
-        """Run the `_DeviceRegionExchange` of `transfers` (built by the callable on first use) on `array`."""
+    def _device_plan(self, key, array, comm, transfers):
+        """The `_DeviceRegionExchange` of `transfers` for `array` (built by the callable on first use)."""
         plans = self.__dict__.setdefault('_device_plans', {})
         plan_key = (key, tuple(array.shape), np.dtype(array.dtype))
         plan = plans.get(plan_key)
         if plan is None:
             plan = plans[plan_key] = _DeviceRegionExchange(comm, transfers(), array.shape, array.dtype)
-        plan(array)
+        return plan
+
+    def _device_exchange(self, key, array, comm, transfers):
+        """Run the `_DeviceRegionExchange` of `transfers` (built by the callable on first use) on `array`."""
+        self._device_plan(key, array, comm, transfers)(array)
 
     def _device_update_ghost_regions_along(self, array, direction):
         """Ghost regions of `direction` (both sides) from the neighbours, as the blocking exchange."""
diff --git a/feectools/ddm/nonblocking_data_exchanger.py b/feectools/ddm/nonblocking_data_exchanger.py
index 29167253a..2eec8fd09 100644
--- a/feectools/ddm/nonblocking_data_exchanger.py
+++ b/feectools/ddm/nonblocking_data_exchanger.py
@@ -76,7 +76,7 @@ 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_update_ghost_regions`)
+        # through contiguous buffers (`_device_ghost_plan`)
         if xp.is_gpu(u):
             return ()
 
@@ -106,7 +106,11 @@ def prepare_communications(self, u):
 
     def start_update_ghost_regions(self, array, requests ):
         if xp.is_gpu(array):
-            self._device_update_ghost_regions(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.
@@ -115,12 +119,14 @@ def start_update_ghost_regions(self, array, 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_update_ghost_regions(self, array):
-        """All ghost regions (every shift, as the persistent requests) of a device array."""
+    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():
@@ -133,7 +139,7 @@ def transfers():
                             self._region(info['recv_starts'], info['buf_shape'], array.shape), info['rank_source']))
             return out
 
-        self._device_exchange('ghost', array, cart._comm_cart, transfers)
+        return self._device_plan('ghost', array, cart._comm_cart, transfers)
 
     # ...
     def start_exchange_assembly_data( self, array ):
diff --git a/feectools/linalg/block.py b/feectools/linalg/block.py
index 3e461202d..3593e080c 100644
--- a/feectools/linalg/block.py
+++ b/feectools/linalg/block.py
@@ -126,6 +126,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,8 +412,16 @@ def update_ghost_regions(self):
 
         req = self.start_update_interface_ghost_regions()
 
-        for vi in self.blocks:
-            vi.update_ghost_regions()
+        if 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)
 
diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py
index 88328f575..8e9aef8cf 100644
--- a/feectools/linalg/solvers.py
+++ b/feectools/linalg/solvers.py
@@ -4,6 +4,7 @@
 
 """
 import cunumpy as xp
+import numpy as np
 from math import sqrt, inf
 
 from feectools.utilities.utils  import is_real
@@ -17,6 +18,7 @@
     'inverse',
     'ConjugateGradient',
     'PConjugateGradient',
+    'PConjugateGradientSingleReduction',
     'BiConjugateGradient',
     'BiConjugateGradientStabilized',
     'PBiConjugateGradientStabilized',
@@ -42,6 +44,39 @@ def _host(s):
     return xp.to_numpy(s)[()] if xp.is_gpu(s) else s
 
 
+def _inner_products(*pairs):
+    """The inner products ``x.inner(y)`` of the ``(x, y)`` pairs, with one reduction.
+
+    Each inner product of distributed vectors is otherwise one Allreduce (one per block for block
+    vectors) and, with device data, one synchronization and copy to the host. Here the local
+    contributions (``space.inner_local``) are collected, copied to the host together (one
+    synchronization) and reduced with a single Allreduce on the spaces' communicator (if any).
+    The results are host scalars, also for device data: the solver's vector updates (axpy) take
+    their coefficient by value, so a device scalar would be copied to the host at every update.
+    Vectors whose spaces do not provide ``inner_local`` / ``inner_comm`` (or whose blocks do not
+    share one communicator) use ``x.inner(y)`` as before.
+    """
+    try:
+        comms = {id(x.space.inner_comm) for x, _ in pairs}
+        if isinstance(pairs[0][0].space, BlockVectorSpace):
+            comms |= {id(Vi.inner_comm) for x, _ in pairs for Vi in x.space.spaces}
+        comm = pairs[0][0].space.inner_comm
+        local = [x.space.inner_local(x, y) for x, y in pairs]
+    except AttributeError:
+        return [_host(x.inner(y)) for x, y in pairs]
+    if len(comms) != 1:
+        return [_host(x.inner(y)) for x, y in pairs]
+
+    if any(xp.is_gpu(v) for v in local):
+        values = xp.to_numpy(xp.stack([xp.asarray(v) for v in local]))  # one copy to the host
+    else:
+        values = np.array(local)
+    if comm is not None and comm.Get_size() > 1:
+        from maybempi import MPI
+        comm.Allreduce(MPI.IN_PLACE, values, op=MPI.SUM)
+    return [values[i] for i in range(len(pairs))]
+
+
 def _sqrt(s):
     """Square root that keeps a device scalar on the device; `math.sqrt` for host scalars."""
     return xp.sqrt(s) if xp.is_gpu(s) else sqrt(s)
@@ -65,7 +100,7 @@ def inverse(A, solver, **kwargs):
         function (e.g. a matrix-vector product A*p).
 
     solver : str
-        Preferred iterative solver. Options are: 'cg', 'pcg', 'bicg',
+        Preferred iterative solver. Options are: 'cg', 'pcg', 'pcg_sr', 'bicg',
         'bicgstab', 'pbicgstab', 'minres', 'lsmr', 'gmres'.
 
     Returns
@@ -81,6 +116,7 @@ def inverse(A, solver, **kwargs):
     solvers_dict = {
         'cg'       : ConjugateGradient,
         'pcg'      : PConjugateGradient,
+        'pcg_sr'   : PConjugateGradientSingleReduction,
         'bicg'     : BiConjugateGradient,
         'bicgstab' : BiConjugateGradientStabilized,
         'pbicgstab': PBiConjugateGradientStabilized,
@@ -374,13 +410,13 @@ def solve(self, b, out=None):
         p = self._tmps["p"]
         s = self._tmps["s"]
 
-        # First values
+        # First values (r.r and s.r with one reduction, as host scalars)
         A.dot(x, out=v)
         b.copy(out=r)
         r       -= v
-        nrmr_sqr = r.inner(r).real
         pc.dot(r, out=s)
-        am       = s.inner(r)
+        nrmr_sqr, am = _inner_products((r, r), (s, r))
+        nrmr_sqr = nrmr_sqr.real
         s.copy(out=p)
 
         tol_sqr  = tol**2
@@ -402,15 +438,15 @@ def solve(self, b, out=None):
                 break
 
             v  = A.dot(p, out=v)
-            l  = am / v.inner(p)
+            l  = am / _inner_products((v, p))[0]
 
             x.mul_iadd(l, p) # this is x += l*p
             r.mul_iadd(-l, v) # this is r -= l*v
 
-            nrmr_sqr = r.inner(r).real
+            # r.r and s.r with one reduction
             pc.dot(r, out=s)
-
-            am1 = s.inner(r)
+            nrmr_sqr, am1 = _inner_products((r, r), (s, r))
+            nrmr_sqr = nrmr_sqr.real
 
             # we are computing p = (am1 / am) * p + s by using axpy on s and exchanging the arrays
             s.mul_iadd((am1/am), p)
@@ -436,6 +472,141 @@ def solve(self, b, out=None):
     def dot(self, b, out=None):
         return self.solve(b, out=out)
 
+#===============================================================================
+class PConjugateGradientSingleReduction(PConjugateGradient):
+    """
+    Preconditioned Conjugate Gradient with one reduction per iteration (Chronopoulos-Gear variant).
+
+    Mathematically the same iterates as :class:`PConjugateGradient`, but all inner products of an
+    iteration, ``(r, u)``, ``(w, u)`` and ``(r, r)`` with ``u = P r`` and ``w = A u``, are computed
+    together, so with distributed vectors each iteration has one Allreduce (and with device data one
+    synchronization) instead of two. ``A p`` is then not computed but updated by the recurrence
+    ``s = w + beta * s``, which costs one more vector update and two more work vectors per
+    iteration. Use it when the global reductions dominate, e.g. many MPI ranks or GPUs; the
+    residual can drift from ``b - A x`` at round-off level, which matters only for tolerances near
+    machine precision.
+
+    Same parameters as :class:`PConjugateGradient`; select it with ``inverse(A, 'pcg_sr', ...)``.
+
+    References
+    ----------
+    A. T. Chronopoulos and C. W. Gear, s-step iterative methods for symmetric linear systems,
+    J. Comput. Appl. Math. 25 (1989), 153-168.
+    """
+    def __init__(self, A, *, pc=None, x0=None, tol=1e-6, maxiter=1000, verbose=False, recycle=False):
+
+        super().__init__(A, pc=pc, x0=x0, tol=tol, maxiter=maxiter, verbose=verbose, recycle=recycle)
+
+        self._tmps["u"] = self.codomain.zeros()
+        self._tmps["w"] = self.domain.zeros()
+
+    def solve(self, b, out=None):
+        """
+        Solve the symmetric positive definite system Ax = b, see :meth:`PConjugateGradient.solve`.
+        Info can be accessed using get_info(), see :func:~`basic.InverseLinearOperator.get_info`.
+
+        Parameters
+        ----------
+        b : feectools.linalg.basic.Vector
+            Right-hand-side vector of linear system.
+
+        out : feectools.linalg.basic.Vector | NoneType
+            The output vector, or None (optional).
+
+        Returns
+        -------
+        x : feectools.linalg.basic.Vector
+            Numerical solution of the linear system.
+        """
+
+        A = self._A
+        options = self._options
+        pc = options["pc"]
+        tol = options["tol"]
+        maxiter = options["maxiter"]
+        verbose = options["verbose"]
+
+        assert isinstance(b, Vector)
+        assert b.space is self._domain
+        assert isinstance(pc, LinearOperator)
+        if out is not None:
+            assert isinstance(out, Vector)
+            assert out.space is self._codomain
+
+        x = options["x0"].copy(out=out)
+
+        # p: search direction, s = A p, u = P r, w = A u
+        r = self._tmps["r"]
+        p = self._tmps["p"]
+        s = self._tmps["s"]
+        u = self._tmps["u"]
+        w = self._tmps["w"]
+        v = self._tmps["v"]
+
+        A.dot(x, out=v)
+        b.copy(out=r)
+        r -= v
+        pc.dot(r, out=u)
+        A.dot(u, out=w)
+        nrmr_sqr, gamma, delta = _inner_products((r, r), (u, r), (w, u))
+        nrmr_sqr = nrmr_sqr.real
+
+        tol_sqr = tol**2
+
+        if verbose:
+            print( "Pre-conditioned CG solver (single reduction):" )
+            print( "+---------+---------------------+")
+            print( "+ Iter. # | L2-norm of residual |")
+            print( "+---------+---------------------+")
+            template = "| {:7d} | {:19.2e} |"
+            print( template.format(1, sqrt(nrmr_sqr)))
+
+        beta = 0.0
+        alpha = 0.0
+        for k in range(2, maxiter+1):
+
+            if nrmr_sqr < tol_sqr:
+                k -= 1
+                break
+
+            if k == 2:
+                alpha = gamma / delta
+                u.copy(out=p)
+                w.copy(out=s)
+            else:
+                alpha = gamma / (delta - beta * gamma / alpha)
+                # p = u + beta * p and s = w + beta * s, by axpy on u and w and exchanging the arrays
+                u.mul_iadd(beta, p)
+                u, p = p, u
+                w.mul_iadd(beta, s)
+                w, s = s, w
+
+            x.mul_iadd(alpha, p)
+            r.mul_iadd(-alpha, s)
+
+            pc.dot(r, out=u)
+            A.dot(u, out=w)
+            gamma_old = gamma
+            nrmr_sqr, gamma, delta = _inner_products((r, r), (u, r), (w, u))
+            nrmr_sqr = nrmr_sqr.real
+            beta = gamma / gamma_old
+
+            if verbose:
+                print( template.format(k, sqrt(nrmr_sqr)))
+
+        if verbose:
+            print( "+---------+---------------------+")
+
+        # keep the exchanged arrays for the next solve
+        self._tmps.update(p=p, s=s, u=u, w=w)
+
+        self._info = {'niter': k, 'success': nrmr_sqr < tol_sqr, 'res_norm': sqrt(nrmr_sqr) }
+
+        if options["recycle"]:
+            x.copy(out=options["x0"])
+
+        return x
+
 #===============================================================================
 class BiConjugateGradient(InverseLinearOperator):
     """
diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py
index 3e446b5c8..58ab4c5c0 100644
--- a/feectools/linalg/stencil.py
+++ b/feectools/linalg/stencil.py
@@ -17,6 +17,7 @@
 
 from maybempi import MPI
 from feectools.linalg.basic  import VectorSpace, Vector, LinearOperator
+from feectools.utilities.cached_launch import CachedLaunch
 from feectools.linalg.memory import stencil_matrix_memory
 from feectools.ddm.cart      import find_mpi_type, CartDecomposition, InterfaceCartDecomposition
 from feectools.ddm.utilities import get_data_exchanger
@@ -210,14 +211,14 @@ def __init__(self, cart, dtype=float):
 
         # Select kernel for AXPY operation
         if self._ndim in [1, 2, 3]:
-            self._axpy_func = stencil_kernels['axpy'][self._ndim]
+            self._axpy_func = CachedLaunch(stencil_kernels['axpy'][self._ndim])
         else:
             self._axpy_func = self._axpy_python
             self._axpy_work = self.zeros()  # work array
 
         # Select kernel for inner product
         if self._ndim in [1, 2, 3]:
-            self._inner_func = stencil_kernels['inner'][self._ndim]
+            self._inner_func = CachedLaunch(stencil_kernels['inner'][self._ndim])
         else:
             self._inner_func = self._inner_python
 
@@ -344,6 +345,31 @@ def inner(self, x, y):
         # a NumPy scalar
         return res[0]
 
+    # ...
+    def inner_local(self, x, y):
+        """The contribution of this process to :meth:`inner` (not reduced over the processes).
+
+        A 0-d device array for device data, a NumPy scalar otherwise. Reduce the contributions
+        of several inner products with one Allreduce on :attr:`inner_comm` to save a reduction
+        (and a device synchronization) per inner product, see
+        :func:`feectools.linalg.solvers._inner_products`.
+        """
+        assert isinstance(x, StencilVector)
+        assert isinstance(y, StencilVector)
+        assert x.space is self
+        assert y.space is self
+
+        res = x._dot_send_data
+        res[0] = 0
+        if x._data.shape[0] != 0:
+            self._inner_func(x._data, y._data, *self._inner_consts, res)
+        return res[0].copy() if xp.is_gpu(res) else res[0]
+
+    @property
+    def inner_comm(self):
+        """The communicator :meth:`inner` reduces over (None if not parallel)."""
+        return self.cart.global_comm if self.parallel else None
+
     # ...
     def axpy(self, a, x, y):
         """
@@ -518,6 +544,82 @@ def set_interface(self, axis, ext, cart):
 
             self._interfaces[axis, ext] = space
 
+#========================================================================
+_SERIAL_GHOST_SRC = r"""
+extern "C" __global__ void serial_ghost_update(double* a, const long long* dst, const long long* src, long long n)
+{
+    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) {
+        long long s = src[i];
+        a[dst[i]] = s >= 0 ? a[s] : 0.0;
+    }
+}
+"""
+
+
+class _SerialGhostPlan:
+    """The serial ghost update (:meth:`StencilVector._update_ghost_regions_serial`) of a device array as one kernel.
+
+    The serial update is up to ``2 * ndim`` strided copies and fills, one launch each. Applied once to an array of
+    flat indices, it gives for every ghost entry the entry its final value comes from (or -1 for zero), corners and
+    the order of the directions included; one kernel then does ``a[dst] = a[src]`` (or 0) for all ghost entries.
+    Valid only if no ghost entry is read (the source entries are all owned), which is checked.
+    """
+
+    def __init__(self, dst, src):
+        import cupy as cp
+        import numpy as np
+
+        self._dst = cp.asarray(dst)
+        self._src = cp.asarray(src)
+        self._n = np.int64(dst.size)
+        self._kernel = cp.RawKernel(_SERIAL_GHOST_SRC, "serial_ghost_update")
+        self._block = (256,)
+        self._grid = ((int(self._n) + 255) // 256,)
+
+    def __call__(self, data):
+        if self._n:
+            self._kernel(self._grid, self._block, (data, self._dst, self._src, self._n))
+
+
+def _serial_ghost_plan(space, data):
+    """The :class:`_SerialGhostPlan` of ``space`` for ``data`` (built on first use), or None if it does not apply.
+
+    Only for C-contiguous float64 device arrays; other data use the slice copies.
+    """
+    import numpy as np
+
+    if not (xp.is_gpu(data) and data.dtype == np.float64 and data.flags.c_contiguous):
+        return None
+    plans = space.__dict__.setdefault('_serial_ghost_plans', {})
+    key = tuple(data.shape)
+    if key in plans:
+        return plans[key]
+
+    # run the serial update on the flat indices (-1: set to zero)
+    idx = np.arange(int(np.prod(key)), dtype=np.int64).reshape(key)
+    for direction in range(space.ndim):
+        p = space.pads[direction] * space.shifts[direction]
+        if p == 0:
+            continue
+        front = [slice(None)] * direction
+        back = [slice(None)] * (idx.ndim - direction - 1)
+        if space.periods[direction]:
+            idx[tuple(front + [slice(-p, None)] + back)] = idx[tuple(front + [slice(p, 2*p)] + back)]
+            idx[tuple(front + [slice(None, p)] + back)] = idx[tuple(front + [slice(-2*p, -p)] + back)]
+        else:
+            idx[tuple(front + [slice(None, p)] + back)] = -1
+            idx[tuple(front + [slice(-p, None)] + back)] = -1
+
+    flat = idx.ravel()
+    dst = np.flatnonzero(flat != np.arange(flat.size))
+    src = flat[dst]
+    # every source must be an entry the update does not write, or the kernel would race
+    plan = None if np.isin(src[src >= 0], dst).any() else _SerialGhostPlan(dst, src)
+    plans[key] = plan
+    return plan
+
+
 #========================================================================
 class StencilVector(Vector):
     """
@@ -835,15 +937,30 @@ def update_ghost_regions(self):
 
         """
 
+        self._start_update_ghost_regions()
+        self._end_update_ghost_regions()
+
+    # ...
+    def _start_update_ghost_regions(self):
+        """First half of :meth:`update_ghost_regions`: start the exchange with the neighbours (if parallel)."""
+        if self.space.parallel and not self.space.cart.is_comm_null:
+            self.space._synchronizer.start_update_ghost_regions(self._data, self._requests)
+
+    # ...
+    def _end_update_ghost_regions(self):
+        """Second half of :meth:`update_ghost_regions`: finish the exchange, update the interfaces, flag in sync."""
         # Update interior ghost regions
         if self.space.parallel:
             if not self.space.cart.is_comm_null:
                 # PARALLEL CASE: fill in ghost regions with data from neighbors
-                self.space._synchronizer.start_update_ghost_regions(self._data, self._requests)
-                self.space._synchronizer.  end_update_ghost_regions(self._data, self._requests)
+                self.space._synchronizer.end_update_ghost_regions(self._data, self._requests)
         else:
             # SERIAL CASE: fill in ghost regions along periodic directions, otherwise set to zero
-            self._update_ghost_regions_serial()
+            plan = _serial_ghost_plan(self._space, self._data)
+            if plan is None:
+                self._update_ghost_regions_serial()
+            else:
+                plan(self._data)
 
         # Update interface ghost regions
         if self.space.parallel:
diff --git a/feectools/linalg/tests/test_solvers.py b/feectools/linalg/tests/test_solvers.py
index e877888f4..332903555 100644
--- a/feectools/linalg/tests/test_solvers.py
+++ b/feectools/linalg/tests/test_solvers.py
@@ -59,7 +59,7 @@ def define_data(n, p, matrix_data, dtype=float):
 @pytest.mark.parametrize( 'n', [5, 10, 13] )
 @pytest.mark.parametrize('p', [2, 3])
 @pytest.mark.parametrize('dtype', [float])
-@pytest.mark.parametrize('solver', ['cg', 'pcg', 'bicg', 'bicgstab', 'pbicgstab', 'minres', 'lsmr', 'gmres'])
+@pytest.mark.parametrize('solver', ['cg', 'pcg', 'pcg_sr', 'bicg', 'bicgstab', 'pbicgstab', 'minres', 'lsmr', 'gmres'])
 
 def test_solver_tridiagonal(n, p, dtype, solver, verbose=False):
 
@@ -82,7 +82,7 @@ def test_solver_tridiagonal(n, p, dtype, solver, verbose=False):
         else:
             diagonals = [-7,-1,-3]
 
-    if solver in ['cg', 'pcg', 'minres']:
+    if solver in ['cg', 'pcg', 'pcg_sr', 'minres']:
         # pcg runs with Jacobi preconditioner
         V, A, xe = define_data_hermitian(n, p, dtype=dtype)
         if solver == 'minres' and dtype == complex:
@@ -106,7 +106,7 @@ def test_solver_tridiagonal(n, p, dtype, solver, verbose=False):
         print()
 
     #Create the solvers
-    if solver in ['pcg', 'pbicgstab']:
+    if solver in ['pcg', 'pcg_sr', 'pbicgstab']:
         pc = A.diagonal(inverse=True)
         solv = inverse(A, solver, pc=pc, tol=1e-13, verbose=verbose, recycle=True)
     else:
@@ -144,7 +144,7 @@ def test_solver_tridiagonal(n, p, dtype, solver, verbose=False):
     assert xp.array_equal(xh.toarray(), solvh_x0.toarray())
     assert xh is not solvh_x0
 
-    if solver != 'pcg':
+    if solver not in ['pcg', 'pcg_sr']:
         # PCG only works with operators with diagonal
         xc = solv2 @ be2
         solv2_x0 = solv2._options["x0"]
@@ -157,7 +157,7 @@ def test_solver_tridiagonal(n, p, dtype, solver, verbose=False):
     b2 = A @ x2
     bt = A.T @ xt
     bh = A.H @ xh
-    if solver != 'pcg':
+    if solver not in ['pcg', 'pcg_sr']:
         bc = A @ A @ xc
 
     err = b - be
@@ -169,7 +169,7 @@ def test_solver_tridiagonal(n, p, dtype, solver, verbose=False):
     errh = bh - beh
     errh_norm = xp.linalg.norm( errh.toarray() )
 
-    if solver != 'pcg': 
+    if solver not in ['pcg', 'pcg_sr']: 
         errc = bc - be2
         errc_norm = xp.linalg.norm( errc.toarray() )
 
@@ -202,7 +202,7 @@ def test_solver_tridiagonal(n, p, dtype, solver, verbose=False):
         assert err2_norm < tol
         assert errt_norm < tol
         assert errh_norm < tol
-        assert solver == 'pcg' or errc_norm < tol
+        assert solver in ['pcg', 'pcg_sr'] or errc_norm < tol
 
 # ===============================================================================
 # SCRIPT FUNCTIONALITY

From 6115208bb191c97f0e09a726271666c2550985a4 Mon Sep 17 00:00:00 2001
From: Max Lindqvist <max.lindqvist@ipp.mpg.de>
Date: Thu, 15 Oct 2026 13:51:08 +0200
Subject: [PATCH 06/10] Optimization

---
 feectools/ddm/basic.py      | 124 ++++++++++++++++++++++++++++++++++++
 feectools/linalg/basic.py   |   5 +-
 feectools/linalg/block.py   |  45 ++++++++++++-
 feectools/linalg/stencil.py |  14 +++-
 4 files changed, 183 insertions(+), 5 deletions(-)

diff --git a/feectools/ddm/basic.py b/feectools/ddm/basic.py
index 346a1d46b..0846a6a31 100644
--- a/feectools/ddm/basic.py
+++ b/feectools/ddm/basic.py
@@ -84,6 +84,10 @@ def layout(groups, local):
             index = np.concatenate(parts) if parts else np.zeros(0, dtype=np.int64)
             return xp.asarray(index), slices, local_slice
 
+        # host index lists, kept for _FusedDeviceExchange (exchanges of several arrays in one)
+        self.comm = comm
+        self.host_regions = (sends, recvs, local_send, local_recv)
+
         self._send_index, send_slices, self._local_send = layout(sends, local_send)
         self._recv_index, recv_slices, self._local_recv = layout(recvs, local_recv)
         self._send_buf = xp.empty(self._send_index.size, dtype=dtype)
@@ -143,6 +147,126 @@ def __call__(self, array):
         self.finish(array, self.start(array))
 
 
+#===============================================================================
+_FUSED_SRC = r"""
+extern "C" __global__ void fused_gather(double* __restrict__ out, const unsigned long long* __restrict__ ptrs,
+                                        const int* __restrict__ comp, const long long* __restrict__ idx, long long n)
+{
+    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) out[i] = ((const double*)ptrs[comp[i]])[idx[i]];
+}
+
+extern "C" __global__ void fused_scatter(const double* __restrict__ in, const unsigned long long* __restrict__ ptrs,
+                                         const int* __restrict__ comp, const long long* __restrict__ idx, long long n)
+{
+    long long i = (long long)blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) ((double*)ptrs[comp[i]])[idx[i]] = in[i];
+}
+"""
+
+
+class _FusedDeviceExchange:
+    """
+    The exchanges of several device arrays (e.g. the components of a block vector) as one.
+
+    Built from one :class:`_DeviceRegionExchange` per array, on the same communicator: the
+    regions that go to the same rank, from every array, travel as one message, gathered by one
+    kernel (and scattered by one) that reads the arrays through a table of their pointers. A
+    block vector of three components thus sends one message per neighbour instead of three, with
+    one gather, one synchronization and one scatter. The arrays must be C-contiguous float64.
+
+    Parameters
+    ----------
+    plans : list of _DeviceRegionExchange
+        One exchange per array, in the order of the arrays passed to :meth:`start`.
+    """
+    TAG = 4243
+
+    def __init__(self, plans):
+        import cupy as cp
+
+        self._comm = plans[0].comm
+
+        def merge(which):
+            """(comp, index) of every region, grouped by rank (sorted), then the local regions; and the slices."""
+            groups = {}
+            local = []
+            for c, plan in enumerate(plans):
+                regions = plan.host_regions
+                for rank, parts in regions[which].items():
+                    groups.setdefault(rank, []).extend((c, p) for p in parts)
+                local.extend((c, p) for p in regions[2 + which])
+            comps, idx, slices, offset = [], [], [], 0
+            for rank in sorted(groups):
+                n = sum(p.size for _, p in groups[rank])
+                for c, p in groups[rank]:
+                    comps.append(np.full(p.size, c, dtype=np.int32))
+                    idx.append(p)
+                slices.append((rank, slice(offset, offset + n)))
+                offset += n
+            local_slice = slice(offset, offset + sum(p.size for _, p in local))
+            for c, p in local:
+                comps.append(np.full(p.size, c, dtype=np.int32))
+                idx.append(p)
+            comp = np.concatenate(comps) if comps else np.zeros(0, dtype=np.int32)
+            index = np.concatenate(idx) if idx else np.zeros(0, dtype=np.int64)
+            return cp.asarray(comp), cp.asarray(index.astype(np.int64)), slices, local_slice
+
+        self._send_comp, self._send_index, send_slices, self._local_send = merge(0)
+        self._recv_comp, self._recv_index, recv_slices, self._local_recv = merge(1)
+        self._send_buf = cp.empty(self._send_index.size, dtype=np.float64)
+        self._recv_buf = cp.empty(self._recv_index.size, dtype=np.float64)
+        self._sends = [(self._send_buf[sl], rank) for rank, sl in send_slices]
+        self._recvs = [(self._recv_buf[sl], rank) for rank, sl in recv_slices]
+        module = cp.RawModule(code=_FUSED_SRC)
+        self._gather = module.get_function("fused_gather")
+        self._scatter = module.get_function("fused_scatter")
+        self._ptr_key = None
+        self._ptrs = None
+
+    def _pointers(self, arrays):
+        import cupy as cp
+
+        key = tuple(int(a.data.ptr) for a in arrays)
+        if key != self._ptr_key:
+            self._ptrs = cp.asarray(np.array(key, dtype=np.uint64))
+            self._ptr_key = key
+        return self._ptrs
+
+    @staticmethod
+    def _launch(kernel, buf, ptrs, comp, idx):
+        n = idx.size
+        if n:
+            kernel(((n + 255) // 256,), (256,), (buf, ptrs, comp, idx, np.int64(n)))
+
+    def start(self, arrays):
+        """Gather the regions of all arrays and post one message per rank; returns the requests."""
+        ptrs = self._pointers(arrays)
+        self._launch(self._gather, self._send_buf, ptrs, self._send_comp, self._send_index)
+        if not (self._sends or self._recvs):
+            return []
+        if get_mpi_cuda_aware() is False:
+            exchange(
+                self._comm,
+                sends=[(buf, rank, self.TAG) for buf, rank in self._sends],
+                receives=[(buf, rank, self.TAG) for buf, rank in self._recvs],
+                cuda_aware=False,
+            )
+            return []
+        synchronize_for_mpi(self._send_buf)
+        requests = [self._comm.Irecv(buf, source=rank, tag=self.TAG) for buf, rank in self._recvs]
+        requests += [self._comm.Isend(buf, dest=rank, tag=self.TAG) for buf, rank in self._sends]
+        return requests
+
+    def finish(self, arrays, requests):
+        """Wait for the messages of :meth:`start` and scatter everything received (and the local copies)."""
+        if requests:
+            MPI.Request.Waitall(requests)
+        if self._local_recv.stop > self._local_recv.start:
+            self._recv_buf[self._local_recv] = self._send_buf[self._local_send]
+        self._launch(self._scatter, self._recv_buf, self._pointers(arrays), self._recv_comp, self._recv_index)
+
+
 #===============================================================================
 class CartDataExchanger(ABC):
     """
diff --git a/feectools/linalg/basic.py b/feectools/linalg/basic.py
index 13eb76e97..a9ea4a25a 100644
--- a/feectools/linalg/basic.py
+++ b/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()
diff --git a/feectools/linalg/block.py b/feectools/linalg/block.py
index 3593e080c..eb3b20e17 100644
--- a/feectools/linalg/block.py
+++ b/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
@@ -412,7 +413,14 @@ def update_ghost_regions(self):
 
         req = self.start_update_interface_ghost_regions()
 
-        if all(hasattr(vi, '_start_update_ghost_regions') and xp.is_gpu(vi._data) for vi in self.blocks):
+        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:
@@ -428,6 +436,41 @@ def update_ghost_regions(self):
         # 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
+        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 = {}
diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py
index 58ab4c5c0..9db8d4265 100644
--- a/feectools/linalg/stencil.py
+++ b/feectools/linalg/stencil.py
@@ -585,12 +585,17 @@ def __call__(self, data):
 def _serial_ghost_plan(space, data):
     """The :class:`_SerialGhostPlan` of ``space`` for ``data`` (built on first use), or None if it does not apply.
 
-    Only for C-contiguous float64 device arrays; other data use the slice copies.
+    Only for C-contiguous float64 device arrays; other data use the slice copies. Also used for the data of a
+    :class:`StencilMatrix` (``space`` its codomain: the slices run over the first ``space.ndim`` axes and take the
+    others whole, as in :meth:`StencilMatrix._update_ghost_regions_serial`), up to ``2**26`` entries: the plan is
+    built from an index array of the size of the data.
     """
     import numpy as np
 
     if not (xp.is_gpu(data) and data.dtype == np.float64 and data.flags.c_contiguous):
         return None
+    if data.size > 2**26:
+        return None
     plans = space.__dict__.setdefault('_serial_ghost_plans', {})
     key = tuple(data.shape)
     if key in plans:
@@ -1636,7 +1641,12 @@ def update_ghost_regions(self):
                 self._synchronizer.end_update_ghost_regions( self._data , self._requests)
         else:
             # SERIAL CASE: fill in ghost regions along periodic directions, otherwise set to zero
-            self._update_ghost_regions_serial()
+            # (on the GPU with one kernel, see _serial_ghost_plan)
+            plan = _serial_ghost_plan(self._codomain, self._data)
+            if plan is None:
+                self._update_ghost_regions_serial()
+            else:
+                plan(self._data)
 
         # Flag ghost regions as up-to-date
         self._sync = True

From fcd8c8acbdffafe6e396bb0fb59591e342a75414 Mon Sep 17 00:00:00 2001
From: Max <max.lindqvist@ipp.mpg.de>
Date: Fri, 16 Oct 2026 09:59:47 +0200
Subject: [PATCH 07/10] Use the stencil kernels directly (the CachedLaunch
 module was never committed)

The import of feectools.utilities.cached_launch made feectools.linalg.stencil, and with it
the whole package, fail to import.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 feectools/linalg/stencil.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py
index 9db8d4265..ee663d15f 100644
--- a/feectools/linalg/stencil.py
+++ b/feectools/linalg/stencil.py
@@ -17,7 +17,6 @@
 
 from maybempi import MPI
 from feectools.linalg.basic  import VectorSpace, Vector, LinearOperator
-from feectools.utilities.cached_launch import CachedLaunch
 from feectools.linalg.memory import stencil_matrix_memory
 from feectools.ddm.cart      import find_mpi_type, CartDecomposition, InterfaceCartDecomposition
 from feectools.ddm.utilities import get_data_exchanger
@@ -211,14 +210,14 @@ def __init__(self, cart, dtype=float):
 
         # Select kernel for AXPY operation
         if self._ndim in [1, 2, 3]:
-            self._axpy_func = CachedLaunch(stencil_kernels['axpy'][self._ndim])
+            self._axpy_func = stencil_kernels['axpy'][self._ndim]
         else:
             self._axpy_func = self._axpy_python
             self._axpy_work = self.zeros()  # work array
 
         # Select kernel for inner product
         if self._ndim in [1, 2, 3]:
-            self._inner_func = CachedLaunch(stencil_kernels['inner'][self._ndim])
+            self._inner_func = stencil_kernels['inner'][self._ndim]
         else:
             self._inner_func = self._inner_python
 

From 8ae781d4db109824a773a7bc6f8f0f6bf114b3ac Mon Sep 17 00:00:00 2001
From: Max <max.lindqvist@ipp.mpg.de>
Date: Fri, 16 Oct 2026 10:06:53 +0200
Subject: [PATCH 08/10] Device paths: no raw kernels on the fake CuPy; PCG with
 one host copy per iteration again

- _serial_ghost_plan and BlockVector._fused_device_exchange return None on cunumpy's fake CuPy
  (tests without a GPU), which cannot launch raw kernels: the slice copies and per-block
  exchanges are used instead.
- PCG computes v.p as a device scalar again. With the fused r.r / s.r reduction it makes one
  host copy per iteration, as on devel-tiny; routing v.p through _inner_products made it two.
- Emulation parity: stencil_dot_3d uses warp shuffles, which serial emulation cannot run; it is
  covered by the GPU parity test.
- test_inner_on_device covers pcg_sr.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 feectools/linalg/block.py                       |  4 ++++
 feectools/linalg/solvers.py                     |  4 +++-
 feectools/linalg/stencil.py                     | 10 ++++++++++
 feectools/linalg/tests/inner_on_device_child.py |  2 +-
 feectools/linalg/tests/test_cuda_emulation.py   | 14 +++++++++++++-
 feectools/linalg/tests/test_inner_on_device.py  |  6 +++---
 6 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/feectools/linalg/block.py b/feectools/linalg/block.py
index eb3b20e17..8e7f686e1 100644
--- a/feectools/linalg/block.py
+++ b/feectools/linalg/block.py
@@ -450,6 +450,10 @@ def _fused_device_exchange(self):
         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
diff --git a/feectools/linalg/solvers.py b/feectools/linalg/solvers.py
index 8e9aef8cf..a3d92c47c 100644
--- a/feectools/linalg/solvers.py
+++ b/feectools/linalg/solvers.py
@@ -438,7 +438,9 @@ def solve(self, b, out=None):
                 break
 
             v  = A.dot(p, out=v)
-            l  = am / _inner_products((v, p))[0]
+            # v.p as a device scalar (no copy to the host: axpy takes it on the device); the host copy of the
+            # iteration is the one of r.r and s.r below
+            l  = am / v.inner(p)
 
             x.mul_iadd(l, p) # this is x += l*p
             r.mul_iadd(-l, v) # this is r -= l*v
diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py
index ee663d15f..9757aad65 100644
--- a/feectools/linalg/stencil.py
+++ b/feectools/linalg/stencil.py
@@ -581,6 +581,13 @@ def __call__(self, data):
             self._kernel(self._grid, self._block, (data, self._dst, self._src, self._n))
 
 
+def _fake_cupy_active():
+    """Whether the device arrays are cunumpy's fake CuPy (NumPy behind the CuPy API), which cannot run raw kernels."""
+    from cunumpy.kernel_testing import fake_cupy_active
+
+    return fake_cupy_active()
+
+
 def _serial_ghost_plan(space, data):
     """The :class:`_SerialGhostPlan` of ``space`` for ``data`` (built on first use), or None if it does not apply.
 
@@ -593,6 +600,9 @@ def _serial_ghost_plan(space, data):
 
     if not (xp.is_gpu(data) and data.dtype == np.float64 and data.flags.c_contiguous):
         return None
+    if _fake_cupy_active():
+        # cunumpy's fake CuPy (tests without a GPU) has no raw kernels: the slice copies
+        return None
     if data.size > 2**26:
         return None
     plans = space.__dict__.setdefault('_serial_ghost_plans', {})
diff --git a/feectools/linalg/tests/inner_on_device_child.py b/feectools/linalg/tests/inner_on_device_child.py
index 6745b730a..fccfda467 100644
--- a/feectools/linalg/tests/inner_on_device_child.py
+++ b/feectools/linalg/tests/inner_on_device_child.py
@@ -76,7 +76,7 @@ def problem():
 
 
 def solver_kwargs(name, A):
-    return {"pc": A.diagonal(inverse=True)} if name in ("pcg", "pbicgstab") else {}
+    return {"pc": A.diagonal(inverse=True)} if name in ("pcg", "pcg_sr", "pbicgstab") else {}
 
 
 def host(a):
diff --git a/feectools/linalg/tests/test_cuda_emulation.py b/feectools/linalg/tests/test_cuda_emulation.py
index 40a1a0be1..65b730abe 100644
--- a/feectools/linalg/tests/test_cuda_emulation.py
+++ b/feectools/linalg/tests/test_cuda_emulation.py
@@ -20,8 +20,20 @@ def arrays(args):
     return [a.copy() for a in args if isinstance(a, np.ndarray)]
 
 
+# kernels that serial emulation cannot run (warp shuffles); the GPU parity test (test_cuda_parity) covers them
+GPU_ONLY = {"stencil_dot_3d": "warp shuffles"}
+
+
 @requires_compiler
-@pytest.mark.parametrize("name", list(PARITY_CASES))
+@pytest.mark.parametrize(
+    "name",
+    [
+        pytest.param(name, marks=pytest.mark.skip(reason=f"uses {GPU_ONLY[name]}, which emulation cannot run"))
+        if name in GPU_ONLY
+        else name
+        for name in PARITY_CASES
+    ],
+)
 def test_emulated_parity(name):
     kernel, spec = CUDA_KERNELS[name], PARITY_CASES[name]
     with xp.use_backend("numpy"):
diff --git a/feectools/linalg/tests/test_inner_on_device.py b/feectools/linalg/tests/test_inner_on_device.py
index 451df1906..cc9ed3435 100644
--- a/feectools/linalg/tests/test_inner_on_device.py
+++ b/feectools/linalg/tests/test_inner_on_device.py
@@ -24,7 +24,7 @@
 from feectools.linalg.stencil import StencilMatrix, StencilVector, StencilVectorSpace
 
 CHILD = Path(__file__).with_name("inner_on_device_child.py")
-DEVICE_SOLVERS = ("cg", "pcg", "bicg", "bicgstab", "pbicgstab")
+DEVICE_SOLVERS = ("cg", "pcg", "pcg_sr", "bicg", "bicgstab", "pbicgstab")
 HOST_SOLVERS = ("minres", "lsmr", "gmres")
 
 
@@ -64,7 +64,7 @@ def test_solver_info_numpy(solver):
     """The convergence information is made of host scalars, as before."""
     with xp.use_backend("numpy"):
         V, A, xe, b = problem()
-        kw = {"pc": A.diagonal(inverse=True)} if solver in ("pcg", "pbicgstab") else {}
+        kw = {"pc": A.diagonal(inverse=True)} if solver in ("pcg", "pcg_sr", "pbicgstab") else {}
         S = inverse(A, solver, tol=1e-12, maxiter=1000, **kw)
         x = S.solve(b)
         info = S.get_info()
@@ -149,7 +149,7 @@ def test_inner_on_gpu():
 def test_one_host_copy_per_iteration_on_gpu(solver):
     with xp.use_backend("cupy"):
         V, A, xe, b = problem()
-        kw = {"pc": A.diagonal(inverse=True)} if solver in ("pcg", "pbicgstab") else {}
+        kw = {"pc": A.diagonal(inverse=True)} if solver in ("pcg", "pcg_sr", "pbicgstab") else {}
         inverse(A, solver, tol=1e-12, maxiter=1000, **kw).solve(b)
         counts = []
         for maxiter in (5, 15):

From b70cb6cf6077dd8ce8dde52ce4dd61d2035d58b0 Mon Sep 17 00:00:00 2001
From: Max <max.lindqvist@ipp.mpg.de>
Date: Fri, 16 Oct 2026 13:18:10 +0200
Subject: [PATCH 09/10] stencil_dot_3d on the fake CuPy: one thread per entry,
 no warp shuffles

The warp-shuffle kernel cannot be emulated one thread at a time, so struphy's tests that emulate
every launch on cunumpy's fake CuPy (test_accumulator_matrix_fake_cupy, the polar and end-to-end
GPU tests) failed at the first stencil product. On the fake CuPy, StencilMatrix uses the previous
one-thread-per-entry kernel (stencil_dot_3d_rows.cu) for device arguments and the usual dispatch
for host arguments. test_emulated_stencil_dot_3d_rows compares it with the pyccel kernel.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .../linalg/kernels/stencil_dot_3d/__init__.py | 27 ++++++++++
 .../stencil_dot_3d/stencil_dot_3d_rows.cu     | 54 +++++++++++++++++++
 feectools/linalg/stencil.py                   |  6 +++
 feectools/linalg/tests/test_cuda_emulation.py | 18 +++++++
 4 files changed, 105 insertions(+)
 create mode 100644 feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_rows.cu

diff --git a/feectools/linalg/kernels/stencil_dot_3d/__init__.py b/feectools/linalg/kernels/stencil_dot_3d/__init__.py
index 3477cdd4d..9c6f64819 100644
--- a/feectools/linalg/kernels/stencil_dot_3d/__init__.py
+++ b/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
 
 # 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)
diff --git a/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_rows.cu b/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_rows.cu
new file mode 100644
index 000000000..45525b47b
--- /dev/null
+++ b/feectools/linalg/kernels/stencil_dot_3d/stencil_dot_3d_rows.cu
@@ -0,0 +1,54 @@
+// One-thread-per-entry CUDA version of stencil_dot_3d (stencil_dot_3d_kernels.py), same arguments in the same order.
+// No warp shuffles: used on cunumpy's fake CuPy, whose CPU emulation runs launches one thread at a time (see
+// __init__.py); on a GPU stencil_dot_3d_cuda.cu is used.
+#include "cunumpy/array_view.cuh"
+#include "cunumpy/index.cuh"
+
+/**
+ * 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 stencil_dot_3d_rows in __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).
+ *
+ * @param mat   matrix data (a 6D view), 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 3), as in the 1D kernel
+ */
+extern "C" __global__ void stencil_dot_3d_rows(Array6D<double> mat, Array3D<double> x, Array3D<double> out,
+                                          const long long* s_in, const long long* p_in, const long long* add,
+                                          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];
+
+    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);
+
+    out(p_out[0] + i1_loc, p_out[1] + i2_loc, p_out[2] + i3_loc) = val;
+}
diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py
index 9757aad65..eb33d7886 100644
--- a/feectools/linalg/stencil.py
+++ b/feectools/linalg/stencil.py
@@ -2182,6 +2182,12 @@ def set_backend(self, backend, precompiled):
 
             # matvec kernel: stencil_dot_<n>d, pyccel or CUDA by backend
             self._func = stencil_kernels['dot'][self._ndim]
+            if self._ndim == 3 and _fake_cupy_active():
+                # the CUDA kernel uses warp shuffles, which the CPU emulation of launches on the fake CuPy
+                # (cunumpy.kernel_testing) cannot run: the same product with one thread per entry
+                from feectools.linalg.kernels.stencil_dot_3d import stencil_dot_3d_fake_cupy
+
+                self._func = stencil_dot_3d_fake_cupy
 
             # parameter for rectangular matrices
             add = [int(end_in >= end_out) for end_in, end_out in zip(self.domain.ends, self.codomain.ends)]
diff --git a/feectools/linalg/tests/test_cuda_emulation.py b/feectools/linalg/tests/test_cuda_emulation.py
index 65b730abe..58189633b 100644
--- a/feectools/linalg/tests/test_cuda_emulation.py
+++ b/feectools/linalg/tests/test_cuda_emulation.py
@@ -48,3 +48,21 @@ def test_emulated_parity(name):
             assert any(not np.array_equal(a, b) for a, b in zip(before, host)), f"{name}{case}: the kernel changed nothing"
             for h, e in zip(host, emulated):
                 np.testing.assert_allclose(e, h, rtol=spec.rtol, atol=spec.atol, err_msg=f"{name}{case}")
+
+
+@requires_compiler
+def test_emulated_stencil_dot_3d_rows():
+    """The one-thread-per-entry stencil_dot_3d (used on the fake CuPy instead of the warp-shuffle kernel) against
+    the pyccel kernel, on the parity cases of stencil_dot_3d."""
+    from feectools.linalg.kernels.stencil_dot_3d import stencil_dot_3d_rows
+
+    kernel, spec = CUDA_KERNELS["stencil_dot_3d"], PARITY_CASES["stencil_dot_3d"]
+    rows = stencil_dot_3d_rows()
+    with xp.use_backend("numpy"):
+        for case in spec.cases:
+            host_args = spec.build(case)
+            emulated_args = spec.build(case)
+            kernel(*host_args)
+            emulate_cuda_kernel(rows, *emulated_args, n_threads=rows.n_threads_from(emulated_args))
+            for h, e in zip(arrays(host_args), arrays(emulated_args)):
+                np.testing.assert_allclose(e, h, rtol=spec.rtol, atol=spec.atol, err_msg=f"stencil_dot_3d_rows{case}")

From 1d8d47c4a3bc8e0b594992cc22944b1ac4230c0f Mon Sep 17 00:00:00 2001
From: Max <max.lindqvist@ipp.mpg.de>
Date: Fri, 16 Oct 2026 17:15:28 +0200
Subject: [PATCH 10/10] CachedLaunch: reuse the prepared launch of a kernel for
 the same buffers

feectools.utilities.cached_launch was imported by stencil.py on this branch but never committed:
the .gitignore rule *cache* ignored it (fcd8c8a dropped its use). The rule now has exceptions for
this module and its test.

A CUDA launch of a cunumpy kernel checks and converts every argument and validates the launch
shape; in a Krylov solver axpy and the inner products are launched with the same buffers every
iteration. CachedLaunch keys a launch by its buffers (pointer, shape, strides, dtype), the types of
its scalars and the launch arguments, prepares it once with the kernel's own checks, and afterwards
converts only the scalars (e.g. the axpy coefficient) with the kernel's checkers and launches the
compiled kernel. Host calls, cunumpy's fake CuPy, argument objects, explicit streams, debug mode
and check_finite go to the wrapped kernel unchanged.

The stencil axpy and inner kernels use it again. test_cached_launch checks reuse, scalar
conversion, re-preparation for new buffers or launch shapes, the pass-through cases and the bound
with a stub kernel (no GPU needed), and that host calls go to the kernel.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 .gitignore                                   |   3 +
 feectools/linalg/stencil.py                  |   5 +-
 feectools/linalg/tests/test_cached_launch.py | 129 +++++++++++++++++
 feectools/utilities/cached_launch.py         | 137 +++++++++++++++++++
 4 files changed, 272 insertions(+), 2 deletions(-)
 create mode 100644 feectools/linalg/tests/test_cached_launch.py
 create mode 100644 feectools/utilities/cached_launch.py

diff --git a/.gitignore b/.gitignore
index 8dcea7023..9db967d3d 100644
--- a/.gitignore
+++ b/.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
diff --git a/feectools/linalg/stencil.py b/feectools/linalg/stencil.py
index eb33d7886..eb65c6473 100644
--- a/feectools/linalg/stencil.py
+++ b/feectools/linalg/stencil.py
@@ -39,6 +39,7 @@
 from feectools.linalg.kernels.stencil_dot_1d       import stencil_dot_1d
 from feectools.linalg.kernels.stencil_dot_2d       import stencil_dot_2d
 from feectools.linalg.kernels.stencil_dot_3d       import stencil_dot_3d
+from feectools.utilities.cached_launch import CachedLaunch
 from feectools.linalg.kernels.stencil_transpose_1d import stencil_transpose_1d
 from feectools.linalg.kernels.stencil_transpose_2d import stencil_transpose_2d
 from feectools.linalg.kernels.stencil_transpose_3d import stencil_transpose_3d
@@ -210,14 +211,14 @@ def __init__(self, cart, dtype=float):
 
         # Select kernel for AXPY operation
         if self._ndim in [1, 2, 3]:
-            self._axpy_func = stencil_kernels['axpy'][self._ndim]
+            self._axpy_func = CachedLaunch(stencil_kernels['axpy'][self._ndim])
         else:
             self._axpy_func = self._axpy_python
             self._axpy_work = self.zeros()  # work array
 
         # Select kernel for inner product
         if self._ndim in [1, 2, 3]:
-            self._inner_func = stencil_kernels['inner'][self._ndim]
+            self._inner_func = CachedLaunch(stencil_kernels['inner'][self._ndim])
         else:
             self._inner_func = self._inner_python
 
diff --git a/feectools/linalg/tests/test_cached_launch.py b/feectools/linalg/tests/test_cached_launch.py
new file mode 100644
index 000000000..d0fbea180
--- /dev/null
+++ b/feectools/linalg/tests/test_cached_launch.py
@@ -0,0 +1,129 @@
+"""CachedLaunch (feectools.utilities.cached_launch): prepared launches are reused only for the same buffers.
+
+Without a GPU the CUDA path is checked with a stub kernel and stub device arrays; host calls go to the wrapped kernel.
+"""
+
+import numpy as np
+import pytest
+
+import feectools.utilities.cached_launch as cached_launch
+from feectools.utilities.cached_launch import CachedLaunch
+
+
+class _Buffer:
+    def __init__(self, ptr):
+        self.ptr = ptr
+
+
+class StubDeviceArray:
+    """Stands in for a cupy.ndarray: pointer, shape, strides, dtype."""
+
+    def __init__(self, ptr, shape=(4,)):
+        self.data = _Buffer(ptr)
+        self.shape = shape
+        self.strides = (8,)
+        self.dtype = np.dtype(float)
+
+
+class StubCudaKernel:
+    """Records the full calls and the raw launches; scalars are converted to float64 by the 'checkers'."""
+
+    def __init__(self, n_args):
+        self.calls = []
+        self.launches = []
+        self.debug = False
+        self.check_finite = False
+        self._signature = object()
+        self._checkers = [np.float64] * n_args
+
+    def debug_active(self):
+        return self.debug
+
+    def __call__(self, *args, **launch):
+        self.calls.append((args, launch))
+
+    def launch_shape(self, n_threads, *, grid=None, block=None, args=None):
+        return (max(1, (n_threads or 1) // 128),), (128,)
+
+    def prepare_args(self, *args):
+        return tuple(a.data.ptr if isinstance(a, StubDeviceArray) else np.float64(a) for a in args)
+
+    def compile(self):
+        def raw(grid, block, values, shared_mem=0):
+            self.launches.append((grid, block, values, shared_mem))
+
+        return raw
+
+
+@pytest.fixture
+def stub(monkeypatch):
+    monkeypatch.setattr(cached_launch, "_real_device_array", lambda a: isinstance(a, StubDeviceArray))
+    kernel = StubCudaKernel(3)
+    monkeypatch.setattr(CachedLaunch, "_cuda_kernel", lambda self, args: kernel if any(
+        isinstance(a, StubDeviceArray) for a in args) else None)
+    return kernel
+
+
+def test_reuses_prepared_launch(stub):
+    launch = CachedLaunch(stub)
+    x, y = StubDeviceArray(100), StubDeviceArray(200)
+
+    launch(2.0, x, y, n_threads=256)
+    assert len(stub.calls) == 1 and launch.n_prepared == 1 and not stub.launches
+
+    # same buffers, another coefficient: no full call, the scalar converted and passed in its slot
+    launch(3.5, x, y, n_threads=256)
+    assert len(stub.calls) == 1 and launch.n_reused == 1
+    grid, block, values, _ = stub.launches[-1]
+    assert values == (np.float64(3.5), 100, 200)
+    assert isinstance(values[0], np.float64)
+    assert grid == (2,) and block == (128,)
+
+
+def test_new_buffers_or_launch_shape_prepare_again(stub):
+    launch = CachedLaunch(stub)
+    x, y = StubDeviceArray(100), StubDeviceArray(200)
+    launch(1.0, x, y, n_threads=256)
+    launch(1.0, x, StubDeviceArray(300), n_threads=256)
+    launch(1.0, x, StubDeviceArray(200, shape=(5,)), n_threads=256)
+    launch(1.0, x, y, n_threads=512)
+    assert launch.n_prepared == 4 and len(stub.calls) == 4 and not stub.launches
+
+
+def test_pass_through(stub):
+    launch = CachedLaunch(stub)
+    x, y = StubDeviceArray(100), StubDeviceArray(200)
+    # an argument object (e.g. a struct), an explicit stream, debug mode, check_finite
+    launch(object(), x, y)
+    launch(1.0, x, y, stream=object())
+    stub.debug = True
+    launch(1.0, x, y)
+    launch(1.0, x, y)
+    stub.debug, stub.check_finite = False, True
+    launch(1.0, x, y)
+    assert len(stub.calls) == 5 and launch.n_prepared == 0 and not stub.launches
+
+
+def test_bounded(stub):
+    launch = CachedLaunch(stub, max_entries=2)
+    a, b, c = (StubDeviceArray(p) for p in (1, 2, 3))
+    for buf in (a, b, c):
+        launch(1.0, buf, buf)
+    assert len(launch._launches) == 2
+    launch(1.0, a, a)  # dropped first: prepared again
+    assert launch.n_prepared == 4
+
+
+def test_host_calls_go_to_the_kernel():
+    """NumPy arguments: the wrapped kernel is called as it is (here the stencil axpy of feectools)."""
+    from feectools.linalg.stencil import stencil_kernels
+
+    axpy = stencil_kernels["axpy"][1]
+    launch = CachedLaunch(axpy)
+    x = np.arange(6.0)
+    y = np.ones(6)
+    expected = y + 2.0 * x
+    launch(2.0, x, y)
+    np.testing.assert_array_equal(y, expected)
+    assert launch.n_prepared == 0 and launch.n_reused == 0
+    assert launch.kernel is axpy
diff --git a/feectools/utilities/cached_launch.py b/feectools/utilities/cached_launch.py
new file mode 100644
index 000000000..f2a1807d4
--- /dev/null
+++ b/feectools/utilities/cached_launch.py
@@ -0,0 +1,137 @@
+"""Kernel launches that reuse their preparation when a kernel is launched with the same buffers again.
+
+A launch of a :class:`cunumpy.kernels.CudaKernel` checks and converts every argument (dtype, contiguity, device,
+array views packed into pointer, shape and strides), computes the launch shape and validates it against the device
+limits. In a Krylov solver the same few kernels (axpy, inner products, banded passes) are launched with the same
+buffers in every iteration, so this preparation is repeated thousands of times for the same result.
+:class:`CachedLaunch` prepares a launch once per set of buffers and launches the compiled kernel directly afterwards.
+"""
+
+from collections import OrderedDict
+
+import numpy as np
+
+__all__ = ("CachedLaunch",)
+
+
+def _real_device_array(value) -> bool:
+    """Whether ``value`` is a CuPy array on a GPU (not cunumpy's fake CuPy)."""
+    module = type(value).__module__
+    if not module.startswith("cupy"):
+        return False
+    from cunumpy.kernel_testing import fake_cupy_active
+
+    return not fake_cupy_active()
+
+
+def _current_device() -> int:
+    """The current CUDA device (0 without CuPy)."""
+    try:
+        import cupy
+    except ImportError:
+        return 0
+    return int(cupy.cuda.runtime.getDevice())
+
+
+class CachedLaunch:
+    """A kernel whose CUDA launches with the same buffers reuse their prepared arguments and launch shape.
+
+    Wraps a :class:`cunumpy.kernels.Kernel` (host or CUDA kernel by its dispatch) or a
+    :class:`cunumpy.kernels.CudaKernel` and is called like it. A CUDA launch is keyed by the buffers (pointer, shape,
+    strides, dtype) and the types of the scalar arguments, and by the launch arguments: the first launch with a key
+    runs the kernel's own checks and conversions and stores them, later ones only convert the scalars (e.g. the
+    coefficient of an axpy, which changes every iteration) with the kernel's checkers and launch the compiled kernel.
+
+    Everything else goes to the wrapped kernel unchanged: host calls, cunumpy's fake CuPy, struct or other argument
+    objects, an explicit stream, and kernels in debug mode or with ``check_finite``.
+
+    Parameters
+    ----------
+    kernel : cunumpy.kernels.Kernel | cunumpy.kernels.CudaKernel
+        The kernel to launch.
+    max_entries : int
+        Number of prepared launches kept (the oldest is dropped first).
+    """
+
+    def __init__(self, kernel, max_entries: int = 64):
+        self._kernel = kernel
+        self._max_entries = int(max_entries)
+        self._launches = OrderedDict()
+        self.n_prepared = 0
+        """Number of launches prepared (cache misses on the GPU)."""
+        self.n_reused = 0
+        """Number of launches that reused a prepared one."""
+
+    @property
+    def kernel(self):
+        """The wrapped kernel."""
+        return self._kernel
+
+    def __getattr__(self, name):
+        return getattr(self._kernel, name)
+
+    def _cuda_kernel(self, args):
+        """The CUDA kernel this call launches on a GPU, or None (the call goes to the wrapped kernel)."""
+        if not any(_real_device_array(a) for a in args):
+            return None
+        from cunumpy.kernels import CudaKernel
+
+        if isinstance(self._kernel, CudaKernel):
+            return self._kernel
+        return getattr(self._kernel, "cuda_kernel", None)
+
+    @staticmethod
+    def _key(args, launch):
+        """The key of a launch, and the positions of its scalar arguments; None if it is not cached."""
+        key, scalars = [], []
+        for i, a in enumerate(args):
+            if _real_device_array(a):
+                key.append(("a", int(a.data.ptr), a.shape, a.strides, a.dtype.str))
+            elif isinstance(a, (bool, int, float, np.generic)) and not isinstance(a, np.void):
+                key.append(("s", type(a)))
+                scalars.append(i)
+            else:
+                return None, None
+        for name in ("n_threads", "grid", "block"):
+            value = launch.get(name)
+            key.append(tuple(value) if isinstance(value, (list, tuple)) else value)
+        key.append(int(launch.get("shared_mem", 0)))
+        key.append(_current_device())
+        return tuple(key), scalars
+
+    def __call__(self, *args, **launch):
+        cuda = self._cuda_kernel(args)
+        if cuda is None or launch.get("stream") is not None or cuda.debug_active() or cuda.check_finite:
+            return self._kernel(*args, **launch)
+        key, scalars = self._key(args, launch)
+        if key is None:
+            return self._kernel(*args, **launch)
+
+        entry = self._launches.get(key)
+        if entry is None:
+            # the first launch with these buffers: the kernel's own checks, conversions and validation
+            result = self._kernel(*args, **launch)
+            grid, block = cuda.launch_shape(
+                launch.get("n_threads"), grid=launch.get("grid"), block=launch.get("block"), args=args
+            )
+            values = [np.asarray(v) if isinstance(v, np.void) else v for v in cuda.prepare_args(*args)]
+            checkers = getattr(cuda, "_checkers", None) if getattr(cuda, "_signature", None) is not None else None
+            if checkers is not None:
+                checkers = {i: checkers[i] for i in scalars}
+            entry = (cuda.compile(), grid, block, values, checkers, int(launch.get("shared_mem", 0)))
+            self._launches[key] = entry
+            if len(self._launches) > self._max_entries:
+                self._launches.popitem(last=False)
+            self.n_prepared += 1
+            return result
+
+        raw, grid, block, values, checkers, shared_mem = entry
+        self._launches.move_to_end(key)
+        if scalars:
+            values = list(values)
+            for i in scalars:
+                values[i] = checkers[i](args[i]) if checkers is not None else args[i]
+        if 0 not in grid:
+            raw(grid, block, tuple(values), shared_mem=shared_mem)
+        self.n_reused += 1
+        return None