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..52347634f 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -19,6 +19,7 @@ import cunumpy as xp 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 @@ -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"]