diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 314027510..462d18830 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -348,10 +348,10 @@ The three spaces (H1vec/Hcurl/Hdiv) differ only in the basis; port one, then the `push_v_sph_pressure`, `push_v_sph_pressure_ideal_gas`, `push_v_viscosity`, `div_u_weak_1form`. Infrastructure that gates the steps, independent of the kernels: mappings are no gate any more (every CUDA kernel -accepts every mapping since PR 19), except for what is still host-only around them: polar splines in `Derham` and MHD -equilibria such as `EQDSKequilibrium` cannot be created on CuPy (see [Open questions](#open-questions)); multi-rank marker sorting without host round trips. Array views with more than 4 dimensions -(all matrix accumulations write 6D stencil matrix data) are in cunumpy since 0.6.1 (`Array6D`); `linear_vlasov_ampere` is -the first kernel that uses them. +accepts every mapping since PR 19), except for what is still host-only around them: polar splines in `Derham` cannot be +created on CuPy (see [Open questions](#open-questions)). Multi-rank marker sorting stays on the device since #712. Array +views with more than 4 dimensions (all matrix accumulations write 6D stencil matrix data) are in cunumpy since 0.6.1 +(`Array6D`); `linear_vlasov_ampere` is the first kernel that uses them. ## Testing @@ -368,11 +368,39 @@ the first kernel that uses them. - **Marker layout.** The markers array is row-major (`n_markers × n_cols`). With one thread per marker, the memory accesses are strided. This is fine for now (each thread reads a few neighbouring columns), but a column-major or struct-of-arrays layout may be faster later. This would affect the CPU code too, so it is out of scope here. The array view represents strides explicitly on the CUDA side. - **MPI + GPUs.** One GPU per MPI rank (`xp.bind_local_device()` before `MPI_Init`, with feectools#86/#87), and GPU-aware MPI for the marker exchange, so markers do not go through the host. The marker exchange in `Particles.mpi_sort_markers` uses device buffers since #698 (see [Marker exchange implementation notes](#marker-exchange-implementation-notes-698)); the SPH ghost-box exchange (`_sendrecv_markers_boxes`) does not yet. - **Single-source alternatives.** Hand-written CUDA stays the default. Generating whole kernels from the Python source (`cupyx.jit`, numba-cuda, or a pyccel CUDA backend) is worth a look before the guiding-center kernels (the largest ones) are ported. Those tools take flat arguments, which `fields` also provides. -- **Polar splines and MHD equilibria on the GPU** (left after PR 19): spline mappings run on the device, but +- **Polar splines on the GPU** (left after PR 19): spline mappings run on the device, but `Derham` with `polar_splines=True` raises on CuPy (`PolarExtractionBlocksC1` builds SciPy sparse matrices, and the - polar extraction operators would apply them to device stencil data; needs `cupyx.scipy.sparse` or kernels), and - `EQDSKequilibrium` cannot be created on CuPy (its SciPy splines get device arrays). A `Tokamak` on CuPy builds its - default equilibrium on the host. Needed once a model with a polar domain or an EQDSK equilibrium runs on the GPU. + polar extraction operators would apply them to device stencil data; needs `cupyx.scipy.sparse` or kernels). Needed + once a model with a polar domain runs on the GPU. (MHD equilibria run on CuPy since #696, see + [MHD equilibria](#mhd-equilibria-on-cupy-696).) +- **Equilibrium splines on the device.** The SciPy splines of `EQDSKequilibrium`, `AdhocTorus` (`q_kind` 1, 2) and + `AdhocTorusQPsi`, and GVEC/DESC, are evaluated on the host with one copy per call (#696). Fine while equilibria are + evaluated only at setup; a device B-spline evaluation of their knots and coefficients would remove the copies. + +## MHD equilibria on CuPy (#696) + +- **What failed.** Creating `AdhocTorus` (`q_kind` 1, 2: SciPy `quad`/`UnivariateSpline`), `AdhocTorusQPsi` (`odeint`, + `fsolve`) and `EQDSKequilibrium` (`RectBivariateSpline` on `xp.linspace`) on CuPy; evaluating `GVECequilibrium` and + `DESCequilibrium` (gvec/DESC got device arrays). The analytic equilibria already worked. +- **Host setup.** `xp.setup_on_host` (cunumpy >= 0.6.2, which also provides `xp.host_call` and `xp.evaluate_on_host`) runs these `__init__`s on the NumPy backend, so the + equilibria hold only host data (floats, NumPy arrays, SciPy splines) on either backend. `Tokamak` no longer builds + its default `EQDSKequilibrium` on the NumPy backend itself; the field-line tracing still runs there. +- **Evaluation follows the arguments.** SciPy spline evaluations go through `xp.host_call`, and the GVEC/DESC `bv`, + `jv`, `p0`, `n0`, `gradB1` through `@xp.evaluate_on_host`: device arguments are copied to the host, evaluated on the NumPy + backend, and the result is copied back, once per call (equilibria are evaluated at setup; the time loop uses the + projected equilibrium). NumPy arguments are evaluated as before. +- **Tests.** `fields_background/tests/test_equils_cupy.py` creates every equilibrium of `equils` on CuPy, evaluates + the methods models call (meshgrid and markers, plus `psi`/`g_tor` with derivatives) and compares with NumPy: on a + GPU, and without one on the fake CuPy, where the four geometry kernels run their pyccel version on the fake arrays' + host buffers (`host_geometry_kernels`), since the fake CuPy cannot launch CUDA kernels. +- **GVEC at markers (#715).** gvec evaluated the markers' `rho`, `theta`, `zeta` as a tensor grid, so `bv`/`jv` at N + markers returned N x N x N arrays and `absB0` failed (on NumPy too). The coordinates are now passed as + `xarray.DataArray`s with one shared dimension, which gvec evaluates point by point. Boozer coordinates + (`use_boozer=True`) are computed per flux surface, so marker evaluation raises there. Test: `test_gvec_equil.py`. +- **Fake-CuPy child processes under MPI.** `run_fake_cupy_child` (`geometry/tests/test_domain.py`) starts the child + only on rank 0 (under `mpirun` every rank used to start the same child at once; on CI one of the two concurrent GVEC + children died with SIGILL and no output), with `faulthandler` and one OpenMP thread, and a failure reports the + signal and the end of stdout and stderr (gvec writes its Fortran messages to stdout). ## PR 10 implementation notes @@ -617,7 +645,7 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a NumPy backend, and the control points are copied to the active backend once. Found and tested without a GPU with cunumpy's fake CuPy (`CUNUMPY_FAKE_CUPY=1`), which rejects host/device mixing like CuPy. - **Still host-only.** `EQDSKequilibrium` cannot be created on CuPy (a `Tokamak` on CuPy takes an equilibrium created - on the NumPy backend, or builds its default one there). Polar splines: `PolarExtractionBlocksC1` builds SciPy sparse + on the NumPy backend, or builds its default one there; fixed in #696). Polar splines: `PolarExtractionBlocksC1` builds SciPy sparse matrices from the control points, and the polar extraction operators apply them to the stencil data, which lives on the device on CuPy; `Derham` raises `NotImplementedError` for `polar_splines=True` on CuPy (see [Open questions](#open-questions)). diff --git a/pyproject.toml b/pyproject.toml index 88547f445..87f42741f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ classifiers = [ ] dependencies = [ "numpy<=2.5.0", - "cunumpy >= 0.6.1, <= 0.6.1", + "cunumpy >= 0.6.2, <= 0.6.2", "pyccel>=2.2.0, <=2.2.3", "maybempi>=0.1.2", "feectools>=0.5.0, <=0.5.0", diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index 34dac281f..2377cebdc 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -874,6 +874,7 @@ def doc_formula(cls): Units are those defned in the parameter file (through :class:`~struphy.io.options.BaseUnits`). """ + @xp.setup_on_host def __init__( self, a: float = 1.0, @@ -1010,7 +1011,8 @@ def psi_r(self, r, der=0): if der == 0: out = -self.params["B0"] * self.params["a"] ** 2 / xp.sqrt(dq * q0 * eps**2 + dq**2) - out *= xp.arctanh( + # not in place: on CuPy, the prefactor is a 0-d device array + out = out * xp.arctanh( xp.sqrt((dq - dq * (r / self.params["R0"]) ** 2) / (q0 * eps**2 + dq)), ) elif der == 1: @@ -1020,7 +1022,7 @@ def psi_r(self, r, der=0): # alternative profile (interpolated) elif self.params["q_kind"] == 1 or self.params["q_kind"] == 2: - out = self._psi_i(r, nu=der) + out = xp.host_call(self._psi_i, r, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(r, (int, float)): @@ -1149,7 +1151,7 @@ def p_r(self, r): # alternative profiles (interpolated) elif self.params["q_kind"] == 1 or self.params["q_kind"] == 2: - pout = self._p_i(r) + pout = xp.host_call(self._p_i, r) # remove all "dimensions" for point-wise evaluation if isinstance(r, (int, float)): @@ -1365,6 +1367,7 @@ def doc_formula(cls): Units are those defned in the parameter file (through :class:`~struphy.io.options.BaseUnits`). """ + @xp.setup_on_host def __init__( self, a: float = 0.361925, @@ -1470,7 +1473,7 @@ def psi_r(self, r, der=0): assert der >= 0 and der <= 2, "Only first and second derivatives available!" - out = self._psi_i(r, nu=der) + out = xp.host_call(self._psi_i, r, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(r, (int, float)): @@ -1665,6 +1668,7 @@ class EQDSKequilibrium(AxisymmMHDequilibrium): Struphy base units. If None, no rescaling of output is performed. """ + @xp.setup_on_host def __init__( self, rel_path: bool = True, @@ -1864,7 +1868,7 @@ def psi_axis_RZ(self): def q_psi(self, psi, der=0): """Safety factor q = q(psi).""" - out = self._q_i(psi, nu=der) + out = xp.host_call(self._q_i, psi, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(psi, (int, float)): @@ -1875,7 +1879,7 @@ def q_psi(self, psi, der=0): def g_psi(self, psi, der=0): """Toroidal field function g = g(psi).""" - out = self._g_i(psi, nu=der) + out = xp.host_call(self._g_i, psi, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(psi, (int, float)): @@ -1886,7 +1890,7 @@ def g_psi(self, psi, der=0): def p_psi(self, psi, der=0): """Pressure profile p = p(psi) in units Pa (as in the EQDSK file).""" - out = self._p_i(psi, nu=der) + out = xp.host_call(self._p_i, psi, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(psi, (int, float)): @@ -1922,7 +1926,7 @@ def psi(self, R, Z, dR=0, dZ=0): is_float = all(isinstance(v, (int, float)) for v in [R, Z]) - out = self._psi_i(R, Z, dx=dR, dy=dZ, grid=False) + out = xp.host_call(self._psi_i, R, Z, dx=dR, dy=dZ, grid=False) # remove all "dimensions" for point-wise evaluation if is_float: @@ -2121,6 +2125,7 @@ def units(self) -> Units: """All Struphy units.""" return self._units + @xp.evaluate_on_host @profile def bv(self, *etas, squeeze_out=False): """Contra-variant (vector field) magnetic field on logical cube [0, 1]^3 in Tesla / meter.""" @@ -2142,6 +2147,7 @@ def bv(self, *etas, squeeze_out=False): return out + @xp.evaluate_on_host @profile def jv(self, *etas, squeeze_out=False): """Contra-variant (vector field) current density (=curl B) on logical cube [0, 1]^3 in Ampere / meter^3.""" @@ -2171,6 +2177,7 @@ def jv(self, *etas, squeeze_out=False): return out + @xp.evaluate_on_host @profile def p0(self, *etas, squeeze_out=False): """0-form equilibrium pressure on logical cube [0, 1]^3.""" @@ -2190,6 +2197,7 @@ def p0(self, *etas, squeeze_out=False): return self.params["p0"] + tmp / self.units.p + @xp.evaluate_on_host @profile def n0(self, *etas, squeeze_out=False): """0-form equilibrium density on logical cube [0, 1]^3.""" @@ -2269,7 +2277,17 @@ def _gvec_evaluations(self, *etas): # evaluate if self.params["use_boozer"]: + if flat_eval: + # the Boozer transform is computed per flux surface, so gvec cannot evaluate it at scattered points + raise NotImplementedError("GVECequilibrium: marker evaluation is not available with use_boozer=True.") ev = gvec.EvaluationsBoozer(rho=rho, theta_B=theta, zeta_B=zeta, state=self.state) + elif flat_eval: + import xarray as xr + + # coordinates sharing one dimension make gvec evaluate point by point instead of on the tensor grid + # rho x theta x zeta (one value per marker, not n_markers**3) + rho, theta, zeta = (xr.DataArray(c, dims="marker") for c in (rho, theta, zeta)) + ev = gvec.Evaluations(rho=rho, theta=theta, zeta=zeta, state=self.state) else: ev = gvec.Evaluations(rho=rho, theta=theta, zeta=zeta, state=self.state) @@ -2407,6 +2425,7 @@ def units(self) -> Units: """All Struphy units.""" return self._units + @xp.evaluate_on_host @profile def bv(self, *etas, squeeze_out=False): """Contra-variant (vector field) magnetic field on logical cube [0, 1]^3 in Tesla / meter.""" @@ -2480,6 +2499,7 @@ def _eval_bv(self, *etas, squeeze_out=False): return out + @xp.evaluate_on_host @profile def jv(self, *etas, squeeze_out=False): """Contra-variant (vector field) current density (=curl B) @@ -2555,6 +2575,7 @@ def _eval_jv(self, *etas, squeeze_out=False): return out + @xp.evaluate_on_host @profile def p0(self, *etas, squeeze_out=False): """0-form equilibrium pressure on logical cube [0, 1]^3 in Pascal.""" @@ -2586,6 +2607,7 @@ def p0(self, *etas, squeeze_out=False): return out + @xp.evaluate_on_host @profile def n0(self, *etas, squeeze_out=False): """0-form equilibrium density on logical cube [0, 1]^3.""" @@ -2610,6 +2632,7 @@ def n0(self, *etas, squeeze_out=False): # density in default units, n=1 --> 10^20 m^(-3) return p0_pascal / (self.params["T_kelvin"] * k_Boltzmann) / self.units.n + @xp.evaluate_on_host @profile def gradB1(self, *etas, squeeze_out=False): """1-form gradient of magnetic field strength on logical cube [0, 1]^3.""" diff --git a/src/struphy/fields_background/tests/test_equils_cupy.py b/src/struphy/fields_background/tests/test_equils_cupy.py new file mode 100644 index 000000000..e7710570d --- /dev/null +++ b/src/struphy/fields_background/tests/test_equils_cupy.py @@ -0,0 +1,269 @@ +"""Fluid and MHD equilibria on the CuPy backend, compared with the NumPy backend. + +The equilibria are created and evaluated on CuPy (with a GPU) or on cunumpy's fake CuPy (without one, in a +subprocess), which rejects host/device mixing like CuPy, and every evaluation is compared with the NumPy backend. +""" + +import contextlib +import importlib.util +import inspect + +import cunumpy +import numpy as np +import pytest +from cunumpy.kernel_testing import requires_cupy + +from struphy.geometry.tests.test_domain import _cupy_installed, run_fake_cupy_child + +# (equilibrium, its parameters, domain, its parameters); domain None for numerical equilibria (own domain). +# Every equilibrium of `equils` is here (`test_all_equilibria_have_cases`); GVEC and DESC need their packages. +EQUIL_CASES = { + "HomogenSlab": ("HomogenSlab", {}, "Cuboid", {}), + "ShearedSlab": ("ShearedSlab", {}, "Cuboid", {"r1": 1.0, "r2": 2 * np.pi, "r3": 2 * np.pi * 10.0}), + "ShearFluid": ("ShearFluid", {}, "Cuboid", {}), + "ScrewPinch": ("ScrewPinch", {}, "HollowCylinder", {"a1": 0.05, "a2": 1.0, "Lz": 2 * np.pi * 5.0}), + "ScrewPinch-q_inf": ( + "ScrewPinch", + {"q0": "inf", "q1": "inf"}, + "HollowCylinder", + {"a1": 0.05, "a2": 1.0, "Lz": 2 * np.pi * 5.0}, + ), + "AdhocTorus-q0p0": ("AdhocTorus", {"q_kind": 0, "p_kind": 0}, "HollowTorus", {"a1": 0.05, "a2": 1.0}), + "AdhocTorus-q1p0": ("AdhocTorus", {"q_kind": 1, "p_kind": 0}, "HollowTorus", {"a1": 0.05, "a2": 1.0}), + "AdhocTorus-q2p1": ("AdhocTorus", {"q_kind": 2, "p_kind": 1}, "HollowTorus", {"a1": 0.05, "a2": 1.0}), + "AdhocTorus-Tokamak": ("AdhocTorus", {"q_kind": 1, "p_kind": 0}, "Tokamak", {"num_elements": (6, 16)}), + "AdhocTorusQPsi": ("AdhocTorusQPsi", {}, "IGAPolarTorus", {"a": 0.361925, "R0": 1.0, "num_elements": (6, 16)}), + "CircularTokamak": ("CircularTokamak", {}, "HollowTorus", {"a1": 0.05, "a2": 1.0, "R0": 2.0}), + "EQDSKequilibrium": ("EQDSKequilibrium", {}, "Tokamak", {"num_elements": (6, 16)}), + "GVECequilibrium": ( + "GVECequilibrium", + { + "dat_file": "run_01/CIRCTOK_State_0000_00000000.dat", + "param_file": "run_01/parameter.ini", + "num_elements": (6, 6, 4), + }, + None, + {}, + ), + "DESCequilibrium": ("DESCequilibrium", {"num_elements": (6, 6, 6)}, None, {}), + "ConstantVelocity": ("ConstantVelocity", {}, "Cuboid", {}), + "HomogenSlabITG": ("HomogenSlabITG", {}, "Cuboid", {}), + "CurrentSheet": ("CurrentSheet", {}, "Cuboid", {}), + "GenericCartesianFluidEquilibrium": ("GenericCartesianFluidEquilibrium", {}, "Cuboid", {}), + "GenericCartesianFluidEquilibriumWithB": ("GenericCartesianFluidEquilibriumWithB", {}, "Cuboid", {}), +} + +# equilibria that need an optional package (skipped without it) +NEEDS = {"GVECequilibrium": "gvec", "DESCequilibrium": "desc"} + +# methods that models and projections call on the logical domain; those an equilibrium does not provide are skipped +METHODS = ( + "absB0", + "absB3", + "p0", + "p3", + "n0", + "n3", + "t0", + "vth0", + "b1", + "b2", + "bv", + "b_cart", + "unit_b1", + "unit_b2", + "unit_bv", + "gradB1", + "gradB2", + "gradBv", + "j1", + "j2", + "jv", + "absJ0", + "curl_unit_b1", + "curl_unit_b2", + "curl_unit_b_dot_b0", + "u1", + "u2", + "uv", + "a1", + "a2", +) + +# methods of axisymmetric equilibria in (R, Z) (flux and toroidal field function with derivatives) +AXISYMM_DERIVATIVES = ((0, 0), (1, 0), (0, 1), (2, 0), (0, 2), (1, 1)) + + +def _build(case): + """Create the equilibrium of ``case`` with its domain on the active backend.""" + from struphy import domains, equils + + name, params, dom_name, dom_params = EQUIL_CASES[case] + equil = getattr(equils, name)(**params) + if dom_name is None: + return equil + if dom_name == "Tokamak": + domain = domains.Tokamak(equilibrium=equil, **dom_params) + else: + domain = getattr(domains, dom_name)(**dom_params) + equil.domain = domain + return equil + + +def _flatten(out): + """Arrays in ``out`` (an array, or nested tuples/lists of arrays).""" + if isinstance(out, (tuple, list)): + return [a for o in out for a in _flatten(o)] + return [out] + + +def _evaluations(equil, args): + """Evaluate all methods of ``equil`` in ``METHODS`` (and psi, g_tor for axisymmetric ones) at ``args``.""" + from struphy.fields_background.base import AxisymmMHDequilibrium + + out = {} + for name in METHODS: + method = getattr(equil, name, None) + if method is None: + continue + try: + out[name] = method(*args) + except (AssertionError, NotImplementedError) as error: + # not available for this equilibrium (vector potential, GVEC gradB1, ...) + out[name] = type(error) + if isinstance(equil, AxisymmMHDequilibrium): + R = 1 * equil.psi_axis_RZ[0] + 0.3 * args[0] + Z = 1 * equil.psi_axis_RZ[1] + 0.2 * args[0] - 0.1 + for dR, dZ in AXISYMM_DERIVATIVES: + out[f"psi_{dR}{dZ}"] = equil.psi(R, Z, dR=dR, dZ=dZ) + for dR, dZ in AXISYMM_DERIVATIVES[:3]: + out[f"g_tor_{dR}{dZ}"] = equil.g_tor(R, Z, dR=dR, dZ=dZ) + return out + + +@contextlib.contextmanager +def host_geometry_kernels(): + """Fake CuPy only: run the geometry kernels of :class:`~struphy.geometry.base.Domain` with their pyccel versions. + + The fake CuPy cannot launch CUDA kernels. Here the four geometry entry kernels run their pyccel version on the + host buffers of the fake device arrays (shared, so the outputs are written in place), with the pyccel + ``DomainArguments`` built from the same buffers. The equilibria themselves run unchanged on the fake CuPy. The + CUDA geometry kernels are checked against pyccel in ``pic/tests/test_cuda_parity.py`` (GPU) and + ``test_cuda_emulation.py`` (CPU emulation). + """ + import struphy.geometry.base as geometry_base + from struphy.kernel_arguments.pusher_args_cuda import CudaDomainArguments + from struphy.kernel_arguments.pusher_args_kernels import DomainArguments + + def host(arg): + if isinstance(arg, CudaDomainArguments): + return DomainArguments(arg.kind_map, *(host(getattr(arg, name)) for name, _ in arg.fields[1:])) + return arg._a if cunumpy.is_gpu(arg) else arg # the host buffer of a fake CuPy array + + def on_host(kernel): + def launch(*args, **launch_options): + return kernel.host_kernel.kernel(*(host(a) for a in args)) + + return launch + + names = ("kernel_evaluate", "kernel_evaluate_pic", "kernel_pullpush", "kernel_pullpush_pic") + kernels = {name: getattr(geometry_base, name) for name in names} + try: + for name, kernel in kernels.items(): + setattr(geometry_base, name, on_host(kernel)) + yield + finally: + for name, kernel in kernels.items(): + setattr(geometry_base, name, kernel) + + +def check_equil_on_cupy(case): + """Create ``case`` on the CuPy backend, evaluate it there and compare with the NumPy backend. + + Every result is a device array (or a scalar for constant parts like ``psi`` derivatives of analytic profiles) and + agrees with the NumPy result. Evaluated on a meshgrid (three 1d arrays) and at markers (one 2d array). + """ + rng = np.random.default_rng(1234) + e1 = np.sort(rng.uniform(0.1, 0.9, 4)) + e2 = np.sort(rng.uniform(0.0, 1.0, 5)) + e3 = np.sort(rng.uniform(0.0, 1.0, 3)) + markers = np.column_stack([rng.uniform(0.1, 0.9, 7), rng.uniform(0.0, 1.0, 7), rng.uniform(0.0, 1.0, 7)]) + + with cunumpy.use_backend("numpy"): + equil = _build(case) + ref_grid = _evaluations(equil, (e1, e2, e3)) + ref_markers = _evaluations(equil, (markers,)) + + with cunumpy.use_backend("cupy"): + equil = _build(case) + out_grid = _evaluations(equil, tuple(cunumpy.asarray(e) for e in (e1, e2, e3))) + out_markers = _evaluations(equil, (cunumpy.asarray(markers),)) + + for ref, out, where in ((ref_grid, out_grid, "meshgrid"), (ref_markers, out_markers, "markers")): + assert ref.keys() == out.keys() + for name in ref: + if isinstance(ref[name], type): # not available on NumPy, not on CuPy either + assert out[name] is ref[name], (case, where, name) + continue + refs, outs = _flatten(ref[name]), _flatten(out[name]) + assert len(refs) == len(outs), (case, where, name) + for r, o in zip(refs, outs): + if isinstance(r, np.ndarray) and r.ndim > 0: + assert cunumpy.is_gpu(o), (case, where, name, type(o)) + o = cunumpy.to_numpy(o) if cunumpy.is_gpu(o) else np.asarray(o) + assert np.allclose(o, r, rtol=1e-12, atol=1e-12, equal_nan=True), (case, where, name) + + +def _check_needs(case): + package = NEEDS.get(EQUIL_CASES[case][0]) + if package is not None and importlib.util.find_spec(package) is None: + pytest.skip(f"{package} is not installed") + + +def test_all_equilibria_have_cases(): + """Every equilibrium class of ``equils`` is checked here.""" + from struphy import equils + from struphy.fields_background.base import FluidEquilibrium + + classes = { + name + for name, cls in vars(equils).items() + if inspect.isclass(cls) + and issubclass(cls, FluidEquilibrium) + and not inspect.isabstract(cls) + and cls.__module__ == equils.__name__ + } + covered = {case[0] for case in EQUIL_CASES.values()} + assert classes == covered + + +def test_host_call_numpy_is_plain_call(): + """On the NumPy backend, :func:`cunumpy.host_call` calls the function directly and returns its result unchanged.""" + a = np.linspace(0.0, 1.0, 5) + with cunumpy.use_backend("numpy"): + out = cunumpy.host_call(lambda x, y=1.0: x * y, a, y=2.0) + assert type(out) is np.ndarray + assert np.array_equal(out, 2 * a) + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +@pytest.mark.parametrize("case", list(EQUIL_CASES)) +def test_equil_fake_cupy(case): + """Without a GPU: :func:`check_equil_on_cupy` with cunumpy's fake CuPy, which rejects host/device mixing. + + Runs in a subprocess because the fake CuPy must be installed before cunumpy is imported. + """ + _check_needs(case) + code = ( + "from struphy.fields_background.tests.test_equils_cupy import check_equil_on_cupy, host_geometry_kernels\n" + f"with host_geometry_kernels(): check_equil_on_cupy({case!r})" + ) + run_fake_cupy_child(code) + + +@requires_cupy +@pytest.mark.parametrize("case", list(EQUIL_CASES)) +def test_equil_cupy(case): + """On a GPU: the equilibria are created and evaluated on CuPy and agree with NumPy.""" + _check_needs(case) + check_equil_on_cupy(case) diff --git a/src/struphy/fields_background/tests/test_gvec_equil.py b/src/struphy/fields_background/tests/test_gvec_equil.py new file mode 100644 index 000000000..40b84e065 --- /dev/null +++ b/src/struphy/fields_background/tests/test_gvec_equil.py @@ -0,0 +1,47 @@ +import importlib.util + +import numpy as np +import pytest + +pytestmark = pytest.mark.skipif(importlib.util.find_spec("gvec") is None, reason="gvec is not installed") + +GVEC_PARAMS = {"dat_file": "run_01/CIRCTOK_State_0000_00000000.dat", "param_file": "run_01/parameter.ini"} + + +def _markers(n=7, seed=3): + """``n`` markers at random logical positions, away from the hole at the axis.""" + rng = np.random.default_rng(seed) + markers = rng.uniform(0.0, 1.0, (n, 3)) + markers[:, 0] = rng.uniform(0.1, 0.9, n) + return markers + + +@pytest.mark.parametrize("name", ["bv", "jv", "p0", "n0", "absB0"]) +def test_gvec_marker_evaluation_matches_grid(name): + """Marker (flat) evaluation gives one value per marker, equal to the grid evaluation at each marker. + + Before, gvec evaluated the markers' rho, theta and zeta as a tensor grid, so ``bv`` and ``jv`` at N markers + returned N x N x N arrays, and ``absB0`` failed. + """ + from struphy.fields_background.equils import GVECequilibrium + + equil = GVECequilibrium(**GVEC_PARAMS) + markers = _markers() + flat = getattr(equil, name)(markers) + flat = flat if isinstance(flat, tuple) else (flat,) + + for ip, (e1, e2, e3) in enumerate(markers): + grid = getattr(equil, name)(np.array([e1]), np.array([e2]), np.array([e3])) + grid = grid if isinstance(grid, tuple) else (grid,) + for f, g in zip(flat, grid): + assert f.shape == (len(markers),) + assert np.isclose(f[ip], g.ravel()[0], rtol=1e-12, atol=1e-14) + + +def test_gvec_marker_evaluation_boozer_raises(): + """gvec computes the Boozer transform per flux surface, so marker evaluation with ``use_boozer=True`` raises.""" + from struphy.fields_background.equils import GVECequilibrium + + equil = GVECequilibrium(**GVEC_PARAMS, use_boozer=True) + with pytest.raises(NotImplementedError, match="use_boozer"): + equil.bv(_markers()) diff --git a/src/struphy/geometry/domains/tokamak/tokamak.py b/src/struphy/geometry/domains/tokamak/tokamak.py index c53bb7882..1fa3d7fcc 100644 --- a/src/struphy/geometry/domains/tokamak/tokamak.py +++ b/src/struphy/geometry/domains/tokamak/tokamak.py @@ -67,11 +67,10 @@ def __init__( ): if r_min != 0.0: r0 = r_min - # The equilibrium and the field-line tracing (SciPy) are host-only setup: they run on the NumPy - # backend, and the control points are copied to the active backend once. + # The field-line tracing (SciPy) is host-only setup: it runs on the NumPy backend, and the control points + # are copied to the active backend once. EQDSKequilibrium does its own setup (file, SciPy splines) on the host. if equilibrium is None: - with xp.use_backend("numpy"): - equilibrium = EQDSKequilibrium() + equilibrium = EQDSKequilibrium() else: assert isinstance(equilibrium, AxisymmMHDequilibrium) diff --git a/src/struphy/geometry/tests/test_domain.py b/src/struphy/geometry/tests/test_domain.py index 495f9546d..466d138d9 100644 --- a/src/struphy/geometry/tests/test_domain.py +++ b/src/struphy/geometry/tests/test_domain.py @@ -3,6 +3,7 @@ import logging import os import pickle +import signal import subprocess import sys @@ -1199,6 +1200,34 @@ def serial_child_env(**extra): return env +def run_fake_cupy_child(code: str): + """Run ``code`` in a serial child Python process on cunumpy's fake CuPy and fail the test if the child fails. + + The check is serial, so under ``mpirun`` only rank 0 starts the child; the other ranks skip. Before, every rank + started the same child at the same time, and on CI the GVEC case crashed (SIGILL) in one of the two concurrent + children with no output. The child runs with ``faulthandler`` (a crash prints the Python traceback) and one OpenMP + thread, and a failure reports the signal or exit code and the end of both output streams (GVEC writes its Fortran + messages to stdout). + """ + from maybempi import MPI + + if MPI.COMM_WORLD.Get_rank() != 0: + pytest.skip("serial check, runs on MPI rank 0") + result = subprocess.run( + [sys.executable, "-X", "faulthandler", "-c", code], + env=serial_child_env(CUNUMPY_FAKE_CUPY="1", OMP_NUM_THREADS="1"), + capture_output=True, + text=True, + ) + if result.returncode != 0: + rc = result.returncode + how = f"signal {signal.Signals(-rc).name}" if rc < 0 else f"exit code {rc}" + pytest.fail( + f"child process failed with {how}\n" + f"--- stdout (end) ---\n{result.stdout[-2000:]}\n--- stderr (end) ---\n{result.stderr[-4000:]}", + ) + + @pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") @pytest.mark.parametrize("mapping", CUDA_DOMAIN_MAPPINGS) def test_cuda_args_domain_fake_cupy(mapping): @@ -1207,10 +1236,7 @@ def test_cuda_args_domain_fake_cupy(mapping): Runs in a subprocess because the fake CuPy must be installed before cunumpy is imported. """ code = f"from struphy.geometry.tests.test_domain import check_cuda_args_domain; check_cuda_args_domain({mapping!r})" - result = subprocess.run( - [sys.executable, "-c", code], env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), capture_output=True, text=True - ) - assert result.returncode == 0, result.stderr[-4000:] + run_fake_cupy_child(code) @requires_cupy diff --git a/src/struphy/pic/tests/test_accum_matrix_cupy.py b/src/struphy/pic/tests/test_accum_matrix_cupy.py index cf0aa62c6..4537e8670 100644 --- a/src/struphy/pic/tests/test_accum_matrix_cupy.py +++ b/src/struphy/pic/tests/test_accum_matrix_cupy.py @@ -6,8 +6,6 @@ cunumpy's fake CuPy with CPU-emulated kernel launches (in a subprocess, see :func:`check_accumulator_matrix_fake_cupy`). """ -import subprocess -import sys from dataclasses import dataclass from typing import Any @@ -16,7 +14,7 @@ import pytest from cunumpy.kernel_testing import emulation_compiler, requires_cupy -from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env +from struphy.geometry.tests.test_domain import _cupy_installed, run_fake_cupy_child ALPHA, KAPPA, VTH, DT = 1.3, 0.7, 0.9, 0.05 N_SCHUR_ITERATIONS = 3 @@ -203,7 +201,4 @@ def test_accumulator_matrix_fake_cupy(): subprocess because the fake CuPy must be installed before cunumpy is imported. """ code = "from struphy.pic.tests.test_accum_matrix_cupy import check_accumulator_matrix_fake_cupy as c; c()" - result = subprocess.run( - [sys.executable, "-c", code], env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), capture_output=True, text=True - ) - assert result.returncode == 0, result.stderr[-4000:] + run_fake_cupy_child(code) diff --git a/src/struphy/simulation/tests/test_compile_cuda_kernels.py b/src/struphy/simulation/tests/test_compile_cuda_kernels.py index 5e7114366..97bec9e23 100644 --- a/src/struphy/simulation/tests/test_compile_cuda_kernels.py +++ b/src/struphy/simulation/tests/test_compile_cuda_kernels.py @@ -5,15 +5,12 @@ (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.geometry.tests.test_domain import _cupy_installed, run_fake_cupy_child 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 @@ -225,13 +222,7 @@ def test_compile_cuda_kernels_on_cupy_backend(tmp_path, monkeypatch): @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:] + run_fake_cupy_child(FAKE_CUPY_CHECK) @requires_cupy