diff --git a/feectools b/feectools index 0c548d17a..c3c7d941e 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 0c548d17a3202bdd0531b0f9d81abccb27d8887e +Subproject commit c3c7d941e97f5adb40dbe6e6704d261fb48c3a11 diff --git a/pyproject.toml b/pyproject.toml index c106ba648..e00ff8276 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index c3c9da0c7..ce27a6eb6 100644 --- a/src/struphy/feec/preconditioner.py +++ b/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: diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 1c3c19948..63a293ce9 100644 --- a/src/struphy/io/options.py +++ b/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[ diff --git a/src/struphy/models/tests/gpu_e2e.py b/src/struphy/models/tests/gpu_e2e.py index ecfa0c9a2..d58e20af4 100644 --- a/src/struphy/models/tests/gpu_e2e.py +++ b/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"): diff --git a/src/struphy/propagators/curl_curl_solve.py b/src/struphy/propagators/curl_curl_solve.py index 7628178ac..f6b8e5b2f 100644 --- a/src/struphy/propagators/curl_curl_solve.py +++ b/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, diff --git a/src/struphy/propagators/gyrokinetic_poisson_solve.py b/src/struphy/propagators/gyrokinetic_poisson_solve.py index d7ab0fb98..a8b0a336a 100644 --- a/src/struphy/propagators/gyrokinetic_poisson_solve.py +++ b/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: diff --git a/src/struphy/propagators/implicit_diffusion.py b/src/struphy/propagators/implicit_diffusion.py index f6638ecf2..20f5bf2eb 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/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, diff --git a/src/struphy/propagators/poisson_solve.py b/src/struphy/propagators/poisson_solve.py index c0b688be7..3e3a80766 100644 --- a/src/struphy/propagators/poisson_solve.py +++ b/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 diff --git a/src/struphy/propagators/shear_alfven_curl_curl_propagator.py b/src/struphy/propagators/shear_alfven_curl_curl_propagator.py index d9b594a42..faad48b17 100644 --- a/src/struphy/propagators/shear_alfven_curl_curl_propagator.py +++ b/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: diff --git a/src/struphy/propagators/shear_alfven_current_coupling_5d.py b/src/struphy/propagators/shear_alfven_current_coupling_5d.py index 8df3bb652..f840d34b8 100644 --- a/src/struphy/propagators/shear_alfven_current_coupling_5d.py +++ b/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. diff --git a/src/struphy/propagators/shear_alfven_propagator.py b/src/struphy/propagators/shear_alfven_propagator.py index 3da449e0e..8841dd2de 100644 --- a/src/struphy/propagators/shear_alfven_propagator.py +++ b/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() diff --git a/src/struphy/propagators/variational_density_evolve.py b/src/struphy/propagators/variational_density_evolve.py index d929f35b4..96e2af052 100644 --- a/src/struphy/propagators/variational_density_evolve.py +++ b/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, diff --git a/src/struphy/propagators/variational_momentum_advection.py b/src/struphy/propagators/variational_momentum_advection.py index 65a6e7cf2..a0b37b897 100644 --- a/src/struphy/propagators/variational_momentum_advection.py +++ b/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,