Polar splines on the CuPy backend

.github/workflows/test-PR-unit-mpi.yml

@@ -82,3 +82,9 @@ jobs:
           source /struphy_fortran_/env_fortran_/bin/activate
           cd ${{ env.STRUPHY_PATH }}
           mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi pic/tests/test_sorting_device.py
+
+      - name: Run polar GPU MPI paths on the fake CuPy
+        run: |
+          source /struphy_fortran_/env_fortran_/bin/activate
+          cd ${{ env.STRUPHY_PATH }}
+          mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi polar/tests/test_polar_cupy.py -k mpi_fake

doc/sections/subsections/numerics-geomFE.rst

@@ -289,4 +289,12 @@ For weighted mass matrices, use :class:`~struphy.feec.preconditioner.PolarMassMa
 (``precond="PolarMassMatrixPreconditioner"`` in the propagator options);
 it is exact for axisymmetric weights and mappings.
 
+On the CuPy backend, extraction, derivatives, Kronecker operators and these solves keep their
+coefficients on the device. Factorization and mass-preconditioner assembly run once on the host;
+cached device factors apply subsequent solves. Polar MPI collectives support CUDA-aware device
+buffers or host staging through ``cunumpy.mpi.mpi_buffer``. The feectools halo exchanges require
+CUDA-aware MPI on real GPUs. ``SplineFunction.__call__`` dispatches CUDA kernels for marker
+and dense/sparse grid evaluation. The fixed local-grid helper ``eval_tp_fixed_loc`` remains
+host-only; distributed field-evaluation MPI handling and coverage are separate follow-ups.
+
 .. image:: ../../pics/polar_derham.png

src/struphy/feec/preconditioner.py

@@ -549,7 +549,7 @@ def _polar_mass_groups(mass_operator: WeightedMassOperator, reduction: str) -> l
         A_pol = (E @ M12 @ E.T).tocsr()
 
         # unweighted 1d mass matrix in eta_3
-        A_3 = _mass_matrix_1d(mass_operator, bases[comps[0]][2], 2, _ones_1d)[0].toarray()
+        A_3 = xp.to_numpy(_mass_matrix_1d(mass_operator, bases[comps[0]][2], 2, _ones_1d)[0].toarray())
 
         groups.append(Kron2Dx1DGroup(comps, A_pol, A_3))
     return groups

src/struphy/feec/projectors.py

@@ -615,7 +615,7 @@ def solve(self, rhs, transposed=False, apply_bc=False, out=None, x0=None):
         krylov = self._polar_krylov[key]
 
         # relative tolerance w.r.t. the norm of the right-hand side
-        rhs_norm = xp.sqrt(rhs.inner(rhs))
+        rhs_norm = float(xp.sqrt(rhs.inner(rhs)))
         krylov.set_options(x0=x0, tol=self._polar_rtol * rhs_norm if rhs_norm > 0.0 else self._polar_rtol)
         x = krylov.solve(rhs, out=out)
 

src/struphy/feec/psydac_derham.py

@@ -619,12 +619,6 @@ def __init__(
             raise NotImplementedError(
                 "Local projectors (DerhamOptions.local_projectors=True) are not supported on the CuPy backend yet."
             )
-        if polar_splines and (xp.get_backend() == "cupy"):
-            # PolarExtractionBlocksC1 builds SciPy sparse matrices from the control points, and the polar
-            # extraction operators apply them to the stencil data, which lives on the device on CuPy.
-            raise NotImplementedError(
-                "Polar splines (DerhamOptions.polar_splines=True) are not supported on the CuPy backend yet."
-            )
 
         # number of elements and spline degrees in each direction
         assert len(num_elements) == 3
@@ -2782,16 +2776,17 @@ def _restart_extraction_op(self):
         n_comps = W.n_comps
 
         # stack blocks of E (polar coeffs x polar rings); incompatible blocks (None) are zero
-        rows = xp.cumsum([0] + list(W.n_polar))
-        cols = xp.cumsum([0] + [n_r * n_2 for n_r, n_2 in zip(W.n_rings, W.n2)])
-        E_full = xp.zeros((rows[-1], cols[-1]), dtype=float)
+        # (host setup data on every backend, like the blocks of E)
+        rows = np.cumsum([0] + list(W.n_polar))
+        cols = np.cumsum([0] + [n_r * n_2 for n_r, n_2 in zip(W.n_rings, W.n2)])
+        E_full = np.zeros((rows[-1], cols[-1]), dtype=float)
         for m in range(n_comps):
             for n in range(n_comps):
                 if E.blocks_ten_to_pol[m][n] is not None:
                     E_full[rows[m] : rows[m + 1], cols[n] : cols[n + 1]] = E.blocks_ten_to_pol[m][n].toarray()
 
         # E^T has full column rank, hence pinv(E^T) @ E^T = identity on polar coeffs
-        L_full = xp.linalg.pinv(E_full.T)
+        L_full = np.linalg.pinv(E_full.T)
 
         blocks = [
             [

src/struphy/feec/tests/test_derham_gpu.py

@@ -67,20 +67,3 @@ def test_local_projectors_not_supported_on_cupy():
     """Local projectors have no device implementation yet; the Derham fails when it is created."""
     with cunumpy.use_backend("cupy"), pytest.raises(NotImplementedError, match="Local projectors"):
         make_derham(local_projectors=True)
-
-
-@requires_cupy
-def test_polar_splines_not_supported_on_cupy():
-    """Polar splines have no device implementation yet; the Derham fails when it is created."""
-    from maybempi import MPI
-
-    from struphy import domains
-    from struphy.feec.psydac_derham import Derham
-    from struphy.io.options import DerhamOptions
-    from struphy.topology.grids import TensorProductGrid
-
-    with cunumpy.use_backend("cupy"):
-        domain = domains.IGAPolarCylinder(num_elements=(8, 6), degree=(2, 3))
-        options = DerhamOptions(degree=(2, 3, 1), bcs=(("dirichlet", "free"), None, None), polar_splines=True)
-        with pytest.raises(NotImplementedError, match="Polar splines"):
-            Derham(TensorProductGrid(num_elements=(8, 6, 4)), options, comm=MPI.COMM_WORLD, domain=domain)

src/struphy/polar/basic.py

@@ -1,6 +1,7 @@
 from typing import TYPE_CHECKING
 
 import cunumpy as xp
+from cunumpy.mpi import mpi_buffer
 from feectools.linalg.basic import Vector, VectorSpace
 from feectools.linalg.block import BlockVector
 from feectools.linalg.stencil import StencilVector
@@ -346,8 +347,9 @@ def toarray(self, allreduce=False):
             out = self.tp.toarray()[self.space.n_rings[0] * self.space.n2[0] * self.space.n3[0] :]
 
             # allreduce tensor-product part
-            if self.space.comm is not None and allreduce:
-                self.space.comm.Allreduce(MPI.IN_PLACE, out, op=MPI.SUM)
+            if self.space.comm is not None and self.space.comm.Get_size() > 1 and allreduce:
+                with mpi_buffer(out, recv=True) as buf:
+                    self.space.comm.Allreduce(MPI.IN_PLACE, buf, op=MPI.SUM)
 
             out = xp.concatenate((self.pol[0].flatten(), out))
 
@@ -357,10 +359,10 @@ def toarray(self, allreduce=False):
             out3 = self.tp[2].toarray()[self.space.n_rings[2] * self.space.n2[2] * self.space.n3[2] :]
 
             # allreduce tensor-product part
-            if self.space.comm is not None and allreduce:
-                self.space.comm.Allreduce(MPI.IN_PLACE, out1, op=MPI.SUM)
-                self.space.comm.Allreduce(MPI.IN_PLACE, out2, op=MPI.SUM)
-                self.space.comm.Allreduce(MPI.IN_PLACE, out3, op=MPI.SUM)
+            if self.space.comm is not None and self.space.comm.Get_size() > 1 and allreduce:
+                for array in (out1, out2, out3):
+                    with mpi_buffer(array, recv=True) as buf:
+                        self.space.comm.Allreduce(MPI.IN_PLACE, buf, op=MPI.SUM)
 
             out = xp.concatenate(
                 (