Compile the CUDA kernels in the setup stage (Simulation.compile_cuda_kernels)

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).

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:
     # --------------------------

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."""

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."""

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
     # -------------------------------------------

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."""

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."""

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

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
+        <struphy.geometry.base.Domain.kernels>`), the Derham complex (:meth:`Derham.kernels
+        <struphy.feec.psydac_derham.Derham.kernels>`), the particles of every kinetic variable
+        (:meth:`Particles.kernels <struphy.pic.base.Particles.kernels>`) and the propagators of the model
+        (:meth:`Propagator.kernels <struphy.propagators.base.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,