From 225bcae42976a7690b7c9ebc97feaaa6b34ab769 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 20:49:16 +0200 Subject: [PATCH 1/2] Compile the CUDA kernels in the setup stage (Simulation.compile_cuda_kernels) Simulation.run() compiles the CUDA kernels of the time stepping after allocate(), in the profiling region 'setup: compile cuda kernels', instead of at their first call. The kernels are listed by the objects that keep them (kernels() on Domain, Derham, Particles, Pusher, KernelSetup, Accumulator, AccumulatorVector and Propagator). On CuPy, kernels without a CUDA version raise NotImplementedError at setup, all listed; on NumPy nothing happens. Solves #704. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 20 +- src/struphy/feec/psydac_derham.py | 22 ++ src/struphy/geometry/base.py | 7 + .../pic/accumulation/particles_to_grid.py | 8 + src/struphy/pic/base.py | 9 + src/struphy/pic/pushing/kernel_setup.py | 4 + src/struphy/pic/pushing/pusher.py | 4 + src/struphy/propagators/base.py | 15 ++ src/struphy/simulation/sim.py | 68 +++++ .../tests/test_compile_cuda_kernels.py | 250 ++++++++++++++++++ src/struphy/utils/kernel_compilation.py | 52 ++++ 11 files changed, 458 insertions(+), 1 deletion(-) create mode 100644 src/struphy/simulation/tests/test_compile_cuda_kernels.py create mode 100644 src/struphy/utils/kernel_compilation.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index daf0d25f0..6d2c60e72 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -500,7 +500,8 @@ through `Kernel(PyccelKernel(...))`, built at the call site where the owner is p cannot be pickled) and kept by `Pusher`, `KernelSetup`, the accumulators, reflection and the spline evaluation; a kernel without a CUDA version raises `NotImplementedError` when it is called on CuPy. `prepare_kernel()` (PR 15) is removed: kernels are no longer checked or compiled at setup, CUDA kernels -compile on their first call. When to compile them (e.g. once at simulation start) is decided later. +compile on their first call. When to compile them (e.g. once at simulation start) is decided later. (Since +struphy#704 they are compiled at setup, see [Compiling at setup](#compiling-the-cuda-kernels-at-setup-struphy704).) Consequence: geometry evaluations (`Domain.__call__`, `jacobian_det`, `pull`/`push`, ...) have no CUDA version yet and raise on CuPy. `Particles` evaluates `jacobian_det` when it initializes weights, so @@ -537,6 +538,23 @@ They are replaced by one test module, `pic/tests/cuda_parity_cases.py`: `PARITY_ `assert_kernels_agree` (one test per case), `test_cuda_emulation.py` to the CPU emulation, and `test_cuda_kernels_have_parity_cases` checks that the CUDA kernels and the parity cases match. +## Compiling the CUDA kernels at setup (struphy#704) + +`Simulation.run()` calls `Simulation.compile_cuda_kernels()` after `allocate()`, in the profiling region +`setup: compile cuda kernels`, so that no CUDA kernel compiles in the first time step. The kernels come from +`Simulation.kernels()`, which collects the `kernels()` of the owners that keep and call them: `Domain` (the four +geometry kernels), `Derham` (the three spline evaluation kernels and feectools' `stencil_{dot,transpose,inner,axpy}_3d`), +`Particles` (`reflect` if an axis reflects), and every propagator of the model. `Propagator.kernels()` collects +by default from the propagator's attributes: `Kernel`s kept as attributes and the `kernels()` of `Pusher` (pusher +kernel and its `KernelSetup`s), `KernelSetup`, `Accumulator` and `AccumulatorVector`, also inside lists, tuples +and dicts (`utils/kernel_compilation.collect_kernels`). A propagator that calls a kernel it does not keep +overrides `kernels()`. Diagnostics kernels and the initial solves are not listed; they run during the setup. + +Fail fast: on CuPy, if a listed kernel has no CUDA version, `compile_cuda_kernels()` raises +`NotImplementedError` naming all of them, before anything is compiled (e.g. `VlasovAmpereOneSpecies`: +`vlasov_maxwell`). On NumPy it does nothing. A test runs a model for one step on NumPy and checks that every +`Kernel` called in the time loop is listed. + ## PR 18 implementation notes Geometry evaluation runs on the GPU for every analytic mapping (`kind_map` 10–12, 20–22, 30–32). diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index dfdb34e4c..2eec857e2 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -22,6 +22,10 @@ from feectools.fem.vector import VectorFemSpace from feectools.linalg.basic import IdentityOperator from feectools.linalg.block import BlockVector, BlockVectorSpace +from feectools.linalg.kernels.stencil_axpy_3d import stencil_axpy_3d +from feectools.linalg.kernels.stencil_dot_3d import stencil_dot_3d +from feectools.linalg.kernels.stencil_inner_3d import stencil_inner_3d +from feectools.linalg.kernels.stencil_transpose_3d import stencil_transpose_3d from feectools.linalg.stencil import StencilVector, StencilVectorSpace from maybempi import MPI, SerialComm @@ -1503,6 +1507,24 @@ def args_derham(self) -> DerhamArguments | CudaDerhamArguments: """Mandatory pusher kernel arguments for the backend used at initialization.""" return self._args_derham + def kernels(self) -> tuple: + """The kernels called on the data of this complex. + + The spline evaluation kernels of :class:`SplineFunction` (at markers, on meshgrids and on + sparse meshgrids) and the feectools kernels of the 3D stencil data (matrix-vector product, + transpose, inner product, axpy), which the FEEC operators and solvers call. Used by + :meth:`~struphy.simulation.sim.Simulation.compile_cuda_kernels`. + """ + return ( + eval_spline_mpi_markers, + eval_spline_mpi_matrix, + eval_spline_mpi_sparse_meshgrid, + stencil_dot_3d, + stencil_transpose_3d, + stencil_inner_3d, + stencil_axpy_3d, + ) + # -------------------------- # methods: # -------------------------- diff --git a/src/struphy/geometry/base.py b/src/struphy/geometry/base.py index 81163bdde..44c48ccd1 100644 --- a/src/struphy/geometry/base.py +++ b/src/struphy/geometry/base.py @@ -469,6 +469,13 @@ def args_domain(self): raise AttributeError("DomainArguments are not available because the domain state is incomplete.") return self._args_domain + def kernels(self) -> tuple: + """The geometry evaluation kernels (evaluation and pull/push, on grids and at markers). + + Used by :meth:`~struphy.simulation.sim.Simulation.compile_cuda_kernels`. + """ + return (kernel_evaluate, kernel_evaluate_pic, kernel_pullpush, kernel_pullpush_pic) + @property def dict_transformations(self): """Dictionary of str->int for pull, push and transformation functions.""" diff --git a/src/struphy/pic/accumulation/particles_to_grid.py b/src/struphy/pic/accumulation/particles_to_grid.py index 4c6688661..560ce65db 100644 --- a/src/struphy/pic/accumulation/particles_to_grid.py +++ b/src/struphy/pic/accumulation/particles_to_grid.py @@ -323,6 +323,10 @@ def kernel(self) -> Kernel: """The accumulation kernel.""" return self._kernel + def kernels(self) -> tuple[Kernel, ...]: + """The kernels this accumulator calls (the accumulation kernel).""" + return (self._kernel,) + @property def derham(self): """Discrete Derham complex on the logical unit cube.""" @@ -608,6 +612,10 @@ def kernel(self) -> Kernel: """The accumulation kernel.""" return self._kernel + def kernels(self) -> tuple[Kernel, ...]: + """The kernels this accumulator calls (the accumulation kernel).""" + return (self._kernel,) + @property def derham(self): """Discrete Derham complex on the logical unit cube.""" diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 1d08e9116..99db7dbfd 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -941,6 +941,15 @@ def args_markers(self) -> MarkerArguments | CudaMarkerArguments: """Arguments for marker kernels, selected when this particle object is created.""" return self._args_markers + def kernels(self) -> tuple: + """The kernels these particles call when they apply their boundary conditions. + + :func:`~struphy.pic.pushing.kernels.reflect.reflect` if an axis has reflecting boundary + conditions, otherwise none. Used by + :meth:`~struphy.simulation.sim.Simulation.compile_cuda_kernels`. + """ + return (reflect,) if self._reflect_axes else () + # ------------------------------------------- # Initial condition and background -> weights # ------------------------------------------- diff --git a/src/struphy/pic/pushing/kernel_setup.py b/src/struphy/pic/pushing/kernel_setup.py index 5577e6976..e0ff50f81 100644 --- a/src/struphy/pic/pushing/kernel_setup.py +++ b/src/struphy/pic/pushing/kernel_setup.py @@ -61,6 +61,10 @@ def __post_init__(self): if not isinstance(self.kernel, Kernel): self.kernel = Kernel(self.kernel) + def kernels(self) -> tuple[Kernel, ...]: + """The kernel of this setup, as a one-element tuple.""" + return (self.kernel,) + @property def name(self) -> str: """Kernel name used in profiling and validation messages.""" diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index 762dd3bd9..c9271883f 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -386,6 +386,10 @@ def kernel(self) -> Kernel: """The pusher kernel for the active backend (pyccel or CUDA).""" return self._kernel + def kernels(self) -> tuple[Kernel, ...]: + """The kernels this pusher calls: the pusher kernel and those of its init and eval kernel setups.""" + return (self._kernel,) + tuple(setup.kernel for setup in self._init_kernels + self._eval_kernels) + @property def init_kernels(self) -> tuple[KernelSetup, ...]: """Ordered setups for evaluations at the initial state.""" diff --git a/src/struphy/propagators/base.py b/src/struphy/propagators/base.py index 56f273d39..811443c72 100644 --- a/src/struphy/propagators/base.py +++ b/src/struphy/propagators/base.py @@ -6,6 +6,7 @@ from typing import Literal import cunumpy as xp +from cunumpy.kernels import Kernel from feectools.linalg.block import BlockVector from feectools.linalg.stencil import StencilVector from scope_profiler import ProfileManager @@ -18,6 +19,7 @@ from struphy.io.options import OptionsBase from struphy.models.variables import FEECVariable, PICVariable, SPHVariable, Variable from struphy.pic.pushing.kernel_setup import KernelSetup +from struphy.utils.kernel_compilation import collect_kernels from struphy.utils.utils import check_option logger = logging.getLogger("struphy") @@ -141,6 +143,19 @@ def update_feec_variables(self, **new_coeffs): return diffs + def kernels(self) -> tuple[Kernel, ...]: + """The :class:`~cunumpy.kernels.Kernel` objects this propagator calls in a time step. + + Used by :meth:`~struphy.simulation.sim.Simulation.compile_cuda_kernels` to compile the CUDA + kernels before the time stepping. The default collects them from the attributes of the + propagator (after :meth:`allocate`): kernels kept as attributes, and the kernels of attributes + with a ``kernels()`` method (``Pusher``, ``KernelSetup``, ``Accumulator``, + ``AccumulatorVector``), also inside lists, tuples and dicts. A propagator that calls a kernel + it does not keep (e.g. a module-level kernel called directly in :meth:`__call__`) overrides this + method and adds it. + """ + return collect_kernels(*vars(self).values()) + @property def init_kernels(self) -> tuple[KernelSetup, ...]: r"""Tuple of initialization kernel setups for evaluation at diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index a3fe85dc2..9c6042fef 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -17,6 +17,7 @@ from typing import TYPE_CHECKING import cunumpy as xp +from cunumpy.kernels import Kernel import h5py import yaml from feectools.linalg.memory import stencil_matrix_memory @@ -77,6 +78,7 @@ from struphy.propagators.base import Propagator from struphy.simulation.base import SimulationBase from struphy.utils.clone_config import CloneConfig +from struphy.utils.kernel_compilation import collect_kernels, missing_cuda from struphy.utils.progress import tqdm from struphy.utils.utils import ruff_autofix_and_format @@ -259,6 +261,69 @@ def allocate(self): logger.debug("... Done.") + def kernels(self) -> tuple[Kernel, ...]: + """The :class:`~cunumpy.kernels.Kernel` objects the time stepping calls, each once. + + Collected after :meth:`allocate` from the domain (:meth:`Domain.kernels + `), the Derham complex (:meth:`Derham.kernels + `), the particles of every kinetic variable + (:meth:`Particles.kernels `) and the propagators of the model + (:meth:`Propagator.kernels `). Kernels that only + diagnostics or the initial solves call are not included; they run during the setup anyway. + """ + owners = [self.domain, self.derham] + for species in self.model.species.values(): + for variable in species.variables.values(): + if isinstance(variable, (PICVariable, SPHVariable)): + owners.append(variable.particles) + owners += self.model.prop_list + return collect_kernels(*owners) + + def compile_cuda_kernels(self) -> tuple[Kernel, ...]: + """Compile the CUDA versions of the kernels of :meth:`kernels` now, before the time stepping. + + A CUDA kernel is otherwise compiled on its first call, which puts the compile time into the + first time step (and its profiling regions). :meth:`run` calls this method after + :meth:`allocate`, in the profiling region ``"setup: compile cuda kernels"``. CuPy caches the + compiled kernels on disk, so a second run on the same machine mostly loads them. + + On the NumPy backend this does nothing (the pyccel kernels are compiled by ``struphy compile``). + + Returns + ------- + tuple[Kernel, ...] + The kernels whose CUDA version was compiled; empty on the NumPy backend. + + Raises + ------ + NotImplementedError + On the CuPy backend, if a kernel the time stepping calls has no CUDA version. All such + kernels are listed, and nothing is compiled; without this check the run would fail at the + first call of the first one, in the time loop. + """ + if xp.get_backend() != "cupy": + return () + + kernels = self.kernels() + missing = missing_cuda(kernels) + if missing: + lines = "\n".join( + f" - {kernel.name}" + (f" (expected at {kernel.cuda_path})" if kernel.cuda_path else "") + for kernel in missing + ) + raise NotImplementedError( + f"Model {self.model_name} cannot run on the CuPy backend: {len(missing)} of the " + f"{len(kernels)} kernels its time stepping calls have no CUDA version:\n{lines}", + ) + + compiled = [] + with ProfileManager.profile_region("setup: compile cuda kernels"): + for kernel in kernels: + if kernel.compile(): + compiled.append(kernel) + logger.info(f"Compiled {len(compiled)} CUDA kernels: {', '.join(k.name for k in compiled)}") + return tuple(compiled) + def estimate_mem(self, print_report: bool = False) -> dict: """Estimate the memory footprint of all model variables and FEEC matrices, in bytes, BEFORE calling :meth:`allocate`. @@ -657,6 +722,9 @@ def run(self, one_time_step: bool = False, profiling_activated: bool | None = No with ProfileManager.profile_region("setup: total"): # equation paramters self.allocate() + + # CUDA kernels (CuPy backend only): compile them now rather than in the first time step + self.compile_cuda_kernels() with ProfileManager.profile_region("setup: run metadata", functions=[self._write_run_metadata]): self._write_run_metadata( one_time_step=one_time_step, diff --git a/src/struphy/simulation/tests/test_compile_cuda_kernels.py b/src/struphy/simulation/tests/test_compile_cuda_kernels.py new file mode 100644 index 000000000..faf5db96e --- /dev/null +++ b/src/struphy/simulation/tests/test_compile_cuda_kernels.py @@ -0,0 +1,250 @@ +"""Tests for ``Simulation.kernels`` and ``Simulation.compile_cuda_kernels`` (struphy#704). + +The CUDA kernels are compiled once at setup instead of at their first call in the time loop. Without a +GPU, the CuPy path is tested with the backend switched by monkeypatching and with cunumpy's fake CuPy +(in a subprocess); the test marked ``requires_cupy`` compiles the kernels of a small model on a GPU. +""" + +import subprocess +import sys + +import cunumpy +import pytest +from cunumpy.kernels import CudaKernel, Kernel + +from struphy import DerhamOptions, EnvironmentOptions, Simulation, Time, domains, equils, grids, maxwellians +from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env +from struphy.models import Vlasov, VlasovAmpereOneSpecies +from struphy.particles.parameters import BoundaryParameters, LoadingParameters +from struphy.pic.pushing.kernels.push_bxu_Hdiv import push_bxu_Hdiv +from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage +from struphy.pic.pushing.kernels.reflect import reflect +from struphy.utils.kernel_compilation import collect_kernels, missing_cuda + +# simulations write to tmp_path, which differs per rank +pytestmark = pytest.mark.mpi_skip + +requires_cupy = pytest.mark.skipif(not cunumpy.cupy_available(), reason="CuPy/GPU not available") + +# kernels of the domain and the Derham complex, listed for every model with a Derham complex +DOMAIN_KERNELS = {"kernel_evaluate", "kernel_evaluate_pic", "kernel_pullpush", "kernel_pullpush_pic"} +DERHAM_KERNELS = { + "eval_spline_mpi_markers", + "eval_spline_mpi_matrix", + "eval_spline_mpi_sparse_meshgrid", + "stencil_dot_3d", + "stencil_transpose_3d", + "stencil_inner_3d", + "stencil_axpy_3d", +} + + +def make_vlasov_ampere(tmp_path, bc=("periodic", "periodic", "periodic")) -> Simulation: + """VlasovAmpereOneSpecies on 8 x 1 x 1 elements with 5 markers per cell, one time step.""" + model = VlasovAmpereOneSpecies(with_B0=False) + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(ppc=5, seed=1234), + boundary_params=BoundaryParameters(bc=list(bc)), + ) + model.propagators.push_eta.options = model.propagators.push_eta.Options() + model.propagators.coupling_va.options = model.propagators.coupling_va.Options() + model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0") + model.kinetic_ions.var.add_background(maxwellians.Maxwellian3D(n=(1.0, None))) + return Simulation( + model=model, + env=EnvironmentOptions(out_folders=str(tmp_path), sim_folder="sim_1"), + time_opts=Time(dt=0.05, Tend=0.05), + domain=domains.Cuboid(r1=12.56), + grid=grids.TensorProductGrid(num_elements=(8, 1, 1)), + derham_opts=DerhamOptions(degree=(2, 1, 1)), + ) + + +def make_vlasov(tmp_path) -> Simulation: + """Vlasov (PushVxB and PushEta) on 4 x 4 x 1 elements, whose kernels all have CUDA versions.""" + model = Vlasov() + model.kinetic_ions.set_markers(loading_params=LoadingParameters(ppc=5, seed=1234)) + model.propagators.push_vxb.options = model.propagators.push_vxb.Options() + model.propagators.push_eta.options = model.propagators.push_eta.Options() + model.kinetic_ions.var.add_background(maxwellians.Maxwellian3D(n=(1.0, None))) + return Simulation( + model=model, + env=EnvironmentOptions(out_folders=str(tmp_path), sim_folder="sim_1"), + time_opts=Time(dt=0.05, Tend=0.05), + domain=domains.Cuboid(), + equil=equils.HomogenSlab(), + grid=grids.TensorProductGrid(num_elements=(4, 4, 1)), + derham_opts=DerhamOptions(degree=(1, 1, 1)), + ) + + +def names(kernels) -> set[str]: + return {kernel.name for kernel in kernels} + + +def test_collect_kernels_walks_owners_and_containers(): + """Kernels, ``kernels()`` methods and containers are walked; each kernel is listed once, in order.""" + + class Owner: + def __init__(self, *kernels): + self._kernels = kernels + + def kernels(self): + return self._kernels + + collected = collect_kernels( + None, + "not an owner", + push_eta_stage, + Owner(push_eta_stage, reflect), + {"a": [Owner(push_bxu_Hdiv)], "b": (reflect,)}, + Owner, # a class with a kernels() method is not an owner + ) + assert collected == (push_eta_stage, reflect, push_bxu_Hdiv) + assert missing_cuda(collected) == [push_bxu_Hdiv] + + +def test_kernels_of_a_small_model(tmp_path): + """The kernels of VlasovAmpereOneSpecies: domain, Derham and its two propagators.""" + sim = make_vlasov_ampere(tmp_path) + sim.allocate() + + kernels = sim.kernels() + assert all(isinstance(kernel, Kernel) for kernel in kernels) + assert len(set(map(id, kernels))) == len(kernels) + assert names(kernels) == DOMAIN_KERNELS | DERHAM_KERNELS | { + "push_eta_stage", + "vlasov_maxwell", + "push_v_with_efield", + } + assert names(sim.model.propagators.push_eta.kernels()) == {"push_eta_stage"} + assert names(sim.model.propagators.coupling_va.kernels()) == {"vlasov_maxwell", "push_v_with_efield"} + # the accumulation kernel vlasov_maxwell has no CUDA version yet: this model cannot run on CuPy + assert names(missing_cuda(kernels)) == {"vlasov_maxwell"} + + +def test_kernels_of_reflecting_particles(tmp_path): + """Particles with a reflecting boundary list the reflection kernel.""" + sim = make_vlasov_ampere(tmp_path, bc=("reflect", "periodic", "periodic")) + sim.allocate() + assert sim.model.kinetic_ions.var.particles.kernels() == (reflect,) + assert reflect in sim.kernels() + + +def test_kernels_cover_the_time_loop(tmp_path, monkeypatch): + """Every kernel called in the time loop of a run is listed by ``Simulation.kernels()``.""" + sim = make_vlasov_ampere(tmp_path) + called = {} + in_loop = [False] + call = Kernel.__call__ + + def recording_call(self, *args, **kwargs): + if in_loop[0]: + called[id(self)] = self + return call(self, *args, **kwargs) + + integrate = sim.model.integrate + + def integrate_in_loop(*args, **kwargs): + in_loop[0] = True + try: + return integrate(*args, **kwargs) + finally: + in_loop[0] = False + + monkeypatch.setattr(Kernel, "__call__", recording_call) + monkeypatch.setattr(sim.model, "integrate", integrate_in_loop) + sim.run(one_time_step=True) + + assert called, "no kernel was called in the time loop" + listed = {id(kernel) for kernel in sim.kernels()} + assert names(k for i, k in called.items() if i not in listed) == set() + + +def test_compile_cuda_kernels_is_a_no_op_on_numpy(tmp_path, monkeypatch): + """On NumPy nothing is compiled or checked, also when the model has kernels without CUDA version.""" + + def fail(self): + raise AssertionError(f"{self.name} compiled on the NumPy backend") + + monkeypatch.setattr(Kernel, "compile", fail) + sim = make_vlasov_ampere(tmp_path) + with cunumpy.use_backend("numpy"): + sim.allocate() + assert sim.compile_cuda_kernels() == () + + +def test_compile_cuda_kernels_on_cupy_backend(tmp_path, monkeypatch): + """On CuPy (backend switched by monkeypatching) every kernel is compiled; missing CUDA versions raise first.""" + sim = make_vlasov_ampere(tmp_path) + compiled = [] + monkeypatch.setattr(CudaKernel, "compile", lambda self, **kwargs: compiled.append(self.name)) + monkeypatch.setattr(cunumpy, "get_backend", lambda: "cupy") + + monkeypatch.setattr(Simulation, "kernels", lambda self: (push_eta_stage, reflect)) + assert sim.compile_cuda_kernels() == (push_eta_stage, reflect) + assert compiled == ["push_eta_stage", "reflect"] + + compiled.clear() + monkeypatch.setattr(Simulation, "kernels", lambda self: (push_eta_stage, push_bxu_Hdiv, reflect)) + with pytest.raises(NotImplementedError, match=r"1 of the\s+3 kernels.*\n - push_bxu_Hdiv \(expected at .*"): + sim.compile_cuda_kernels() + assert compiled == [] + + +FAKE_CUPY_CHECK = """ +import cunumpy +from cunumpy.kernels import CudaKernel + +from struphy import Simulation +from struphy.models import Maxwell +from struphy.pic.pushing.kernels.push_bxu_Hdiv import push_bxu_Hdiv +from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage +from struphy.pic.pushing.kernels.reflect import reflect + +compiled = [] +CudaKernel.compile = lambda self, **kwargs: compiled.append(self.name) +sim = Simulation(model=Maxwell()) +with cunumpy.use_backend("numpy"): + Simulation.kernels = lambda self: (push_eta_stage, push_bxu_Hdiv) + assert sim.compile_cuda_kernels() == () +with cunumpy.use_backend("cupy"): + try: + sim.compile_cuda_kernels() + except NotImplementedError as error: + assert "push_bxu_Hdiv" in str(error) and "push_eta_stage" not in str(error), error + else: + raise AssertionError("a kernel without CUDA version did not raise") + assert compiled == [] + Simulation.kernels = lambda self: (push_eta_stage, reflect) + assert sim.compile_cuda_kernels() == (push_eta_stage, reflect) + assert compiled == ["push_eta_stage", "reflect"] +""" + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +def test_compile_cuda_kernels_fake_cupy(): + """With cunumpy's fake CuPy the CuPy backend is active: kernels are checked and compiled (compile recorded).""" + result = subprocess.run( + [sys.executable, "-c", FAKE_CUPY_CHECK], + env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stderr[-4000:] + + +@requires_cupy +def test_compile_cuda_kernels_on_gpu(tmp_path): + """On a GPU: the kernels of Vlasov are compiled at setup, VlasovAmpereOneSpecies fails fast.""" + with cunumpy.use_backend("cupy"): + sim = make_vlasov(tmp_path) + sim.allocate() + compiled = sim.compile_cuda_kernels() + assert compiled == sim.kernels() + assert {"push_vxb_analytic", "push_eta_stage"} <= names(compiled) + + sim = make_vlasov_ampere(tmp_path / "va") + sim.allocate() + with pytest.raises(NotImplementedError, match="vlasov_maxwell"): + sim.compile_cuda_kernels() diff --git a/src/struphy/utils/kernel_compilation.py b/src/struphy/utils/kernel_compilation.py new file mode 100644 index 000000000..93e428912 --- /dev/null +++ b/src/struphy/utils/kernel_compilation.py @@ -0,0 +1,52 @@ +"""Collect the kernels a simulation calls and compile their CUDA versions before the time stepping. + +A CUDA kernel is otherwise compiled on its first call, which then shows up in the profile of the first +time step. The objects that keep and call :class:`cunumpy.kernels.Kernel` objects (``Pusher``, +``KernelSetup``, ``Accumulator``, ``AccumulatorVector``, ``Particles``, ``Domain``, ``Derham``, the +propagators) list them with a ``kernels()`` method; :func:`collect_kernels` gathers them and +:meth:`~struphy.simulation.sim.Simulation.compile_cuda_kernels` compiles them once at setup. +""" + +from collections.abc import Iterable, Iterator, Mapping + +from cunumpy.kernels import Kernel + + +def _walk(value, seen: dict[int, object]) -> Iterator[Kernel]: + # `seen` keeps the visited objects alive, so that the id of a temporary (e.g. a tuple returned + # by ``kernels()``) is not reused by another object during the walk + if id(value) in seen: + return + seen[id(value)] = value + if isinstance(value, Kernel): + yield value + elif isinstance(value, Mapping): + for item in value.values(): + yield from _walk(item, seen) + elif isinstance(value, (list, tuple, set, frozenset)): + for item in value: + yield from _walk(item, seen) + elif not isinstance(value, type) and callable(getattr(value, "kernels", None)): + for item in value.kernels(): + yield from _walk(item, seen) + + +def collect_kernels(*owners) -> tuple[Kernel, ...]: + """The kernels of `owners`, each once, in the order they are found. + + An owner is a :class:`~cunumpy.kernels.Kernel`, an object with a ``kernels()`` method (whose + result is collected the same way), or a list, tuple, set or dict of owners. Anything else is + ignored, so ``collect_kernels(*vars(obj).values())`` collects the kernels of the attributes of + ``obj`` without descending into other objects. + """ + seen: dict[int, object] = {} + kernels: dict[int, Kernel] = {} + for owner in owners: + for kernel in _walk(owner, seen): + kernels.setdefault(id(kernel), kernel) + return tuple(kernels.values()) + + +def missing_cuda(kernels: Iterable[Kernel]) -> list[Kernel]: + """The kernels that would raise on the CuPy backend because they have no CUDA version.""" + return [k for k in kernels if not k.has_cuda and k.missing_cuda == "raise"] From d412fc298699aeeefb8eccb3110a5d2358baeb05 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 8 Oct 2026 09:22:24 +0200 Subject: [PATCH 2/2] formatting --- src/struphy/simulation/sim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 9c6042fef..52347634f 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -17,9 +17,9 @@ from typing import TYPE_CHECKING import cunumpy as xp -from cunumpy.kernels import Kernel import h5py import yaml +from cunumpy.kernels import Kernel from feectools.linalg.memory import stencil_matrix_memory from feectools.linalg.stencil import StencilVector from line_profiler import profile