Cuda push v in efield only

.gitignore

@@ -101,6 +101,8 @@ src/struphy/io/out/
 src/struphy/state.yml
 src/struphy/io/inp/params_*
 *.bin
+bench_gpu/out_*
+bench_gpu/
 
 # models list
 bin/

pyproject.toml

@@ -63,6 +63,9 @@ phys = [
     "gvec>=1.1.0, <=1.5.0",
     "desc-opt<=0.17.1",
 ]
+cuda = [
+    "cupy-cuda12x==13.6.0",
+]
 dev = [
     "struphy[mpi]",
     "notebook",
@@ -137,6 +140,7 @@ struphy = "struphy.console.main:struphy"
 ]
 struphy = [
     "compile_struphy.mk",
+    "**/*.cu",
 ]
 
 [tool.autopep8]
@@ -182,4 +186,5 @@ markers = [
     "hybrid",
     "single",
     "mpi_pic",
+    "needs_host_kernels",
 ]

setup/modules.pitagora.sh

@@ -5,12 +5,22 @@ python/3.11.7"
 
 # openmpi/4.1.6--gcc--12.3.0
 MODULES_GCC="gcc/12.3.0 \
+openmpi/4.1.6--gcc--12.3.0 \
 python/3.11.7 \
 hdf5/1.14.3--gcc--12.3.0 \
 cmake/3.27.9 \
 netcdf-fortran/4.6.1--gcc--12.3.0 \
 netlib-scalapack/2.2.0--openmpi--4.1.6--gcc--12.3.0-ucx1.20"
 
+# On the Booster (GPU) partition, ARRAY_BACKEND=cupy runs need libnvrtc.so.12 for
+# cupy's RawKernel/JIT compilation -- otherwise every cupy import fails as soon as
+# it touches the GPU (e.g. `xp.tri()` at struphy import time). SLURM_JOB_PARTITION
+# is only set inside a submitted job, so this is a no-op on the DCGP (CPU) partition
+# or outside SLURM.
+if [[ "${SLURM_JOB_PARTITION:-}" == *boost* ]]; then
+    MODULES_INTEL="$MODULES_INTEL cuda/12.6"
+    MODULES_GCC="$MODULES_GCC cuda/12.6"
+fi
 
 # For GVEC
 # Should be fixed so it works with both gcc and intel

src/struphy/__init__.py

@@ -4,6 +4,33 @@
 import logging.config
 import os
 
+# HDF5's file locking relies on flock(), which is unreliable/unsupported on
+# parallel filesystems such as Lustre or GPFS (common on HPC clusters) and
+# causes spurious `BlockingIOError: Unable to synchronously open file` errors.
+# Disable it unless the user has explicitly configured it.
+os.environ.setdefault("HDF5_USE_FILE_LOCKING", "FALSE")
+
+# mpi4py defaults to requesting MPI_THREAD_MULTIPLE (thread level 3) from
+# MPI_Init_thread. On at least one cluster this repo runs on (Pitagora's Booster
+# partition, OpenMPI 4.1.6 + UCX 1.20), the UCX worker does not support that level,
+# which OpenMPI reports at every multi-rank run ("UCP worker does not support
+# MPI_THREAD_MULTIPLE" / "failed to init ucx" / hcoll init failure) and works
+# around by making hcoll (its GPU-aware collective component) fail to initialize,
+# silently falling back to a different, working collective implementation. Struphy
+# only ever calls MPI from the main Python thread (CuPy's internal CUDA driver
+# threads don't touch MPI), so requesting the weaker MPI_THREAD_FUNNELED guarantee
+# instead is sufficient and avoids the warnings -- but it ALSO lets hcoll
+# successfully initialize where it previously failed to, and hcoll's own
+# Alltoallv implementation on this cluster then segfaults
+# (hmca_bcol_ucx_p2p_alltoallv_pairwise_chunk_progress) the first time it's
+# actually used, something the failed init was silently protecting us from.
+# hcoll must therefore stay disabled explicitly alongside the thread-level
+# change, not just left to fail its own init. Both must be set before mpi4py.MPI
+# is imported anywhere (thread level can't change after MPI_Init), and only if
+# the user hasn't already configured them themselves.
+os.environ.setdefault("MPI4PY_RC_THREAD_LEVEL", "funneled")
+os.environ.setdefault("OMPI_MCA_coll_hcoll_enable", "0")
+
 from feectools.ddm.mpi import mpi as MPI
 
 from struphy.utils.mpi_launch import launched_under_mpi

src/struphy/conftest.py

@@ -1,9 +1,17 @@
 import logging
