GPU stack 1/8: update feectools (device ghost exchange, single-reduction PCG); pcg_sr solver option

pyproject.toml

@@ -27,7 +27,7 @@ dependencies = [
     "cunumpy>=0.6.3, <=0.7.0",
     "pyccel>=2.2.0, <=2.2.3",
     "maybempi>=0.1.2",
-    "feectools>=0.9.0, <=0.9.0",
+    "feectools>=0.10.0, <=0.10.0",
     "scipy<=1.18.0",
     "h5py<=3.16.0",
     "h5netcdf<=1.8.1",

src/struphy/feec/preconditioner.py

@@ -1944,7 +1944,7 @@ def supported(derham) -> bool:
     @staticmethod
     def resolve(precond: str, derham, algo: str, solver: str, u_space: str) -> str:
         """The preconditioner to use for the option ``precond`` of a shear Alfven propagator:
-        ``"StiffnessPreconditioner"`` only if its requirements hold (implicit algorithm with ``solver="pcg"``,
+        ``"StiffnessPreconditioner"`` only if its requirements hold (implicit algorithm with ``solver="pcg"`` or ``"pcg_sr"``,
         ``u_space="Hdiv"`` and :meth:`supported`), else the mass matrix preconditioner of the velocity space
         (``"PolarMassMatrixPreconditioner"`` with polar splines), with a warning. Other options are returned as
         they are."""
@@ -1953,8 +1953,8 @@ def resolve(precond: str, derham, algo: str, solver: str, u_space: str) -> str:
         missing = []
         if algo != "implicit":
             missing += ['algo="implicit"']
-        if solver != "pcg":
-            missing += ['solver="pcg"']
+        if solver not in ("pcg", "pcg_sr"):
+            missing += ['solver="pcg" or "pcg_sr"']
         if u_space != "Hdiv":
             missing += ['u_space="Hdiv"']
         if derham.polar_splines:

src/struphy/io/options.py

@@ -74,7 +74,7 @@ class LiteralOptions:
     GivenInBasis = Literal["0", "1", "2", "3", "v", "physical", "physical_at_eta", "norm", None]
 
     # solvers
-    OptsSymmSolver = Literal["pcg", "cg"]
+    OptsSymmSolver = Literal["pcg", "pcg_sr", "cg"]
     OptsGenSolver = Literal["pbicgstab", "bicgstab", "gmres"]
     OptsMassPrecond = Literal["MassMatrixPreconditioner", "PolarMassMatrixPreconditioner", None]
     OptsShearAlfvenPrecond = Literal[

src/struphy/models/tests/gpu_e2e.py

@@ -66,13 +66,18 @@ def device_context():
         yield
 
 
+# most inner products a feectools solver iteration copies to the host at once (pipelined PCG: 3)
+MAX_FUSED_REDUCTIONS = 3
+
+
 @dataclass
 class TransferBudget:
     """Transfers recorded per phase of the time loop (the setup before it is not counted).
 
-    ``integrate`` (the propagators) may only copy scalars to the host: the convergence test of the Krylov solvers
-    reads the 8-byte residual norm once per iteration (feectools keeps everything else of the solve on the device).
-    No arrays, no host-to-device copies, no host kernels. ``diagnostics`` (scalars, saved markers, binned
+    ``integrate`` (the propagators) may only copy scalars to the host: the Krylov solvers copy the inner products
+    of one iteration (residual norm and step coefficients, at most ``MAX_FUSED_REDUCTIONS`` of them) to the host
+    together, once per iteration (feectools keeps everything else of the solve on the device). No other arrays, no
+    host-to-device copies, no host kernels. ``diagnostics`` (scalars, saved markers, binned
     distribution functions) may copy small results to the host: at most one scalar per tracked scalar and call,
     never a host-to-device copy or a host kernel. ``output`` (``DataContainer.save_data``) may copy each saved
     dataset to the host once, nothing more.
@@ -99,18 +104,21 @@ def check(self, n_scalars: int):
         report = self.report()
         assert self.n_integrate_calls == self.n_steps, report
 
-        # 1. the propagators: only scalars to the host (solver convergence tests), nothing else
+        # 1. the propagators: only (fused) scalars to the host (solver inner products), nothing else
         for e in self.integrate.events:
             if e.kind == "to_host":
-                assert e.nbytes is not None and e.nbytes <= 8, report
+                assert e.nbytes is not None and e.nbytes <= 8 * MAX_FUSED_REDUCTIONS, report
             else:
                 assert e.kind not in HOST_DEVICE_KINDS, report
 
         # 2. diagnostics: no host-to-device copies, no host kernels, only scalar-sized downloads
         for kind in ("to_device", "kernel_conversion", "fallback"):
             assert self.diagnostics.count(kind) == 0, report
-        assert self.diagnostics.to_host <= n_scalars * self.n_diagnostics_calls, report
-        assert self.diagnostics.bytes_to_host <= 8 * n_scalars * self.n_diagnostics_calls, report
+        # MPI scalar reductions copy one value per scalar, followed by the consolidated scalar buffer.
+        max_diagnostic_copies = (n_scalars + 1) * self.n_diagnostics_calls
+        assert self.diagnostics.to_host <= max_diagnostic_copies, report
+        max_diagnostic_bytes = 16 * n_scalars * self.n_diagnostics_calls
+        assert self.diagnostics.bytes_to_host <= max_diagnostic_bytes, report
 
         # 3. output: each saved dataset once to the host
         for kind in ("to_device", "kernel_conversion", "fallback"):

src/struphy/propagators/curl_curl_solve.py

@@ -191,7 +191,7 @@ def __post_init__(self):
             check_option(self.solver, LiteralOptions.OptsSymmSolver)
             check_option(self.precond, LiteralOptions.OptsCurlCurlPrecond)
             if self.precond is not None:
-                assert self.solver == "pcg", f"precond={self.precond!r} requires solver='pcg'."
+                assert self.solver in ("pcg", "pcg_sr"), f"precond={self.precond!r} requires solver='pcg' or 'pcg_sr'."
 
             # defaults
             if self.precond_params is None:
@@ -323,8 +323,8 @@ def verify_rhs(j) -> Vector | FEECVariable | AccumulatorVector:
             pc_class = getattr(preconditioner, precond)
             pc = pc_class(self.mass_ops.M1, **self.options.precond_params)
 
-        # (only "pcg" takes a preconditioner)
-        pc_kwargs = {"pc": pc} if self.options.solver == "pcg" else {}
+        # (only "pcg" and "pcg_sr" take a preconditioner)
+        pc_kwargs = {"pc": pc} if self.options.solver in ("pcg", "pcg_sr") else {}
         self._solver = inverse(
             lhs,
             self.options.solver,

src/struphy/propagators/gyrokinetic_poisson_solve.py

@@ -178,7 +178,7 @@ def __post_init__(self):
             check_option(self.precond, LiteralOptions.OptsDiffusionPrecond)
             check_option(self.which_geometry, self.OptsGeometry)
             if self.precond is not None:
-                assert self.solver == "pcg", f"precond={self.precond!r} requires solver='pcg'."
+                assert self.solver in ("pcg", "pcg_sr"), f"precond={self.precond!r} requires solver='pcg' or 'pcg_sr'."
 
             # defaults
             if self.precond_params is None:

src/struphy/propagators/implicit_diffusion.py

@@ -217,7 +217,7 @@ class Options(OptionsBase):
             - ``None``: no preconditioner.
             - ``"MultiGrid"``: geometric multigrid V-cycle
               (:class:`~struphy.linear_algebra.multigrid.preconditioner.MultiGridPreconditioner`,
-              requires ``solver="pcg"``).
+              requires ``solver="pcg"`` or ``"pcg_sr"``).
             - ``"MassMatrixPreconditioner"``: Kronecker
               approximation of the inverse of the 0-form mass matrix ``M0``
               (see :mod:`struphy.feec.preconditioner`).
@@ -300,7 +300,7 @@ def __post_init__(self):
             check_option(self.solver, LiteralOptions.OptsSymmSolver)
             check_option(self.precond, LiteralOptions.OptsDiffusionPrecond)
             if self.precond is not None:
-                assert self.solver == "pcg", f"precond={self.precond!r} requires solver='pcg'."
+                assert self.solver in ("pcg", "pcg_sr"), f"precond={self.precond!r} requires solver='pcg' or 'pcg_sr'."
             if self.diffusion_mat is None:
                 assert self.precond not in ("MultiGrid", "StiffnessPreconditioner"), (
                     f"precond={self.precond!r} needs a diffusion term (diffusion_mat=None)."
@@ -470,8 +470,8 @@ def verify_rhs(rho) -> StencilVector | FEECVariable | AccumulatorVector:
             pc = pc_class(self.mass_ops.M0, **self.options.precond_params)
 
         # solver just with A_2, but will be set during call with dt
-        # (only "pcg" takes a preconditioner)
-        pc_kwargs = {"pc": pc} if self.options.solver == "pcg" else {}
+        # (only "pcg" and "pcg_sr" take a preconditioner)
+        pc_kwargs = {"pc": pc} if self.options.solver in ("pcg", "pcg_sr") else {}
         self._solver = inverse(
             stab_mat if self._diffusion_op is None else self._diffusion_op,
             self.options.solver,

src/struphy/propagators/poisson_solve.py

@@ -145,7 +145,7 @@ def __post_init__(self):
             check_option(self.solver, LiteralOptions.OptsSymmSolver)
             check_option(self.precond, LiteralOptions.OptsDiffusionPrecond)
             if self.precond is not None:
-                assert self.solver == "pcg", f"precond={self.precond!r} requires solver='pcg'."
+                assert self.solver in ("pcg", "pcg_sr"), f"precond={self.precond!r} requires solver='pcg' or 'pcg_sr'."
             _parse_directions(self.stab_average)
 
             # defaults

src/struphy/propagators/shear_alfven_curl_curl_propagator.py

@@ -187,7 +187,7 @@ def __post_init__(self):
                 if self.algo == "explicit":
                     self.precond = "MassMatrixPreconditioner"
                 else:
-                    assert self.solver == "pcg", 'precond="HiptmairXu" requires solver="pcg".'
+                    assert self.solver in ("pcg", "pcg_sr"), 'precond="HiptmairXu" requires solver="pcg" or "pcg_sr".'
 
             # defaults
             if self.solver_params is None:

src/struphy/propagators/shear_alfven_current_coupling_5d.py

@@ -134,7 +134,7 @@ class Options(OptionsBase):
             :class:`~struphy.feec.preconditioner.ShearAlfvenSchurPreconditioner`, built from the equilibrium part
             :math:`\mathbb T` of the basis projection operator; with ``nonlinear=True`` it approximates the Schur
             operator with :math:`\mathbb T + \mathbb T^B`, since :math:`|\tilde{\mathbf B}| \ll |\mathbf B_0|`). The
-            Kronecker preconditioner needs ``algo="implicit"``, ``solver="pcg"``, ``u_space="Hdiv"``, tensor-product
+            Kronecker preconditioner needs ``algo="implicit"``, ``solver="pcg"`` or ``"pcg_sr"``, ``u_space="Hdiv"``, tensor-product
             splines (no polar splines) and the global commuting projectors; otherwise it falls back (with a warning) to
             ``"MassMatrixPreconditioner"`` (``"PolarMassMatrixPreconditioner"`` with polar splines). ``"MultiGrid"`` is
             not available here.

src/struphy/propagators/shear_alfven_propagator.py

@@ -99,7 +99,7 @@ class Options(OptionsBase):
             (fast diagonalization) preconditioner (``"StiffnessPreconditioner"``, the default,
             :class:`~struphy.feec.preconditioner.ShearAlfvenSchurPreconditioner`: block Jacobi in the velocity
             components with the auxiliary space correction of the :math:`\mathbf E \times \mathbf B` flows if
-            ``exb_correction``, see there). Both need the implicit algorithm with ``solver="pcg"``. The Kronecker
+            ``exb_correction``, see there). Both need the implicit algorithm with ``solver="pcg"`` or ``"pcg_sr"``. The Kronecker
             preconditioner also needs ``u_space="Hdiv"``, tensor-product splines (no polar splines) and the global
             commuting projectors; otherwise ``"StiffnessPreconditioner"`` falls back (with a warning) to
             ``"MassMatrixPreconditioner"`` (``"PolarMassMatrixPreconditioner"`` with polar splines). The Schur operator
@@ -147,7 +147,7 @@ def __post_init__(self):
             check_option(self.precond, LiteralOptions.OptsShearAlfvenPrecond)
             if self.precond == "MultiGrid":
                 assert self.algo == "implicit", f'precond="{self.precond}" requires algo="implicit".'
-                assert self.solver == "pcg", f'precond="{self.precond}" requires solver="pcg".'
+                assert self.solver in ("pcg", "pcg_sr"), f'precond="{self.precond}" requires solver="pcg" or "pcg_sr".'
             if self.precond == "MultiGrid" and self.multigrid is None:
                 self.multigrid = MultiGridOptions()
 

src/struphy/propagators/variational_density_evolve.py

@@ -223,7 +223,9 @@ def allocate(self):
         self._Mrho = kinetic_metric(self.derham, self.mass_ops, self.basis_ops, self.divdiv_alpha)
         # "cg" takes no preconditioner; the density-weighted mass matrix also preconditions the regularized metric
         self._Mrho_pc = (
-            MassMatrixPreconditioner(self.mass_ops.WMMnew, dim_reduce=None) if self.options.solver == "pcg" else None
+            MassMatrixPreconditioner(self.mass_ops.WMMnew, dim_reduce=None)
+            if self.options.solver in ("pcg", "pcg_sr")
+            else None
         )
         self._Mrho_inv = inverse(
             self._Mrho,
@@ -442,7 +444,7 @@ def _initialize_projectors_and_mass(self):
             self.mass_ops.Mv,
             self.options.solver,
             # "cg" takes no preconditioner
-            **({"pc": self.pc_Mv} if self.options.solver == "pcg" else {}),
+            **({"pc": self.pc_Mv} if self.options.solver in ("pcg", "pcg_sr") else {}),
             tol=1e-10,
             maxiter=1000,
             verbose=False,

src/struphy/propagators/variational_momentum_advection.py

@@ -150,7 +150,9 @@ def allocate(self):
             )
         # "cg" takes no preconditioner
         self._Mrho_pc = (
-            MassMatrixPreconditioner(self.mass_ops.WMMnew, dim_reduce=None) if self.options.solver == "pcg" else None
+            MassMatrixPreconditioner(self.mass_ops.WMMnew, dim_reduce=None)
+            if self.options.solver in ("pcg", "pcg_sr")
+            else None
         )
         self._Mrho_inv = inverse(
             self._Mrho,
@@ -308,7 +310,7 @@ def _initialize_mass(self):
             self.mass_ops.Mv,
             self.options.solver,
             # "cg" takes no preconditioner
-            **({"pc": self.pc_Mv} if self.options.solver == "pcg" else {}),
+            **({"pc": self.pc_Mv} if self.options.solver in ("pcg", "pcg_sr") else {}),
             tol=1e-10,
             maxiter=1000,
             verbose=False,