From afe80fc806673d55212f691b0f7dc86da40d32f5 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 09:59:59 +0200 Subject: [PATCH 01/12] Update feectools (device ghost exchange plans, single-reduction PCG); add the pcg_sr solver option feectools: struphy-hub/feectools#104 (fix-cuda-aware-mpi-issues into devel-tiny). pcg_sr (PConjugateGradientSingleReduction) is accepted wherever pcg is, including the preconditioned solvers of the Poisson, diffusion, curl-curl and variational propagators. Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- src/struphy/io/options.py | 2 +- src/struphy/propagators/curl_curl_solve.py | 6 +++--- src/struphy/propagators/gyrokinetic_poisson_solve.py | 2 +- src/struphy/propagators/implicit_diffusion.py | 6 +++--- src/struphy/propagators/poisson_solve.py | 2 +- src/struphy/propagators/variational_density_evolve.py | 4 ++-- src/struphy/propagators/variational_momentum_advection.py | 4 ++-- 8 files changed, 14 insertions(+), 14 deletions(-) diff --git a/feectools b/feectools index 0c548d17a..8ae781d4d 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 0c548d17a3202bdd0531b0f9d81abccb27d8887e +Subproject commit 8ae781d4db109824a773a7bc6f8f0f6bf114b3ac diff --git a/src/struphy/io/options.py b/src/struphy/io/options.py index 4f8919ed6..dadcc4800 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["MassMatrixPreconditioner", "PolarMassMatrixPreconditioner", "MultiGrid", None] 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 b5bc5a89e..18ab5eb5c 100644 --- a/src/struphy/propagators/gyrokinetic_poisson_solve.py +++ b/src/struphy/propagators/gyrokinetic_poisson_solve.py @@ -190,7 +190,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 c6e01c065..46143adc6 100644 --- a/src/struphy/propagators/implicit_diffusion.py +++ b/src/struphy/propagators/implicit_diffusion.py @@ -280,7 +280,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'." preconditioner._parse_directions(self.stab_average) # defaults @@ -436,8 +436,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( self._diffusion_op, self.options.solver, diff --git a/src/struphy/propagators/poisson_solve.py b/src/struphy/propagators/poisson_solve.py index 79137400c..7a765b811 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/variational_density_evolve.py b/src/struphy/propagators/variational_density_evolve.py index 2b54296f7..cd98c37f9 100644 --- a/src/struphy/propagators/variational_density_evolve.py +++ b/src/struphy/propagators/variational_density_evolve.py @@ -207,7 +207,7 @@ def allocate(self): self._Mrho, self.options.solver, # "cg" takes no preconditioner - **({"pc": pc} if self.options.solver == "pcg" else {}), + **({"pc": pc} if self.options.solver in ("pcg", "pcg_sr") else {}), tol=1e-16, maxiter=500, recycle=True, @@ -416,7 +416,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-16, maxiter=1000, verbose=False, diff --git a/src/struphy/propagators/variational_momentum_advection.py b/src/struphy/propagators/variational_momentum_advection.py index 7fc6e64ce..4ba1a5b65 100644 --- a/src/struphy/propagators/variational_momentum_advection.py +++ b/src/struphy/propagators/variational_momentum_advection.py @@ -135,7 +135,7 @@ def allocate(self): self._Mrho, self.options.solver, # "cg" takes no preconditioner - **({"pc": pc} if self.options.solver == "pcg" else {}), + **({"pc": pc} if self.options.solver in ("pcg", "pcg_sr") else {}), tol=1e-16, maxiter=500, recycle=True, @@ -284,7 +284,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-16, maxiter=1000, verbose=False, From aaf14a7f7733c80fb3399f9d73bccaa7cec9efa7 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 13:18:19 +0200 Subject: [PATCH 02/12] Update feectools: stencil_dot_3d without warp shuffles on the fake CuPy feectools b70cb6c (struphy-hub/feectools#104): the emulated GPU tests (fake CuPy) failed at the first stencil product. Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index 8ae781d4d..b70cb6cf6 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 8ae781d4db109824a773a7bc6f8f0f6bf114b3ac +Subproject commit b70cb6cf6077dd8ce8dde52ce4dd61d2035d58b0 From 8c7d0c62af43d690aecaf72394b864c8c9409a87 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:15:37 +0200 Subject: [PATCH 03/12] Update feectools: CachedLaunch for the stencil axpy and inner kernels feectools 1d8d47c (struphy-hub/feectools#104). Co-Authored-By: Claude Opus 5.5 --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index b70cb6cf6..1d8d47c4a 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit b70cb6cf6077dd8ce8dde52ce4dd61d2035d58b0 +Subproject commit 1d8d47c4a3bc8e0b594992cc22944b1ac4230c0f From 0166e833806bbb71af1c7339a1b2f3c2aa4c1803 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 16 Oct 2026 17:41:05 +0200 Subject: [PATCH 04/12] pcg_sr also for the multigrid and Hiptmair-Xu preconditioners of the shear Alfven solves ShearAlfven and ShearAlfvenCurlCurl (#778, #780) accepted these preconditioners only with solver='pcg'; pcg_sr takes a preconditioner as well, as in the other propagators. Co-Authored-By: Claude Opus 5.5 --- src/struphy/propagators/shear_alfven_curl_curl_propagator.py | 2 +- src/struphy/propagators/shear_alfven_propagator.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/struphy/propagators/shear_alfven_curl_curl_propagator.py b/src/struphy/propagators/shear_alfven_curl_curl_propagator.py index 80d21555e..729ae0f95 100644 --- a/src/struphy/propagators/shear_alfven_curl_curl_propagator.py +++ b/src/struphy/propagators/shear_alfven_curl_curl_propagator.py @@ -185,7 +185,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_propagator.py b/src/struphy/propagators/shear_alfven_propagator.py index 70d4cd293..cc6004ec4 100644 --- a/src/struphy/propagators/shear_alfven_propagator.py +++ b/src/struphy/propagators/shear_alfven_propagator.py @@ -148,7 +148,7 @@ def __post_init__(self): check_option(self.precond, LiteralOptions.OptsShearAlfvenPrecond) if self.precond in ("MultiGrid", "StiffnessPreconditioner"): 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 == "StiffnessPreconditioner": assert self.u_space == "Hdiv", 'precond="StiffnessPreconditioner" requires u_space="Hdiv".' if self.precond == "MultiGrid" and self.multigrid is None: From 2f091cdcf4a887b96849c84c907e1bfcd37d14bd Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 09:17:35 +0200 Subject: [PATCH 05/12] Bump feectools version number --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index 1d8d47c4a..e75b4c766 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 1d8d47c4a3bc8e0b594992cc22944b1ac4230c0f +Subproject commit e75b4c76616ef93f0091d17e752e4155f430c5d2 From 20b77a56c939d11274fd490ec8cea0beafe03a82 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 13:59:36 +0200 Subject: [PATCH 06/12] Allow fused solver inner products in the GPU transfer budget --- src/struphy/models/tests/gpu_e2e.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/src/struphy/models/tests/gpu_e2e.py b/src/struphy/models/tests/gpu_e2e.py index ecfa0c9a2..46dbe2765 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,10 +104,10 @@ 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 From 9e6eeec20284c765afe6852284d359f00ac8115e Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 14:56:25 +0200 Subject: [PATCH 07/12] polar MPI decomposition parity requires more than one rank --- src/struphy/polar/tests/test_polar_cupy.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/struphy/polar/tests/test_polar_cupy.py b/src/struphy/polar/tests/test_polar_cupy.py index 2e7dc2c5f..6f23408a6 100644 --- a/src/struphy/polar/tests/test_polar_cupy.py +++ b/src/struphy/polar/tests/test_polar_cupy.py @@ -342,9 +342,13 @@ def check_polar_mpi(cuda_aware, mpi_dims_mask, fake=True): """Backend parity with both Kronecker redistributions and ranks away from the pole.""" from contextlib import nullcontext + from maybempi import MPI from struphy.pic.tests.cuda_emulation import emulated_launches from struphy.polar.linear_operators import set_device_sparse_module + if MPI.COMM_WORLD.Get_size() == 1: + pytest.skip("polar MPI decomposition parity requires more than one rank") + previous_mpi = cunumpy.mpi.get_mpi_cuda_aware() previous_sparse = set_device_sparse_module(DenseDeviceSparse) if fake else None try: From 3f090ceb8b21bee6c1a1300e8052247089a31115 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 15:30:29 +0200 Subject: [PATCH 08/12] Fix fake CuPy MPI and solver paths Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- feectools | 2 +- src/struphy/polar/tests/test_polar_cupy.py | 4 ---- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/feectools b/feectools index e75b4c766..717edc98b 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit e75b4c76616ef93f0091d17e752e4155f430c5d2 +Subproject commit 717edc98b09464358b9865dfbe91cea0d00ca5a5 diff --git a/src/struphy/polar/tests/test_polar_cupy.py b/src/struphy/polar/tests/test_polar_cupy.py index 6f23408a6..2e7dc2c5f 100644 --- a/src/struphy/polar/tests/test_polar_cupy.py +++ b/src/struphy/polar/tests/test_polar_cupy.py @@ -342,13 +342,9 @@ def check_polar_mpi(cuda_aware, mpi_dims_mask, fake=True): """Backend parity with both Kronecker redistributions and ranks away from the pole.""" from contextlib import nullcontext - from maybempi import MPI from struphy.pic.tests.cuda_emulation import emulated_launches from struphy.polar.linear_operators import set_device_sparse_module - if MPI.COMM_WORLD.Get_size() == 1: - pytest.skip("polar MPI decomposition parity requires more than one rank") - previous_mpi = cunumpy.mpi.get_mpi_cuda_aware() previous_sparse = set_device_sparse_module(DenseDeviceSparse) if fake else None try: From 29e918008adb0095bbe6dda82eb25490825010f2 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 15:40:11 +0200 Subject: [PATCH 09/12] Account for MPI scalar diagnostic transfers Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- src/struphy/models/tests/gpu_e2e.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/struphy/models/tests/gpu_e2e.py b/src/struphy/models/tests/gpu_e2e.py index 46dbe2765..d58e20af4 100644 --- a/src/struphy/models/tests/gpu_e2e.py +++ b/src/struphy/models/tests/gpu_e2e.py @@ -114,8 +114,11 @@ def check(self, n_scalars: int): # 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"): From 080cead991ddb4d6d1cd225f564feedf1e1439b2 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 23 Oct 2026 16:20:02 +0200 Subject: [PATCH 10/12] Fix fake CuPy block operations and serial inner products --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index 717edc98b..b3de2933b 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 717edc98b09464358b9865dfbe91cea0d00ca5a5 +Subproject commit b3de2933b5071a6762aa0432838c39f201cdc73a From 352a4bb73513a98f3e2d0905b06837299a704956 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 24 Oct 2026 13:39:58 +0200 Subject: [PATCH 11/12] Set feectools to devel-tiny --- feectools | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/feectools b/feectools index b3de2933b..c3c7d941e 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit b3de2933b5071a6762aa0432838c39f201cdc73a +Subproject commit c3c7d941e97f5adb40dbe6e6704d261fb48c3a11 From b3d5f2ecfe73b497a1812599039a207a157de1b3 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 24 Oct 2026 13:40:24 +0200 Subject: [PATCH 12/12] Update feectools dependency to 0.10.0 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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",