+import os
 
 import pytest
 
 from struphy import set_logging_level
 
+# Tests marked "needs_host_kernels" call a Pyccel kernel that is not yet
+# ported for the CuPy backend (e.g. FEEC mass-matrix/basis-projection
+# assembly, particle-to-grid accumulation). Under ARRAY_BACKEND=cupy they are
+# skipped rather than run to failure, so a GPU CI run reports the state of
+# the actually-ported code paths instead of drowning in known gaps.
+_NEEDS_HOST_KERNELS_SKIP_REASON = "needs_host_kernels: not yet ported to the CuPy backend (ARRAY_BACKEND=cupy)"
+
 
 def set_logging_level_pytest(config):
     level_name = str(config.getoption("--logging-level")).upper()
@@ -25,6 +33,15 @@ def pytest_configure(config):
     set_logging_level_pytest(config)
 
 
+def pytest_collection_modifyitems(config, items):
+    if os.environ.get("ARRAY_BACKEND") != "cupy":
+        return
+    skip_host_only = pytest.mark.skip(reason=_NEEDS_HOST_KERNELS_SKIP_REASON)
+    for item in items:
+        if "needs_host_kernels" in item.keywords:
+            item.add_marker(skip_host_only)
+
+
 def pytest_addoption(parser):
     parser.addoption("--with-desc", action="store_true")
     parser.addoption("--vrbose", action="store_true")

src/struphy/console/format.py

@@ -1211,7 +1211,15 @@ def confirm_formatting(python_files, linters, yes):
     )
     print("\n")
     if not yes:
-        ans = input("Format files (Y/n)?\n")
+        try:
+            ans = input("Format files (Y/n)?\n")
+        except EOFError:
+            # stdin isn't an interactive terminal (piped, scripted, non-tty
+            # shell, ...) -- input() can't prompt, so there's no way to get
+            # a real answer. Fail safe (don't format) instead of crashing
+            # with a raw traceback, and point at the flag that avoids this.
+            print("\nNo interactive terminal to confirm on. Exiting... (use --yes/-y to skip this prompt)")
+            sys.exit(1)
         if ans.lower() not in ("y", "yes", ""):
             print("Exiting...")
             sys.exit(1)

src/struphy/cuda.py

@@ -0,0 +1,63 @@
+"""Helpers for loading CUDA C sources used by CuPy RawKernel wrappers."""
+
+from functools import lru_cache
+from pathlib import Path
+from typing import Sequence
+
+
+@lru_cache(maxsize=None)
+def load_cuda_source(module_file: str, source_name: str) -> str:
+    """Load a CUDA C source fragment stored alongside its Python wrapper."""
+    path = Path(module_file).with_name("cuda") / source_name
+    return path.read_text(encoding="utf-8")
+
+
+class CudaKernel:
+    """A lazily-compiled, cached ``cupy.RawKernel``.
+
+    Every hand-written CUDA replacement in ``struphy.pic.*_cuda`` /
+    ``struphy.feec.*_cuda`` used to repeat the same boilerplate at each call
+    site: a module-level ``_foo_kernel = None`` sentinel, a
+    ``_get_foo_kernel()`` function that imports ``cupy`` and compiles the
+    ``RawKernel`` the first time it's needed (so importing these modules
+    under ``ARRAY_BACKEND=numpy`` never touches CuPy), and caches it back into
+    the global. This class replaces that boilerplate with one declaration:
+
+        _foo_kernel = CudaKernel(_FOO_SRC, "foo_cuda")
+
+    made once at module level, right next to the ``*_gpu`` function it backs.
+    Compilation is still deferred to first call (``import cupy`` only happens
+    inside :meth:`__call__`), and the compiled kernel is cached on the
+    instance -- identical behavior to the old pattern, but now every real GPU
+    kernel a module launches shows up as a ``CudaKernel(...)`` at module
+    scope, so `grep -n "CudaKernel("` (or just reading the top of the file)
+    tells you exactly which functions do device work and which don't.
+    """
+
+    __slots__ = ("_source", "_name", "_kernel")
+
+    def __init__(self, source: str, name: str) -> None:
+        self._source = source
+        self._name = name
+        self._kernel = None
+
+    def __call__(self, grid, block, args) -> None:
+        self._compiled()(grid, block, args)
+
+    def _compiled(self):
+        if self._kernel is None:
+            import cupy as cp
+
+            kernel = cp.RawKernel(self._source, self._name)
+            kernel.compile()
+            self._kernel = kernel
+        return self._kernel
+
+
+def launch_1d(kernel: CudaKernel, n: int, args: Sequence, threads: int = 256) -> None:
+    """Launch ``kernel`` over a 1-D grid with one thread per element of a
+    length-``n`` array (markers, indices, quadrature points, ...) -- the
+    launch geometry shared by every kernel in ``struphy.pic.*_cuda`` /
+    ``struphy.feec.*_cuda``."""
+    blocks = (n + threads - 1) // threads
+    kernel((blocks,), (threads,), tuple(args))

src/struphy/feec/preconditioner.py

@@ -1,6 +1,7 @@
 import logging
 
 import cunumpy as xp
+import numpy as np
 from feectools.api.essential_bc import apply_essential_bc_stencil
 from feectools.ddm.cart import CartDecomposition, DomainDecomposition
 from feectools.ddm.mpi import MockComm
@@ -260,7 +261,7 @@ def fun(e):
 
                 M_local = StencilMatrix(V_local, V_local)
 
-                row_indices, col_indices = xp.nonzero(M_arr)
+                row_indices, col_indices = np.nonzero(M_arr)
 
                 for row_i, col_i in zip(row_indices, col_indices):
                     # only consider row indices on process
@@ -273,7 +274,7 @@ def fun(e):
                         ] = M_arr[row_i, col_i]
 
                 # check if stencil matrix was built correctly
-                assert xp.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1])
+                assert np.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1])
 
                 matrixcells += [M_local.copy()]
                 # =======================================================================================================
@@ -625,7 +626,7 @@ def __init__(self, mass_operator, apply_bc=True):
 
                 M_local = StencilMatrix(V_local, V_local)
 
-                row_indices, col_indices = xp.nonzero(M_arr)
+                row_indices, col_indices = np.nonzero(M_arr)
 
                 for row_i, col_i in zip(row_indices, col_indices):
                     # only consider row indices on process
@@ -638,7 +639,7 @@ def __init__(self, mass_operator, apply_bc=True):
                         ] = M_arr[row_i, col_i]
 
                 # check if stencil matrix was built correctly
-                assert xp.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1])
+                assert np.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1])
 
                 matrixcells += [M_local.copy()]
                 # =======================================================================================================
@@ -911,10 +912,12 @@ class FFTSolver(BandedSolver):
     """
 
     def __init__(self, circmat):
-        assert isinstance(circmat, xp.ndarray)
+        # circmat comes from StencilMatrix.toarray(), always host numpy; the
+        # underlying solve() also calls scipy's solve_circulant, host-only.
+        assert isinstance(circmat, np.ndarray)
         assert is_circulant(circmat)
 
-        self._space = xp.ndarray
+        self._space = np.ndarray
         self._column = circmat[:, 0]
 
     # --------------------------------------
@@ -979,13 +982,15 @@ def is_circulant(mat):
         Whether the matrix is circulant (=True) or not (=False).
     """
 
-    assert isinstance(mat, xp.ndarray)
+    # mat comes from StencilMatrix.toarray(), which always densifies to a
+    # host numpy.ndarray regardless of the active backend.
+    assert isinstance(mat, np.ndarray)
     assert len(mat.shape) == 2
     assert mat.shape[0] == mat.shape[1]
 
     if mat.shape[0] > 1:
         for i in range(mat.shape[0] - 1):
-            circulant = xp.allclose(mat[i, :], xp.roll(mat[i + 1, :], -1))
+            circulant = np.allclose(mat[i, :], np.roll(mat[i + 1, :], -1))
             if not circulant:
                 return circulant
     else:

src/struphy/feec/psydac_derham.py

@@ -1925,11 +1925,12 @@ def _get_domain_array(self):
         else:
             nproc = 1
 
-        # send buffer
-        dom_arr_loc = xp.zeros(9, dtype=float)
+        # send buffer -- host-resident: mpi4py's Allgather needs a real
+        # buffer-protocol array, not a CuPy array, regardless of backend.
+        dom_arr_loc = np.zeros(9, dtype=float)
 
         # main array (receive buffers)
-        dom_arr = xp.zeros(nproc * 9, dtype=float)
+        dom_arr = np.zeros(nproc * 9, dtype=float)
 
         # Get global starts and ends of domain decomposition
         gl_s = self.domain_decomposition.starts

src/struphy/kernel_arguments/pusher_args_kernels.py

@@ -108,7 +108,6 @@ def __init__(
         self.bd2 = np.empty(int(pn[1]), dtype=float)
         self.bd3 = np.empty(int(pn[2]), dtype=float)
 
-
 class DomainArguments:
     """Holds the mandatory arguments pertaining to :class:`~struphy.geometry.base.Domain` passed to particle pusher kernels.