From 1ad4cd6fba1146bbf0ebbd1364bd8218b935723a Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 19:12:45 +0200 Subject: [PATCH 01/50] CUDA version of the linear_vlasov_ampere accumulation - linear_vlasov_ampere_cuda.cu: same arguments as pyccel, one thread per marker row, matrix blocks as Array6D views (cunumpy 0.6.1) - device fill_mat, fill_mat_vec (filler_kernels.cuh), m_v_fill_b_v1_symm (particle_to_mat_kernels.cuh) and outer (linalg_kernels.cuh), with atomic adds - parity cases (Cuboid, Colella, HollowTorus, 3d spline) and a device-helper test of m_v_fill_b_v1_symm, on a GPU and by CPU emulation - test_accum_matrix_cupy.py: the matrix path of Accumulator on CuPy vs NumPy (device data, ghost regions, transposed blocks, Schur operator and solve), without a GPU on the fake CuPy with emulated kernel launches (emulated_launches() in pic/tests/cuda_emulation.py) - CUDA_STRATEGY.md: porting order and implementation notes Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 60 ++++- src/struphy/linear_algebra/linalg_kernels.cuh | 11 + .../pic/accumulation/filler_kernels.cuh | 125 ++++++++++- .../linear_vlasov_ampere_cuda.cu | 90 ++++++++ .../accumulation/particle_to_mat_kernels.cuh | 69 ++++++ src/struphy/pic/tests/cuda_emulation.py | 66 +++++- src/struphy/pic/tests/cuda_parity_cases.py | 17 ++ src/struphy/pic/tests/kernel_test_args.py | 27 ++- .../pic/tests/test_accum_matrix_cupy.py | 209 ++++++++++++++++++ src/struphy/pic/tests/test_cuda_emulation.py | 12 + src/struphy/pic/tests/test_cuda_parity.py | 4 +- src/struphy/pic/tests/test_device_helpers.py | 61 +++++ 12 files changed, 731 insertions(+), 20 deletions(-) create mode 100644 src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu create mode 100644 src/struphy/pic/tests/test_accum_matrix_cupy.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index daf0d25f0..d4f29568d 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -26,7 +26,8 @@ Progress is tracked in struphy-hub/struphy#650. - [ ] **PR 17: struphy on the feectools CUDA stack** — the `feectools` submodule points at the top of the feectools CUDA stack (`cuda-4-device-kernels`, [feectools#88](https://github.com/struphy-hub/feectools/pull/88)) instead of `devel-tiny`, so the struphy CUDA PRs run against feectools with device stencil operations, MPI with device buffers and one GPU per rank. Moves along with the stack; before the struphy stack is merged into `devel`, the feectools stack is merged into `devel-tiny` and the submodule points there again (see [feectools](#feectools)). - [x] **PR 18: Geometry evaluation on the GPU for all analytic mappings** (#681, open) — CUDA versions of the four geometry entry kernels (`kernel_evaluate_pic`, `kernel_evaluate`, `kernel_pullpush_pic`, `kernel_pullpush` in `geometry/kernels/`), built on device versions of the whole metric chain (`f`, `df`, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`/`push`/`tran`) for every analytic mapping (`kind_map` 10–12, 20–22, 30–32). Restores CuPy particle runs (weight initialization evaluates `jacobian_det`), and removes the Cuboid-only checks in `Pusher`, the accumulators and `reflect`. Parity arguments cover every analytic mapping (see [PR 18](#pr-18-geometry-evaluation-for-all-analytic-mappings) and the [PR 18 implementation notes](#pr-18-implementation-notes)). - [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines and `EQDSKequilibrium` stay host-only (see [PR 19](#pr-19-spline-mappings) and the [PR 19 implementation notes](#pr-19-implementation-notes)). -- [ ] **Next (order to be confirmed)**: 6D array views in cunumpy and the blocked matrix accumulations `vlasov_maxwell`, `linear_vlasov_ampere` (steps 1–2 of the [porting order](#porting-order)); one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, including when to compile the kernels, which needs the [feectools stack](#feectools) merged and released first; a decision on hand-written CUDA vs. code generation before the guiding-center kernels (step 3). +- [x] **`linear_vlasov_ampere` on the GPU** (part of #688) — the first matrix accumulation: CUDA version of `linear_vlasov_ampere` (`EfieldWeightsCoupling`) on cunumpy's `Array6D` views (cunumpy 0.6.1), device versions of `fill_mat`, `fill_mat_vec`, `m_v_fill_b_v1_symm` and `outer`, and the matrix path of `Accumulator` checked on the CuPy backend (device data, ghost regions, transposed blocks, Schur operator) without a GPU through cunumpy's fake CuPy with emulated kernel launches (see the [implementation notes](#linear_vlasov_ampere-implementation-notes)). +- [ ] **Next (order to be confirmed)**: the blocked matrix accumulation `vlasov_maxwell` (step 1 of the [porting order](#porting-order)); one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, including when to compile the kernels, which needs the [feectools stack](#feectools) merged and released first; a decision on hand-written CUDA vs. code generation before the guiding-center kernels (step 3). - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)). - [x] **PR 13: Move the kernel infrastructure to cunumpy** — `Kernel`, `KernelCatalog`, `CudaKernel` and `Argument` are replaced by their cunumpy counterparts; each owner has a single `args_*` object on both backends, and `pusher_args.cuh` is generated (see [Moving to cunumpy](#moving-to-cunumpy-pr-13)). Kernels and device helpers only change their includes. - [ ] **CI**: a GPU runner that runs the CUDA tests (can happen any time; until then the GPU tests are run by hand on an H100 before each PR that touches CUDA code is merged). @@ -64,7 +65,7 @@ CUDA kernels can be added one by one. If the code runs on the GPU and needs a ke | `src/struphy/utils/cuda_arguments.py` | `CUDA_STRUCTS`, `CUDA_OPTIONS`, `write_pusher_header()` and `write_local_projectors_header()` | | `src/struphy/kernel_arguments/pusher_args.cuh`, `local_projectors_args.cuh` | the C structs `MarkerArgs`, `DerhamArgs`, `DomainArgs` and `LocalProjectorsArgs`, generated from the CUDA classes, using cunumpy array views | | `src/struphy/geometry/base.py`, `src/struphy/pic/base.py`, `src/struphy/feec/psydac_derham.py` | `Domain.args_domain`, `Particles.args_markers` and `Derham.args_derham` are the pyccel class on NumPy and the CUDA class on CuPy, chosen once at construction; every kernel call goes through a `Kernel` | -| `src/struphy/*/kernels/` (`pic/pushing`, `pic/accumulation`, `pic/diagnostics`, `pic/sph`, `bsplines`, `geometry`, `feec`, `feec/local_projectors`) | every kernel called from Python with argument objects, one folder each: 44 pusher/evaluation (incl. `reflect`), 16 accumulation, 10 marker diagnostics, 4 SPH evaluation, 3 spline evaluation, 4 geometry, 1 FEEC utility and 8 local projector kernels. Each folder's `__init__.py` declares its `Kernel`; the code imports it (`from struphy.geometry.kernels.kernel_evaluate import kernel_evaluate`). Fifteen have CUDA versions, with parity cases in `pic/tests/cuda_parity_cases.py`, among them the four geometry kernels (PR 18) | +| `src/struphy/*/kernels/` (`pic/pushing`, `pic/accumulation`, `pic/diagnostics`, `pic/sph`, `bsplines`, `geometry`, `feec`, `feec/local_projectors`) | every kernel called from Python with argument objects, one folder each: 44 pusher/evaluation (incl. `reflect`), 16 accumulation, 10 marker diagnostics, 4 SPH evaluation, 3 spline evaluation, 4 geometry, 1 FEEC utility and 8 local projector kernels. Each folder's `__init__.py` declares its `Kernel`; the code imports it (`from struphy.geometry.kernels.kernel_evaluate import kernel_evaluate`). Sixteen have CUDA versions (`linear_vlasov_ampere` the latest), with parity cases in `pic/tests/cuda_parity_cases.py`, among them the four geometry kernels (PR 18) | | `src/struphy/geometry/evaluation_kernels.cuh`, `transform_kernels.cuh`, `spline_mappings_kernels.cuh`, `geometry/domains//_cuda.cuh` | device versions of the mapping helpers for every mapping: `f`/`df` per analytic domain (`kind_map` 10–12, 20–22, 30–32, PR 18) and the spline mappings `spline_3d`, `spline_2d_straight`, `spline_2d_torus` (`kind_map` 0–2, PR 19), the `kind_map` switch, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`, `push`, `tran` | | `src/struphy/pic/tests/test_cuda_parity.py`, `test_cuda_emulation.py` | parity of every CUDA kernel with pyccel on the cases of `pic/tests/cuda_parity_cases.py` (`PARITY_CASES`): on a GPU (`assert_kernels_agree`), and without one by CPU emulation | | `src/struphy/pic/tests/test_kernel_backends.py` | the generated-header tests, the check that each CUDA class mirrors its pyccel class, signature check of all kernels, owner class selection, and (on a GPU) struct layout and copy/pickle checks | @@ -277,7 +278,7 @@ from the struphy tracking issue ([#650](https://github.com/struphy-hub/struphy/i Kernels are ported in the order that completes one model after the other on the GPU. Each step lists only the kernels it adds; P = pusher, A = accumulation, E = marker evaluation. Status: -done (✓), in this PR (PR 14), blocked (⏸). +done (✓), in a PR (PR 14), blocked (⏸). **Step 0 – Vlasov** (PRs 11–12, ✓): `push_eta_stage`, `push_vxb_analytic`, `push_vxb_implicit`, `reflect`. @@ -289,14 +290,14 @@ done (✓), in this PR (PR 14), blocked (⏸). |---|---|---|---|---| | 1 | `push_v_with_efield` | P | `VlasovAmpereCoupling`, `PushVinForceField` | PR 14 | | 2 | `charge_density_0form` | A (vector) | initial Poisson solve of these models | PR 14 | -| 3 | `vlasov_maxwell` | A (matrix + vector) | `VlasovAmpereCoupling` | ⏸ 6D views | +| 3 | `vlasov_maxwell` | A (matrix + vector) | `VlasovAmpereCoupling` | next (6D views available) | **Step 2 – LinearVlasovAmpère/Maxwell (δf)** | # | Kernel | Type | Used by | Status | |---|---|---|---|---| | 4 | `push_weights_with_efield_lin_va` | P | `EfieldWeightsCoupling` | PR 14 | -| 5 | `linear_vlasov_ampere` | A (matrix + vector) | `EfieldWeightsCoupling` | ⏸ 6D views | +| 5 | `linear_vlasov_ampere` | A (matrix + vector) | `EfieldWeightsCoupling` | ✓ (#688) | **Step 3 – Guiding center: ToyDrift, then DriftKineticElectrostaticAdiabatic.** The default algorithm is `discrete_gradient_1st_order`; its chain comes first. @@ -348,9 +349,9 @@ The three spaces (H1vec/Hcurl/Hdiv) differ only in the basis; port one, then the 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, -and array views with more than 4 dimensions in cunumpy (all matrix accumulations write 6D stencil -matrix data). +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. ## Testing @@ -626,3 +627,46 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a The GPU tests have not run on an H100; the CPU emulation and the CPU regression tests are the gate. - **Parity cases.** The spline mappings are added to the geometry kernels' entries of `pic/tests/cuda_parity_cases.py` (the per-folder `_test_args.py` modules were replaced by that module in PR 16). + +## `linear_vlasov_ampere` implementation notes + +The first matrix accumulation on the GPU (part of #688, tracked in #650): `linear_vlasov_ampere`, the accumulation of +`EfieldWeightsCoupling` in `LinearVlasovAmpereOneSpecies` (step 2 of the [porting order](#porting-order)). + +- **6D views.** The six blocks `mat11, mat12, mat13, mat22, mat23, mat33` are `Array6D` (cunumpy 0.6.1, strided, + like the `Array3D` vectors); `cunumpy.kernels.CudaKernel` packs the `StencilMatrix._data` blocks without + copying them. `f0_values` is a `const double*`, as in `push_weights_with_efield_lin_va`. +- **Device helpers, one to one.** `filler_kernels.cuh` gets `fill_mat` and `fill_mat_vec` next to `fill_vec`, + `particle_to_mat_kernels.cuh` gets `m_v_fill_b_v1_symm` next to `vec_fill_b_v0`, and `linalg_kernels.cuh` gets + `outer`; same names, arguments and loop order as pyccel (`starts` and `pads` are `const long long*`). Every addition + into matrix or vector data is a `cunumpy_atomic_add`, as in `charge_density_0form`; the spline values pyccel keeps + in `args_derham.bn1, ..., bd3` are in the thread-local `SplineScratch`. Of the remaining `m_v_fill_*`/`mat_fill_*` + fillers, each one is ported with the first kernel that calls it. +- **Kernel.** `linear_vlasov_ampere_cuda.cu`: same arguments in the same order as pyccel, one thread per marker row, + holes and boundary particles skipped; `df`, `matrix_inv`, `matrix_vector`, `outer`, then `m_v_fill_b_v1_symm`. The + 3 x 3 `filling_m` is nine row-major entries; the CUDA-only `factor` is pyccel's `f0_values[ip] / (markers[ip, 7] * + n_markers_tot)`, computed once as pyccel does. +- **`Accumulator` on CuPy.** Nothing in `Accumulator` had to change: on the CuPy backend `StencilMatrix._data` and + `StencilVector._data` are device arrays (feectools allocates them with `xp`), the accumulator zeroes them in place and + passes the same arrays to the kernel; `exchange_assembly_data`, `update_ghost_regions`, the transposed blocks of the + symmetric matrix (`StencilMatrix.transpose`, a CUDA kernel in feectools), the scaled operator of `SchurSolver` and its + `dot` all work on them. `test_accum_matrix_cupy.py` checks this against NumPy: the kernel data are the blocks' data + (no copies), the nine arrays, the dense blocks, the accumulated vector, `BC.dot(x)` and three CG iterations of the + Schur solve agree, and nothing is copied to the host from the accumulation to `BC.dot(x)`. On a GPU it runs with + `assert_no_transfers` (not run yet); without one it runs in a subprocess on cunumpy's fake CuPy, with every + `CudaKernel` launch emulated on the CPU in place on the fake device arrays (`emulated_launches()` in + `pic/tests/cuda_emulation.py`; struphy's and feectools' kernels, struct arguments included). This would fit into + cunumpy (fake CuPy with emulated launches). +- **Found, not fixed here.** feectools' `StencilVector.inner` returns its result on the host (`xp.to_numpy`), so each + inner product of the CG solve copies a scalar from the device; the Schur solve is therefore outside the transfer + guard. `MassMatrixPreconditioner`, the default preconditioner of `EfieldWeightsCoupling`, cannot be created on CuPy + (`is_circulant` asserts a host array; the device Kronecker solve is a separate PR). `LinearVlasovAmpereOneSpecies` + end to end on the GPU (#689) also needs the feectools stencil 6D views and device Kronecker solve PRs and the first + H100 run (#687). +- **Tests.** Parity cases `PARITY_CASES["linear_vlasov_ampere"]`: 129 markers (a hole and a boundary particle) in + Cuboid, Colella, HollowTorus and a 3d spline mapping, data shaped like the real accumulation matrices of clamped + splines of degrees 2, 3, 1 on 8 cells (`v1_symm_accumulation_data()` in `pic/tests/kernel_test_args.py`); + `rtol = 1e-12`, `atol = 1e-10` (atomics change the summation order; the entries reach 1e4). The CPU emulation agrees + with pyccel to round-off and catches a transposed `DF`. `m_v_fill_b_v1_symm` is compared with pyccel through a + wrapper kernel (`test_v1_symm_filler` on a GPU, `test_emulated_v1_symm_filler` by emulation). GPU tests have not run + on an H100. diff --git a/src/struphy/linear_algebra/linalg_kernels.cuh b/src/struphy/linear_algebra/linalg_kernels.cuh index 1b72c25fe..4636f2ea3 100644 --- a/src/struphy/linear_algebra/linalg_kernels.cuh +++ b/src/struphy/linear_algebra/linalg_kernels.cuh @@ -100,4 +100,15 @@ __device__ inline void transpose(const double* a, double* b) { for (int i = 0; i < 3; ++i) for (int j = 0; j < 3; ++j) b[3 * i + j] = a[3 * j + i]; } +/** + * Outer product of two vectors of length 3, as in linalg_kernels.outer. + * + * @param a First input vector with three entries. + * @param b Second input vector with three entries. + * @param c Output matrix c_ij = a_i * b_j with nine row-major entries; must not overlap a or b. + */ +__device__ inline void outer(const double* a, const double* b, double* c) { + for (int i = 0; i < 3; ++i) + for (int j = 0; j < 3; ++j) c[3 * i + j] = a[i] * b[j]; +} } diff --git a/src/struphy/pic/accumulation/filler_kernels.cuh b/src/struphy/pic/accumulation/filler_kernels.cuh index 37f3867d9..8a33d7773 100644 --- a/src/struphy/pic/accumulation/filler_kernels.cuh +++ b/src/struphy/pic/accumulation/filler_kernels.cuh @@ -1,5 +1,5 @@ -// Device versions of the fillers in filler_kernels.py: add the contribution of one marker to a stencil vector. -// Many markers write to the same entries, so every addition is atomic. +// Device versions of the fillers in filler_kernels.py: add the contribution of one marker to a stencil vector +// or to a block of a stencil matrix. Many markers write to the same entries, so every addition is atomic. #pragma once #include "cunumpy/array_view.cuh" #include "cunumpy/atomic.cuh" @@ -40,4 +40,125 @@ __device__ inline void fill_vec(int pi1, int pi2, int pi3, const double* bi1, co } } } + +/** + * Add basis functions times filling to a block of a stencil matrix, as in filler_kernels.fill_mat. + * + * @param pi1 Spline degree of the codomain (row indices) along the first axis. + * @param pi2 Spline degree of the codomain along the second axis. + * @param pi3 Spline degree of the codomain along the third axis. + * @param pj1 Spline degree of the domain (column indices) along the first axis. + * @param pj2 Spline degree of the domain along the second axis. + * @param pj3 Spline degree of the domain along the third axis. + * @param bi1 The pi1 + 1 non-vanishing N- or D-spline values of the codomain along the first axis. + * @param bi2 The pi2 + 1 non-vanishing spline values of the codomain along the second axis. + * @param bi3 The pi3 + 1 non-vanishing spline values of the codomain along the third axis. + * @param bj1 The pj1 + 1 non-vanishing spline values of the domain along the first axis. + * @param bj2 The pj2 + 1 non-vanishing spline values of the domain along the second axis. + * @param bj3 The pj3 + 1 non-vanishing spline values of the domain along the third axis. + * @param span1 Knot span index along the first axis. + * @param span2 Knot span index along the second axis. + * @param span3 Knot span index along the third axis. + * @param starts Start indices of the codomain (three entries). + * @param pads Paddings of the codomain (three entries). + * @param mat Stencil matrix data (the _data of a StencilMatrix, shape (n1, n2, n3, 2 * pads + 1)), any strides; + * written atomically. + * @param filling Number multiplied by the basis functions. + * + * The summation order of concurrent threads (and the last bits of the result) depends on the thread schedule. + */ +__device__ inline void fill_mat(int pi1, int pi2, int pi3, int pj1, int pj2, int pj3, const double* bi1, + const double* bi2, const double* bi3, const double* bj1, const double* bj2, + const double* bj3, int span1, int span2, int span3, const long long* starts, + const long long* pads, Array6D mat, double filling) { + for (int il1 = 0; il1 <= pi1; ++il1) { + long long i1 = span1 + il1 - starts[0]; + double b1 = bi1[il1] * filling; + for (int il2 = 0; il2 <= pi2; ++il2) { + long long i2 = span2 + il2 - starts[1]; + double b2 = b1 * bi2[il2]; + for (int il3 = 0; il3 <= pi3; ++il3) { + long long i3 = span3 + il3 - starts[2]; + double b3 = b2 * bi3[il3]; + + for (int jl1 = 0; jl1 <= pj1; ++jl1) { + long long j1 = pads[0] + jl1 - il1; + double b4 = b3 * bj1[jl1]; + for (int jl2 = 0; jl2 <= pj2; ++jl2) { + long long j2 = pads[1] + jl2 - il2; + double b5 = b4 * bj2[jl2]; + for (int jl3 = 0; jl3 <= pj3; ++jl3) { + long long j3 = pads[2] + jl3 - il3; + double b6 = b5 * bj3[jl3]; + + cunumpy_atomic_add(&mat(i1, i2, i3, j1, j2, j3), b6); + } + } + } + } + } + } +} + +/** + * Add basis functions times fillings to a block of a stencil matrix and to a stencil vector, as in + * filler_kernels.fill_mat_vec. + * + * @param pi1 Spline degree of the codomain (row indices) along the first axis. + * @param pi2 Spline degree of the codomain along the second axis. + * @param pi3 Spline degree of the codomain along the third axis. + * @param pj1 Spline degree of the domain (column indices) along the first axis. + * @param pj2 Spline degree of the domain along the second axis. + * @param pj3 Spline degree of the domain along the third axis. + * @param bi1 The pi1 + 1 non-vanishing N- or D-spline values of the codomain along the first axis. + * @param bi2 The pi2 + 1 non-vanishing spline values of the codomain along the second axis. + * @param bi3 The pi3 + 1 non-vanishing spline values of the codomain along the third axis. + * @param bj1 The pj1 + 1 non-vanishing spline values of the domain along the first axis. + * @param bj2 The pj2 + 1 non-vanishing spline values of the domain along the second axis. + * @param bj3 The pj3 + 1 non-vanishing spline values of the domain along the third axis. + * @param span1 Knot span index along the first axis. + * @param span2 Knot span index along the second axis. + * @param span3 Knot span index along the third axis. + * @param starts Start indices of the codomain (three entries). + * @param pads Paddings of the codomain (three entries). + * @param mat Stencil matrix data (the _data of a StencilMatrix), any strides; written atomically. + * @param filling_mat Number multiplied by the basis functions and added to mat. + * @param vec Stencil vector data of the codomain (the _data of a StencilVector), any strides; written atomically. + * @param filling_vec Number multiplied by the basis functions of the codomain and added to vec. + */ +__device__ inline void fill_mat_vec(int pi1, int pi2, int pi3, int pj1, int pj2, int pj3, const double* bi1, + const double* bi2, const double* bi3, const double* bj1, const double* bj2, + const double* bj3, int span1, int span2, int span3, const long long* starts, + const long long* pads, Array6D mat, double filling_mat, + Array3D vec, double filling_vec) { + for (int il1 = 0; il1 <= pi1; ++il1) { + long long i1 = span1 + il1 - starts[0]; + double b1 = bi1[il1]; + for (int il2 = 0; il2 <= pi2; ++il2) { + long long i2 = span2 + il2 - starts[1]; + double b2 = b1 * bi2[il2]; + for (int il3 = 0; il3 <= pi3; ++il3) { + long long i3 = span3 + il3 - starts[2]; + double b3 = b2 * bi3[il3]; + + cunumpy_atomic_add(&vec(i1, i2, i3), b3 * filling_vec); + + for (int jl1 = 0; jl1 <= pj1; ++jl1) { + long long j1 = pads[0] + jl1 - il1; + double b4 = b3 * bj1[jl1] * filling_mat; + for (int jl2 = 0; jl2 <= pj2; ++jl2) { + long long j2 = pads[1] + jl2 - il2; + double b5 = b4 * bj2[jl2]; + for (int jl3 = 0; jl3 <= pj3; ++jl3) { + long long j3 = pads[2] + jl3 - il3; + double b6 = b5 * bj3[jl3]; + + cunumpy_atomic_add(&mat(i1, i2, i3, j1, j2, j3), b6); + } + } + } + } + } + } +} } diff --git a/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu b/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu new file mode 100644 index 000000000..387013af2 --- /dev/null +++ b/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu @@ -0,0 +1,90 @@ +#include "struphy/geometry/evaluation_kernels.cuh" +#include "struphy/linear_algebra/linalg_kernels.cuh" +#include "struphy/pic/accumulation/particle_to_mat_kernels.cuh" + +/** + * Accumulate the matrix and vector of the linear Vlasov-Ampere system (delta-f) into V1, as in + * linear_vlasov_ampere_kernels.linear_vlasov_ampere. + * + * Filling functions (the factors alpha^2 kappa^2 / v_th^2 and alpha^2 kappa are applied by the propagator): + * A_p^{mu, nu} = f_0(eta_p, v_p) / (N s_0) [DF^{-1}(eta_p) v_p]_mu [DF^{-1}(eta_p) v_p]_nu, + * B_p^mu = w_p [DF^{-1}(eta_p) v_p]_mu; + * one thread per marker row. + * + * @param args_markers Marker buffer (n_markers x n_cols, row-major) and the total marker number Np; the weight is in + * column 6 and s_0 in column 7, as in pyccel. + * @param args_derham Spline degrees (1 to 8), knots and start indices. + * @param args_domain Mapping arguments; every spline and analytic mapping. + * @param mat11 (1,1)-block data of the symmetric V1 -> V1 block matrix (the _data of a StencilMatrix, 6D, any + * strides); written with atomic additions. + * @param mat12 (1,2)-block data; written atomically. + * @param mat13 (1,3)-block data; written atomically. + * @param mat22 (2,2)-block data; written atomically. + * @param mat23 (2,3)-block data; written atomically. + * @param mat33 (3,3)-block data; written atomically. + * @param vec1 First component of the V1 vector data; written atomically. + * @param vec2 Second component of the V1 vector data; written atomically. + * @param vec3 Third component of the V1 vector data; written atomically. + * @param f0_values Value of f0 for each marker row (n_markers entries). + * + * Holes (markers[ip, 0] == -1) and boundary particles (markers[ip, -1] == -2) are skipped. The pyccel loop runs over + * all rows (shape(markers)[0]), which is n_markers here. Many threads add to the same entries, so the summation + * order (and the last bits of the result) differs from the serial pyccel loop. + */ +extern "C" __global__ void linear_vlasov_ampere(MarkerArgs args_markers, DerhamArgs args_derham, DomainArgs args_domain, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + Array3D vec1, Array3D vec2, Array3D vec3, + const double* f0_values) { + // CUDA-only: the marker row of this thread replaces the pyccel loop variable + int ip = blockDim.x * blockIdx.x + threadIdx.x; + + // get number of markers + int n_markers = args_markers.n_markers; + if (ip >= n_markers) return; + + // total number of markers (weights are w_p = delta f_p / (N * s_0)) + int n_markers_tot = args_markers.Np; + + // only do something if particle is a "true" particle (i.e. not a hole) + if (args_markers.markers(ip, 0) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) + return; + + // allocate for metric coeffs (3 x 3 matrices as nine row-major entries) + double dfm[9], df_inv[9]; + + // allocate for filling + double v[3], df_inv_v[3], filling_m[9], filling_v[3]; + + // marker positions + double eta1 = args_markers.markers(ip, 0); + double eta2 = args_markers.markers(ip, 1); + double eta3 = args_markers.markers(ip, 2); + + // get velocity + v[0] = args_markers.markers(ip, 3); + v[1] = args_markers.markers(ip, 4); + v[2] = args_markers.markers(ip, 5); + + // evaluate Jacobian, result in dfm + struphy_cuda::df(eta1, eta2, eta3, args_domain, dfm); + + // invert Jacobian matrix + struphy_cuda::matrix_inv(dfm, df_inv); + + // compute DF^{-1} v + struphy_cuda::matrix_vector(df_inv, v, df_inv_v); + + // filling_m = alpha^2 * kappa^2 * f0 / (N * s_0 * v_th^2) * (DF^{-1} v_p)_mu * (DF^{-1} v_p)_nu + struphy_cuda::outer(df_inv_v, df_inv_v, filling_m); + double factor = f0_values[ip] / (args_markers.markers(ip, 7) * n_markers_tot); + for (int k = 0; k < 9; ++k) filling_m[k] *= factor; + + // filling_v = alpha^2 * kappa * w_p * DL^{-1} * v_p + for (int k = 0; k < 3; ++k) filling_v[k] = args_markers.markers(ip, 6) * df_inv_v[k]; + + // call the appropriate matvec filler + struphy_cuda::m_v_fill_b_v1_symm(args_derham, eta1, eta2, eta3, mat11, mat12, mat13, mat22, mat23, mat33, + filling_m[0], filling_m[1], filling_m[2], filling_m[4], filling_m[5], + filling_m[8], vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2]); +} diff --git a/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh b/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh index 7cabaee45..ade422b97 100644 --- a/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh +++ b/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh @@ -28,4 +28,73 @@ __device__ inline void vec_fill_b_v0(const DerhamArgs& args_derham, double eta1, fill_vec(pn1, pn2, pn3, scratch.bn1, scratch.bn2, scratch.bn3, span1, span2, span3, args_derham.starts, vec, fill); } + +/** + * Add one marker's contribution to the symmetric blocks (1,1), (1,2), (1,3), (2,2), (2,3), (3,3) of a V1 -> V1 + * block matrix and to a V1 vector, as in particle_to_mat_kernels.m_v_fill_b_v1_symm. + * + * @param args_derham Spline degrees (1 to 8), knots and start indices; pn is also the padding of the blocks. + * @param eta1 Logical position along the first axis. + * @param eta2 Logical position along the second axis. + * @param eta3 Logical position along the third axis. + * @param mat11 (mu=1, nu=1)-block data (D N N x D N N), shape (n1, n2, n3, 2 * pn + 1); written atomically. + * @param mat12 (mu=1, nu=2)-block data (D N N x N D N); written atomically. + * @param mat13 (mu=1, nu=3)-block data (D N N x N N D); written atomically. + * @param mat22 (mu=2, nu=2)-block data (N D N x N D N); written atomically. + * @param mat23 (mu=2, nu=3)-block data (N D N x N N D); written atomically. + * @param mat33 (mu=3, nu=3)-block data (N N D x N N D); written atomically. + * @param fill11 Number multiplied by the basis functions and added to mat11. + * @param fill12 Number multiplied by the basis functions and added to mat12. + * @param fill13 Number multiplied by the basis functions and added to mat13. + * @param fill22 Number multiplied by the basis functions and added to mat22. + * @param fill23 Number multiplied by the basis functions and added to mat23. + * @param fill33 Number multiplied by the basis functions and added to mat33. + * @param vec1 mu=1 component of the V1 vector data (D N N); written atomically. + * @param vec2 mu=2 component of the V1 vector data (N D N); written atomically. + * @param vec3 mu=3 component of the V1 vector data (N N D); written atomically. + * @param fill1 Number multiplied by the basis functions and added to vec1. + * @param fill2 Number multiplied by the basis functions and added to vec2. + * @param fill3 Number multiplied by the basis functions and added to vec3. + * + * Pyccel keeps the spline values in args_derham.bn1, ..., bd3; here they are in a thread-local SplineScratch. + */ +__device__ inline void m_v_fill_b_v1_symm(const DerhamArgs& args_derham, double eta1, double eta2, double eta3, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + double fill11, double fill12, double fill13, double fill22, + double fill23, double fill33, Array3D vec1, Array3D vec2, + Array3D vec3, double fill1, double fill2, double fill3) { + // degrees of the basis functions : B-splines (pn) and D-splines (pd) + int pn1 = args_derham.pn[0]; + int pn2 = args_derham.pn[1]; + int pn3 = args_derham.pn[2]; + + int pd1 = pn1 - 1; + int pd2 = pn2 - 1; + int pd3 = pn3 - 1; + + // CUDA-only scratch holds the spans and the spline values pyccel keeps in args_derham + SplineScratch scratch; + get_spans(eta1, eta2, eta3, args_derham, scratch); + int span1 = scratch.span1, span2 = scratch.span2, span3 = scratch.span3; + + // fill matrix entries + fill_mat_vec(pd1, pn2, pn3, pd1, pn2, pn3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bd1, scratch.bn2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat11, fill11, vec1, fill1); + + fill_mat_vec(pn1, pd2, pn3, pn1, pd2, pn3, scratch.bn1, scratch.bd2, scratch.bn3, scratch.bn1, scratch.bd2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat22, fill22, vec2, fill2); + + fill_mat_vec(pn1, pn2, pd3, pn1, pn2, pd3, scratch.bn1, scratch.bn2, scratch.bd3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat33, fill33, vec3, fill3); + + fill_mat(pd1, pn2, pn3, pn1, pd2, pn3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bn1, scratch.bd2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat12, fill12); + + fill_mat(pd1, pn2, pn3, pn1, pn2, pd3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat13, fill13); + + fill_mat(pn1, pd2, pn3, pn1, pn2, pd3, scratch.bn1, scratch.bd2, scratch.bn3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat23, fill23); +} } diff --git a/src/struphy/pic/tests/cuda_emulation.py b/src/struphy/pic/tests/cuda_emulation.py index 0e8db4f12..ca3c76d01 100644 --- a/src/struphy/pic/tests/cuda_emulation.py +++ b/src/struphy/pic/tests/cuda_emulation.py @@ -6,8 +6,10 @@ the fields of each struct as separate parameters, rebuilds the structs and calls the real kernel. """ +import contextlib + import numpy as np -from cunumpy.kernel_testing import emulate_cuda_kernel +from cunumpy.kernel_testing import emulate_cuda_kernel, fake_cupy_active from cunumpy.kernels import CudaKernel from struphy.utils.cuda_arguments import CUDA_OPTIONS @@ -55,11 +57,12 @@ def _struct_field(value, name): return getattr(value, name) -def emulate_struct_kernel(kernel: CudaKernel, *args, n_threads=None): +def emulate_struct_kernel(kernel: CudaKernel, *args, n_threads=None, grid=None, block=None): """Emulate `kernel` with host arguments; struct arguments are the pyccel argument objects (NumPy arrays). - Arrays (also those inside the argument objects) are updated in place, as by a launch. Without `n_threads`, - one thread per row of the first array is launched, as in a real launch. + Arrays (also those inside the argument objects) are updated in place, as by a launch. Without `n_threads` or + `grid`, one thread per row of the first array is launched, as in a real launch; `grid` and `block` give the + launch shape explicitly. """ flat = [] for p, value in zip(kernel.signature, args): @@ -67,9 +70,58 @@ def emulate_struct_kernel(kernel: CudaKernel, *args, n_threads=None): flat.append(value) else: flat += [_struct_field(value, f.name) for f in p.struct.fields] - if n_threads is None: - n_threads = next(value.shape[0] for value in flat if isinstance(value, np.ndarray)) - grid, block = kernel.launch_shape(n_threads) + if grid is None: + if n_threads is None: + n_threads = next(value.shape[0] for value in flat if isinstance(value, np.ndarray)) + grid, block = kernel.launch_shape(n_threads, block=block) source, name = _wrapper(kernel) wrapper = CudaKernel(source, name, source_dir=kernel.source_dir, **CUDA_OPTIONS) emulate_cuda_kernel(wrapper, *flat, grid=grid, block=block, options=EMULATION_OPTIONS) + + +def _host(value): + """The host buffer behind a fake CuPy array (the same memory, no copy); other values unchanged.""" + if isinstance(value, np.ndarray) or not hasattr(value, "__cuda_array_interface__"): + return value + return object.__getattribute__(value, "_a") + + +class _HostFields: + """A CUDA argument object (``Cuda*Arguments``) whose array attributes are their fake CuPy host buffers.""" + + def __init__(self, args): + self._args = args + + def __getattr__(self, name): + return _host(getattr(self._args, name)) + + +@contextlib.contextmanager +def emulated_launches(): + """On cunumpy's fake CuPy, run every ``CudaKernel`` launch by CPU emulation, in place on the fake device arrays. + + The fake CuPy (``CUNUMPY_FAKE_CUPY=1``) keeps device arrays in host memory and rejects host/device mixing, but + cannot launch kernels. Inside this context a launch emulates the kernel (struct arguments included, through + :func:`emulate_struct_kernel`) on the host buffers of the fake arrays, so code that launches CUDA kernels (struphy's + and feectools') runs end to end on the CuPy backend without a GPU. Each launch compiles the kernel, so this is slow. + """ + if not fake_cupy_active(): + raise RuntimeError("emulated_launches() needs cunumpy's fake CuPy (CUNUMPY_FAKE_CUPY=1)") + original = CudaKernel.__call__ + + def launch(self, *args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None): + grid, block = self.launch_shape(n_threads, grid=grid, block=block, args=args) + if any(p.struct is not None for p in self.signature): + host_args = [_host(a) if p.struct is None else _HostFields(a) for p, a in zip(self.signature, args)] + emulate_struct_kernel(self, *host_args, grid=grid, block=block) + else: + host_args = [_host(a) for a in args] + emulate_cuda_kernel( + self, *host_args, grid=grid, block=block, shared_mem=shared_mem, options=EMULATION_OPTIONS + ) + + CudaKernel.__call__ = launch + try: + yield + finally: + CudaKernel.__call__ = original diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index e2ce0a5d5..85380b6e2 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -29,6 +29,7 @@ marker_arguments, spline_coefficients, spline_evaluation_arguments, + v1_symm_accumulation_data, ) @@ -79,6 +80,19 @@ def charge_density_0form_args(bc): return (args_markers, derham_arguments(), args_domain, xp.zeros((18, 20, 16))) +def linear_vlasov_ampere_args(domain): + """129 markers (a hole and a boundary particle) in domain `domain` of geometry_domain(); zeroed V1 data.""" + args_markers, _ = marker_arguments((0, 0, 0)) + f0_values = xp.asarray(np.random.default_rng(14).random(N_MARKERS)) + return ( + args_markers, + derham_arguments(), + geometry_domain(domain).args_domain, + *v1_symm_accumulation_data(), + f0_values, + ) + + # ---------------------------------------------------------------- spline evaluation @@ -267,6 +281,9 @@ def size_of(index, per_thread=1): "reflect": ParityCases((0, 1, 2), reflect_args, n_threads=size_of(2)), # atomics add in another order than the serial loop "charge_density_0form": ParityCases(((0, 0, 0), (2, 0, 1)), charge_density_0form_args, rtol=1e-12, atol=1e-13), + # Cuboid, Colella and HollowTorus (non-diagonal DF) and a 3d spline mapping, by index into geometry_domain(); + # atomics add in another order than the serial loop, and the entries reach 1e4 (atol for cancelling sums) + "linear_vlasov_ampere": ParityCases((0, 2, 5, 12), linear_vlasov_ampere_args, rtol=1e-12, atol=1e-10), "eval_spline_mpi_markers": ParityCases(SPLINE_KINDS, eval_spline_mpi_markers_args), "eval_spline_mpi_matrix": ParityCases(SPLINE_KINDS, eval_spline_mpi_grid_args(sparse=False), n_threads=size_of(-1)), "eval_spline_mpi_sparse_meshgrid": ParityCases( diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index d413b5380..eaed48adf 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -49,14 +49,37 @@ def butcher_arguments(method): return xp.asarray(butcher.a_stage), xp.asarray(butcher.b), xp.asarray(butcher.c), butcher.n_stages +DERHAM_DEGREE = (2, 3, 1) +DERHAM_CELLS = 8 + + def derham_arguments(): """Splines of degrees 2, 3, 1 on 8 cells, starting at index 0.""" - degree = np.array([2, 3, 1], dtype=np.int64) - knots = [np.r_[np.zeros(p), np.linspace(0, 1, 9), np.ones(p)] for p in degree] + degree = np.array(DERHAM_DEGREE, dtype=np.int64) + knots = [np.r_[np.zeros(p), np.linspace(0, 1, DERHAM_CELLS + 1), np.ones(p)] for p in degree] args_class = CudaDerhamArguments if xp.get_backend() == "cupy" else DerhamArguments return args_class(xp.asarray(degree), *(xp.asarray(t) for t in knots), xp.zeros(3, dtype=np.int64)) +def v1_symm_accumulation_data(): + """Zeroed data of a symmetric V1 -> V1 accumulation matrix and a V1 vector, for :func:`derham_arguments`. + + Returns the six independent blocks ``mat11, mat12, mat13, mat22, mat23, mat33`` and the three vector components + ``vec1, vec2, vec3`` on the active backend, shaped like the ``_data`` of the ``StencilMatrix`` blocks and + ``StencilVector`` components that ``Accumulator(..., "Hcurl", ..., symmetry="symm", add_vector=True)`` passes on + one process for clamped splines of degrees 2, 3, 1 on 8 cells: component ``mu`` has D-splines (``n - 1`` + functions) along axis ``mu`` and N-splines (``n = 8 + p``) along the others, every axis is padded by ``p`` on both + sides, and the matrix blocks have ``2 p + 1`` diagonals per axis (the paddings are the degrees of V0). + """ + degree = np.array(DERHAM_DEGREE) + n_spline = DERHAM_CELLS + degree + rows = [tuple(n_spline - (np.arange(3) == mu) + 2 * degree) for mu in range(3)] + diagonals = tuple(2 * degree + 1) + mats = tuple(xp.zeros(rows[mu] + diagonals) for mu, nu in ((0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2))) + vecs = tuple(xp.zeros(rows[mu]) for mu in range(3)) + return mats + vecs + + def spline_coefficients(n=3, seed=11): """Random coefficients, large enough for every span of :func:`derham_arguments`.""" rng = np.random.default_rng(seed) diff --git a/src/struphy/pic/tests/test_accum_matrix_cupy.py b/src/struphy/pic/tests/test_accum_matrix_cupy.py new file mode 100644 index 000000000..cf0aa62c6 --- /dev/null +++ b/src/struphy/pic/tests/test_accum_matrix_cupy.py @@ -0,0 +1,209 @@ +"""The matrix path of :class:`~struphy.pic.accumulation.particles_to_grid.Accumulator` on the CuPy backend. + +``linear_vlasov_ampere`` accumulates the symmetric V1 -> V1 block matrix and the V1 vector of ``EfieldWeightsCoupling``; +the matrix data (``StencilMatrix._data``) must be device arrays that are zeroed, passed to the CUDA kernel and turned +into the Schur operator without host copies. Checked against the NumPy backend on a GPU and, without one, on +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 + +import cunumpy as xp +import numpy as np +import pytest +from cunumpy.kernel_testing import emulation_compiler, requires_cupy + +from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env + +ALPHA, KAPPA, VTH, DT = 1.3, 0.7, 0.9, 0.05 +N_SCHUR_ITERATIONS = 3 + + +@dataclass +class MarkerOwner: + """What the accumulators read from ``Particles``: the marker arguments, no clones. + + The same markers on both backends (``Particles`` draws them with the backend's random generator). + """ + + args_markers: Any + clone_config: Any = None + + +def accumulate_linear_vlasov_ampere(guard=None): + """Accumulate ``linear_vlasov_ampere`` like ``EfieldWeightsCoupling`` on the active backend, then solve with it. + + Colella mapping (non-diagonal Jacobian), clamped splines in eta1 and periodic ones in eta2, eta3, so the three + V1 components have different sizes. `guard` (a context manager factory, e.g. one that fails on host transfers) + encloses the accumulation, the scaled operator and its application, as in a time step of the propagator; the + Schur solve is not guarded (see below). + + Returns a dict of host arrays: the nine kernel data arrays (``mat11`` ... ``vec3``) after the accumulation (ghost + regions exchanged), the assembled block matrix as dense blocks, the accumulated vector after the boundary + operators, ``BC.dot(x)`` with the operator scaled as in ``EfieldWeightsCoupling`` and the Schur solution after + ``N_SCHUR_ITERATIONS`` iterations of the conjugate gradient method. + """ + import contextlib + + from feectools.linalg.stencil import StencilMatrix + + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.psydac_derham import Derham + from struphy.geometry.domains import Colella + from struphy.io.options import DerhamOptions + from struphy.linear_algebra.schur_solver import SchurSolver + from struphy.linear_algebra.solver import SolverParameters + from struphy.pic.accumulation.kernels.linear_vlasov_ampere import linear_vlasov_ampere + from struphy.pic.accumulation.particles_to_grid import Accumulator + from struphy.pic.tests.kernel_test_args import N_MARKERS, marker_arguments + from struphy.topology.grids import TensorProductGrid + + guard = guard or contextlib.nullcontext + derham = Derham( + TensorProductGrid(num_elements=(6, 5, 4)), + DerhamOptions(degree=(2, 3, 1), bcs=(("free", "free"), None, None)), + comm=None, + ) + domain = Colella(Lx=2.0, Ly=3.0, alpha=0.05, Lz=4.0) + mass_ops = WeightedMassOperators(derham, domain) + args_markers, _ = marker_arguments((0, 0, 0)) + accum = Accumulator( + MarkerOwner(args_markers), + "Hcurl", + linear_vlasov_ampere, + mass_ops, + domain.args_domain, + add_vector=True, + symmetry="symm", + ) + matrix = accum.operators[0].matrix + + # the kernel writes into the data of the operator's blocks (and vector), not into copies + blocks = [matrix[a, b] for a in range(3) for b in range(a, 3)] + assert all(isinstance(block, StencilMatrix) for block in blocks) + assert all(data is block._data for data, block in zip(accum._args_data[:6], blocks)) + assert all(data is vec._data for data, vec in zip(accum._args_data[6:], accum._vectors[0].blocks)) + assert all(xp.is_gpu(data) == (xp.get_backend() == "cupy") for data in accum._args_data) + + rng = np.random.default_rng(16) + f0_values = xp.asarray(rng.random(N_MARKERS)) + for data in accum._args_data: + data[...] = 1.0 # left over from a previous step: the accumulator zeroes it + + x = accum._vectors[0].space.zeros() + for block in x.blocks: + block._data[...] = xp.asarray(rng.normal(size=block._data.shape)) + x.update_ghost_regions() + + schur_solver = SchurSolver( + mass_ops.M1, + ALPHA**2 * KAPPA**2 * accum.operators[0] / (4 * VTH**2), + "pcg", + # a few iterations (each one launches nine stencil dot kernels, which the emulation compiles one by one) + solver_params=SolverParameters(tol=1e-30, maxiter=N_SCHUR_ITERATIONS), + ) + + with guard(): + accum(f0_values) + # as in EfieldWeightsCoupling.__call__ + schur_solver.BC = accum.operators[0] + schur_solver.BC *= (-1) * ALPHA**2 * KAPPA**2 / (4 * VTH**2) + bc_x = schur_solver.BC.dot(x) + byn = accum.vectors[0].copy() + byn *= ALPHA**2 * KAPPA / 2.0 + # not guarded: feectools' StencilVector.inner returns its result on the host (xp.to_numpy), so every inner + # product of the conjugate gradient method copies one scalar from the device + solution, info = schur_solver(xn=x, Byn=byn, dt=DT) + + out = {name: xp.to_numpy(data) for name, data in zip(MATRIX_DATA_NAMES, accum._args_data)} + for a in range(3): + for b in range(3): + out[f"block{a}{b}"] = xp.to_numpy(matrix[a, b].toarray()) + out["vector"] = xp.to_numpy(accum.vectors[0].toarray()) + out["bc_x"] = xp.to_numpy(bc_x.toarray()) + out["solution"] = xp.to_numpy(solution.toarray()) + out["niter"] = np.array(info["niter"]) + return out + + +MATRIX_DATA_NAMES = ("mat11", "mat12", "mat13", "mat22", "mat23", "mat33", "vec1", "vec2", "vec3") + + +def compare_with_numpy(result): + """Compare a CuPy-backend `result` of :func:`accumulate_linear_vlasov_ampere` with the NumPy backend's.""" + with xp.use_backend("numpy"): + expected = accumulate_linear_vlasov_ampere() + assert expected["niter"] == result["niter"] == N_SCHUR_ITERATIONS + assert np.any(expected["mat12"] != 0.0) and np.any(expected["vec3"] != 0.0) + # the transposed blocks of the symmetric matrix are filled from the accumulated ones + np.testing.assert_allclose(expected["block10"], expected["block01"].T, rtol=1e-14, atol=1e-14) + for name, value in expected.items(): + scale = max(np.max(np.abs(value)), 1.0) + # atomics add in another order than the serial loop; the solve repeats it + np.testing.assert_allclose(result[name], value, rtol=1e-10, atol=1e-12 * scale, err_msg=name) + + +@requires_cupy +def test_accumulator_matrix_on_cupy(): + """On a GPU: the accumulated matrix and vector, the Schur operator and solve agree with NumPy; the accumulation and + the operator run without host/device transfers.""" + from cunumpy.profiling import assert_no_transfers + + with xp.use_backend("cupy"): + accumulate_linear_vlasov_ampere() # warm up (NVRTC compilation) + result = accumulate_linear_vlasov_ampere(guard=assert_no_transfers) + compare_with_numpy(result) + + +def fake_cupy_transfer_guard(): + """On cunumpy's fake CuPy: fail on every explicit host copy of a device array (``.get()``, ``cupy.asnumpy``). + + The fake CuPy already rejects implicit conversions (``numpy.asarray`` of a device array, mixing host and device + arrays); this closes the explicit ones, which are legal but are transfers inside a time step. + """ + import contextlib + + import cupy + + @contextlib.contextmanager + def guard(): + def fail(*args, **kwargs): + raise AssertionError("host copy of a device array during the accumulation") + + saved = cupy.ndarray.get, cupy.asnumpy + cupy.ndarray.get, cupy.asnumpy = fail, fail + try: + yield + finally: + cupy.ndarray.get, cupy.asnumpy = saved + + return guard + + +def check_accumulator_matrix_fake_cupy(): + """Run in a child process with ``CUNUMPY_FAKE_CUPY=1``: the CuPy backend with emulated kernel launches.""" + from struphy.pic.tests.cuda_emulation import emulated_launches + + with emulated_launches(), xp.use_backend("cupy"): + result = accumulate_linear_vlasov_ampere(guard=fake_cupy_transfer_guard()) + compare_with_numpy(result) + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +@pytest.mark.skipif(emulation_compiler() is None, reason="no C++ compiler") +def test_accumulator_matrix_fake_cupy(): + """Without a GPU: the matrix path of the Accumulator on the CuPy backend (fake CuPy, emulated launches). + + Device data are zeroed and passed to the CUDA kernel, the ghost regions, the transposed blocks, the scaled + operator and the Schur solve run on device arrays, and no array is copied to the host from the accumulation to + the operator's application. Runs in a + 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:] diff --git a/src/struphy/pic/tests/test_cuda_emulation.py b/src/struphy/pic/tests/test_cuda_emulation.py index 793bb9525..f2ab111a2 100644 --- a/src/struphy/pic/tests/test_cuda_emulation.py +++ b/src/struphy/pic/tests/test_cuda_emulation.py @@ -113,3 +113,15 @@ def test_emulated_transform_helpers(domain_index): out = np.zeros(len(expected)) emulate_struct_kernel(kernel, *inputs, args_domain, out, len(out), n_threads=len(out)) np.testing.assert_allclose(out, expected, rtol=1e-10, atol=1e-10) + + +@requires_compiler +def test_emulated_v1_symm_filler(): + """test_device_helpers.test_v1_symm_filler by CPU emulation: the device m_v_fill_b_v1_symm agrees with pyccel.""" + from struphy.pic.tests.test_device_helpers import v1_symm_filler_case + + kernel, (eta, fills), args_derham, data, expected = v1_symm_filler_case() + emulate_struct_kernel(kernel, eta, fills, args_derham, *data, len(eta), n_threads=len(eta)) + for result, reference in zip(data, expected): + assert np.any(reference != 0.0) + np.testing.assert_allclose(result, reference, rtol=1e-12, atol=1e-13) diff --git a/src/struphy/pic/tests/test_cuda_parity.py b/src/struphy/pic/tests/test_cuda_parity.py index 75d53245f..0b0a14c9e 100644 --- a/src/struphy/pic/tests/test_cuda_parity.py +++ b/src/struphy/pic/tests/test_cuda_parity.py @@ -10,6 +10,7 @@ from cunumpy.profiling import assert_no_transfers from struphy.pic.accumulation.kernels.charge_density_0form import charge_density_0form +from struphy.pic.accumulation.kernels.linear_vlasov_ampere import linear_vlasov_ampere from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage from struphy.pic.pushing.kernels.push_v_with_efield import push_v_with_efield from struphy.pic.pushing.kernels.push_vxb_analytic import push_vxb_analytic @@ -108,7 +109,7 @@ def device_only(value, *args, **kwargs): def test_vlasov_kernel_coverage(): - """Vlasov (PushEta, both PushVxB algorithms) and the ported Vlasov-Ampere kernels have device kernels.""" + """Vlasov (PushEta, both PushVxB algorithms) and the ported (linear) Vlasov-Ampere kernels have device kernels.""" for kernel in ( push_eta_stage, push_vxb_analytic, @@ -116,5 +117,6 @@ def test_vlasov_kernel_coverage(): push_v_with_efield, push_weights_with_efield_lin_va, charge_density_0form, + linear_vlasov_ampere, ): assert kernel.has_cuda, kernel.name diff --git a/src/struphy/pic/tests/test_device_helpers.py b/src/struphy/pic/tests/test_device_helpers.py index a2390ef4d..b771d8efc 100644 --- a/src/struphy/pic/tests/test_device_helpers.py +++ b/src/struphy/pic/tests/test_device_helpers.py @@ -391,3 +391,64 @@ def test_transform_helpers(domain_index): kernel, inputs, _, expected = transform_helper_case(domain_index) result = _run_on_device(kernel, inputs, domain_index, len(expected)) np.testing.assert_allclose(result, expected, rtol=1e-10, atol=1e-10) + + +# CUDA-only wrapper: one thread per point calls m_v_fill_b_v1_symm with the point's nine fillings +V1_SYMM_FILLER_SOURCE = r""" +#include "struphy/pic/accumulation/particle_to_mat_kernels.cuh" +extern "C" __global__ void fill_v1_symm(const double* eta, const double* fills, DerhamArgs args_derham, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + Array3D vec1, Array3D vec2, Array3D vec3, int n) { + int ip = blockDim.x * blockIdx.x + threadIdx.x; + if (ip >= n) return; + const double* fill = fills + 9 * ip; + struphy_cuda::m_v_fill_b_v1_symm(args_derham, eta[3 * ip], eta[3 * ip + 1], eta[3 * ip + 2], mat11, mat12, mat13, + mat22, mat23, mat33, fill[0], fill[1], fill[2], fill[3], fill[4], fill[5], vec1, + vec2, vec3, fill[6], fill[7], fill[8]); +} +""" + + +def v1_symm_filler_case(): + """Wrapper kernel, host inputs and pyccel reference for the device m_v_fill_b_v1_symm. + + 65 points (the corners of the unit cube among them) with random fillings, accumulated into zeroed V1 data shaped + like the real accumulation matrices (:func:`~struphy.pic.tests.kernel_test_args.v1_symm_accumulation_data`). + Used by :func:`test_v1_symm_filler` (GPU) and by its CPU emulation in ``test_cuda_emulation.py``. Returns + ``kernel, (eta, fills), args_derham, data, expected``: ``args_derham`` is the pyccel class and ``data`` the nine + zeroed host arrays in kernel order. + """ + from struphy.pic.accumulation import particle_to_mat_kernels + from struphy.pic.tests.kernel_test_args import derham_arguments, v1_symm_accumulation_data + + rng = np.random.default_rng(15) + eta = np.r_[np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0], [0.0, 1.0, 0.5]]), rng.random((62, 3))] + fills = rng.normal(size=(len(eta), 9)) + with cunumpy.use_backend("numpy"): + args_derham = derham_arguments() + data = v1_symm_accumulation_data() + expected = tuple(a.copy() for a in data) + for point, fill in zip(eta, fills): + particle_to_mat_kernels.m_v_fill_b_v1_symm( + args_derham, *point, *expected[:6], *fill[:6], *expected[6:], *fill[6:] + ) + kernel = CudaKernel(V1_SYMM_FILLER_SOURCE, "fill_v1_symm", **CUDA_OPTIONS) + return kernel, (eta, fills), args_derham, data, expected + + +@requires_cupy +def test_v1_symm_filler(): + """m_v_fill_b_v1_symm (and the fill_mat/fill_mat_vec it calls) on the device agrees with pyccel.""" + import cupy as cp + + from struphy.pic.tests.kernel_test_args import derham_arguments + + kernel, (eta, fills), _, data, expected = v1_symm_filler_case() + with cunumpy.use_backend("cupy"): + args_derham = derham_arguments() + device_data = [cp.asarray(a) for a in data] + kernel(cp.asarray(eta), cp.asarray(fills), args_derham, *device_data, len(eta), n_threads=len(eta)) + for result, reference in zip(device_data, expected): + assert np.any(reference != 0.0) + np.testing.assert_allclose(result.get(), reference, rtol=1e-12, atol=1e-13) From 7de72859bedf05720d8c5e3ce3b8dcf30e3cacf2 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 7 Oct 2026 19:24:11 +0200 Subject: [PATCH 02/50] Fuse the pusher preamble into one CUDA kernel --- src/struphy/pic/pushing/prepare_push_cuda.cu | 28 ++++++++++++++ src/struphy/pic/pushing/pusher.py | 39 +++++++++++++++----- 2 files changed, 58 insertions(+), 9 deletions(-) create mode 100644 src/struphy/pic/pushing/prepare_push_cuda.cu diff --git a/src/struphy/pic/pushing/prepare_push_cuda.cu b/src/struphy/pic/pushing/prepare_push_cuda.cu new file mode 100644 index 000000000..5c1dbd28e --- /dev/null +++ b/src/struphy/pic/pushing/prepare_push_cuda.cu @@ -0,0 +1,28 @@ +/** + * Prepare the marker buffer for a push, as the three column-slice assignments at the start of Pusher._push: + * + * markers[:, first_init_idx:first_shift_idx] = markers[:, :n_init] + * markers[:, first_shift_idx:residual_idx] = 0.0 + * markers[:, residual_idx:-2] = 0.0 + * + * One thread per entry of the contiguous column range [first_init_idx, n_cols - 2) of every row, so that + * neighbouring threads write neighbouring addresses of the row-major buffer and the whole range is written + * in a single pass (instead of one strided pass over the buffer per slice). + * + * @param markers Marker buffer (n_rows x n_cols, row-major), every row including holes. + * @param n_rows Number of rows of the buffer. + * @param n_cols Number of columns of the buffer. + * @param first_init_idx First column of the saved initial phase space coordinates. + * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_init_idx + n_init). + */ +extern "C" __global__ void prepare_push(double* markers, long long n_rows, int n_cols, int first_init_idx, + int n_init) { + long long width = n_cols - 2 - first_init_idx; + long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; + if (t >= n_rows * width) return; + + long long ip = t / width; + int col = first_init_idx + (int)(t - ip * width); + double* row = markers + ip * n_cols; + row[col] = col < first_init_idx + n_init ? row[col - first_init_idx] : 0.0; +} diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index 7f3d64961..762dd3bd9 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -1,9 +1,11 @@ "Accelerated particle pushing." import logging +from functools import cache +from pathlib import Path import cunumpy as xp -from cunumpy.kernels import Kernel, PyccelKernel +from cunumpy.kernels import CudaKernel, Kernel, PyccelKernel from line_profiler import profile from maybempi import MPI from scope_profiler import ProfileManager @@ -16,6 +18,13 @@ logger = logging.getLogger("struphy") +@cache +def _prepare_push_cuda() -> CudaKernel: + """The CUDA kernel that fuses the column-slice assignments at the start of :meth:`Pusher._push` + (compiled on first call).""" + return CudaKernel.from_file(Path(__file__).with_name("prepare_push_cuda.cu"), "prepare_push") + + class Pusher: r""" Class for solving particle ODEs @@ -208,17 +217,29 @@ def _push(self, dt: float): logger.debug(f"{residual_idx =}") logger.debug(f"{self.particles.n_cols =}") - init_slice = slice(first_pusher_idx, first_shift_idx) - shift_slice = slice(first_shift_idx, residual_idx) + if self._cuda: + # the three assignments below in one pass over the buffer (each slice is a strided pass on its own) + n_rows, n_cols = markers.shape + _prepare_push_cuda()( + markers, + n_rows, + n_cols, + first_pusher_idx, + 3 + vdim, + n_threads=n_rows * (n_cols - 2 - first_pusher_idx), + ) + else: + init_slice = slice(first_pusher_idx, first_shift_idx) + shift_slice = slice(first_shift_idx, residual_idx) - # save initial phase space coordinates - markers[:, init_slice] = markers[:, : 3 + vdim] + # save initial phase space coordinates + markers[:, init_slice] = markers[:, : 3 + vdim] - # set boundary shifts to zero - markers[:, shift_slice] = 0.0 + # set boundary shifts to zero + markers[:, shift_slice] = 0.0 - # clear buffer columns starting from residual index, dont clear ID (last column) and loc_box - markers[:, residual_idx:-2] = 0.0 + # clear buffer columns starting from residual index, dont clear ID (last column) and loc_box + markers[:, residual_idx:-2] = 0.0 rank = self.particles.mpi_rank logger.debug(f"rank {rank}: starting {self.kernel} ...") From 48c446b8230d81801706def57a9ea1422089f7b8 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 7 Oct 2026 19:46:18 +0200 Subject: [PATCH 03/50] cupy array fixes --- src/struphy/feec/psydac_derham.py | 3 ++- src/struphy/pic/sorting.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index a15ac1591..e5d12d3de 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -2302,7 +2302,8 @@ def __init__( if xp.get_backend() == "cupy" and np.any((degree < 1) | (degree > 8)): raise ValueError("CUDA spline degrees must be between 1 and 8.") pn = xp.asarray(degree) - knots = tuple(xp.ascontiguousarray(t, dtype=float) for t in derham.V0fem.knots) + # the knots are host arrays; cupy.ascontiguousarray does not take NumPy arrays, xp.asarray copies them + knots = tuple(xp.asarray(np.ascontiguousarray(t, dtype=float)) for t in derham.V0fem.knots) starts = (self.starts,) if isinstance(self._vector_stencil, StencilVector) else self.starts kinds = derham.spline_attributes[self.space_key].spline_types_pyccel self._args_eval = tuple( diff --git a/src/struphy/pic/sorting.py b/src/struphy/pic/sorting.py index 2d331d6c1..74db8c182 100644 --- a/src/struphy/pic/sorting.py +++ b/src/struphy/pic/sorting.py @@ -1,4 +1,5 @@ import logging +import math from typing import TYPE_CHECKING import cunumpy as xp @@ -226,7 +227,7 @@ def _set_boxes(self): n_particles = self._markers_shape[0] n_mkr = int(n_particles / n_box_in) + 1 n_cols = round( - n_mkr * (1 + 1 / xp.sqrt(n_mkr) + self._box_bufsize), + n_mkr * (1 + 1 / math.sqrt(n_mkr) + self._box_bufsize), ) # cartesian boxes (extra last row stores holes/outside particles) From 0352dd29a07949fa3821642cc32f5a5c6d2469da Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 20:49:16 +0200 Subject: [PATCH 04/50] 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 a15ac1591..4118101d8 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 @@ -1494,6 +1498,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 7f3d64961..d55b267e4 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -365,6 +365,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 8d1c824c53de78aa45aa62df499eeadaa6363bfc Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 20:49:16 +0200 Subject: [PATCH 05/50] Marker exchange between ranks on the device (CUDA-aware MPI) On the CuPy backend, Particles._sendrecv_markers hands the marker rows to Isend/Irecv as device buffers through cunumpy.mpi.mpi_buffer (directly after synchronize_for_mpi with CUDA-aware MPI, staged through host memory otherwise). Rows from each rank arrive in a contiguous slice of one device receive buffer, and one device scatter fills the holes. The per-rank counts are small NumPy arrays on every backend (they are host integers already), and a uniform alpha is a Python scalar, so the exchange makes no host/device copy and no implicit host read of a device array. The NumPy path is unchanged. Adds pic/tests/test_sorting_device.py (fake CuPy under mpirun, and a GPU variant) and a CI step running it on the fake CuPy. Solves #698, part of #650. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/test-PR-unit-mpi.yml | 6 + CUDA_STRATEGY.md | 20 +- src/struphy/pic/base.py | 129 ++++++++++++- src/struphy/pic/tests/test_sorting_device.py | 190 +++++++++++++++++++ 4 files changed, 338 insertions(+), 7 deletions(-) create mode 100644 src/struphy/pic/tests/test_sorting_device.py diff --git a/.github/workflows/test-PR-unit-mpi.yml b/.github/workflows/test-PR-unit-mpi.yml index e3a13dbc5..7bb53fbe2 100644 --- a/.github/workflows/test-PR-unit-mpi.yml +++ b/.github/workflows/test-PR-unit-mpi.yml @@ -76,3 +76,9 @@ jobs: source /struphy_fortran_/env_fortran_/bin/activate cd ${{ env.STRUPHY_PATH }} mpirun --oversubscribe -n ${{ matrix.n-procs }} pytest -v --with-mpi -m mpi_pic pic/tests + + - name: Run the marker exchange on the fake CuPy + run: | + source /struphy_fortran_/env_fortran_/bin/activate + cd ${{ env.STRUPHY_PATH }} + mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi pic/tests/test_sorting_device.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index daf0d25f0..2b85a5ef2 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -365,7 +365,7 @@ matrix data). - **Kernel launch configuration.** One thread per marker with a fixed block size for now. Kernels over grid points (accumulation, FEEC) need their own launch sizes; `CudaKernel` needs a `grid`/`block` override (PR 12+), or cunumpy's. - **Accumulation strategy.** Atomics vs. sort-then-reduce; see PR 12+. Decided by measurement on the first accumulation kernel. - **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. `xp.mpi_is_cuda_aware()` detects it; the exchange in `Particles.mpi_sort_markers` has to be checked for host staging buffers. +- **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 `Derham` with `polar_splines=True` raises on CuPy (`PolarExtractionBlocksC1` builds SciPy sparse matrices, and the @@ -626,3 +626,21 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a The GPU tests have not run on an H100; the CPU emulation and the CPU regression tests are the gate. - **Parity cases.** The spline mappings are added to the geometry kernels' entries of `pic/tests/cuda_parity_cases.py` (the per-folder `_test_args.py` modules were replaced by that module in PR 16). + +## Marker exchange implementation notes (#698) + +- `Particles._sendrecv_markers` on CuPy (`_sendrecv_markers_device`): the rows to send (gathered on the device per + destination) and one device receive buffer go to `Isend`/`Irecv` through `xp.mpi.mpi_buffer`, i.e. directly after + `synchronize_for_mpi` with CUDA-aware MPI, staged through (pinned) host memory otherwise (with a warning, once). + Rows from rank `i` arrive in a contiguous slice of the receive buffer; one device scatter fills the holes. + Zero-count messages are skipped (both sides know the counts). The NumPy path is unchanged. +- Whether MPI is CUDA-aware: the answer recorded by `xp.mpi.mpi_is_cuda_aware()`/`set_mpi_cuda_aware()`, else a + collective probe of `mpi_comm` at the first exchange (`mpi_sort_markers` is collective). +- Counts: `send_info`/`recv_info` are small NumPy arrays on every backend. The send counts are sizes of `nonzero` + results, which are host integers already (CuPy synchronizes to size them), so the `Alltoall` of the counts needs no + copy. Uniform `alpha` is a Python scalar (no host-to-device copy of a 3-vector). +- Tests: `pic/tests/test_sorting_device.py` sorts the same markers on NumPy and on CuPy over all ranks and compares + the marker arrays; with CUDA-aware MPI no transfer (`count_transfers`) and no implicit host read of a device array + during the exchange, staged: one `to_host` per non-empty send and one `to_device`. Without a GPU it runs on the fake + CuPy (`mpirun -n N env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi pic/tests/test_sorting_device.py`, in the PIC MPI + workflow); the GPU variant has not run (no GPU here). diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 1d08e9116..85c33293e 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2,6 +2,7 @@ import os import warnings from abc import ABCMeta, abstractmethod +from contextlib import ExitStack from typing import TYPE_CHECKING import cunumpy as xp @@ -1748,6 +1749,12 @@ def mpi_sort_markers( remove_ghost : bool Remove ghost particles before send. + + Notes + ----- + On the CuPy backend, the marker rows are exchanged as device buffers with CUDA-aware MPI + (staged through host memory otherwise), see :meth:`_sendrecv_markers_device`. The + per-rank counts are small host arrays on every backend (:meth:`_sendrecv_all_to_all`). """ if remove_ghost: self._remove_ghost_particles() @@ -4354,9 +4361,19 @@ def _sendrecv_determine_mtbs( """ # position that determines the sorting (including periodic shift of boundary conditions) if not isinstance(alpha, xp.ndarray): - alpha = xp.array(alpha, dtype=float) - assert alpha.size == 3 - assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) + # check a tuple/list on the host: on CuPy, asserting on a device comparison reads it back + alpha_host = np.array(alpha, dtype=float) + assert alpha_host.size == 3 + assert np.all(alpha_host >= 0.0) and np.all(alpha_host <= 1.0) + if alpha_host[0] == alpha_host[1] == alpha_host[2]: + # the common case (alpha = 1 or alpha = 0.5 in every direction): a Python scalar, + # which needs no host-to-device copy on CuPy and gives the same values on NumPy + alpha = float(alpha_host[0]) + else: + alpha = xp.asarray(alpha_host) + else: + assert alpha.size == 3 + assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) bi = self.first_pusher_idx self._sorting_etas[:] = ( alpha * (self.markers[:, :3] + self.markers[:, bi + 3 + self.vdim : bi + 3 + self.vdim + 3]) @@ -4456,8 +4473,11 @@ def _sendrecv_get_destinations(self, send_inds): Amount of particles sent to i-th process. """ - # One entry for each process - send_info = xp.zeros(self.mpi_size, dtype=int) + # One entry for each process. The counts are host integers on every backend: the + # size of a nonzero() result is known on the host (CuPy synchronizes to size it), so + # a small NumPy array holds them without any host/device copy, and the Alltoall of + # the counts (_sendrecv_all_to_all) passes host buffers to MPI. + send_info = np.zeros(self.mpi_size, dtype=int) # Gathered once and reused for every rank below, instead of re-gathering # self.markers[send_inds] and self._sorting_etas[send_inds] fresh on every @@ -4521,9 +4541,16 @@ def _sendrecv_all_to_all(self, send_info): ------- recv_info : array[int] Amount of marticles to be received from i-th process. + + Notes + ----- + ``send_info`` and ``recv_info`` are small NumPy arrays (one integer per rank) on every + backend: the counts size the receive buffers and the holes to fill, which Python needs + on the host anyway. On CuPy, only the marker rows themselves are device buffers + (:meth:`_sendrecv_markers`). """ - recv_info = xp.zeros(self.mpi_size, dtype=int) + recv_info = np.zeros(self.mpi_size, dtype=int) self.mpi_comm.Alltoall(send_info, recv_info) @@ -4542,6 +4569,10 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): Indices of empty rows in markers after send. """ + if xp.get_array_backend(self._markers) == "cupy": + self._sendrecv_markers_device(recv_info, hole_inds_after_send) + return + # i-th entry holds the number (not the index) of the first hole to be filled by data from process i first_hole = xp.cumsum(recv_info) - recv_info @@ -4585,6 +4616,92 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): for req in send_reqs: req.Wait() + def _mpi_cuda_aware(self) -> bool: + """Whether MPI takes device (CuPy) buffers, for :meth:`_sendrecv_markers_device`. + + Uses the answer recorded by :func:`cunumpy.mpi.mpi_is_cuda_aware` or + :func:`cunumpy.mpi.set_mpi_cuda_aware` (e.g. by feectools or at startup). If nothing + was recorded, probes :attr:`mpi_comm` once; the probe is collective, which is safe + here because :meth:`mpi_sort_markers` is called by every rank of :attr:`mpi_comm`. + """ + cuda_aware = xp.mpi.get_mpi_cuda_aware() + if cuda_aware is None: + cuda_aware = xp.mpi.mpi_is_cuda_aware(self.mpi_comm) + if not cuda_aware and not Particles._warned_mpi_staging: + Particles._warned_mpi_staging = True + logger.warning( + "MPI is not CUDA-aware: marker rows exchanged between ranks are staged through host memory.", + ) + return cuda_aware + + _warned_mpi_staging = False + + def _sendrecv_markers_device(self, recv_info, hole_inds_after_send): + """:meth:`_sendrecv_markers` on the CuPy backend: marker rows go to MPI as device buffers. + + With CUDA-aware MPI (:meth:`_mpi_cuda_aware`), the rows to send (:attr:`_send_list`, + gathered on the device by :meth:`_sendrecv_get_destinations`) and one device receive + buffer are handed to ``Isend``/``Irecv`` directly, after + :func:`cunumpy.mpi.synchronize_for_mpi` (inside :func:`cunumpy.mpi.mpi_buffer`), so MPI + reads what the kernels wrote. The rows from rank ``i`` arrive in a contiguous slice of + the receive buffer, and one device scatter puts all of them into the holes. Without + CUDA-aware MPI, :func:`~cunumpy.mpi.mpi_buffer` stages the same buffers through host + memory (counted by :func:`cunumpy.profiling.count_transfers`). + + The bookkeeping stays on the host without host/device copies: the counts are host + integers (``recv_info`` from the Alltoall, the send counts are the row counts of the + send buffers), and so is the number of holes. Ranks exchange no message when the + count is zero (both sides know the count, so they skip consistently). + + Parameters + ---------- + recv_info : array[int] + Amount of markers to be received from i-th process (NumPy array). + + hole_inds_after_send : array[int] + Indices of empty rows in markers after send (device array). + """ + recv_info = np.asarray(recv_info) + n_recv = int(recv_info.sum()) + if hole_inds_after_send.size < n_recv: + warnings.warn( + f'Strong load imbalance detected: \ +number of holes ({hole_inds_after_send.size}) on rank {self.mpi_rank} \ +is smaller than number of incoming particles ({n_recv}). \ +Increasing the value of "bufsize" in the markers parameters for the next run.', + ) + self.mpi_comm.Abort() + + cuda_aware = self._mpi_cuda_aware() + + # rows from rank i land in recvbuf[first_row[i] : first_row[i] + recv_info[i]] + first_row = np.cumsum(recv_info) - recv_info + recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=float) + + with ExitStack() as buffers: + if n_recv > 0: + recv_mpi = buffers.enter_context( + xp.mpi.mpi_buffer(recvbuf, send=False, recv=True, cuda_aware=cuda_aware), + ) + reqs = [] + for i in range(self.mpi_size): + if i == self.mpi_rank: + continue + if recv_info[i] > 0: + rows = recv_mpi[first_row[i] : first_row[i] + recv_info[i]] + reqs.append(self.mpi_comm.Irecv(rows, source=i, tag=i)) + data = self._send_list[i] + if data.shape[0] > 0: + send_mpi = buffers.enter_context(xp.mpi.mpi_buffer(data, cuda_aware=cuda_aware)) + reqs.append(self.mpi_comm.Isend(send_mpi, dest=i, tag=self.mpi_rank)) + # wait inside the block: the buffers must stay valid until MPI is done, and staged + # receive rows are copied to the device when the block ends + for req in reqs: + req.Wait() + + if n_recv > 0: + self._markers[hole_inds_after_send[:n_recv]] = recvbuf + class Tesselation: """ diff --git a/src/struphy/pic/tests/test_sorting_device.py b/src/struphy/pic/tests/test_sorting_device.py new file mode 100644 index 000000000..55c5d09f5 --- /dev/null +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -0,0 +1,190 @@ +"""Marker exchange between ranks on the CuPy backend (:meth:`Particles.mpi_sort_markers`). + +The marker rows are exchanged as device buffers with CUDA-aware MPI, and staged through +host memory otherwise. Without a GPU, the tests run on cunumpy's fake CuPy, which must be +installed before Python imports cunumpy, so they are skipped unless the fake is active:: + + mpirun -n 2 env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi src/struphy/pic/tests/test_sorting_device.py + +Fake CuPy arrays live in host memory, so a non-CUDA-aware MPI library takes them, and the +CUDA-aware path (the probe of :func:`cunumpy.mpi.mpi_is_cuda_aware` succeeds) is what runs by +default. The staged path is forced with :func:`cunumpy.mpi.set_mpi_cuda_aware`. +""" + +import sys +from contextlib import contextmanager + +import cunumpy as xp +import numpy as np +import pytest +from cunumpy.kernel_testing import requires_cupy +from maybempi import MPI + +from struphy import BoundaryParameters, LoadingParameters +from struphy.pic.particles import Particles6D + + +def _fake_cupy_active() -> bool: + return bool(getattr(sys.modules.get("cupy"), "__cunumpy_fake__", False)) + + +requires_fake_cupy = pytest.mark.skipif( + not _fake_cupy_active(), + reason="needs the fake CuPy: run under `mpirun -n N env CUNUMPY_FAKE_CUPY=1 pytest ...`", +) + + +@contextmanager +def _count_host_reads(): + """Count reads of fake CuPy arrays by the host (``int()``, ``bool()``, iteration, ``.get()``, ...). + + :func:`cunumpy.profiling.count_transfers` sees the copies made through cunumpy; this also sees + the implicit ones (a device scalar in an ``if``, ``list(device_array)``), which synchronize + and copy on a GPU. + """ + import cupy + + reads = [] + names = ("__int__", "__float__", "__bool__", "__index__", "__iter__", "get") + originals = {name: getattr(cupy.ndarray, name) for name in names} + original_getattr = cupy.ndarray.__getattr__ + + def counting(name, method): + def wrapper(self, *args, **kwargs): + reads.append(name) + return method(self, *args, **kwargs) + + return wrapper + + def counting_getattr(self, name): + if name in ("item", "tolist"): + reads.append(name) + return original_getattr(self, name) + + for name, method in originals.items(): + setattr(cupy.ndarray, name, counting(name, method)) + cupy.ndarray.__getattr__ = counting_getattr + try: + yield reads + finally: + for name, method in originals.items(): + setattr(cupy.ndarray, name, method) + cupy.ndarray.__getattr__ = original_getattr + + +def _make_particles(comm, bc): + return Particles6D( + comm_world=comm, + loading_params=LoadingParameters(Np=2000, seed=1234), + boundary_params=BoundaryParameters(bc=(bc, bc, bc)), + ) + + +def _shift_positions(markers: np.ndarray, valid: np.ndarray, step: int) -> np.ndarray: + """New positions in [0, 1) for the valid rows, far enough away that most markers change rank.""" + new = markers.copy() + shift = np.array([0.37, 0.21, 0.53]) * (step + 1) + new[valid, :3] = (markers[valid, :3] + shift) % 1.0 + return new + + +def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", count_reads: bool = True): + """Sort the same markers on NumPy and CuPy across all ranks and compare the marker arrays. + + ``cuda_aware=None`` uses the probe (:func:`cunumpy.mpi.mpi_is_cuda_aware`); ``True``/``False`` + force the device-buffer or the host-staged exchange. + """ + comm = MPI.COMM_WORLD + rank, size = comm.Get_rank(), comm.Get_size() + + ref = _make_particles(comm, bc) + ref.draw_markers(sort=False) + ref.mpi_sort_markers() + + # the reference is sorted on NumPy, the device markers inside `use_backend("cupy")` blocks + on_device = lambda: xp.use_backend("cupy") # noqa: E731 + + previous = xp.mpi.get_mpi_cuda_aware() + try: + with on_device(): + if cuda_aware is not None: + xp.mpi.set_mpi_cuda_aware(cuda_aware) + else: + xp.mpi.set_mpi_cuda_aware(None) + cuda_aware = xp.mpi.mpi_is_cuda_aware(comm) + + dev = _make_particles(comm, bc) + assert xp.get_array_backend(dev.markers) == "cupy" + assert dev.markers.shape == ref.markers.shape + + for step in range(2): + new = _shift_positions(ref.markers, ref.valid_mks, step) + + # number of rows this rank has to send, from the host reference + etas = new[ref.valid_mks, :3] + left, right = ref.domain_array[rank, 0::3], ref.domain_array[rank, 1::3] + right = np.where(right == 1.0, np.nextafter(right, 2.0), right) + n_send = int(np.count_nonzero(~np.all((etas >= left) & (etas < right), axis=1))) + n_send_total = comm.allreduce(n_send) + assert size == 1 or n_send_total > 0 + + ref.markers[:] = new + ref.update_holes() + n_before = ref.n_mks_loc + ref.mpi_sort_markers(do_test=True) + n_recv = ref.n_mks_loc - (n_before - n_send) + + with on_device(): + dev.markers[:] = xp.to_cupy(new) + dev.update_holes() + with xp.profiling.count_transfers() as transfers, _count_host_reads() as reads: + dev.mpi_sort_markers() + + if cuda_aware: + assert transfers.total == 0, transfers.report() + else: + # one staged copy per non-empty send buffer, one copy back of the receive buffer + n_dest = sum(1 for i in range(size) if i != rank and dev._send_list[i].shape[0] > 0) + assert transfers.to_host == n_dest, transfers.report() + assert transfers.to_device == (1 if n_recv > 0 else 0), transfers.report() + sent_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_host") + assert sent_bytes == n_send * ref.markers.shape[1] * 8 + + if count_reads and cuda_aware: + # (the staged exchange reads the send buffers with .get(), counted above) + assert reads == [], f"host reads of device arrays during the exchange: {reads}" + + with on_device(): + assert np.array_equal(xp.to_numpy(dev.markers), ref.markers) + assert np.array_equal(xp.to_numpy(dev.valid_mks), ref.valid_mks) + n_dev = dev.n_mks_loc + assert n_dev == ref.n_mks_loc + + with on_device(): + # all markers are on the right rank now: nothing moves + dev.mpi_sort_markers(do_test=True) + assert np.array_equal(xp.to_numpy(dev.markers), ref.markers) + finally: + xp.mpi.set_mpi_cuda_aware(previous) + + +@requires_fake_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("cuda_aware", [None, True, False]) +@pytest.mark.parametrize("bc", ["periodic", "remove"]) +def test_mpi_sort_markers_fake_cupy(cuda_aware, bc): + """On the fake CuPy: the exchange gives the NumPy result, without host/device copies when MPI is CUDA-aware.""" + check_mpi_sort_markers_device(cuda_aware, bc=bc) + + +@requires_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("bc", ["periodic", "remove"]) +def test_mpi_sort_markers_gpu(bc): + """On a GPU (``mpirun -n 2``, one GPU per rank): the exchange gives the NumPy result; with + CUDA-aware MPI without host/device copies.""" + check_mpi_sort_markers_device(None, bc=bc, count_reads=False) + + +if __name__ == "__main__": + check_mpi_sort_markers_device(None) From e6930958a113ef72368ae5872e0da1b68a341312 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 20:49:17 +0200 Subject: [PATCH 06/50] CUDA version of the vlasov_maxwell accumulation Port the accumulation kernel of VlasovAmpereCoupling (VlasovAmpereOneSpecies, VlasovMaxwellOneSpecies) to CUDA, 1:1 with pyccel: Array6D matrix blocks, atomic adds, holes skipped and boundary particles accumulated as in pyccel. The device helpers fill_mat, fill_mat_vec and m_v_fill_b_v1_symm, v1_symm_accumulation_data() and the m_v_fill_b_v1_symm wrapper tests are byte-identical to the linear_vlasov_ampere PR (#705). Parity cases in Cuboid, Colella, HollowTorus, ShafranovDshapedCylinder and a 3d spline mapping; the CPU emulation agrees with pyccel and fails for a transposed DF. Part of #688. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 24 +++- .../pic/accumulation/filler_kernels.cuh | 125 +++++++++++++++++- .../vlasov_maxwell/vlasov_maxwell_cuda.cu | 78 +++++++++++ .../accumulation/particle_to_mat_kernels.cuh | 69 ++++++++++ src/struphy/pic/tests/cuda_parity_cases.py | 13 ++ src/struphy/pic/tests/kernel_test_args.py | 27 +++- src/struphy/pic/tests/test_cuda_emulation.py | 12 ++ src/struphy/pic/tests/test_cuda_parity.py | 2 + src/struphy/pic/tests/test_device_helpers.py | 61 +++++++++ 9 files changed, 406 insertions(+), 5 deletions(-) create mode 100644 src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_cuda.cu diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index daf0d25f0..c0f7edd49 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -289,7 +289,7 @@ done (✓), in this PR (PR 14), blocked (⏸). |---|---|---|---|---| | 1 | `push_v_with_efield` | P | `VlasovAmpereCoupling`, `PushVinForceField` | PR 14 | | 2 | `charge_density_0form` | A (vector) | initial Poisson solve of these models | PR 14 | -| 3 | `vlasov_maxwell` | A (matrix + vector) | `VlasovAmpereCoupling` | ⏸ 6D views | +| 3 | `vlasov_maxwell` | A (matrix + vector) | `VlasovAmpereCoupling` | ✓ (#688, see [notes](#vlasov_maxwell-implementation-notes)) | **Step 2 – LinearVlasovAmpère/Maxwell (δf)** @@ -626,3 +626,25 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a The GPU tests have not run on an H100; the CPU emulation and the CPU regression tests are the gate. - **Parity cases.** The spline mappings are added to the geometry kernels' entries of `pic/tests/cuda_parity_cases.py` (the per-folder `_test_args.py` modules were replaced by that module in PR 16). + +## `vlasov_maxwell` implementation notes + +The accumulation of `VlasovAmpereCoupling` in `VlasovAmpereOneSpecies` and `VlasovMaxwellOneSpecies` (step 1, #3 of +the [porting order](#porting-order); part of #688). + +- **Kernel.** `vlasov_maxwell_cuda.cu`: same arguments in the same order as pyccel, one thread per marker row; `df`, + `matrix_inv`, `transpose`, `matrix_matrix`, `matrix_vector`, then `m_v_fill_b_v1_symm` with + `A_p = w_p DF^{-1} DF^{-T}` and `B_p = w_p DF^{-1} v_p`. As in pyccel, only holes are skipped: boundary particles + (`markers[ip, -1] == -2`) are accumulated (unlike `linear_vlasov_ampere`). +- **Device helpers.** `fill_mat`, `fill_mat_vec` (`filler_kernels.cuh`) and `m_v_fill_b_v1_symm` + (`particle_to_mat_kernels.cuh`) with atomic adds into `Array6D`/`Array3D` data, byte-identical to the + ones of the `linear_vlasov_ampere` PR (#705), as are `v1_symm_accumulation_data()` and the `m_v_fill_b_v1_symm` + wrapper tests. +- **Tests.** `PARITY_CASES["vlasov_maxwell"]`: 129 markers (a hole, a boundary particle) in Cuboid, Colella, + HollowTorus, ShafranovDshapedCylinder and a 3d spline mapping; the CPU emulation agrees with pyccel to 7e-12 at + entries up to 6e5 (spline mapping) and fails for a transposed `DF`, skipped boundary particles or unskipped holes. + `rtol = 1e-12`, `atol = 1e-8` for the GPU's atomic summation order. GPU tests have not run on an H100. +- **Still missing for the models end to end on the GPU.** `MassMatrixPreconditioner` (the default of + `VlasovAmpereCoupling`, `MaxwellWeakAmpere` and the initial `PoissonSolve`) cannot be created on CuPy, CG inner products + return host scalars, the FEEC operators (`curl`, `grad`, mass matrices, Schur solves) need the feectools CUDA stack, + multi-rank marker sorting goes through the host, and the first H100 run (#687) is pending. diff --git a/src/struphy/pic/accumulation/filler_kernels.cuh b/src/struphy/pic/accumulation/filler_kernels.cuh index 37f3867d9..8a33d7773 100644 --- a/src/struphy/pic/accumulation/filler_kernels.cuh +++ b/src/struphy/pic/accumulation/filler_kernels.cuh @@ -1,5 +1,5 @@ -// Device versions of the fillers in filler_kernels.py: add the contribution of one marker to a stencil vector. -// Many markers write to the same entries, so every addition is atomic. +// Device versions of the fillers in filler_kernels.py: add the contribution of one marker to a stencil vector +// or to a block of a stencil matrix. Many markers write to the same entries, so every addition is atomic. #pragma once #include "cunumpy/array_view.cuh" #include "cunumpy/atomic.cuh" @@ -40,4 +40,125 @@ __device__ inline void fill_vec(int pi1, int pi2, int pi3, const double* bi1, co } } } + +/** + * Add basis functions times filling to a block of a stencil matrix, as in filler_kernels.fill_mat. + * + * @param pi1 Spline degree of the codomain (row indices) along the first axis. + * @param pi2 Spline degree of the codomain along the second axis. + * @param pi3 Spline degree of the codomain along the third axis. + * @param pj1 Spline degree of the domain (column indices) along the first axis. + * @param pj2 Spline degree of the domain along the second axis. + * @param pj3 Spline degree of the domain along the third axis. + * @param bi1 The pi1 + 1 non-vanishing N- or D-spline values of the codomain along the first axis. + * @param bi2 The pi2 + 1 non-vanishing spline values of the codomain along the second axis. + * @param bi3 The pi3 + 1 non-vanishing spline values of the codomain along the third axis. + * @param bj1 The pj1 + 1 non-vanishing spline values of the domain along the first axis. + * @param bj2 The pj2 + 1 non-vanishing spline values of the domain along the second axis. + * @param bj3 The pj3 + 1 non-vanishing spline values of the domain along the third axis. + * @param span1 Knot span index along the first axis. + * @param span2 Knot span index along the second axis. + * @param span3 Knot span index along the third axis. + * @param starts Start indices of the codomain (three entries). + * @param pads Paddings of the codomain (three entries). + * @param mat Stencil matrix data (the _data of a StencilMatrix, shape (n1, n2, n3, 2 * pads + 1)), any strides; + * written atomically. + * @param filling Number multiplied by the basis functions. + * + * The summation order of concurrent threads (and the last bits of the result) depends on the thread schedule. + */ +__device__ inline void fill_mat(int pi1, int pi2, int pi3, int pj1, int pj2, int pj3, const double* bi1, + const double* bi2, const double* bi3, const double* bj1, const double* bj2, + const double* bj3, int span1, int span2, int span3, const long long* starts, + const long long* pads, Array6D mat, double filling) { + for (int il1 = 0; il1 <= pi1; ++il1) { + long long i1 = span1 + il1 - starts[0]; + double b1 = bi1[il1] * filling; + for (int il2 = 0; il2 <= pi2; ++il2) { + long long i2 = span2 + il2 - starts[1]; + double b2 = b1 * bi2[il2]; + for (int il3 = 0; il3 <= pi3; ++il3) { + long long i3 = span3 + il3 - starts[2]; + double b3 = b2 * bi3[il3]; + + for (int jl1 = 0; jl1 <= pj1; ++jl1) { + long long j1 = pads[0] + jl1 - il1; + double b4 = b3 * bj1[jl1]; + for (int jl2 = 0; jl2 <= pj2; ++jl2) { + long long j2 = pads[1] + jl2 - il2; + double b5 = b4 * bj2[jl2]; + for (int jl3 = 0; jl3 <= pj3; ++jl3) { + long long j3 = pads[2] + jl3 - il3; + double b6 = b5 * bj3[jl3]; + + cunumpy_atomic_add(&mat(i1, i2, i3, j1, j2, j3), b6); + } + } + } + } + } + } +} + +/** + * Add basis functions times fillings to a block of a stencil matrix and to a stencil vector, as in + * filler_kernels.fill_mat_vec. + * + * @param pi1 Spline degree of the codomain (row indices) along the first axis. + * @param pi2 Spline degree of the codomain along the second axis. + * @param pi3 Spline degree of the codomain along the third axis. + * @param pj1 Spline degree of the domain (column indices) along the first axis. + * @param pj2 Spline degree of the domain along the second axis. + * @param pj3 Spline degree of the domain along the third axis. + * @param bi1 The pi1 + 1 non-vanishing N- or D-spline values of the codomain along the first axis. + * @param bi2 The pi2 + 1 non-vanishing spline values of the codomain along the second axis. + * @param bi3 The pi3 + 1 non-vanishing spline values of the codomain along the third axis. + * @param bj1 The pj1 + 1 non-vanishing spline values of the domain along the first axis. + * @param bj2 The pj2 + 1 non-vanishing spline values of the domain along the second axis. + * @param bj3 The pj3 + 1 non-vanishing spline values of the domain along the third axis. + * @param span1 Knot span index along the first axis. + * @param span2 Knot span index along the second axis. + * @param span3 Knot span index along the third axis. + * @param starts Start indices of the codomain (three entries). + * @param pads Paddings of the codomain (three entries). + * @param mat Stencil matrix data (the _data of a StencilMatrix), any strides; written atomically. + * @param filling_mat Number multiplied by the basis functions and added to mat. + * @param vec Stencil vector data of the codomain (the _data of a StencilVector), any strides; written atomically. + * @param filling_vec Number multiplied by the basis functions of the codomain and added to vec. + */ +__device__ inline void fill_mat_vec(int pi1, int pi2, int pi3, int pj1, int pj2, int pj3, const double* bi1, + const double* bi2, const double* bi3, const double* bj1, const double* bj2, + const double* bj3, int span1, int span2, int span3, const long long* starts, + const long long* pads, Array6D mat, double filling_mat, + Array3D vec, double filling_vec) { + for (int il1 = 0; il1 <= pi1; ++il1) { + long long i1 = span1 + il1 - starts[0]; + double b1 = bi1[il1]; + for (int il2 = 0; il2 <= pi2; ++il2) { + long long i2 = span2 + il2 - starts[1]; + double b2 = b1 * bi2[il2]; + for (int il3 = 0; il3 <= pi3; ++il3) { + long long i3 = span3 + il3 - starts[2]; + double b3 = b2 * bi3[il3]; + + cunumpy_atomic_add(&vec(i1, i2, i3), b3 * filling_vec); + + for (int jl1 = 0; jl1 <= pj1; ++jl1) { + long long j1 = pads[0] + jl1 - il1; + double b4 = b3 * bj1[jl1] * filling_mat; + for (int jl2 = 0; jl2 <= pj2; ++jl2) { + long long j2 = pads[1] + jl2 - il2; + double b5 = b4 * bj2[jl2]; + for (int jl3 = 0; jl3 <= pj3; ++jl3) { + long long j3 = pads[2] + jl3 - il3; + double b6 = b5 * bj3[jl3]; + + cunumpy_atomic_add(&mat(i1, i2, i3, j1, j2, j3), b6); + } + } + } + } + } + } +} } diff --git a/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_cuda.cu b/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_cuda.cu new file mode 100644 index 000000000..e6abf9bae --- /dev/null +++ b/src/struphy/pic/accumulation/kernels/vlasov_maxwell/vlasov_maxwell_cuda.cu @@ -0,0 +1,78 @@ +#include "struphy/geometry/evaluation_kernels.cuh" +#include "struphy/linear_algebra/linalg_kernels.cuh" +#include "struphy/pic/accumulation/particle_to_mat_kernels.cuh" + +/** + * Accumulate the matrix and vector of the full-orbit Vlasov-Maxwell (Vlasov-Ampere) coupling into V1, as in + * vlasov_maxwell_kernels.vlasov_maxwell. + * + * Filling functions: + * A_p^{mu, nu} = w_p G^{-1}_{mu, nu}(eta_p), G^{-1} = DF^{-1} DF^{-T}, + * B_p^mu = w_p [DF^{-1}(eta_p) v_p]_mu; + * one thread per marker row. + * + * @param args_markers Marker buffer (n_markers x n_cols, row-major); the weight is in column 6, as in pyccel. + * @param args_derham Spline degrees (1 to 8), knots and start indices. + * @param args_domain Mapping arguments; every spline and analytic mapping. + * @param mat11 (1,1)-block data of the symmetric V1 -> V1 block matrix (the _data of a StencilMatrix, 6D, any + * strides); written with atomic additions. + * @param mat12 (1,2)-block data; written atomically. + * @param mat13 (1,3)-block data; written atomically. + * @param mat22 (2,2)-block data; written atomically. + * @param mat23 (2,3)-block data; written atomically. + * @param mat33 (3,3)-block data; written atomically. + * @param vec1 First component of the V1 vector data; written atomically. + * @param vec2 Second component of the V1 vector data; written atomically. + * @param vec3 Third component of the V1 vector data; written atomically. + * + * Holes (markers[ip, 0] == -1) are skipped; boundary particles are not, as in pyccel. The pyccel loop runs over all + * rows (shape(markers)[0]), which is n_markers here. Many threads add to the same entries, so the summation order + * (and the last bits of the result) differs from the serial pyccel loop. + */ +extern "C" __global__ void vlasov_maxwell(MarkerArgs args_markers, DerhamArgs args_derham, DomainArgs args_domain, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + Array3D vec1, Array3D vec2, Array3D vec3) { + // CUDA-only: the marker row of this thread replaces the pyccel loop variable + int ip = blockDim.x * blockIdx.x + threadIdx.x; + if (ip >= args_markers.n_markers) return; + + // allocate for metric coeffs (3 x 3 matrices as nine row-major entries) + double dfm[9], df_inv[9], df_inv_t[9], g_inv[9]; + + // allocate for filling + double v[3], df_inv_times_v[3], filling_m[9], filling_v[3]; + + // only do something if particle is a "true" particle (i.e. not a hole) + if (args_markers.markers(ip, 0) == -1.) return; + + // marker positions + double eta1 = args_markers.markers(ip, 0); + double eta2 = args_markers.markers(ip, 1); + double eta3 = args_markers.markers(ip, 2); + + // evaluate Jacobian, result in dfm + struphy_cuda::df(eta1, eta2, eta3, args_domain, dfm); + + // compute shifted and stretched velocity + v[0] = args_markers.markers(ip, 3); + v[1] = args_markers.markers(ip, 4); + v[2] = args_markers.markers(ip, 5); + + // filling functions + struphy_cuda::matrix_inv(dfm, df_inv); + struphy_cuda::transpose(df_inv, df_inv_t); + struphy_cuda::matrix_matrix(df_inv, df_inv_t, g_inv); + struphy_cuda::matrix_vector(df_inv, v, df_inv_times_v); + + // filling_m = w_p * DF^{-1} * DF^{-T} + for (int k = 0; k < 9; ++k) filling_m[k] = args_markers.markers(ip, 6) * g_inv[k]; + + // filling_v = w_p * DF^{-1} * \V + for (int k = 0; k < 3; ++k) filling_v[k] = args_markers.markers(ip, 6) * df_inv_times_v[k]; + + // call the appropriate matvec filler + struphy_cuda::m_v_fill_b_v1_symm(args_derham, eta1, eta2, eta3, mat11, mat12, mat13, mat22, mat23, mat33, + filling_m[0], filling_m[1], filling_m[2], filling_m[4], filling_m[5], + filling_m[8], vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2]); +} diff --git a/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh b/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh index 7cabaee45..ade422b97 100644 --- a/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh +++ b/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh @@ -28,4 +28,73 @@ __device__ inline void vec_fill_b_v0(const DerhamArgs& args_derham, double eta1, fill_vec(pn1, pn2, pn3, scratch.bn1, scratch.bn2, scratch.bn3, span1, span2, span3, args_derham.starts, vec, fill); } + +/** + * Add one marker's contribution to the symmetric blocks (1,1), (1,2), (1,3), (2,2), (2,3), (3,3) of a V1 -> V1 + * block matrix and to a V1 vector, as in particle_to_mat_kernels.m_v_fill_b_v1_symm. + * + * @param args_derham Spline degrees (1 to 8), knots and start indices; pn is also the padding of the blocks. + * @param eta1 Logical position along the first axis. + * @param eta2 Logical position along the second axis. + * @param eta3 Logical position along the third axis. + * @param mat11 (mu=1, nu=1)-block data (D N N x D N N), shape (n1, n2, n3, 2 * pn + 1); written atomically. + * @param mat12 (mu=1, nu=2)-block data (D N N x N D N); written atomically. + * @param mat13 (mu=1, nu=3)-block data (D N N x N N D); written atomically. + * @param mat22 (mu=2, nu=2)-block data (N D N x N D N); written atomically. + * @param mat23 (mu=2, nu=3)-block data (N D N x N N D); written atomically. + * @param mat33 (mu=3, nu=3)-block data (N N D x N N D); written atomically. + * @param fill11 Number multiplied by the basis functions and added to mat11. + * @param fill12 Number multiplied by the basis functions and added to mat12. + * @param fill13 Number multiplied by the basis functions and added to mat13. + * @param fill22 Number multiplied by the basis functions and added to mat22. + * @param fill23 Number multiplied by the basis functions and added to mat23. + * @param fill33 Number multiplied by the basis functions and added to mat33. + * @param vec1 mu=1 component of the V1 vector data (D N N); written atomically. + * @param vec2 mu=2 component of the V1 vector data (N D N); written atomically. + * @param vec3 mu=3 component of the V1 vector data (N N D); written atomically. + * @param fill1 Number multiplied by the basis functions and added to vec1. + * @param fill2 Number multiplied by the basis functions and added to vec2. + * @param fill3 Number multiplied by the basis functions and added to vec3. + * + * Pyccel keeps the spline values in args_derham.bn1, ..., bd3; here they are in a thread-local SplineScratch. + */ +__device__ inline void m_v_fill_b_v1_symm(const DerhamArgs& args_derham, double eta1, double eta2, double eta3, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + double fill11, double fill12, double fill13, double fill22, + double fill23, double fill33, Array3D vec1, Array3D vec2, + Array3D vec3, double fill1, double fill2, double fill3) { + // degrees of the basis functions : B-splines (pn) and D-splines (pd) + int pn1 = args_derham.pn[0]; + int pn2 = args_derham.pn[1]; + int pn3 = args_derham.pn[2]; + + int pd1 = pn1 - 1; + int pd2 = pn2 - 1; + int pd3 = pn3 - 1; + + // CUDA-only scratch holds the spans and the spline values pyccel keeps in args_derham + SplineScratch scratch; + get_spans(eta1, eta2, eta3, args_derham, scratch); + int span1 = scratch.span1, span2 = scratch.span2, span3 = scratch.span3; + + // fill matrix entries + fill_mat_vec(pd1, pn2, pn3, pd1, pn2, pn3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bd1, scratch.bn2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat11, fill11, vec1, fill1); + + fill_mat_vec(pn1, pd2, pn3, pn1, pd2, pn3, scratch.bn1, scratch.bd2, scratch.bn3, scratch.bn1, scratch.bd2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat22, fill22, vec2, fill2); + + fill_mat_vec(pn1, pn2, pd3, pn1, pn2, pd3, scratch.bn1, scratch.bn2, scratch.bd3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat33, fill33, vec3, fill3); + + fill_mat(pd1, pn2, pn3, pn1, pd2, pn3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bn1, scratch.bd2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat12, fill12); + + fill_mat(pd1, pn2, pn3, pn1, pn2, pd3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat13, fill13); + + fill_mat(pn1, pd2, pn3, pn1, pn2, pd3, scratch.bn1, scratch.bd2, scratch.bn3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat23, fill23); +} } diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index e2ce0a5d5..a13d4189a 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -29,6 +29,7 @@ marker_arguments, spline_coefficients, spline_evaluation_arguments, + v1_symm_accumulation_data, ) @@ -79,6 +80,14 @@ def charge_density_0form_args(bc): return (args_markers, derham_arguments(), args_domain, xp.zeros((18, 20, 16))) +def vlasov_maxwell_args(domain): + """129 markers (a hole in row 0, a boundary particle in row 1, which is accumulated as in pyccel) in domain + `domain` of geometry_domain(); zeroed V1 data.""" + args_markers, _ = marker_arguments((0, 0, 0)) + args_markers.markers[0, 0] = -1.0 + return (args_markers, derham_arguments(), geometry_domain(domain).args_domain, *v1_symm_accumulation_data()) + + # ---------------------------------------------------------------- spline evaluation @@ -267,6 +276,10 @@ def size_of(index, per_thread=1): "reflect": ParityCases((0, 1, 2), reflect_args, n_threads=size_of(2)), # atomics add in another order than the serial loop "charge_density_0form": ParityCases(((0, 0, 0), (2, 0, 1)), charge_density_0form_args, rtol=1e-12, atol=1e-13), + # Cuboid, then Colella, HollowTorus, ShafranovDshapedCylinder (non-diagonal DF) and a 3d spline mapping, by index into + # geometry_domain(); atomics add in another order than the serial loop, and G^{-1} of the spline mapping makes + # entries of 6e5 (atol for cancelling sums) + "vlasov_maxwell": ParityCases((0, 2, 5, 9, 12), vlasov_maxwell_args, rtol=1e-12, atol=1e-8), "eval_spline_mpi_markers": ParityCases(SPLINE_KINDS, eval_spline_mpi_markers_args), "eval_spline_mpi_matrix": ParityCases(SPLINE_KINDS, eval_spline_mpi_grid_args(sparse=False), n_threads=size_of(-1)), "eval_spline_mpi_sparse_meshgrid": ParityCases( diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index d413b5380..eaed48adf 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -49,14 +49,37 @@ def butcher_arguments(method): return xp.asarray(butcher.a_stage), xp.asarray(butcher.b), xp.asarray(butcher.c), butcher.n_stages +DERHAM_DEGREE = (2, 3, 1) +DERHAM_CELLS = 8 + + def derham_arguments(): """Splines of degrees 2, 3, 1 on 8 cells, starting at index 0.""" - degree = np.array([2, 3, 1], dtype=np.int64) - knots = [np.r_[np.zeros(p), np.linspace(0, 1, 9), np.ones(p)] for p in degree] + degree = np.array(DERHAM_DEGREE, dtype=np.int64) + knots = [np.r_[np.zeros(p), np.linspace(0, 1, DERHAM_CELLS + 1), np.ones(p)] for p in degree] args_class = CudaDerhamArguments if xp.get_backend() == "cupy" else DerhamArguments return args_class(xp.asarray(degree), *(xp.asarray(t) for t in knots), xp.zeros(3, dtype=np.int64)) +def v1_symm_accumulation_data(): + """Zeroed data of a symmetric V1 -> V1 accumulation matrix and a V1 vector, for :func:`derham_arguments`. + + Returns the six independent blocks ``mat11, mat12, mat13, mat22, mat23, mat33`` and the three vector components + ``vec1, vec2, vec3`` on the active backend, shaped like the ``_data`` of the ``StencilMatrix`` blocks and + ``StencilVector`` components that ``Accumulator(..., "Hcurl", ..., symmetry="symm", add_vector=True)`` passes on + one process for clamped splines of degrees 2, 3, 1 on 8 cells: component ``mu`` has D-splines (``n - 1`` + functions) along axis ``mu`` and N-splines (``n = 8 + p``) along the others, every axis is padded by ``p`` on both + sides, and the matrix blocks have ``2 p + 1`` diagonals per axis (the paddings are the degrees of V0). + """ + degree = np.array(DERHAM_DEGREE) + n_spline = DERHAM_CELLS + degree + rows = [tuple(n_spline - (np.arange(3) == mu) + 2 * degree) for mu in range(3)] + diagonals = tuple(2 * degree + 1) + mats = tuple(xp.zeros(rows[mu] + diagonals) for mu, nu in ((0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2))) + vecs = tuple(xp.zeros(rows[mu]) for mu in range(3)) + return mats + vecs + + def spline_coefficients(n=3, seed=11): """Random coefficients, large enough for every span of :func:`derham_arguments`.""" rng = np.random.default_rng(seed) diff --git a/src/struphy/pic/tests/test_cuda_emulation.py b/src/struphy/pic/tests/test_cuda_emulation.py index 793bb9525..f2ab111a2 100644 --- a/src/struphy/pic/tests/test_cuda_emulation.py +++ b/src/struphy/pic/tests/test_cuda_emulation.py @@ -113,3 +113,15 @@ def test_emulated_transform_helpers(domain_index): out = np.zeros(len(expected)) emulate_struct_kernel(kernel, *inputs, args_domain, out, len(out), n_threads=len(out)) np.testing.assert_allclose(out, expected, rtol=1e-10, atol=1e-10) + + +@requires_compiler +def test_emulated_v1_symm_filler(): + """test_device_helpers.test_v1_symm_filler by CPU emulation: the device m_v_fill_b_v1_symm agrees with pyccel.""" + from struphy.pic.tests.test_device_helpers import v1_symm_filler_case + + kernel, (eta, fills), args_derham, data, expected = v1_symm_filler_case() + emulate_struct_kernel(kernel, eta, fills, args_derham, *data, len(eta), n_threads=len(eta)) + for result, reference in zip(data, expected): + assert np.any(reference != 0.0) + np.testing.assert_allclose(result, reference, rtol=1e-12, atol=1e-13) diff --git a/src/struphy/pic/tests/test_cuda_parity.py b/src/struphy/pic/tests/test_cuda_parity.py index 75d53245f..7e39bdbef 100644 --- a/src/struphy/pic/tests/test_cuda_parity.py +++ b/src/struphy/pic/tests/test_cuda_parity.py @@ -10,6 +10,7 @@ from cunumpy.profiling import assert_no_transfers from struphy.pic.accumulation.kernels.charge_density_0form import charge_density_0form +from struphy.pic.accumulation.kernels.vlasov_maxwell import vlasov_maxwell from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage from struphy.pic.pushing.kernels.push_v_with_efield import push_v_with_efield from struphy.pic.pushing.kernels.push_vxb_analytic import push_vxb_analytic @@ -116,5 +117,6 @@ def test_vlasov_kernel_coverage(): push_v_with_efield, push_weights_with_efield_lin_va, charge_density_0form, + vlasov_maxwell, ): assert kernel.has_cuda, kernel.name diff --git a/src/struphy/pic/tests/test_device_helpers.py b/src/struphy/pic/tests/test_device_helpers.py index a2390ef4d..b771d8efc 100644 --- a/src/struphy/pic/tests/test_device_helpers.py +++ b/src/struphy/pic/tests/test_device_helpers.py @@ -391,3 +391,64 @@ def test_transform_helpers(domain_index): kernel, inputs, _, expected = transform_helper_case(domain_index) result = _run_on_device(kernel, inputs, domain_index, len(expected)) np.testing.assert_allclose(result, expected, rtol=1e-10, atol=1e-10) + + +# CUDA-only wrapper: one thread per point calls m_v_fill_b_v1_symm with the point's nine fillings +V1_SYMM_FILLER_SOURCE = r""" +#include "struphy/pic/accumulation/particle_to_mat_kernels.cuh" +extern "C" __global__ void fill_v1_symm(const double* eta, const double* fills, DerhamArgs args_derham, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + Array3D vec1, Array3D vec2, Array3D vec3, int n) { + int ip = blockDim.x * blockIdx.x + threadIdx.x; + if (ip >= n) return; + const double* fill = fills + 9 * ip; + struphy_cuda::m_v_fill_b_v1_symm(args_derham, eta[3 * ip], eta[3 * ip + 1], eta[3 * ip + 2], mat11, mat12, mat13, + mat22, mat23, mat33, fill[0], fill[1], fill[2], fill[3], fill[4], fill[5], vec1, + vec2, vec3, fill[6], fill[7], fill[8]); +} +""" + + +def v1_symm_filler_case(): + """Wrapper kernel, host inputs and pyccel reference for the device m_v_fill_b_v1_symm. + + 65 points (the corners of the unit cube among them) with random fillings, accumulated into zeroed V1 data shaped + like the real accumulation matrices (:func:`~struphy.pic.tests.kernel_test_args.v1_symm_accumulation_data`). + Used by :func:`test_v1_symm_filler` (GPU) and by its CPU emulation in ``test_cuda_emulation.py``. Returns + ``kernel, (eta, fills), args_derham, data, expected``: ``args_derham`` is the pyccel class and ``data`` the nine + zeroed host arrays in kernel order. + """ + from struphy.pic.accumulation import particle_to_mat_kernels + from struphy.pic.tests.kernel_test_args import derham_arguments, v1_symm_accumulation_data + + rng = np.random.default_rng(15) + eta = np.r_[np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0], [0.0, 1.0, 0.5]]), rng.random((62, 3))] + fills = rng.normal(size=(len(eta), 9)) + with cunumpy.use_backend("numpy"): + args_derham = derham_arguments() + data = v1_symm_accumulation_data() + expected = tuple(a.copy() for a in data) + for point, fill in zip(eta, fills): + particle_to_mat_kernels.m_v_fill_b_v1_symm( + args_derham, *point, *expected[:6], *fill[:6], *expected[6:], *fill[6:] + ) + kernel = CudaKernel(V1_SYMM_FILLER_SOURCE, "fill_v1_symm", **CUDA_OPTIONS) + return kernel, (eta, fills), args_derham, data, expected + + +@requires_cupy +def test_v1_symm_filler(): + """m_v_fill_b_v1_symm (and the fill_mat/fill_mat_vec it calls) on the device agrees with pyccel.""" + import cupy as cp + + from struphy.pic.tests.kernel_test_args import derham_arguments + + kernel, (eta, fills), _, data, expected = v1_symm_filler_case() + with cunumpy.use_backend("cupy"): + args_derham = derham_arguments() + device_data = [cp.asarray(a) for a in data] + kernel(cp.asarray(eta), cp.asarray(fills), args_derham, *device_data, len(eta), n_threads=len(eta)) + for result, reference in zip(device_data, expected): + assert np.any(reference != 0.0) + np.testing.assert_allclose(result.get(), reference, rtol=1e-12, atol=1e-13) From 14510d21e6fc44dd6acc4ea6593d6e4d775aef75 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 21:11:34 +0200 Subject: [PATCH 07/50] MHD equilibria on the CuPy backend Equilibria with host-only setup (AdhocTorus q_kind 1/2, AdhocTorusQPsi, EQDSKequilibrium) run their __init__ on the NumPy backend (setup_on_host) and hold only host data. SciPy spline evaluations go through host_call, GVEC/DESC evaluations through @evaluate_on_host: device arguments are copied to the host and the result back, once per call; NumPy arguments are evaluated as before. Tokamak no longer builds its default EQDSKequilibrium on the NumPy backend itself. New test_equils_cupy.py compares every equilibrium of equils on CuPy (GPU, or cunumpy's fake CuPy with host geometry kernels) with NumPy. Solves #696. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 34 ++- src/struphy/fields_background/base.py | 80 +++++ src/struphy/fields_background/equils.py | 32 +- .../tests/test_equils_cupy.py | 276 ++++++++++++++++++ .../geometry/domains/tokamak/tokamak.py | 7 +- 5 files changed, 410 insertions(+), 19 deletions(-) create mode 100644 src/struphy/fields_background/tests/test_equils_cupy.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index daf0d25f0..68c596ea3 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -347,8 +347,8 @@ 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, +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 without host round trips, and array views with more than 4 dimensions in cunumpy (all matrix accumulations write 6D stencil matrix data). @@ -367,11 +367,31 @@ matrix data). - **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. `xp.mpi_is_cuda_aware()` detects it; the exchange in `Particles.mpi_sort_markers` has to be checked for host staging buffers. - **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.** `setup_on_host` (`fields_background/base.py`) 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 `host_call`, and the GVEC/DESC `bv`, `jv`, + `p0`, `n0`, `gradB1` through `@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. ## PR 10 implementation notes @@ -598,7 +618,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/src/struphy/fields_background/base.py b/src/struphy/fields_background/base.py index 426fd8f28..7e894b366 100644 --- a/src/struphy/fields_background/base.py +++ b/src/struphy/fields_background/base.py @@ -1,9 +1,11 @@ "Base classes for MHD equilibria." +import functools import logging from abc import ABCMeta, abstractmethod import cunumpy as xp +import numpy as np from pyevtk.hl import gridToVTK from struphy.geometry.base import Domain @@ -18,6 +20,84 @@ logger = logging.getLogger("struphy") +def _on_device(arg) -> bool: + """Whether ``arg`` is (or contains, for tuples and lists) a CuPy array.""" + if isinstance(arg, (tuple, list)): + return any(_on_device(a) for a in arg) + return xp.is_gpu(arg) + + +def _to_host(arg): + """Copy CuPy arrays in ``arg`` (an array, or a tuple/list of them) to the host; leave everything else alone.""" + if isinstance(arg, (tuple, list)): + return type(arg)(_to_host(a) for a in arg) + return xp.to_numpy(arg) if xp.is_gpu(arg) else arg + + +def _to_device(arg): + """Copy NumPy arrays in ``arg`` (an array, or a tuple/list of them) to the device; leave scalars alone.""" + if isinstance(arg, (tuple, list)): + return type(arg)(_to_device(a) for a in arg) + return xp.to_cupy(arg) if isinstance(arg, np.ndarray) else arg + + +def host_call(fun, *args, **kwargs): + """Call a host-only function (SciPy spline, external equilibrium code) with arguments of any backend. + + Device (CuPy) arrays among the arguments are copied to the host, ``fun`` runs on the NumPy backend, and its array + results are copied back to the device, once per call. Without device arguments, ``fun`` is called directly on the + NumPy backend (on the NumPy backend, this is a plain call), so the result lives where the arguments live. + + The copies cost one host round trip per call; equilibria are evaluated at setup (initial conditions, projections + onto the FEEC spaces, mappings), not in the time loop, which uses the projected equilibrium. + + Parameters + ---------- + fun : callable + The host-only function. + + *args, **kwargs + Arguments of ``fun``; arrays (or tuples/lists of arrays) of either backend, or scalars. + + Returns + ------- + The result of ``fun`` (an array, a tuple/list of arrays or a scalar), with arrays on the device if any argument + was on the device. + """ + device = _on_device(args) or _on_device(list(kwargs.values())) + if xp.get_backend() == "numpy" and not device: + return fun(*args, **kwargs) + with xp.use_backend("numpy"): + out = fun(*_to_host(args), **{k: _to_host(v) for k, v in kwargs.items()}) + return _to_device(out) if device else out + + +def evaluate_on_host(method): + """Decorator for equilibrium methods that can only be evaluated on the host (see :func:`host_call`).""" + + @functools.wraps(method) + def wrapper(self, *args, **kwargs): + return host_call(method, self, *args, **kwargs) + + return wrapper + + +def setup_on_host(init): + """Decorator for ``__init__`` of equilibria whose setup is host-only (file reading, ODE solves, SciPy fits). + + ``__init__`` runs on the NumPy backend, so the equilibrium holds only host data (NumPy arrays, SciPy splines, + floats) on either backend. Its evaluation follows the backend of the arguments; host-only parts of it go through + :func:`host_call`. + """ + + @functools.wraps(init) + def wrapper(self, *args, **kwargs): + with xp.use_backend("numpy"): + init(self, *args, **kwargs) + + return wrapper + + class FluidEquilibrium(metaclass=ABCMeta): """ Abstract base class for callable fluid equilibria on arbitrary domains. diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index 34dac281f..576c918c0 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -27,6 +27,9 @@ NumericalFluidEquilibrium, NumericalFluidEquilibriumWithB, NumericalMHDequilibrium, + evaluate_on_host, + host_call, + setup_on_host, ) from struphy.fields_background.mhd_equil.eqdsk import readeqdsk from struphy.io.options import BaseUnits @@ -874,6 +877,7 @@ def doc_formula(cls): Units are those defned in the parameter file (through :class:`~struphy.io.options.BaseUnits`). """ + @setup_on_host def __init__( self, a: float = 1.0, @@ -1010,7 +1014,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 +1025,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 = host_call(self._psi_i, r, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(r, (int, float)): @@ -1149,7 +1154,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 = host_call(self._p_i, r) # remove all "dimensions" for point-wise evaluation if isinstance(r, (int, float)): @@ -1365,6 +1370,7 @@ def doc_formula(cls): Units are those defned in the parameter file (through :class:`~struphy.io.options.BaseUnits`). """ + @setup_on_host def __init__( self, a: float = 0.361925, @@ -1470,7 +1476,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 = host_call(self._psi_i, r, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(r, (int, float)): @@ -1665,6 +1671,7 @@ class EQDSKequilibrium(AxisymmMHDequilibrium): Struphy base units. If None, no rescaling of output is performed. """ + @setup_on_host def __init__( self, rel_path: bool = True, @@ -1864,7 +1871,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 = host_call(self._q_i, psi, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(psi, (int, float)): @@ -1875,7 +1882,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 = host_call(self._g_i, psi, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(psi, (int, float)): @@ -1886,7 +1893,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 = host_call(self._p_i, psi, nu=der) # remove all "dimensions" for point-wise evaluation if isinstance(psi, (int, float)): @@ -1922,7 +1929,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 = 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 +2128,7 @@ def units(self) -> Units: """All Struphy units.""" return self._units + @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 +2150,7 @@ def bv(self, *etas, squeeze_out=False): return out + @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 +2180,7 @@ def jv(self, *etas, squeeze_out=False): return out + @evaluate_on_host @profile def p0(self, *etas, squeeze_out=False): """0-form equilibrium pressure on logical cube [0, 1]^3.""" @@ -2190,6 +2200,7 @@ def p0(self, *etas, squeeze_out=False): return self.params["p0"] + tmp / self.units.p + @evaluate_on_host @profile def n0(self, *etas, squeeze_out=False): """0-form equilibrium density on logical cube [0, 1]^3.""" @@ -2407,6 +2418,7 @@ def units(self) -> Units: """All Struphy units.""" return self._units + @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 +2492,7 @@ def _eval_bv(self, *etas, squeeze_out=False): return out + @evaluate_on_host @profile def jv(self, *etas, squeeze_out=False): """Contra-variant (vector field) current density (=curl B) @@ -2555,6 +2568,7 @@ def _eval_jv(self, *etas, squeeze_out=False): return out + @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 +2600,7 @@ def p0(self, *etas, squeeze_out=False): return out + @evaluate_on_host @profile def n0(self, *etas, squeeze_out=False): """0-form equilibrium density on logical cube [0, 1]^3.""" @@ -2610,6 +2625,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 + @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..f7ced7543 --- /dev/null +++ b/src/struphy/fields_background/tests/test_equils_cupy.py @@ -0,0 +1,276 @@ +"""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 subprocess +import sys + +import cunumpy +import numpy as np +import pytest +from cunumpy.kernel_testing import requires_cupy + +from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env + +# (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:`host_call` calls the function directly and returns its result unchanged.""" + from struphy.fields_background.base import host_call + + a = np.linspace(0.0, 1.0, 5) + with cunumpy.use_backend("numpy"): + out = 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})" + ) + 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:] + + +@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/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) From b51e11f276d602b3c137e159bda0f3638b3d0353 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 22:13:48 +0200 Subject: [PATCH 08/50] Hand-written CUDA vs code generation: port bstar_parallel_3form both ways (decision for #690) - bstar_parallel_3form_cuda.cu: hand-written CUDA version, 1:1 with pyccel, on the existing device helpers; new device helper eval_0form_spline_mpi in evaluation_kernels_3d.cuh. - 14 parity cases (every mapping, alpha = 1/0/mixed, wrap-around in mod, hole) in cuda_parity_cases.py; passes the CPU emulation. - pic/tests/codegen_spike: prototype generator translating the pyccel source (kernel and all helpers) into numba.cuda; agrees with pyccel in the numba.cuda simulator and under njit. Test-only, skipped without numba. - CUDA_STRATEGY.md: decision section (hand-written stays), comparison and what would change it. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 91 +++- .../bsplines/evaluation_kernels_3d.cuh | 18 + .../bstar_parallel_3form_cuda.cu | 75 +++ .../pic/tests/codegen_spike/__init__.py | 1 + .../pic/tests/codegen_spike/compare.py | 61 +++ .../pic/tests/codegen_spike/numba_codegen.py | 458 ++++++++++++++++++ .../tests/codegen_spike/test_numba_codegen.py | 29 ++ src/struphy/pic/tests/cuda_parity_cases.py | 32 ++ 8 files changed, 762 insertions(+), 3 deletions(-) create mode 100644 src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu create mode 100644 src/struphy/pic/tests/codegen_spike/__init__.py create mode 100644 src/struphy/pic/tests/codegen_spike/compare.py create mode 100644 src/struphy/pic/tests/codegen_spike/numba_codegen.py create mode 100644 src/struphy/pic/tests/codegen_spike/test_numba_codegen.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index daf0d25f0..2a621ce88 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -26,7 +26,7 @@ Progress is tracked in struphy-hub/struphy#650. - [ ] **PR 17: struphy on the feectools CUDA stack** — the `feectools` submodule points at the top of the feectools CUDA stack (`cuda-4-device-kernels`, [feectools#88](https://github.com/struphy-hub/feectools/pull/88)) instead of `devel-tiny`, so the struphy CUDA PRs run against feectools with device stencil operations, MPI with device buffers and one GPU per rank. Moves along with the stack; before the struphy stack is merged into `devel`, the feectools stack is merged into `devel-tiny` and the submodule points there again (see [feectools](#feectools)). - [x] **PR 18: Geometry evaluation on the GPU for all analytic mappings** (#681, open) — CUDA versions of the four geometry entry kernels (`kernel_evaluate_pic`, `kernel_evaluate`, `kernel_pullpush_pic`, `kernel_pullpush` in `geometry/kernels/`), built on device versions of the whole metric chain (`f`, `df`, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`/`push`/`tran`) for every analytic mapping (`kind_map` 10–12, 20–22, 30–32). Restores CuPy particle runs (weight initialization evaluates `jacobian_det`), and removes the Cuboid-only checks in `Pusher`, the accumulators and `reflect`. Parity arguments cover every analytic mapping (see [PR 18](#pr-18-geometry-evaluation-for-all-analytic-mappings) and the [PR 18 implementation notes](#pr-18-implementation-notes)). - [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines and `EQDSKequilibrium` stay host-only (see [PR 19](#pr-19-spline-mappings) and the [PR 19 implementation notes](#pr-19-implementation-notes)). -- [ ] **Next (order to be confirmed)**: 6D array views in cunumpy and the blocked matrix accumulations `vlasov_maxwell`, `linear_vlasov_ampere` (steps 1–2 of the [porting order](#porting-order)); one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, including when to compile the kernels, which needs the [feectools stack](#feectools) merged and released first; a decision on hand-written CUDA vs. code generation before the guiding-center kernels (step 3). +- [ ] **Next (order to be confirmed)**: 6D array views in cunumpy and the blocked matrix accumulations `vlasov_maxwell`, `linear_vlasov_ampere` (steps 1–2 of the [porting order](#porting-order)); one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, including when to compile the kernels, which needs the [feectools stack](#feectools) merged and released first; the guiding-center kernels (step 3), by hand (decided in #690, see [Hand-written CUDA vs. code generation](#hand-written-cuda-vs-code-generation-decision-for-690)). - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)). - [x] **PR 13: Move the kernel infrastructure to cunumpy** — `Kernel`, `KernelCatalog`, `CudaKernel` and `Argument` are replaced by their cunumpy counterparts; each owner has a single `args_*` object on both backends, and `pusher_args.cuh` is generated (see [Moving to cunumpy](#moving-to-cunumpy-pr-13)). Kernels and device helpers only change their includes. - [ ] **CI**: a GPU runner that runs the CUDA tests (can happen any time; until then the GPU tests are run by hand on an H100 before each PR that touches CUDA code is merged). @@ -315,7 +315,7 @@ algorithm is `discrete_gradient_1st_order`; its chain comes first. | 16–17 | `push_gc_bxEstar_discrete_gradient_2nd_order`, `push_gc_Bstar_discrete_gradient_2nd_order` | P | | 18–19 | `push_gc_bxEstar_explicit_multistage`, `push_gc_Bstar_explicit_multistage` | P | -These are the largest kernels; evaluate code generation (see [Open questions](#open-questions)) before porting 13–19 by hand. +These are the largest kernels. They are ported by hand like the others (decision for #690, see [Hand-written CUDA vs. code generation](#hand-written-cuda-vs-code-generation-decision-for-690)); `bstar_parallel_3form` (7) has its CUDA version. **Step 4 – Hybrid MHD–kinetic 6D (current and pressure coupling)** @@ -366,13 +366,98 @@ matrix data). - **Accumulation strategy.** Atomics vs. sort-then-reduce; see PR 12+. Decided by measurement on the first accumulation kernel. - **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. `xp.mpi_is_cuda_aware()` detects it; the exchange in `Particles.mpi_sort_markers` has to be checked for host staging buffers. -- **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. +- **Single-source alternatives.** Decided for now (#690): hand-written CUDA stays the default; see [Hand-written CUDA vs. code generation](#hand-written-cuda-vs-code-generation-decision-for-690) for the comparison and for what would change the decision. - **Polar splines and MHD equilibria 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. +## Hand-written CUDA vs. code generation (decision for #690) + +**Decision: the guiding-center and hybrid kernels (steps 3–5 of the [porting order](#porting-order)) are ported by +hand, as `_cuda.cu` next to the pyccel kernel, like every kernel so far.** None of the generators is adopted. +This is decided again if a condition under [What would change the decision](#what-would-change-the-decision) +comes true. + +To decide, one guiding-center kernel, `bstar_parallel_3form` (step 3, #7: drift-kinetic B*∥ as a 3-form; mapping +Jacobian, two 0-form spline evaluations, weighted average of old and new marker positions), was ported in two ways: + +- **By hand:** `pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu`, one to one with pyccel, with + 14 parity cases in `pic/tests/cuda_parity_cases.py`: one per mapping (all analytic and four spline mappings), + with α = 1, 0 and mixed, so the evaluation point wraps around in `mod(eta, 1)`, a negative `eta_n`, a hole + and two output columns. This port stays in the code. +- **Generated:** a prototype generator, `pic/tests/codegen_spike/numba_codegen.py`, that translates the pyccel + source of the kernel *and of every helper it calls* into one numba.cuda module (`compare.py` runs it, + `test_numba_codegen.py` tests it; skipped without numba). It only exists for this comparison and can be + deleted. There is no `kernels` in the module names, so `struphy compile` does not pick it up. + +### Results + +| | Hand-written CUDA | Generated (numba.cuda) | +|---|---|---| +| Code for this kernel | 75 lines `.cu` (36 without comments), plus one new 5-line helper `eval_0form_spline_mpi` in `bsplines/evaluation_kernels_3d.cuh`; everything else is reused (`df` for 14 mappings, `det`, `get_spans`, `eval_spline_mpi_kernel`) | 0 lines by hand, but a 458-line generator. The generated module has 490 lines and contains its own copy of the whole helper chain (24 device functions); every generated kernel repeats it | +| Agreement with pyccel (14 cases) | CPU emulation of the real `.cu` source: all agree, max abs. difference 5.7·10⁻¹⁴. Removing the `mod` from the CUDA kernel makes the emulation test fail (115 of 3225 values) | numba.cuda **simulator**: all agree, max 6.3·10⁻¹³; the same translation compiled with `numba.njit`: the same values. The simulator interprets the Python one thread at a time and compiles nothing | +| Compiles for a GPU? | Not checked here (no GPU); the same NVRTC path as the 15 kernels before it | Not checked (numba.cuda needs NVVM from the CUDA toolkit, which macOS does not have). Not checked either: named tuples of local arrays in device code, and the 38 flat kernel parameters | +| Argument classes | Takes `MarkerArgs`, `DomainArgs`, `DerhamArgs` by value, like every kernel; cunumpy's `CudaKernel` checks the call against the signature | numba does not accept the C structs. The generator turns each argument object into one kernel parameter per struct field (38 parameters) and rebuilds the objects as named tuples inside the kernel. Calling it needs a second launch path next to cunumpy's (flattening, numba's array conversion, its own compile cache) | +| 1:1 with pyccel | Names, arguments and order as in pyccel, kept by review | Exactly 1:1 for the parts the generator understands; whatever it does not understand fails with `NotImplementedError` (see the last row) | +| Parity tests, emulation | The existing tests run unchanged: `PARITY_CASES` entry, GPU parity with `assert_kernels_agree`, CPU emulation, which compiles the `.cu` itself | The parity cases are reused, but the simulator checks the Python, not the code the GPU will run. cunumpy's emulation and `assert_kernels_agree` do not apply to numba kernels | +| Debugging | A C++ compiler and printf in emulation; `-lineinfo`, cunumpy's debug mode, compute-sanitizer and Nsight on the GPU; the source you read is the source you run | pdb in the simulator, where Python float arithmetic hides bugs specific to the GPU. On the GPU, errors point into generated code. The generator's own mistakes are a new kind of bug (see below) | +| Expected performance (no GPU numbers possible here) | Per-thread arrays with fixed size (`MAX_SPLINE_DEGREE`), `__device__ inline` helpers; atomics where needed (`cunumpy_atomic_add`) | About the same local memory (the same `MAX_LOCAL` bounds; `b_der_splines_slim` has the same 9×9 `values`), plus loops that fill `zeros` arrays. numba applies Python's wrap-around for negative indices to every array access (one extra compare/select each), passes shape, strides and data pointer per array, and has a higher launch overhead than `RawKernel`. The generator cannot tell which writes need atomics: a translated accumulation kernel would race | +| Dependencies | None beyond CuPy | numba + llvmlite + numba-cuda (≈175 MB installed), and numba only supports a range of NumPy versions (numba 0.68 was needed for NumPy 2.4). The numba.cuda that comes with numba is deprecated in favour of NVIDIA's numba-cuda | +| The other 59 pusher and accumulation kernels | Each is a transcription onto existing helpers; new helpers as needed (e.g. 1-/2-form evaluations exist, `eval_0form_spline_mpi` was added here) | Tried on all 60 entry kernels (`.claude/scripts/codegen_probe.py`, not committed): 31 translate, 29 do not (2D slice assignments such as `c[:, :] = ...` in the linear-algebra helpers, `numpy.abs` in the SPH helpers). "Translates" does not mean "is correct": none of them ran, and the accumulation kernels would race | + +What the generator had to change in the pyccel code, and the risks that come with it: + +- `empty`/`zeros` with a size known only at run time (`zeros(pn)`) become per-thread arrays of `MAX_LOCAL` + entries. This silently changes `len(x)`/`x.shape` of these arrays, e.g. `diff[:] = diff * pn` in + `b_der_splines_slim` now loops over 9 entries instead of `pn` (harmless there, because the extra entries are + zero). +- Array expressions on slices (`eta[:] = alpha[:3] * eta_k + ...`) become loops over the elements, because numba.cuda + cannot allocate the temporaries. Expressions with 2D slices are not supported yet. +- `numpy.mod` has to be translated exactly as pyccel compiles it (Fortran `MODULO`, `a - floor(a / b) * b`). +- The scratch arrays of `DerhamArguments` (`bn1`, ..., `bd3`), which every marker shares in pyccel, have to become + per-thread arrays; the hand-written `SplineScratch` does the same. +- The marker loop becomes one thread per marker; `continue` becomes `return`. + +**The other generators.** `cupyx.jit` (CuPy's experimental Python-to-CUDA transpiler) would need the same +translation of the pyccel source: its reference lists elementary scalar operations, thread indices, shared memory +and atomics, but no per-thread arrays and no struct arguments, and "CuPy's array operations ... are not supported". +It adds no dependency, but it was not run here (no CuPy). **pyccel 2.2.1 has no CUDA backend** (`--language` is +Fortran, C or Python). The CUDA work happens in the separate fork [pyccel-cuda](https://github.com/pyccel/pyccel-cuda) +(`@kernel`/`@device`, last commit January 2025); its open issues still include thread indexing, passing arrays as +arguments and device arrays, and nothing there handles class arguments like `MarkerArguments`. + +### Reasoning + +1. **The shared part is already written.** What a generator would save is mostly the helper chain (mappings, + metric, splines, linear algebra, boundary conditions), and PRs 10–19 have already ported it by hand and tested + it, every mapping included. What is left per kernel is its body: here 36 lines that read like the pyccel code. +2. **A generator is a second compiler to maintain.** The prototype has 458 lines for a subset of what the kernels + use, and its mistakes (wrong array lengths, missing atomics) are silent. pyccel already is the compiler for the + CPU; a second one written by ourselves only pays off for many more kernels than ~45. +3. **It does not fit the infrastructure.** Everything since PR 8 (structs, signature checks in `CudaKernel`, + compile cache, debug mode, `assert_kernels_agree`, CPU emulation of the real source, no host transfers) is + built for CUDA C++ source. numba kernels would need a second launch and test path in cunumpy and struphy. +4. **It is harder to test.** The CPU emulation compiles the `.cu` that will run on the GPU; the numba simulator + only interprets Python. The hand-written port is gated better without a GPU than the generated one. +5. **Drift between pyccel and CUDA** (the main argument for one source) is caught by the parity cases: a change + to a pyccel kernel that is not made in its `.cu` fails `test_emulated_parity` in the normal CI. + +### What would change the decision + +- **pyccel gets a CUDA backend** in a release (not the fork) that accepts our argument classes and arrays. Then the + `.cu` files could be generated by pyccel in place, behind the same `Kernel`, structs and tests; the parity cases + would check them unchanged. This is the most attractive single-source route; check pyccel's releases before + step 4. +- **Porting by hand becomes the bottleneck**, e.g. new models add many kernels, or the pyccel kernels change often. + Then the next step is a generator that writes **`.cu` text** (C++) instead of numba: the AST transformations of + the prototype (helper calls, local arrays, slice loops, marker loop to threads, `mod`) carry over, and the output + goes through the existing structs, `CudaKernel`, emulation and parity tests. Accumulation kernels would still need + their atomics marked by hand. +- **GPU measurements** that show a generated kernel to be as fast as the hand-written one would remove the + performance argument; they do not touch points 2–4. + ## PR 10 implementation notes Device headers now provide Cuboid Jacobians, matrix inversion/vector products, diff --git a/src/struphy/bsplines/evaluation_kernels_3d.cuh b/src/struphy/bsplines/evaluation_kernels_3d.cuh index ccb10b6f7..3c4c353ad 100644 --- a/src/struphy/bsplines/evaluation_kernels_3d.cuh +++ b/src/struphy/bsplines/evaluation_kernels_3d.cuh @@ -71,6 +71,24 @@ __device__ inline double eval_spline_mpi_kernel(int p1, int p2, int p3, const do return spline_value; } +/** + * Evaluate a 0-form spline, as in evaluation_kernels_3d.eval_0form_spline_mpi. + * + * @param span1 Knot span index along the first axis (from get_spans). + * @param span2 Knot span index along the second axis. + * @param span3 Knot span index along the third axis. + * @param args_derham Spline degrees and start indices. + * @param scratch N-spline values from get_spans; pyccel reads them from args_derham.bn1, bn2, bn3. + * @param form_coeffs Coefficients of the 0-form (N N N). + * @return out, the value of the 0-form at the point of get_spans. + */ +__device__ inline double eval_0form_spline_mpi(int span1, int span2, int span3, const DerhamArgs& args_derham, + const SplineScratch& scratch, Array3D form_coeffs) { + double out = eval_spline_mpi_kernel(args_derham.pn[0], args_derham.pn[1], args_derham.pn[2], scratch.bn1, + scratch.bn2, scratch.bn3, span1, span2, span3, form_coeffs, args_derham.starts); + return out; +} + /** * Evaluate the three components of a 1-form spline, as in evaluation_kernels_3d.eval_1form_spline_mpi. * diff --git a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu new file mode 100644 index 000000000..ec7389284 --- /dev/null +++ b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu @@ -0,0 +1,75 @@ +#include "struphy/bsplines/evaluation_kernels_3d.cuh" +#include "struphy/geometry/evaluation_kernels.cuh" +#include "struphy/linear_algebra/linalg_kernels.cuh" + +/** + * Evaluate B*_parallel as a 3-form, as in bstar_parallel_3form_kernels.bstar_parallel_3form. + * + * B^{*3}_parallel(Z_p) = sqrt(g) (B . b + epsilon v_parallel,p ((curl b_0) . b_0)(eta_p)) at the weighted average + * Z_{p,i} = alpha_i Z^{n+1,k}_{p,i} + (1 - alpha_i) Z^n_{p,i} (i = 1, ..., 4); one thread per marker row. + * + * @param alpha Weights of the average (at least four entries): alpha[0:3] for eta, alpha[3] for v_parallel. + * @param output_indices Marker column of the result in output_indices[0]; -1 skips the output. + * @param args_markers Marker buffer (n_markers x n_cols, row-major) and column indices. + * @param args_domain Mapping arguments; every spline and analytic mapping. + * @param args_derham Spline degrees (1 to 8), knots and start indices. + * @param epsilon Scaling parameter epsilon in front of v_parallel. + * @param B_dot_b_coeffs Coefficients of the 0-form B . b. + * @param curl_unit_b_dot_b0 Coefficients of the 0-form (curl b_0) . b_0. + * + * Holes (markers[ip, 0] == -1) are skipped, as in pyccel (which does not use valid_mks here). + */ +extern "C" __global__ void bstar_parallel_3form(const double* alpha, const long long* output_indices, + MarkerArgs args_markers, DomainArgs args_domain, + DerhamArgs args_derham, double epsilon, + Array3D B_dot_b_coeffs, Array3D curl_unit_b_dot_b0) { + // CUDA-only: the marker row of this thread replaces the pyccel loop variable + int ip = blockDim.x * blockIdx.x + threadIdx.x; + int n_markers = args_markers.n_markers; + if (ip >= n_markers) return; + + // allocate stack arrays + double eta_k[3], eta_n[3], eta[3], dfm[9]; + + // get marker arguments + int first_init_idx = args_markers.first_init_idx; + int first_shift_idx = args_markers.first_shift_idx; + + // only do something if particle is a "true" particle (i.e. not a hole) + if (args_markers.markers(ip, 0) == -1.0) return; + + for (int j = 0; j < 3; ++j) { + eta_k[j] = args_markers.markers(ip, j) + args_markers.markers(ip, first_shift_idx + j); + eta_n[j] = args_markers.markers(ip, first_init_idx + j); + + eta[j] = alpha[j] * eta_k[j] + (1.0 - alpha[j]) * eta_n[j]; + // numpy.mod(eta, 1.0) (Fortran MODULO): the result has the sign of the divisor + eta[j] = eta[j] - floor(eta[j]); + } + + double v_k = args_markers.markers(ip, 3); + double v_n = args_markers.markers(ip, first_init_idx + 3); + double v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n; + + // evaluate Jacobian, result in dfm + struphy_cuda::df(eta[0], eta[1], eta[2], args_domain, dfm); + + double det_df = struphy_cuda::det(dfm); + + // spline evaluation; CUDA-only scratch holds the spline values pyccel keeps in args_derham + struphy_cuda::SplineScratch scratch; + struphy_cuda::get_spans(eta[0], eta[1], eta[2], args_derham, scratch); + int span1 = scratch.span1, span2 = scratch.span2, span3 = scratch.span3; + + // compute B*_parallel + double B_dot_b = struphy_cuda::eval_0form_spline_mpi(span1, span2, span3, args_derham, scratch, B_dot_b_coeffs); + + double b_star_parallel = + struphy_cuda::eval_0form_spline_mpi(span1, span2, span3, args_derham, scratch, curl_unit_b_dot_b0); + + b_star_parallel *= epsilon * v; + b_star_parallel += B_dot_b; + b_star_parallel *= det_df; + + if (output_indices[0] >= 0) args_markers.markers(ip, output_indices[0]) = b_star_parallel; +} diff --git a/src/struphy/pic/tests/codegen_spike/__init__.py b/src/struphy/pic/tests/codegen_spike/__init__.py new file mode 100644 index 000000000..d1b189c74 --- /dev/null +++ b/src/struphy/pic/tests/codegen_spike/__init__.py @@ -0,0 +1 @@ +"""Prototype for issue #690 (hand-written CUDA vs. code generation); not used by struphy, see ``numba_codegen``.""" diff --git a/src/struphy/pic/tests/codegen_spike/compare.py b/src/struphy/pic/tests/codegen_spike/compare.py new file mode 100644 index 000000000..ae72c1a6f --- /dev/null +++ b/src/struphy/pic/tests/codegen_spike/compare.py @@ -0,0 +1,61 @@ +"""Run the numba kernel generated from ``bstar_parallel_3form`` on its parity cases and compare it with pyccel. + +Usage (numba is not a struphy dependency):: + + NUMBA_ENABLE_CUDASIM=1 python -m struphy.pic.tests.codegen_spike.compare cuda # numba.cuda simulator + python -m struphy.pic.tests.codegen_spike.compare cpu # same translation, numba.njit + +The simulator interprets the generated Python thread by thread (it does not compile anything); ``cpu`` compiles the +same translation in nopython mode. Prints one line per case and a summary; exits with 1 if a case disagrees. +""" + +import sys +import time + +import cunumpy as xp +import numpy as np + +from struphy.pic.tests.codegen_spike.numba_codegen import flatten, load + +MODULE = "struphy.pic.pushing.kernels.bstar_parallel_3form.bstar_parallel_3form_kernels" +NAME = "bstar_parallel_3form" +BLOCK = 128 + + +def compare(target): + from struphy.pic.pushing.kernels.bstar_parallel_3form import bstar_parallel_3form + from struphy.pic.tests.cuda_parity_cases import PARITY_CASES + from struphy.pic.tests.test_cuda_emulation import arrays + + kernel, path = load(MODULE, NAME, target) + print(f"generated {path} ({len(path.read_text().splitlines())} lines)") + spec = PARITY_CASES[NAME] + failures, bitwise = 0, 0 + with xp.use_backend("numpy"): + for case in spec.cases: + reference_args, generated_args = spec.build(case), spec.build(case) + bstar_parallel_3form(*reference_args) + flat = flatten(generated_args) + start = time.perf_counter() + if target == "cuda": + n_markers = generated_args[2].n_markers + kernel[(n_markers + BLOCK - 1) // BLOCK, BLOCK](*flat) + else: + kernel(*flat) + seconds = time.perf_counter() - start + reference, generated = arrays(reference_args), arrays(generated_args) + diff = max( + float(np.max(np.abs(r.astype(float) - g.astype(float)), initial=0.0)) + for r, g in zip(reference, generated) + ) + same = all(np.array_equal(r, g, equal_nan=True) for r, g in zip(reference, generated)) + ok = all(np.allclose(g, r, rtol=spec.rtol, atol=spec.atol) for r, g in zip(reference, generated)) + failures += not ok + bitwise += same + print(f"domain {case[0]:2d}: max |diff| = {diff:.2e}, bitwise equal: {same}, {seconds:.2f} s") + print(f"{target}: {len(spec.cases) - failures}/{len(spec.cases)} cases agree, {bitwise} bitwise") + return failures + + +if __name__ == "__main__": + sys.exit(1 if compare(sys.argv[1] if len(sys.argv) > 1 else "cpu") else 0) diff --git a/src/struphy/pic/tests/codegen_spike/numba_codegen.py b/src/struphy/pic/tests/codegen_spike/numba_codegen.py new file mode 100644 index 000000000..062df877d --- /dev/null +++ b/src/struphy/pic/tests/codegen_spike/numba_codegen.py @@ -0,0 +1,458 @@ +"""Prototype for issue #690: generate a ``numba.cuda`` kernel from the Python source of a pyccel kernel. + +Not used by struphy. It exists to compare code generation with the hand-written ``_cuda.cu`` files on one +real kernel (``bstar_parallel_3form``); the result of the comparison is in ``CUDA_STRATEGY.md`` (decision for +#690). The module name contains no ``kernels``, so ``struphy compile`` ignores it. + +The generator works on the source files (the imported kernel modules are usually the compiled pyccel extensions): + +- it parses the entry kernel's module, follows every call into struphy modules (``import a.b as x`` and + ``from a.b import f``) and translates each reached function once, as a numba device function named + ``__``; +- it strips the pyccel decorators and type annotations, maps the NumPy scalar functions to ``math``, and + ``numpy.mod`` to the Fortran ``MODULO`` pyccel generates; +- ``empty``/``zeros`` become ``cuda.local.array`` with a constant size (a run-time size such as ``zeros(pn)`` gets + ``MAX_LOCAL`` entries, like the hand-written ``MAX_SPLINE_DEGREE`` arrays); +- array expressions on slices (``eta[:] = alpha[:3] * eta_k + ...``, ``bn[:] = 0.0``) become element loops, since + numba.cuda cannot allocate the temporaries; +- the entry kernel's marker loop ``for ip in range(n): ...`` becomes ``ip = cuda.grid(1)`` and ``continue`` becomes + ``return``; +- argument objects (``MarkerArguments``, ``DomainArguments``, ``DerhamArguments``) are passed as the fields of + their CUDA structs (``CudaMarkerArguments.struct.fields``, one kernel parameter each) and rebuilt as named tuples + at the top of the kernel; ``DerhamArguments`` gets its scratch arrays ``bn1``, ..., ``bd3`` as per-thread local + arrays there, so the helpers can stay one to one with pyccel. + +With ``target="cpu"`` the same translation is compiled with ``numba.njit`` (NumPy allocations, the marker loop kept), +which checks that the generated code types in nopython mode; numba.cuda itself needs a GPU (or NVVM) to compile, and +its simulator (``NUMBA_ENABLE_CUDASIM=1``) only interprets the Python. +""" + +import ast +import importlib.util +import re +import sys +import tempfile +from pathlib import Path + +from struphy.kernel_arguments.pusher_args_cuda import CudaDerhamArguments, CudaDomainArguments, CudaMarkerArguments + +MAX_LOCAL = 9 # MAX_SPLINE_DEGREE + 1 of bsplines_kernels.cuh + +STRUCT_CLASSES = { + "MarkerArguments": CudaMarkerArguments, + "DerhamArguments": CudaDerhamArguments, + "DomainArguments": CudaDomainArguments, +} +# pyccel scratch arrays of the argument classes: per-thread local arrays in the generated kernel +SCRATCH = {"DerhamArguments": ("bn1", "bn2", "bn3", "bd1", "bd2", "bd3")} + +NUMPY_MATH = { + "cos": "cos", + "sin": "sin", + "tan": "tan", + "arctan": "atan", + "arcsin": "asin", + "arccos": "acos", + "arctan2": "atan2", + "sqrt": "sqrt", + "exp": "exp", + "log": "log", + "floor": "floor", +} +NUMPY_ALLOC = ("empty", "zeros") + + +def _source_file(module): + """The ``.py`` file of `module` (the import system would find the compiled pyccel extension first).""" + spec = importlib.util.find_spec(module) + return Path(spec.origin).parent / (module.rsplit(".", 1)[-1] + ".py") + + +class _Module: + """Functions and imports of one source module.""" + + def __init__(self, name): + self.name = name + self.short = name.rsplit(".", 1)[-1] + tree = ast.parse(_source_file(name).read_text()) + self.functions = {n.name: n for n in tree.body if isinstance(n, ast.FunctionDef)} + self.module_aliases = {} # alias -> module + self.imported = {} # name -> (module, name) + for node in tree.body: + if isinstance(node, ast.Import): + for a in node.names: + self.module_aliases[a.asname or a.name] = a.name + elif isinstance(node, ast.ImportFrom): + for a in node.names: + self.imported[a.asname or a.name] = (node.module, a.name) + + def resolve(self, func): + """``(module, name)`` of a called function node, or ``None``.""" + if isinstance(func, ast.Name): + if func.id in self.functions: + return self.name, func.id + return self.imported.get(func.id) + if isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name): + module = self.module_aliases.get(func.value.id) + if module is not None: + return module, func.attr + return None + + +def _array_ndim(annotation): + """Number of dimensions of a pyccel array annotation such as ``"float[:,:]"``, or 0.""" + if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str): + match = re.search(r"\[([:,\s]+)\]", annotation.value) + if match: + return match.group(1).count(":") + return 0 + + +def _struct_type(annotation): + if isinstance(annotation, ast.Constant) and annotation.value in STRUCT_CLASSES: + return annotation.value + if isinstance(annotation, ast.Name) and annotation.id in STRUCT_CLASSES: + return annotation.id + return None + + +def _expr(source): + return ast.parse(source, mode="eval").body + + +def _stmt(source): + return ast.parse(source).body + + +class _Elementwise(ast.NodeTransformer): + """Index a slice expression at element ``index``: ``a[lo:hi]`` -> ``a[lo + index]``, 1D array ``x`` -> ``x[index]``.""" + + def __init__(self, index, arrays): + self.index = index + self.arrays = arrays + + def _element(self, slc): + lower = slc.lower if slc.lower is not None else ast.Constant(0) + return ast.BinOp(lower, ast.Add(), ast.Name(self.index, ast.Load())) + + def visit_Subscript(self, node): + if isinstance(node.slice, ast.Slice): + return ast.Subscript(node.value, self._element(node.slice), node.ctx) + if isinstance(node.slice, ast.Tuple) and any(isinstance(e, ast.Slice) for e in node.slice.elts): + elts = [self._element(e) if isinstance(e, ast.Slice) else e for e in node.slice.elts] + return ast.Subscript(node.value, ast.Tuple(elts, ast.Load()), node.ctx) + return node + + def visit_Name(self, node): + if self.arrays.get(node.id) == 1: + return ast.Subscript(ast.Name(node.id, ast.Load()), ast.Name(self.index, ast.Load()), node.ctx) + if node.id in self.arrays and self.arrays[node.id] > 1: + raise NotImplementedError(f"array expression on the {self.arrays[node.id]}D array {node.id}") + return node + + +def _slice_length(target): + """Number of elements written by a sliced assignment target (one slice per target).""" + slices = target.slice.elts if isinstance(target.slice, ast.Tuple) else [target.slice] + positions = [i for i, s in enumerate(slices) if isinstance(s, ast.Slice)] + if len(positions) != 1: + raise NotImplementedError(f"assignment to {ast.unparse(target)}: exactly one slice is supported") + slc = slices[positions[0]] + if slc.step is not None: + raise NotImplementedError(f"assignment to {ast.unparse(target)}: slice steps are not supported") + if slc.upper is None: + upper = ast.Subscript(ast.Attribute(target.value, "shape", ast.Load()), ast.Constant(positions[0]), ast.Load()) + else: + upper = slc.upper + if slc.lower is None: + return upper + return ast.BinOp(upper, ast.Sub(), slc.lower) + + +def _has_slice(target): + if not isinstance(target, ast.Subscript): + return False + s = target.slice + return isinstance(s, ast.Slice) or (isinstance(s, ast.Tuple) and any(isinstance(e, ast.Slice) for e in s.elts)) + + +class _ContinueToReturn(ast.NodeTransformer): + """``continue`` of the marker loop -> ``return`` (nested loops keep theirs).""" + + def visit_Continue(self, node): + return ast.Return(None) + + def visit_For(self, node): + return node + + def visit_While(self, node): + return node + + +class Generator: + """Translate a pyccel entry kernel and every struphy function it calls into one numba module.""" + + def __init__(self, target="cuda"): + assert target in ("cuda", "cpu") + self.target = target + self.modules = {} + self.translated = {} # (module, name) -> mangled name + self.sources = [] # translated functions, callees first + self.counter = 0 + + def _module(self, name): + if name not in self.modules: + self.modules[name] = _Module(name) + return self.modules[name] + + def _is_struphy_function(self, ref): + if ref is None or not ref[0].startswith("struphy."): + return False + try: + return ref[1] in self._module(ref[0]).functions + except (ModuleNotFoundError, FileNotFoundError, AttributeError): + return False + + def _numpy_name(self, module, node): + """The NumPy name a call or name node refers to in `module`, or ``None``.""" + if isinstance(node, ast.Name): + ref = module.imported.get(node.id) + if ref is not None and ref[0] == "numpy": + return ref[1] + return None + + def _fresh(self, prefix): + self.counter += 1 + return f"{prefix}{self.counter}" + + # ------------------------------------------------------------------ allocations + + def _allocation(self, module, target, call, arrays): + kind = self._numpy_name(module, call.func) + shape = call.args[0] + ndim = len(shape.elts) if isinstance(shape, ast.Tuple) else 1 + dtype = "int64" + for kw in call.keywords: + if kw.arg == "dtype" and isinstance(kw.value, ast.Name) and kw.value.id == "float": + dtype = "float64" + if len(call.args) > 1 and isinstance(call.args[1], ast.Name) and call.args[1].id == "float": + dtype = "float64" + if not call.keywords and len(call.args) == 1: + dtype = "float64" + arrays[target] = ndim + if self.target == "cpu": + return _stmt(f"{target} = np.{kind}({ast.unparse(shape)}, dtype=np.{dtype})") + elts = shape.elts if isinstance(shape, ast.Tuple) else [shape] + sizes = [e.value if isinstance(e, ast.Constant) else MAX_LOCAL for e in elts] + size = sizes[0] if ndim == 1 else "(" + ", ".join(map(str, sizes)) + ")" + stmts = _stmt(f"{target} = cuda.local.array({size}, {dtype})") + if kind == "zeros": + loops = "".join(f"{' ' * d}for _z{d} in range({s}):\n" for d, s in enumerate(sizes)) + index = ", ".join(f"_z{d}" for d in range(ndim)) + stmts += _stmt(f"{loops}{' ' * ndim}{target}[{index}] = 0") + return stmts + + # ------------------------------------------------------------------ functions + + def translate(self, module_name, name, entry=False): + """Translate function `name` of `module_name`; returns its name in the generated module.""" + key = (module_name, name) + if key in self.translated: + return self.translated[key] + module = self._module(module_name) + mangled = f"{module.short}__{name}" + self.translated[key] = mangled + node = ast.parse(ast.unparse(module.functions[name])).body[0] # a fresh copy + + # parameters: array dimensions and argument objects from the pyccel annotations + arrays = {a.arg: _array_ndim(a.annotation) for a in node.args.args if _array_ndim(a.annotation)} + structs = {a.arg: _struct_type(a.annotation) for a in node.args.args if _struct_type(a.annotation)} + for a in node.args.args: + a.annotation = None + node.returns = None + node.decorator_list = [] + node.name = mangled + if node.body and isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant): + node.body = node.body[1:] or [ast.Pass()] + + node.body = self._statements(module, node.body, arrays) + + if entry: + self._make_entry(node, structs) + decorator = "@cuda.jit" if self.target == "cuda" else "@njit" + else: + decorator = "@cuda.jit(device=True)" if self.target == "cuda" else "@njit" + self.sources.append(f"{decorator}\n{ast.unparse(ast.fix_missing_locations(node))}\n") + return mangled + + def _statements(self, module, body, arrays): + out = [] + for stmt in body: + out += self._statement(module, stmt, arrays) + return out + + def _statement(self, module, stmt, arrays): + # local arrays + if ( + isinstance(stmt, ast.Assign) + and isinstance(stmt.value, ast.Call) + and self._numpy_name(module, stmt.value.func) in NUMPY_ALLOC + ): + (target,) = stmt.targets + return self._allocation(module, target.id, stmt.value, arrays) + # array expressions on slices -> element loops + if isinstance(stmt, (ast.Assign, ast.AugAssign)): + target = stmt.targets[0] if isinstance(stmt, ast.Assign) else stmt.target + if _has_slice(target): + index = self._fresh("_k") + length = _slice_length(target) + element = _Elementwise(index, arrays) + value = self._expressions(module, element.visit(stmt.value)) + new_target = element.visit(target) + if isinstance(stmt, ast.Assign): + inner = ast.Assign([new_target], value) + else: + inner = ast.AugAssign(new_target, stmt.op, value) + loop = ast.For( + ast.Name(index, ast.Store()), + ast.Call(ast.Name("range", ast.Load()), [length], []), + [inner], + [], + lineno=stmt.lineno, + ) + return [loop] + # compound statements: recurse into their bodies + for field in ("body", "orelse"): + if hasattr(stmt, field) and isinstance(getattr(stmt, field), list): + setattr(stmt, field, self._statements(module, getattr(stmt, field), arrays)) + for field in ("value", "test", "iter"): + if hasattr(stmt, field) and isinstance(getattr(stmt, field), ast.AST): + setattr(stmt, field, self._expressions(module, getattr(stmt, field))) + if isinstance(stmt, ast.Assign): + stmt.targets = [self._expressions(module, t) for t in stmt.targets] + if isinstance(stmt, ast.AugAssign): + stmt.target = self._expressions(module, stmt.target) + return [stmt] + + def _expressions(self, module, node): + generator = self + + class Calls(ast.NodeTransformer): + def visit_Call(self, call): + self.generic_visit(call) + ref = module.resolve(call.func) + numpy_name = generator._numpy_name(module, call.func) + if generator._is_struphy_function(ref): + call.func = ast.Name(generator.translate(*ref), ast.Load()) + elif numpy_name in NUMPY_MATH: + call.func = _expr(f"math.{NUMPY_MATH[numpy_name]}") + elif numpy_name == "mod": + call.func = ast.Name("_numpy_mod", ast.Load()) + elif numpy_name == "shape": + return ast.Attribute(call.args[0], "shape", ast.Load()) + elif numpy_name is not None: + raise NotImplementedError(f"numpy.{numpy_name} in {module.name}") + return call + + def visit_Name(self, name): + if generator._numpy_name(module, name) == "pi": + return _expr("math.pi") + return name + + return Calls().visit(node) + + def _make_entry(self, node, structs): + """Flatten the argument objects into struct fields, rebuild them, and map the marker loop to threads.""" + params, prologue = [], [] + for a in node.args.args: + kind = structs.get(a.arg) + if kind is None: + params.append(a.arg) + continue + fields = [f.name for f in STRUCT_CLASSES[kind].struct.fields] + params += [f"{a.arg}__{f}" for f in fields] + values = [f"{a.arg}__{f}" for f in fields] + if self.target == "cuda": + values += [f"cuda.local.array({MAX_LOCAL}, float64)" for _ in SCRATCH.get(kind, ())] + else: + values += [f"np.empty({MAX_LOCAL})" for _ in SCRATCH.get(kind, ())] + prologue += _stmt(f"{a.arg} = {kind}_T({', '.join(values)})") + node.args.args = [ast.arg(p) for p in params] + + body = prologue + node.body + if self.target == "cuda": + for i, stmt in enumerate(body): + if ( + isinstance(stmt, ast.For) + and isinstance(stmt.iter, ast.Call) + and getattr(stmt.iter.func, "id", None) == "range" + and len(stmt.iter.args) == 1 + ): + ip, n = stmt.target.id, ast.unparse(stmt.iter.args[0]) + loop_body = [_ContinueToReturn().visit(s) for s in stmt.body] + body = body[:i] + _stmt(f"{ip} = cuda.grid(1)\nif {ip} >= {n}:\n return") + loop_body + break + else: + raise NotImplementedError("no marker loop 'for ip in range(n)' in the entry kernel") + node.body = body + + def module_source(self): + decorator = "@cuda.jit(device=True)" if self.target == "cuda" else "@njit" + header = [ + f"# Generated by struphy.pic.tests.codegen_spike.numba_codegen (target={self.target}); do not edit.", + "import math", + "from collections import namedtuple", + "", + "import numpy as np", + "from numba import cuda, float64, int64, njit", + "", + ] + for kind, cls in STRUCT_CLASSES.items(): + names = [f.name for f in cls.struct.fields] + list(SCRATCH.get(kind, ())) + header.append(f"{kind}_T = namedtuple({kind + '_T'!r}, {names!r})") + header += [ + "", + "", + decorator, + "def _numpy_mod(a, b):", + " # numpy.mod as pyccel compiles it (Fortran MODULO)", + " return a - math.floor(a / b) * b", + "", + "", + ] + return "\n".join(header) + "\n\n".join(self.sources) + + +def generate(module_name, name, target="cuda"): + """Source of the numba module for the pyccel kernel `name` in `module_name` (and every function it calls).""" + generator = Generator(target) + generator.translate(module_name, name, entry=True) + return generator.module_source() + + +def load(module_name, name, target="cuda", directory=None): + """Generate, write to `directory` (a temporary one by default) and import; returns the kernel function.""" + source = generate(module_name, name, target) + directory = Path(directory or tempfile.mkdtemp(prefix="struphy_numba_codegen_")) + stem = f"generated_{name}_{target}" + path = directory / f"{stem}.py" + path.write_text(source) + spec = importlib.util.spec_from_file_location(stem, path) + module = importlib.util.module_from_spec(spec) + sys.modules[stem] = module + spec.loader.exec_module(module) + return getattr(module, f"{module_name.rsplit('.', 1)[-1]}__{name}"), path + + +def flatten(args): + """Kernel arguments with every argument object replaced by the fields of its CUDA struct (host arrays).""" + from struphy.pic.tests.cuda_emulation import _struct_field + + flat = [] + for value in args: + kind = type(value).__name__ + if kind in STRUCT_CLASSES: + flat += [_struct_field(value, f.name) for f in STRUCT_CLASSES[kind].struct.fields] + else: + flat.append(value) + return flat diff --git a/src/struphy/pic/tests/codegen_spike/test_numba_codegen.py b/src/struphy/pic/tests/codegen_spike/test_numba_codegen.py new file mode 100644 index 000000000..7648440b6 --- /dev/null +++ b/src/struphy/pic/tests/codegen_spike/test_numba_codegen.py @@ -0,0 +1,29 @@ +"""The numba kernel generated from ``bstar_parallel_3form`` agrees with pyccel (issue #690 prototype). + +Skipped unless numba (with numba-cuda for the simulator) is installed; numba is not a struphy dependency. +Runs in a subprocess because ``NUMBA_ENABLE_CUDASIM`` must be set before numba is imported. +""" + +import importlib.util +import subprocess +import sys + +import pytest + +from struphy.geometry.tests.test_domain import serial_child_env + +requires_numba = pytest.mark.skipif(importlib.util.find_spec("numba") is None, reason="numba is not installed") + + +@requires_numba +@pytest.mark.parametrize("target", ["cpu", "cuda"]) +def test_generated_kernel_agrees_with_pyccel(target): + """``cpu``: the translation compiled with numba.njit; ``cuda``: the numba.cuda kernel in the CUDA simulator.""" + extra = {"NUMBA_ENABLE_CUDASIM": "1"} if target == "cuda" else {} + result = subprocess.run( + [sys.executable, "-m", "struphy.pic.tests.codegen_spike.compare", target], + env=serial_child_env(**extra), + capture_output=True, + text=True, + ) + assert result.returncode == 0, result.stdout + result.stderr diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index e2ce0a5d5..7ee91c0c0 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -71,6 +71,31 @@ def reflect_args(axis): return (xp.asarray(markers), Cuboid(r1=2.0, r2=3.0, r3=4.0).args_domain, xp.asarray(outside_inds), axis) +# guiding-centre evaluation kernels: every mapping, the evaluation point at the new position (alpha = 1), the old one +# (alpha = 0) and in between; the output column alternates. (index into geometry_domain(), alpha, output column) +BSTAR_PARALLEL_ALPHAS = ((1.0,) * 6, (0.0,) * 6, (0.3, 0.6, 0.9, 0.5, 0.0, 0.0)) +BSTAR_PARALLEL_3FORM_CASES = tuple( + (domain, BSTAR_PARALLEL_ALPHAS[domain % 3], 19 if domain % 2 == 0 else 24) for domain in range(N_GEOMETRY_DOMAINS) +) + + +def bstar_parallel_3form_args(case): + """Markers with shifts in [0, 1), so eta + shift wraps around, a negative eta_n in row 0 and a hole in row 2.""" + domain, alpha, output_column = case + args_markers, _ = marker_arguments((0, 0, 0)) + args_markers.markers[0, 8] = -0.75 # eta_n[0] = mod(-0.75, 1.0) = 0.25 (not 0, where some Jacobians are singular) + args_markers.markers[2, 0] = -1.0 # a hole: pyccel tests markers[ip, 0] == -1, not valid_mks + return ( + xp.asarray(alpha), + xp.asarray([output_column], dtype=np.int64), + args_markers, + geometry_domain(domain).args_domain, + derham_arguments(), + 0.37, + *spline_coefficients(n=2, seed=23), + ) + + # ---------------------------------------------------------------- accumulation @@ -265,6 +290,13 @@ def size_of(index, per_thread=1): BOUNDARY_CONDITIONS, push_weights_with_efield_lin_va_args, **PUSHER_TOLERANCES ), "reflect": ParityCases((0, 1, 2), reflect_args, n_threads=size_of(2)), + # one thread per marker row; the first array, alpha, does not tell the launch size + "bstar_parallel_3form": ParityCases( + BSTAR_PARALLEL_3FORM_CASES, + bstar_parallel_3form_args, + n_threads=lambda args: args[2].n_markers, + **GEOMETRY_TOLERANCES, + ), # atomics add in another order than the serial loop "charge_density_0form": ParityCases(((0, 0, 0), (2, 0, 1)), charge_density_0form_args, rtol=1e-12, atol=1e-13), "eval_spline_mpi_markers": ParityCases(SPLINE_KINDS, eval_spline_mpi_markers_args), From c55990f03c49888705c403b6840f97a7bf34f60a Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 22:50:46 +0200 Subject: [PATCH 09/50] Polar splines on the CuPy backend Build the polar blocks (PolarExtractionBlocksC1) on the host from a host copy of the control points on every backend, and apply device copies of them (cupyx.scipy.sparse via xp.scipy.sparse, cached per block list) in PolarExtractionOperator/PolarLinearOperator when the CuPy backend is active. Synchronize device buffers before the ring Allreduces. Remove the NotImplementedError for polar_splines=True on CuPy. The restart left inverse of SplineFunction is built on the host. New test_polar_cupy.py compares NumPy and CuPy (fake CuPy without a GPU) for IGAPolarCylinder and IGAPolarTorus. Solves #695, part of #650. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 37 ++- src/struphy/feec/psydac_derham.py | 15 +- src/struphy/feec/tests/test_derham_gpu.py | 17 -- src/struphy/polar/basic.py | 3 + src/struphy/polar/extraction_operators.py | 165 ++++++------- src/struphy/polar/linear_operators.py | 145 ++++++++++-- src/struphy/polar/tests/test_polar_cupy.py | 255 +++++++++++++++++++++ 7 files changed, 508 insertions(+), 129 deletions(-) create mode 100644 src/struphy/polar/tests/test_polar_cupy.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index daf0d25f0..d0e612a02 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -145,7 +145,7 @@ from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage # a Kerne - The CUDA argument objects hold references. If an owner reallocates an array, it must rebuild its arguments at the same place, exactly like for the pyccel arguments. - `Domain` (PR 5): arrays that already have the dtype and layout the kernels expect (`float64`/`int64`, C-contiguous) are referenced, not copied; otherwise one device copy is made when the arguments are built (e.g. `degree`, a tuple). Arguments are recreated after deepcopy or unpickling. Spline mappings (e.g. `IGAPolarCylinder`) cannot be created on CuPy yet (`interp_mapping` passes CuPy arrays to `scipy.sparse`), so CUDA `args_domain` exists only for analytic mappings for now. - `Particles` (PR 6): particle arrays, validity masks and boundary-condition codes are allocated through `xp`. Scalar MPI gathers use small NumPy buffers (the own entry filled first, see `MockComm` above) and copy the result back to the active backend. -- `Derham` (PR 7): depends on feectools running on the CuPy backend (struphy-hub/feectools#85, merged). Data that describes the spline spaces (knots, quadrature and projection grids, `spline_types_pyccel`) is host data on every backend; only coefficients and stencil matrices live on the device. `domain_array`, `index_array(_N/_D)` and `neighbours` are gathered with NumPy MPI buffers and converted with `xp.asarray`. `args_derham` is built from the host knots, degrees and starts on NumPy, and from one device copy of them on CuPy. The pyccel scratch arrays (`bn1`, ..., `bd3`) are not part of the CUDA arguments; they become per-thread local arrays in CUDA (PR 10). Not supported on CuPy yet: local projectors (`NotImplementedError` at construction), polar splines (`NotImplementedError` at construction since PR 19, see the [PR 19 implementation notes](#pr-19-implementation-notes)), and field evaluation (`SplineFunction.__call__` still calls pyccel kernels with device coefficients). +- `Derham` (PR 7): depends on feectools running on the CuPy backend (struphy-hub/feectools#85, merged). Data that describes the spline spaces (knots, quadrature and projection grids, `spline_types_pyccel`) is host data on every backend; only coefficients and stencil matrices live on the device. `domain_array`, `index_array(_N/_D)` and `neighbours` are gathered with NumPy MPI buffers and converted with `xp.asarray`. `args_derham` is built from the host knots, degrees and starts on NumPy, and from one device copy of them on CuPy. The pyccel scratch arrays (`bn1`, ..., `bd3`) are not part of the CUDA arguments; they become per-thread local arrays in CUDA (PR 10). Polar splines run on CuPy since #695 (see [Polar splines on CuPy](#polar-splines-on-cupy-695)). Not supported on CuPy yet: local projectors (`NotImplementedError` at construction) and field evaluation (`SplineFunction.__call__` still calls pyccel kernels with device coefficients). ### PR 8: Argument structs in a shared header (complete, open) @@ -207,7 +207,7 @@ For each kernel: add `_cuda.cu`; the parity test from PR 11 picks it up; ` ### PR 19: Spline mappings - `DomainArgs`: `t1..t3` become `Array1D` (`find_span` needs the number of knots), `ind1..3` `Array2D` and `cx/cy/cz` `Array3D` (the device needs their shapes); `CudaDomainArguments.fields` and the generated `pusher_args.cuh` change, the pyccel `DomainArguments` does not. -- Spline-mapped domains (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC) can be created on the CuPy backend: the control-point fitting (SciPy) runs on the host and the control points are copied to the device once. Polar splines (`Derham` with `polar_splines=True`) stay host-only and raise on CuPy. +- Spline-mapped domains (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC) can be created on the CuPy backend: the control-point fitting (SciPy) runs on the host and the control points are copied to the device once. Polar splines (`Derham` with `polar_splines=True`) stayed host-only and raised on CuPy (on CuPy since #695). - Device versions of `spline_3d(_df)`, `spline_2d_straight(_df)`, `spline_2d_torus(_df)` in `geometry/spline_mappings_kernels.cuh`, using the B-spline helpers of PR 10; the parity arguments of the geometry kernels add the spline mappings. ## Moving to cunumpy (PR 13) @@ -347,7 +347,7 @@ 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 +accepts every mapping since PR 19), except for what is still host-only around them: MHD equilibria such as `EQDSKequilibrium` cannot be created on CuPy (see [Open questions](#open-questions)); multi-rank marker sorting without host round trips, and array views with more than 4 dimensions in cunumpy (all matrix accumulations write 6D stencil matrix data). @@ -367,10 +367,8 @@ matrix data). - **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. `xp.mpi_is_cuda_aware()` detects it; the exchange in `Particles.mpi_sort_markers` has to be checked for host staging buffers. - **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 - `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 +- **MHD equilibria on the GPU** (left after PR 19; polar splines run on CuPy since #695): spline mappings run on + the device, but `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. ## PR 10 implementation notes @@ -600,8 +598,8 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a - **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 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)). + the device on CuPy; `Derham` raised `NotImplementedError` for `polar_splines=True` on CuPy (resolved in #695, see + [Polar splines on CuPy](#polar-splines-on-cupy-695)). - **Device helpers.** `b_splines_slim`, `b_der_splines_slim` (`bsplines/bsplines_kernels.cuh`), `evaluation_kernel_2d` (`bsplines/evaluation_kernels_2d.cuh`), `evaluation_kernel_3d` (`bsplines/evaluation_kernels_3d.cuh`) and the six spline mappings in `geometry/spline_mappings_kernels.cuh`, with the pyccel names and arguments (`const DomainArgs&` @@ -626,3 +624,24 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a The GPU tests have not run on an H100; the CPU emulation and the CPU regression tests are the gate. - **Parity cases.** The spline mappings are added to the geometry kernels' entries of `pic/tests/cuda_parity_cases.py` (the per-folder `_test_args.py` modules were replaced by that module in PR 16). + +## Polar splines on CuPy (#695) + +- **Blocks on the host.** `PolarExtractionBlocksC1` builds every block (extraction, DOF extraction, polar + derivatives) with NumPy/SciPy from a host copy of the control points (`xp.to_numpy(domain.cx)`), on every backend: + they are setup data, and `.toarray()`, `.T` and the restart left inverse (`SplineFunction._restart_extraction_op`) + keep working on them. +- **Products on the device.** `PolarExtractionOperator` and `PolarLinearOperator` keep their `blocks_*` as host + matrices; on CuPy, `dot` uses device copies (`DeviceSparseMatrix`, a `cupyx.scipy.sparse.csr_matrix` via + `xp.scipy.sparse`), made once per block list and dropped by the setters. The blocks act only on the first two or + three radial rings (at most `3 n2 x 3 n2` in `eta1-eta2`, `n3 x n3` in `eta3`), so they are small; blocks without + non-zeros are not copied. On NumPy nothing changes (the SciPy blocks are applied as before). +- **MPI.** The ring reductions (`dot_inner_tp_rings`, `dot_parts_of_polar`, `PolarVector.toarray`) call + `cunumpy.mpi.synchronize_for_mpi` before `Allreduce`, which gets device buffers on CuPy (CUDA-aware MPI). +- **Tests.** `polar/tests/test_polar_cupy.py` compares NumPy and CuPy for `IGAPolarCylinder` and `IGAPolarTorus`: + extraction and DOF extraction operators and their transposes, `grad`/`curl`/`div` and transposes, `PolarVector` + arithmetic, the polar projectors, polar mass matrices (GPU only) and polar `SplineFunction` coefficients + (`extract_coeffs`, restart). Without a GPU it runs on the fake CuPy with two substitutes: a dense stand-in for + `cupyx.scipy.sparse` (`set_device_sparse_module`) and the host versions of the feectools kernels (the fake cannot + launch CUDA kernels). The mass matrices need struphy's CUDA geometry kernels and are compared on the GPU only. + Field evaluation (`SplineFunction.__call__`) is not covered: it is not ported to CuPy for any space yet. diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index a15ac1591..ea51d5e8c 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -614,12 +614,6 @@ def __init__( raise NotImplementedError( "Local projectors (DerhamOptions.local_projectors=True) are not supported on the CuPy backend yet." ) - if polar_splines and (xp.get_backend() == "cupy"): - # 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. - raise NotImplementedError( - "Polar splines (DerhamOptions.polar_splines=True) are not supported on the CuPy backend yet." - ) # number of elements and spline degrees in each direction assert len(num_elements) == 3 @@ -2712,16 +2706,17 @@ def _restart_extraction_op(self): n_comps = W.n_comps # stack blocks of E (polar coeffs x polar rings); incompatible blocks (None) are zero - rows = xp.cumsum([0] + list(W.n_polar)) - cols = xp.cumsum([0] + [n_r * n_2 for n_r, n_2 in zip(W.n_rings, W.n2)]) - E_full = xp.zeros((rows[-1], cols[-1]), dtype=float) + # (host setup data on every backend, like the blocks of E) + rows = np.cumsum([0] + list(W.n_polar)) + cols = np.cumsum([0] + [n_r * n_2 for n_r, n_2 in zip(W.n_rings, W.n2)]) + E_full = np.zeros((rows[-1], cols[-1]), dtype=float) for m in range(n_comps): for n in range(n_comps): if E.blocks_ten_to_pol[m][n] is not None: E_full[rows[m] : rows[m + 1], cols[n] : cols[n + 1]] = E.blocks_ten_to_pol[m][n].toarray() # E^T has full column rank, hence pinv(E^T) @ E^T = identity on polar coeffs - L_full = xp.linalg.pinv(E_full.T) + L_full = np.linalg.pinv(E_full.T) blocks = [ [ diff --git a/src/struphy/feec/tests/test_derham_gpu.py b/src/struphy/feec/tests/test_derham_gpu.py index 64c55a0d8..d8f3df46a 100644 --- a/src/struphy/feec/tests/test_derham_gpu.py +++ b/src/struphy/feec/tests/test_derham_gpu.py @@ -67,20 +67,3 @@ def test_local_projectors_not_supported_on_cupy(): """Local projectors have no device implementation yet; the Derham fails when it is created.""" with cunumpy.use_backend("cupy"), pytest.raises(NotImplementedError, match="Local projectors"): make_derham(local_projectors=True) - - -@requires_cupy -def test_polar_splines_not_supported_on_cupy(): - """Polar splines have no device implementation yet; the Derham fails when it is created.""" - from maybempi import MPI - - from struphy import domains - from struphy.feec.psydac_derham import Derham - from struphy.io.options import DerhamOptions - from struphy.topology.grids import TensorProductGrid - - with cunumpy.use_backend("cupy"): - domain = domains.IGAPolarCylinder(num_elements=(8, 6), degree=(2, 3)) - options = DerhamOptions(degree=(2, 3, 1), bcs=(("dirichlet", "free"), None, None), polar_splines=True) - with pytest.raises(NotImplementedError, match="Polar splines"): - Derham(TensorProductGrid(num_elements=(8, 6, 4)), options, comm=MPI.COMM_WORLD, domain=domain) diff --git a/src/struphy/polar/basic.py b/src/struphy/polar/basic.py index 23398555c..e181f542d 100644 --- a/src/struphy/polar/basic.py +++ b/src/struphy/polar/basic.py @@ -1,6 +1,7 @@ from typing import TYPE_CHECKING import cunumpy as xp +from cunumpy.mpi import synchronize_for_mpi from feectools.linalg.basic import Vector, VectorSpace from feectools.linalg.block import BlockVector from feectools.linalg.stencil import StencilVector @@ -329,6 +330,7 @@ def toarray(self, allreduce=False): # allreduce tensor-product part if self.space.comm is not None and allreduce: + synchronize_for_mpi(out) self.space.comm.Allreduce(MPI.IN_PLACE, out, op=MPI.SUM) out = xp.concatenate((self.pol[0].flatten(), out)) @@ -340,6 +342,7 @@ def toarray(self, allreduce=False): # allreduce tensor-product part if self.space.comm is not None and allreduce: + synchronize_for_mpi(out1, out2, out3) self.space.comm.Allreduce(MPI.IN_PLACE, out1, op=MPI.SUM) self.space.comm.Allreduce(MPI.IN_PLACE, out2, op=MPI.SUM) self.space.comm.Allreduce(MPI.IN_PLACE, out3, op=MPI.SUM) diff --git a/src/struphy/polar/extraction_operators.py b/src/struphy/polar/extraction_operators.py index 0de114d15..a6caa9e14 100644 --- a/src/struphy/polar/extraction_operators.py +++ b/src/struphy/polar/extraction_operators.py @@ -1,4 +1,5 @@ import cunumpy as xp +import numpy as np from struphy.eigenvalue_solvers.derivatives import grad_1d_matrix @@ -26,29 +27,33 @@ class PolarExtractionBlocksC1: A linear combination of n_rings rings is used to construct n_polar polar basis functions/DOFs and "first tp ring" basis functions/DOFs. + The blocks are setup data: they are built on the host (NumPy, SciPy sparse) on every array backend, + from a host copy of the control points. On the CuPy backend the polar operators + (:mod:`struphy.polar.linear_operators`) copy them to the device once and apply them there. + Parameters ---------- - cx : array-like - Control points defining the x-component of a 2D B-spline mapping. + domain : struphy.geometry.base.Domain + Spline mapping (IGAPolarCylinder, IGAPolarTorus, ...) whose control points define the polar splines. - cy : array-like - Control points defining the y-component of a 2D B-spline mapping. + derham : struphy.feec.psydac_derham.Derham + Discrete Derham complex. """ def __init__(self, domain, derham): from scipy.sparse import csr_matrix as csr - # get control points - cx = domain.cx[:, :, 0] - cy = domain.cy[:, :, 0] + # get control points (host copies; they live on the device on the CuPy backend) + cx = xp.to_numpy(domain.cx[:, :, 0]) + cy = xp.to_numpy(domain.cy[:, :, 0]) self._cx = cx self._cy = cy self._pole = (cx[0, 0], cy[0, 0]) - assert xp.all(cx[0] == self.pole[0]) - assert xp.all(cy[0] == self.pole[1]) + assert np.all(cx[0] == self.pole[0]) + assert np.all(cy[0] == self.pole[1]) self._n0 = cx.shape[0] self._n1 = cx.shape[1] @@ -70,14 +75,14 @@ def __init__(self, domain, derham): self._tau = max( [ ((self.cx[1] - self.pole[0]) * (-2)).max(), - ((self.cx[1] - self.pole[0]) - xp.sqrt(3) * (self.cy[1] - self.pole[1])).max(), - ((self.cx[1] - self.pole[0]) + xp.sqrt(3) * (self.cy[1] - self.pole[1])).max(), + ((self.cx[1] - self.pole[0]) - np.sqrt(3) * (self.cy[1] - self.pole[1])).max(), + ((self.cx[1] - self.pole[0]) + np.sqrt(3) * (self.cy[1] - self.pole[1])).max(), ], ) # barycentric coordinates - self._xi_0 = xp.zeros((3, self.n1), dtype=float) - self._xi_1 = xp.zeros((3, self.n1), dtype=float) + self._xi_0 = np.zeros((3, self.n1), dtype=float) + self._xi_1 = np.zeros((3, self.n1), dtype=float) self._xi_0[:, :] = 1 / 3 @@ -85,12 +90,12 @@ def __init__(self, domain, derham): self._xi_1[1, :] = ( 1 / 3 - 1 / (3 * self.tau) * (self.cx[1] - self.pole[0]) - + xp.sqrt(3) / (3 * self.tau) * (self.cy[1] - self.pole[1]) + + np.sqrt(3) / (3 * self.tau) * (self.cy[1] - self.pole[1]) ) self._xi_1[2, :] = ( 1 / 3 - 1 / (3 * self.tau) * (self.cx[1] - self.pole[0]) - - xp.sqrt(3) / (3 * self.tau) * (self.cy[1] - self.pole[1]) + - np.sqrt(3) / (3 * self.tau) * (self.cy[1] - self.pole[1]) ) # remove small values @@ -102,17 +107,17 @@ def __init__(self, domain, derham): # ============= basis extraction operator for discrete 0-forms ================ # first n_rings tp rings --> "polar coeffs" - e0_blocks_ten_to_pol = xp.block([self.xi_0, self.xi_1]) + e0_blocks_ten_to_pol = np.block([self.xi_0, self.xi_1]) self._e_ten_to_pol["0"] = [[csr(e0_blocks_ten_to_pol)]] # ============ basis extraction operator for discrete 1-forms (Hcurl) ========= # first n_rings tp rings --> "polar coeffs" - e1_11_blocks_ten_to_pol = xp.zeros((self.n_polar[1][0], self.n_rings[1][0] * self.n1), dtype=float) - e1_12_blocks_ten_to_pol = xp.zeros((self.n_polar[1][0], self.n_rings[1][1] * self.d1), dtype=float) + e1_11_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][0] * self.n1), dtype=float) + e1_12_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][1] * self.d1), dtype=float) - e1_21_blocks_ten_to_pol = xp.zeros((self.n_polar[1][1], self.n_rings[1][0] * self.n1), dtype=float) - e1_22_blocks_ten_to_pol = xp.zeros((self.n_polar[1][1], self.n_rings[1][1] * self.d1), dtype=float) + e1_21_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][0] * self.n1), dtype=float) + e1_22_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][1] * self.d1), dtype=float) # 1st component for l in range(2): @@ -135,7 +140,7 @@ def __init__(self, domain, derham): # =============== basis extraction operator for discrete 1-forms (Hdiv) ========= # first n_rings tp rings --> "polar coeffs" - e3_blocks_ten_to_pol = xp.zeros((self.n_polar[3][0], self.n_rings[3][0] * self.d1), dtype=float) + e3_blocks_ten_to_pol = np.zeros((self.n_polar[3][0], self.n_rings[3][0] * self.d1), dtype=float) self._e_ten_to_pol["2"] = [ [csr(e1_22_blocks_ten_to_pol), csr(-e1_21_blocks_ten_to_pol), None], @@ -161,7 +166,7 @@ def __init__(self, domain, derham): self._p_ten_to_ten = {} # first n_rings tp rings --> "polar coeffs" - p0_blocks_ten_to_pol = xp.zeros((self.n_polar[0][0], self.n_rings[0][0] * self.n1), dtype=float) + p0_blocks_ten_to_pol = np.zeros((self.n_polar[0][0], self.n_rings[0][0] * self.n1), dtype=float) # !! NOTE: for odd spline degrees and periodic splines the first Greville point sometimes does NOT start at zero!! if domain.degree[1] % 2 != 0 and not (abs(derham.V0fem.spaces[1].interpolation_grid[0]) < 1e-14): @@ -176,15 +181,15 @@ def __init__(self, domain, derham): self._p_ten_to_pol["0"] = [[csr(p0_blocks_ten_to_pol)]] # first n_rings + 1 tp rings --> "first tp ring" - p0_blocks_ten_to_ten = xp.block([0 * xp.identity(self.n1)] * self.n_rings[0][0] + [xp.identity(self.n1)]) + p0_blocks_ten_to_ten = np.block([0 * np.identity(self.n1)] * self.n_rings[0][0] + [np.identity(self.n1)]) self._p_ten_to_ten["0"] = [[csr(p0_blocks_ten_to_ten)]] # =========== projection extraction operator for discrete 1-forms (Hcurl) ======== # first n_rings tp rings --> "polar coeffs" - p1_11_blocks_ten_to_pol = xp.zeros((self.n_polar[1][0], self.n_rings[1][0] * self.n1), dtype=float) - p1_22_blocks_ten_to_pol = xp.zeros((self.n_polar[1][1], self.n_rings[1][1] * self.d1), dtype=float) + p1_11_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][0] * self.n1), dtype=float) + p1_22_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][1] * self.d1), dtype=float) # !! NOTE: PSYDAC's first integration interval sometimes start at < 0 !! if derham.V3fem.spaces[1].histopolation_grid[0] < -1e-14: @@ -196,8 +201,8 @@ def __init__(self, domain, derham): p1_22_blocks_ten_to_pol[1, (self.d1 + 0 * self.d1 // 3) : (self.d1 + 1 * self.d1 // 3)] = 1.0 p1_22_blocks_ten_to_pol[1, (self.d1 + 1 * self.d1 // 3) : (self.d1 + 2 * self.d1 // 3)] = 1.0 - p1_12_blocks_ten_to_pol = xp.zeros((self.n_polar[1][0], self.n_rings[1][1] * self.d1), dtype=float) - p1_21_blocks_ten_to_pol = xp.zeros((self.n_polar[1][1], self.n_rings[1][0] * self.d1), dtype=float) + p1_12_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][1] * self.d1), dtype=float) + p1_21_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][0] * self.d1), dtype=float) self._p_ten_to_pol["1"] = [ [csr(p1_11_blocks_ten_to_pol), csr(p1_12_blocks_ten_to_pol), None], @@ -206,26 +211,26 @@ def __init__(self, domain, derham): ] # first n_rings + 1 tp rings --> "first tp ring" - p1_11_blocks_ten_to_ten = xp.zeros((self.n1, self.n1), dtype=float) + p1_11_blocks_ten_to_ten = np.zeros((self.n1, self.n1), dtype=float) # !! NOTE: for odd spline degrees and periodic splines the first Greville point sometimes does NOT start at zero!! if domain.degree[1] % 2 != 0 and not (abs(derham.V0fem.spaces[1].interpolation_grid[0]) < 1e-14): - p1_11_blocks_ten_to_ten[:, 3 * self.n1 // 3 - 1] = -xp.roll(self.xi_1[0], -1) - p1_11_blocks_ten_to_ten[:, 1 * self.n1 // 3 - 1] = -xp.roll(self.xi_1[1], -1) - p1_11_blocks_ten_to_ten[:, 2 * self.n1 // 3 - 1] = -xp.roll(self.xi_1[2], -1) + p1_11_blocks_ten_to_ten[:, 3 * self.n1 // 3 - 1] = -np.roll(self.xi_1[0], -1) + p1_11_blocks_ten_to_ten[:, 1 * self.n1 // 3 - 1] = -np.roll(self.xi_1[1], -1) + p1_11_blocks_ten_to_ten[:, 2 * self.n1 // 3 - 1] = -np.roll(self.xi_1[2], -1) else: p1_11_blocks_ten_to_ten[:, 0 * self.n1 // 3] = -self.xi_1[0] p1_11_blocks_ten_to_ten[:, 1 * self.n1 // 3] = -self.xi_1[1] p1_11_blocks_ten_to_ten[:, 2 * self.n1 // 3] = -self.xi_1[2] - p1_11_blocks_ten_to_ten += xp.identity(self.n1) + p1_11_blocks_ten_to_ten += np.identity(self.n1) - p1_11_blocks_ten_to_ten = xp.block([p1_11_blocks_ten_to_ten, xp.identity(self.n1)]) + p1_11_blocks_ten_to_ten = np.block([p1_11_blocks_ten_to_ten, np.identity(self.n1)]) - p1_22_blocks_ten_to_ten = xp.block([0 * xp.identity(self.d1)] * self.n_rings[1][1] + [xp.identity(self.d1)]) + p1_22_blocks_ten_to_ten = np.block([0 * np.identity(self.d1)] * self.n_rings[1][1] + [np.identity(self.d1)]) - p1_12_blocks_ten_to_ten = xp.zeros((self.d1, (self.n_rings[1][1] + 1) * self.d1), dtype=float) - p1_21_blocks_ten_to_ten = xp.zeros((self.n1, (self.n_rings[1][0] + 1) * self.n1), dtype=float) + p1_12_blocks_ten_to_ten = np.zeros((self.d1, (self.n_rings[1][1] + 1) * self.d1), dtype=float) + p1_21_blocks_ten_to_ten = np.zeros((self.n1, (self.n_rings[1][0] + 1) * self.n1), dtype=float) self._p_ten_to_ten["1"] = [ [csr(p1_11_blocks_ten_to_ten), csr(p1_12_blocks_ten_to_ten), None], @@ -236,7 +241,7 @@ def __init__(self, domain, derham): # ========== projection extraction operator for discrete 1-forms (Hdiv) ========== # first n_rings tp rings --> "polar coeffs" - p3_blocks_ten_to_pol = xp.zeros((self.n_polar[3][0], self.n_rings[3][0] * self.d1), dtype=float) + p3_blocks_ten_to_pol = np.zeros((self.n_polar[3][0], self.n_rings[3][0] * self.d1), dtype=float) self._p_ten_to_pol["2"] = [ [csr(p1_22_blocks_ten_to_pol), csr(p1_21_blocks_ten_to_pol), None], @@ -245,24 +250,24 @@ def __init__(self, domain, derham): ] # first n_rings + 1 tp rings --> "first tp ring" - p3_blocks_ten_to_ten = xp.zeros((self.d1, self.d1), dtype=float) + p3_blocks_ten_to_ten = np.zeros((self.d1, self.d1), dtype=float) - a0 = xp.diff(self.xi_1[1], append=self.xi_1[1, 0]) - a1 = xp.diff(self.xi_1[2], append=self.xi_1[2, 0]) + a0 = np.diff(self.xi_1[1], append=self.xi_1[1, 0]) + a1 = np.diff(self.xi_1[2], append=self.xi_1[2, 0]) # !! NOTE: PSYDAC's first integration interval sometimes start at < 0 !! if derham.V3fem.spaces[1].histopolation_grid[0] < -1e-14: p3_blocks_ten_to_ten[:, (0 * self.n1 // 3 + 1) : (1 * self.n1 // 3 + 1)] = ( - -xp.roll(a0, +1)[:, None] - xp.roll(a1, +1)[:, None] + -np.roll(a0, +1)[:, None] - np.roll(a1, +1)[:, None] ) - p3_blocks_ten_to_ten[:, (1 * self.n1 // 3 + 1) : (2 * self.n1 // 3 + 1)] = -xp.roll(a1, +1)[:, None] + p3_blocks_ten_to_ten[:, (1 * self.n1 // 3 + 1) : (2 * self.n1 // 3 + 1)] = -np.roll(a1, +1)[:, None] else: p3_blocks_ten_to_ten[:, 0 * self.n1 // 3 : 1 * self.n1 // 3] = -a0[:, None] - a1[:, None] p3_blocks_ten_to_ten[:, 1 * self.n1 // 3 : 2 * self.n1 // 3] = -a1[:, None] - p3_blocks_ten_to_ten += xp.identity(self.d1) + p3_blocks_ten_to_ten += np.identity(self.d1) - p3_blocks_ten_to_ten = xp.block([p3_blocks_ten_to_ten, xp.identity(self.d1)]) + p3_blocks_ten_to_ten = np.block([p3_blocks_ten_to_ten, np.identity(self.d1)]) self._p_ten_to_ten["2"] = [ [csr(p1_22_blocks_ten_to_ten), csr(p1_21_blocks_ten_to_ten), None], @@ -295,39 +300,39 @@ def __init__(self, domain, derham): # ======================= discrete gradient ====================================== # "polar coeffs" to "polar coeffs" - grad_pol_to_pol_1 = xp.zeros((self.n_polar[1][0], self.n_polar[0][0]), dtype=float) - grad_pol_to_pol_2 = xp.array([[-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]]) - grad_pol_to_pol_3 = xp.identity(self.n_polar[0][0], dtype=float) + grad_pol_to_pol_1 = np.zeros((self.n_polar[1][0], self.n_polar[0][0]), dtype=float) + grad_pol_to_pol_2 = np.array([[-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]]) + grad_pol_to_pol_3 = np.identity(self.n_polar[0][0], dtype=float) self._grad_pol_to_pol = [[csr(grad_pol_to_pol_1)], [csr(grad_pol_to_pol_2)], [csr(grad_pol_to_pol_3)]] # "polar coeffs" to "first tp ring" - grad_pol_to_ten_1 = xp.zeros(((self.n_rings[1][0] + 1) * self.n1, self.n_polar[0][0])) - grad_pol_to_ten_2 = xp.zeros(((self.n_rings[1][1] + 1) * self.d1, self.n_polar[0][0])) - grad_pol_to_ten_3 = xp.zeros(((self.n_rings[0][0] + 1) * self.n1, self.n_polar[0][0])) + grad_pol_to_ten_1 = np.zeros(((self.n_rings[1][0] + 1) * self.n1, self.n_polar[0][0])) + grad_pol_to_ten_2 = np.zeros(((self.n_rings[1][1] + 1) * self.d1, self.n_polar[0][0])) + grad_pol_to_ten_3 = np.zeros(((self.n_rings[0][0] + 1) * self.n1, self.n_polar[0][0])) grad_pol_to_ten_1[-self.n1 :, :] = -self.xi_1.T self._grad_pol_to_ten = [[csr(grad_pol_to_ten_1)], [csr(grad_pol_to_ten_2)], [csr(grad_pol_to_ten_3)]] # eta_3 direction - grad_e3_1 = xp.identity(self.n2, dtype=float) - grad_e3_2 = xp.identity(self.n2, dtype=float) - grad_e3_3 = grad_1d_matrix(derham.spl_kind[2], self.n2) + grad_e3_1 = np.identity(self.n2, dtype=float) + grad_e3_2 = np.identity(self.n2, dtype=float) + grad_e3_3 = xp.to_numpy(grad_1d_matrix(derham.spl_kind[2], self.n2)) self._grad_e3 = [[csr(grad_e3_1)], [csr(grad_e3_2)], [csr(grad_e3_3)]] # =========================== discrete curl ====================================== # "polar coeffs" to "polar coeffs" - curl_pol_to_pol_12 = xp.identity(self.n_polar[1][1], dtype=float) - curl_pol_to_pol_13 = xp.array([[-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]]) + curl_pol_to_pol_12 = np.identity(self.n_polar[1][1], dtype=float) + curl_pol_to_pol_13 = np.array([[-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]]) - curl_pol_to_pol_21 = xp.identity(self.n_polar[1][0], dtype=float) - curl_pol_to_pol_23 = xp.zeros((self.n_polar[2][1], self.n_polar[0][0]), dtype=float) + curl_pol_to_pol_21 = np.identity(self.n_polar[1][0], dtype=float) + curl_pol_to_pol_23 = np.zeros((self.n_polar[2][1], self.n_polar[0][0]), dtype=float) - curl_pol_to_pol_31 = xp.zeros((self.n_polar[3][0], self.n_polar[1][0]), dtype=float) - curl_pol_to_pol_32 = xp.zeros((self.n_polar[3][0], self.n_polar[1][1]), dtype=float) + curl_pol_to_pol_31 = np.zeros((self.n_polar[3][0], self.n_polar[1][0]), dtype=float) + curl_pol_to_pol_32 = np.zeros((self.n_polar[3][0], self.n_polar[1][1]), dtype=float) self._curl_pol_to_pol = [ [None, csr(-curl_pol_to_pol_12), csr(curl_pol_to_pol_13)], @@ -336,14 +341,14 @@ def __init__(self, domain, derham): ] # "polar coeffs" to "first tp ring" - curl_pol_to_ten_12 = xp.zeros(((self.n_rings[2][0] + 1) * self.d1, self.n_polar[1][1])) - curl_pol_to_ten_13 = xp.zeros(((self.n_rings[2][0] + 1) * self.d1, self.n_polar[0][0])) + curl_pol_to_ten_12 = np.zeros(((self.n_rings[2][0] + 1) * self.d1, self.n_polar[1][1])) + curl_pol_to_ten_13 = np.zeros(((self.n_rings[2][0] + 1) * self.d1, self.n_polar[0][0])) - curl_pol_to_ten_21 = xp.zeros(((self.n_rings[2][1] + 1) * self.n1, self.n_polar[1][0])) - curl_pol_to_ten_23 = xp.zeros(((self.n_rings[2][1] + 1) * self.n1, self.n_polar[0][0])) + curl_pol_to_ten_21 = np.zeros(((self.n_rings[2][1] + 1) * self.n1, self.n_polar[1][0])) + curl_pol_to_ten_23 = np.zeros(((self.n_rings[2][1] + 1) * self.n1, self.n_polar[0][0])) - curl_pol_to_ten_31 = xp.zeros(((self.n_rings[3][0] + 1) * self.n1, self.n_polar[1][0])) - curl_pol_to_ten_32 = xp.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[1][1])) + curl_pol_to_ten_31 = np.zeros(((self.n_rings[3][0] + 1) * self.n1, self.n_polar[1][0])) + curl_pol_to_ten_32 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[1][1])) curl_pol_to_ten_23[-self.n1 :, :] = -self.xi_1.T @@ -360,14 +365,14 @@ def __init__(self, domain, derham): ] # eta_3 direction - curl_e3_12 = grad_1d_matrix(derham.spl_kind[2], self.n2) - curl_e3_13 = xp.identity(self.d2) + curl_e3_12 = xp.to_numpy(grad_1d_matrix(derham.spl_kind[2], self.n2)) + curl_e3_13 = np.identity(self.d2) - curl_e3_21 = grad_1d_matrix(derham.spl_kind[2], self.n2) - curl_e3_23 = xp.identity(self.d2) + curl_e3_21 = xp.to_numpy(grad_1d_matrix(derham.spl_kind[2], self.n2)) + curl_e3_23 = np.identity(self.d2) - curl_e3_31 = xp.identity(self.n2) - curl_e3_32 = xp.identity(self.n2) + curl_e3_31 = np.identity(self.n2) + curl_e3_32 = np.identity(self.n2) self._curl_e3 = [ [None, csr(curl_e3_12), csr(curl_e3_13)], @@ -378,16 +383,16 @@ def __init__(self, domain, derham): # =========================== discrete div ====================================== # "polar coeffs" to "polar coeffs" - div_pol_to_pol_1 = xp.zeros((self.n_polar[3][0], self.n_polar[2][0]), dtype=float) - div_pol_to_pol_2 = xp.zeros((self.n_polar[3][0], self.n_polar[2][1]), dtype=float) - div_pol_to_pol_3 = xp.identity(self.n_polar[3][0], dtype=float) + div_pol_to_pol_1 = np.zeros((self.n_polar[3][0], self.n_polar[2][0]), dtype=float) + div_pol_to_pol_2 = np.zeros((self.n_polar[3][0], self.n_polar[2][1]), dtype=float) + div_pol_to_pol_3 = np.identity(self.n_polar[3][0], dtype=float) self._div_pol_to_pol = [[csr(div_pol_to_pol_1), csr(div_pol_to_pol_2), csr(div_pol_to_pol_3)]] # "polar coeffs" to "first tp ring" - div_pol_to_ten_1 = xp.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[2][0])) - div_pol_to_ten_2 = xp.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[2][1])) - div_pol_to_ten_3 = xp.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[3][0])) + div_pol_to_ten_1 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[2][0])) + div_pol_to_ten_2 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[2][1])) + div_pol_to_ten_3 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[3][0])) for l in range(2): for j in range(self.d1, 2 * self.d1): @@ -398,9 +403,9 @@ def __init__(self, domain, derham): self._div_pol_to_ten = [[csr(div_pol_to_ten_1), csr(div_pol_to_ten_2), csr(div_pol_to_ten_3)]] # eta_3 direction - div_e3_1 = xp.identity(self.d2, dtype=float) - div_e3_2 = xp.identity(self.d2, dtype=float) - div_e3_3 = grad_1d_matrix(derham.spl_kind[2], self.n2) + div_e3_1 = np.identity(self.d2, dtype=float) + div_e3_2 = np.identity(self.d2, dtype=float) + div_e3_3 = xp.to_numpy(grad_1d_matrix(derham.spl_kind[2], self.n2)) self._div_e3 = [[csr(div_e3_1), csr(div_e3_2), csr(div_e3_3)]] diff --git a/src/struphy/polar/linear_operators.py b/src/struphy/polar/linear_operators.py index 6b4784032..46ba3851f 100644 --- a/src/struphy/polar/linear_operators.py +++ b/src/struphy/polar/linear_operators.py @@ -1,4 +1,5 @@ import cunumpy as xp +from cunumpy.mpi import synchronize_for_mpi from feectools.linalg.basic import LinearOperator from feectools.linalg.block import BlockVector, BlockVectorSpace from feectools.linalg.stencil import StencilVector, StencilVectorSpace @@ -8,8 +9,113 @@ from struphy.linear_algebra.linalg_kron import kron_matvec_2d from struphy.polar.basic import PolarDerhamSpace, PolarVector +# Sparse module used to put the polar blocks on the device on the CuPy backend; ``None`` means ``xp.scipy.sparse`` +# (``cupyx.scipy.sparse``). Tests without a GPU replace it with :func:`set_device_sparse_module`. +_DEVICE_SPARSE_MODULE = None -class PolarExtractionOperator(LinearOperator): + +def set_device_sparse_module(module): + """Set the sparse module that builds the device copies of the polar blocks on the CuPy backend. + + Parameters + ---------- + module : module | None + Must provide ``csr_matrix(host_csr_matrix)`` returning an object with ``.dot(device_array)``. + ``None`` restores the default, ``xp.scipy.sparse`` (``cupyx.scipy.sparse`` on the CuPy backend). + + Returns + ------- + previous : module | None + The module that was set before. + """ + global _DEVICE_SPARSE_MODULE + previous = _DEVICE_SPARSE_MODULE + _DEVICE_SPARSE_MODULE = module + return previous + + +def device_sparse_module(): + """The sparse module for the device copies of the polar blocks (``cupyx.scipy.sparse`` by default).""" + return xp.scipy.sparse if _DEVICE_SPARSE_MODULE is None else _DEVICE_SPARSE_MODULE + + +class DeviceSparseMatrix: + """Device copy of a (small) host matrix of the polar operators, for products with device arrays. + + The polar blocks (:class:`~struphy.polar.extraction_operators.PolarExtractionBlocksC1`) are setup data built + with SciPy on the host. On the CuPy backend the polar operators apply them to coefficients that live on the + device; this class holds a CSR copy of a block on the device (``cupyx.scipy.sparse``) and computes ``A @ x`` + there (SpMV/SpMM on the GPU). Blocks without non-zero entries (or with an empty dimension) are not copied, + their products are zeros. + + Parameters + ---------- + host : scipy.sparse matrix | numpy.ndarray + The host matrix. + """ + + def __init__(self, host): + host = csr_matrix(host) + self._shape = host.shape + self._dtype = host.dtype + if host.nnz == 0 or 0 in host.shape: + self._device = None + else: + self._device = device_sparse_module().csr_matrix(host) + + @property + def shape(self): + return self._shape + + @property + def dtype(self): + return self._dtype + + def dot(self, x): + """Product with a device array x of shape (shape[1],) or (shape[1], k).""" + assert x.shape[0] == self.shape[1], f"{x.shape =} does not match {self.shape =}" + if self._device is None: + return xp.zeros((self.shape[0],) + tuple(x.shape[1:]), dtype=float) + # kron_matvec_2d passes transposed views; cuSPARSE takes a contiguous dense operand + return self._device.dot(xp.ascontiguousarray(x)) + + +def device_blocks(blocks): + """2D nested list of host blocks (``None`` allowed) -> same list with :class:`DeviceSparseMatrix` copies.""" + if blocks is None: + return None + return [[None if blk is None else DeviceSparseMatrix(blk) for blk in row] for row in blocks] + + +class _BlocksOnBackend: + """Mixin: the blocks of a polar operator for the active array backend. + + The blocks themselves (properties ``blocks_*``) stay host matrices on every backend. On the NumPy backend + :meth:`_backend_blocks` returns them unchanged; on the CuPy backend it returns device copies + (:func:`device_blocks`), made once per block list and dropped when the list is replaced. + """ + + def _backend_blocks(self, name): + blocks = getattr(self, name) + if blocks is None or xp.get_backend() == "numpy": + return blocks + cache = self.__dict__.setdefault("_device_blocks_cache", {}) + if name not in cache: + cache[name] = device_blocks(blocks) + return cache[name] + + def _drop_device_blocks(self, name): + self.__dict__.get("_device_blocks_cache", {}).pop(name, None) + + +def _allreduce_sum(comm, buf): + """In-place sum over the ranks of comm (buf may live on the device: CUDA-aware MPI).""" + if comm is not None: + synchronize_for_mpi(buf) + comm.Allreduce(MPI.IN_PLACE, buf, op=MPI.SUM) + + +class PolarExtractionOperator(_BlocksOnBackend, LinearOperator): """ Linear operator mapping from Stencil-/BlockVectorSpace (V) to PolarDerhamSpace (W). @@ -149,6 +255,7 @@ def blocks_ten_to_pol(self, blocks): check_blocks(blocks, self.blocks_ten_to_pol_shapes) self._blocks_ten_to_pol = blocks + self._drop_device_blocks("blocks_ten_to_pol") @property def blocks_ten_to_ten_shapes(self): @@ -168,6 +275,7 @@ def blocks_ten_to_ten(self, blocks): check_blocks(blocks, self.blocks_ten_to_ten_shapes) self._blocks_ten_to_ten = blocks + self._drop_device_blocks("blocks_ten_to_ten") @property def blocks_e3(self): @@ -204,14 +312,16 @@ def dot(self, v, out=None): assert out.space == self._codomain v.tp.copy(out=out) + blocks_e3 = self._backend_blocks("blocks_e3") + # 2. map "first tp ring" to "polar rings" + "first tp ring" if self.blocks_ten_to_ten is not None: - dot_parts_of_polar(self.blocks_ten_to_ten, self.blocks_e3, v, out, map_from_tp=True) + dot_parts_of_polar(self._backend_blocks("blocks_ten_to_ten"), blocks_e3, v, out, map_from_tp=True) # 3. map polar coeffs to "polar rings" if self.blocks_ten_to_pol is not None: out2 = out.space.zeros() - dot_parts_of_polar(self.blocks_ten_to_pol, self.blocks_e3, v, out2, map_from_tp=False) + dot_parts_of_polar(self._backend_blocks("blocks_ten_to_pol"), blocks_e3, v, out2, map_from_tp=False) # add contributions to "polar rings" out += out2 @@ -230,13 +340,15 @@ def dot(self, v, out=None): # 1. identity operation on outer tp zone out.tp = v + blocks_e3 = self._backend_blocks("blocks_e3") + # 2. map from "polar rings" to polar coeffs if self.blocks_ten_to_pol is not None: - dot_inner_tp_rings(self.blocks_ten_to_pol, self.blocks_e3, v, out, map_to_tp=False) + dot_inner_tp_rings(self._backend_blocks("blocks_ten_to_pol"), blocks_e3, v, out, map_to_tp=False) # 3. map to "polar rings" + "first tp ring" to "first tp ring" if self.blocks_ten_to_ten is not None: - dot_inner_tp_rings(self.blocks_ten_to_ten, self.blocks_e3, v, out, map_to_tp=True) + dot_inner_tp_rings(self._backend_blocks("blocks_ten_to_ten"), blocks_e3, v, out, map_to_tp=True) return out @@ -271,7 +383,7 @@ def transpose(self, conjugate=False): ) -class PolarLinearOperator(LinearOperator): +class PolarLinearOperator(_BlocksOnBackend, LinearOperator): """ Linear operator mapping from PolarDerhamSpace (V) to PolarDerhamSpace (W). @@ -431,6 +543,8 @@ def blocks_pol_to_ten(self, blocks): check_blocks(blocks, self.blocks_pol_to_ten_shapes) self._blocks_pol_to_ten = blocks + self._drop_device_blocks("blocks_pol_to_ten") + @property def tp_blocks_shapes(self): return self._tp_blocks_shapes @@ -471,6 +585,8 @@ def blocks_pol_to_pol(self, blocks): check_blocks(blocks, self.blocks_pol_to_pol_shapes) self._blocks_pol_to_pol = blocks + self._drop_device_blocks("blocks_pol_to_pol") + @property def blocks_e3_shapes(self): return self._blocks_e3_shapes @@ -491,6 +607,8 @@ def blocks_e3(self, blocks): check_blocks(blocks, self.blocks_e3_shapes) self._blocks_e3 = blocks + self._drop_device_blocks("blocks_e3") + def dot(self, v, out=None): """ Dot product mapping from PolarVector to PolarVector @@ -521,20 +639,23 @@ def dot(self, v, out=None): # 1. total tp zone to total tp zone (stencil) out.tp = self.tp_operator.dot(v.tp) + blocks_e3 = self._backend_blocks("blocks_e3") + blocks_pol_to_ten = self._backend_blocks("blocks_pol_to_ten") + # 2. polar to polar (dense) - dot_pol_pol(self.blocks_pol_to_pol, self.blocks_e3, v, out) + dot_pol_pol(self._backend_blocks("blocks_pol_to_pol"), blocks_e3, v, out) out2 = PolarVector(self.codomain) # transposed operator if self.transposed: # 3. "first tp ring" to polar - dot_inner_tp_rings(self.blocks_pol_to_ten, self.blocks_e3, v.tp, out2, map_to_tp=False) + dot_inner_tp_rings(blocks_pol_to_ten, blocks_e3, v.tp, out2, map_to_tp=False) # "standard" operator else: # 3. polar to "first tp ring" - dot_parts_of_polar(self.blocks_pol_to_ten, self.blocks_e3, v, out2.tp, map_from_tp=False) + dot_parts_of_polar(blocks_pol_to_ten, blocks_e3, v, out2.tp, map_from_tp=False) # sum up contributions out += out2 @@ -674,8 +795,7 @@ def dot_inner_tp_rings(blocks_e1_e2, blocks_e3, v, out, map_to_tp): res += kron_matvec_2d([block_e1_e2, block_e3], tmp.reshape(n_rings_in[n] * n2, n3_in[n])) # sum up local dot products - if polar_space.comm is not None: - polar_space.comm.Allreduce(MPI.IN_PLACE, res, op=MPI.SUM) + _allreduce_sum(polar_space.comm, res) # write result to output polar vector (in-place) # (only the rank owning the first eta_1 coeffs holds the "first tp ring", see PolarDerhamSpace) @@ -796,8 +916,7 @@ def dot_parts_of_polar(blocks_e1_e2, blocks_e3, v, out, map_from_tp): res += kron_matvec_2d([block_e1_e2, block_e3], v.pol[n]).reshape(n_rings_out[m], n2, n3_out[m]) if map_from_tp: - if polar_space.comm is not None: - polar_space.comm.Allreduce(MPI.IN_PLACE, res, op=MPI.SUM) + _allreduce_sum(polar_space.comm, res) if out_starts[m][0] == 0: s1, s2, s3 = out_starts[m] diff --git a/src/struphy/polar/tests/test_polar_cupy.py b/src/struphy/polar/tests/test_polar_cupy.py new file mode 100644 index 000000000..4b26bac3a --- /dev/null +++ b/src/struphy/polar/tests/test_polar_cupy.py @@ -0,0 +1,255 @@ +"""Polar splines on the CuPy backend (struphy-hub/struphy#695, CUDA_STRATEGY.md). + +The polar blocks (:class:`~struphy.polar.extraction_operators.PolarExtractionBlocksC1`) are built on the host on +every backend; on CuPy the polar operators apply device copies of them (``cupyx.scipy.sparse``) to the device +coefficients. :func:`check_polar_backends` builds the same polar Derham on the NumPy and the CuPy backend and +compares extraction operators, derivatives, ``PolarVector`` arithmetic, mass matrices, projectors and polar +``SplineFunction`` coefficients. + +Without a GPU it runs on cunumpy's fake CuPy (which rejects host/device mixing like CuPy), in a subprocess, with +two substitutes for what the fake cannot do: the device sparse module (``cupyx.scipy.sparse``) is replaced by +:class:`DenseDeviceSparse` (dense products of device arrays), and kernels are called in their host version (the +fake cannot launch CUDA kernels; the feectools stencil kernels then copy their arrays to the host and back). +""" + +import subprocess +import sys + +import cunumpy +import numpy as np +import pytest + +from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env + +requires_cupy = pytest.mark.skipif(not cunumpy.cupy_available(), reason="CuPy/GPU not available") + +MAPPINGS = ("IGAPolarCylinder", "IGAPolarTorus") +FORMS = ("0", "1", "2", "3", "v") +SPACE_IDS = {"0": "H1", "1": "Hcurl", "2": "Hdiv", "3": "L2", "v": "H1vec"} + + +class DenseDeviceSparse: + """Stand-in for ``cupyx.scipy.sparse`` on the fake CuPy: the block as a dense array of the active backend.""" + + class csr_matrix: + def __init__(self, host): + self._array = cunumpy.asarray(host.toarray()) + self.shape = host.shape + + def dot(self, x): + return self._array @ x + + +def make_domain(mapping): + from struphy import domains + + if mapping == "IGAPolarCylinder": + return domains.IGAPolarCylinder(num_elements=(6, 6), degree=(2, 3), a=1.2, Lz=3.0) + if mapping == "IGAPolarTorus": + return domains.IGAPolarTorus(num_elements=(6, 6), degree=(3, 2), a=0.8, R0=3.0, sfl=True) + raise ValueError(mapping) + + +def make_polar_derham(mapping, with_mass=True): + """Polar Derham (with projectors), its domain and mass matrices (or None), on the active backend.""" + from maybempi import MPI + + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + domain = make_domain(mapping) + options = DerhamOptions( + degree=(domain.degree[0], domain.degree[1], 1), + bcs=(("dirichlet", "free"), None, None), + nquads_proj=(4, 4, 2), + polar_splines=True, + ) + derham = Derham(TensorProductGrid(num_elements=(6, 6, 4)), options, comm=MPI.COMM_WORLD, domain=domain) + return domain, derham, WeightedMassOperators(derham, domain) if with_mass else None + + +def _stencils(vec): + """The Stencil- (or Block-) vector part of vec as a list of StencilVectors.""" + from feectools.linalg.stencil import StencilVector + + tp = vec.tp if hasattr(vec, "tp") else vec + return [tp] if isinstance(tp, StencilVector) else list(tp.blocks) + + +def to_host(vec): + """Data of a Stencil-/Block-/PolarVector as a list of NumPy arrays (polar coeffs first).""" + pol = list(vec.pol) if hasattr(vec, "pol") else [] + return [cunumpy.to_numpy(a) for a in pol] + [cunumpy.to_numpy(s._data) for s in _stencils(vec)] + + +def fill(vec, seed): + """Fill vec (on the active backend) with reproducible random values (same on every backend).""" + rng = np.random.default_rng(seed) + if hasattr(vec, "pol"): + vec.pol = [cunumpy.asarray(rng.random(p.shape)) for p in vec.pol] + for s in _stencils(vec): + s._data[...] = cunumpy.asarray(rng.random(s._data.shape)) + if hasattr(vec, "pol"): + vec.set_tp_coeffs_to_zero() + vec.update_ghost_regions() + return vec + + +def assert_on_backend(vec, backend): + """All data of vec lives on the given backend.""" + pol = list(vec.pol) if hasattr(vec, "pol") else [] + for a in pol + [s._data for s in _stencils(vec)]: + assert cunumpy.is_gpu(a) == (backend == "cupy"), type(a) + + +def polar_results(mapping, backend, with_mass=True): + """Results of the polar operations on `backend`, as host arrays (dict name -> list of arrays).""" + from struphy.polar.basic import PolarVector + + res = {} + with cunumpy.use_backend(backend): + domain, derham, mass = make_polar_derham(mapping, with_mass) + + def record(name, vec): + assert_on_backend(vec, backend) + res[name] = to_host(vec) + + for n, form in enumerate(FORMS): + E = derham.extraction_ops[form] + x = fill(E.domain.zeros(), seed=10 + n) + y = fill(PolarVector(E.codomain), seed=20 + n) + record(f"E{form}", E.dot(x)) + record(f"E{form}T", E.transpose().dot(y)) + P = derham.dofs_extraction_ops[form] + record(f"P{form}", P.dot(x)) + record(f"P{form}T", P.transpose().dot(y)) + + # PolarVector arithmetic + z = fill(PolarVector(E.codomain), seed=30 + n) + record(f"add{form}", y + z) + record(f"sub{form}", y - z) + record(f"mul{form}", y * 1.5) + record(f"rmul{form}", 0.5 * y) + record(f"neg{form}", -y) + w = y.copy() + w += z + w -= 2.0 * z + w *= 3.0 + record(f"inplace{form}", w) + res[f"dot{form}"] = [np.asarray(float(y.dot(z)))] + res[f"toarray{form}"] = [cunumpy.to_numpy(y.toarray())] + + # mass matrices (E M E^T with the polar extraction operators) + if with_mass: + record(f"M{form}", getattr(mass, "M" + form).dot(y)) + + for name, op in (("grad", derham.grad), ("curl", derham.curl), ("div", derham.div)): + x = fill(PolarVector(op.domain), seed=40) + y = fill(PolarVector(op.codomain), seed=41) + record(name, op.dot(x)) + record(name + "T", op.transpose().dot(y)) + + # commuting projectors (polar case: iterative solve with the polar extraction operators) + def f(e1, e2, e3): + return cunumpy.sin(np.pi * e1) * cunumpy.cos(2 * np.pi * e2) * cunumpy.cos(2 * np.pi * e3) + e1**2 + + record("Pi0", derham.P0(f)) + record("Pi1", derham.P1([f, f, f])) + record("Pi2", derham.P2([f, f, f])) + record("Pi3", derham.P3(f)) + + # polar SplineFunction: polar coefficients, extracted tensor-product coefficients + for form in ("0", "2"): + field = derham.create_spline_function("f" + form, SPACE_IDS[form]) + assert isinstance(field.vector, PolarVector) + fill(field.vector, seed=50) + field.extract_coeffs() + record("field" + form, field.vector) + record("field_stencil" + form, field.vector_stencil) + # restart: polar coefficients recovered from the extracted tensor-product coefficients + restored = field._restart_extraction_op().dot(field.vector_stencil) + record("restart" + form, restored) + np.testing.assert_allclose(to_host(restored)[0], to_host(field.vector)[0], rtol=1e-10, atol=1e-12) + + return res + + +def check_polar_backends(mapping, with_mass=True): + """The polar operations agree on the NumPy and the CuPy backend.""" + host = polar_results(mapping, "numpy", with_mass) + device = polar_results(mapping, "cupy", with_mass) + assert host.keys() == device.keys() + for name in host: + assert len(host[name]) == len(device[name]), name + for a, b in zip(host[name], device[name]): + np.testing.assert_allclose(b, a, rtol=1e-9, atol=1e-11, err_msg=name) + + +def check_polar_backends_fake_cupy(mapping): + """:func:`check_polar_backends` on the fake CuPy (see the module docstring), without the mass matrices. + + The mass matrix weights are evaluated by struphy's CUDA geometry kernels, which take CUDA argument objects + and cannot fall back to their pyccel versions; the polar mass matrices are compared on the GPU only. + """ + from cunumpy._dispatch import Kernel + + from struphy.polar import linear_operators + + Kernel.get_kernel = lambda self: self._host_kernel + linear_operators.set_device_sparse_module(DenseDeviceSparse) + check_polar_backends(mapping, with_mass=False) + + +def test_device_blocks_numpy_backend(): + """On the NumPy backend the polar operators apply their host (SciPy) blocks, as before.""" + import scipy.sparse as sp + + from struphy.polar.linear_operators import DeviceSparseMatrix, set_device_sparse_module + + with cunumpy.use_backend("numpy"): + domain, derham, _ = make_polar_derham("IGAPolarCylinder", with_mass=False) + E = derham.extraction_ops["1"] + assert E._backend_blocks("blocks_ten_to_pol") is E.blocks_ten_to_pol + assert E._backend_blocks("blocks_e3") is E.blocks_e3 + assert all(sp.issparse(b) for row in E.blocks_ten_to_pol for b in row if b is not None) + + # the device copy (here with the dense stand-in on the NumPy backend) computes the same products + previous = set_device_sparse_module(DenseDeviceSparse) + try: + rng = np.random.default_rng(0) + for row in derham.dofs_extraction_ops["1"].blocks_ten_to_ten + E.blocks_ten_to_pol: + for blk in row: + if blk is None: + continue + x = rng.random((blk.shape[1], 3)) + np.testing.assert_allclose(DeviceSparseMatrix(blk).dot(x), blk.dot(x), rtol=1e-14) + empty = DeviceSparseMatrix(sp.csr_matrix((0, 4))) + assert empty.dot(np.ones((4, 2))).shape == (0, 2) + finally: + set_device_sparse_module(previous) + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +@pytest.mark.parametrize("mapping", MAPPINGS) +def test_polar_fake_cupy(mapping): + """Without a GPU: polar Derham on cunumpy's fake CuPy agrees with the NumPy backend. + + Runs in a subprocess because the fake CuPy must be installed before cunumpy is imported. + """ + code = ( + "from struphy.polar.tests.test_polar_cupy import check_polar_backends_fake_cupy; " + f"check_polar_backends_fake_cupy({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:] + + +@requires_cupy +@pytest.mark.parametrize("mapping", MAPPINGS) +def test_polar_on_cupy(mapping): + """Polar Derham on the GPU (cupyx.scipy.sparse for the polar blocks) agrees with the NumPy backend.""" + check_polar_backends(mapping) From 1979ea91fa3ddb3a47448f6c17f10d674cdb464f Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 23:35:08 +0200 Subject: [PATCH 10/50] Mark every column of the test hole so accumulation kernels skip it The hole in marker_arguments() was only marked in column 8 (first_init_idx), which the pushers test. Accumulation kernels test column 0, so the linear_vlasov_ampere parity case had no real hole. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/kernel_test_args.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index eaed48adf..24e597e89 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -29,7 +29,9 @@ def marker_arguments(bc): markers[:, 3:6] = rng.uniform(-2, 2, (N_MARKERS, 3)) markers[:, 8:14] = markers[:, :6] markers[:, 18:21] = 0.0 - markers[0, 8] = -1.0 + # Row 0 is a hole. Struphy marks every column of a hole with -1: pushers test column first_init_idx (8), + # accumulation kernels test column 0. + markers[0, :] = -1.0 markers[1, -1] = -2.0 valid = np.ones(N_MARKERS, dtype=bool) valid[:2] = False From 256d02d1c8bb0dd1df17940540c4ec90353b256e Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 8 Oct 2026 09:22:24 +0200 Subject: [PATCH 11/50] 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 From f2e8676f1553a37a532d4a6221e8040e4490844c Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 20:49:16 +0200 Subject: [PATCH 12/50] Marker exchange between ranks on the device (CUDA-aware MPI) On the CuPy backend, Particles._sendrecv_markers hands the marker rows to Isend/Irecv as device buffers through cunumpy.mpi.mpi_buffer (directly after synchronize_for_mpi with CUDA-aware MPI, staged through host memory otherwise). Rows from each rank arrive in a contiguous slice of one device receive buffer, and one device scatter fills the holes. The per-rank counts are small NumPy arrays on every backend (they are host integers already), and a uniform alpha is a Python scalar, so the exchange makes no host/device copy and no implicit host read of a device array. The NumPy path is unchanged. Adds pic/tests/test_sorting_device.py (fake CuPy under mpirun, and a GPU variant) and a CI step running it on the fake CuPy. Solves #698, part of #650. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/test-PR-unit-mpi.yml | 6 + CUDA_STRATEGY.md | 20 +- src/struphy/pic/base.py | 129 ++++++++++++- src/struphy/pic/tests/test_sorting_device.py | 190 +++++++++++++++++++ 4 files changed, 338 insertions(+), 7 deletions(-) create mode 100644 src/struphy/pic/tests/test_sorting_device.py diff --git a/.github/workflows/test-PR-unit-mpi.yml b/.github/workflows/test-PR-unit-mpi.yml index e3a13dbc5..7bb53fbe2 100644 --- a/.github/workflows/test-PR-unit-mpi.yml +++ b/.github/workflows/test-PR-unit-mpi.yml @@ -76,3 +76,9 @@ jobs: source /struphy_fortran_/env_fortran_/bin/activate cd ${{ env.STRUPHY_PATH }} mpirun --oversubscribe -n ${{ matrix.n-procs }} pytest -v --with-mpi -m mpi_pic pic/tests + + - name: Run the marker exchange on the fake CuPy + run: | + source /struphy_fortran_/env_fortran_/bin/activate + cd ${{ env.STRUPHY_PATH }} + mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi pic/tests/test_sorting_device.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 6d2c60e72..68035f68c 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -365,7 +365,7 @@ matrix data). - **Kernel launch configuration.** One thread per marker with a fixed block size for now. Kernels over grid points (accumulation, FEEC) need their own launch sizes; `CudaKernel` needs a `grid`/`block` override (PR 12+), or cunumpy's. - **Accumulation strategy.** Atomics vs. sort-then-reduce; see PR 12+. Decided by measurement on the first accumulation kernel. - **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. `xp.mpi_is_cuda_aware()` detects it; the exchange in `Particles.mpi_sort_markers` has to be checked for host staging buffers. +- **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 `Derham` with `polar_splines=True` raises on CuPy (`PolarExtractionBlocksC1` builds SciPy sparse matrices, and the @@ -644,3 +644,21 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a The GPU tests have not run on an H100; the CPU emulation and the CPU regression tests are the gate. - **Parity cases.** The spline mappings are added to the geometry kernels' entries of `pic/tests/cuda_parity_cases.py` (the per-folder `_test_args.py` modules were replaced by that module in PR 16). + +## Marker exchange implementation notes (#698) + +- `Particles._sendrecv_markers` on CuPy (`_sendrecv_markers_device`): the rows to send (gathered on the device per + destination) and one device receive buffer go to `Isend`/`Irecv` through `xp.mpi.mpi_buffer`, i.e. directly after + `synchronize_for_mpi` with CUDA-aware MPI, staged through (pinned) host memory otherwise (with a warning, once). + Rows from rank `i` arrive in a contiguous slice of the receive buffer; one device scatter fills the holes. + Zero-count messages are skipped (both sides know the counts). The NumPy path is unchanged. +- Whether MPI is CUDA-aware: the answer recorded by `xp.mpi.mpi_is_cuda_aware()`/`set_mpi_cuda_aware()`, else a + collective probe of `mpi_comm` at the first exchange (`mpi_sort_markers` is collective). +- Counts: `send_info`/`recv_info` are small NumPy arrays on every backend. The send counts are sizes of `nonzero` + results, which are host integers already (CuPy synchronizes to size them), so the `Alltoall` of the counts needs no + copy. Uniform `alpha` is a Python scalar (no host-to-device copy of a 3-vector). +- Tests: `pic/tests/test_sorting_device.py` sorts the same markers on NumPy and on CuPy over all ranks and compares + the marker arrays; with CUDA-aware MPI no transfer (`count_transfers`) and no implicit host read of a device array + during the exchange, staged: one `to_host` per non-empty send and one `to_device`. Without a GPU it runs on the fake + CuPy (`mpirun -n N env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi pic/tests/test_sorting_device.py`, in the PIC MPI + workflow); the GPU variant has not run (no GPU here). diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 99db7dbfd..cae11c476 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2,6 +2,7 @@ import os import warnings from abc import ABCMeta, abstractmethod +from contextlib import ExitStack from typing import TYPE_CHECKING import cunumpy as xp @@ -1757,6 +1758,12 @@ def mpi_sort_markers( remove_ghost : bool Remove ghost particles before send. + + Notes + ----- + On the CuPy backend, the marker rows are exchanged as device buffers with CUDA-aware MPI + (staged through host memory otherwise), see :meth:`_sendrecv_markers_device`. The + per-rank counts are small host arrays on every backend (:meth:`_sendrecv_all_to_all`). """ if remove_ghost: self._remove_ghost_particles() @@ -4363,9 +4370,19 @@ def _sendrecv_determine_mtbs( """ # position that determines the sorting (including periodic shift of boundary conditions) if not isinstance(alpha, xp.ndarray): - alpha = xp.array(alpha, dtype=float) - assert alpha.size == 3 - assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) + # check a tuple/list on the host: on CuPy, asserting on a device comparison reads it back + alpha_host = np.array(alpha, dtype=float) + assert alpha_host.size == 3 + assert np.all(alpha_host >= 0.0) and np.all(alpha_host <= 1.0) + if alpha_host[0] == alpha_host[1] == alpha_host[2]: + # the common case (alpha = 1 or alpha = 0.5 in every direction): a Python scalar, + # which needs no host-to-device copy on CuPy and gives the same values on NumPy + alpha = float(alpha_host[0]) + else: + alpha = xp.asarray(alpha_host) + else: + assert alpha.size == 3 + assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) bi = self.first_pusher_idx self._sorting_etas[:] = ( alpha * (self.markers[:, :3] + self.markers[:, bi + 3 + self.vdim : bi + 3 + self.vdim + 3]) @@ -4465,8 +4482,11 @@ def _sendrecv_get_destinations(self, send_inds): Amount of particles sent to i-th process. """ - # One entry for each process - send_info = xp.zeros(self.mpi_size, dtype=int) + # One entry for each process. The counts are host integers on every backend: the + # size of a nonzero() result is known on the host (CuPy synchronizes to size it), so + # a small NumPy array holds them without any host/device copy, and the Alltoall of + # the counts (_sendrecv_all_to_all) passes host buffers to MPI. + send_info = np.zeros(self.mpi_size, dtype=int) # Gathered once and reused for every rank below, instead of re-gathering # self.markers[send_inds] and self._sorting_etas[send_inds] fresh on every @@ -4530,9 +4550,16 @@ def _sendrecv_all_to_all(self, send_info): ------- recv_info : array[int] Amount of marticles to be received from i-th process. + + Notes + ----- + ``send_info`` and ``recv_info`` are small NumPy arrays (one integer per rank) on every + backend: the counts size the receive buffers and the holes to fill, which Python needs + on the host anyway. On CuPy, only the marker rows themselves are device buffers + (:meth:`_sendrecv_markers`). """ - recv_info = xp.zeros(self.mpi_size, dtype=int) + recv_info = np.zeros(self.mpi_size, dtype=int) self.mpi_comm.Alltoall(send_info, recv_info) @@ -4551,6 +4578,10 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): Indices of empty rows in markers after send. """ + if xp.get_array_backend(self._markers) == "cupy": + self._sendrecv_markers_device(recv_info, hole_inds_after_send) + return + # i-th entry holds the number (not the index) of the first hole to be filled by data from process i first_hole = xp.cumsum(recv_info) - recv_info @@ -4594,6 +4625,92 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): for req in send_reqs: req.Wait() + def _mpi_cuda_aware(self) -> bool: + """Whether MPI takes device (CuPy) buffers, for :meth:`_sendrecv_markers_device`. + + Uses the answer recorded by :func:`cunumpy.mpi.mpi_is_cuda_aware` or + :func:`cunumpy.mpi.set_mpi_cuda_aware` (e.g. by feectools or at startup). If nothing + was recorded, probes :attr:`mpi_comm` once; the probe is collective, which is safe + here because :meth:`mpi_sort_markers` is called by every rank of :attr:`mpi_comm`. + """ + cuda_aware = xp.mpi.get_mpi_cuda_aware() + if cuda_aware is None: + cuda_aware = xp.mpi.mpi_is_cuda_aware(self.mpi_comm) + if not cuda_aware and not Particles._warned_mpi_staging: + Particles._warned_mpi_staging = True + logger.warning( + "MPI is not CUDA-aware: marker rows exchanged between ranks are staged through host memory.", + ) + return cuda_aware + + _warned_mpi_staging = False + + def _sendrecv_markers_device(self, recv_info, hole_inds_after_send): + """:meth:`_sendrecv_markers` on the CuPy backend: marker rows go to MPI as device buffers. + + With CUDA-aware MPI (:meth:`_mpi_cuda_aware`), the rows to send (:attr:`_send_list`, + gathered on the device by :meth:`_sendrecv_get_destinations`) and one device receive + buffer are handed to ``Isend``/``Irecv`` directly, after + :func:`cunumpy.mpi.synchronize_for_mpi` (inside :func:`cunumpy.mpi.mpi_buffer`), so MPI + reads what the kernels wrote. The rows from rank ``i`` arrive in a contiguous slice of + the receive buffer, and one device scatter puts all of them into the holes. Without + CUDA-aware MPI, :func:`~cunumpy.mpi.mpi_buffer` stages the same buffers through host + memory (counted by :func:`cunumpy.profiling.count_transfers`). + + The bookkeeping stays on the host without host/device copies: the counts are host + integers (``recv_info`` from the Alltoall, the send counts are the row counts of the + send buffers), and so is the number of holes. Ranks exchange no message when the + count is zero (both sides know the count, so they skip consistently). + + Parameters + ---------- + recv_info : array[int] + Amount of markers to be received from i-th process (NumPy array). + + hole_inds_after_send : array[int] + Indices of empty rows in markers after send (device array). + """ + recv_info = np.asarray(recv_info) + n_recv = int(recv_info.sum()) + if hole_inds_after_send.size < n_recv: + warnings.warn( + f'Strong load imbalance detected: \ +number of holes ({hole_inds_after_send.size}) on rank {self.mpi_rank} \ +is smaller than number of incoming particles ({n_recv}). \ +Increasing the value of "bufsize" in the markers parameters for the next run.', + ) + self.mpi_comm.Abort() + + cuda_aware = self._mpi_cuda_aware() + + # rows from rank i land in recvbuf[first_row[i] : first_row[i] + recv_info[i]] + first_row = np.cumsum(recv_info) - recv_info + recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=float) + + with ExitStack() as buffers: + if n_recv > 0: + recv_mpi = buffers.enter_context( + xp.mpi.mpi_buffer(recvbuf, send=False, recv=True, cuda_aware=cuda_aware), + ) + reqs = [] + for i in range(self.mpi_size): + if i == self.mpi_rank: + continue + if recv_info[i] > 0: + rows = recv_mpi[first_row[i] : first_row[i] + recv_info[i]] + reqs.append(self.mpi_comm.Irecv(rows, source=i, tag=i)) + data = self._send_list[i] + if data.shape[0] > 0: + send_mpi = buffers.enter_context(xp.mpi.mpi_buffer(data, cuda_aware=cuda_aware)) + reqs.append(self.mpi_comm.Isend(send_mpi, dest=i, tag=self.mpi_rank)) + # wait inside the block: the buffers must stay valid until MPI is done, and staged + # receive rows are copied to the device when the block ends + for req in reqs: + req.Wait() + + if n_recv > 0: + self._markers[hole_inds_after_send[:n_recv]] = recvbuf + class Tesselation: """ diff --git a/src/struphy/pic/tests/test_sorting_device.py b/src/struphy/pic/tests/test_sorting_device.py new file mode 100644 index 000000000..55c5d09f5 --- /dev/null +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -0,0 +1,190 @@ +"""Marker exchange between ranks on the CuPy backend (:meth:`Particles.mpi_sort_markers`). + +The marker rows are exchanged as device buffers with CUDA-aware MPI, and staged through +host memory otherwise. Without a GPU, the tests run on cunumpy's fake CuPy, which must be +installed before Python imports cunumpy, so they are skipped unless the fake is active:: + + mpirun -n 2 env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi src/struphy/pic/tests/test_sorting_device.py + +Fake CuPy arrays live in host memory, so a non-CUDA-aware MPI library takes them, and the +CUDA-aware path (the probe of :func:`cunumpy.mpi.mpi_is_cuda_aware` succeeds) is what runs by +default. The staged path is forced with :func:`cunumpy.mpi.set_mpi_cuda_aware`. +""" + +import sys +from contextlib import contextmanager + +import cunumpy as xp +import numpy as np +import pytest +from cunumpy.kernel_testing import requires_cupy +from maybempi import MPI + +from struphy import BoundaryParameters, LoadingParameters +from struphy.pic.particles import Particles6D + + +def _fake_cupy_active() -> bool: + return bool(getattr(sys.modules.get("cupy"), "__cunumpy_fake__", False)) + + +requires_fake_cupy = pytest.mark.skipif( + not _fake_cupy_active(), + reason="needs the fake CuPy: run under `mpirun -n N env CUNUMPY_FAKE_CUPY=1 pytest ...`", +) + + +@contextmanager +def _count_host_reads(): + """Count reads of fake CuPy arrays by the host (``int()``, ``bool()``, iteration, ``.get()``, ...). + + :func:`cunumpy.profiling.count_transfers` sees the copies made through cunumpy; this also sees + the implicit ones (a device scalar in an ``if``, ``list(device_array)``), which synchronize + and copy on a GPU. + """ + import cupy + + reads = [] + names = ("__int__", "__float__", "__bool__", "__index__", "__iter__", "get") + originals = {name: getattr(cupy.ndarray, name) for name in names} + original_getattr = cupy.ndarray.__getattr__ + + def counting(name, method): + def wrapper(self, *args, **kwargs): + reads.append(name) + return method(self, *args, **kwargs) + + return wrapper + + def counting_getattr(self, name): + if name in ("item", "tolist"): + reads.append(name) + return original_getattr(self, name) + + for name, method in originals.items(): + setattr(cupy.ndarray, name, counting(name, method)) + cupy.ndarray.__getattr__ = counting_getattr + try: + yield reads + finally: + for name, method in originals.items(): + setattr(cupy.ndarray, name, method) + cupy.ndarray.__getattr__ = original_getattr + + +def _make_particles(comm, bc): + return Particles6D( + comm_world=comm, + loading_params=LoadingParameters(Np=2000, seed=1234), + boundary_params=BoundaryParameters(bc=(bc, bc, bc)), + ) + + +def _shift_positions(markers: np.ndarray, valid: np.ndarray, step: int) -> np.ndarray: + """New positions in [0, 1) for the valid rows, far enough away that most markers change rank.""" + new = markers.copy() + shift = np.array([0.37, 0.21, 0.53]) * (step + 1) + new[valid, :3] = (markers[valid, :3] + shift) % 1.0 + return new + + +def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", count_reads: bool = True): + """Sort the same markers on NumPy and CuPy across all ranks and compare the marker arrays. + + ``cuda_aware=None`` uses the probe (:func:`cunumpy.mpi.mpi_is_cuda_aware`); ``True``/``False`` + force the device-buffer or the host-staged exchange. + """ + comm = MPI.COMM_WORLD + rank, size = comm.Get_rank(), comm.Get_size() + + ref = _make_particles(comm, bc) + ref.draw_markers(sort=False) + ref.mpi_sort_markers() + + # the reference is sorted on NumPy, the device markers inside `use_backend("cupy")` blocks + on_device = lambda: xp.use_backend("cupy") # noqa: E731 + + previous = xp.mpi.get_mpi_cuda_aware() + try: + with on_device(): + if cuda_aware is not None: + xp.mpi.set_mpi_cuda_aware(cuda_aware) + else: + xp.mpi.set_mpi_cuda_aware(None) + cuda_aware = xp.mpi.mpi_is_cuda_aware(comm) + + dev = _make_particles(comm, bc) + assert xp.get_array_backend(dev.markers) == "cupy" + assert dev.markers.shape == ref.markers.shape + + for step in range(2): + new = _shift_positions(ref.markers, ref.valid_mks, step) + + # number of rows this rank has to send, from the host reference + etas = new[ref.valid_mks, :3] + left, right = ref.domain_array[rank, 0::3], ref.domain_array[rank, 1::3] + right = np.where(right == 1.0, np.nextafter(right, 2.0), right) + n_send = int(np.count_nonzero(~np.all((etas >= left) & (etas < right), axis=1))) + n_send_total = comm.allreduce(n_send) + assert size == 1 or n_send_total > 0 + + ref.markers[:] = new + ref.update_holes() + n_before = ref.n_mks_loc + ref.mpi_sort_markers(do_test=True) + n_recv = ref.n_mks_loc - (n_before - n_send) + + with on_device(): + dev.markers[:] = xp.to_cupy(new) + dev.update_holes() + with xp.profiling.count_transfers() as transfers, _count_host_reads() as reads: + dev.mpi_sort_markers() + + if cuda_aware: + assert transfers.total == 0, transfers.report() + else: + # one staged copy per non-empty send buffer, one copy back of the receive buffer + n_dest = sum(1 for i in range(size) if i != rank and dev._send_list[i].shape[0] > 0) + assert transfers.to_host == n_dest, transfers.report() + assert transfers.to_device == (1 if n_recv > 0 else 0), transfers.report() + sent_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_host") + assert sent_bytes == n_send * ref.markers.shape[1] * 8 + + if count_reads and cuda_aware: + # (the staged exchange reads the send buffers with .get(), counted above) + assert reads == [], f"host reads of device arrays during the exchange: {reads}" + + with on_device(): + assert np.array_equal(xp.to_numpy(dev.markers), ref.markers) + assert np.array_equal(xp.to_numpy(dev.valid_mks), ref.valid_mks) + n_dev = dev.n_mks_loc + assert n_dev == ref.n_mks_loc + + with on_device(): + # all markers are on the right rank now: nothing moves + dev.mpi_sort_markers(do_test=True) + assert np.array_equal(xp.to_numpy(dev.markers), ref.markers) + finally: + xp.mpi.set_mpi_cuda_aware(previous) + + +@requires_fake_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("cuda_aware", [None, True, False]) +@pytest.mark.parametrize("bc", ["periodic", "remove"]) +def test_mpi_sort_markers_fake_cupy(cuda_aware, bc): + """On the fake CuPy: the exchange gives the NumPy result, without host/device copies when MPI is CUDA-aware.""" + check_mpi_sort_markers_device(cuda_aware, bc=bc) + + +@requires_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("bc", ["periodic", "remove"]) +def test_mpi_sort_markers_gpu(bc): + """On a GPU (``mpirun -n 2``, one GPU per rank): the exchange gives the NumPy result; with + CUDA-aware MPI without host/device copies.""" + check_mpi_sort_markers_device(None, bc=bc, count_reads=False) + + +if __name__ == "__main__": + check_mpi_sort_markers_device(None) From 1636822f3ddb7c5e37d8de2348ea99b488d96eda Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 19:12:45 +0200 Subject: [PATCH 13/50] CUDA version of the linear_vlasov_ampere accumulation - linear_vlasov_ampere_cuda.cu: same arguments as pyccel, one thread per marker row, matrix blocks as Array6D views (cunumpy 0.6.1) - device fill_mat, fill_mat_vec (filler_kernels.cuh), m_v_fill_b_v1_symm (particle_to_mat_kernels.cuh) and outer (linalg_kernels.cuh), with atomic adds - parity cases (Cuboid, Colella, HollowTorus, 3d spline) and a device-helper test of m_v_fill_b_v1_symm, on a GPU and by CPU emulation - test_accum_matrix_cupy.py: the matrix path of Accumulator on CuPy vs NumPy (device data, ghost regions, transposed blocks, Schur operator and solve), without a GPU on the fake CuPy with emulated kernel launches (emulated_launches() in pic/tests/cuda_emulation.py) - CUDA_STRATEGY.md: porting order and implementation notes Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 59 ++++- src/struphy/linear_algebra/linalg_kernels.cuh | 11 + .../pic/accumulation/filler_kernels.cuh | 125 ++++++++++- .../linear_vlasov_ampere_cuda.cu | 90 ++++++++ .../accumulation/particle_to_mat_kernels.cuh | 69 ++++++ src/struphy/pic/tests/cuda_emulation.py | 66 +++++- src/struphy/pic/tests/cuda_parity_cases.py | 17 ++ src/struphy/pic/tests/kernel_test_args.py | 27 ++- .../pic/tests/test_accum_matrix_cupy.py | 209 ++++++++++++++++++ src/struphy/pic/tests/test_cuda_emulation.py | 12 + src/struphy/pic/tests/test_cuda_parity.py | 4 +- src/struphy/pic/tests/test_device_helpers.py | 61 +++++ 12 files changed, 730 insertions(+), 20 deletions(-) create mode 100644 src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu create mode 100644 src/struphy/pic/tests/test_accum_matrix_cupy.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 68035f68c..5f8635336 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -26,7 +26,8 @@ Progress is tracked in struphy-hub/struphy#650. - [ ] **PR 17: struphy on the feectools CUDA stack** — the `feectools` submodule points at the top of the feectools CUDA stack (`cuda-4-device-kernels`, [feectools#88](https://github.com/struphy-hub/feectools/pull/88)) instead of `devel-tiny`, so the struphy CUDA PRs run against feectools with device stencil operations, MPI with device buffers and one GPU per rank. Moves along with the stack; before the struphy stack is merged into `devel`, the feectools stack is merged into `devel-tiny` and the submodule points there again (see [feectools](#feectools)). - [x] **PR 18: Geometry evaluation on the GPU for all analytic mappings** (#681, open) — CUDA versions of the four geometry entry kernels (`kernel_evaluate_pic`, `kernel_evaluate`, `kernel_pullpush_pic`, `kernel_pullpush` in `geometry/kernels/`), built on device versions of the whole metric chain (`f`, `df`, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`/`push`/`tran`) for every analytic mapping (`kind_map` 10–12, 20–22, 30–32). Restores CuPy particle runs (weight initialization evaluates `jacobian_det`), and removes the Cuboid-only checks in `Pusher`, the accumulators and `reflect`. Parity arguments cover every analytic mapping (see [PR 18](#pr-18-geometry-evaluation-for-all-analytic-mappings) and the [PR 18 implementation notes](#pr-18-implementation-notes)). - [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines and `EQDSKequilibrium` stay host-only (see [PR 19](#pr-19-spline-mappings) and the [PR 19 implementation notes](#pr-19-implementation-notes)). -- [ ] **Next (order to be confirmed)**: 6D array views in cunumpy and the blocked matrix accumulations `vlasov_maxwell`, `linear_vlasov_ampere` (steps 1–2 of the [porting order](#porting-order)); one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, including when to compile the kernels, which needs the [feectools stack](#feectools) merged and released first; a decision on hand-written CUDA vs. code generation before the guiding-center kernels (step 3). +- [x] **`linear_vlasov_ampere` on the GPU** (part of #688) — the first matrix accumulation: CUDA version of `linear_vlasov_ampere` (`EfieldWeightsCoupling`) on cunumpy's `Array6D` views (cunumpy 0.6.1), device versions of `fill_mat`, `fill_mat_vec`, `m_v_fill_b_v1_symm` and `outer`, and the matrix path of `Accumulator` checked on the CuPy backend (device data, ghost regions, transposed blocks, Schur operator) without a GPU through cunumpy's fake CuPy with emulated kernel launches (see the [implementation notes](#linear_vlasov_ampere-implementation-notes)). +- [ ] **Next (order to be confirmed)**: the blocked matrix accumulation `vlasov_maxwell` (step 1 of the [porting order](#porting-order)); one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, including when to compile the kernels, which needs the [feectools stack](#feectools) merged and released first; a decision on hand-written CUDA vs. code generation before the guiding-center kernels (step 3). - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)). - [x] **PR 13: Move the kernel infrastructure to cunumpy** — `Kernel`, `KernelCatalog`, `CudaKernel` and `Argument` are replaced by their cunumpy counterparts; each owner has a single `args_*` object on both backends, and `pusher_args.cuh` is generated (see [Moving to cunumpy](#moving-to-cunumpy-pr-13)). Kernels and device helpers only change their includes. - [ ] **CI**: a GPU runner that runs the CUDA tests (can happen any time; until then the GPU tests are run by hand on an H100 before each PR that touches CUDA code is merged). @@ -64,7 +65,7 @@ CUDA kernels can be added one by one. If the code runs on the GPU and needs a ke | `src/struphy/utils/cuda_arguments.py` | `CUDA_STRUCTS`, `CUDA_OPTIONS`, `write_pusher_header()` and `write_local_projectors_header()` | | `src/struphy/kernel_arguments/pusher_args.cuh`, `local_projectors_args.cuh` | the C structs `MarkerArgs`, `DerhamArgs`, `DomainArgs` and `LocalProjectorsArgs`, generated from the CUDA classes, using cunumpy array views | | `src/struphy/geometry/base.py`, `src/struphy/pic/base.py`, `src/struphy/feec/psydac_derham.py` | `Domain.args_domain`, `Particles.args_markers` and `Derham.args_derham` are the pyccel class on NumPy and the CUDA class on CuPy, chosen once at construction; every kernel call goes through a `Kernel` | -| `src/struphy/*/kernels/` (`pic/pushing`, `pic/accumulation`, `pic/diagnostics`, `pic/sph`, `bsplines`, `geometry`, `feec`, `feec/local_projectors`) | every kernel called from Python with argument objects, one folder each: 44 pusher/evaluation (incl. `reflect`), 16 accumulation, 10 marker diagnostics, 4 SPH evaluation, 3 spline evaluation, 4 geometry, 1 FEEC utility and 8 local projector kernels. Each folder's `__init__.py` declares its `Kernel`; the code imports it (`from struphy.geometry.kernels.kernel_evaluate import kernel_evaluate`). Fifteen have CUDA versions, with parity cases in `pic/tests/cuda_parity_cases.py`, among them the four geometry kernels (PR 18) | +| `src/struphy/*/kernels/` (`pic/pushing`, `pic/accumulation`, `pic/diagnostics`, `pic/sph`, `bsplines`, `geometry`, `feec`, `feec/local_projectors`) | every kernel called from Python with argument objects, one folder each: 44 pusher/evaluation (incl. `reflect`), 16 accumulation, 10 marker diagnostics, 4 SPH evaluation, 3 spline evaluation, 4 geometry, 1 FEEC utility and 8 local projector kernels. Each folder's `__init__.py` declares its `Kernel`; the code imports it (`from struphy.geometry.kernels.kernel_evaluate import kernel_evaluate`). Sixteen have CUDA versions (`linear_vlasov_ampere` the latest), with parity cases in `pic/tests/cuda_parity_cases.py`, among them the four geometry kernels (PR 18) | | `src/struphy/geometry/evaluation_kernels.cuh`, `transform_kernels.cuh`, `spline_mappings_kernels.cuh`, `geometry/domains//_cuda.cuh` | device versions of the mapping helpers for every mapping: `f`/`df` per analytic domain (`kind_map` 10–12, 20–22, 30–32, PR 18) and the spline mappings `spline_3d`, `spline_2d_straight`, `spline_2d_torus` (`kind_map` 0–2, PR 19), the `kind_map` switch, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`, `push`, `tran` | | `src/struphy/pic/tests/test_cuda_parity.py`, `test_cuda_emulation.py` | parity of every CUDA kernel with pyccel on the cases of `pic/tests/cuda_parity_cases.py` (`PARITY_CASES`): on a GPU (`assert_kernels_agree`), and without one by CPU emulation | | `src/struphy/pic/tests/test_kernel_backends.py` | the generated-header tests, the check that each CUDA class mirrors its pyccel class, signature check of all kernels, owner class selection, and (on a GPU) struct layout and copy/pickle checks | @@ -277,7 +278,7 @@ from the struphy tracking issue ([#650](https://github.com/struphy-hub/struphy/i Kernels are ported in the order that completes one model after the other on the GPU. Each step lists only the kernels it adds; P = pusher, A = accumulation, E = marker evaluation. Status: -done (✓), in this PR (PR 14), blocked (⏸). +done (✓), in a PR (PR 14), blocked (⏸). **Step 0 – Vlasov** (PRs 11–12, ✓): `push_eta_stage`, `push_vxb_analytic`, `push_vxb_implicit`, `reflect`. @@ -289,14 +290,14 @@ done (✓), in this PR (PR 14), blocked (⏸). |---|---|---|---|---| | 1 | `push_v_with_efield` | P | `VlasovAmpereCoupling`, `PushVinForceField` | PR 14 | | 2 | `charge_density_0form` | A (vector) | initial Poisson solve of these models | PR 14 | -| 3 | `vlasov_maxwell` | A (matrix + vector) | `VlasovAmpereCoupling` | ⏸ 6D views | +| 3 | `vlasov_maxwell` | A (matrix + vector) | `VlasovAmpereCoupling` | next (6D views available) | **Step 2 – LinearVlasovAmpère/Maxwell (δf)** | # | Kernel | Type | Used by | Status | |---|---|---|---|---| | 4 | `push_weights_with_efield_lin_va` | P | `EfieldWeightsCoupling` | PR 14 | -| 5 | `linear_vlasov_ampere` | A (matrix + vector) | `EfieldWeightsCoupling` | ⏸ 6D views | +| 5 | `linear_vlasov_ampere` | A (matrix + vector) | `EfieldWeightsCoupling` | ✓ (#688) | **Step 3 – Guiding center: ToyDrift, then DriftKineticElectrostaticAdiabatic.** The default algorithm is `discrete_gradient_1st_order`; its chain comes first. @@ -348,9 +349,9 @@ The three spaces (H1vec/Hcurl/Hdiv) differ only in the basis; port one, then the 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, -and array views with more than 4 dimensions in cunumpy (all matrix accumulations write 6D stencil -matrix data). +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. ## Testing @@ -662,3 +663,45 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a during the exchange, staged: one `to_host` per non-empty send and one `to_device`. Without a GPU it runs on the fake CuPy (`mpirun -n N env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi pic/tests/test_sorting_device.py`, in the PIC MPI workflow); the GPU variant has not run (no GPU here). +## `linear_vlasov_ampere` implementation notes + +The first matrix accumulation on the GPU (part of #688, tracked in #650): `linear_vlasov_ampere`, the accumulation of +`EfieldWeightsCoupling` in `LinearVlasovAmpereOneSpecies` (step 2 of the [porting order](#porting-order)). + +- **6D views.** The six blocks `mat11, mat12, mat13, mat22, mat23, mat33` are `Array6D` (cunumpy 0.6.1, strided, + like the `Array3D` vectors); `cunumpy.kernels.CudaKernel` packs the `StencilMatrix._data` blocks without + copying them. `f0_values` is a `const double*`, as in `push_weights_with_efield_lin_va`. +- **Device helpers, one to one.** `filler_kernels.cuh` gets `fill_mat` and `fill_mat_vec` next to `fill_vec`, + `particle_to_mat_kernels.cuh` gets `m_v_fill_b_v1_symm` next to `vec_fill_b_v0`, and `linalg_kernels.cuh` gets + `outer`; same names, arguments and loop order as pyccel (`starts` and `pads` are `const long long*`). Every addition + into matrix or vector data is a `cunumpy_atomic_add`, as in `charge_density_0form`; the spline values pyccel keeps + in `args_derham.bn1, ..., bd3` are in the thread-local `SplineScratch`. Of the remaining `m_v_fill_*`/`mat_fill_*` + fillers, each one is ported with the first kernel that calls it. +- **Kernel.** `linear_vlasov_ampere_cuda.cu`: same arguments in the same order as pyccel, one thread per marker row, + holes and boundary particles skipped; `df`, `matrix_inv`, `matrix_vector`, `outer`, then `m_v_fill_b_v1_symm`. The + 3 x 3 `filling_m` is nine row-major entries; the CUDA-only `factor` is pyccel's `f0_values[ip] / (markers[ip, 7] * + n_markers_tot)`, computed once as pyccel does. +- **`Accumulator` on CuPy.** Nothing in `Accumulator` had to change: on the CuPy backend `StencilMatrix._data` and + `StencilVector._data` are device arrays (feectools allocates them with `xp`), the accumulator zeroes them in place and + passes the same arrays to the kernel; `exchange_assembly_data`, `update_ghost_regions`, the transposed blocks of the + symmetric matrix (`StencilMatrix.transpose`, a CUDA kernel in feectools), the scaled operator of `SchurSolver` and its + `dot` all work on them. `test_accum_matrix_cupy.py` checks this against NumPy: the kernel data are the blocks' data + (no copies), the nine arrays, the dense blocks, the accumulated vector, `BC.dot(x)` and three CG iterations of the + Schur solve agree, and nothing is copied to the host from the accumulation to `BC.dot(x)`. On a GPU it runs with + `assert_no_transfers` (not run yet); without one it runs in a subprocess on cunumpy's fake CuPy, with every + `CudaKernel` launch emulated on the CPU in place on the fake device arrays (`emulated_launches()` in + `pic/tests/cuda_emulation.py`; struphy's and feectools' kernels, struct arguments included). This would fit into + cunumpy (fake CuPy with emulated launches). +- **Found, not fixed here.** feectools' `StencilVector.inner` returns its result on the host (`xp.to_numpy`), so each + inner product of the CG solve copies a scalar from the device; the Schur solve is therefore outside the transfer + guard. `MassMatrixPreconditioner`, the default preconditioner of `EfieldWeightsCoupling`, cannot be created on CuPy + (`is_circulant` asserts a host array; the device Kronecker solve is a separate PR). `LinearVlasovAmpereOneSpecies` + end to end on the GPU (#689) also needs the feectools stencil 6D views and device Kronecker solve PRs and the first + H100 run (#687). +- **Tests.** Parity cases `PARITY_CASES["linear_vlasov_ampere"]`: 129 markers (a hole and a boundary particle) in + Cuboid, Colella, HollowTorus and a 3d spline mapping, data shaped like the real accumulation matrices of clamped + splines of degrees 2, 3, 1 on 8 cells (`v1_symm_accumulation_data()` in `pic/tests/kernel_test_args.py`); + `rtol = 1e-12`, `atol = 1e-10` (atomics change the summation order; the entries reach 1e4). The CPU emulation agrees + with pyccel to round-off and catches a transposed `DF`. `m_v_fill_b_v1_symm` is compared with pyccel through a + wrapper kernel (`test_v1_symm_filler` on a GPU, `test_emulated_v1_symm_filler` by emulation). GPU tests have not run + on an H100. diff --git a/src/struphy/linear_algebra/linalg_kernels.cuh b/src/struphy/linear_algebra/linalg_kernels.cuh index 1b72c25fe..4636f2ea3 100644 --- a/src/struphy/linear_algebra/linalg_kernels.cuh +++ b/src/struphy/linear_algebra/linalg_kernels.cuh @@ -100,4 +100,15 @@ __device__ inline void transpose(const double* a, double* b) { for (int i = 0; i < 3; ++i) for (int j = 0; j < 3; ++j) b[3 * i + j] = a[3 * j + i]; } +/** + * Outer product of two vectors of length 3, as in linalg_kernels.outer. + * + * @param a First input vector with three entries. + * @param b Second input vector with three entries. + * @param c Output matrix c_ij = a_i * b_j with nine row-major entries; must not overlap a or b. + */ +__device__ inline void outer(const double* a, const double* b, double* c) { + for (int i = 0; i < 3; ++i) + for (int j = 0; j < 3; ++j) c[3 * i + j] = a[i] * b[j]; +} } diff --git a/src/struphy/pic/accumulation/filler_kernels.cuh b/src/struphy/pic/accumulation/filler_kernels.cuh index 37f3867d9..8a33d7773 100644 --- a/src/struphy/pic/accumulation/filler_kernels.cuh +++ b/src/struphy/pic/accumulation/filler_kernels.cuh @@ -1,5 +1,5 @@ -// Device versions of the fillers in filler_kernels.py: add the contribution of one marker to a stencil vector. -// Many markers write to the same entries, so every addition is atomic. +// Device versions of the fillers in filler_kernels.py: add the contribution of one marker to a stencil vector +// or to a block of a stencil matrix. Many markers write to the same entries, so every addition is atomic. #pragma once #include "cunumpy/array_view.cuh" #include "cunumpy/atomic.cuh" @@ -40,4 +40,125 @@ __device__ inline void fill_vec(int pi1, int pi2, int pi3, const double* bi1, co } } } + +/** + * Add basis functions times filling to a block of a stencil matrix, as in filler_kernels.fill_mat. + * + * @param pi1 Spline degree of the codomain (row indices) along the first axis. + * @param pi2 Spline degree of the codomain along the second axis. + * @param pi3 Spline degree of the codomain along the third axis. + * @param pj1 Spline degree of the domain (column indices) along the first axis. + * @param pj2 Spline degree of the domain along the second axis. + * @param pj3 Spline degree of the domain along the third axis. + * @param bi1 The pi1 + 1 non-vanishing N- or D-spline values of the codomain along the first axis. + * @param bi2 The pi2 + 1 non-vanishing spline values of the codomain along the second axis. + * @param bi3 The pi3 + 1 non-vanishing spline values of the codomain along the third axis. + * @param bj1 The pj1 + 1 non-vanishing spline values of the domain along the first axis. + * @param bj2 The pj2 + 1 non-vanishing spline values of the domain along the second axis. + * @param bj3 The pj3 + 1 non-vanishing spline values of the domain along the third axis. + * @param span1 Knot span index along the first axis. + * @param span2 Knot span index along the second axis. + * @param span3 Knot span index along the third axis. + * @param starts Start indices of the codomain (three entries). + * @param pads Paddings of the codomain (three entries). + * @param mat Stencil matrix data (the _data of a StencilMatrix, shape (n1, n2, n3, 2 * pads + 1)), any strides; + * written atomically. + * @param filling Number multiplied by the basis functions. + * + * The summation order of concurrent threads (and the last bits of the result) depends on the thread schedule. + */ +__device__ inline void fill_mat(int pi1, int pi2, int pi3, int pj1, int pj2, int pj3, const double* bi1, + const double* bi2, const double* bi3, const double* bj1, const double* bj2, + const double* bj3, int span1, int span2, int span3, const long long* starts, + const long long* pads, Array6D mat, double filling) { + for (int il1 = 0; il1 <= pi1; ++il1) { + long long i1 = span1 + il1 - starts[0]; + double b1 = bi1[il1] * filling; + for (int il2 = 0; il2 <= pi2; ++il2) { + long long i2 = span2 + il2 - starts[1]; + double b2 = b1 * bi2[il2]; + for (int il3 = 0; il3 <= pi3; ++il3) { + long long i3 = span3 + il3 - starts[2]; + double b3 = b2 * bi3[il3]; + + for (int jl1 = 0; jl1 <= pj1; ++jl1) { + long long j1 = pads[0] + jl1 - il1; + double b4 = b3 * bj1[jl1]; + for (int jl2 = 0; jl2 <= pj2; ++jl2) { + long long j2 = pads[1] + jl2 - il2; + double b5 = b4 * bj2[jl2]; + for (int jl3 = 0; jl3 <= pj3; ++jl3) { + long long j3 = pads[2] + jl3 - il3; + double b6 = b5 * bj3[jl3]; + + cunumpy_atomic_add(&mat(i1, i2, i3, j1, j2, j3), b6); + } + } + } + } + } + } +} + +/** + * Add basis functions times fillings to a block of a stencil matrix and to a stencil vector, as in + * filler_kernels.fill_mat_vec. + * + * @param pi1 Spline degree of the codomain (row indices) along the first axis. + * @param pi2 Spline degree of the codomain along the second axis. + * @param pi3 Spline degree of the codomain along the third axis. + * @param pj1 Spline degree of the domain (column indices) along the first axis. + * @param pj2 Spline degree of the domain along the second axis. + * @param pj3 Spline degree of the domain along the third axis. + * @param bi1 The pi1 + 1 non-vanishing N- or D-spline values of the codomain along the first axis. + * @param bi2 The pi2 + 1 non-vanishing spline values of the codomain along the second axis. + * @param bi3 The pi3 + 1 non-vanishing spline values of the codomain along the third axis. + * @param bj1 The pj1 + 1 non-vanishing spline values of the domain along the first axis. + * @param bj2 The pj2 + 1 non-vanishing spline values of the domain along the second axis. + * @param bj3 The pj3 + 1 non-vanishing spline values of the domain along the third axis. + * @param span1 Knot span index along the first axis. + * @param span2 Knot span index along the second axis. + * @param span3 Knot span index along the third axis. + * @param starts Start indices of the codomain (three entries). + * @param pads Paddings of the codomain (three entries). + * @param mat Stencil matrix data (the _data of a StencilMatrix), any strides; written atomically. + * @param filling_mat Number multiplied by the basis functions and added to mat. + * @param vec Stencil vector data of the codomain (the _data of a StencilVector), any strides; written atomically. + * @param filling_vec Number multiplied by the basis functions of the codomain and added to vec. + */ +__device__ inline void fill_mat_vec(int pi1, int pi2, int pi3, int pj1, int pj2, int pj3, const double* bi1, + const double* bi2, const double* bi3, const double* bj1, const double* bj2, + const double* bj3, int span1, int span2, int span3, const long long* starts, + const long long* pads, Array6D mat, double filling_mat, + Array3D vec, double filling_vec) { + for (int il1 = 0; il1 <= pi1; ++il1) { + long long i1 = span1 + il1 - starts[0]; + double b1 = bi1[il1]; + for (int il2 = 0; il2 <= pi2; ++il2) { + long long i2 = span2 + il2 - starts[1]; + double b2 = b1 * bi2[il2]; + for (int il3 = 0; il3 <= pi3; ++il3) { + long long i3 = span3 + il3 - starts[2]; + double b3 = b2 * bi3[il3]; + + cunumpy_atomic_add(&vec(i1, i2, i3), b3 * filling_vec); + + for (int jl1 = 0; jl1 <= pj1; ++jl1) { + long long j1 = pads[0] + jl1 - il1; + double b4 = b3 * bj1[jl1] * filling_mat; + for (int jl2 = 0; jl2 <= pj2; ++jl2) { + long long j2 = pads[1] + jl2 - il2; + double b5 = b4 * bj2[jl2]; + for (int jl3 = 0; jl3 <= pj3; ++jl3) { + long long j3 = pads[2] + jl3 - il3; + double b6 = b5 * bj3[jl3]; + + cunumpy_atomic_add(&mat(i1, i2, i3, j1, j2, j3), b6); + } + } + } + } + } + } +} } diff --git a/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu b/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu new file mode 100644 index 000000000..387013af2 --- /dev/null +++ b/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu @@ -0,0 +1,90 @@ +#include "struphy/geometry/evaluation_kernels.cuh" +#include "struphy/linear_algebra/linalg_kernels.cuh" +#include "struphy/pic/accumulation/particle_to_mat_kernels.cuh" + +/** + * Accumulate the matrix and vector of the linear Vlasov-Ampere system (delta-f) into V1, as in + * linear_vlasov_ampere_kernels.linear_vlasov_ampere. + * + * Filling functions (the factors alpha^2 kappa^2 / v_th^2 and alpha^2 kappa are applied by the propagator): + * A_p^{mu, nu} = f_0(eta_p, v_p) / (N s_0) [DF^{-1}(eta_p) v_p]_mu [DF^{-1}(eta_p) v_p]_nu, + * B_p^mu = w_p [DF^{-1}(eta_p) v_p]_mu; + * one thread per marker row. + * + * @param args_markers Marker buffer (n_markers x n_cols, row-major) and the total marker number Np; the weight is in + * column 6 and s_0 in column 7, as in pyccel. + * @param args_derham Spline degrees (1 to 8), knots and start indices. + * @param args_domain Mapping arguments; every spline and analytic mapping. + * @param mat11 (1,1)-block data of the symmetric V1 -> V1 block matrix (the _data of a StencilMatrix, 6D, any + * strides); written with atomic additions. + * @param mat12 (1,2)-block data; written atomically. + * @param mat13 (1,3)-block data; written atomically. + * @param mat22 (2,2)-block data; written atomically. + * @param mat23 (2,3)-block data; written atomically. + * @param mat33 (3,3)-block data; written atomically. + * @param vec1 First component of the V1 vector data; written atomically. + * @param vec2 Second component of the V1 vector data; written atomically. + * @param vec3 Third component of the V1 vector data; written atomically. + * @param f0_values Value of f0 for each marker row (n_markers entries). + * + * Holes (markers[ip, 0] == -1) and boundary particles (markers[ip, -1] == -2) are skipped. The pyccel loop runs over + * all rows (shape(markers)[0]), which is n_markers here. Many threads add to the same entries, so the summation + * order (and the last bits of the result) differs from the serial pyccel loop. + */ +extern "C" __global__ void linear_vlasov_ampere(MarkerArgs args_markers, DerhamArgs args_derham, DomainArgs args_domain, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + Array3D vec1, Array3D vec2, Array3D vec3, + const double* f0_values) { + // CUDA-only: the marker row of this thread replaces the pyccel loop variable + int ip = blockDim.x * blockIdx.x + threadIdx.x; + + // get number of markers + int n_markers = args_markers.n_markers; + if (ip >= n_markers) return; + + // total number of markers (weights are w_p = delta f_p / (N * s_0)) + int n_markers_tot = args_markers.Np; + + // only do something if particle is a "true" particle (i.e. not a hole) + if (args_markers.markers(ip, 0) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) + return; + + // allocate for metric coeffs (3 x 3 matrices as nine row-major entries) + double dfm[9], df_inv[9]; + + // allocate for filling + double v[3], df_inv_v[3], filling_m[9], filling_v[3]; + + // marker positions + double eta1 = args_markers.markers(ip, 0); + double eta2 = args_markers.markers(ip, 1); + double eta3 = args_markers.markers(ip, 2); + + // get velocity + v[0] = args_markers.markers(ip, 3); + v[1] = args_markers.markers(ip, 4); + v[2] = args_markers.markers(ip, 5); + + // evaluate Jacobian, result in dfm + struphy_cuda::df(eta1, eta2, eta3, args_domain, dfm); + + // invert Jacobian matrix + struphy_cuda::matrix_inv(dfm, df_inv); + + // compute DF^{-1} v + struphy_cuda::matrix_vector(df_inv, v, df_inv_v); + + // filling_m = alpha^2 * kappa^2 * f0 / (N * s_0 * v_th^2) * (DF^{-1} v_p)_mu * (DF^{-1} v_p)_nu + struphy_cuda::outer(df_inv_v, df_inv_v, filling_m); + double factor = f0_values[ip] / (args_markers.markers(ip, 7) * n_markers_tot); + for (int k = 0; k < 9; ++k) filling_m[k] *= factor; + + // filling_v = alpha^2 * kappa * w_p * DL^{-1} * v_p + for (int k = 0; k < 3; ++k) filling_v[k] = args_markers.markers(ip, 6) * df_inv_v[k]; + + // call the appropriate matvec filler + struphy_cuda::m_v_fill_b_v1_symm(args_derham, eta1, eta2, eta3, mat11, mat12, mat13, mat22, mat23, mat33, + filling_m[0], filling_m[1], filling_m[2], filling_m[4], filling_m[5], + filling_m[8], vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2]); +} diff --git a/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh b/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh index 7cabaee45..ade422b97 100644 --- a/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh +++ b/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh @@ -28,4 +28,73 @@ __device__ inline void vec_fill_b_v0(const DerhamArgs& args_derham, double eta1, fill_vec(pn1, pn2, pn3, scratch.bn1, scratch.bn2, scratch.bn3, span1, span2, span3, args_derham.starts, vec, fill); } + +/** + * Add one marker's contribution to the symmetric blocks (1,1), (1,2), (1,3), (2,2), (2,3), (3,3) of a V1 -> V1 + * block matrix and to a V1 vector, as in particle_to_mat_kernels.m_v_fill_b_v1_symm. + * + * @param args_derham Spline degrees (1 to 8), knots and start indices; pn is also the padding of the blocks. + * @param eta1 Logical position along the first axis. + * @param eta2 Logical position along the second axis. + * @param eta3 Logical position along the third axis. + * @param mat11 (mu=1, nu=1)-block data (D N N x D N N), shape (n1, n2, n3, 2 * pn + 1); written atomically. + * @param mat12 (mu=1, nu=2)-block data (D N N x N D N); written atomically. + * @param mat13 (mu=1, nu=3)-block data (D N N x N N D); written atomically. + * @param mat22 (mu=2, nu=2)-block data (N D N x N D N); written atomically. + * @param mat23 (mu=2, nu=3)-block data (N D N x N N D); written atomically. + * @param mat33 (mu=3, nu=3)-block data (N N D x N N D); written atomically. + * @param fill11 Number multiplied by the basis functions and added to mat11. + * @param fill12 Number multiplied by the basis functions and added to mat12. + * @param fill13 Number multiplied by the basis functions and added to mat13. + * @param fill22 Number multiplied by the basis functions and added to mat22. + * @param fill23 Number multiplied by the basis functions and added to mat23. + * @param fill33 Number multiplied by the basis functions and added to mat33. + * @param vec1 mu=1 component of the V1 vector data (D N N); written atomically. + * @param vec2 mu=2 component of the V1 vector data (N D N); written atomically. + * @param vec3 mu=3 component of the V1 vector data (N N D); written atomically. + * @param fill1 Number multiplied by the basis functions and added to vec1. + * @param fill2 Number multiplied by the basis functions and added to vec2. + * @param fill3 Number multiplied by the basis functions and added to vec3. + * + * Pyccel keeps the spline values in args_derham.bn1, ..., bd3; here they are in a thread-local SplineScratch. + */ +__device__ inline void m_v_fill_b_v1_symm(const DerhamArgs& args_derham, double eta1, double eta2, double eta3, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + double fill11, double fill12, double fill13, double fill22, + double fill23, double fill33, Array3D vec1, Array3D vec2, + Array3D vec3, double fill1, double fill2, double fill3) { + // degrees of the basis functions : B-splines (pn) and D-splines (pd) + int pn1 = args_derham.pn[0]; + int pn2 = args_derham.pn[1]; + int pn3 = args_derham.pn[2]; + + int pd1 = pn1 - 1; + int pd2 = pn2 - 1; + int pd3 = pn3 - 1; + + // CUDA-only scratch holds the spans and the spline values pyccel keeps in args_derham + SplineScratch scratch; + get_spans(eta1, eta2, eta3, args_derham, scratch); + int span1 = scratch.span1, span2 = scratch.span2, span3 = scratch.span3; + + // fill matrix entries + fill_mat_vec(pd1, pn2, pn3, pd1, pn2, pn3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bd1, scratch.bn2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat11, fill11, vec1, fill1); + + fill_mat_vec(pn1, pd2, pn3, pn1, pd2, pn3, scratch.bn1, scratch.bd2, scratch.bn3, scratch.bn1, scratch.bd2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat22, fill22, vec2, fill2); + + fill_mat_vec(pn1, pn2, pd3, pn1, pn2, pd3, scratch.bn1, scratch.bn2, scratch.bd3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat33, fill33, vec3, fill3); + + fill_mat(pd1, pn2, pn3, pn1, pd2, pn3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bn1, scratch.bd2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat12, fill12); + + fill_mat(pd1, pn2, pn3, pn1, pn2, pd3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat13, fill13); + + fill_mat(pn1, pd2, pn3, pn1, pn2, pd3, scratch.bn1, scratch.bd2, scratch.bn3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat23, fill23); +} } diff --git a/src/struphy/pic/tests/cuda_emulation.py b/src/struphy/pic/tests/cuda_emulation.py index 0e8db4f12..ca3c76d01 100644 --- a/src/struphy/pic/tests/cuda_emulation.py +++ b/src/struphy/pic/tests/cuda_emulation.py @@ -6,8 +6,10 @@ the fields of each struct as separate parameters, rebuilds the structs and calls the real kernel. """ +import contextlib + import numpy as np -from cunumpy.kernel_testing import emulate_cuda_kernel +from cunumpy.kernel_testing import emulate_cuda_kernel, fake_cupy_active from cunumpy.kernels import CudaKernel from struphy.utils.cuda_arguments import CUDA_OPTIONS @@ -55,11 +57,12 @@ def _struct_field(value, name): return getattr(value, name) -def emulate_struct_kernel(kernel: CudaKernel, *args, n_threads=None): +def emulate_struct_kernel(kernel: CudaKernel, *args, n_threads=None, grid=None, block=None): """Emulate `kernel` with host arguments; struct arguments are the pyccel argument objects (NumPy arrays). - Arrays (also those inside the argument objects) are updated in place, as by a launch. Without `n_threads`, - one thread per row of the first array is launched, as in a real launch. + Arrays (also those inside the argument objects) are updated in place, as by a launch. Without `n_threads` or + `grid`, one thread per row of the first array is launched, as in a real launch; `grid` and `block` give the + launch shape explicitly. """ flat = [] for p, value in zip(kernel.signature, args): @@ -67,9 +70,58 @@ def emulate_struct_kernel(kernel: CudaKernel, *args, n_threads=None): flat.append(value) else: flat += [_struct_field(value, f.name) for f in p.struct.fields] - if n_threads is None: - n_threads = next(value.shape[0] for value in flat if isinstance(value, np.ndarray)) - grid, block = kernel.launch_shape(n_threads) + if grid is None: + if n_threads is None: + n_threads = next(value.shape[0] for value in flat if isinstance(value, np.ndarray)) + grid, block = kernel.launch_shape(n_threads, block=block) source, name = _wrapper(kernel) wrapper = CudaKernel(source, name, source_dir=kernel.source_dir, **CUDA_OPTIONS) emulate_cuda_kernel(wrapper, *flat, grid=grid, block=block, options=EMULATION_OPTIONS) + + +def _host(value): + """The host buffer behind a fake CuPy array (the same memory, no copy); other values unchanged.""" + if isinstance(value, np.ndarray) or not hasattr(value, "__cuda_array_interface__"): + return value + return object.__getattribute__(value, "_a") + + +class _HostFields: + """A CUDA argument object (``Cuda*Arguments``) whose array attributes are their fake CuPy host buffers.""" + + def __init__(self, args): + self._args = args + + def __getattr__(self, name): + return _host(getattr(self._args, name)) + + +@contextlib.contextmanager +def emulated_launches(): + """On cunumpy's fake CuPy, run every ``CudaKernel`` launch by CPU emulation, in place on the fake device arrays. + + The fake CuPy (``CUNUMPY_FAKE_CUPY=1``) keeps device arrays in host memory and rejects host/device mixing, but + cannot launch kernels. Inside this context a launch emulates the kernel (struct arguments included, through + :func:`emulate_struct_kernel`) on the host buffers of the fake arrays, so code that launches CUDA kernels (struphy's + and feectools') runs end to end on the CuPy backend without a GPU. Each launch compiles the kernel, so this is slow. + """ + if not fake_cupy_active(): + raise RuntimeError("emulated_launches() needs cunumpy's fake CuPy (CUNUMPY_FAKE_CUPY=1)") + original = CudaKernel.__call__ + + def launch(self, *args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None): + grid, block = self.launch_shape(n_threads, grid=grid, block=block, args=args) + if any(p.struct is not None for p in self.signature): + host_args = [_host(a) if p.struct is None else _HostFields(a) for p, a in zip(self.signature, args)] + emulate_struct_kernel(self, *host_args, grid=grid, block=block) + else: + host_args = [_host(a) for a in args] + emulate_cuda_kernel( + self, *host_args, grid=grid, block=block, shared_mem=shared_mem, options=EMULATION_OPTIONS + ) + + CudaKernel.__call__ = launch + try: + yield + finally: + CudaKernel.__call__ = original diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index e2ce0a5d5..85380b6e2 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -29,6 +29,7 @@ marker_arguments, spline_coefficients, spline_evaluation_arguments, + v1_symm_accumulation_data, ) @@ -79,6 +80,19 @@ def charge_density_0form_args(bc): return (args_markers, derham_arguments(), args_domain, xp.zeros((18, 20, 16))) +def linear_vlasov_ampere_args(domain): + """129 markers (a hole and a boundary particle) in domain `domain` of geometry_domain(); zeroed V1 data.""" + args_markers, _ = marker_arguments((0, 0, 0)) + f0_values = xp.asarray(np.random.default_rng(14).random(N_MARKERS)) + return ( + args_markers, + derham_arguments(), + geometry_domain(domain).args_domain, + *v1_symm_accumulation_data(), + f0_values, + ) + + # ---------------------------------------------------------------- spline evaluation @@ -267,6 +281,9 @@ def size_of(index, per_thread=1): "reflect": ParityCases((0, 1, 2), reflect_args, n_threads=size_of(2)), # atomics add in another order than the serial loop "charge_density_0form": ParityCases(((0, 0, 0), (2, 0, 1)), charge_density_0form_args, rtol=1e-12, atol=1e-13), + # Cuboid, Colella and HollowTorus (non-diagonal DF) and a 3d spline mapping, by index into geometry_domain(); + # atomics add in another order than the serial loop, and the entries reach 1e4 (atol for cancelling sums) + "linear_vlasov_ampere": ParityCases((0, 2, 5, 12), linear_vlasov_ampere_args, rtol=1e-12, atol=1e-10), "eval_spline_mpi_markers": ParityCases(SPLINE_KINDS, eval_spline_mpi_markers_args), "eval_spline_mpi_matrix": ParityCases(SPLINE_KINDS, eval_spline_mpi_grid_args(sparse=False), n_threads=size_of(-1)), "eval_spline_mpi_sparse_meshgrid": ParityCases( diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index d413b5380..eaed48adf 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -49,14 +49,37 @@ def butcher_arguments(method): return xp.asarray(butcher.a_stage), xp.asarray(butcher.b), xp.asarray(butcher.c), butcher.n_stages +DERHAM_DEGREE = (2, 3, 1) +DERHAM_CELLS = 8 + + def derham_arguments(): """Splines of degrees 2, 3, 1 on 8 cells, starting at index 0.""" - degree = np.array([2, 3, 1], dtype=np.int64) - knots = [np.r_[np.zeros(p), np.linspace(0, 1, 9), np.ones(p)] for p in degree] + degree = np.array(DERHAM_DEGREE, dtype=np.int64) + knots = [np.r_[np.zeros(p), np.linspace(0, 1, DERHAM_CELLS + 1), np.ones(p)] for p in degree] args_class = CudaDerhamArguments if xp.get_backend() == "cupy" else DerhamArguments return args_class(xp.asarray(degree), *(xp.asarray(t) for t in knots), xp.zeros(3, dtype=np.int64)) +def v1_symm_accumulation_data(): + """Zeroed data of a symmetric V1 -> V1 accumulation matrix and a V1 vector, for :func:`derham_arguments`. + + Returns the six independent blocks ``mat11, mat12, mat13, mat22, mat23, mat33`` and the three vector components + ``vec1, vec2, vec3`` on the active backend, shaped like the ``_data`` of the ``StencilMatrix`` blocks and + ``StencilVector`` components that ``Accumulator(..., "Hcurl", ..., symmetry="symm", add_vector=True)`` passes on + one process for clamped splines of degrees 2, 3, 1 on 8 cells: component ``mu`` has D-splines (``n - 1`` + functions) along axis ``mu`` and N-splines (``n = 8 + p``) along the others, every axis is padded by ``p`` on both + sides, and the matrix blocks have ``2 p + 1`` diagonals per axis (the paddings are the degrees of V0). + """ + degree = np.array(DERHAM_DEGREE) + n_spline = DERHAM_CELLS + degree + rows = [tuple(n_spline - (np.arange(3) == mu) + 2 * degree) for mu in range(3)] + diagonals = tuple(2 * degree + 1) + mats = tuple(xp.zeros(rows[mu] + diagonals) for mu, nu in ((0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2))) + vecs = tuple(xp.zeros(rows[mu]) for mu in range(3)) + return mats + vecs + + def spline_coefficients(n=3, seed=11): """Random coefficients, large enough for every span of :func:`derham_arguments`.""" rng = np.random.default_rng(seed) diff --git a/src/struphy/pic/tests/test_accum_matrix_cupy.py b/src/struphy/pic/tests/test_accum_matrix_cupy.py new file mode 100644 index 000000000..cf0aa62c6 --- /dev/null +++ b/src/struphy/pic/tests/test_accum_matrix_cupy.py @@ -0,0 +1,209 @@ +"""The matrix path of :class:`~struphy.pic.accumulation.particles_to_grid.Accumulator` on the CuPy backend. + +``linear_vlasov_ampere`` accumulates the symmetric V1 -> V1 block matrix and the V1 vector of ``EfieldWeightsCoupling``; +the matrix data (``StencilMatrix._data``) must be device arrays that are zeroed, passed to the CUDA kernel and turned +into the Schur operator without host copies. Checked against the NumPy backend on a GPU and, without one, on +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 + +import cunumpy as xp +import numpy as np +import pytest +from cunumpy.kernel_testing import emulation_compiler, requires_cupy + +from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env + +ALPHA, KAPPA, VTH, DT = 1.3, 0.7, 0.9, 0.05 +N_SCHUR_ITERATIONS = 3 + + +@dataclass +class MarkerOwner: + """What the accumulators read from ``Particles``: the marker arguments, no clones. + + The same markers on both backends (``Particles`` draws them with the backend's random generator). + """ + + args_markers: Any + clone_config: Any = None + + +def accumulate_linear_vlasov_ampere(guard=None): + """Accumulate ``linear_vlasov_ampere`` like ``EfieldWeightsCoupling`` on the active backend, then solve with it. + + Colella mapping (non-diagonal Jacobian), clamped splines in eta1 and periodic ones in eta2, eta3, so the three + V1 components have different sizes. `guard` (a context manager factory, e.g. one that fails on host transfers) + encloses the accumulation, the scaled operator and its application, as in a time step of the propagator; the + Schur solve is not guarded (see below). + + Returns a dict of host arrays: the nine kernel data arrays (``mat11`` ... ``vec3``) after the accumulation (ghost + regions exchanged), the assembled block matrix as dense blocks, the accumulated vector after the boundary + operators, ``BC.dot(x)`` with the operator scaled as in ``EfieldWeightsCoupling`` and the Schur solution after + ``N_SCHUR_ITERATIONS`` iterations of the conjugate gradient method. + """ + import contextlib + + from feectools.linalg.stencil import StencilMatrix + + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.psydac_derham import Derham + from struphy.geometry.domains import Colella + from struphy.io.options import DerhamOptions + from struphy.linear_algebra.schur_solver import SchurSolver + from struphy.linear_algebra.solver import SolverParameters + from struphy.pic.accumulation.kernels.linear_vlasov_ampere import linear_vlasov_ampere + from struphy.pic.accumulation.particles_to_grid import Accumulator + from struphy.pic.tests.kernel_test_args import N_MARKERS, marker_arguments + from struphy.topology.grids import TensorProductGrid + + guard = guard or contextlib.nullcontext + derham = Derham( + TensorProductGrid(num_elements=(6, 5, 4)), + DerhamOptions(degree=(2, 3, 1), bcs=(("free", "free"), None, None)), + comm=None, + ) + domain = Colella(Lx=2.0, Ly=3.0, alpha=0.05, Lz=4.0) + mass_ops = WeightedMassOperators(derham, domain) + args_markers, _ = marker_arguments((0, 0, 0)) + accum = Accumulator( + MarkerOwner(args_markers), + "Hcurl", + linear_vlasov_ampere, + mass_ops, + domain.args_domain, + add_vector=True, + symmetry="symm", + ) + matrix = accum.operators[0].matrix + + # the kernel writes into the data of the operator's blocks (and vector), not into copies + blocks = [matrix[a, b] for a in range(3) for b in range(a, 3)] + assert all(isinstance(block, StencilMatrix) for block in blocks) + assert all(data is block._data for data, block in zip(accum._args_data[:6], blocks)) + assert all(data is vec._data for data, vec in zip(accum._args_data[6:], accum._vectors[0].blocks)) + assert all(xp.is_gpu(data) == (xp.get_backend() == "cupy") for data in accum._args_data) + + rng = np.random.default_rng(16) + f0_values = xp.asarray(rng.random(N_MARKERS)) + for data in accum._args_data: + data[...] = 1.0 # left over from a previous step: the accumulator zeroes it + + x = accum._vectors[0].space.zeros() + for block in x.blocks: + block._data[...] = xp.asarray(rng.normal(size=block._data.shape)) + x.update_ghost_regions() + + schur_solver = SchurSolver( + mass_ops.M1, + ALPHA**2 * KAPPA**2 * accum.operators[0] / (4 * VTH**2), + "pcg", + # a few iterations (each one launches nine stencil dot kernels, which the emulation compiles one by one) + solver_params=SolverParameters(tol=1e-30, maxiter=N_SCHUR_ITERATIONS), + ) + + with guard(): + accum(f0_values) + # as in EfieldWeightsCoupling.__call__ + schur_solver.BC = accum.operators[0] + schur_solver.BC *= (-1) * ALPHA**2 * KAPPA**2 / (4 * VTH**2) + bc_x = schur_solver.BC.dot(x) + byn = accum.vectors[0].copy() + byn *= ALPHA**2 * KAPPA / 2.0 + # not guarded: feectools' StencilVector.inner returns its result on the host (xp.to_numpy), so every inner + # product of the conjugate gradient method copies one scalar from the device + solution, info = schur_solver(xn=x, Byn=byn, dt=DT) + + out = {name: xp.to_numpy(data) for name, data in zip(MATRIX_DATA_NAMES, accum._args_data)} + for a in range(3): + for b in range(3): + out[f"block{a}{b}"] = xp.to_numpy(matrix[a, b].toarray()) + out["vector"] = xp.to_numpy(accum.vectors[0].toarray()) + out["bc_x"] = xp.to_numpy(bc_x.toarray()) + out["solution"] = xp.to_numpy(solution.toarray()) + out["niter"] = np.array(info["niter"]) + return out + + +MATRIX_DATA_NAMES = ("mat11", "mat12", "mat13", "mat22", "mat23", "mat33", "vec1", "vec2", "vec3") + + +def compare_with_numpy(result): + """Compare a CuPy-backend `result` of :func:`accumulate_linear_vlasov_ampere` with the NumPy backend's.""" + with xp.use_backend("numpy"): + expected = accumulate_linear_vlasov_ampere() + assert expected["niter"] == result["niter"] == N_SCHUR_ITERATIONS + assert np.any(expected["mat12"] != 0.0) and np.any(expected["vec3"] != 0.0) + # the transposed blocks of the symmetric matrix are filled from the accumulated ones + np.testing.assert_allclose(expected["block10"], expected["block01"].T, rtol=1e-14, atol=1e-14) + for name, value in expected.items(): + scale = max(np.max(np.abs(value)), 1.0) + # atomics add in another order than the serial loop; the solve repeats it + np.testing.assert_allclose(result[name], value, rtol=1e-10, atol=1e-12 * scale, err_msg=name) + + +@requires_cupy +def test_accumulator_matrix_on_cupy(): + """On a GPU: the accumulated matrix and vector, the Schur operator and solve agree with NumPy; the accumulation and + the operator run without host/device transfers.""" + from cunumpy.profiling import assert_no_transfers + + with xp.use_backend("cupy"): + accumulate_linear_vlasov_ampere() # warm up (NVRTC compilation) + result = accumulate_linear_vlasov_ampere(guard=assert_no_transfers) + compare_with_numpy(result) + + +def fake_cupy_transfer_guard(): + """On cunumpy's fake CuPy: fail on every explicit host copy of a device array (``.get()``, ``cupy.asnumpy``). + + The fake CuPy already rejects implicit conversions (``numpy.asarray`` of a device array, mixing host and device + arrays); this closes the explicit ones, which are legal but are transfers inside a time step. + """ + import contextlib + + import cupy + + @contextlib.contextmanager + def guard(): + def fail(*args, **kwargs): + raise AssertionError("host copy of a device array during the accumulation") + + saved = cupy.ndarray.get, cupy.asnumpy + cupy.ndarray.get, cupy.asnumpy = fail, fail + try: + yield + finally: + cupy.ndarray.get, cupy.asnumpy = saved + + return guard + + +def check_accumulator_matrix_fake_cupy(): + """Run in a child process with ``CUNUMPY_FAKE_CUPY=1``: the CuPy backend with emulated kernel launches.""" + from struphy.pic.tests.cuda_emulation import emulated_launches + + with emulated_launches(), xp.use_backend("cupy"): + result = accumulate_linear_vlasov_ampere(guard=fake_cupy_transfer_guard()) + compare_with_numpy(result) + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +@pytest.mark.skipif(emulation_compiler() is None, reason="no C++ compiler") +def test_accumulator_matrix_fake_cupy(): + """Without a GPU: the matrix path of the Accumulator on the CuPy backend (fake CuPy, emulated launches). + + Device data are zeroed and passed to the CUDA kernel, the ghost regions, the transposed blocks, the scaled + operator and the Schur solve run on device arrays, and no array is copied to the host from the accumulation to + the operator's application. Runs in a + 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:] diff --git a/src/struphy/pic/tests/test_cuda_emulation.py b/src/struphy/pic/tests/test_cuda_emulation.py index 793bb9525..f2ab111a2 100644 --- a/src/struphy/pic/tests/test_cuda_emulation.py +++ b/src/struphy/pic/tests/test_cuda_emulation.py @@ -113,3 +113,15 @@ def test_emulated_transform_helpers(domain_index): out = np.zeros(len(expected)) emulate_struct_kernel(kernel, *inputs, args_domain, out, len(out), n_threads=len(out)) np.testing.assert_allclose(out, expected, rtol=1e-10, atol=1e-10) + + +@requires_compiler +def test_emulated_v1_symm_filler(): + """test_device_helpers.test_v1_symm_filler by CPU emulation: the device m_v_fill_b_v1_symm agrees with pyccel.""" + from struphy.pic.tests.test_device_helpers import v1_symm_filler_case + + kernel, (eta, fills), args_derham, data, expected = v1_symm_filler_case() + emulate_struct_kernel(kernel, eta, fills, args_derham, *data, len(eta), n_threads=len(eta)) + for result, reference in zip(data, expected): + assert np.any(reference != 0.0) + np.testing.assert_allclose(result, reference, rtol=1e-12, atol=1e-13) diff --git a/src/struphy/pic/tests/test_cuda_parity.py b/src/struphy/pic/tests/test_cuda_parity.py index 75d53245f..0b0a14c9e 100644 --- a/src/struphy/pic/tests/test_cuda_parity.py +++ b/src/struphy/pic/tests/test_cuda_parity.py @@ -10,6 +10,7 @@ from cunumpy.profiling import assert_no_transfers from struphy.pic.accumulation.kernels.charge_density_0form import charge_density_0form +from struphy.pic.accumulation.kernels.linear_vlasov_ampere import linear_vlasov_ampere from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage from struphy.pic.pushing.kernels.push_v_with_efield import push_v_with_efield from struphy.pic.pushing.kernels.push_vxb_analytic import push_vxb_analytic @@ -108,7 +109,7 @@ def device_only(value, *args, **kwargs): def test_vlasov_kernel_coverage(): - """Vlasov (PushEta, both PushVxB algorithms) and the ported Vlasov-Ampere kernels have device kernels.""" + """Vlasov (PushEta, both PushVxB algorithms) and the ported (linear) Vlasov-Ampere kernels have device kernels.""" for kernel in ( push_eta_stage, push_vxb_analytic, @@ -116,5 +117,6 @@ def test_vlasov_kernel_coverage(): push_v_with_efield, push_weights_with_efield_lin_va, charge_density_0form, + linear_vlasov_ampere, ): assert kernel.has_cuda, kernel.name diff --git a/src/struphy/pic/tests/test_device_helpers.py b/src/struphy/pic/tests/test_device_helpers.py index a2390ef4d..b771d8efc 100644 --- a/src/struphy/pic/tests/test_device_helpers.py +++ b/src/struphy/pic/tests/test_device_helpers.py @@ -391,3 +391,64 @@ def test_transform_helpers(domain_index): kernel, inputs, _, expected = transform_helper_case(domain_index) result = _run_on_device(kernel, inputs, domain_index, len(expected)) np.testing.assert_allclose(result, expected, rtol=1e-10, atol=1e-10) + + +# CUDA-only wrapper: one thread per point calls m_v_fill_b_v1_symm with the point's nine fillings +V1_SYMM_FILLER_SOURCE = r""" +#include "struphy/pic/accumulation/particle_to_mat_kernels.cuh" +extern "C" __global__ void fill_v1_symm(const double* eta, const double* fills, DerhamArgs args_derham, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + Array3D vec1, Array3D vec2, Array3D vec3, int n) { + int ip = blockDim.x * blockIdx.x + threadIdx.x; + if (ip >= n) return; + const double* fill = fills + 9 * ip; + struphy_cuda::m_v_fill_b_v1_symm(args_derham, eta[3 * ip], eta[3 * ip + 1], eta[3 * ip + 2], mat11, mat12, mat13, + mat22, mat23, mat33, fill[0], fill[1], fill[2], fill[3], fill[4], fill[5], vec1, + vec2, vec3, fill[6], fill[7], fill[8]); +} +""" + + +def v1_symm_filler_case(): + """Wrapper kernel, host inputs and pyccel reference for the device m_v_fill_b_v1_symm. + + 65 points (the corners of the unit cube among them) with random fillings, accumulated into zeroed V1 data shaped + like the real accumulation matrices (:func:`~struphy.pic.tests.kernel_test_args.v1_symm_accumulation_data`). + Used by :func:`test_v1_symm_filler` (GPU) and by its CPU emulation in ``test_cuda_emulation.py``. Returns + ``kernel, (eta, fills), args_derham, data, expected``: ``args_derham`` is the pyccel class and ``data`` the nine + zeroed host arrays in kernel order. + """ + from struphy.pic.accumulation import particle_to_mat_kernels + from struphy.pic.tests.kernel_test_args import derham_arguments, v1_symm_accumulation_data + + rng = np.random.default_rng(15) + eta = np.r_[np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0], [0.0, 1.0, 0.5]]), rng.random((62, 3))] + fills = rng.normal(size=(len(eta), 9)) + with cunumpy.use_backend("numpy"): + args_derham = derham_arguments() + data = v1_symm_accumulation_data() + expected = tuple(a.copy() for a in data) + for point, fill in zip(eta, fills): + particle_to_mat_kernels.m_v_fill_b_v1_symm( + args_derham, *point, *expected[:6], *fill[:6], *expected[6:], *fill[6:] + ) + kernel = CudaKernel(V1_SYMM_FILLER_SOURCE, "fill_v1_symm", **CUDA_OPTIONS) + return kernel, (eta, fills), args_derham, data, expected + + +@requires_cupy +def test_v1_symm_filler(): + """m_v_fill_b_v1_symm (and the fill_mat/fill_mat_vec it calls) on the device agrees with pyccel.""" + import cupy as cp + + from struphy.pic.tests.kernel_test_args import derham_arguments + + kernel, (eta, fills), _, data, expected = v1_symm_filler_case() + with cunumpy.use_backend("cupy"): + args_derham = derham_arguments() + device_data = [cp.asarray(a) for a in data] + kernel(cp.asarray(eta), cp.asarray(fills), args_derham, *device_data, len(eta), n_threads=len(eta)) + for result, reference in zip(device_data, expected): + assert np.any(reference != 0.0) + np.testing.assert_allclose(result.get(), reference, rtol=1e-12, atol=1e-13) From 9da31a6a458842071d406ed59a08b7f235e9280f Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 23:35:08 +0200 Subject: [PATCH 14/50] Mark every column of the test hole so accumulation kernels skip it The hole in marker_arguments() was only marked in column 8 (first_init_idx), which the pushers test. Accumulation kernels test column 0, so the linear_vlasov_ampere parity case had no real hole. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/kernel_test_args.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index eaed48adf..24e597e89 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -29,7 +29,9 @@ def marker_arguments(bc): markers[:, 3:6] = rng.uniform(-2, 2, (N_MARKERS, 3)) markers[:, 8:14] = markers[:, :6] markers[:, 18:21] = 0.0 - markers[0, 8] = -1.0 + # Row 0 is a hole. Struphy marks every column of a hole with -1: pushers test column first_init_idx (8), + # accumulation kernels test column 0. + markers[0, :] = -1.0 markers[1, -1] = -2.0 valid = np.ones(N_MARKERS, dtype=bool) valid[:2] = False From b11ee1869896f114955733bed3f583012e07ca74 Mon Sep 17 00:00:00 2001 From: Max Date: Thu, 8 Oct 2026 13:16:00 +0200 Subject: [PATCH 15/50] GVEC: marker evaluation, cunumpy 0.6.2 host helpers, robust fake-CuPy child tests - GVECequilibrium evaluates markers point by point: the coordinates go to gvec as xarray DataArrays with a shared dimension, so bv/jv at N markers no longer return N x N x N arrays and absB0 at markers works (also on NumPy). With use_boozer=True, marker evaluation raises, since gvec computes the Boozer transform per flux surface. New test_gvec_equil.py. - Use xp.host_call, xp.evaluate_on_host and xp.setup_on_host from cunumpy 0.6.2 instead of the copies in fields_background/base.py; pin cunumpy 0.6.2. - run_fake_cupy_child: under mpirun only rank 0 starts the fake-CuPy child (every rank started the same child at once, and on CI one of two concurrent GVEC children died with SIGILL and no output); the child runs with faulthandler and one OpenMP thread, and failures report the signal and the end of stdout/stderr. Used by the four fake-CuPy subprocess tests. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 14 +++- pyproject.toml | 2 +- src/struphy/fields_background/base.py | 80 ------------------- src/struphy/fields_background/equils.py | 51 +++++++----- .../tests/test_equils_cupy.py | 15 +--- .../tests/test_gvec_equil.py | 47 +++++++++++ src/struphy/geometry/tests/test_domain.py | 34 +++++++- .../pic/tests/test_accum_matrix_cupy.py | 9 +-- .../tests/test_compile_cuda_kernels.py | 13 +-- 9 files changed, 126 insertions(+), 139 deletions(-) create mode 100644 src/struphy/fields_background/tests/test_gvec_equil.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 8fcc4368b..462d18830 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -382,17 +382,25 @@ views with more than 4 dimensions (all matrix accumulations write 6D stencil mat - **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.** `setup_on_host` (`fields_background/base.py`) runs these `__init__`s on the NumPy backend, so the +- **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 `host_call`, and the GVEC/DESC `bv`, `jv`, - `p0`, `n0`, `gradB1` through `@evaluate_on_host`: device arguments are copied to the host, evaluated on the NumPy +- **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 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/base.py b/src/struphy/fields_background/base.py index 7e894b366..426fd8f28 100644 --- a/src/struphy/fields_background/base.py +++ b/src/struphy/fields_background/base.py @@ -1,11 +1,9 @@ "Base classes for MHD equilibria." -import functools import logging from abc import ABCMeta, abstractmethod import cunumpy as xp -import numpy as np from pyevtk.hl import gridToVTK from struphy.geometry.base import Domain @@ -20,84 +18,6 @@ logger = logging.getLogger("struphy") -def _on_device(arg) -> bool: - """Whether ``arg`` is (or contains, for tuples and lists) a CuPy array.""" - if isinstance(arg, (tuple, list)): - return any(_on_device(a) for a in arg) - return xp.is_gpu(arg) - - -def _to_host(arg): - """Copy CuPy arrays in ``arg`` (an array, or a tuple/list of them) to the host; leave everything else alone.""" - if isinstance(arg, (tuple, list)): - return type(arg)(_to_host(a) for a in arg) - return xp.to_numpy(arg) if xp.is_gpu(arg) else arg - - -def _to_device(arg): - """Copy NumPy arrays in ``arg`` (an array, or a tuple/list of them) to the device; leave scalars alone.""" - if isinstance(arg, (tuple, list)): - return type(arg)(_to_device(a) for a in arg) - return xp.to_cupy(arg) if isinstance(arg, np.ndarray) else arg - - -def host_call(fun, *args, **kwargs): - """Call a host-only function (SciPy spline, external equilibrium code) with arguments of any backend. - - Device (CuPy) arrays among the arguments are copied to the host, ``fun`` runs on the NumPy backend, and its array - results are copied back to the device, once per call. Without device arguments, ``fun`` is called directly on the - NumPy backend (on the NumPy backend, this is a plain call), so the result lives where the arguments live. - - The copies cost one host round trip per call; equilibria are evaluated at setup (initial conditions, projections - onto the FEEC spaces, mappings), not in the time loop, which uses the projected equilibrium. - - Parameters - ---------- - fun : callable - The host-only function. - - *args, **kwargs - Arguments of ``fun``; arrays (or tuples/lists of arrays) of either backend, or scalars. - - Returns - ------- - The result of ``fun`` (an array, a tuple/list of arrays or a scalar), with arrays on the device if any argument - was on the device. - """ - device = _on_device(args) or _on_device(list(kwargs.values())) - if xp.get_backend() == "numpy" and not device: - return fun(*args, **kwargs) - with xp.use_backend("numpy"): - out = fun(*_to_host(args), **{k: _to_host(v) for k, v in kwargs.items()}) - return _to_device(out) if device else out - - -def evaluate_on_host(method): - """Decorator for equilibrium methods that can only be evaluated on the host (see :func:`host_call`).""" - - @functools.wraps(method) - def wrapper(self, *args, **kwargs): - return host_call(method, self, *args, **kwargs) - - return wrapper - - -def setup_on_host(init): - """Decorator for ``__init__`` of equilibria whose setup is host-only (file reading, ODE solves, SciPy fits). - - ``__init__`` runs on the NumPy backend, so the equilibrium holds only host data (NumPy arrays, SciPy splines, - floats) on either backend. Its evaluation follows the backend of the arguments; host-only parts of it go through - :func:`host_call`. - """ - - @functools.wraps(init) - def wrapper(self, *args, **kwargs): - with xp.use_backend("numpy"): - init(self, *args, **kwargs) - - return wrapper - - class FluidEquilibrium(metaclass=ABCMeta): """ Abstract base class for callable fluid equilibria on arbitrary domains. diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index 576c918c0..2377cebdc 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -27,9 +27,6 @@ NumericalFluidEquilibrium, NumericalFluidEquilibriumWithB, NumericalMHDequilibrium, - evaluate_on_host, - host_call, - setup_on_host, ) from struphy.fields_background.mhd_equil.eqdsk import readeqdsk from struphy.io.options import BaseUnits @@ -877,7 +874,7 @@ def doc_formula(cls): Units are those defned in the parameter file (through :class:`~struphy.io.options.BaseUnits`). """ - @setup_on_host + @xp.setup_on_host def __init__( self, a: float = 1.0, @@ -1025,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 = host_call(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)): @@ -1154,7 +1151,7 @@ def p_r(self, r): # alternative profiles (interpolated) elif self.params["q_kind"] == 1 or self.params["q_kind"] == 2: - pout = host_call(self._p_i, r) + pout = xp.host_call(self._p_i, r) # remove all "dimensions" for point-wise evaluation if isinstance(r, (int, float)): @@ -1370,7 +1367,7 @@ def doc_formula(cls): Units are those defned in the parameter file (through :class:`~struphy.io.options.BaseUnits`). """ - @setup_on_host + @xp.setup_on_host def __init__( self, a: float = 0.361925, @@ -1476,7 +1473,7 @@ def psi_r(self, r, der=0): assert der >= 0 and der <= 2, "Only first and second derivatives available!" - out = host_call(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)): @@ -1671,7 +1668,7 @@ class EQDSKequilibrium(AxisymmMHDequilibrium): Struphy base units. If None, no rescaling of output is performed. """ - @setup_on_host + @xp.setup_on_host def __init__( self, rel_path: bool = True, @@ -1871,7 +1868,7 @@ def psi_axis_RZ(self): def q_psi(self, psi, der=0): """Safety factor q = q(psi).""" - out = host_call(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)): @@ -1882,7 +1879,7 @@ def q_psi(self, psi, der=0): def g_psi(self, psi, der=0): """Toroidal field function g = g(psi).""" - out = host_call(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)): @@ -1893,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 = host_call(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)): @@ -1929,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 = host_call(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: @@ -2128,7 +2125,7 @@ def units(self) -> Units: """All Struphy units.""" return self._units - @evaluate_on_host + @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.""" @@ -2150,7 +2147,7 @@ def bv(self, *etas, squeeze_out=False): return out - @evaluate_on_host + @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.""" @@ -2180,7 +2177,7 @@ def jv(self, *etas, squeeze_out=False): return out - @evaluate_on_host + @xp.evaluate_on_host @profile def p0(self, *etas, squeeze_out=False): """0-form equilibrium pressure on logical cube [0, 1]^3.""" @@ -2200,7 +2197,7 @@ def p0(self, *etas, squeeze_out=False): return self.params["p0"] + tmp / self.units.p - @evaluate_on_host + @xp.evaluate_on_host @profile def n0(self, *etas, squeeze_out=False): """0-form equilibrium density on logical cube [0, 1]^3.""" @@ -2280,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) @@ -2418,7 +2425,7 @@ def units(self) -> Units: """All Struphy units.""" return self._units - @evaluate_on_host + @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.""" @@ -2492,7 +2499,7 @@ def _eval_bv(self, *etas, squeeze_out=False): return out - @evaluate_on_host + @xp.evaluate_on_host @profile def jv(self, *etas, squeeze_out=False): """Contra-variant (vector field) current density (=curl B) @@ -2568,7 +2575,7 @@ def _eval_jv(self, *etas, squeeze_out=False): return out - @evaluate_on_host + @xp.evaluate_on_host @profile def p0(self, *etas, squeeze_out=False): """0-form equilibrium pressure on logical cube [0, 1]^3 in Pascal.""" @@ -2600,7 +2607,7 @@ def p0(self, *etas, squeeze_out=False): return out - @evaluate_on_host + @xp.evaluate_on_host @profile def n0(self, *etas, squeeze_out=False): """0-form equilibrium density on logical cube [0, 1]^3.""" @@ -2625,7 +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 - @evaluate_on_host + @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 index f7ced7543..e7710570d 100644 --- a/src/struphy/fields_background/tests/test_equils_cupy.py +++ b/src/struphy/fields_background/tests/test_equils_cupy.py @@ -7,15 +7,13 @@ import contextlib import importlib.util import inspect -import subprocess -import sys import cunumpy import numpy as np import pytest from cunumpy.kernel_testing import 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 # (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. @@ -240,12 +238,10 @@ def test_all_equilibria_have_cases(): def test_host_call_numpy_is_plain_call(): - """On the NumPy backend, :func:`host_call` calls the function directly and returns its result unchanged.""" - from struphy.fields_background.base import host_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 = host_call(lambda x, y=1.0: x * y, a, y=2.0) + 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) @@ -262,10 +258,7 @@ def test_equil_fake_cupy(case): "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})" ) - 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/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/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 From 7d18d8cb01cec3e862f47719a57c97a5f8659128 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 11:29:47 +0200 Subject: [PATCH 16/50] Remove the numba code-generation prototype (#690 decided: hand-written CUDA only) The decision for #690 is made: kernels are not generated from the pyccel source. Remove pic/tests/codegen_spike and shorten the decision section in CUDA_STRATEGY.md. The hand-written port of bstar_parallel_3form and its parity cases stay. Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 108 +---- .../pic/tests/codegen_spike/__init__.py | 1 - .../pic/tests/codegen_spike/compare.py | 61 --- .../pic/tests/codegen_spike/numba_codegen.py | 458 ------------------ .../tests/codegen_spike/test_numba_codegen.py | 29 -- 5 files changed, 26 insertions(+), 631 deletions(-) delete mode 100644 src/struphy/pic/tests/codegen_spike/__init__.py delete mode 100644 src/struphy/pic/tests/codegen_spike/compare.py delete mode 100644 src/struphy/pic/tests/codegen_spike/numba_codegen.py delete mode 100644 src/struphy/pic/tests/codegen_spike/test_numba_codegen.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 8988497d1..775e97382 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -27,7 +27,7 @@ Progress is tracked in struphy-hub/struphy#650. - [x] **PR 18: Geometry evaluation on the GPU for all analytic mappings** (#681, open) — CUDA versions of the four geometry entry kernels (`kernel_evaluate_pic`, `kernel_evaluate`, `kernel_pullpush_pic`, `kernel_pullpush` in `geometry/kernels/`), built on device versions of the whole metric chain (`f`, `df`, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`/`push`/`tran`) for every analytic mapping (`kind_map` 10–12, 20–22, 30–32). Restores CuPy particle runs (weight initialization evaluates `jacobian_det`), and removes the Cuboid-only checks in `Pusher`, the accumulators and `reflect`. Parity arguments cover every analytic mapping (see [PR 18](#pr-18-geometry-evaluation-for-all-analytic-mappings) and the [PR 18 implementation notes](#pr-18-implementation-notes)). - [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines and `EQDSKequilibrium` stay host-only (see [PR 19](#pr-19-spline-mappings) and the [PR 19 implementation notes](#pr-19-implementation-notes)). - [x] **`linear_vlasov_ampere` on the GPU** (part of #688) — the first matrix accumulation: CUDA version of `linear_vlasov_ampere` (`EfieldWeightsCoupling`) on cunumpy's `Array6D` views (cunumpy 0.6.1), device versions of `fill_mat`, `fill_mat_vec`, `m_v_fill_b_v1_symm` and `outer`, and the matrix path of `Accumulator` checked on the CuPy backend (device data, ghost regions, transposed blocks, Schur operator) without a GPU through cunumpy's fake CuPy with emulated kernel launches (see the [implementation notes](#linear_vlasov_ampere-implementation-notes)). -- [ ] **Next (order to be confirmed)**: one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, which needs the [feectools stack](#feectools) merged and released first; the guiding-center kernels (step 3), by hand (decided in #690, see [Hand-written CUDA vs. code generation](#hand-written-cuda-vs-code-generation-decision-for-690)). +- [ ] **Next (order to be confirmed)**: one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, which needs the [feectools stack](#feectools) merged and released first; the guiding-center kernels (step 3), by hand (decided in #690, see [Hand-written CUDA](#hand-written-cuda-decision-for-690)). - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)). - [x] **PR 13: Move the kernel infrastructure to cunumpy** — `Kernel`, `KernelCatalog`, `CudaKernel` and `Argument` are replaced by their cunumpy counterparts; each owner has a single `args_*` object on both backends, and `pusher_args.cuh` is generated (see [Moving to cunumpy](#moving-to-cunumpy-pr-13)). Kernels and device helpers only change their includes. - [ ] **CI**: a GPU runner that runs the CUDA tests (can happen any time; until then the GPU tests are run by hand on an H100 before each PR that touches CUDA code is merged). @@ -316,7 +316,7 @@ algorithm is `discrete_gradient_1st_order`; its chain comes first. | 16–17 | `push_gc_bxEstar_discrete_gradient_2nd_order`, `push_gc_Bstar_discrete_gradient_2nd_order` | P | | 18–19 | `push_gc_bxEstar_explicit_multistage`, `push_gc_Bstar_explicit_multistage` | P | -These are the largest kernels. They are ported by hand like the others (decision for #690, see [Hand-written CUDA vs. code generation](#hand-written-cuda-vs-code-generation-decision-for-690)); `bstar_parallel_3form` (7) has its CUDA version. +These are the largest kernels. They are ported by hand like the others (decision for #690, see [Hand-written CUDA](#hand-written-cuda-decision-for-690)); `bstar_parallel_3form` (7) has its CUDA version. **Step 4 – Hybrid MHD–kinetic 6D (current and pressure coupling)** @@ -367,7 +367,7 @@ views with more than 4 dimensions (all matrix accumulations write 6D stencil mat - **Accumulation strategy.** Atomics vs. sort-then-reduce; see PR 12+. Decided by measurement on the first accumulation kernel. - **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.** Decided for now (#690): hand-written CUDA stays the default; see [Hand-written CUDA vs. code generation](#hand-written-cuda-vs-code-generation-decision-for-690) for the comparison and for what would change the decision. +- **Single-source alternatives.** Decided (#690): no code generation; every CUDA kernel is written by hand, see [Hand-written CUDA](#hand-written-cuda-decision-for-690). - **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). Needed @@ -402,90 +402,34 @@ views with more than 4 dimensions (all matrix accumulations write 6D stencil mat 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). -## Hand-written CUDA vs. code generation (decision for #690) +## Hand-written CUDA (decision for #690) -**Decision: the guiding-center and hybrid kernels (steps 3–5 of the [porting order](#porting-order)) are ported by -hand, as `_cuda.cu` next to the pyccel kernel, like every kernel so far.** None of the generators is adopted. -This is decided again if a condition under [What would change the decision](#what-would-change-the-decision) -comes true. +**Decision: all CUDA kernels are written by hand, as `_cuda.cu` next to the pyccel kernel, like every kernel +so far. Kernels are not generated from the pyccel source (no numba, `cupyx.jit` or other code generation).** To decide, one guiding-center kernel, `bstar_parallel_3form` (step 3, #7: drift-kinetic B*∥ as a 3-form; mapping -Jacobian, two 0-form spline evaluations, weighted average of old and new marker positions), was ported in two ways: - -- **By hand:** `pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu`, one to one with pyccel, with - 14 parity cases in `pic/tests/cuda_parity_cases.py`: one per mapping (all analytic and four spline mappings), - with α = 1, 0 and mixed, so the evaluation point wraps around in `mod(eta, 1)`, a negative `eta_n`, a hole - and two output columns. This port stays in the code. -- **Generated:** a prototype generator, `pic/tests/codegen_spike/numba_codegen.py`, that translates the pyccel - source of the kernel *and of every helper it calls* into one numba.cuda module (`compare.py` runs it, - `test_numba_codegen.py` tests it; skipped without numba). It only exists for this comparison and can be - deleted. There is no `kernels` in the module names, so `struphy compile` does not pick it up. - -### Results - -| | Hand-written CUDA | Generated (numba.cuda) | -|---|---|---| -| Code for this kernel | 75 lines `.cu` (36 without comments), plus one new 5-line helper `eval_0form_spline_mpi` in `bsplines/evaluation_kernels_3d.cuh`; everything else is reused (`df` for 14 mappings, `det`, `get_spans`, `eval_spline_mpi_kernel`) | 0 lines by hand, but a 458-line generator. The generated module has 490 lines and contains its own copy of the whole helper chain (24 device functions); every generated kernel repeats it | -| Agreement with pyccel (14 cases) | CPU emulation of the real `.cu` source: all agree, max abs. difference 5.7·10⁻¹⁴. Removing the `mod` from the CUDA kernel makes the emulation test fail (115 of 3225 values) | numba.cuda **simulator**: all agree, max 6.3·10⁻¹³; the same translation compiled with `numba.njit`: the same values. The simulator interprets the Python one thread at a time and compiles nothing | -| Compiles for a GPU? | Not checked here (no GPU); the same NVRTC path as the 15 kernels before it | Not checked (numba.cuda needs NVVM from the CUDA toolkit, which macOS does not have). Not checked either: named tuples of local arrays in device code, and the 38 flat kernel parameters | -| Argument classes | Takes `MarkerArgs`, `DomainArgs`, `DerhamArgs` by value, like every kernel; cunumpy's `CudaKernel` checks the call against the signature | numba does not accept the C structs. The generator turns each argument object into one kernel parameter per struct field (38 parameters) and rebuilds the objects as named tuples inside the kernel. Calling it needs a second launch path next to cunumpy's (flattening, numba's array conversion, its own compile cache) | -| 1:1 with pyccel | Names, arguments and order as in pyccel, kept by review | Exactly 1:1 for the parts the generator understands; whatever it does not understand fails with `NotImplementedError` (see the last row) | -| Parity tests, emulation | The existing tests run unchanged: `PARITY_CASES` entry, GPU parity with `assert_kernels_agree`, CPU emulation, which compiles the `.cu` itself | The parity cases are reused, but the simulator checks the Python, not the code the GPU will run. cunumpy's emulation and `assert_kernels_agree` do not apply to numba kernels | -| Debugging | A C++ compiler and printf in emulation; `-lineinfo`, cunumpy's debug mode, compute-sanitizer and Nsight on the GPU; the source you read is the source you run | pdb in the simulator, where Python float arithmetic hides bugs specific to the GPU. On the GPU, errors point into generated code. The generator's own mistakes are a new kind of bug (see below) | -| Expected performance (no GPU numbers possible here) | Per-thread arrays with fixed size (`MAX_SPLINE_DEGREE`), `__device__ inline` helpers; atomics where needed (`cunumpy_atomic_add`) | About the same local memory (the same `MAX_LOCAL` bounds; `b_der_splines_slim` has the same 9×9 `values`), plus loops that fill `zeros` arrays. numba applies Python's wrap-around for negative indices to every array access (one extra compare/select each), passes shape, strides and data pointer per array, and has a higher launch overhead than `RawKernel`. The generator cannot tell which writes need atomics: a translated accumulation kernel would race | -| Dependencies | None beyond CuPy | numba + llvmlite + numba-cuda (≈175 MB installed), and numba only supports a range of NumPy versions (numba 0.68 was needed for NumPy 2.4). The numba.cuda that comes with numba is deprecated in favour of NVIDIA's numba-cuda | -| The other 59 pusher and accumulation kernels | Each is a transcription onto existing helpers; new helpers as needed (e.g. 1-/2-form evaluations exist, `eval_0form_spline_mpi` was added here) | Tried on all 60 entry kernels (`.claude/scripts/codegen_probe.py`, not committed): 31 translate, 29 do not (2D slice assignments such as `c[:, :] = ...` in the linear-algebra helpers, `numpy.abs` in the SPH helpers). "Translates" does not mean "is correct": none of them ran, and the accumulation kernels would race | - -What the generator had to change in the pyccel code, and the risks that come with it: - -- `empty`/`zeros` with a size known only at run time (`zeros(pn)`) become per-thread arrays of `MAX_LOCAL` - entries. This silently changes `len(x)`/`x.shape` of these arrays, e.g. `diff[:] = diff * pn` in - `b_der_splines_slim` now loops over 9 entries instead of `pn` (harmless there, because the extra entries are - zero). -- Array expressions on slices (`eta[:] = alpha[:3] * eta_k + ...`) become loops over the elements, because numba.cuda - cannot allocate the temporaries. Expressions with 2D slices are not supported yet. -- `numpy.mod` has to be translated exactly as pyccel compiles it (Fortran `MODULO`, `a - floor(a / b) * b`). -- The scratch arrays of `DerhamArguments` (`bn1`, ..., `bd3`), which every marker shares in pyccel, have to become - per-thread arrays; the hand-written `SplineScratch` does the same. -- The marker loop becomes one thread per marker; `continue` becomes `return`. - -**The other generators.** `cupyx.jit` (CuPy's experimental Python-to-CUDA transpiler) would need the same -translation of the pyccel source: its reference lists elementary scalar operations, thread indices, shared memory -and atomics, but no per-thread arrays and no struct arguments, and "CuPy's array operations ... are not supported". -It adds no dependency, but it was not run here (no CuPy). **pyccel 2.2.1 has no CUDA backend** (`--language` is -Fortran, C or Python). The CUDA work happens in the separate fork [pyccel-cuda](https://github.com/pyccel/pyccel-cuda) -(`@kernel`/`@device`, last commit January 2025); its open issues still include thread indexing, passing arrays as -arguments and device arrays, and nothing there handles class arguments like `MarkerArguments`. - -### Reasoning +Jacobian, two 0-form spline evaluations, weighted average of old and new marker positions), was ported by hand +(`pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu`, one to one with pyccel) and through a +numba.cuda prototype generator. The prototype was removed after the decision. The reasons: 1. **The shared part is already written.** What a generator would save is mostly the helper chain (mappings, - metric, splines, linear algebra, boundary conditions), and PRs 10–19 have already ported it by hand and tested - it, every mapping included. What is left per kernel is its body: here 36 lines that read like the pyccel code. -2. **A generator is a second compiler to maintain.** The prototype has 458 lines for a subset of what the kernels - use, and its mistakes (wrong array lengths, missing atomics) are silent. pyccel already is the compiler for the - CPU; a second one written by ourselves only pays off for many more kernels than ~45. -3. **It does not fit the infrastructure.** Everything since PR 8 (structs, signature checks in `CudaKernel`, - compile cache, debug mode, `assert_kernels_agree`, CPU emulation of the real source, no host transfers) is - built for CUDA C++ source. numba kernels would need a second launch and test path in cunumpy and struphy. -4. **It is harder to test.** The CPU emulation compiles the `.cu` that will run on the GPU; the numba simulator - only interprets Python. The hand-written port is gated better without a GPU than the generated one. -5. **Drift between pyccel and CUDA** (the main argument for one source) is caught by the parity cases: a change - to a pyccel kernel that is not made in its `.cu` fails `test_emulated_parity` in the normal CI. - -### What would change the decision - -- **pyccel gets a CUDA backend** in a release (not the fork) that accepts our argument classes and arrays. Then the - `.cu` files could be generated by pyccel in place, behind the same `Kernel`, structs and tests; the parity cases - would check them unchanged. This is the most attractive single-source route; check pyccel's releases before - step 4. -- **Porting by hand becomes the bottleneck**, e.g. new models add many kernels, or the pyccel kernels change often. - Then the next step is a generator that writes **`.cu` text** (C++) instead of numba: the AST transformations of - the prototype (helper calls, local arrays, slice loops, marker loop to threads, `mod`) carry over, and the output - goes through the existing structs, `CudaKernel`, emulation and parity tests. Accumulation kernels would still need - their atomics marked by hand. -- **GPU measurements** that show a generated kernel to be as fast as the hand-written one would remove the - performance argument; they do not touch points 2–4. + metric, splines, linear algebra, boundary conditions), which PRs 10–19 have already ported and tested by hand. + What is left per kernel is its body: for `bstar_parallel_3form` 36 lines that read like the pyccel code, plus one + 5-line helper (`eval_0form_spline_mpi`). +2. **A generator is a second compiler to maintain**, and its mistakes are silent (array lengths that change when a + size is known only at run time; it cannot tell which writes need atomics, so accumulation kernels would race). + The prototype translated only 31 of the 60 pusher and accumulation entry kernels. +3. **It does not fit the infrastructure.** Structs, signature checks in `CudaKernel`, the compile cache, debug + mode, `assert_kernels_agree`, the CPU emulation of the real source and the transfer checks are all built for + CUDA C++ source. numba does not accept the C structs and would need a second launch and test path, plus about + 175 MB of dependencies. +4. **Drift between pyccel and CUDA** (the main argument for one source) is caught by the parity cases: a change to a + pyccel kernel that is not made in its `.cu` fails `test_emulated_parity` in the normal CI. + +The port of `bstar_parallel_3form` stays, with 14 parity cases in `pic/tests/cuda_parity_cases.py`: one per mapping +(all analytic and four spline mappings), α = 1, 0 and mixed, so the evaluation point wraps around in `mod(eta, 1)` +(Fortran `MODULO`, `eta - floor(eta)` in CUDA), a negative `eta_n`, a hole and two output columns. The CPU emulation +agrees with pyccel to 5.7·10⁻¹⁴. ## PR 10 implementation notes diff --git a/src/struphy/pic/tests/codegen_spike/__init__.py b/src/struphy/pic/tests/codegen_spike/__init__.py deleted file mode 100644 index d1b189c74..000000000 --- a/src/struphy/pic/tests/codegen_spike/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Prototype for issue #690 (hand-written CUDA vs. code generation); not used by struphy, see ``numba_codegen``.""" diff --git a/src/struphy/pic/tests/codegen_spike/compare.py b/src/struphy/pic/tests/codegen_spike/compare.py deleted file mode 100644 index ae72c1a6f..000000000 --- a/src/struphy/pic/tests/codegen_spike/compare.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Run the numba kernel generated from ``bstar_parallel_3form`` on its parity cases and compare it with pyccel. - -Usage (numba is not a struphy dependency):: - - NUMBA_ENABLE_CUDASIM=1 python -m struphy.pic.tests.codegen_spike.compare cuda # numba.cuda simulator - python -m struphy.pic.tests.codegen_spike.compare cpu # same translation, numba.njit - -The simulator interprets the generated Python thread by thread (it does not compile anything); ``cpu`` compiles the -same translation in nopython mode. Prints one line per case and a summary; exits with 1 if a case disagrees. -""" - -import sys -import time - -import cunumpy as xp -import numpy as np - -from struphy.pic.tests.codegen_spike.numba_codegen import flatten, load - -MODULE = "struphy.pic.pushing.kernels.bstar_parallel_3form.bstar_parallel_3form_kernels" -NAME = "bstar_parallel_3form" -BLOCK = 128 - - -def compare(target): - from struphy.pic.pushing.kernels.bstar_parallel_3form import bstar_parallel_3form - from struphy.pic.tests.cuda_parity_cases import PARITY_CASES - from struphy.pic.tests.test_cuda_emulation import arrays - - kernel, path = load(MODULE, NAME, target) - print(f"generated {path} ({len(path.read_text().splitlines())} lines)") - spec = PARITY_CASES[NAME] - failures, bitwise = 0, 0 - with xp.use_backend("numpy"): - for case in spec.cases: - reference_args, generated_args = spec.build(case), spec.build(case) - bstar_parallel_3form(*reference_args) - flat = flatten(generated_args) - start = time.perf_counter() - if target == "cuda": - n_markers = generated_args[2].n_markers - kernel[(n_markers + BLOCK - 1) // BLOCK, BLOCK](*flat) - else: - kernel(*flat) - seconds = time.perf_counter() - start - reference, generated = arrays(reference_args), arrays(generated_args) - diff = max( - float(np.max(np.abs(r.astype(float) - g.astype(float)), initial=0.0)) - for r, g in zip(reference, generated) - ) - same = all(np.array_equal(r, g, equal_nan=True) for r, g in zip(reference, generated)) - ok = all(np.allclose(g, r, rtol=spec.rtol, atol=spec.atol) for r, g in zip(reference, generated)) - failures += not ok - bitwise += same - print(f"domain {case[0]:2d}: max |diff| = {diff:.2e}, bitwise equal: {same}, {seconds:.2f} s") - print(f"{target}: {len(spec.cases) - failures}/{len(spec.cases)} cases agree, {bitwise} bitwise") - return failures - - -if __name__ == "__main__": - sys.exit(1 if compare(sys.argv[1] if len(sys.argv) > 1 else "cpu") else 0) diff --git a/src/struphy/pic/tests/codegen_spike/numba_codegen.py b/src/struphy/pic/tests/codegen_spike/numba_codegen.py deleted file mode 100644 index 062df877d..000000000 --- a/src/struphy/pic/tests/codegen_spike/numba_codegen.py +++ /dev/null @@ -1,458 +0,0 @@ -"""Prototype for issue #690: generate a ``numba.cuda`` kernel from the Python source of a pyccel kernel. - -Not used by struphy. It exists to compare code generation with the hand-written ``_cuda.cu`` files on one -real kernel (``bstar_parallel_3form``); the result of the comparison is in ``CUDA_STRATEGY.md`` (decision for -#690). The module name contains no ``kernels``, so ``struphy compile`` ignores it. - -The generator works on the source files (the imported kernel modules are usually the compiled pyccel extensions): - -- it parses the entry kernel's module, follows every call into struphy modules (``import a.b as x`` and - ``from a.b import f``) and translates each reached function once, as a numba device function named - ``__``; -- it strips the pyccel decorators and type annotations, maps the NumPy scalar functions to ``math``, and - ``numpy.mod`` to the Fortran ``MODULO`` pyccel generates; -- ``empty``/``zeros`` become ``cuda.local.array`` with a constant size (a run-time size such as ``zeros(pn)`` gets - ``MAX_LOCAL`` entries, like the hand-written ``MAX_SPLINE_DEGREE`` arrays); -- array expressions on slices (``eta[:] = alpha[:3] * eta_k + ...``, ``bn[:] = 0.0``) become element loops, since - numba.cuda cannot allocate the temporaries; -- the entry kernel's marker loop ``for ip in range(n): ...`` becomes ``ip = cuda.grid(1)`` and ``continue`` becomes - ``return``; -- argument objects (``MarkerArguments``, ``DomainArguments``, ``DerhamArguments``) are passed as the fields of - their CUDA structs (``CudaMarkerArguments.struct.fields``, one kernel parameter each) and rebuilt as named tuples - at the top of the kernel; ``DerhamArguments`` gets its scratch arrays ``bn1``, ..., ``bd3`` as per-thread local - arrays there, so the helpers can stay one to one with pyccel. - -With ``target="cpu"`` the same translation is compiled with ``numba.njit`` (NumPy allocations, the marker loop kept), -which checks that the generated code types in nopython mode; numba.cuda itself needs a GPU (or NVVM) to compile, and -its simulator (``NUMBA_ENABLE_CUDASIM=1``) only interprets the Python. -""" - -import ast -import importlib.util -import re -import sys -import tempfile -from pathlib import Path - -from struphy.kernel_arguments.pusher_args_cuda import CudaDerhamArguments, CudaDomainArguments, CudaMarkerArguments - -MAX_LOCAL = 9 # MAX_SPLINE_DEGREE + 1 of bsplines_kernels.cuh - -STRUCT_CLASSES = { - "MarkerArguments": CudaMarkerArguments, - "DerhamArguments": CudaDerhamArguments, - "DomainArguments": CudaDomainArguments, -} -# pyccel scratch arrays of the argument classes: per-thread local arrays in the generated kernel -SCRATCH = {"DerhamArguments": ("bn1", "bn2", "bn3", "bd1", "bd2", "bd3")} - -NUMPY_MATH = { - "cos": "cos", - "sin": "sin", - "tan": "tan", - "arctan": "atan", - "arcsin": "asin", - "arccos": "acos", - "arctan2": "atan2", - "sqrt": "sqrt", - "exp": "exp", - "log": "log", - "floor": "floor", -} -NUMPY_ALLOC = ("empty", "zeros") - - -def _source_file(module): - """The ``.py`` file of `module` (the import system would find the compiled pyccel extension first).""" - spec = importlib.util.find_spec(module) - return Path(spec.origin).parent / (module.rsplit(".", 1)[-1] + ".py") - - -class _Module: - """Functions and imports of one source module.""" - - def __init__(self, name): - self.name = name - self.short = name.rsplit(".", 1)[-1] - tree = ast.parse(_source_file(name).read_text()) - self.functions = {n.name: n for n in tree.body if isinstance(n, ast.FunctionDef)} - self.module_aliases = {} # alias -> module - self.imported = {} # name -> (module, name) - for node in tree.body: - if isinstance(node, ast.Import): - for a in node.names: - self.module_aliases[a.asname or a.name] = a.name - elif isinstance(node, ast.ImportFrom): - for a in node.names: - self.imported[a.asname or a.name] = (node.module, a.name) - - def resolve(self, func): - """``(module, name)`` of a called function node, or ``None``.""" - if isinstance(func, ast.Name): - if func.id in self.functions: - return self.name, func.id - return self.imported.get(func.id) - if isinstance(func, ast.Attribute) and isinstance(func.value, ast.Name): - module = self.module_aliases.get(func.value.id) - if module is not None: - return module, func.attr - return None - - -def _array_ndim(annotation): - """Number of dimensions of a pyccel array annotation such as ``"float[:,:]"``, or 0.""" - if isinstance(annotation, ast.Constant) and isinstance(annotation.value, str): - match = re.search(r"\[([:,\s]+)\]", annotation.value) - if match: - return match.group(1).count(":") - return 0 - - -def _struct_type(annotation): - if isinstance(annotation, ast.Constant) and annotation.value in STRUCT_CLASSES: - return annotation.value - if isinstance(annotation, ast.Name) and annotation.id in STRUCT_CLASSES: - return annotation.id - return None - - -def _expr(source): - return ast.parse(source, mode="eval").body - - -def _stmt(source): - return ast.parse(source).body - - -class _Elementwise(ast.NodeTransformer): - """Index a slice expression at element ``index``: ``a[lo:hi]`` -> ``a[lo + index]``, 1D array ``x`` -> ``x[index]``.""" - - def __init__(self, index, arrays): - self.index = index - self.arrays = arrays - - def _element(self, slc): - lower = slc.lower if slc.lower is not None else ast.Constant(0) - return ast.BinOp(lower, ast.Add(), ast.Name(self.index, ast.Load())) - - def visit_Subscript(self, node): - if isinstance(node.slice, ast.Slice): - return ast.Subscript(node.value, self._element(node.slice), node.ctx) - if isinstance(node.slice, ast.Tuple) and any(isinstance(e, ast.Slice) for e in node.slice.elts): - elts = [self._element(e) if isinstance(e, ast.Slice) else e for e in node.slice.elts] - return ast.Subscript(node.value, ast.Tuple(elts, ast.Load()), node.ctx) - return node - - def visit_Name(self, node): - if self.arrays.get(node.id) == 1: - return ast.Subscript(ast.Name(node.id, ast.Load()), ast.Name(self.index, ast.Load()), node.ctx) - if node.id in self.arrays and self.arrays[node.id] > 1: - raise NotImplementedError(f"array expression on the {self.arrays[node.id]}D array {node.id}") - return node - - -def _slice_length(target): - """Number of elements written by a sliced assignment target (one slice per target).""" - slices = target.slice.elts if isinstance(target.slice, ast.Tuple) else [target.slice] - positions = [i for i, s in enumerate(slices) if isinstance(s, ast.Slice)] - if len(positions) != 1: - raise NotImplementedError(f"assignment to {ast.unparse(target)}: exactly one slice is supported") - slc = slices[positions[0]] - if slc.step is not None: - raise NotImplementedError(f"assignment to {ast.unparse(target)}: slice steps are not supported") - if slc.upper is None: - upper = ast.Subscript(ast.Attribute(target.value, "shape", ast.Load()), ast.Constant(positions[0]), ast.Load()) - else: - upper = slc.upper - if slc.lower is None: - return upper - return ast.BinOp(upper, ast.Sub(), slc.lower) - - -def _has_slice(target): - if not isinstance(target, ast.Subscript): - return False - s = target.slice - return isinstance(s, ast.Slice) or (isinstance(s, ast.Tuple) and any(isinstance(e, ast.Slice) for e in s.elts)) - - -class _ContinueToReturn(ast.NodeTransformer): - """``continue`` of the marker loop -> ``return`` (nested loops keep theirs).""" - - def visit_Continue(self, node): - return ast.Return(None) - - def visit_For(self, node): - return node - - def visit_While(self, node): - return node - - -class Generator: - """Translate a pyccel entry kernel and every struphy function it calls into one numba module.""" - - def __init__(self, target="cuda"): - assert target in ("cuda", "cpu") - self.target = target - self.modules = {} - self.translated = {} # (module, name) -> mangled name - self.sources = [] # translated functions, callees first - self.counter = 0 - - def _module(self, name): - if name not in self.modules: - self.modules[name] = _Module(name) - return self.modules[name] - - def _is_struphy_function(self, ref): - if ref is None or not ref[0].startswith("struphy."): - return False - try: - return ref[1] in self._module(ref[0]).functions - except (ModuleNotFoundError, FileNotFoundError, AttributeError): - return False - - def _numpy_name(self, module, node): - """The NumPy name a call or name node refers to in `module`, or ``None``.""" - if isinstance(node, ast.Name): - ref = module.imported.get(node.id) - if ref is not None and ref[0] == "numpy": - return ref[1] - return None - - def _fresh(self, prefix): - self.counter += 1 - return f"{prefix}{self.counter}" - - # ------------------------------------------------------------------ allocations - - def _allocation(self, module, target, call, arrays): - kind = self._numpy_name(module, call.func) - shape = call.args[0] - ndim = len(shape.elts) if isinstance(shape, ast.Tuple) else 1 - dtype = "int64" - for kw in call.keywords: - if kw.arg == "dtype" and isinstance(kw.value, ast.Name) and kw.value.id == "float": - dtype = "float64" - if len(call.args) > 1 and isinstance(call.args[1], ast.Name) and call.args[1].id == "float": - dtype = "float64" - if not call.keywords and len(call.args) == 1: - dtype = "float64" - arrays[target] = ndim - if self.target == "cpu": - return _stmt(f"{target} = np.{kind}({ast.unparse(shape)}, dtype=np.{dtype})") - elts = shape.elts if isinstance(shape, ast.Tuple) else [shape] - sizes = [e.value if isinstance(e, ast.Constant) else MAX_LOCAL for e in elts] - size = sizes[0] if ndim == 1 else "(" + ", ".join(map(str, sizes)) + ")" - stmts = _stmt(f"{target} = cuda.local.array({size}, {dtype})") - if kind == "zeros": - loops = "".join(f"{' ' * d}for _z{d} in range({s}):\n" for d, s in enumerate(sizes)) - index = ", ".join(f"_z{d}" for d in range(ndim)) - stmts += _stmt(f"{loops}{' ' * ndim}{target}[{index}] = 0") - return stmts - - # ------------------------------------------------------------------ functions - - def translate(self, module_name, name, entry=False): - """Translate function `name` of `module_name`; returns its name in the generated module.""" - key = (module_name, name) - if key in self.translated: - return self.translated[key] - module = self._module(module_name) - mangled = f"{module.short}__{name}" - self.translated[key] = mangled - node = ast.parse(ast.unparse(module.functions[name])).body[0] # a fresh copy - - # parameters: array dimensions and argument objects from the pyccel annotations - arrays = {a.arg: _array_ndim(a.annotation) for a in node.args.args if _array_ndim(a.annotation)} - structs = {a.arg: _struct_type(a.annotation) for a in node.args.args if _struct_type(a.annotation)} - for a in node.args.args: - a.annotation = None - node.returns = None - node.decorator_list = [] - node.name = mangled - if node.body and isinstance(node.body[0], ast.Expr) and isinstance(node.body[0].value, ast.Constant): - node.body = node.body[1:] or [ast.Pass()] - - node.body = self._statements(module, node.body, arrays) - - if entry: - self._make_entry(node, structs) - decorator = "@cuda.jit" if self.target == "cuda" else "@njit" - else: - decorator = "@cuda.jit(device=True)" if self.target == "cuda" else "@njit" - self.sources.append(f"{decorator}\n{ast.unparse(ast.fix_missing_locations(node))}\n") - return mangled - - def _statements(self, module, body, arrays): - out = [] - for stmt in body: - out += self._statement(module, stmt, arrays) - return out - - def _statement(self, module, stmt, arrays): - # local arrays - if ( - isinstance(stmt, ast.Assign) - and isinstance(stmt.value, ast.Call) - and self._numpy_name(module, stmt.value.func) in NUMPY_ALLOC - ): - (target,) = stmt.targets - return self._allocation(module, target.id, stmt.value, arrays) - # array expressions on slices -> element loops - if isinstance(stmt, (ast.Assign, ast.AugAssign)): - target = stmt.targets[0] if isinstance(stmt, ast.Assign) else stmt.target - if _has_slice(target): - index = self._fresh("_k") - length = _slice_length(target) - element = _Elementwise(index, arrays) - value = self._expressions(module, element.visit(stmt.value)) - new_target = element.visit(target) - if isinstance(stmt, ast.Assign): - inner = ast.Assign([new_target], value) - else: - inner = ast.AugAssign(new_target, stmt.op, value) - loop = ast.For( - ast.Name(index, ast.Store()), - ast.Call(ast.Name("range", ast.Load()), [length], []), - [inner], - [], - lineno=stmt.lineno, - ) - return [loop] - # compound statements: recurse into their bodies - for field in ("body", "orelse"): - if hasattr(stmt, field) and isinstance(getattr(stmt, field), list): - setattr(stmt, field, self._statements(module, getattr(stmt, field), arrays)) - for field in ("value", "test", "iter"): - if hasattr(stmt, field) and isinstance(getattr(stmt, field), ast.AST): - setattr(stmt, field, self._expressions(module, getattr(stmt, field))) - if isinstance(stmt, ast.Assign): - stmt.targets = [self._expressions(module, t) for t in stmt.targets] - if isinstance(stmt, ast.AugAssign): - stmt.target = self._expressions(module, stmt.target) - return [stmt] - - def _expressions(self, module, node): - generator = self - - class Calls(ast.NodeTransformer): - def visit_Call(self, call): - self.generic_visit(call) - ref = module.resolve(call.func) - numpy_name = generator._numpy_name(module, call.func) - if generator._is_struphy_function(ref): - call.func = ast.Name(generator.translate(*ref), ast.Load()) - elif numpy_name in NUMPY_MATH: - call.func = _expr(f"math.{NUMPY_MATH[numpy_name]}") - elif numpy_name == "mod": - call.func = ast.Name("_numpy_mod", ast.Load()) - elif numpy_name == "shape": - return ast.Attribute(call.args[0], "shape", ast.Load()) - elif numpy_name is not None: - raise NotImplementedError(f"numpy.{numpy_name} in {module.name}") - return call - - def visit_Name(self, name): - if generator._numpy_name(module, name) == "pi": - return _expr("math.pi") - return name - - return Calls().visit(node) - - def _make_entry(self, node, structs): - """Flatten the argument objects into struct fields, rebuild them, and map the marker loop to threads.""" - params, prologue = [], [] - for a in node.args.args: - kind = structs.get(a.arg) - if kind is None: - params.append(a.arg) - continue - fields = [f.name for f in STRUCT_CLASSES[kind].struct.fields] - params += [f"{a.arg}__{f}" for f in fields] - values = [f"{a.arg}__{f}" for f in fields] - if self.target == "cuda": - values += [f"cuda.local.array({MAX_LOCAL}, float64)" for _ in SCRATCH.get(kind, ())] - else: - values += [f"np.empty({MAX_LOCAL})" for _ in SCRATCH.get(kind, ())] - prologue += _stmt(f"{a.arg} = {kind}_T({', '.join(values)})") - node.args.args = [ast.arg(p) for p in params] - - body = prologue + node.body - if self.target == "cuda": - for i, stmt in enumerate(body): - if ( - isinstance(stmt, ast.For) - and isinstance(stmt.iter, ast.Call) - and getattr(stmt.iter.func, "id", None) == "range" - and len(stmt.iter.args) == 1 - ): - ip, n = stmt.target.id, ast.unparse(stmt.iter.args[0]) - loop_body = [_ContinueToReturn().visit(s) for s in stmt.body] - body = body[:i] + _stmt(f"{ip} = cuda.grid(1)\nif {ip} >= {n}:\n return") + loop_body - break - else: - raise NotImplementedError("no marker loop 'for ip in range(n)' in the entry kernel") - node.body = body - - def module_source(self): - decorator = "@cuda.jit(device=True)" if self.target == "cuda" else "@njit" - header = [ - f"# Generated by struphy.pic.tests.codegen_spike.numba_codegen (target={self.target}); do not edit.", - "import math", - "from collections import namedtuple", - "", - "import numpy as np", - "from numba import cuda, float64, int64, njit", - "", - ] - for kind, cls in STRUCT_CLASSES.items(): - names = [f.name for f in cls.struct.fields] + list(SCRATCH.get(kind, ())) - header.append(f"{kind}_T = namedtuple({kind + '_T'!r}, {names!r})") - header += [ - "", - "", - decorator, - "def _numpy_mod(a, b):", - " # numpy.mod as pyccel compiles it (Fortran MODULO)", - " return a - math.floor(a / b) * b", - "", - "", - ] - return "\n".join(header) + "\n\n".join(self.sources) - - -def generate(module_name, name, target="cuda"): - """Source of the numba module for the pyccel kernel `name` in `module_name` (and every function it calls).""" - generator = Generator(target) - generator.translate(module_name, name, entry=True) - return generator.module_source() - - -def load(module_name, name, target="cuda", directory=None): - """Generate, write to `directory` (a temporary one by default) and import; returns the kernel function.""" - source = generate(module_name, name, target) - directory = Path(directory or tempfile.mkdtemp(prefix="struphy_numba_codegen_")) - stem = f"generated_{name}_{target}" - path = directory / f"{stem}.py" - path.write_text(source) - spec = importlib.util.spec_from_file_location(stem, path) - module = importlib.util.module_from_spec(spec) - sys.modules[stem] = module - spec.loader.exec_module(module) - return getattr(module, f"{module_name.rsplit('.', 1)[-1]}__{name}"), path - - -def flatten(args): - """Kernel arguments with every argument object replaced by the fields of its CUDA struct (host arrays).""" - from struphy.pic.tests.cuda_emulation import _struct_field - - flat = [] - for value in args: - kind = type(value).__name__ - if kind in STRUCT_CLASSES: - flat += [_struct_field(value, f.name) for f in STRUCT_CLASSES[kind].struct.fields] - else: - flat.append(value) - return flat diff --git a/src/struphy/pic/tests/codegen_spike/test_numba_codegen.py b/src/struphy/pic/tests/codegen_spike/test_numba_codegen.py deleted file mode 100644 index 7648440b6..000000000 --- a/src/struphy/pic/tests/codegen_spike/test_numba_codegen.py +++ /dev/null @@ -1,29 +0,0 @@ -"""The numba kernel generated from ``bstar_parallel_3form`` agrees with pyccel (issue #690 prototype). - -Skipped unless numba (with numba-cuda for the simulator) is installed; numba is not a struphy dependency. -Runs in a subprocess because ``NUMBA_ENABLE_CUDASIM`` must be set before numba is imported. -""" - -import importlib.util -import subprocess -import sys - -import pytest - -from struphy.geometry.tests.test_domain import serial_child_env - -requires_numba = pytest.mark.skipif(importlib.util.find_spec("numba") is None, reason="numba is not installed") - - -@requires_numba -@pytest.mark.parametrize("target", ["cpu", "cuda"]) -def test_generated_kernel_agrees_with_pyccel(target): - """``cpu``: the translation compiled with numba.njit; ``cuda``: the numba.cuda kernel in the CUDA simulator.""" - extra = {"NUMBA_ENABLE_CUDASIM": "1"} if target == "cuda" else {} - result = subprocess.run( - [sys.executable, "-m", "struphy.pic.tests.codegen_spike.compare", target], - env=serial_child_env(**extra), - capture_output=True, - text=True, - ) - assert result.returncode == 0, result.stdout + result.stderr From 0e74e15bb8cfc0abc28741280c740ce50b047359 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 15:02:41 +0200 Subject: [PATCH 17/50] Rename first_init_idx to first_pusher_idx The pusher's scratch columns were called first_pusher_idx in Particles and first_init_idx in the kernel arguments; use first_pusher_idx everywhere. Co-Authored-By: Claude Opus 5.5 --- src/struphy/kernel_arguments/pusher_args.cuh | 2 +- src/struphy/kernel_arguments/pusher_args_cuda.py | 4 ++-- .../kernel_arguments/pusher_args_kernels.py | 2 +- .../cc_lin_mhd_5d_gradB_kernels.py | 4 ++-- .../cc_lin_mhd_5d_gradB_dg_kernels.py | 6 +++--- .../eval_gradB_ediff/eval_gradB_ediff_kernels.py | 6 +++--- .../kernels/bstar_2form/bstar_2form_kernels.py | 6 +++--- .../bstar_parallel_3form_kernels.py | 6 +++--- .../driftkinetic_hamiltonian_kernels.py | 6 +++--- .../grad_driftkinetic_hamiltonian_kernels.py | 4 ++-- ...push_deterministic_diffusion_stage_kernels.py | 4 ++-- .../push_eta_stage/push_eta_stage_cuda.cu | 8 ++++---- .../push_eta_stage/push_eta_stage_kernels.py | 6 +++--- ..._Bstar_discrete_gradient_1st_order_kernels.py | 8 ++++---- ...discrete_gradient_1st_order_newton_kernels.py | 8 ++++---- ..._Bstar_discrete_gradient_2nd_order_kernels.py | 8 ++++---- .../push_gc_Bstar_explicit_multistage_kernels.py | 10 ++++++---- ...ush_gc_bxEstar_dg_1st_order_newton_kernels.py | 6 +++--- ...xEstar_discrete_gradient_1st_order_kernels.py | 6 +++--- ...xEstar_discrete_gradient_2nd_order_kernels.py | 6 +++--- ...ush_gc_bxEstar_explicit_multistage_kernels.py | 6 +++--- .../push_gc_cc_J2_dg_Hdiv_kernels.py | 6 +++--- .../push_gc_cc_J2_dg_init_Hdiv_kernels.py | 4 ++-- .../push_gc_cc_J2_stage_H1vec_kernels.py | 6 +++--- .../push_gc_cc_J2_stage_Hdiv_kernels.py | 6 +++--- .../push_pc_eta_stage_H1vec_kernels.py | 4 ++-- .../push_pc_eta_stage_Hcurl_kernels.py | 4 ++-- .../push_pc_eta_stage_Hdiv_kernels.py | 4 ++-- .../push_vxb_analytic/push_vxb_analytic_cuda.cu | 4 ++-- .../push_vxb_analytic_kernels.py | 4 ++-- .../push_vxb_implicit/push_vxb_implicit_cuda.cu | 4 ++-- .../push_vxb_implicit_kernels.py | 4 ++-- .../kernels/unit_b_1form/unit_b_1form_kernels.py | 4 ++-- src/struphy/pic/pushing/prepare_push_cuda.cu | 16 ++++++++-------- .../pic/pushing/pusher_utilities_kernels.cuh | 6 +++--- .../pic/pushing/pusher_utilities_kernels.py | 6 +++--- .../propagators/current_coupling_5d_gradb.py | 12 ++++++------ 37 files changed, 109 insertions(+), 107 deletions(-) diff --git a/src/struphy/kernel_arguments/pusher_args.cuh b/src/struphy/kernel_arguments/pusher_args.cuh index 5c85f3d45..59815fced 100644 --- a/src/struphy/kernel_arguments/pusher_args.cuh +++ b/src/struphy/kernel_arguments/pusher_args.cuh @@ -12,7 +12,7 @@ struct MarkerArgs { int vdim; int weight_idx; int first_diagnostics_idx; - int first_init_idx; + int first_pusher_idx; int first_shift_idx; int residual_idx; int first_free_idx; diff --git a/src/struphy/kernel_arguments/pusher_args_cuda.py b/src/struphy/kernel_arguments/pusher_args_cuda.py index 0c9911d72..08fefff46 100644 --- a/src/struphy/kernel_arguments/pusher_args_cuda.py +++ b/src/struphy/kernel_arguments/pusher_args_cuda.py @@ -39,7 +39,7 @@ class CudaMarkerArguments(CudaStructArguments): ("vdim", "int"), ("weight_idx", "int"), ("first_diagnostics_idx", "int"), - ("first_init_idx", "int"), + ("first_pusher_idx", "int"), ("first_shift_idx", "int"), ("residual_idx", "int"), ("first_free_idx", "int"), @@ -71,7 +71,7 @@ def __init__( self.weight_idx = weight_idx self.n_markers = markers.shape[0] self.first_diagnostics_idx = first_diagnostics_idx - self.first_init_idx = first_pusher_idx + self.first_pusher_idx = first_pusher_idx self.first_shift_idx = first_shift_idx self.residual_idx = residual_idx self.first_free_idx = first_free_idx diff --git a/src/struphy/kernel_arguments/pusher_args_kernels.py b/src/struphy/kernel_arguments/pusher_args_kernels.py index c5b3f52d8..2e66dc58e 100644 --- a/src/struphy/kernel_arguments/pusher_args_kernels.py +++ b/src/struphy/kernel_arguments/pusher_args_kernels.py @@ -71,7 +71,7 @@ def __init__( # useful indices self.first_diagnostics_idx = first_diagnostics_idx - self.first_init_idx = first_pusher_idx + self.first_pusher_idx = first_pusher_idx self.first_shift_idx = first_shift_idx # starting idx for eta-shifts due to boundary conditions self.residual_idx = residual_idx # residual in iterative solvers self.first_free_idx = first_free_idx # index after which auxiliary saving is possible diff --git a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py index 4428c130d..e50a5e7a4 100644 --- a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py +++ b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py @@ -99,7 +99,7 @@ def cc_lin_mhd_5d_gradB( markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx mu_idx = args_markers.mu_idx # allocate for magnetic field evaluation @@ -130,7 +130,7 @@ def cc_lin_mhd_5d_gradB( continue # if particle is refilled - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue # marker positions diff --git a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py index 7705768f5..015b4506a 100644 --- a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py +++ b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py @@ -71,7 +71,7 @@ def cc_lin_mhd_5d_gradB_dg( markers = args_markers.markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx # allocate for magnetic field evaluation eta_diff = empty(3, dtype=float) @@ -108,10 +108,10 @@ def cc_lin_mhd_5d_gradB_dg( continue # marker positions, mid point - eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_init_idx : first_init_idx + 3]) / 2.0 + eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_pusher_idx : first_pusher_idx + 3]) / 2.0 eta_mid[:] = mod(eta_mid[:], 1.0) - eta_diff[:] = markers[ip, 0:3] - markers[ip, first_init_idx : first_init_idx + 3] + eta_diff[:] = markers[ip, 0:3] - markers[ip, first_pusher_idx : first_pusher_idx + 3] # marker weight and velocity weight = markers[ip, 5] diff --git a/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py index c21807b7c..e33183e72 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py @@ -42,7 +42,7 @@ def eval_gradB_ediff( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): @@ -51,10 +51,10 @@ def eval_gradB_ediff( continue # marker positions, mid point - eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_init_idx : first_init_idx + 3]) / 2.0 + eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_pusher_idx : first_pusher_idx + 3]) / 2.0 eta_mid[:] = mod(eta_mid[:], 1.0) - eta_diff = markers[ip, 0:3] - markers[ip, first_init_idx : first_init_idx + 3] + eta_diff = markers[ip, 0:3] - markers[ip, first_pusher_idx : first_pusher_idx + 3] # marker weight and velocity weight = markers[ip, 5] diff --git a/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py b/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py index ea1ba6ad6..42bed58ef 100644 --- a/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py +++ b/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py @@ -50,7 +50,7 @@ def bstar_2form( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -59,13 +59,13 @@ def bstar_2form( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n # spline evaluation diff --git a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py index f72014dd2..b6356a1b0 100644 --- a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py +++ b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py @@ -45,7 +45,7 @@ def bstar_parallel_3form( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -54,13 +54,13 @@ def bstar_parallel_3form( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n # evaluate Jacobian, result in dfm diff --git a/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py b/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py index 26acb7d20..d622133b3 100644 --- a/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py +++ b/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py @@ -44,7 +44,7 @@ def driftkinetic_hamiltonian( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -53,13 +53,13 @@ def driftkinetic_hamiltonian( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n mu = markers[ip, mu_idx] diff --git a/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py b/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py index 1e89b6e36..417d384e7 100644 --- a/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py +++ b/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py @@ -58,7 +58,7 @@ def grad_driftkinetic_hamiltonian( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -67,7 +67,7 @@ def grad_driftkinetic_hamiltonian( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) diff --git a/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py b/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py index a12243267..3b891be49 100644 --- a/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py @@ -49,7 +49,7 @@ def push_deterministic_diffusion_stage( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -121,7 +121,7 @@ def push_deterministic_diffusion_stage( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * a[stage] * k + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu index 978a5f1d5..0f5a1938c 100644 --- a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu @@ -14,7 +14,7 @@ * @param c Butcher nodes c (n_stages entries; unused, as in the pyccel kernel). * @param n_stages Number of stages; a raw device pointer carries no length. * - * Holes (first_init_idx == -1) and boundary particles (last column == -2) are skipped. + * Holes (first_pusher_idx == -1) and boundary particles (last column == -2) are skipped. * The kinetic boundary conditions are applied per marker after the update. */ extern "C" __global__ void push_eta_stage(double dt, int stage, MarkerArgs args_markers, DomainArgs args_domain, @@ -24,12 +24,12 @@ extern "C" __global__ void push_eta_stage(double dt, int stage, MarkerArgs args_ int n_markers = args_markers.n_markers; if (ip >= n_markers) return; - int first_init_idx = args_markers.first_init_idx; + int first_pusher_idx = args_markers.first_pusher_idx; int first_free_idx = args_markers.first_free_idx; double last = stage == n_stages - 1 ? 1. : 0.; // check if marker is a hole or a boundary particle - if (args_markers.markers(ip, first_init_idx) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) + if (args_markers.markers(ip, first_pusher_idx) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) return; double e1 = args_markers.markers(ip, 0); @@ -51,7 +51,7 @@ extern "C" __global__ void push_eta_stage(double dt, int stage, MarkerArgs args_ // accumulation for last stage args_markers.markers(ip, first_free_idx + j) += dt * b[stage] * k[j]; // update positions for intermediate stages or last stage - args_markers.markers(ip, j) = args_markers.markers(ip, first_init_idx + j) + dt * a[stage] * k[j] + + args_markers.markers(ip, j) = args_markers.markers(ip, first_pusher_idx + j) + dt * a[stage] * k[j] + last * args_markers.markers(ip, first_free_idx + j); } diff --git a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py index b037d4746..e5f0245f1 100644 --- a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py @@ -43,7 +43,7 @@ def push_eta_stage( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx if stage == n_stages - 1: @@ -57,7 +57,7 @@ def push_eta_stage( # fmt: on for ip in range(n_markers): # check if marker is a hole or a boundary particle - if markers[ip, first_init_idx] == -1.0 or markers[ip, -1] == -2.0: + if markers[ip, first_pusher_idx] == -1.0 or markers[ip, -1] == -2.0: continue e1 = markers[ip, 0] @@ -85,7 +85,7 @@ def push_eta_stage( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * a[stage] * k + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py index 32e0b3d64..4f0a4ac73 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py @@ -82,25 +82,25 @@ def push_gc_Bstar_discrete_gradient_1st_order( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) eta_diff[:] = eta_k - eta_n v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v_mid = (v_k + v_n) / 2.0 v_diff = v_k - v_n diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py index fb88521fc..8f296d666 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py @@ -143,23 +143,23 @@ def push_gc_Bstar_discrete_gradient_1st_order_newton( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_k_shifted[:] = eta_k + markers[ip, first_shift_idx : first_shift_idx + 3] eta_diff[:] = eta_k_shifted - eta_n v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v_diff = v_k - v_n mu = markers[ip, mu_idx] diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py index 5fe3aadac..3c6b8d82e 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py @@ -108,25 +108,25 @@ def push_gc_Bstar_discrete_gradient_2nd_order( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) eta_diff[:] = eta_k - eta_n v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v_mid = (v_k + v_n) / 2.0 v_diff = v_k - v_n diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py index f47a562a9..03ac967b6 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py @@ -80,7 +80,7 @@ def push_gc_Bstar_explicit_multistage( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -93,7 +93,7 @@ def push_gc_Bstar_explicit_multistage( for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue # if stage == 0.: @@ -207,11 +207,13 @@ def push_gc_Bstar_explicit_multistage( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * a[stage] * k + last * markers[ip, first_free_idx : first_free_idx + 3] ) - markers[ip, 3] = markers[ip, first_init_idx + 3] + dt * a[stage] * k_v + last * markers[ip, first_free_idx + 3] + markers[ip, 3] = ( + markers[ip, first_pusher_idx + 3] + dt * a[stage] * k_v + last * markers[ip, first_free_idx + 3] + ) # apply kinetic boundary conditions pusher_utilities_kernels.apply_kinetic_bc_marker(ip, args_markers, args_domain, False) diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py index f383d4329..3d9c1a3e2 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py @@ -129,18 +129,18 @@ def push_gc_bxEstar_dg_1st_order_newton( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_k_shifted[:] = eta_k + markers[ip, first_shift_idx : first_shift_idx + 3] eta_diff[:] = eta_k_shifted - eta_n diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py index ca3507264..7c26a86dd 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py @@ -75,18 +75,18 @@ def push_gc_bxEstar_discrete_gradient_1st_order( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py index 619512ae4..bd0457365 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py @@ -91,18 +91,18 @@ def push_gc_bxEstar_discrete_gradient_2nd_order( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py index c68033f5d..905bfb8e0 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py @@ -66,7 +66,7 @@ def push_gc_bxEstar_explicit_multistage( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -79,7 +79,7 @@ def push_gc_bxEstar_explicit_multistage( for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta1 = markers[ip, 0] @@ -171,7 +171,7 @@ def push_gc_bxEstar_explicit_multistage( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * a[stage] * k + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py index 624ab42b5..c10e36e0a 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py @@ -80,7 +80,7 @@ def push_gc_cc_J2_dg_Hdiv( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): @@ -90,7 +90,7 @@ def push_gc_cc_J2_dg_Hdiv( # marker positions, mid point eta_old[:] = markers[ip, 0:3] - eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_init_idx : first_init_idx + 3]) / 2.0 + eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_pusher_idx : first_pusher_idx + 3]) / 2.0 eta_mid[:] = mod(eta_mid[:], 1.0) v = markers[ip, 3] @@ -204,6 +204,6 @@ def push_gc_cc_J2_dg_Hdiv( e /= abs_b_star_para e /= det_df - markers[ip, 0:3] = markers[ip, first_init_idx : first_init_idx + 3] - dt * e + markers[ip, 0:3] = markers[ip, first_pusher_idx : first_pusher_idx + 3] - dt * e markers[ip, 0:3] *= alpha markers[ip, 0:3] += eta_old * (1.0 - alpha) diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py index 8fcaafd1c..388c069cc 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py @@ -67,12 +67,12 @@ def push_gc_cc_J2_dg_init_Hdiv( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta1 = markers[ip, 0] diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py index a3cb0e4e5..d40c24fc7 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py @@ -85,7 +85,7 @@ def push_gc_cc_J2_stage_H1vec( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -98,7 +98,7 @@ def push_gc_cc_J2_stage_H1vec( for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta1 = markers[ip, 0] @@ -203,7 +203,7 @@ def push_gc_cc_J2_stage_H1vec( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] - dt * a[stage] * e + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py index ab5e5f89d..c2a82a5f5 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py @@ -81,7 +81,7 @@ def push_gc_cc_J2_stage_Hdiv( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -98,7 +98,7 @@ def push_gc_cc_J2_stage_Hdiv( continue # if particle is refilled - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta1 = markers[ip, 0] @@ -204,7 +204,7 @@ def push_gc_cc_J2_stage_Hdiv( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] - dt * a[stage] * e + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py index 5946bdf4b..e9fe1c0c5 100644 --- a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py @@ -66,7 +66,7 @@ def push_pc_eta_stage_H1vec( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -131,7 +131,7 @@ def push_pc_eta_stage_H1vec( # update markers for the next stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * k * a[stage] + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py index 33004306d..55e02b6ff 100644 --- a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py @@ -59,7 +59,7 @@ def push_pc_eta_stage_Hcurl( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -127,7 +127,7 @@ def push_pc_eta_stage_Hcurl( # update markers for the next stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * k * a[stage] + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py index 6ced6a2f9..0cae6016c 100644 --- a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py @@ -59,7 +59,7 @@ def push_pc_eta_stage_Hdiv( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -128,7 +128,7 @@ def push_pc_eta_stage_Hdiv( # update markers for the next stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * k * a[stage] + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu index bc96bfd77..d14f81ec6 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu @@ -28,10 +28,10 @@ extern "C" __global__ void push_vxb_analytic(double dt, int stage, MarkerArgs ar int n_markers = args_markers.n_markers; if (ip >= n_markers) return; - int first_init_idx = args_markers.first_init_idx; + int first_pusher_idx = args_markers.first_pusher_idx; // check if marker is a hole - if (args_markers.markers(ip, first_init_idx) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) + if (args_markers.markers(ip, first_pusher_idx) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) return; double e1 = args_markers.markers(ip, 0); diff --git a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py index 3341e4794..defa85cc3 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py @@ -54,7 +54,7 @@ def push_vxb_analytic( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx # fmt: off #$ omp parallel private (ip, e1, e2, e3, v, dfm, det_df, span1, span2, span3, b_form, b_cart, b_abs, b_norm, vpar, vxb_norm, vperp, b_normxvperp) @@ -62,7 +62,7 @@ def push_vxb_analytic( # fmt: on for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0 or markers[ip, -1] == -2.0: + if markers[ip, first_pusher_idx] == -1.0 or markers[ip, -1] == -2.0: continue e1 = markers[ip, 0] diff --git a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu index 1190f3e25..8fbb4a7e2 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu @@ -27,10 +27,10 @@ extern "C" __global__ void push_vxb_implicit(double dt, int stage, MarkerArgs ar int n_markers = args_markers.n_markers; if (ip >= n_markers) return; - int first_init_idx = args_markers.first_init_idx; + int first_pusher_idx = args_markers.first_pusher_idx; // check if marker is a hole - if (args_markers.markers(ip, first_init_idx) == -1.) return; + if (args_markers.markers(ip, first_pusher_idx) == -1.) return; double e1 = args_markers.markers(ip, 0); double e2 = args_markers.markers(ip, 1); diff --git a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py index 505959518..2b1c5facf 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py @@ -65,7 +65,7 @@ def push_vxb_implicit( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx # fmt: off #$ omp parallel firstprivate(b_prod) private (ip, e1, e2, e3, v, dfm, det_df, span1, span2, span3, b_form, b_cart, rhs, lhs, lhs_inv, vec, res) @@ -73,7 +73,7 @@ def push_vxb_implicit( # fmt: on for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue e1 = markers[ip, 0] diff --git a/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py b/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py index 07460ad4e..77cd4aa0a 100644 --- a/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py +++ b/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py @@ -39,7 +39,7 @@ def unit_b_1form( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -48,7 +48,7 @@ def unit_b_1form( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) diff --git a/src/struphy/pic/pushing/prepare_push_cuda.cu b/src/struphy/pic/pushing/prepare_push_cuda.cu index 5c1dbd28e..71c02556d 100644 --- a/src/struphy/pic/pushing/prepare_push_cuda.cu +++ b/src/struphy/pic/pushing/prepare_push_cuda.cu @@ -1,28 +1,28 @@ /** * Prepare the marker buffer for a push, as the three column-slice assignments at the start of Pusher._push: * - * markers[:, first_init_idx:first_shift_idx] = markers[:, :n_init] + * markers[:, first_pusher_idx:first_shift_idx] = markers[:, :n_init] * markers[:, first_shift_idx:residual_idx] = 0.0 * markers[:, residual_idx:-2] = 0.0 * - * One thread per entry of the contiguous column range [first_init_idx, n_cols - 2) of every row, so that + * One thread per entry of the contiguous column range [first_pusher_idx, n_cols - 2) of every row, so that * neighbouring threads write neighbouring addresses of the row-major buffer and the whole range is written * in a single pass (instead of one strided pass over the buffer per slice). * * @param markers Marker buffer (n_rows x n_cols, row-major), every row including holes. * @param n_rows Number of rows of the buffer. * @param n_cols Number of columns of the buffer. - * @param first_init_idx First column of the saved initial phase space coordinates. - * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_init_idx + n_init). + * @param first_pusher_idx First column of the saved initial phase space coordinates. + * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_pusher_idx + n_init). */ -extern "C" __global__ void prepare_push(double* markers, long long n_rows, int n_cols, int first_init_idx, +extern "C" __global__ void prepare_push(double* markers, long long n_rows, int n_cols, int first_pusher_idx, int n_init) { - long long width = n_cols - 2 - first_init_idx; + long long width = n_cols - 2 - first_pusher_idx; long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; if (t >= n_rows * width) return; long long ip = t / width; - int col = first_init_idx + (int)(t - ip * width); + int col = first_pusher_idx + (int)(t - ip * width); double* row = markers + ip * n_cols; - row[col] = col < first_init_idx + n_init ? row[col - first_init_idx] : 0.0; + row[col] = col < first_pusher_idx + n_init ? row[col - first_pusher_idx] : 0.0; } diff --git a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh index a40dc406a..48b06ef19 100644 --- a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh +++ b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh @@ -38,7 +38,7 @@ __device__ inline void reflect_velocity(long long ip, Array2D markers, * @param newton Accumulate periodic shifts if true; overwrite them otherwise. * * Call after updating the position. Boundary type 0 wraps periodically, - * type 1 mirrors position and logical velocity and sets first_init_idx to -1, + * type 1 mirrors position and logical velocity and sets first_pusher_idx to -1, * type 2 clears all columns except the ID to -1, and type 3 is handled on * the host. The marker view carries shape and strides; j is a CUDA-only * column index replacing Pyccel array slices. @@ -46,7 +46,7 @@ __device__ inline void reflect_velocity(long long ip, Array2D markers, __device__ inline void apply_kinetic_bc_marker(int ip, const MarkerArgs& args_markers, const DomainArgs& args_domain, bool newton) { const long long* bc_type = args_markers.bc_type; - int first_init_idx = args_markers.first_init_idx; + int first_pusher_idx = args_markers.first_pusher_idx; int first_shift_idx = args_markers.first_shift_idx; // Remove markers before applying periodic or reflecting boundaries. @@ -92,7 +92,7 @@ __device__ inline void apply_kinetic_bc_marker(int ip, const MarkerArgs& args_ma } } if (n_reflected == 0) return; - args_markers.markers(ip, first_init_idx) = -1.; + args_markers.markers(ip, first_pusher_idx) = -1.; for (int axis = 0; axis < 3; ++axis) { if (reflected[axis]) reflect_velocity(ip, args_markers.markers, args_domain, axis); } diff --git a/src/struphy/pic/pushing/pusher_utilities_kernels.py b/src/struphy/pic/pushing/pusher_utilities_kernels.py index 78b9c7d70..b21c1bfb0 100644 --- a/src/struphy/pic/pushing/pusher_utilities_kernels.py +++ b/src/struphy/pic/pushing/pusher_utilities_kernels.py @@ -27,7 +27,7 @@ def apply_kinetic_bc_marker( * 0 (periodic): wrap :math:`\eta_i` into :math:`[0, 1)` and store the shift :math:`\pm 1` at ``first_shift_idx + axis`` (set for ``newton=False``, accumulated for ``newton=True``). * 1 (reflect): mirror :math:`\eta_i` at the boundary, flip the logical velocity component - and mark the marker as done for multi-stage/iterative pushers (``first_init_idx = -1``). + and mark the marker as done for multi-stage/iterative pushers (``first_pusher_idx = -1``). * 2 (remove): turn the marker into a hole (all columns except the ID set to -1). * 3 (refill): skipped here, handled by :meth:`~struphy.pic.base.Particles.apply_kinetic_bc`. @@ -48,7 +48,7 @@ def apply_kinetic_bc_marker( markers = args_markers.markers bc_type = args_markers.bc_type - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx # remove @@ -94,7 +94,7 @@ def apply_kinetic_bc_marker( if n_reflected == 0: return - markers[ip, first_init_idx] = -1.0 + markers[ip, first_pusher_idx] = -1.0 # reflect velocities (Jacobian at the reflected position) dfm = empty((3, 3), dtype=float) diff --git a/src/struphy/propagators/current_coupling_5d_gradb.py b/src/struphy/propagators/current_coupling_5d_gradb.py index 9fac1ec42..a596074df 100644 --- a/src/struphy/propagators/current_coupling_5d_gradb.py +++ b/src/struphy/propagators/current_coupling_5d_gradb.py @@ -459,14 +459,14 @@ def __call__(self, dt): holes = particles.holes args_markers = particles.args_markers markers = args_markers.markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # clear buffer - markers[:, first_init_idx:-2] = 0.0 + markers[:, first_pusher_idx:-2] = 0.0 # save old marker positions - markers[:, first_init_idx : first_init_idx + 3] = markers[:, :3] + markers[:, first_pusher_idx : first_pusher_idx + 3] = markers[:, :3] # sum up total magnetic field b_full1 = b_eq + b_tilde (in-place) b_full = self._b2.copy(out=self._b_full) @@ -530,7 +530,7 @@ def __call__(self, dt): diffs = self.update_feec_variables(u=u_new) # clear the buffer - markers[:, first_init_idx:-2] = 0.0 + markers[:, first_pusher_idx:-2] = 0.0 # update_weights if self.variables.energetic_ions.species.weights_params.control_variate: @@ -661,7 +661,7 @@ def __call__(self, dt): sum_u_diff_loc = xp.sum((u_diff.toarray() ** 2)) sum_H_diff_loc = xp.sum( - (markers[~holes, :3] - markers[~holes, first_init_idx : first_init_idx + 3]) ** 2, + (markers[~holes, :3] - markers[~holes, first_pusher_idx : first_pusher_idx + 3]) ** 2, ) buffer_array = xp.array([sum_u_diff_loc]) @@ -853,7 +853,7 @@ def __call__(self, dt): diffs = self.update_feec_variables(u=u_new) # clear the buffer - markers[:, first_init_idx:-2] = 0.0 + markers[:, first_pusher_idx:-2] = 0.0 # update_weights if self.variables.energetic_ions.species.weights_params.control_variate: From 7e922969e3e4c94c0f3daeed45167a6618f1bb51 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 15:11:12 +0200 Subject: [PATCH 18/50] Rename first_init_idx to first_pusher_idx in the code new on this branch Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/kernel_test_args.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index 24e597e89..abe1ef445 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -29,7 +29,7 @@ def marker_arguments(bc): markers[:, 3:6] = rng.uniform(-2, 2, (N_MARKERS, 3)) markers[:, 8:14] = markers[:, :6] markers[:, 18:21] = 0.0 - # Row 0 is a hole. Struphy marks every column of a hole with -1: pushers test column first_init_idx (8), + # Row 0 is a hole. Struphy marks every column of a hole with -1: pushers test column first_pusher_idx (8), # accumulation kernels test column 0. markers[0, :] = -1.0 markers[1, -1] = -2.0 From 8223abfaff756a844fa20ba7f5736cbb567d40e8 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 15:13:18 +0200 Subject: [PATCH 19/50] Skip the new GVEC tests by default (with_gvec=False), like the other GVEC tests GVEC dies with SIGILL on the CI runners (in gvec's library init), which took down the serial and the 2-rank unit test jobs. Co-Authored-By: Claude Opus 5.5 --- .../fields_background/tests/test_equils_cupy.py | 12 +++++++----- .../fields_background/tests/test_gvec_equil.py | 10 ++++++++-- 2 files changed, 15 insertions(+), 7 deletions(-) diff --git a/src/struphy/fields_background/tests/test_equils_cupy.py b/src/struphy/fields_background/tests/test_equils_cupy.py index e7710570d..28614e65a 100644 --- a/src/struphy/fields_background/tests/test_equils_cupy.py +++ b/src/struphy/fields_background/tests/test_equils_cupy.py @@ -214,7 +214,9 @@ def check_equil_on_cupy(case): assert np.allclose(o, r, rtol=1e-12, atol=1e-12, equal_nan=True), (case, where, name) -def _check_needs(case): +def _check_needs(case, with_gvec=False): + if EQUIL_CASES[case][0] == "GVECequilibrium" and not with_gvec: + pytest.skip("GVEC not tested here (with_gvec=False), like the other GVEC tests") 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") @@ -248,12 +250,12 @@ def test_host_call_numpy_is_plain_call(): @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): +def test_equil_fake_cupy(case, with_gvec=False): """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) + _check_needs(case, with_gvec) 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})" @@ -263,7 +265,7 @@ def test_equil_fake_cupy(case): @requires_cupy @pytest.mark.parametrize("case", list(EQUIL_CASES)) -def test_equil_cupy(case): +def test_equil_cupy(case, with_gvec=False): """On a GPU: the equilibria are created and evaluated on CuPy and agree with NumPy.""" - _check_needs(case) + _check_needs(case, with_gvec) 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 index 40b84e065..d21017ede 100644 --- a/src/struphy/fields_background/tests/test_gvec_equil.py +++ b/src/struphy/fields_background/tests/test_gvec_equil.py @@ -17,12 +17,15 @@ def _markers(n=7, seed=3): @pytest.mark.parametrize("name", ["bv", "jv", "p0", "n0", "absB0"]) -def test_gvec_marker_evaluation_matches_grid(name): +def test_gvec_marker_evaluation_matches_grid(name, with_gvec=False): """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. """ + if not with_gvec: + pytest.skip("GVEC not tested here (with_gvec=False), like the other GVEC tests") + from struphy.fields_background.equils import GVECequilibrium equil = GVECequilibrium(**GVEC_PARAMS) @@ -38,8 +41,11 @@ def test_gvec_marker_evaluation_matches_grid(name): assert np.isclose(f[ip], g.ravel()[0], rtol=1e-12, atol=1e-14) -def test_gvec_marker_evaluation_boozer_raises(): +def test_gvec_marker_evaluation_boozer_raises(with_gvec=False): """gvec computes the Boozer transform per flux surface, so marker evaluation with ``use_boozer=True`` raises.""" + if not with_gvec: + pytest.skip("GVEC not tested here (with_gvec=False), like the other GVEC tests") + from struphy.fields_background.equils import GVECequilibrium equil = GVECequilibrium(**GVEC_PARAMS, use_boozer=True) From 703a94717ab344da55617a04165cc37efe39e1d5 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 15:13:19 +0200 Subject: [PATCH 20/50] Rename first_init_idx to first_pusher_idx in the code new on this branch Co-Authored-By: Claude Opus 5.5 --- .../bstar_parallel_3form/bstar_parallel_3form_cuda.cu | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu index ec7389284..7a473ed1b 100644 --- a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu +++ b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_cuda.cu @@ -32,7 +32,7 @@ extern "C" __global__ void bstar_parallel_3form(const double* alpha, const long double eta_k[3], eta_n[3], eta[3], dfm[9]; // get marker arguments - int first_init_idx = args_markers.first_init_idx; + int first_pusher_idx = args_markers.first_pusher_idx; int first_shift_idx = args_markers.first_shift_idx; // only do something if particle is a "true" particle (i.e. not a hole) @@ -40,7 +40,7 @@ extern "C" __global__ void bstar_parallel_3form(const double* alpha, const long for (int j = 0; j < 3; ++j) { eta_k[j] = args_markers.markers(ip, j) + args_markers.markers(ip, first_shift_idx + j); - eta_n[j] = args_markers.markers(ip, first_init_idx + j); + eta_n[j] = args_markers.markers(ip, first_pusher_idx + j); eta[j] = alpha[j] * eta_k[j] + (1.0 - alpha[j]) * eta_n[j]; // numpy.mod(eta, 1.0) (Fortran MODULO): the result has the sign of the divisor @@ -48,7 +48,7 @@ extern "C" __global__ void bstar_parallel_3form(const double* alpha, const long } double v_k = args_markers.markers(ip, 3); - double v_n = args_markers.markers(ip, first_init_idx + 3); + double v_n = args_markers.markers(ip, first_pusher_idx + 3); double v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n; // evaluate Jacobian, result in dfm From 05bc8d95e6d55b82224157915b37ee39728e74fc Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 9 Oct 2026 16:21:37 +0200 Subject: [PATCH 21/50] Use xp.mpi.exchange --- pyproject.toml | 2 +- src/struphy/pic/base.py | 85 +++++++++----------- src/struphy/pic/tests/test_sorting_device.py | 7 +- 3 files changed, 42 insertions(+), 52 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d14c35a11..e5474e8a2 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 @ git+https://github.com/max-models/cunumpy.git@26a640b1d47d5677663785e49b1c44d343540df1", "pyccel>=2.2.0, <=2.2.3", "maybempi>=0.1.2", "feectools>=0.7.0, <=0.7.0", diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index cae11c476..5e1d217b2 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2,7 +2,6 @@ import os import warnings from abc import ABCMeta, abstractmethod -from contextlib import ExitStack from typing import TYPE_CHECKING import cunumpy as xp @@ -4646,67 +4645,55 @@ def _mpi_cuda_aware(self) -> bool: _warned_mpi_staging = False def _sendrecv_markers_device(self, recv_info, hole_inds_after_send): - """:meth:`_sendrecv_markers` on the CuPy backend: marker rows go to MPI as device buffers. - - With CUDA-aware MPI (:meth:`_mpi_cuda_aware`), the rows to send (:attr:`_send_list`, - gathered on the device by :meth:`_sendrecv_get_destinations`) and one device receive - buffer are handed to ``Isend``/``Irecv`` directly, after - :func:`cunumpy.mpi.synchronize_for_mpi` (inside :func:`cunumpy.mpi.mpi_buffer`), so MPI - reads what the kernels wrote. The rows from rank ``i`` arrive in a contiguous slice of - the receive buffer, and one device scatter puts all of them into the holes. Without - CUDA-aware MPI, :func:`~cunumpy.mpi.mpi_buffer` stages the same buffers through host - memory (counted by :func:`cunumpy.profiling.count_transfers`). - - The bookkeeping stays on the host without host/device copies: the counts are host - integers (``recv_info`` from the Alltoall, the send counts are the row counts of the - send buffers), and so is the number of holes. Ranks exchange no message when the - count is zero (both sides know the count, so they skip consistently). + """Exchange marker rows and scatter incoming rows into empty slots. + + :func:`cunumpy.mpi.exchange` handles GPU synchronization, optional host + staging, and completion of all MPI transfers. Parameters ---------- - recv_info : array[int] - Amount of markers to be received from i-th process (NumPy array). - - hole_inds_after_send : array[int] - Indices of empty rows in markers after send (device array). + recv_info : array[int] + Number of markers to receive from each process (NumPy array). + hole_inds_after_send : array[int] + Indices of empty rows in markers after send (device array). """ recv_info = np.asarray(recv_info) n_recv = int(recv_info.sum()) + if hole_inds_after_send.size < n_recv: warnings.warn( - f'Strong load imbalance detected: \ -number of holes ({hole_inds_after_send.size}) on rank {self.mpi_rank} \ -is smaller than number of incoming particles ({n_recv}). \ -Increasing the value of "bufsize" in the markers parameters for the next run.', + f"Strong load imbalance detected: " + f"number of holes ({hole_inds_after_send.size}) " + f"on rank {self.mpi_rank} is smaller than number of " + f"incoming particles ({n_recv}). " + 'Increase "bufsize" in the markers parameters for the next run.', ) self.mpi_comm.Abort() - cuda_aware = self._mpi_cuda_aware() - - # rows from rank i land in recvbuf[first_row[i] : first_row[i] + recv_info[i]] + recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=self.markers.dtype) first_row = np.cumsum(recv_info) - recv_info - recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=float) - with ExitStack() as buffers: - if n_recv > 0: - recv_mpi = buffers.enter_context( - xp.mpi.mpi_buffer(recvbuf, send=False, recv=True, cuda_aware=cuda_aware), - ) - reqs = [] - for i in range(self.mpi_size): - if i == self.mpi_rank: - continue - if recv_info[i] > 0: - rows = recv_mpi[first_row[i] : first_row[i] + recv_info[i]] - reqs.append(self.mpi_comm.Irecv(rows, source=i, tag=i)) - data = self._send_list[i] - if data.shape[0] > 0: - send_mpi = buffers.enter_context(xp.mpi.mpi_buffer(data, cuda_aware=cuda_aware)) - reqs.append(self.mpi_comm.Isend(send_mpi, dest=i, tag=self.mpi_rank)) - # wait inside the block: the buffers must stay valid until MPI is done, and staged - # receive rows are copied to the device when the block ends - for req in reqs: - req.Wait() + sends = [] + receives = [] + for rank in range(self.mpi_size): + if rank == self.mpi_rank: + continue + + if recv_info[rank] > 0: + start = int(first_row[rank]) + stop = start + int(recv_info[rank]) + receives.append((recvbuf[start:stop], rank, rank)) + + data = self._send_list[rank] + if data.shape[0] > 0: + sends.append((data, rank, self.mpi_rank)) + + xp.mpi.exchange( + self.mpi_comm, + sends=sends, + receives=receives, + cuda_aware=self._mpi_cuda_aware(), + ) if n_recv > 0: self._markers[hole_inds_after_send[:n_recv]] = recvbuf diff --git a/src/struphy/pic/tests/test_sorting_device.py b/src/struphy/pic/tests/test_sorting_device.py index 55c5d09f5..dc8fa9462 100644 --- a/src/struphy/pic/tests/test_sorting_device.py +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -143,12 +143,15 @@ def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", if cuda_aware: assert transfers.total == 0, transfers.report() else: - # one staged copy per non-empty send buffer, one copy back of the receive buffer + # exchange stages each non-empty send and receive buffer separately n_dest = sum(1 for i in range(size) if i != rank and dev._send_list[i].shape[0] > 0) + n_sources = sum(1 for i in range(size) if i != rank and ref._recvbufs[i].shape[0] > 0) assert transfers.to_host == n_dest, transfers.report() - assert transfers.to_device == (1 if n_recv > 0 else 0), transfers.report() + assert transfers.to_device == n_sources, transfers.report() sent_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_host") assert sent_bytes == n_send * ref.markers.shape[1] * 8 + received_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_device") + assert received_bytes == n_recv * ref.markers.shape[1] * 8 if count_reads and cuda_aware: # (the staged exchange reads the send buffers with .get(), counted above) From a6e1b19faae844cb43d997cfa15ee69926103a9c Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 7 Oct 2026 19:24:11 +0200 Subject: [PATCH 22/50] Fuse the pusher preamble into one CUDA kernel --- src/struphy/pic/pushing/prepare_push_cuda.cu | 28 ++++++++++++++ src/struphy/pic/pushing/pusher.py | 39 +++++++++++++++----- 2 files changed, 58 insertions(+), 9 deletions(-) create mode 100644 src/struphy/pic/pushing/prepare_push_cuda.cu diff --git a/src/struphy/pic/pushing/prepare_push_cuda.cu b/src/struphy/pic/pushing/prepare_push_cuda.cu new file mode 100644 index 000000000..5c1dbd28e --- /dev/null +++ b/src/struphy/pic/pushing/prepare_push_cuda.cu @@ -0,0 +1,28 @@ +/** + * Prepare the marker buffer for a push, as the three column-slice assignments at the start of Pusher._push: + * + * markers[:, first_init_idx:first_shift_idx] = markers[:, :n_init] + * markers[:, first_shift_idx:residual_idx] = 0.0 + * markers[:, residual_idx:-2] = 0.0 + * + * One thread per entry of the contiguous column range [first_init_idx, n_cols - 2) of every row, so that + * neighbouring threads write neighbouring addresses of the row-major buffer and the whole range is written + * in a single pass (instead of one strided pass over the buffer per slice). + * + * @param markers Marker buffer (n_rows x n_cols, row-major), every row including holes. + * @param n_rows Number of rows of the buffer. + * @param n_cols Number of columns of the buffer. + * @param first_init_idx First column of the saved initial phase space coordinates. + * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_init_idx + n_init). + */ +extern "C" __global__ void prepare_push(double* markers, long long n_rows, int n_cols, int first_init_idx, + int n_init) { + long long width = n_cols - 2 - first_init_idx; + long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; + if (t >= n_rows * width) return; + + long long ip = t / width; + int col = first_init_idx + (int)(t - ip * width); + double* row = markers + ip * n_cols; + row[col] = col < first_init_idx + n_init ? row[col - first_init_idx] : 0.0; +} diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index 7f3d64961..762dd3bd9 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -1,9 +1,11 @@ "Accelerated particle pushing." import logging +from functools import cache +from pathlib import Path import cunumpy as xp -from cunumpy.kernels import Kernel, PyccelKernel +from cunumpy.kernels import CudaKernel, Kernel, PyccelKernel from line_profiler import profile from maybempi import MPI from scope_profiler import ProfileManager @@ -16,6 +18,13 @@ logger = logging.getLogger("struphy") +@cache +def _prepare_push_cuda() -> CudaKernel: + """The CUDA kernel that fuses the column-slice assignments at the start of :meth:`Pusher._push` + (compiled on first call).""" + return CudaKernel.from_file(Path(__file__).with_name("prepare_push_cuda.cu"), "prepare_push") + + class Pusher: r""" Class for solving particle ODEs @@ -208,17 +217,29 @@ def _push(self, dt: float): logger.debug(f"{residual_idx =}") logger.debug(f"{self.particles.n_cols =}") - init_slice = slice(first_pusher_idx, first_shift_idx) - shift_slice = slice(first_shift_idx, residual_idx) + if self._cuda: + # the three assignments below in one pass over the buffer (each slice is a strided pass on its own) + n_rows, n_cols = markers.shape + _prepare_push_cuda()( + markers, + n_rows, + n_cols, + first_pusher_idx, + 3 + vdim, + n_threads=n_rows * (n_cols - 2 - first_pusher_idx), + ) + else: + init_slice = slice(first_pusher_idx, first_shift_idx) + shift_slice = slice(first_shift_idx, residual_idx) - # save initial phase space coordinates - markers[:, init_slice] = markers[:, : 3 + vdim] + # save initial phase space coordinates + markers[:, init_slice] = markers[:, : 3 + vdim] - # set boundary shifts to zero - markers[:, shift_slice] = 0.0 + # set boundary shifts to zero + markers[:, shift_slice] = 0.0 - # clear buffer columns starting from residual index, dont clear ID (last column) and loc_box - markers[:, residual_idx:-2] = 0.0 + # clear buffer columns starting from residual index, dont clear ID (last column) and loc_box + markers[:, residual_idx:-2] = 0.0 rank = self.particles.mpi_rank logger.debug(f"rank {rank}: starting {self.kernel} ...") From c1e299ba20a942b62e8da8f617d3f34d0974a0f7 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 15:02:41 +0200 Subject: [PATCH 23/50] Rename first_init_idx to first_pusher_idx The pusher's scratch columns were called first_pusher_idx in Particles and first_init_idx in the kernel arguments; use first_pusher_idx everywhere. Co-Authored-By: Claude Opus 5.5 --- src/struphy/kernel_arguments/pusher_args.cuh | 2 +- src/struphy/kernel_arguments/pusher_args_cuda.py | 4 ++-- .../kernel_arguments/pusher_args_kernels.py | 2 +- .../cc_lin_mhd_5d_gradB_kernels.py | 4 ++-- .../cc_lin_mhd_5d_gradB_dg_kernels.py | 6 +++--- .../eval_gradB_ediff/eval_gradB_ediff_kernels.py | 6 +++--- .../kernels/bstar_2form/bstar_2form_kernels.py | 6 +++--- .../bstar_parallel_3form_kernels.py | 6 +++--- .../driftkinetic_hamiltonian_kernels.py | 6 +++--- .../grad_driftkinetic_hamiltonian_kernels.py | 4 ++-- ...push_deterministic_diffusion_stage_kernels.py | 4 ++-- .../push_eta_stage/push_eta_stage_cuda.cu | 8 ++++---- .../push_eta_stage/push_eta_stage_kernels.py | 6 +++--- ..._Bstar_discrete_gradient_1st_order_kernels.py | 8 ++++---- ...discrete_gradient_1st_order_newton_kernels.py | 8 ++++---- ..._Bstar_discrete_gradient_2nd_order_kernels.py | 8 ++++---- .../push_gc_Bstar_explicit_multistage_kernels.py | 10 ++++++---- ...ush_gc_bxEstar_dg_1st_order_newton_kernels.py | 6 +++--- ...xEstar_discrete_gradient_1st_order_kernels.py | 6 +++--- ...xEstar_discrete_gradient_2nd_order_kernels.py | 6 +++--- ...ush_gc_bxEstar_explicit_multistage_kernels.py | 6 +++--- .../push_gc_cc_J2_dg_Hdiv_kernels.py | 6 +++--- .../push_gc_cc_J2_dg_init_Hdiv_kernels.py | 4 ++-- .../push_gc_cc_J2_stage_H1vec_kernels.py | 6 +++--- .../push_gc_cc_J2_stage_Hdiv_kernels.py | 6 +++--- .../push_pc_eta_stage_H1vec_kernels.py | 4 ++-- .../push_pc_eta_stage_Hcurl_kernels.py | 4 ++-- .../push_pc_eta_stage_Hdiv_kernels.py | 4 ++-- .../push_vxb_analytic/push_vxb_analytic_cuda.cu | 4 ++-- .../push_vxb_analytic_kernels.py | 4 ++-- .../push_vxb_implicit/push_vxb_implicit_cuda.cu | 4 ++-- .../push_vxb_implicit_kernels.py | 4 ++-- .../kernels/unit_b_1form/unit_b_1form_kernels.py | 4 ++-- src/struphy/pic/pushing/prepare_push_cuda.cu | 16 ++++++++-------- .../pic/pushing/pusher_utilities_kernels.cuh | 6 +++--- .../pic/pushing/pusher_utilities_kernels.py | 6 +++--- .../propagators/current_coupling_5d_gradb.py | 12 ++++++------ 37 files changed, 109 insertions(+), 107 deletions(-) diff --git a/src/struphy/kernel_arguments/pusher_args.cuh b/src/struphy/kernel_arguments/pusher_args.cuh index 5c85f3d45..59815fced 100644 --- a/src/struphy/kernel_arguments/pusher_args.cuh +++ b/src/struphy/kernel_arguments/pusher_args.cuh @@ -12,7 +12,7 @@ struct MarkerArgs { int vdim; int weight_idx; int first_diagnostics_idx; - int first_init_idx; + int first_pusher_idx; int first_shift_idx; int residual_idx; int first_free_idx; diff --git a/src/struphy/kernel_arguments/pusher_args_cuda.py b/src/struphy/kernel_arguments/pusher_args_cuda.py index 0c9911d72..08fefff46 100644 --- a/src/struphy/kernel_arguments/pusher_args_cuda.py +++ b/src/struphy/kernel_arguments/pusher_args_cuda.py @@ -39,7 +39,7 @@ class CudaMarkerArguments(CudaStructArguments): ("vdim", "int"), ("weight_idx", "int"), ("first_diagnostics_idx", "int"), - ("first_init_idx", "int"), + ("first_pusher_idx", "int"), ("first_shift_idx", "int"), ("residual_idx", "int"), ("first_free_idx", "int"), @@ -71,7 +71,7 @@ def __init__( self.weight_idx = weight_idx self.n_markers = markers.shape[0] self.first_diagnostics_idx = first_diagnostics_idx - self.first_init_idx = first_pusher_idx + self.first_pusher_idx = first_pusher_idx self.first_shift_idx = first_shift_idx self.residual_idx = residual_idx self.first_free_idx = first_free_idx diff --git a/src/struphy/kernel_arguments/pusher_args_kernels.py b/src/struphy/kernel_arguments/pusher_args_kernels.py index c5b3f52d8..2e66dc58e 100644 --- a/src/struphy/kernel_arguments/pusher_args_kernels.py +++ b/src/struphy/kernel_arguments/pusher_args_kernels.py @@ -71,7 +71,7 @@ def __init__( # useful indices self.first_diagnostics_idx = first_diagnostics_idx - self.first_init_idx = first_pusher_idx + self.first_pusher_idx = first_pusher_idx self.first_shift_idx = first_shift_idx # starting idx for eta-shifts due to boundary conditions self.residual_idx = residual_idx # residual in iterative solvers self.first_free_idx = first_free_idx # index after which auxiliary saving is possible diff --git a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py index 4428c130d..e50a5e7a4 100644 --- a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py +++ b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py @@ -99,7 +99,7 @@ def cc_lin_mhd_5d_gradB( markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx mu_idx = args_markers.mu_idx # allocate for magnetic field evaluation @@ -130,7 +130,7 @@ def cc_lin_mhd_5d_gradB( continue # if particle is refilled - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue # marker positions diff --git a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py index 7705768f5..015b4506a 100644 --- a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py +++ b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py @@ -71,7 +71,7 @@ def cc_lin_mhd_5d_gradB_dg( markers = args_markers.markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx # allocate for magnetic field evaluation eta_diff = empty(3, dtype=float) @@ -108,10 +108,10 @@ def cc_lin_mhd_5d_gradB_dg( continue # marker positions, mid point - eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_init_idx : first_init_idx + 3]) / 2.0 + eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_pusher_idx : first_pusher_idx + 3]) / 2.0 eta_mid[:] = mod(eta_mid[:], 1.0) - eta_diff[:] = markers[ip, 0:3] - markers[ip, first_init_idx : first_init_idx + 3] + eta_diff[:] = markers[ip, 0:3] - markers[ip, first_pusher_idx : first_pusher_idx + 3] # marker weight and velocity weight = markers[ip, 5] diff --git a/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py index c21807b7c..e33183e72 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py @@ -42,7 +42,7 @@ def eval_gradB_ediff( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): @@ -51,10 +51,10 @@ def eval_gradB_ediff( continue # marker positions, mid point - eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_init_idx : first_init_idx + 3]) / 2.0 + eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_pusher_idx : first_pusher_idx + 3]) / 2.0 eta_mid[:] = mod(eta_mid[:], 1.0) - eta_diff = markers[ip, 0:3] - markers[ip, first_init_idx : first_init_idx + 3] + eta_diff = markers[ip, 0:3] - markers[ip, first_pusher_idx : first_pusher_idx + 3] # marker weight and velocity weight = markers[ip, 5] diff --git a/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py b/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py index ea1ba6ad6..42bed58ef 100644 --- a/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py +++ b/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py @@ -50,7 +50,7 @@ def bstar_2form( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -59,13 +59,13 @@ def bstar_2form( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n # spline evaluation diff --git a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py index f72014dd2..b6356a1b0 100644 --- a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py +++ b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py @@ -45,7 +45,7 @@ def bstar_parallel_3form( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -54,13 +54,13 @@ def bstar_parallel_3form( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n # evaluate Jacobian, result in dfm diff --git a/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py b/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py index 26acb7d20..d622133b3 100644 --- a/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py +++ b/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py @@ -44,7 +44,7 @@ def driftkinetic_hamiltonian( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -53,13 +53,13 @@ def driftkinetic_hamiltonian( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n mu = markers[ip, mu_idx] diff --git a/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py b/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py index 1e89b6e36..417d384e7 100644 --- a/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py +++ b/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py @@ -58,7 +58,7 @@ def grad_driftkinetic_hamiltonian( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -67,7 +67,7 @@ def grad_driftkinetic_hamiltonian( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) diff --git a/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py b/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py index a12243267..3b891be49 100644 --- a/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py @@ -49,7 +49,7 @@ def push_deterministic_diffusion_stage( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -121,7 +121,7 @@ def push_deterministic_diffusion_stage( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * a[stage] * k + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu index 978a5f1d5..0f5a1938c 100644 --- a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu @@ -14,7 +14,7 @@ * @param c Butcher nodes c (n_stages entries; unused, as in the pyccel kernel). * @param n_stages Number of stages; a raw device pointer carries no length. * - * Holes (first_init_idx == -1) and boundary particles (last column == -2) are skipped. + * Holes (first_pusher_idx == -1) and boundary particles (last column == -2) are skipped. * The kinetic boundary conditions are applied per marker after the update. */ extern "C" __global__ void push_eta_stage(double dt, int stage, MarkerArgs args_markers, DomainArgs args_domain, @@ -24,12 +24,12 @@ extern "C" __global__ void push_eta_stage(double dt, int stage, MarkerArgs args_ int n_markers = args_markers.n_markers; if (ip >= n_markers) return; - int first_init_idx = args_markers.first_init_idx; + int first_pusher_idx = args_markers.first_pusher_idx; int first_free_idx = args_markers.first_free_idx; double last = stage == n_stages - 1 ? 1. : 0.; // check if marker is a hole or a boundary particle - if (args_markers.markers(ip, first_init_idx) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) + if (args_markers.markers(ip, first_pusher_idx) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) return; double e1 = args_markers.markers(ip, 0); @@ -51,7 +51,7 @@ extern "C" __global__ void push_eta_stage(double dt, int stage, MarkerArgs args_ // accumulation for last stage args_markers.markers(ip, first_free_idx + j) += dt * b[stage] * k[j]; // update positions for intermediate stages or last stage - args_markers.markers(ip, j) = args_markers.markers(ip, first_init_idx + j) + dt * a[stage] * k[j] + + args_markers.markers(ip, j) = args_markers.markers(ip, first_pusher_idx + j) + dt * a[stage] * k[j] + last * args_markers.markers(ip, first_free_idx + j); } diff --git a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py index b037d4746..e5f0245f1 100644 --- a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py @@ -43,7 +43,7 @@ def push_eta_stage( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx if stage == n_stages - 1: @@ -57,7 +57,7 @@ def push_eta_stage( # fmt: on for ip in range(n_markers): # check if marker is a hole or a boundary particle - if markers[ip, first_init_idx] == -1.0 or markers[ip, -1] == -2.0: + if markers[ip, first_pusher_idx] == -1.0 or markers[ip, -1] == -2.0: continue e1 = markers[ip, 0] @@ -85,7 +85,7 @@ def push_eta_stage( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * a[stage] * k + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py index 32e0b3d64..4f0a4ac73 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py @@ -82,25 +82,25 @@ def push_gc_Bstar_discrete_gradient_1st_order( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) eta_diff[:] = eta_k - eta_n v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v_mid = (v_k + v_n) / 2.0 v_diff = v_k - v_n diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py index fb88521fc..8f296d666 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py @@ -143,23 +143,23 @@ def push_gc_Bstar_discrete_gradient_1st_order_newton( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_k_shifted[:] = eta_k + markers[ip, first_shift_idx : first_shift_idx + 3] eta_diff[:] = eta_k_shifted - eta_n v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v_diff = v_k - v_n mu = markers[ip, mu_idx] diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py index 5fe3aadac..3c6b8d82e 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py @@ -108,25 +108,25 @@ def push_gc_Bstar_discrete_gradient_2nd_order( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) eta_diff[:] = eta_k - eta_n v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v_mid = (v_k + v_n) / 2.0 v_diff = v_k - v_n diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py index f47a562a9..03ac967b6 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py @@ -80,7 +80,7 @@ def push_gc_Bstar_explicit_multistage( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -93,7 +93,7 @@ def push_gc_Bstar_explicit_multistage( for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue # if stage == 0.: @@ -207,11 +207,13 @@ def push_gc_Bstar_explicit_multistage( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * a[stage] * k + last * markers[ip, first_free_idx : first_free_idx + 3] ) - markers[ip, 3] = markers[ip, first_init_idx + 3] + dt * a[stage] * k_v + last * markers[ip, first_free_idx + 3] + markers[ip, 3] = ( + markers[ip, first_pusher_idx + 3] + dt * a[stage] * k_v + last * markers[ip, first_free_idx + 3] + ) # apply kinetic boundary conditions pusher_utilities_kernels.apply_kinetic_bc_marker(ip, args_markers, args_domain, False) diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py index f383d4329..3d9c1a3e2 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py @@ -129,18 +129,18 @@ def push_gc_bxEstar_dg_1st_order_newton( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_k_shifted[:] = eta_k + markers[ip, first_shift_idx : first_shift_idx + 3] eta_diff[:] = eta_k_shifted - eta_n diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py index ca3507264..7c26a86dd 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py @@ -75,18 +75,18 @@ def push_gc_bxEstar_discrete_gradient_1st_order( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py index 619512ae4..bd0457365 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py @@ -91,18 +91,18 @@ def push_gc_bxEstar_discrete_gradient_2nd_order( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py index c68033f5d..905bfb8e0 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py @@ -66,7 +66,7 @@ def push_gc_bxEstar_explicit_multistage( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -79,7 +79,7 @@ def push_gc_bxEstar_explicit_multistage( for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta1 = markers[ip, 0] @@ -171,7 +171,7 @@ def push_gc_bxEstar_explicit_multistage( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * a[stage] * k + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py index 624ab42b5..c10e36e0a 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py @@ -80,7 +80,7 @@ def push_gc_cc_J2_dg_Hdiv( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): @@ -90,7 +90,7 @@ def push_gc_cc_J2_dg_Hdiv( # marker positions, mid point eta_old[:] = markers[ip, 0:3] - eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_init_idx : first_init_idx + 3]) / 2.0 + eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_pusher_idx : first_pusher_idx + 3]) / 2.0 eta_mid[:] = mod(eta_mid[:], 1.0) v = markers[ip, 3] @@ -204,6 +204,6 @@ def push_gc_cc_J2_dg_Hdiv( e /= abs_b_star_para e /= det_df - markers[ip, 0:3] = markers[ip, first_init_idx : first_init_idx + 3] - dt * e + markers[ip, 0:3] = markers[ip, first_pusher_idx : first_pusher_idx + 3] - dt * e markers[ip, 0:3] *= alpha markers[ip, 0:3] += eta_old * (1.0 - alpha) diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py index 8fcaafd1c..388c069cc 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py @@ -67,12 +67,12 @@ def push_gc_cc_J2_dg_init_Hdiv( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta1 = markers[ip, 0] diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py index a3cb0e4e5..d40c24fc7 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py @@ -85,7 +85,7 @@ def push_gc_cc_J2_stage_H1vec( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -98,7 +98,7 @@ def push_gc_cc_J2_stage_H1vec( for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta1 = markers[ip, 0] @@ -203,7 +203,7 @@ def push_gc_cc_J2_stage_H1vec( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] - dt * a[stage] * e + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py index ab5e5f89d..c2a82a5f5 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py @@ -81,7 +81,7 @@ def push_gc_cc_J2_stage_Hdiv( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -98,7 +98,7 @@ def push_gc_cc_J2_stage_Hdiv( continue # if particle is refilled - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta1 = markers[ip, 0] @@ -204,7 +204,7 @@ def push_gc_cc_J2_stage_Hdiv( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] - dt * a[stage] * e + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py index 5946bdf4b..e9fe1c0c5 100644 --- a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py @@ -66,7 +66,7 @@ def push_pc_eta_stage_H1vec( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -131,7 +131,7 @@ def push_pc_eta_stage_H1vec( # update markers for the next stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * k * a[stage] + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py index 33004306d..55e02b6ff 100644 --- a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py @@ -59,7 +59,7 @@ def push_pc_eta_stage_Hcurl( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -127,7 +127,7 @@ def push_pc_eta_stage_Hcurl( # update markers for the next stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * k * a[stage] + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py index 6ced6a2f9..0cae6016c 100644 --- a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py @@ -59,7 +59,7 @@ def push_pc_eta_stage_Hdiv( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -128,7 +128,7 @@ def push_pc_eta_stage_Hdiv( # update markers for the next stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * k * a[stage] + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu index bc96bfd77..d14f81ec6 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu @@ -28,10 +28,10 @@ extern "C" __global__ void push_vxb_analytic(double dt, int stage, MarkerArgs ar int n_markers = args_markers.n_markers; if (ip >= n_markers) return; - int first_init_idx = args_markers.first_init_idx; + int first_pusher_idx = args_markers.first_pusher_idx; // check if marker is a hole - if (args_markers.markers(ip, first_init_idx) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) + if (args_markers.markers(ip, first_pusher_idx) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) return; double e1 = args_markers.markers(ip, 0); diff --git a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py index 3341e4794..defa85cc3 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py @@ -54,7 +54,7 @@ def push_vxb_analytic( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx # fmt: off #$ omp parallel private (ip, e1, e2, e3, v, dfm, det_df, span1, span2, span3, b_form, b_cart, b_abs, b_norm, vpar, vxb_norm, vperp, b_normxvperp) @@ -62,7 +62,7 @@ def push_vxb_analytic( # fmt: on for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0 or markers[ip, -1] == -2.0: + if markers[ip, first_pusher_idx] == -1.0 or markers[ip, -1] == -2.0: continue e1 = markers[ip, 0] diff --git a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu index 1190f3e25..8fbb4a7e2 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu @@ -27,10 +27,10 @@ extern "C" __global__ void push_vxb_implicit(double dt, int stage, MarkerArgs ar int n_markers = args_markers.n_markers; if (ip >= n_markers) return; - int first_init_idx = args_markers.first_init_idx; + int first_pusher_idx = args_markers.first_pusher_idx; // check if marker is a hole - if (args_markers.markers(ip, first_init_idx) == -1.) return; + if (args_markers.markers(ip, first_pusher_idx) == -1.) return; double e1 = args_markers.markers(ip, 0); double e2 = args_markers.markers(ip, 1); diff --git a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py index 505959518..2b1c5facf 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py @@ -65,7 +65,7 @@ def push_vxb_implicit( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx # fmt: off #$ omp parallel firstprivate(b_prod) private (ip, e1, e2, e3, v, dfm, det_df, span1, span2, span3, b_form, b_cart, rhs, lhs, lhs_inv, vec, res) @@ -73,7 +73,7 @@ def push_vxb_implicit( # fmt: on for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue e1 = markers[ip, 0] diff --git a/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py b/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py index 07460ad4e..77cd4aa0a 100644 --- a/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py +++ b/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py @@ -39,7 +39,7 @@ def unit_b_1form( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -48,7 +48,7 @@ def unit_b_1form( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) diff --git a/src/struphy/pic/pushing/prepare_push_cuda.cu b/src/struphy/pic/pushing/prepare_push_cuda.cu index 5c1dbd28e..71c02556d 100644 --- a/src/struphy/pic/pushing/prepare_push_cuda.cu +++ b/src/struphy/pic/pushing/prepare_push_cuda.cu @@ -1,28 +1,28 @@ /** * Prepare the marker buffer for a push, as the three column-slice assignments at the start of Pusher._push: * - * markers[:, first_init_idx:first_shift_idx] = markers[:, :n_init] + * markers[:, first_pusher_idx:first_shift_idx] = markers[:, :n_init] * markers[:, first_shift_idx:residual_idx] = 0.0 * markers[:, residual_idx:-2] = 0.0 * - * One thread per entry of the contiguous column range [first_init_idx, n_cols - 2) of every row, so that + * One thread per entry of the contiguous column range [first_pusher_idx, n_cols - 2) of every row, so that * neighbouring threads write neighbouring addresses of the row-major buffer and the whole range is written * in a single pass (instead of one strided pass over the buffer per slice). * * @param markers Marker buffer (n_rows x n_cols, row-major), every row including holes. * @param n_rows Number of rows of the buffer. * @param n_cols Number of columns of the buffer. - * @param first_init_idx First column of the saved initial phase space coordinates. - * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_init_idx + n_init). + * @param first_pusher_idx First column of the saved initial phase space coordinates. + * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_pusher_idx + n_init). */ -extern "C" __global__ void prepare_push(double* markers, long long n_rows, int n_cols, int first_init_idx, +extern "C" __global__ void prepare_push(double* markers, long long n_rows, int n_cols, int first_pusher_idx, int n_init) { - long long width = n_cols - 2 - first_init_idx; + long long width = n_cols - 2 - first_pusher_idx; long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; if (t >= n_rows * width) return; long long ip = t / width; - int col = first_init_idx + (int)(t - ip * width); + int col = first_pusher_idx + (int)(t - ip * width); double* row = markers + ip * n_cols; - row[col] = col < first_init_idx + n_init ? row[col - first_init_idx] : 0.0; + row[col] = col < first_pusher_idx + n_init ? row[col - first_pusher_idx] : 0.0; } diff --git a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh index a40dc406a..48b06ef19 100644 --- a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh +++ b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh @@ -38,7 +38,7 @@ __device__ inline void reflect_velocity(long long ip, Array2D markers, * @param newton Accumulate periodic shifts if true; overwrite them otherwise. * * Call after updating the position. Boundary type 0 wraps periodically, - * type 1 mirrors position and logical velocity and sets first_init_idx to -1, + * type 1 mirrors position and logical velocity and sets first_pusher_idx to -1, * type 2 clears all columns except the ID to -1, and type 3 is handled on * the host. The marker view carries shape and strides; j is a CUDA-only * column index replacing Pyccel array slices. @@ -46,7 +46,7 @@ __device__ inline void reflect_velocity(long long ip, Array2D markers, __device__ inline void apply_kinetic_bc_marker(int ip, const MarkerArgs& args_markers, const DomainArgs& args_domain, bool newton) { const long long* bc_type = args_markers.bc_type; - int first_init_idx = args_markers.first_init_idx; + int first_pusher_idx = args_markers.first_pusher_idx; int first_shift_idx = args_markers.first_shift_idx; // Remove markers before applying periodic or reflecting boundaries. @@ -92,7 +92,7 @@ __device__ inline void apply_kinetic_bc_marker(int ip, const MarkerArgs& args_ma } } if (n_reflected == 0) return; - args_markers.markers(ip, first_init_idx) = -1.; + args_markers.markers(ip, first_pusher_idx) = -1.; for (int axis = 0; axis < 3; ++axis) { if (reflected[axis]) reflect_velocity(ip, args_markers.markers, args_domain, axis); } diff --git a/src/struphy/pic/pushing/pusher_utilities_kernels.py b/src/struphy/pic/pushing/pusher_utilities_kernels.py index 78b9c7d70..b21c1bfb0 100644 --- a/src/struphy/pic/pushing/pusher_utilities_kernels.py +++ b/src/struphy/pic/pushing/pusher_utilities_kernels.py @@ -27,7 +27,7 @@ def apply_kinetic_bc_marker( * 0 (periodic): wrap :math:`\eta_i` into :math:`[0, 1)` and store the shift :math:`\pm 1` at ``first_shift_idx + axis`` (set for ``newton=False``, accumulated for ``newton=True``). * 1 (reflect): mirror :math:`\eta_i` at the boundary, flip the logical velocity component - and mark the marker as done for multi-stage/iterative pushers (``first_init_idx = -1``). + and mark the marker as done for multi-stage/iterative pushers (``first_pusher_idx = -1``). * 2 (remove): turn the marker into a hole (all columns except the ID set to -1). * 3 (refill): skipped here, handled by :meth:`~struphy.pic.base.Particles.apply_kinetic_bc`. @@ -48,7 +48,7 @@ def apply_kinetic_bc_marker( markers = args_markers.markers bc_type = args_markers.bc_type - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx # remove @@ -94,7 +94,7 @@ def apply_kinetic_bc_marker( if n_reflected == 0: return - markers[ip, first_init_idx] = -1.0 + markers[ip, first_pusher_idx] = -1.0 # reflect velocities (Jacobian at the reflected position) dfm = empty((3, 3), dtype=float) diff --git a/src/struphy/propagators/current_coupling_5d_gradb.py b/src/struphy/propagators/current_coupling_5d_gradb.py index 9fac1ec42..a596074df 100644 --- a/src/struphy/propagators/current_coupling_5d_gradb.py +++ b/src/struphy/propagators/current_coupling_5d_gradb.py @@ -459,14 +459,14 @@ def __call__(self, dt): holes = particles.holes args_markers = particles.args_markers markers = args_markers.markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # clear buffer - markers[:, first_init_idx:-2] = 0.0 + markers[:, first_pusher_idx:-2] = 0.0 # save old marker positions - markers[:, first_init_idx : first_init_idx + 3] = markers[:, :3] + markers[:, first_pusher_idx : first_pusher_idx + 3] = markers[:, :3] # sum up total magnetic field b_full1 = b_eq + b_tilde (in-place) b_full = self._b2.copy(out=self._b_full) @@ -530,7 +530,7 @@ def __call__(self, dt): diffs = self.update_feec_variables(u=u_new) # clear the buffer - markers[:, first_init_idx:-2] = 0.0 + markers[:, first_pusher_idx:-2] = 0.0 # update_weights if self.variables.energetic_ions.species.weights_params.control_variate: @@ -661,7 +661,7 @@ def __call__(self, dt): sum_u_diff_loc = xp.sum((u_diff.toarray() ** 2)) sum_H_diff_loc = xp.sum( - (markers[~holes, :3] - markers[~holes, first_init_idx : first_init_idx + 3]) ** 2, + (markers[~holes, :3] - markers[~holes, first_pusher_idx : first_pusher_idx + 3]) ** 2, ) buffer_array = xp.array([sum_u_diff_loc]) @@ -853,7 +853,7 @@ def __call__(self, dt): diffs = self.update_feec_variables(u=u_new) # clear the buffer - markers[:, first_init_idx:-2] = 0.0 + markers[:, first_pusher_idx:-2] = 0.0 # update_weights if self.variables.energetic_ions.species.weights_params.control_variate: From c3ea6406b6279053b75bab7cd4226fe8a73a33f6 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 20:49:16 +0200 Subject: [PATCH 24/50] 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 e5d12d3de..38843c763 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 @@ -1494,6 +1498,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 958c0133d105ab3625c90b5c9b3783f7faeea174 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Thu, 8 Oct 2026 09:22:24 +0200 Subject: [PATCH 25/50] 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 From 08665fef07f8b0075f661f78cae43ddce5abe16e Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 20:49:16 +0200 Subject: [PATCH 26/50] Marker exchange between ranks on the device (CUDA-aware MPI) On the CuPy backend, Particles._sendrecv_markers hands the marker rows to Isend/Irecv as device buffers through cunumpy.mpi.mpi_buffer (directly after synchronize_for_mpi with CUDA-aware MPI, staged through host memory otherwise). Rows from each rank arrive in a contiguous slice of one device receive buffer, and one device scatter fills the holes. The per-rank counts are small NumPy arrays on every backend (they are host integers already), and a uniform alpha is a Python scalar, so the exchange makes no host/device copy and no implicit host read of a device array. The NumPy path is unchanged. Adds pic/tests/test_sorting_device.py (fake CuPy under mpirun, and a GPU variant) and a CI step running it on the fake CuPy. Solves #698, part of #650. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/test-PR-unit-mpi.yml | 6 + CUDA_STRATEGY.md | 20 +- src/struphy/pic/base.py | 129 ++++++++++++- src/struphy/pic/tests/test_sorting_device.py | 190 +++++++++++++++++++ 4 files changed, 338 insertions(+), 7 deletions(-) create mode 100644 src/struphy/pic/tests/test_sorting_device.py diff --git a/.github/workflows/test-PR-unit-mpi.yml b/.github/workflows/test-PR-unit-mpi.yml index e3a13dbc5..7bb53fbe2 100644 --- a/.github/workflows/test-PR-unit-mpi.yml +++ b/.github/workflows/test-PR-unit-mpi.yml @@ -76,3 +76,9 @@ jobs: source /struphy_fortran_/env_fortran_/bin/activate cd ${{ env.STRUPHY_PATH }} mpirun --oversubscribe -n ${{ matrix.n-procs }} pytest -v --with-mpi -m mpi_pic pic/tests + + - name: Run the marker exchange on the fake CuPy + run: | + source /struphy_fortran_/env_fortran_/bin/activate + cd ${{ env.STRUPHY_PATH }} + mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi pic/tests/test_sorting_device.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 6d2c60e72..68035f68c 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -365,7 +365,7 @@ matrix data). - **Kernel launch configuration.** One thread per marker with a fixed block size for now. Kernels over grid points (accumulation, FEEC) need their own launch sizes; `CudaKernel` needs a `grid`/`block` override (PR 12+), or cunumpy's. - **Accumulation strategy.** Atomics vs. sort-then-reduce; see PR 12+. Decided by measurement on the first accumulation kernel. - **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. `xp.mpi_is_cuda_aware()` detects it; the exchange in `Particles.mpi_sort_markers` has to be checked for host staging buffers. +- **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 `Derham` with `polar_splines=True` raises on CuPy (`PolarExtractionBlocksC1` builds SciPy sparse matrices, and the @@ -644,3 +644,21 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a The GPU tests have not run on an H100; the CPU emulation and the CPU regression tests are the gate. - **Parity cases.** The spline mappings are added to the geometry kernels' entries of `pic/tests/cuda_parity_cases.py` (the per-folder `_test_args.py` modules were replaced by that module in PR 16). + +## Marker exchange implementation notes (#698) + +- `Particles._sendrecv_markers` on CuPy (`_sendrecv_markers_device`): the rows to send (gathered on the device per + destination) and one device receive buffer go to `Isend`/`Irecv` through `xp.mpi.mpi_buffer`, i.e. directly after + `synchronize_for_mpi` with CUDA-aware MPI, staged through (pinned) host memory otherwise (with a warning, once). + Rows from rank `i` arrive in a contiguous slice of the receive buffer; one device scatter fills the holes. + Zero-count messages are skipped (both sides know the counts). The NumPy path is unchanged. +- Whether MPI is CUDA-aware: the answer recorded by `xp.mpi.mpi_is_cuda_aware()`/`set_mpi_cuda_aware()`, else a + collective probe of `mpi_comm` at the first exchange (`mpi_sort_markers` is collective). +- Counts: `send_info`/`recv_info` are small NumPy arrays on every backend. The send counts are sizes of `nonzero` + results, which are host integers already (CuPy synchronizes to size them), so the `Alltoall` of the counts needs no + copy. Uniform `alpha` is a Python scalar (no host-to-device copy of a 3-vector). +- Tests: `pic/tests/test_sorting_device.py` sorts the same markers on NumPy and on CuPy over all ranks and compares + the marker arrays; with CUDA-aware MPI no transfer (`count_transfers`) and no implicit host read of a device array + during the exchange, staged: one `to_host` per non-empty send and one `to_device`. Without a GPU it runs on the fake + CuPy (`mpirun -n N env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi pic/tests/test_sorting_device.py`, in the PIC MPI + workflow); the GPU variant has not run (no GPU here). diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 99db7dbfd..cae11c476 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2,6 +2,7 @@ import os import warnings from abc import ABCMeta, abstractmethod +from contextlib import ExitStack from typing import TYPE_CHECKING import cunumpy as xp @@ -1757,6 +1758,12 @@ def mpi_sort_markers( remove_ghost : bool Remove ghost particles before send. + + Notes + ----- + On the CuPy backend, the marker rows are exchanged as device buffers with CUDA-aware MPI + (staged through host memory otherwise), see :meth:`_sendrecv_markers_device`. The + per-rank counts are small host arrays on every backend (:meth:`_sendrecv_all_to_all`). """ if remove_ghost: self._remove_ghost_particles() @@ -4363,9 +4370,19 @@ def _sendrecv_determine_mtbs( """ # position that determines the sorting (including periodic shift of boundary conditions) if not isinstance(alpha, xp.ndarray): - alpha = xp.array(alpha, dtype=float) - assert alpha.size == 3 - assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) + # check a tuple/list on the host: on CuPy, asserting on a device comparison reads it back + alpha_host = np.array(alpha, dtype=float) + assert alpha_host.size == 3 + assert np.all(alpha_host >= 0.0) and np.all(alpha_host <= 1.0) + if alpha_host[0] == alpha_host[1] == alpha_host[2]: + # the common case (alpha = 1 or alpha = 0.5 in every direction): a Python scalar, + # which needs no host-to-device copy on CuPy and gives the same values on NumPy + alpha = float(alpha_host[0]) + else: + alpha = xp.asarray(alpha_host) + else: + assert alpha.size == 3 + assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) bi = self.first_pusher_idx self._sorting_etas[:] = ( alpha * (self.markers[:, :3] + self.markers[:, bi + 3 + self.vdim : bi + 3 + self.vdim + 3]) @@ -4465,8 +4482,11 @@ def _sendrecv_get_destinations(self, send_inds): Amount of particles sent to i-th process. """ - # One entry for each process - send_info = xp.zeros(self.mpi_size, dtype=int) + # One entry for each process. The counts are host integers on every backend: the + # size of a nonzero() result is known on the host (CuPy synchronizes to size it), so + # a small NumPy array holds them without any host/device copy, and the Alltoall of + # the counts (_sendrecv_all_to_all) passes host buffers to MPI. + send_info = np.zeros(self.mpi_size, dtype=int) # Gathered once and reused for every rank below, instead of re-gathering # self.markers[send_inds] and self._sorting_etas[send_inds] fresh on every @@ -4530,9 +4550,16 @@ def _sendrecv_all_to_all(self, send_info): ------- recv_info : array[int] Amount of marticles to be received from i-th process. + + Notes + ----- + ``send_info`` and ``recv_info`` are small NumPy arrays (one integer per rank) on every + backend: the counts size the receive buffers and the holes to fill, which Python needs + on the host anyway. On CuPy, only the marker rows themselves are device buffers + (:meth:`_sendrecv_markers`). """ - recv_info = xp.zeros(self.mpi_size, dtype=int) + recv_info = np.zeros(self.mpi_size, dtype=int) self.mpi_comm.Alltoall(send_info, recv_info) @@ -4551,6 +4578,10 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): Indices of empty rows in markers after send. """ + if xp.get_array_backend(self._markers) == "cupy": + self._sendrecv_markers_device(recv_info, hole_inds_after_send) + return + # i-th entry holds the number (not the index) of the first hole to be filled by data from process i first_hole = xp.cumsum(recv_info) - recv_info @@ -4594,6 +4625,92 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): for req in send_reqs: req.Wait() + def _mpi_cuda_aware(self) -> bool: + """Whether MPI takes device (CuPy) buffers, for :meth:`_sendrecv_markers_device`. + + Uses the answer recorded by :func:`cunumpy.mpi.mpi_is_cuda_aware` or + :func:`cunumpy.mpi.set_mpi_cuda_aware` (e.g. by feectools or at startup). If nothing + was recorded, probes :attr:`mpi_comm` once; the probe is collective, which is safe + here because :meth:`mpi_sort_markers` is called by every rank of :attr:`mpi_comm`. + """ + cuda_aware = xp.mpi.get_mpi_cuda_aware() + if cuda_aware is None: + cuda_aware = xp.mpi.mpi_is_cuda_aware(self.mpi_comm) + if not cuda_aware and not Particles._warned_mpi_staging: + Particles._warned_mpi_staging = True + logger.warning( + "MPI is not CUDA-aware: marker rows exchanged between ranks are staged through host memory.", + ) + return cuda_aware + + _warned_mpi_staging = False + + def _sendrecv_markers_device(self, recv_info, hole_inds_after_send): + """:meth:`_sendrecv_markers` on the CuPy backend: marker rows go to MPI as device buffers. + + With CUDA-aware MPI (:meth:`_mpi_cuda_aware`), the rows to send (:attr:`_send_list`, + gathered on the device by :meth:`_sendrecv_get_destinations`) and one device receive + buffer are handed to ``Isend``/``Irecv`` directly, after + :func:`cunumpy.mpi.synchronize_for_mpi` (inside :func:`cunumpy.mpi.mpi_buffer`), so MPI + reads what the kernels wrote. The rows from rank ``i`` arrive in a contiguous slice of + the receive buffer, and one device scatter puts all of them into the holes. Without + CUDA-aware MPI, :func:`~cunumpy.mpi.mpi_buffer` stages the same buffers through host + memory (counted by :func:`cunumpy.profiling.count_transfers`). + + The bookkeeping stays on the host without host/device copies: the counts are host + integers (``recv_info`` from the Alltoall, the send counts are the row counts of the + send buffers), and so is the number of holes. Ranks exchange no message when the + count is zero (both sides know the count, so they skip consistently). + + Parameters + ---------- + recv_info : array[int] + Amount of markers to be received from i-th process (NumPy array). + + hole_inds_after_send : array[int] + Indices of empty rows in markers after send (device array). + """ + recv_info = np.asarray(recv_info) + n_recv = int(recv_info.sum()) + if hole_inds_after_send.size < n_recv: + warnings.warn( + f'Strong load imbalance detected: \ +number of holes ({hole_inds_after_send.size}) on rank {self.mpi_rank} \ +is smaller than number of incoming particles ({n_recv}). \ +Increasing the value of "bufsize" in the markers parameters for the next run.', + ) + self.mpi_comm.Abort() + + cuda_aware = self._mpi_cuda_aware() + + # rows from rank i land in recvbuf[first_row[i] : first_row[i] + recv_info[i]] + first_row = np.cumsum(recv_info) - recv_info + recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=float) + + with ExitStack() as buffers: + if n_recv > 0: + recv_mpi = buffers.enter_context( + xp.mpi.mpi_buffer(recvbuf, send=False, recv=True, cuda_aware=cuda_aware), + ) + reqs = [] + for i in range(self.mpi_size): + if i == self.mpi_rank: + continue + if recv_info[i] > 0: + rows = recv_mpi[first_row[i] : first_row[i] + recv_info[i]] + reqs.append(self.mpi_comm.Irecv(rows, source=i, tag=i)) + data = self._send_list[i] + if data.shape[0] > 0: + send_mpi = buffers.enter_context(xp.mpi.mpi_buffer(data, cuda_aware=cuda_aware)) + reqs.append(self.mpi_comm.Isend(send_mpi, dest=i, tag=self.mpi_rank)) + # wait inside the block: the buffers must stay valid until MPI is done, and staged + # receive rows are copied to the device when the block ends + for req in reqs: + req.Wait() + + if n_recv > 0: + self._markers[hole_inds_after_send[:n_recv]] = recvbuf + class Tesselation: """ diff --git a/src/struphy/pic/tests/test_sorting_device.py b/src/struphy/pic/tests/test_sorting_device.py new file mode 100644 index 000000000..55c5d09f5 --- /dev/null +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -0,0 +1,190 @@ +"""Marker exchange between ranks on the CuPy backend (:meth:`Particles.mpi_sort_markers`). + +The marker rows are exchanged as device buffers with CUDA-aware MPI, and staged through +host memory otherwise. Without a GPU, the tests run on cunumpy's fake CuPy, which must be +installed before Python imports cunumpy, so they are skipped unless the fake is active:: + + mpirun -n 2 env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi src/struphy/pic/tests/test_sorting_device.py + +Fake CuPy arrays live in host memory, so a non-CUDA-aware MPI library takes them, and the +CUDA-aware path (the probe of :func:`cunumpy.mpi.mpi_is_cuda_aware` succeeds) is what runs by +default. The staged path is forced with :func:`cunumpy.mpi.set_mpi_cuda_aware`. +""" + +import sys +from contextlib import contextmanager + +import cunumpy as xp +import numpy as np +import pytest +from cunumpy.kernel_testing import requires_cupy +from maybempi import MPI + +from struphy import BoundaryParameters, LoadingParameters +from struphy.pic.particles import Particles6D + + +def _fake_cupy_active() -> bool: + return bool(getattr(sys.modules.get("cupy"), "__cunumpy_fake__", False)) + + +requires_fake_cupy = pytest.mark.skipif( + not _fake_cupy_active(), + reason="needs the fake CuPy: run under `mpirun -n N env CUNUMPY_FAKE_CUPY=1 pytest ...`", +) + + +@contextmanager +def _count_host_reads(): + """Count reads of fake CuPy arrays by the host (``int()``, ``bool()``, iteration, ``.get()``, ...). + + :func:`cunumpy.profiling.count_transfers` sees the copies made through cunumpy; this also sees + the implicit ones (a device scalar in an ``if``, ``list(device_array)``), which synchronize + and copy on a GPU. + """ + import cupy + + reads = [] + names = ("__int__", "__float__", "__bool__", "__index__", "__iter__", "get") + originals = {name: getattr(cupy.ndarray, name) for name in names} + original_getattr = cupy.ndarray.__getattr__ + + def counting(name, method): + def wrapper(self, *args, **kwargs): + reads.append(name) + return method(self, *args, **kwargs) + + return wrapper + + def counting_getattr(self, name): + if name in ("item", "tolist"): + reads.append(name) + return original_getattr(self, name) + + for name, method in originals.items(): + setattr(cupy.ndarray, name, counting(name, method)) + cupy.ndarray.__getattr__ = counting_getattr + try: + yield reads + finally: + for name, method in originals.items(): + setattr(cupy.ndarray, name, method) + cupy.ndarray.__getattr__ = original_getattr + + +def _make_particles(comm, bc): + return Particles6D( + comm_world=comm, + loading_params=LoadingParameters(Np=2000, seed=1234), + boundary_params=BoundaryParameters(bc=(bc, bc, bc)), + ) + + +def _shift_positions(markers: np.ndarray, valid: np.ndarray, step: int) -> np.ndarray: + """New positions in [0, 1) for the valid rows, far enough away that most markers change rank.""" + new = markers.copy() + shift = np.array([0.37, 0.21, 0.53]) * (step + 1) + new[valid, :3] = (markers[valid, :3] + shift) % 1.0 + return new + + +def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", count_reads: bool = True): + """Sort the same markers on NumPy and CuPy across all ranks and compare the marker arrays. + + ``cuda_aware=None`` uses the probe (:func:`cunumpy.mpi.mpi_is_cuda_aware`); ``True``/``False`` + force the device-buffer or the host-staged exchange. + """ + comm = MPI.COMM_WORLD + rank, size = comm.Get_rank(), comm.Get_size() + + ref = _make_particles(comm, bc) + ref.draw_markers(sort=False) + ref.mpi_sort_markers() + + # the reference is sorted on NumPy, the device markers inside `use_backend("cupy")` blocks + on_device = lambda: xp.use_backend("cupy") # noqa: E731 + + previous = xp.mpi.get_mpi_cuda_aware() + try: + with on_device(): + if cuda_aware is not None: + xp.mpi.set_mpi_cuda_aware(cuda_aware) + else: + xp.mpi.set_mpi_cuda_aware(None) + cuda_aware = xp.mpi.mpi_is_cuda_aware(comm) + + dev = _make_particles(comm, bc) + assert xp.get_array_backend(dev.markers) == "cupy" + assert dev.markers.shape == ref.markers.shape + + for step in range(2): + new = _shift_positions(ref.markers, ref.valid_mks, step) + + # number of rows this rank has to send, from the host reference + etas = new[ref.valid_mks, :3] + left, right = ref.domain_array[rank, 0::3], ref.domain_array[rank, 1::3] + right = np.where(right == 1.0, np.nextafter(right, 2.0), right) + n_send = int(np.count_nonzero(~np.all((etas >= left) & (etas < right), axis=1))) + n_send_total = comm.allreduce(n_send) + assert size == 1 or n_send_total > 0 + + ref.markers[:] = new + ref.update_holes() + n_before = ref.n_mks_loc + ref.mpi_sort_markers(do_test=True) + n_recv = ref.n_mks_loc - (n_before - n_send) + + with on_device(): + dev.markers[:] = xp.to_cupy(new) + dev.update_holes() + with xp.profiling.count_transfers() as transfers, _count_host_reads() as reads: + dev.mpi_sort_markers() + + if cuda_aware: + assert transfers.total == 0, transfers.report() + else: + # one staged copy per non-empty send buffer, one copy back of the receive buffer + n_dest = sum(1 for i in range(size) if i != rank and dev._send_list[i].shape[0] > 0) + assert transfers.to_host == n_dest, transfers.report() + assert transfers.to_device == (1 if n_recv > 0 else 0), transfers.report() + sent_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_host") + assert sent_bytes == n_send * ref.markers.shape[1] * 8 + + if count_reads and cuda_aware: + # (the staged exchange reads the send buffers with .get(), counted above) + assert reads == [], f"host reads of device arrays during the exchange: {reads}" + + with on_device(): + assert np.array_equal(xp.to_numpy(dev.markers), ref.markers) + assert np.array_equal(xp.to_numpy(dev.valid_mks), ref.valid_mks) + n_dev = dev.n_mks_loc + assert n_dev == ref.n_mks_loc + + with on_device(): + # all markers are on the right rank now: nothing moves + dev.mpi_sort_markers(do_test=True) + assert np.array_equal(xp.to_numpy(dev.markers), ref.markers) + finally: + xp.mpi.set_mpi_cuda_aware(previous) + + +@requires_fake_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("cuda_aware", [None, True, False]) +@pytest.mark.parametrize("bc", ["periodic", "remove"]) +def test_mpi_sort_markers_fake_cupy(cuda_aware, bc): + """On the fake CuPy: the exchange gives the NumPy result, without host/device copies when MPI is CUDA-aware.""" + check_mpi_sort_markers_device(cuda_aware, bc=bc) + + +@requires_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("bc", ["periodic", "remove"]) +def test_mpi_sort_markers_gpu(bc): + """On a GPU (``mpirun -n 2``, one GPU per rank): the exchange gives the NumPy result; with + CUDA-aware MPI without host/device copies.""" + check_mpi_sort_markers_device(None, bc=bc, count_reads=False) + + +if __name__ == "__main__": + check_mpi_sort_markers_device(None) From d4abdf6b0eeb9ef882f3bddfecbc95480bac3c90 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 9 Oct 2026 16:21:37 +0200 Subject: [PATCH 27/50] Use xp.mpi.exchange --- pyproject.toml | 2 +- src/struphy/pic/base.py | 85 +++++++++----------- src/struphy/pic/tests/test_sorting_device.py | 7 +- 3 files changed, 42 insertions(+), 52 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c75808225..fdeda2015 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 @ git+https://github.com/max-models/cunumpy.git@26a640b1d47d5677663785e49b1c44d343540df1", "pyccel>=2.2.0, <=2.2.3", "maybempi>=0.1.2", "feectools>=0.8.0, <=0.8.0", diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index cae11c476..5e1d217b2 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2,7 +2,6 @@ import os import warnings from abc import ABCMeta, abstractmethod -from contextlib import ExitStack from typing import TYPE_CHECKING import cunumpy as xp @@ -4646,67 +4645,55 @@ def _mpi_cuda_aware(self) -> bool: _warned_mpi_staging = False def _sendrecv_markers_device(self, recv_info, hole_inds_after_send): - """:meth:`_sendrecv_markers` on the CuPy backend: marker rows go to MPI as device buffers. - - With CUDA-aware MPI (:meth:`_mpi_cuda_aware`), the rows to send (:attr:`_send_list`, - gathered on the device by :meth:`_sendrecv_get_destinations`) and one device receive - buffer are handed to ``Isend``/``Irecv`` directly, after - :func:`cunumpy.mpi.synchronize_for_mpi` (inside :func:`cunumpy.mpi.mpi_buffer`), so MPI - reads what the kernels wrote. The rows from rank ``i`` arrive in a contiguous slice of - the receive buffer, and one device scatter puts all of them into the holes. Without - CUDA-aware MPI, :func:`~cunumpy.mpi.mpi_buffer` stages the same buffers through host - memory (counted by :func:`cunumpy.profiling.count_transfers`). - - The bookkeeping stays on the host without host/device copies: the counts are host - integers (``recv_info`` from the Alltoall, the send counts are the row counts of the - send buffers), and so is the number of holes. Ranks exchange no message when the - count is zero (both sides know the count, so they skip consistently). + """Exchange marker rows and scatter incoming rows into empty slots. + + :func:`cunumpy.mpi.exchange` handles GPU synchronization, optional host + staging, and completion of all MPI transfers. Parameters ---------- - recv_info : array[int] - Amount of markers to be received from i-th process (NumPy array). - - hole_inds_after_send : array[int] - Indices of empty rows in markers after send (device array). + recv_info : array[int] + Number of markers to receive from each process (NumPy array). + hole_inds_after_send : array[int] + Indices of empty rows in markers after send (device array). """ recv_info = np.asarray(recv_info) n_recv = int(recv_info.sum()) + if hole_inds_after_send.size < n_recv: warnings.warn( - f'Strong load imbalance detected: \ -number of holes ({hole_inds_after_send.size}) on rank {self.mpi_rank} \ -is smaller than number of incoming particles ({n_recv}). \ -Increasing the value of "bufsize" in the markers parameters for the next run.', + f"Strong load imbalance detected: " + f"number of holes ({hole_inds_after_send.size}) " + f"on rank {self.mpi_rank} is smaller than number of " + f"incoming particles ({n_recv}). " + 'Increase "bufsize" in the markers parameters for the next run.', ) self.mpi_comm.Abort() - cuda_aware = self._mpi_cuda_aware() - - # rows from rank i land in recvbuf[first_row[i] : first_row[i] + recv_info[i]] + recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=self.markers.dtype) first_row = np.cumsum(recv_info) - recv_info - recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=float) - with ExitStack() as buffers: - if n_recv > 0: - recv_mpi = buffers.enter_context( - xp.mpi.mpi_buffer(recvbuf, send=False, recv=True, cuda_aware=cuda_aware), - ) - reqs = [] - for i in range(self.mpi_size): - if i == self.mpi_rank: - continue - if recv_info[i] > 0: - rows = recv_mpi[first_row[i] : first_row[i] + recv_info[i]] - reqs.append(self.mpi_comm.Irecv(rows, source=i, tag=i)) - data = self._send_list[i] - if data.shape[0] > 0: - send_mpi = buffers.enter_context(xp.mpi.mpi_buffer(data, cuda_aware=cuda_aware)) - reqs.append(self.mpi_comm.Isend(send_mpi, dest=i, tag=self.mpi_rank)) - # wait inside the block: the buffers must stay valid until MPI is done, and staged - # receive rows are copied to the device when the block ends - for req in reqs: - req.Wait() + sends = [] + receives = [] + for rank in range(self.mpi_size): + if rank == self.mpi_rank: + continue + + if recv_info[rank] > 0: + start = int(first_row[rank]) + stop = start + int(recv_info[rank]) + receives.append((recvbuf[start:stop], rank, rank)) + + data = self._send_list[rank] + if data.shape[0] > 0: + sends.append((data, rank, self.mpi_rank)) + + xp.mpi.exchange( + self.mpi_comm, + sends=sends, + receives=receives, + cuda_aware=self._mpi_cuda_aware(), + ) if n_recv > 0: self._markers[hole_inds_after_send[:n_recv]] = recvbuf diff --git a/src/struphy/pic/tests/test_sorting_device.py b/src/struphy/pic/tests/test_sorting_device.py index 55c5d09f5..dc8fa9462 100644 --- a/src/struphy/pic/tests/test_sorting_device.py +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -143,12 +143,15 @@ def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", if cuda_aware: assert transfers.total == 0, transfers.report() else: - # one staged copy per non-empty send buffer, one copy back of the receive buffer + # exchange stages each non-empty send and receive buffer separately n_dest = sum(1 for i in range(size) if i != rank and dev._send_list[i].shape[0] > 0) + n_sources = sum(1 for i in range(size) if i != rank and ref._recvbufs[i].shape[0] > 0) assert transfers.to_host == n_dest, transfers.report() - assert transfers.to_device == (1 if n_recv > 0 else 0), transfers.report() + assert transfers.to_device == n_sources, transfers.report() sent_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_host") assert sent_bytes == n_send * ref.markers.shape[1] * 8 + received_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_device") + assert received_bytes == n_recv * ref.markers.shape[1] * 8 if count_reads and cuda_aware: # (the staged exchange reads the send buffers with .get(), counted above) From 22ed9e050803a569a399b5696cd7ea3fcee724b6 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 19:12:45 +0200 Subject: [PATCH 28/50] CUDA version of the linear_vlasov_ampere accumulation - linear_vlasov_ampere_cuda.cu: same arguments as pyccel, one thread per marker row, matrix blocks as Array6D views (cunumpy 0.6.1) - device fill_mat, fill_mat_vec (filler_kernels.cuh), m_v_fill_b_v1_symm (particle_to_mat_kernels.cuh) and outer (linalg_kernels.cuh), with atomic adds - parity cases (Cuboid, Colella, HollowTorus, 3d spline) and a device-helper test of m_v_fill_b_v1_symm, on a GPU and by CPU emulation - test_accum_matrix_cupy.py: the matrix path of Accumulator on CuPy vs NumPy (device data, ghost regions, transposed blocks, Schur operator and solve), without a GPU on the fake CuPy with emulated kernel launches (emulated_launches() in pic/tests/cuda_emulation.py) - CUDA_STRATEGY.md: porting order and implementation notes Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 59 ++++- src/struphy/linear_algebra/linalg_kernels.cuh | 11 + .../pic/accumulation/filler_kernels.cuh | 125 ++++++++++- .../linear_vlasov_ampere_cuda.cu | 90 ++++++++ .../accumulation/particle_to_mat_kernels.cuh | 69 ++++++ src/struphy/pic/tests/cuda_emulation.py | 66 +++++- src/struphy/pic/tests/cuda_parity_cases.py | 17 ++ src/struphy/pic/tests/kernel_test_args.py | 27 ++- .../pic/tests/test_accum_matrix_cupy.py | 209 ++++++++++++++++++ src/struphy/pic/tests/test_cuda_emulation.py | 12 + src/struphy/pic/tests/test_cuda_parity.py | 4 +- src/struphy/pic/tests/test_device_helpers.py | 61 +++++ 12 files changed, 730 insertions(+), 20 deletions(-) create mode 100644 src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu create mode 100644 src/struphy/pic/tests/test_accum_matrix_cupy.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 68035f68c..5f8635336 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -26,7 +26,8 @@ Progress is tracked in struphy-hub/struphy#650. - [ ] **PR 17: struphy on the feectools CUDA stack** — the `feectools` submodule points at the top of the feectools CUDA stack (`cuda-4-device-kernels`, [feectools#88](https://github.com/struphy-hub/feectools/pull/88)) instead of `devel-tiny`, so the struphy CUDA PRs run against feectools with device stencil operations, MPI with device buffers and one GPU per rank. Moves along with the stack; before the struphy stack is merged into `devel`, the feectools stack is merged into `devel-tiny` and the submodule points there again (see [feectools](#feectools)). - [x] **PR 18: Geometry evaluation on the GPU for all analytic mappings** (#681, open) — CUDA versions of the four geometry entry kernels (`kernel_evaluate_pic`, `kernel_evaluate`, `kernel_pullpush_pic`, `kernel_pullpush` in `geometry/kernels/`), built on device versions of the whole metric chain (`f`, `df`, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`/`push`/`tran`) for every analytic mapping (`kind_map` 10–12, 20–22, 30–32). Restores CuPy particle runs (weight initialization evaluates `jacobian_det`), and removes the Cuboid-only checks in `Pusher`, the accumulators and `reflect`. Parity arguments cover every analytic mapping (see [PR 18](#pr-18-geometry-evaluation-for-all-analytic-mappings) and the [PR 18 implementation notes](#pr-18-implementation-notes)). - [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines and `EQDSKequilibrium` stay host-only (see [PR 19](#pr-19-spline-mappings) and the [PR 19 implementation notes](#pr-19-implementation-notes)). -- [ ] **Next (order to be confirmed)**: 6D array views in cunumpy and the blocked matrix accumulations `vlasov_maxwell`, `linear_vlasov_ampere` (steps 1–2 of the [porting order](#porting-order)); one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, including when to compile the kernels, which needs the [feectools stack](#feectools) merged and released first; a decision on hand-written CUDA vs. code generation before the guiding-center kernels (step 3). +- [x] **`linear_vlasov_ampere` on the GPU** (part of #688) — the first matrix accumulation: CUDA version of `linear_vlasov_ampere` (`EfieldWeightsCoupling`) on cunumpy's `Array6D` views (cunumpy 0.6.1), device versions of `fill_mat`, `fill_mat_vec`, `m_v_fill_b_v1_symm` and `outer`, and the matrix path of `Accumulator` checked on the CuPy backend (device data, ghost regions, transposed blocks, Schur operator) without a GPU through cunumpy's fake CuPy with emulated kernel launches (see the [implementation notes](#linear_vlasov_ampere-implementation-notes)). +- [ ] **Next (order to be confirmed)**: the blocked matrix accumulation `vlasov_maxwell` (step 1 of the [porting order](#porting-order)); one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, including when to compile the kernels, which needs the [feectools stack](#feectools) merged and released first; a decision on hand-written CUDA vs. code generation before the guiding-center kernels (step 3). - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)). - [x] **PR 13: Move the kernel infrastructure to cunumpy** — `Kernel`, `KernelCatalog`, `CudaKernel` and `Argument` are replaced by their cunumpy counterparts; each owner has a single `args_*` object on both backends, and `pusher_args.cuh` is generated (see [Moving to cunumpy](#moving-to-cunumpy-pr-13)). Kernels and device helpers only change their includes. - [ ] **CI**: a GPU runner that runs the CUDA tests (can happen any time; until then the GPU tests are run by hand on an H100 before each PR that touches CUDA code is merged). @@ -64,7 +65,7 @@ CUDA kernels can be added one by one. If the code runs on the GPU and needs a ke | `src/struphy/utils/cuda_arguments.py` | `CUDA_STRUCTS`, `CUDA_OPTIONS`, `write_pusher_header()` and `write_local_projectors_header()` | | `src/struphy/kernel_arguments/pusher_args.cuh`, `local_projectors_args.cuh` | the C structs `MarkerArgs`, `DerhamArgs`, `DomainArgs` and `LocalProjectorsArgs`, generated from the CUDA classes, using cunumpy array views | | `src/struphy/geometry/base.py`, `src/struphy/pic/base.py`, `src/struphy/feec/psydac_derham.py` | `Domain.args_domain`, `Particles.args_markers` and `Derham.args_derham` are the pyccel class on NumPy and the CUDA class on CuPy, chosen once at construction; every kernel call goes through a `Kernel` | -| `src/struphy/*/kernels/` (`pic/pushing`, `pic/accumulation`, `pic/diagnostics`, `pic/sph`, `bsplines`, `geometry`, `feec`, `feec/local_projectors`) | every kernel called from Python with argument objects, one folder each: 44 pusher/evaluation (incl. `reflect`), 16 accumulation, 10 marker diagnostics, 4 SPH evaluation, 3 spline evaluation, 4 geometry, 1 FEEC utility and 8 local projector kernels. Each folder's `__init__.py` declares its `Kernel`; the code imports it (`from struphy.geometry.kernels.kernel_evaluate import kernel_evaluate`). Fifteen have CUDA versions, with parity cases in `pic/tests/cuda_parity_cases.py`, among them the four geometry kernels (PR 18) | +| `src/struphy/*/kernels/` (`pic/pushing`, `pic/accumulation`, `pic/diagnostics`, `pic/sph`, `bsplines`, `geometry`, `feec`, `feec/local_projectors`) | every kernel called from Python with argument objects, one folder each: 44 pusher/evaluation (incl. `reflect`), 16 accumulation, 10 marker diagnostics, 4 SPH evaluation, 3 spline evaluation, 4 geometry, 1 FEEC utility and 8 local projector kernels. Each folder's `__init__.py` declares its `Kernel`; the code imports it (`from struphy.geometry.kernels.kernel_evaluate import kernel_evaluate`). Sixteen have CUDA versions (`linear_vlasov_ampere` the latest), with parity cases in `pic/tests/cuda_parity_cases.py`, among them the four geometry kernels (PR 18) | | `src/struphy/geometry/evaluation_kernels.cuh`, `transform_kernels.cuh`, `spline_mappings_kernels.cuh`, `geometry/domains//_cuda.cuh` | device versions of the mapping helpers for every mapping: `f`/`df` per analytic domain (`kind_map` 10–12, 20–22, 30–32, PR 18) and the spline mappings `spline_3d`, `spline_2d_straight`, `spline_2d_torus` (`kind_map` 0–2, PR 19), the `kind_map` switch, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`, `push`, `tran` | | `src/struphy/pic/tests/test_cuda_parity.py`, `test_cuda_emulation.py` | parity of every CUDA kernel with pyccel on the cases of `pic/tests/cuda_parity_cases.py` (`PARITY_CASES`): on a GPU (`assert_kernels_agree`), and without one by CPU emulation | | `src/struphy/pic/tests/test_kernel_backends.py` | the generated-header tests, the check that each CUDA class mirrors its pyccel class, signature check of all kernels, owner class selection, and (on a GPU) struct layout and copy/pickle checks | @@ -277,7 +278,7 @@ from the struphy tracking issue ([#650](https://github.com/struphy-hub/struphy/i Kernels are ported in the order that completes one model after the other on the GPU. Each step lists only the kernels it adds; P = pusher, A = accumulation, E = marker evaluation. Status: -done (✓), in this PR (PR 14), blocked (⏸). +done (✓), in a PR (PR 14), blocked (⏸). **Step 0 – Vlasov** (PRs 11–12, ✓): `push_eta_stage`, `push_vxb_analytic`, `push_vxb_implicit`, `reflect`. @@ -289,14 +290,14 @@ done (✓), in this PR (PR 14), blocked (⏸). |---|---|---|---|---| | 1 | `push_v_with_efield` | P | `VlasovAmpereCoupling`, `PushVinForceField` | PR 14 | | 2 | `charge_density_0form` | A (vector) | initial Poisson solve of these models | PR 14 | -| 3 | `vlasov_maxwell` | A (matrix + vector) | `VlasovAmpereCoupling` | ⏸ 6D views | +| 3 | `vlasov_maxwell` | A (matrix + vector) | `VlasovAmpereCoupling` | next (6D views available) | **Step 2 – LinearVlasovAmpère/Maxwell (δf)** | # | Kernel | Type | Used by | Status | |---|---|---|---|---| | 4 | `push_weights_with_efield_lin_va` | P | `EfieldWeightsCoupling` | PR 14 | -| 5 | `linear_vlasov_ampere` | A (matrix + vector) | `EfieldWeightsCoupling` | ⏸ 6D views | +| 5 | `linear_vlasov_ampere` | A (matrix + vector) | `EfieldWeightsCoupling` | ✓ (#688) | **Step 3 – Guiding center: ToyDrift, then DriftKineticElectrostaticAdiabatic.** The default algorithm is `discrete_gradient_1st_order`; its chain comes first. @@ -348,9 +349,9 @@ The three spaces (H1vec/Hcurl/Hdiv) differ only in the basis; port one, then the 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, -and array views with more than 4 dimensions in cunumpy (all matrix accumulations write 6D stencil -matrix data). +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. ## Testing @@ -662,3 +663,45 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a during the exchange, staged: one `to_host` per non-empty send and one `to_device`. Without a GPU it runs on the fake CuPy (`mpirun -n N env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi pic/tests/test_sorting_device.py`, in the PIC MPI workflow); the GPU variant has not run (no GPU here). +## `linear_vlasov_ampere` implementation notes + +The first matrix accumulation on the GPU (part of #688, tracked in #650): `linear_vlasov_ampere`, the accumulation of +`EfieldWeightsCoupling` in `LinearVlasovAmpereOneSpecies` (step 2 of the [porting order](#porting-order)). + +- **6D views.** The six blocks `mat11, mat12, mat13, mat22, mat23, mat33` are `Array6D` (cunumpy 0.6.1, strided, + like the `Array3D` vectors); `cunumpy.kernels.CudaKernel` packs the `StencilMatrix._data` blocks without + copying them. `f0_values` is a `const double*`, as in `push_weights_with_efield_lin_va`. +- **Device helpers, one to one.** `filler_kernels.cuh` gets `fill_mat` and `fill_mat_vec` next to `fill_vec`, + `particle_to_mat_kernels.cuh` gets `m_v_fill_b_v1_symm` next to `vec_fill_b_v0`, and `linalg_kernels.cuh` gets + `outer`; same names, arguments and loop order as pyccel (`starts` and `pads` are `const long long*`). Every addition + into matrix or vector data is a `cunumpy_atomic_add`, as in `charge_density_0form`; the spline values pyccel keeps + in `args_derham.bn1, ..., bd3` are in the thread-local `SplineScratch`. Of the remaining `m_v_fill_*`/`mat_fill_*` + fillers, each one is ported with the first kernel that calls it. +- **Kernel.** `linear_vlasov_ampere_cuda.cu`: same arguments in the same order as pyccel, one thread per marker row, + holes and boundary particles skipped; `df`, `matrix_inv`, `matrix_vector`, `outer`, then `m_v_fill_b_v1_symm`. The + 3 x 3 `filling_m` is nine row-major entries; the CUDA-only `factor` is pyccel's `f0_values[ip] / (markers[ip, 7] * + n_markers_tot)`, computed once as pyccel does. +- **`Accumulator` on CuPy.** Nothing in `Accumulator` had to change: on the CuPy backend `StencilMatrix._data` and + `StencilVector._data` are device arrays (feectools allocates them with `xp`), the accumulator zeroes them in place and + passes the same arrays to the kernel; `exchange_assembly_data`, `update_ghost_regions`, the transposed blocks of the + symmetric matrix (`StencilMatrix.transpose`, a CUDA kernel in feectools), the scaled operator of `SchurSolver` and its + `dot` all work on them. `test_accum_matrix_cupy.py` checks this against NumPy: the kernel data are the blocks' data + (no copies), the nine arrays, the dense blocks, the accumulated vector, `BC.dot(x)` and three CG iterations of the + Schur solve agree, and nothing is copied to the host from the accumulation to `BC.dot(x)`. On a GPU it runs with + `assert_no_transfers` (not run yet); without one it runs in a subprocess on cunumpy's fake CuPy, with every + `CudaKernel` launch emulated on the CPU in place on the fake device arrays (`emulated_launches()` in + `pic/tests/cuda_emulation.py`; struphy's and feectools' kernels, struct arguments included). This would fit into + cunumpy (fake CuPy with emulated launches). +- **Found, not fixed here.** feectools' `StencilVector.inner` returns its result on the host (`xp.to_numpy`), so each + inner product of the CG solve copies a scalar from the device; the Schur solve is therefore outside the transfer + guard. `MassMatrixPreconditioner`, the default preconditioner of `EfieldWeightsCoupling`, cannot be created on CuPy + (`is_circulant` asserts a host array; the device Kronecker solve is a separate PR). `LinearVlasovAmpereOneSpecies` + end to end on the GPU (#689) also needs the feectools stencil 6D views and device Kronecker solve PRs and the first + H100 run (#687). +- **Tests.** Parity cases `PARITY_CASES["linear_vlasov_ampere"]`: 129 markers (a hole and a boundary particle) in + Cuboid, Colella, HollowTorus and a 3d spline mapping, data shaped like the real accumulation matrices of clamped + splines of degrees 2, 3, 1 on 8 cells (`v1_symm_accumulation_data()` in `pic/tests/kernel_test_args.py`); + `rtol = 1e-12`, `atol = 1e-10` (atomics change the summation order; the entries reach 1e4). The CPU emulation agrees + with pyccel to round-off and catches a transposed `DF`. `m_v_fill_b_v1_symm` is compared with pyccel through a + wrapper kernel (`test_v1_symm_filler` on a GPU, `test_emulated_v1_symm_filler` by emulation). GPU tests have not run + on an H100. diff --git a/src/struphy/linear_algebra/linalg_kernels.cuh b/src/struphy/linear_algebra/linalg_kernels.cuh index 1b72c25fe..4636f2ea3 100644 --- a/src/struphy/linear_algebra/linalg_kernels.cuh +++ b/src/struphy/linear_algebra/linalg_kernels.cuh @@ -100,4 +100,15 @@ __device__ inline void transpose(const double* a, double* b) { for (int i = 0; i < 3; ++i) for (int j = 0; j < 3; ++j) b[3 * i + j] = a[3 * j + i]; } +/** + * Outer product of two vectors of length 3, as in linalg_kernels.outer. + * + * @param a First input vector with three entries. + * @param b Second input vector with three entries. + * @param c Output matrix c_ij = a_i * b_j with nine row-major entries; must not overlap a or b. + */ +__device__ inline void outer(const double* a, const double* b, double* c) { + for (int i = 0; i < 3; ++i) + for (int j = 0; j < 3; ++j) c[3 * i + j] = a[i] * b[j]; +} } diff --git a/src/struphy/pic/accumulation/filler_kernels.cuh b/src/struphy/pic/accumulation/filler_kernels.cuh index 37f3867d9..8a33d7773 100644 --- a/src/struphy/pic/accumulation/filler_kernels.cuh +++ b/src/struphy/pic/accumulation/filler_kernels.cuh @@ -1,5 +1,5 @@ -// Device versions of the fillers in filler_kernels.py: add the contribution of one marker to a stencil vector. -// Many markers write to the same entries, so every addition is atomic. +// Device versions of the fillers in filler_kernels.py: add the contribution of one marker to a stencil vector +// or to a block of a stencil matrix. Many markers write to the same entries, so every addition is atomic. #pragma once #include "cunumpy/array_view.cuh" #include "cunumpy/atomic.cuh" @@ -40,4 +40,125 @@ __device__ inline void fill_vec(int pi1, int pi2, int pi3, const double* bi1, co } } } + +/** + * Add basis functions times filling to a block of a stencil matrix, as in filler_kernels.fill_mat. + * + * @param pi1 Spline degree of the codomain (row indices) along the first axis. + * @param pi2 Spline degree of the codomain along the second axis. + * @param pi3 Spline degree of the codomain along the third axis. + * @param pj1 Spline degree of the domain (column indices) along the first axis. + * @param pj2 Spline degree of the domain along the second axis. + * @param pj3 Spline degree of the domain along the third axis. + * @param bi1 The pi1 + 1 non-vanishing N- or D-spline values of the codomain along the first axis. + * @param bi2 The pi2 + 1 non-vanishing spline values of the codomain along the second axis. + * @param bi3 The pi3 + 1 non-vanishing spline values of the codomain along the third axis. + * @param bj1 The pj1 + 1 non-vanishing spline values of the domain along the first axis. + * @param bj2 The pj2 + 1 non-vanishing spline values of the domain along the second axis. + * @param bj3 The pj3 + 1 non-vanishing spline values of the domain along the third axis. + * @param span1 Knot span index along the first axis. + * @param span2 Knot span index along the second axis. + * @param span3 Knot span index along the third axis. + * @param starts Start indices of the codomain (three entries). + * @param pads Paddings of the codomain (three entries). + * @param mat Stencil matrix data (the _data of a StencilMatrix, shape (n1, n2, n3, 2 * pads + 1)), any strides; + * written atomically. + * @param filling Number multiplied by the basis functions. + * + * The summation order of concurrent threads (and the last bits of the result) depends on the thread schedule. + */ +__device__ inline void fill_mat(int pi1, int pi2, int pi3, int pj1, int pj2, int pj3, const double* bi1, + const double* bi2, const double* bi3, const double* bj1, const double* bj2, + const double* bj3, int span1, int span2, int span3, const long long* starts, + const long long* pads, Array6D mat, double filling) { + for (int il1 = 0; il1 <= pi1; ++il1) { + long long i1 = span1 + il1 - starts[0]; + double b1 = bi1[il1] * filling; + for (int il2 = 0; il2 <= pi2; ++il2) { + long long i2 = span2 + il2 - starts[1]; + double b2 = b1 * bi2[il2]; + for (int il3 = 0; il3 <= pi3; ++il3) { + long long i3 = span3 + il3 - starts[2]; + double b3 = b2 * bi3[il3]; + + for (int jl1 = 0; jl1 <= pj1; ++jl1) { + long long j1 = pads[0] + jl1 - il1; + double b4 = b3 * bj1[jl1]; + for (int jl2 = 0; jl2 <= pj2; ++jl2) { + long long j2 = pads[1] + jl2 - il2; + double b5 = b4 * bj2[jl2]; + for (int jl3 = 0; jl3 <= pj3; ++jl3) { + long long j3 = pads[2] + jl3 - il3; + double b6 = b5 * bj3[jl3]; + + cunumpy_atomic_add(&mat(i1, i2, i3, j1, j2, j3), b6); + } + } + } + } + } + } +} + +/** + * Add basis functions times fillings to a block of a stencil matrix and to a stencil vector, as in + * filler_kernels.fill_mat_vec. + * + * @param pi1 Spline degree of the codomain (row indices) along the first axis. + * @param pi2 Spline degree of the codomain along the second axis. + * @param pi3 Spline degree of the codomain along the third axis. + * @param pj1 Spline degree of the domain (column indices) along the first axis. + * @param pj2 Spline degree of the domain along the second axis. + * @param pj3 Spline degree of the domain along the third axis. + * @param bi1 The pi1 + 1 non-vanishing N- or D-spline values of the codomain along the first axis. + * @param bi2 The pi2 + 1 non-vanishing spline values of the codomain along the second axis. + * @param bi3 The pi3 + 1 non-vanishing spline values of the codomain along the third axis. + * @param bj1 The pj1 + 1 non-vanishing spline values of the domain along the first axis. + * @param bj2 The pj2 + 1 non-vanishing spline values of the domain along the second axis. + * @param bj3 The pj3 + 1 non-vanishing spline values of the domain along the third axis. + * @param span1 Knot span index along the first axis. + * @param span2 Knot span index along the second axis. + * @param span3 Knot span index along the third axis. + * @param starts Start indices of the codomain (three entries). + * @param pads Paddings of the codomain (three entries). + * @param mat Stencil matrix data (the _data of a StencilMatrix), any strides; written atomically. + * @param filling_mat Number multiplied by the basis functions and added to mat. + * @param vec Stencil vector data of the codomain (the _data of a StencilVector), any strides; written atomically. + * @param filling_vec Number multiplied by the basis functions of the codomain and added to vec. + */ +__device__ inline void fill_mat_vec(int pi1, int pi2, int pi3, int pj1, int pj2, int pj3, const double* bi1, + const double* bi2, const double* bi3, const double* bj1, const double* bj2, + const double* bj3, int span1, int span2, int span3, const long long* starts, + const long long* pads, Array6D mat, double filling_mat, + Array3D vec, double filling_vec) { + for (int il1 = 0; il1 <= pi1; ++il1) { + long long i1 = span1 + il1 - starts[0]; + double b1 = bi1[il1]; + for (int il2 = 0; il2 <= pi2; ++il2) { + long long i2 = span2 + il2 - starts[1]; + double b2 = b1 * bi2[il2]; + for (int il3 = 0; il3 <= pi3; ++il3) { + long long i3 = span3 + il3 - starts[2]; + double b3 = b2 * bi3[il3]; + + cunumpy_atomic_add(&vec(i1, i2, i3), b3 * filling_vec); + + for (int jl1 = 0; jl1 <= pj1; ++jl1) { + long long j1 = pads[0] + jl1 - il1; + double b4 = b3 * bj1[jl1] * filling_mat; + for (int jl2 = 0; jl2 <= pj2; ++jl2) { + long long j2 = pads[1] + jl2 - il2; + double b5 = b4 * bj2[jl2]; + for (int jl3 = 0; jl3 <= pj3; ++jl3) { + long long j3 = pads[2] + jl3 - il3; + double b6 = b5 * bj3[jl3]; + + cunumpy_atomic_add(&mat(i1, i2, i3, j1, j2, j3), b6); + } + } + } + } + } + } +} } diff --git a/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu b/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu new file mode 100644 index 000000000..387013af2 --- /dev/null +++ b/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu @@ -0,0 +1,90 @@ +#include "struphy/geometry/evaluation_kernels.cuh" +#include "struphy/linear_algebra/linalg_kernels.cuh" +#include "struphy/pic/accumulation/particle_to_mat_kernels.cuh" + +/** + * Accumulate the matrix and vector of the linear Vlasov-Ampere system (delta-f) into V1, as in + * linear_vlasov_ampere_kernels.linear_vlasov_ampere. + * + * Filling functions (the factors alpha^2 kappa^2 / v_th^2 and alpha^2 kappa are applied by the propagator): + * A_p^{mu, nu} = f_0(eta_p, v_p) / (N s_0) [DF^{-1}(eta_p) v_p]_mu [DF^{-1}(eta_p) v_p]_nu, + * B_p^mu = w_p [DF^{-1}(eta_p) v_p]_mu; + * one thread per marker row. + * + * @param args_markers Marker buffer (n_markers x n_cols, row-major) and the total marker number Np; the weight is in + * column 6 and s_0 in column 7, as in pyccel. + * @param args_derham Spline degrees (1 to 8), knots and start indices. + * @param args_domain Mapping arguments; every spline and analytic mapping. + * @param mat11 (1,1)-block data of the symmetric V1 -> V1 block matrix (the _data of a StencilMatrix, 6D, any + * strides); written with atomic additions. + * @param mat12 (1,2)-block data; written atomically. + * @param mat13 (1,3)-block data; written atomically. + * @param mat22 (2,2)-block data; written atomically. + * @param mat23 (2,3)-block data; written atomically. + * @param mat33 (3,3)-block data; written atomically. + * @param vec1 First component of the V1 vector data; written atomically. + * @param vec2 Second component of the V1 vector data; written atomically. + * @param vec3 Third component of the V1 vector data; written atomically. + * @param f0_values Value of f0 for each marker row (n_markers entries). + * + * Holes (markers[ip, 0] == -1) and boundary particles (markers[ip, -1] == -2) are skipped. The pyccel loop runs over + * all rows (shape(markers)[0]), which is n_markers here. Many threads add to the same entries, so the summation + * order (and the last bits of the result) differs from the serial pyccel loop. + */ +extern "C" __global__ void linear_vlasov_ampere(MarkerArgs args_markers, DerhamArgs args_derham, DomainArgs args_domain, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + Array3D vec1, Array3D vec2, Array3D vec3, + const double* f0_values) { + // CUDA-only: the marker row of this thread replaces the pyccel loop variable + int ip = blockDim.x * blockIdx.x + threadIdx.x; + + // get number of markers + int n_markers = args_markers.n_markers; + if (ip >= n_markers) return; + + // total number of markers (weights are w_p = delta f_p / (N * s_0)) + int n_markers_tot = args_markers.Np; + + // only do something if particle is a "true" particle (i.e. not a hole) + if (args_markers.markers(ip, 0) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) + return; + + // allocate for metric coeffs (3 x 3 matrices as nine row-major entries) + double dfm[9], df_inv[9]; + + // allocate for filling + double v[3], df_inv_v[3], filling_m[9], filling_v[3]; + + // marker positions + double eta1 = args_markers.markers(ip, 0); + double eta2 = args_markers.markers(ip, 1); + double eta3 = args_markers.markers(ip, 2); + + // get velocity + v[0] = args_markers.markers(ip, 3); + v[1] = args_markers.markers(ip, 4); + v[2] = args_markers.markers(ip, 5); + + // evaluate Jacobian, result in dfm + struphy_cuda::df(eta1, eta2, eta3, args_domain, dfm); + + // invert Jacobian matrix + struphy_cuda::matrix_inv(dfm, df_inv); + + // compute DF^{-1} v + struphy_cuda::matrix_vector(df_inv, v, df_inv_v); + + // filling_m = alpha^2 * kappa^2 * f0 / (N * s_0 * v_th^2) * (DF^{-1} v_p)_mu * (DF^{-1} v_p)_nu + struphy_cuda::outer(df_inv_v, df_inv_v, filling_m); + double factor = f0_values[ip] / (args_markers.markers(ip, 7) * n_markers_tot); + for (int k = 0; k < 9; ++k) filling_m[k] *= factor; + + // filling_v = alpha^2 * kappa * w_p * DL^{-1} * v_p + for (int k = 0; k < 3; ++k) filling_v[k] = args_markers.markers(ip, 6) * df_inv_v[k]; + + // call the appropriate matvec filler + struphy_cuda::m_v_fill_b_v1_symm(args_derham, eta1, eta2, eta3, mat11, mat12, mat13, mat22, mat23, mat33, + filling_m[0], filling_m[1], filling_m[2], filling_m[4], filling_m[5], + filling_m[8], vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2]); +} diff --git a/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh b/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh index 7cabaee45..ade422b97 100644 --- a/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh +++ b/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh @@ -28,4 +28,73 @@ __device__ inline void vec_fill_b_v0(const DerhamArgs& args_derham, double eta1, fill_vec(pn1, pn2, pn3, scratch.bn1, scratch.bn2, scratch.bn3, span1, span2, span3, args_derham.starts, vec, fill); } + +/** + * Add one marker's contribution to the symmetric blocks (1,1), (1,2), (1,3), (2,2), (2,3), (3,3) of a V1 -> V1 + * block matrix and to a V1 vector, as in particle_to_mat_kernels.m_v_fill_b_v1_symm. + * + * @param args_derham Spline degrees (1 to 8), knots and start indices; pn is also the padding of the blocks. + * @param eta1 Logical position along the first axis. + * @param eta2 Logical position along the second axis. + * @param eta3 Logical position along the third axis. + * @param mat11 (mu=1, nu=1)-block data (D N N x D N N), shape (n1, n2, n3, 2 * pn + 1); written atomically. + * @param mat12 (mu=1, nu=2)-block data (D N N x N D N); written atomically. + * @param mat13 (mu=1, nu=3)-block data (D N N x N N D); written atomically. + * @param mat22 (mu=2, nu=2)-block data (N D N x N D N); written atomically. + * @param mat23 (mu=2, nu=3)-block data (N D N x N N D); written atomically. + * @param mat33 (mu=3, nu=3)-block data (N N D x N N D); written atomically. + * @param fill11 Number multiplied by the basis functions and added to mat11. + * @param fill12 Number multiplied by the basis functions and added to mat12. + * @param fill13 Number multiplied by the basis functions and added to mat13. + * @param fill22 Number multiplied by the basis functions and added to mat22. + * @param fill23 Number multiplied by the basis functions and added to mat23. + * @param fill33 Number multiplied by the basis functions and added to mat33. + * @param vec1 mu=1 component of the V1 vector data (D N N); written atomically. + * @param vec2 mu=2 component of the V1 vector data (N D N); written atomically. + * @param vec3 mu=3 component of the V1 vector data (N N D); written atomically. + * @param fill1 Number multiplied by the basis functions and added to vec1. + * @param fill2 Number multiplied by the basis functions and added to vec2. + * @param fill3 Number multiplied by the basis functions and added to vec3. + * + * Pyccel keeps the spline values in args_derham.bn1, ..., bd3; here they are in a thread-local SplineScratch. + */ +__device__ inline void m_v_fill_b_v1_symm(const DerhamArgs& args_derham, double eta1, double eta2, double eta3, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + double fill11, double fill12, double fill13, double fill22, + double fill23, double fill33, Array3D vec1, Array3D vec2, + Array3D vec3, double fill1, double fill2, double fill3) { + // degrees of the basis functions : B-splines (pn) and D-splines (pd) + int pn1 = args_derham.pn[0]; + int pn2 = args_derham.pn[1]; + int pn3 = args_derham.pn[2]; + + int pd1 = pn1 - 1; + int pd2 = pn2 - 1; + int pd3 = pn3 - 1; + + // CUDA-only scratch holds the spans and the spline values pyccel keeps in args_derham + SplineScratch scratch; + get_spans(eta1, eta2, eta3, args_derham, scratch); + int span1 = scratch.span1, span2 = scratch.span2, span3 = scratch.span3; + + // fill matrix entries + fill_mat_vec(pd1, pn2, pn3, pd1, pn2, pn3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bd1, scratch.bn2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat11, fill11, vec1, fill1); + + fill_mat_vec(pn1, pd2, pn3, pn1, pd2, pn3, scratch.bn1, scratch.bd2, scratch.bn3, scratch.bn1, scratch.bd2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat22, fill22, vec2, fill2); + + fill_mat_vec(pn1, pn2, pd3, pn1, pn2, pd3, scratch.bn1, scratch.bn2, scratch.bd3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat33, fill33, vec3, fill3); + + fill_mat(pd1, pn2, pn3, pn1, pd2, pn3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bn1, scratch.bd2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat12, fill12); + + fill_mat(pd1, pn2, pn3, pn1, pn2, pd3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat13, fill13); + + fill_mat(pn1, pd2, pn3, pn1, pn2, pd3, scratch.bn1, scratch.bd2, scratch.bn3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat23, fill23); +} } diff --git a/src/struphy/pic/tests/cuda_emulation.py b/src/struphy/pic/tests/cuda_emulation.py index 0e8db4f12..ca3c76d01 100644 --- a/src/struphy/pic/tests/cuda_emulation.py +++ b/src/struphy/pic/tests/cuda_emulation.py @@ -6,8 +6,10 @@ the fields of each struct as separate parameters, rebuilds the structs and calls the real kernel. """ +import contextlib + import numpy as np -from cunumpy.kernel_testing import emulate_cuda_kernel +from cunumpy.kernel_testing import emulate_cuda_kernel, fake_cupy_active from cunumpy.kernels import CudaKernel from struphy.utils.cuda_arguments import CUDA_OPTIONS @@ -55,11 +57,12 @@ def _struct_field(value, name): return getattr(value, name) -def emulate_struct_kernel(kernel: CudaKernel, *args, n_threads=None): +def emulate_struct_kernel(kernel: CudaKernel, *args, n_threads=None, grid=None, block=None): """Emulate `kernel` with host arguments; struct arguments are the pyccel argument objects (NumPy arrays). - Arrays (also those inside the argument objects) are updated in place, as by a launch. Without `n_threads`, - one thread per row of the first array is launched, as in a real launch. + Arrays (also those inside the argument objects) are updated in place, as by a launch. Without `n_threads` or + `grid`, one thread per row of the first array is launched, as in a real launch; `grid` and `block` give the + launch shape explicitly. """ flat = [] for p, value in zip(kernel.signature, args): @@ -67,9 +70,58 @@ def emulate_struct_kernel(kernel: CudaKernel, *args, n_threads=None): flat.append(value) else: flat += [_struct_field(value, f.name) for f in p.struct.fields] - if n_threads is None: - n_threads = next(value.shape[0] for value in flat if isinstance(value, np.ndarray)) - grid, block = kernel.launch_shape(n_threads) + if grid is None: + if n_threads is None: + n_threads = next(value.shape[0] for value in flat if isinstance(value, np.ndarray)) + grid, block = kernel.launch_shape(n_threads, block=block) source, name = _wrapper(kernel) wrapper = CudaKernel(source, name, source_dir=kernel.source_dir, **CUDA_OPTIONS) emulate_cuda_kernel(wrapper, *flat, grid=grid, block=block, options=EMULATION_OPTIONS) + + +def _host(value): + """The host buffer behind a fake CuPy array (the same memory, no copy); other values unchanged.""" + if isinstance(value, np.ndarray) or not hasattr(value, "__cuda_array_interface__"): + return value + return object.__getattribute__(value, "_a") + + +class _HostFields: + """A CUDA argument object (``Cuda*Arguments``) whose array attributes are their fake CuPy host buffers.""" + + def __init__(self, args): + self._args = args + + def __getattr__(self, name): + return _host(getattr(self._args, name)) + + +@contextlib.contextmanager +def emulated_launches(): + """On cunumpy's fake CuPy, run every ``CudaKernel`` launch by CPU emulation, in place on the fake device arrays. + + The fake CuPy (``CUNUMPY_FAKE_CUPY=1``) keeps device arrays in host memory and rejects host/device mixing, but + cannot launch kernels. Inside this context a launch emulates the kernel (struct arguments included, through + :func:`emulate_struct_kernel`) on the host buffers of the fake arrays, so code that launches CUDA kernels (struphy's + and feectools') runs end to end on the CuPy backend without a GPU. Each launch compiles the kernel, so this is slow. + """ + if not fake_cupy_active(): + raise RuntimeError("emulated_launches() needs cunumpy's fake CuPy (CUNUMPY_FAKE_CUPY=1)") + original = CudaKernel.__call__ + + def launch(self, *args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None): + grid, block = self.launch_shape(n_threads, grid=grid, block=block, args=args) + if any(p.struct is not None for p in self.signature): + host_args = [_host(a) if p.struct is None else _HostFields(a) for p, a in zip(self.signature, args)] + emulate_struct_kernel(self, *host_args, grid=grid, block=block) + else: + host_args = [_host(a) for a in args] + emulate_cuda_kernel( + self, *host_args, grid=grid, block=block, shared_mem=shared_mem, options=EMULATION_OPTIONS + ) + + CudaKernel.__call__ = launch + try: + yield + finally: + CudaKernel.__call__ = original diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index e2ce0a5d5..85380b6e2 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -29,6 +29,7 @@ marker_arguments, spline_coefficients, spline_evaluation_arguments, + v1_symm_accumulation_data, ) @@ -79,6 +80,19 @@ def charge_density_0form_args(bc): return (args_markers, derham_arguments(), args_domain, xp.zeros((18, 20, 16))) +def linear_vlasov_ampere_args(domain): + """129 markers (a hole and a boundary particle) in domain `domain` of geometry_domain(); zeroed V1 data.""" + args_markers, _ = marker_arguments((0, 0, 0)) + f0_values = xp.asarray(np.random.default_rng(14).random(N_MARKERS)) + return ( + args_markers, + derham_arguments(), + geometry_domain(domain).args_domain, + *v1_symm_accumulation_data(), + f0_values, + ) + + # ---------------------------------------------------------------- spline evaluation @@ -267,6 +281,9 @@ def size_of(index, per_thread=1): "reflect": ParityCases((0, 1, 2), reflect_args, n_threads=size_of(2)), # atomics add in another order than the serial loop "charge_density_0form": ParityCases(((0, 0, 0), (2, 0, 1)), charge_density_0form_args, rtol=1e-12, atol=1e-13), + # Cuboid, Colella and HollowTorus (non-diagonal DF) and a 3d spline mapping, by index into geometry_domain(); + # atomics add in another order than the serial loop, and the entries reach 1e4 (atol for cancelling sums) + "linear_vlasov_ampere": ParityCases((0, 2, 5, 12), linear_vlasov_ampere_args, rtol=1e-12, atol=1e-10), "eval_spline_mpi_markers": ParityCases(SPLINE_KINDS, eval_spline_mpi_markers_args), "eval_spline_mpi_matrix": ParityCases(SPLINE_KINDS, eval_spline_mpi_grid_args(sparse=False), n_threads=size_of(-1)), "eval_spline_mpi_sparse_meshgrid": ParityCases( diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index d413b5380..eaed48adf 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -49,14 +49,37 @@ def butcher_arguments(method): return xp.asarray(butcher.a_stage), xp.asarray(butcher.b), xp.asarray(butcher.c), butcher.n_stages +DERHAM_DEGREE = (2, 3, 1) +DERHAM_CELLS = 8 + + def derham_arguments(): """Splines of degrees 2, 3, 1 on 8 cells, starting at index 0.""" - degree = np.array([2, 3, 1], dtype=np.int64) - knots = [np.r_[np.zeros(p), np.linspace(0, 1, 9), np.ones(p)] for p in degree] + degree = np.array(DERHAM_DEGREE, dtype=np.int64) + knots = [np.r_[np.zeros(p), np.linspace(0, 1, DERHAM_CELLS + 1), np.ones(p)] for p in degree] args_class = CudaDerhamArguments if xp.get_backend() == "cupy" else DerhamArguments return args_class(xp.asarray(degree), *(xp.asarray(t) for t in knots), xp.zeros(3, dtype=np.int64)) +def v1_symm_accumulation_data(): + """Zeroed data of a symmetric V1 -> V1 accumulation matrix and a V1 vector, for :func:`derham_arguments`. + + Returns the six independent blocks ``mat11, mat12, mat13, mat22, mat23, mat33`` and the three vector components + ``vec1, vec2, vec3`` on the active backend, shaped like the ``_data`` of the ``StencilMatrix`` blocks and + ``StencilVector`` components that ``Accumulator(..., "Hcurl", ..., symmetry="symm", add_vector=True)`` passes on + one process for clamped splines of degrees 2, 3, 1 on 8 cells: component ``mu`` has D-splines (``n - 1`` + functions) along axis ``mu`` and N-splines (``n = 8 + p``) along the others, every axis is padded by ``p`` on both + sides, and the matrix blocks have ``2 p + 1`` diagonals per axis (the paddings are the degrees of V0). + """ + degree = np.array(DERHAM_DEGREE) + n_spline = DERHAM_CELLS + degree + rows = [tuple(n_spline - (np.arange(3) == mu) + 2 * degree) for mu in range(3)] + diagonals = tuple(2 * degree + 1) + mats = tuple(xp.zeros(rows[mu] + diagonals) for mu, nu in ((0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2))) + vecs = tuple(xp.zeros(rows[mu]) for mu in range(3)) + return mats + vecs + + def spline_coefficients(n=3, seed=11): """Random coefficients, large enough for every span of :func:`derham_arguments`.""" rng = np.random.default_rng(seed) diff --git a/src/struphy/pic/tests/test_accum_matrix_cupy.py b/src/struphy/pic/tests/test_accum_matrix_cupy.py new file mode 100644 index 000000000..cf0aa62c6 --- /dev/null +++ b/src/struphy/pic/tests/test_accum_matrix_cupy.py @@ -0,0 +1,209 @@ +"""The matrix path of :class:`~struphy.pic.accumulation.particles_to_grid.Accumulator` on the CuPy backend. + +``linear_vlasov_ampere`` accumulates the symmetric V1 -> V1 block matrix and the V1 vector of ``EfieldWeightsCoupling``; +the matrix data (``StencilMatrix._data``) must be device arrays that are zeroed, passed to the CUDA kernel and turned +into the Schur operator without host copies. Checked against the NumPy backend on a GPU and, without one, on +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 + +import cunumpy as xp +import numpy as np +import pytest +from cunumpy.kernel_testing import emulation_compiler, requires_cupy + +from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env + +ALPHA, KAPPA, VTH, DT = 1.3, 0.7, 0.9, 0.05 +N_SCHUR_ITERATIONS = 3 + + +@dataclass +class MarkerOwner: + """What the accumulators read from ``Particles``: the marker arguments, no clones. + + The same markers on both backends (``Particles`` draws them with the backend's random generator). + """ + + args_markers: Any + clone_config: Any = None + + +def accumulate_linear_vlasov_ampere(guard=None): + """Accumulate ``linear_vlasov_ampere`` like ``EfieldWeightsCoupling`` on the active backend, then solve with it. + + Colella mapping (non-diagonal Jacobian), clamped splines in eta1 and periodic ones in eta2, eta3, so the three + V1 components have different sizes. `guard` (a context manager factory, e.g. one that fails on host transfers) + encloses the accumulation, the scaled operator and its application, as in a time step of the propagator; the + Schur solve is not guarded (see below). + + Returns a dict of host arrays: the nine kernel data arrays (``mat11`` ... ``vec3``) after the accumulation (ghost + regions exchanged), the assembled block matrix as dense blocks, the accumulated vector after the boundary + operators, ``BC.dot(x)`` with the operator scaled as in ``EfieldWeightsCoupling`` and the Schur solution after + ``N_SCHUR_ITERATIONS`` iterations of the conjugate gradient method. + """ + import contextlib + + from feectools.linalg.stencil import StencilMatrix + + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.psydac_derham import Derham + from struphy.geometry.domains import Colella + from struphy.io.options import DerhamOptions + from struphy.linear_algebra.schur_solver import SchurSolver + from struphy.linear_algebra.solver import SolverParameters + from struphy.pic.accumulation.kernels.linear_vlasov_ampere import linear_vlasov_ampere + from struphy.pic.accumulation.particles_to_grid import Accumulator + from struphy.pic.tests.kernel_test_args import N_MARKERS, marker_arguments + from struphy.topology.grids import TensorProductGrid + + guard = guard or contextlib.nullcontext + derham = Derham( + TensorProductGrid(num_elements=(6, 5, 4)), + DerhamOptions(degree=(2, 3, 1), bcs=(("free", "free"), None, None)), + comm=None, + ) + domain = Colella(Lx=2.0, Ly=3.0, alpha=0.05, Lz=4.0) + mass_ops = WeightedMassOperators(derham, domain) + args_markers, _ = marker_arguments((0, 0, 0)) + accum = Accumulator( + MarkerOwner(args_markers), + "Hcurl", + linear_vlasov_ampere, + mass_ops, + domain.args_domain, + add_vector=True, + symmetry="symm", + ) + matrix = accum.operators[0].matrix + + # the kernel writes into the data of the operator's blocks (and vector), not into copies + blocks = [matrix[a, b] for a in range(3) for b in range(a, 3)] + assert all(isinstance(block, StencilMatrix) for block in blocks) + assert all(data is block._data for data, block in zip(accum._args_data[:6], blocks)) + assert all(data is vec._data for data, vec in zip(accum._args_data[6:], accum._vectors[0].blocks)) + assert all(xp.is_gpu(data) == (xp.get_backend() == "cupy") for data in accum._args_data) + + rng = np.random.default_rng(16) + f0_values = xp.asarray(rng.random(N_MARKERS)) + for data in accum._args_data: + data[...] = 1.0 # left over from a previous step: the accumulator zeroes it + + x = accum._vectors[0].space.zeros() + for block in x.blocks: + block._data[...] = xp.asarray(rng.normal(size=block._data.shape)) + x.update_ghost_regions() + + schur_solver = SchurSolver( + mass_ops.M1, + ALPHA**2 * KAPPA**2 * accum.operators[0] / (4 * VTH**2), + "pcg", + # a few iterations (each one launches nine stencil dot kernels, which the emulation compiles one by one) + solver_params=SolverParameters(tol=1e-30, maxiter=N_SCHUR_ITERATIONS), + ) + + with guard(): + accum(f0_values) + # as in EfieldWeightsCoupling.__call__ + schur_solver.BC = accum.operators[0] + schur_solver.BC *= (-1) * ALPHA**2 * KAPPA**2 / (4 * VTH**2) + bc_x = schur_solver.BC.dot(x) + byn = accum.vectors[0].copy() + byn *= ALPHA**2 * KAPPA / 2.0 + # not guarded: feectools' StencilVector.inner returns its result on the host (xp.to_numpy), so every inner + # product of the conjugate gradient method copies one scalar from the device + solution, info = schur_solver(xn=x, Byn=byn, dt=DT) + + out = {name: xp.to_numpy(data) for name, data in zip(MATRIX_DATA_NAMES, accum._args_data)} + for a in range(3): + for b in range(3): + out[f"block{a}{b}"] = xp.to_numpy(matrix[a, b].toarray()) + out["vector"] = xp.to_numpy(accum.vectors[0].toarray()) + out["bc_x"] = xp.to_numpy(bc_x.toarray()) + out["solution"] = xp.to_numpy(solution.toarray()) + out["niter"] = np.array(info["niter"]) + return out + + +MATRIX_DATA_NAMES = ("mat11", "mat12", "mat13", "mat22", "mat23", "mat33", "vec1", "vec2", "vec3") + + +def compare_with_numpy(result): + """Compare a CuPy-backend `result` of :func:`accumulate_linear_vlasov_ampere` with the NumPy backend's.""" + with xp.use_backend("numpy"): + expected = accumulate_linear_vlasov_ampere() + assert expected["niter"] == result["niter"] == N_SCHUR_ITERATIONS + assert np.any(expected["mat12"] != 0.0) and np.any(expected["vec3"] != 0.0) + # the transposed blocks of the symmetric matrix are filled from the accumulated ones + np.testing.assert_allclose(expected["block10"], expected["block01"].T, rtol=1e-14, atol=1e-14) + for name, value in expected.items(): + scale = max(np.max(np.abs(value)), 1.0) + # atomics add in another order than the serial loop; the solve repeats it + np.testing.assert_allclose(result[name], value, rtol=1e-10, atol=1e-12 * scale, err_msg=name) + + +@requires_cupy +def test_accumulator_matrix_on_cupy(): + """On a GPU: the accumulated matrix and vector, the Schur operator and solve agree with NumPy; the accumulation and + the operator run without host/device transfers.""" + from cunumpy.profiling import assert_no_transfers + + with xp.use_backend("cupy"): + accumulate_linear_vlasov_ampere() # warm up (NVRTC compilation) + result = accumulate_linear_vlasov_ampere(guard=assert_no_transfers) + compare_with_numpy(result) + + +def fake_cupy_transfer_guard(): + """On cunumpy's fake CuPy: fail on every explicit host copy of a device array (``.get()``, ``cupy.asnumpy``). + + The fake CuPy already rejects implicit conversions (``numpy.asarray`` of a device array, mixing host and device + arrays); this closes the explicit ones, which are legal but are transfers inside a time step. + """ + import contextlib + + import cupy + + @contextlib.contextmanager + def guard(): + def fail(*args, **kwargs): + raise AssertionError("host copy of a device array during the accumulation") + + saved = cupy.ndarray.get, cupy.asnumpy + cupy.ndarray.get, cupy.asnumpy = fail, fail + try: + yield + finally: + cupy.ndarray.get, cupy.asnumpy = saved + + return guard + + +def check_accumulator_matrix_fake_cupy(): + """Run in a child process with ``CUNUMPY_FAKE_CUPY=1``: the CuPy backend with emulated kernel launches.""" + from struphy.pic.tests.cuda_emulation import emulated_launches + + with emulated_launches(), xp.use_backend("cupy"): + result = accumulate_linear_vlasov_ampere(guard=fake_cupy_transfer_guard()) + compare_with_numpy(result) + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +@pytest.mark.skipif(emulation_compiler() is None, reason="no C++ compiler") +def test_accumulator_matrix_fake_cupy(): + """Without a GPU: the matrix path of the Accumulator on the CuPy backend (fake CuPy, emulated launches). + + Device data are zeroed and passed to the CUDA kernel, the ghost regions, the transposed blocks, the scaled + operator and the Schur solve run on device arrays, and no array is copied to the host from the accumulation to + the operator's application. Runs in a + 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:] diff --git a/src/struphy/pic/tests/test_cuda_emulation.py b/src/struphy/pic/tests/test_cuda_emulation.py index 793bb9525..f2ab111a2 100644 --- a/src/struphy/pic/tests/test_cuda_emulation.py +++ b/src/struphy/pic/tests/test_cuda_emulation.py @@ -113,3 +113,15 @@ def test_emulated_transform_helpers(domain_index): out = np.zeros(len(expected)) emulate_struct_kernel(kernel, *inputs, args_domain, out, len(out), n_threads=len(out)) np.testing.assert_allclose(out, expected, rtol=1e-10, atol=1e-10) + + +@requires_compiler +def test_emulated_v1_symm_filler(): + """test_device_helpers.test_v1_symm_filler by CPU emulation: the device m_v_fill_b_v1_symm agrees with pyccel.""" + from struphy.pic.tests.test_device_helpers import v1_symm_filler_case + + kernel, (eta, fills), args_derham, data, expected = v1_symm_filler_case() + emulate_struct_kernel(kernel, eta, fills, args_derham, *data, len(eta), n_threads=len(eta)) + for result, reference in zip(data, expected): + assert np.any(reference != 0.0) + np.testing.assert_allclose(result, reference, rtol=1e-12, atol=1e-13) diff --git a/src/struphy/pic/tests/test_cuda_parity.py b/src/struphy/pic/tests/test_cuda_parity.py index 75d53245f..0b0a14c9e 100644 --- a/src/struphy/pic/tests/test_cuda_parity.py +++ b/src/struphy/pic/tests/test_cuda_parity.py @@ -10,6 +10,7 @@ from cunumpy.profiling import assert_no_transfers from struphy.pic.accumulation.kernels.charge_density_0form import charge_density_0form +from struphy.pic.accumulation.kernels.linear_vlasov_ampere import linear_vlasov_ampere from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage from struphy.pic.pushing.kernels.push_v_with_efield import push_v_with_efield from struphy.pic.pushing.kernels.push_vxb_analytic import push_vxb_analytic @@ -108,7 +109,7 @@ def device_only(value, *args, **kwargs): def test_vlasov_kernel_coverage(): - """Vlasov (PushEta, both PushVxB algorithms) and the ported Vlasov-Ampere kernels have device kernels.""" + """Vlasov (PushEta, both PushVxB algorithms) and the ported (linear) Vlasov-Ampere kernels have device kernels.""" for kernel in ( push_eta_stage, push_vxb_analytic, @@ -116,5 +117,6 @@ def test_vlasov_kernel_coverage(): push_v_with_efield, push_weights_with_efield_lin_va, charge_density_0form, + linear_vlasov_ampere, ): assert kernel.has_cuda, kernel.name diff --git a/src/struphy/pic/tests/test_device_helpers.py b/src/struphy/pic/tests/test_device_helpers.py index a2390ef4d..b771d8efc 100644 --- a/src/struphy/pic/tests/test_device_helpers.py +++ b/src/struphy/pic/tests/test_device_helpers.py @@ -391,3 +391,64 @@ def test_transform_helpers(domain_index): kernel, inputs, _, expected = transform_helper_case(domain_index) result = _run_on_device(kernel, inputs, domain_index, len(expected)) np.testing.assert_allclose(result, expected, rtol=1e-10, atol=1e-10) + + +# CUDA-only wrapper: one thread per point calls m_v_fill_b_v1_symm with the point's nine fillings +V1_SYMM_FILLER_SOURCE = r""" +#include "struphy/pic/accumulation/particle_to_mat_kernels.cuh" +extern "C" __global__ void fill_v1_symm(const double* eta, const double* fills, DerhamArgs args_derham, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + Array3D vec1, Array3D vec2, Array3D vec3, int n) { + int ip = blockDim.x * blockIdx.x + threadIdx.x; + if (ip >= n) return; + const double* fill = fills + 9 * ip; + struphy_cuda::m_v_fill_b_v1_symm(args_derham, eta[3 * ip], eta[3 * ip + 1], eta[3 * ip + 2], mat11, mat12, mat13, + mat22, mat23, mat33, fill[0], fill[1], fill[2], fill[3], fill[4], fill[5], vec1, + vec2, vec3, fill[6], fill[7], fill[8]); +} +""" + + +def v1_symm_filler_case(): + """Wrapper kernel, host inputs and pyccel reference for the device m_v_fill_b_v1_symm. + + 65 points (the corners of the unit cube among them) with random fillings, accumulated into zeroed V1 data shaped + like the real accumulation matrices (:func:`~struphy.pic.tests.kernel_test_args.v1_symm_accumulation_data`). + Used by :func:`test_v1_symm_filler` (GPU) and by its CPU emulation in ``test_cuda_emulation.py``. Returns + ``kernel, (eta, fills), args_derham, data, expected``: ``args_derham`` is the pyccel class and ``data`` the nine + zeroed host arrays in kernel order. + """ + from struphy.pic.accumulation import particle_to_mat_kernels + from struphy.pic.tests.kernel_test_args import derham_arguments, v1_symm_accumulation_data + + rng = np.random.default_rng(15) + eta = np.r_[np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0], [0.0, 1.0, 0.5]]), rng.random((62, 3))] + fills = rng.normal(size=(len(eta), 9)) + with cunumpy.use_backend("numpy"): + args_derham = derham_arguments() + data = v1_symm_accumulation_data() + expected = tuple(a.copy() for a in data) + for point, fill in zip(eta, fills): + particle_to_mat_kernels.m_v_fill_b_v1_symm( + args_derham, *point, *expected[:6], *fill[:6], *expected[6:], *fill[6:] + ) + kernel = CudaKernel(V1_SYMM_FILLER_SOURCE, "fill_v1_symm", **CUDA_OPTIONS) + return kernel, (eta, fills), args_derham, data, expected + + +@requires_cupy +def test_v1_symm_filler(): + """m_v_fill_b_v1_symm (and the fill_mat/fill_mat_vec it calls) on the device agrees with pyccel.""" + import cupy as cp + + from struphy.pic.tests.kernel_test_args import derham_arguments + + kernel, (eta, fills), _, data, expected = v1_symm_filler_case() + with cunumpy.use_backend("cupy"): + args_derham = derham_arguments() + device_data = [cp.asarray(a) for a in data] + kernel(cp.asarray(eta), cp.asarray(fills), args_derham, *device_data, len(eta), n_threads=len(eta)) + for result, reference in zip(device_data, expected): + assert np.any(reference != 0.0) + np.testing.assert_allclose(result.get(), reference, rtol=1e-12, atol=1e-13) From 6dd36dbac4b9e172a2be5a29e25b5666531b2ddf Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 23:35:08 +0200 Subject: [PATCH 29/50] Mark every column of the test hole so accumulation kernels skip it The hole in marker_arguments() was only marked in column 8 (first_init_idx), which the pushers test. Accumulation kernels test column 0, so the linear_vlasov_ampere parity case had no real hole. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/kernel_test_args.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index eaed48adf..24e597e89 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -29,7 +29,9 @@ def marker_arguments(bc): markers[:, 3:6] = rng.uniform(-2, 2, (N_MARKERS, 3)) markers[:, 8:14] = markers[:, :6] markers[:, 18:21] = 0.0 - markers[0, 8] = -1.0 + # Row 0 is a hole. Struphy marks every column of a hole with -1: pushers test column first_init_idx (8), + # accumulation kernels test column 0. + markers[0, :] = -1.0 markers[1, -1] = -2.0 valid = np.ones(N_MARKERS, dtype=bool) valid[:2] = False From 52bc2b93c4af624cbb27d14c5fc9886bbc2fe41e Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 15:11:12 +0200 Subject: [PATCH 30/50] Rename first_init_idx to first_pusher_idx in the code new on this branch Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/kernel_test_args.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index 24e597e89..abe1ef445 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -29,7 +29,7 @@ def marker_arguments(bc): markers[:, 3:6] = rng.uniform(-2, 2, (N_MARKERS, 3)) markers[:, 8:14] = markers[:, :6] markers[:, 18:21] = 0.0 - # Row 0 is a hole. Struphy marks every column of a hole with -1: pushers test column first_init_idx (8), + # Row 0 is a hole. Struphy marks every column of a hole with -1: pushers test column first_pusher_idx (8), # accumulation kernels test column 0. markers[0, :] = -1.0 markers[1, -1] = -2.0 From 1cad34ae0fe7bdc983ff3f24435fe68586433b5d Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 10 Oct 2026 11:34:09 +0200 Subject: [PATCH 31/50] Evaluate the spline-based equilibria on the device; batched field-line tracing Implements items 1-4 and 6 of the plan in #714: - AdhocTorus (q_kind 1, 2), AdhocTorusQPsi, EQDSKequilibrium: the splines are fitted on the active backend with xp.scipy.interpolate and evaluated there, without host round trips. Only the scalar setup steps (quad, odeint, fsolve, the EQDSK file, minimize on a host copy of the flux spline) run on the host. RectBivariateSpline, which cupyx.scipy lacks, is replaced by make_interp_spline along R and Z as an NdBSpline, with the points clamped to the domain of the spline. - field_line_tracing: one batched xp.optimize.newton over all rays per flux surface instead of one scalar Newton per point; Tokamak runs its setup on the active backend. - test_equils_cupy: the evaluation of all equilibria except GVEC/DESC runs under xp.profiling.assert_no_transfers(). - cunumpy pinned to fb0f7c0a (#93: fake cupyx.scipy, xp.optimize.newton). Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 36 ++- pyproject.toml | 2 +- src/struphy/fields_background/equils.py | 234 +++++++-------- .../tests/test_equils_cupy.py | 18 +- .../geometry/domains/tokamak/tokamak.py | 75 +++-- src/struphy/geometry/utilities.py | 280 +++++++++--------- 6 files changed, 324 insertions(+), 321 deletions(-) diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 462d18830..64ffb0759 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -373,26 +373,40 @@ views with more than 4 dimensions (all matrix accumulations write 6D stencil mat 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. +- **GVEC/DESC fields on the device.** `GVECequilibrium` and `DESCequilibrium` still evaluate B, J, p and n through + gvec/DESC on the host, one copy per call (#696). Plan: evaluate them once at setup on a grid, interpolate onto 3d + tensor-product splines (`make_interp_spline` per axis + `NdBSpline`) and evaluate those on the device, with native + evaluation kept as an option (`evaluate="spline" | "native"`), since this trades exactness for speed. Separate PR. ## 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. +- **Splines on the device.** The splines of `AdhocTorus` (`q_kind` 1, 2), `AdhocTorusQPsi` and `EQDSKequilibrium` + are fitted on the active backend with `xp.scipy.interpolate` (SciPy or `cupyx.scipy`, looked up at the point of use) + and evaluated there, without host round trips (also at run time: `pic/base.py` evaluates `b_cart` and + `pic/particles.py` `psi_r` at markers). Only the scalar setup steps run on the host, under + `xp.use_backend("numpy")`: the `quad` loops, `odeint`/`fsolve`, the EQDSK file and the `minimize` for the magnetic + axis (on a host copy of the flux spline, so the optimizer steps don't evaluate on the device). + `RectBivariateSpline` is not in `cupyx.scipy`; with `s=0` it is the same spline as `make_interp_spline` along R, then + along Z, as an `NdBSpline` (same knots, values and derivatives equal to round-off). `NdBSpline` extrapolates, so the + points are clamped to the domain of the spline first, as `RectBivariateSpline` does. +- **Field-line tracing (`Tokamak`).** One batched Newton (secant) solve over all rays per flux surface + (`xp.optimize.newton`, cunumpy#93) instead of one scalar `scipy.optimize.newton` per point; `psi` is evaluated on the + active backend, so `Tokamak` runs its setup there. The small spline interpolation (`splu`, `kron_lusolve_2d`) and the + other xi-parametrizations (a SciPy `root` per flux surface) run on the host. The control points agree with the + scalar tracing to ~1e-12, and the tracing is 6-15x faster on the CPU too. +- **GVEC/DESC** `bv`, `jv`, `p0`, `n0`, `gradB1` go 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 (see the open question + [GVEC/DESC fields on the device](#open-questions)). - **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. + host buffers (`host_geometry_kernels`), since the fake CuPy cannot launch CUDA kernels. The fake `cupyx.scipy` + (cunumpy#93) has only the `interpolate` names CuPy has, and the evaluation of every equilibrium except GVEC/DESC runs + under `xp.profiling.assert_no_transfers()`. Still to check on a GPU: that CuPy's `UnivariateSpline` and `NdBSpline` + agree with SciPy to the test's 1e-12 (the fake runs SciPy itself). - **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 diff --git a/pyproject.toml b/pyproject.toml index fdeda2015..3e637a235 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ classifiers = [ ] dependencies = [ "numpy<=2.5.0", - "cunumpy @ git+https://github.com/max-models/cunumpy.git@26a640b1d47d5677663785e49b1c44d343540df1", + "cunumpy @ git+https://github.com/max-models/cunumpy.git@fb0f7c0afaac7994dd62147c1fa39e1d9618318a", "pyccel>=2.2.0, <=2.2.3", "maybempi>=0.1.2", "feectools>=0.8.0, <=0.8.0", diff --git a/src/struphy/fields_background/equils.py b/src/struphy/fields_background/equils.py index 2377cebdc..bac148c33 100644 --- a/src/struphy/fields_background/equils.py +++ b/src/struphy/fields_background/equils.py @@ -39,6 +39,53 @@ logger = logging.getLogger("struphy") +# Splines of the numerical profiles are fitted on the active backend (SciPy or cupyx.scipy, through xp.scipy, looked up +# at the point of use), so that they are evaluated on the device on the CuPy backend. Only the scalar setup steps +# (quad, odeint, fsolve, minimize, file reading) run on the host. + + +def _univariate_spline(x, y, k: int): + """Interpolating spline through the host data ``(x, y)`` on the active backend, constant outside ``[x[0], x[-1]]``. + + The same as ``scipy.interpolate.UnivariateSpline(x, y, k=k, s=0.0, ext=3)``. + """ + return xp.scipy.interpolate.UnivariateSpline(xp.to_cunumpy(x), xp.to_cunumpy(y), k=k, s=0.0, ext=3) + + +def _evaluate_1d(spline, x, nu: int = 0): + """Evaluate a 1d spline of :func:`_univariate_spline` at ``x``; a Python float for a scalar ``x``.""" + out = spline(xp.asarray(x), nu=nu) + # remove all "dimensions" for point-wise evaluation + if isinstance(x, (int, float)): + assert out.ndim == 0 + out = out.item() + return out + + +def _rect_spline(x, y, z, kx: int, ky: int): + """Interpolating tensor-product spline through the host data ``z[i, j]`` at ``(x[i], y[j])`` on the active backend. + + The same spline as ``scipy.interpolate.RectBivariateSpline(x, y, z, kx=kx, ky=ky, s=0.0)`` (same knots, values and + derivatives equal to round-off), which cupyx.scipy does not have. Unlike it, the returned ``NdBSpline`` extrapolates: + evaluate it with :func:`_evaluate_rect`, which clamps the points like ``RectBivariateSpline``. + """ + interpolate = xp.scipy.interpolate + bx = interpolate.make_interp_spline(xp.to_cunumpy(x), xp.to_cunumpy(z), k=kx) + by = interpolate.make_interp_spline(xp.to_cunumpy(y), bx.c.T, k=ky) + return interpolate.NdBSpline((bx.t, by.t), by.c.T, (kx, ky)) + + +def _evaluate_rect(spline, bounds, x, y, dx: int = 0, dy: int = 0): + """Evaluate a spline of :func:`_rect_spline` (or its derivative) at the points ``(x, y)``, clamped to ``bounds``. + + ``bounds`` is ``((x_min, x_max), (y_min, y_max))``, the domain of the spline. + """ + x, y = xp.broadcast_arrays(xp.asarray(x, dtype=float), xp.asarray(y, dtype=float)) + x = xp.clip(x, *bounds[0]) + y = xp.clip(y, *bounds[1]) + return spline(xp.stack([x, y], axis=-1), nu=(dx, dy)) + + class HomogenSlab(CartesianMHDequilibrium): r""" Homogeneous MHD equilibrium. @@ -874,7 +921,6 @@ 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, @@ -899,7 +945,6 @@ def __init__( self.params = copy.deepcopy(locals()) from scipy.integrate import quad - from scipy.interpolate import UnivariateSpline # plasma boundary contour ths = xp.linspace(0.0, 2 * xp.pi, 201) @@ -916,27 +961,10 @@ def __init__( self._p_i = None elif self.params["q_kind"] == 1 or self.params["q_kind"] == 2: - r_i = xp.linspace(0.0, self.params["a"], self.params["psi_nel"] + 1) def dpsi_dr(r): return self.params["B0"] * r / (self.q_r(r) * xp.sqrt(1 - r**2 / self.params["R0"] ** 2)) - psis = xp.zeros_like(r_i) - - for i, rr in enumerate(r_i): - psis[i] = quad(dpsi_dr, 0.0, rr)[0] - - self._psi_i = UnivariateSpline( - r_i, - psis, - k=self.params["psi_k"], - s=0.0, - ext=3, - ) - - self._psi0 = 0.0 - self._psi1 = self.psi(self.params["R0"] + self.params["a"], 0.0) - def dp_dr(r): return ( -(self.params["B0"] ** 2 * r) @@ -944,18 +972,17 @@ def dp_dr(r): * (2 * self.q_r(r) - r * self.q_r(r, der=1)) ) - ps = xp.zeros_like(r_i) + # integrate the profiles on the host (scalar quad) + with xp.use_backend("numpy"): + r_i = xp.linspace(0.0, self.params["a"], self.params["psi_nel"] + 1) + psis = xp.array([quad(dpsi_dr, 0.0, rr)[0] for rr in r_i]) + ps = xp.array([quad(dp_dr, 0.0, rr)[0] for rr in r_i]) - for i, rr in enumerate(r_i): - ps[i] = quad(dp_dr, 0.0, rr)[0] + self._psi_i = _univariate_spline(r_i, psis, self.params["psi_k"]) + self._p_i = _univariate_spline(r_i, ps - ps[-1], self.params["psi_k"]) - self._p_i = UnivariateSpline( - r_i, - ps - ps[-1], - k=self.params["psi_k"], - s=0.0, - ext=3, - ) + self._psi0 = 0.0 + self._psi1 = self.psi(self.params["R0"] + self.params["a"], 0.0) @property def boundary_pts_R(self): @@ -1022,12 +1049,7 @@ def psi_r(self, r, der=0): # alternative profile (interpolated) elif self.params["q_kind"] == 1 or self.params["q_kind"] == 2: - out = xp.host_call(self._psi_i, r, nu=der) - - # remove all "dimensions" for point-wise evaluation - if isinstance(r, (int, float)): - assert out.ndim == 0 - out = out.item() + out = _evaluate_1d(self._psi_i, r, nu=der) return out @@ -1151,12 +1173,7 @@ def p_r(self, r): # alternative profiles (interpolated) elif self.params["q_kind"] == 1 or self.params["q_kind"] == 2: - pout = xp.host_call(self._p_i, r) - - # remove all "dimensions" for point-wise evaluation - if isinstance(r, (int, float)): - assert pout.ndim == 0 - pout = pout.item() + pout = _evaluate_1d(self._p_i, r) # ad-hoc profile elif self.params["p_kind"] == 1: @@ -1367,7 +1384,6 @@ 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, @@ -1389,7 +1405,6 @@ def __init__( self.params = copy.deepcopy(locals()) from scipy.integrate import odeint - from scipy.interpolate import UnivariateSpline from scipy.optimize import fsolve # plasma boundary contour @@ -1421,24 +1436,20 @@ def dpsi_dr(psi, r, psi1): return out - # solve differential equation and fix boundary flux - r_i = xp.linspace(0.0, self.params["a"], self.params["psi_nel"] + 1) + # solve differential equation and fix boundary flux (on the host) + with xp.use_backend("numpy"): + r_i = xp.linspace(0.0, self.params["a"], self.params["psi_nel"] + 1) - def fun(psi1): - out = odeint(dpsi_dr, self._psi0, r_i, args=(psi1,)).flatten() + def fun(psi1): + out = odeint(dpsi_dr, self._psi0, r_i, args=(psi1,)).flatten() - return out[-1] - psi1 + return out[-1] - psi1 - self._psi1 = fsolve(fun, -9.5)[0] + self._psi1 = float(fsolve(fun, -9.5)[0]) + psis = odeint(dpsi_dr, self._psi0, r_i, args=(self._psi1,)).flatten() # interpolate flux function - self._psi_i = UnivariateSpline( - r_i, - odeint(dpsi_dr, self._psi0, r_i, args=(self._psi1,)).flatten(), - k=self.params["psi_k"], - s=0.0, - ext=3, - ) + self._psi_i = _univariate_spline(r_i, psis, self.params["psi_k"]) @property def boundary_pts_R(self): @@ -1473,14 +1484,7 @@ def psi_r(self, r, der=0): assert der >= 0 and der <= 2, "Only first and second derivatives available!" - out = xp.host_call(self._psi_i, r, nu=der) - - # remove all "dimensions" for point-wise evaluation - if isinstance(r, (int, float)): - assert out.ndim == 0 - out = out.item() - - return out + return _evaluate_1d(self._psi_i, r, nu=der) def q_psi(self, psi): """Safety factor profile q = q(psi).""" @@ -1668,7 +1672,6 @@ 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, @@ -1686,7 +1689,6 @@ def __init__( # use params setter self.params = copy.deepcopy(locals()) - from scipy.interpolate import RectBivariateSpline, UnivariateSpline from scipy.optimize import minimize # default input file @@ -1755,31 +1757,32 @@ def __init__( self._r_range = [rleft, rleft + rdim] self._z_range = [zmid - zdim / 2, zmid + zdim / 2] - R = xp.linspace(self._r_range[0], self._r_range[1], nR) - Z = xp.linspace(self._z_range[0], self._z_range[1], nZ) - smooth_steps = [ int(1 / (self.params["psi_resolution"][0] * 0.01)), int(1 / (self.params["psi_resolution"][1] * 0.01)), ] - self._psi_i = RectBivariateSpline( - R[:: smooth_steps[0]], - Z[:: smooth_steps[1]], - psi[:: smooth_steps[0], :: smooth_steps[1]], - kx=self.params["degree_for_psi"][0], - ky=self.params["degree_for_psi"][1], - s=0.0, - ) - - # find minimum of interpolated flux function (is not the same as (R_at_axis, Z_at_axis) and psi.min()!) - self._psi_i_min = minimize( - lambda x: self.psi( - x[0], - x[1], - ), - x0=[R_at_axis, Z_at_axis], - ) + # the file data and the subsampled grid stay on the host + with xp.use_backend("numpy"): + R = xp.linspace(self._r_range[0], self._r_range[1], nR)[:: smooth_steps[0]] + Z = xp.linspace(self._z_range[0], self._z_range[1], nZ)[:: smooth_steps[1]] + psi = psi[:: smooth_steps[0], :: smooth_steps[1]] + kx, ky = self.params["degree_for_psi"] + + # the spline is clamped to its data, like RectBivariateSpline (the subsampled grid may end before the range) + self._psi_bounds = ((float(R[0]), float(R[-1])), (float(Z[0]), float(Z[-1]))) + self._psi_i = _rect_spline(R, Z, psi, kx, ky) + + # find minimum of interpolated flux function (is not the same as (R_at_axis, Z_at_axis) and psi.min()!), with + # a host copy of the spline, so that the optimizer steps don't evaluate on the device + with xp.use_backend("numpy"): + psi_i_host = _rect_spline(R, Z, psi, kx, ky) + self._psi_i_min = minimize( + lambda x: ( + float(_evaluate_rect(psi_i_host, self._psi_bounds, x[0], x[1])) / (self.units.B * self.units.x**2) + ), + x0=[R_at_axis, Z_at_axis], + ) # set on-axis and boundary fluxes # (psi_edge is read from file in Weber/rad; rescale like self.psi()) @@ -1787,31 +1790,15 @@ def __init__( self._psi1 = psi_edge / (self.units.B * self.units.x**2) # interpolate toroidal field function, pressure profile and q-profile on unifrom flux grid from axis to boundary - flux_grid = xp.linspace(self._psi0, self._psi1, g_profile.size) + with xp.use_backend("numpy"): + flux_grid = xp.linspace(self._psi0, self._psi1, g_profile.size) smooth_step = int(1 / (self.params["flux_resolution"] * 0.01)) + k = self.params["degree_for_flux"] - self._g_i = UnivariateSpline( - flux_grid[::smooth_step], - g_profile[::smooth_step], - k=self.params["degree_for_flux"], - s=0.0, - ext=3, - ) - self._p_i = UnivariateSpline( - flux_grid[::smooth_step], - p_profile[::smooth_step], - k=self.params["degree_for_flux"], - s=0.0, - ext=3, - ) - self._q_i = UnivariateSpline( - flux_grid[::smooth_step], - q_profile[::smooth_step], - k=self.params["degree_for_flux"], - s=0.0, - ext=3, - ) + self._g_i = _univariate_spline(flux_grid[::smooth_step], g_profile[::smooth_step], k) + self._p_i = _univariate_spline(flux_grid[::smooth_step], p_profile[::smooth_step], k) + self._q_i = _univariate_spline(flux_grid[::smooth_step], q_profile[::smooth_step], k) @property def units(self) -> Units: @@ -1868,36 +1855,15 @@ def psi_axis_RZ(self): def q_psi(self, psi, der=0): """Safety factor q = q(psi).""" - out = xp.host_call(self._q_i, psi, nu=der) - - # remove all "dimensions" for point-wise evaluation - if isinstance(psi, (int, float)): - assert out.ndim == 0 - out = out.item() - - return out + return _evaluate_1d(self._q_i, psi, nu=der) def g_psi(self, psi, der=0): """Toroidal field function g = g(psi).""" - out = xp.host_call(self._g_i, psi, nu=der) - - # remove all "dimensions" for point-wise evaluation - if isinstance(psi, (int, float)): - assert out.ndim == 0 - out = out.item() - - return out + return _evaluate_1d(self._g_i, psi, nu=der) def p_psi(self, psi, der=0): """Pressure profile p = p(psi) in units Pa (as in the EQDSK file).""" - out = xp.host_call(self._p_i, psi, nu=der) - - # remove all "dimensions" for point-wise evaluation - if isinstance(psi, (int, float)): - assert out.ndim == 0 - out = out.item() - - return out + return _evaluate_1d(self._p_i, psi, nu=der) def n_psi(self, psi, der=0): """Number density profile n = n(psi).""" @@ -1926,11 +1892,11 @@ def psi(self, R, Z, dR=0, dZ=0): is_float = all(isinstance(v, (int, float)) for v in [R, Z]) - out = xp.host_call(self._psi_i, R, Z, dx=dR, dy=dZ, grid=False) + out = _evaluate_rect(self._psi_i, self._psi_bounds, R, Z, dR, dZ) # remove all "dimensions" for point-wise evaluation if is_float: - assert out.size == 1 # scipy >= 1.18 returns shape (1,) for a single point + assert out.size == 1 out = out.item() # rescale to Struphy units diff --git a/src/struphy/fields_background/tests/test_equils_cupy.py b/src/struphy/fields_background/tests/test_equils_cupy.py index 28614e65a..8c51674bc 100644 --- a/src/struphy/fields_background/tests/test_equils_cupy.py +++ b/src/struphy/fields_background/tests/test_equils_cupy.py @@ -56,6 +56,10 @@ # equilibria that need an optional package (skipped without it) NEEDS = {"GVECequilibrium": "gvec", "DESCequilibrium": "desc"} +# equilibria evaluated through their native code on the host (one round trip per call); all others are evaluated on +# the device only, without host/device transfers +HOST_EVALUATION = {"GVECequilibrium", "DESCequilibrium"} + # methods that models and projections call on the logical domain; those an equilibrium does not provide are skipped METHODS = ( "absB0", @@ -181,7 +185,8 @@ 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). + agrees with the NumPy result. Evaluated on a meshgrid (three 1d arrays) and at markers (one 2d array), without + host/device transfers except for the equilibria in ``HOST_EVALUATION``. """ rng = np.random.default_rng(1234) e1 = np.sort(rng.uniform(0.1, 0.9, 4)) @@ -196,8 +201,15 @@ def check_equil_on_cupy(case): 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),)) + grid = tuple(cunumpy.asarray(e) for e in (e1, e2, e3)) + markers_device = cunumpy.asarray(markers) + if EQUIL_CASES[case][0] in HOST_EVALUATION: + no_transfers = contextlib.nullcontext() + else: + no_transfers = cunumpy.profiling.assert_no_transfers() + with no_transfers: + out_grid = _evaluations(equil, grid) + out_markers = _evaluations(equil, (markers_device,)) for ref, out, where in ((ref_grid, out_grid, "meshgrid"), (ref_markers, out_markers, "markers")): assert ref.keys() == out.keys() diff --git a/src/struphy/geometry/domains/tokamak/tokamak.py b/src/struphy/geometry/domains/tokamak/tokamak.py index 1fa3d7fcc..7c3c788ea 100644 --- a/src/struphy/geometry/domains/tokamak/tokamak.py +++ b/src/struphy/geometry/domains/tokamak/tokamak.py @@ -67,8 +67,8 @@ def __init__( ): if r_min != 0.0: r0 = r_min - # 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. + # The field-line tracing evaluates psi on the active backend (batched Newton over all rays per flux surface); + # its small spline interpolation runs on the host, and the control points are copied to the active backend once. if equilibrium is None: equilibrium = EQDSKequilibrium() else: @@ -77,48 +77,47 @@ def __init__( # use the params setter self.params = copy.deepcopy(locals()) - with xp.use_backend("numpy"): - # get control points via field tracing between fluxes [psi_s, psi_e] - psi0, psi1 = equilibrium.psi_range[0], equilibrium.psi_range[1] + # get control points via field tracing between fluxes [psi_s, psi_e] + psi0, psi1 = equilibrium.psi_range[0], equilibrium.psi_range[1] - assert r_min >= 0.0, f"Inner radius must be non-negative, got {r_min = }." + assert r_min >= 0.0, f"Inner radius must be non-negative, got {r_min = }." - if r_min == 0.0: - # Default behaviour: keep exactly the historical psi_shifts logic. - psi_s = psi0 + psi_shifts[0] * 0.01 * (psi1 - psi0) - else: - # Annular domain: eta1=0 is the flux surface crossing the outboard - # midplane at distance r_min from the magnetic axis. - psi_s = equilibrium.psi( - equilibrium.psi_axis_RZ[0] + r_min, - equilibrium.psi_axis_RZ[1], - ) + if r_min == 0.0: + # Default behaviour: keep exactly the historical psi_shifts logic. + psi_s = psi0 + psi_shifts[0] * 0.01 * (psi1 - psi0) + else: + # Annular domain: eta1=0 is the flux surface crossing the outboard + # midplane at distance r_min from the magnetic axis. + psi_s = equilibrium.psi( + equilibrium.psi_axis_RZ[0] + r_min, + equilibrium.psi_axis_RZ[1], + ) - psi_e = psi1 - psi_shifts[1] * 0.01 * (psi1 - psi0) + psi_e = psi1 - psi_shifts[1] * 0.01 * (psi1 - psi0) - assert (psi_s - psi0) * (psi_s - psi1) <= 0.0, ( - f"Inner radius gives a flux outside equilibrium.psi_range: " - f"{r_min = }, {psi_s = }, {equilibrium.psi_range = }." - ) + assert (psi_s - psi0) * (psi_s - psi1) <= 0.0, ( + f"Inner radius gives a flux outside equilibrium.psi_range: " + f"{r_min = }, {psi_s = }, {equilibrium.psi_range = }." + ) - assert (psi_e - psi_s) * (psi1 - psi0) > 0.0, ( - f"Invalid radial interval: {psi_s = }, {psi_e = }, {equilibrium.psi_range = }." - ) + assert (psi_e - psi_s) * (psi1 - psi0) > 0.0, ( + f"Invalid radial interval: {psi_s = }, {psi_e = }, {equilibrium.psi_range = }." + ) - cx, cy = field_line_tracing( - equilibrium.psi, - equilibrium.psi_axis_RZ[0], - equilibrium.psi_axis_RZ[1], - psi_s, - psi_e, - num_elements, - degree, - psi_power=psi_power, - xi_param=xi_param, - num_elements_pre=num_elements_pre, - p_pre=p_pre, - r0=r0, - ) + cx, cy = field_line_tracing( + equilibrium.psi, + equilibrium.psi_axis_RZ[0], + equilibrium.psi_axis_RZ[1], + psi_s, + psi_e, + num_elements, + degree, + psi_power=psi_power, + xi_param=xi_param, + num_elements_pre=num_elements_pre, + p_pre=p_pre, + r0=r0, + ) cx, cy = xp.to_cunumpy(cx), xp.to_cunumpy(cy) diff --git a/src/struphy/geometry/utilities.py b/src/struphy/geometry/utilities.py index e09c7f78b..c4802af89 100644 --- a/src/struphy/geometry/utilities.py +++ b/src/struphy/geometry/utilities.py @@ -150,9 +150,19 @@ def field_line_tracing( Control points (2d) of flux aligned spline mapping (Z-component). """ - from scipy.optimize import newton, root + from scipy.optimize import root from scipy.sparse.linalg import splu + # The roots of the equal-angle tracing are found on the active backend: one batched Newton (secant) solve over + # all rays per flux surface, with psi evaluated on the device on the CuPy backend. The knots, collocation matrices + # and their (small) LU solves, and the other xi-parametrizations (a SciPy root per flux surface) run on the host. + backend = xp.get_backend() + + def psi_host(R, Z, **kwargs): + """psi at host points, for the host parts: evaluated on the caller's backend.""" + with xp.use_backend(backend): + return xp.to_numpy(psi(xp.to_cunumpy(R), xp.to_cunumpy(Z), **kwargs)) + # for equal_angle one mapping is enough if xi_param == "equal_angle": ns, nx = num_elements @@ -161,70 +171,70 @@ def field_line_tracing( ns, nx = num_elements_pre ps, px = p_pre - # spline knots - Ts = bsp.make_knots(xp.linspace(0.0, 1.0, ns + 1), ps, False) - Tx = bsp.make_knots(xp.linspace(0.0, 1.0, nx + 1), px, True) + with xp.use_backend("numpy"): + # spline knots + Ts = bsp.make_knots(xp.linspace(0.0, 1.0, ns + 1), ps, False) + Tx = bsp.make_knots(xp.linspace(0.0, 1.0, nx + 1), px, True) - # interpolation (Greville) points - s_gr = bsp.greville(Ts, ps, False) - x_gr = bsp.greville(Tx, px, True) + # interpolation (Greville) points + s_gr = bsp.greville(Ts, ps, False) + x_gr = bsp.greville(Tx, px, True) - if px % 2 == 1: - assert x_gr[0] == 0.0 + if px % 2 == 1: + assert x_gr[0] == 0.0 - # collocation matrices - Is = bsp.collocation_matrix(Ts, ps, s_gr, False) - Ix = bsp.collocation_matrix(Tx, px, x_gr, True) + # collocation matrices + Is = bsp.collocation_matrix(Ts, ps, s_gr, False) + Ix = bsp.collocation_matrix(Tx, px, x_gr, True) - ILUs = [ - splu(csc_matrix(Is)), - splu(csc_matrix(Ix)), - ] + ILUs = [ + splu(csc_matrix(Is)), + splu(csc_matrix(Ix)), + ] # check if pole is included - if xp.abs(psi(psi_axis_R, psi_axis_Z) - psi0) < 1e-14: + if abs(float(psi(psi_axis_R, psi_axis_Z)) - psi0) < 1e-14: pole = True else: pole = False + # the rays from the magnetic axis + cos_x = xp.cos(2 * xp.pi * xp.to_cunumpy(x_gr)) + sin_x = xp.sin(2 * xp.pi * xp.to_cunumpy(x_gr)) + R = xp.zeros((s_gr.size, x_gr.size), dtype=float) Z = xp.zeros((s_gr.size, x_gr.size), dtype=float) r_flux_surface = None - # function whose root must be found - for j, x in enumerate(x_gr): - for i, s in enumerate(s_gr): - if pole and i == 0: - R[i, j] = psi_axis_R - Z[i, j] = psi_axis_Z - continue - - if i < s_gr.size // 2: - r_guess = 1 * r0 - else: - r_guess = 1 * r_flux_surface + for i, s in enumerate(s_gr): + if pole and i == 0: + R[i] = psi_axis_R + Z[i] = psi_axis_Z + continue - # function whose root must be found - def f(r): - _R = psi_axis_R + r * xp.cos(2 * xp.pi * x) - _Z = psi_axis_Z + r * xp.sin(2 * xp.pi * x) + # function whose roots must be found: one equation per ray + def f(r, s=float(s)): + psi_norm = (psi(psi_axis_R + r * cos_x, psi_axis_Z + r * sin_x) - psi0) / (psi1 - psi0) + return xp.sign(psi_norm) * xp.abs(psi_norm) ** psi_power - s - psi_norm = (psi(_R, _Z) - psi0) / (psi1 - psi0) - - if psi_norm < 0.0: - return -((-psi_norm) ** psi_power) - s - else: - return psi_norm**psi_power - s + # the inner surfaces start from r0, the outer ones from the previous surface on the same ray + if i < s_gr.size // 2: + r_guess = xp.full(x_gr.size, float(r0)) + else: + r_guess = r_flux_surface - r_flux_surface = newton(f, x0=r_guess) + tracing = xp.optimize.newton(f, r_guess, full_output=True) + assert bool(xp.all(tracing.converged)), f"Field-line tracing did not converge on flux surface {s = }." + r_flux_surface = tracing.root - R[i, j] = psi_axis_R + r_flux_surface * xp.cos(2 * xp.pi * x) - Z[i, j] = psi_axis_Z + r_flux_surface * xp.sin(2 * xp.pi * x) + R[i] = psi_axis_R + r_flux_surface * cos_x + Z[i] = psi_axis_Z + r_flux_surface * sin_x - # get control points - cR_equal_angle = kron_lusolve_2d(ILUs, R) - cZ_equal_angle = kron_lusolve_2d(ILUs, Z) + # get control points (on the host) + with xp.use_backend("numpy"): + cR_equal_angle = kron_lusolve_2d(ILUs, xp.to_numpy(R)) + cZ_equal_angle = kron_lusolve_2d(ILUs, xp.to_numpy(Z)) if pole: cR_equal_angle[0, :] = psi_axis_R @@ -236,125 +246,127 @@ def f(r): # for all other parametrizations continue else: - logger.info("Calculation of pre-mapping successful! Start angle parametrization " + xi_param + ".") + # on the host (a SciPy root and a pyccel kernel per flux surface) + with xp.use_backend("numpy"): + logger.info("Calculation of pre-mapping successful! Start angle parametrization " + xi_param + ".") - # create temporary domain - domain_eq_angle = PoloidalSplineTorus( - num_elements=num_elements_pre, degree=p_pre, cx=cR_equal_angle, cy=cZ_equal_angle - ) + # create temporary domain + domain_eq_angle = PoloidalSplineTorus( + num_elements=num_elements_pre, degree=p_pre, cx=cR_equal_angle, cy=cZ_equal_angle + ) - # create new interpolation data - ns, nx = num_elements - ps, px = degree + # create new interpolation data + ns, nx = num_elements + ps, px = degree - # spline knots - Ts = bsp.make_knots(xp.linspace(0.0, 1.0, ns + 1), ps, False) - Tx = bsp.make_knots(xp.linspace(0.0, 1.0, nx + 1), px, True) + # spline knots + Ts = bsp.make_knots(xp.linspace(0.0, 1.0, ns + 1), ps, False) + Tx = bsp.make_knots(xp.linspace(0.0, 1.0, nx + 1), px, True) - # interpolation (Greville) points - s_gr = bsp.greville(Ts, ps, False) - x_gr = bsp.greville(Tx, px, True) + # interpolation (Greville) points + s_gr = bsp.greville(Ts, ps, False) + x_gr = bsp.greville(Tx, px, True) - if px % 2 == 1: - assert x_gr[0] == 0.0 + if px % 2 == 1: + assert x_gr[0] == 0.0 - # collocation matrices - Is = bsp.collocation_matrix(Ts, ps, s_gr, False) - Ix = bsp.collocation_matrix(Tx, px, x_gr, True) + # collocation matrices + Is = bsp.collocation_matrix(Ts, ps, s_gr, False) + Ix = bsp.collocation_matrix(Tx, px, x_gr, True) - ILUs = [ - splu(csc_matrix(Is)), - splu(csc_matrix(Ix)), - ] + ILUs = [ + splu(csc_matrix(Is)), + splu(csc_matrix(Ix)), + ] - xi_param_dict = { - "equal_arc_length": 1, - "sfl": 2, - "equal_area": 3, - "equal_volume": 4, - } + xi_param_dict = { + "equal_arc_length": 1, + "sfl": 2, + "equal_area": 3, + "equal_volume": 4, + } - # target function for xi parametrization - def f_angles(xis, s_val): - assert xp.all(xp.logical_and(xis > 0.0, xis < 1.0)) + # target function for xi parametrization + def f_angles(xis, s_val): + assert xp.all(xp.logical_and(xis > 0.0, xis < 1.0)) - # add 0 and 1 to angles array - xis_extended = xp.array([0.0] + list(xis) + [1.0]) + # add 0 and 1 to angles array + xis_extended = xp.array([0.0] + list(xis) + [1.0]) - # compute (R, Z) coordinates for given xis on fixed flux surface corresponding to s_val - _RZ = domain_eq_angle(s_val, xis_extended, 0.0, squeeze_out=True) + # compute (R, Z) coordinates for given xis on fixed flux surface corresponding to s_val + _RZ = domain_eq_angle(s_val, xis_extended, 0.0, squeeze_out=True) - _R = _RZ[0] - _Z = _RZ[2] + _R = _RZ[0] + _Z = _RZ[2] - # |grad(psi)| at xis - gp = xp.sqrt(psi(_R, _Z, dR=1) ** 2 + psi(_R, _Z, dZ=1) ** 2) + # |grad(psi)| at xis + gp = xp.sqrt(psi_host(_R, _Z, dR=1) ** 2 + psi_host(_R, _Z, dZ=1) ** 2) - # compute weighted arc_lengths between two successive points in xis_extended array - dl = xp.zeros(xis_extended.size - 1, dtype=float) - weighted_arc_lengths_flux_surface(_R, _Z, gp, dl, xi_param_dict[xi_param]) + # compute weighted arc_lengths between two successive points in xis_extended array + dl = xp.zeros(xis_extended.size - 1, dtype=float) + weighted_arc_lengths_flux_surface(_R, _Z, gp, dl, xi_param_dict[xi_param]) - # total length of the flux surface - l = xp.sum(dl) + # total length of the flux surface + l = xp.sum(dl) - # cumulative sum of arc lengths, start with 0! - l_cum = xp.cumsum(dl) + # cumulative sum of arc lengths, start with 0! + l_cum = xp.cumsum(dl) - # odd spline degree - if px % 2 == 1: - xi_diff = l_cum[:-1] / l - x_gr[1:] - # even spline degree - else: - xi_diff = l_cum[:-1] / l - x_gr + # odd spline degree + if px % 2 == 1: + xi_diff = l_cum[:-1] / l - x_gr[1:] + # even spline degree + else: + xi_diff = l_cum[:-1] / l - x_gr - return xi_diff + return xi_diff - # loop over flux surfaces and find xi parametrization - R = xp.zeros((s_gr.size, x_gr.size), dtype=float) - Z = xp.zeros((s_gr.size, x_gr.size), dtype=float) + # loop over flux surfaces and find xi parametrization + R = xp.zeros((s_gr.size, x_gr.size), dtype=float) + Z = xp.zeros((s_gr.size, x_gr.size), dtype=float) - if px % 2 == 1: - xis0 = x_gr[1:].copy() - else: - xis0 = x_gr.copy() + if px % 2 == 1: + xis0 = x_gr[1:].copy() + else: + xis0 = x_gr.copy() - # loop over flux surfaces and finds roots of F_single - for i in range(s_gr.size): - s_flux = s_gr[i] + # loop over flux surfaces and finds roots of F_single + for i in range(s_gr.size): + s_flux = s_gr[i] - if i == 0 and pole: - R[i, :] = psi_axis_R - Z[i, :] = psi_axis_Z - continue + if i == 0 and pole: + R[i, :] = psi_axis_R + Z[i, :] = psi_axis_Z + continue - # find root of target function and check for convergence - tracing = root(f_angles, x0=xis0, args=(s_flux,), method="hybr") - assert tracing["success"] + # find root of target function and check for convergence + tracing = root(f_angles, x0=xis0, args=(s_flux,), method="hybr") + assert tracing["success"] - # set new initial guess - xis0 = tracing["x"] + # set new initial guess + xis0 = tracing["x"] - # add zero angle for odd degree - if px % 2 == 1: - R[i, 1:] = domain_eq_angle(s_flux, tracing["x"], 0.0, squeeze_out=True)[0] - Z[i, 1:] = domain_eq_angle(s_flux, tracing["x"], 0.0, squeeze_out=True)[2] + # add zero angle for odd degree + if px % 2 == 1: + R[i, 1:] = domain_eq_angle(s_flux, tracing["x"], 0.0, squeeze_out=True)[0] + Z[i, 1:] = domain_eq_angle(s_flux, tracing["x"], 0.0, squeeze_out=True)[2] - R[i, 0] = domain_eq_angle(s_flux, 0.0, 0.0, squeeze_out=True)[0] - Z[i, 0] = domain_eq_angle(s_flux, 0.0, 0.0, squeeze_out=True)[2] + R[i, 0] = domain_eq_angle(s_flux, 0.0, 0.0, squeeze_out=True)[0] + Z[i, 0] = domain_eq_angle(s_flux, 0.0, 0.0, squeeze_out=True)[2] - else: - R[i, :] = domain_eq_angle(s_flux, tracing["x"], 0.0, squeeze_out=True)[0] - Z[i, :] = domain_eq_angle(s_flux, tracing["x"], 0.0, squeeze_out=True)[2] + else: + R[i, :] = domain_eq_angle(s_flux, tracing["x"], 0.0, squeeze_out=True)[0] + Z[i, :] = domain_eq_angle(s_flux, tracing["x"], 0.0, squeeze_out=True)[2] - # get control points - cR = kron_lusolve_2d(ILUs, R) - cZ = kron_lusolve_2d(ILUs, Z) + # get control points + cR = kron_lusolve_2d(ILUs, R) + cZ = kron_lusolve_2d(ILUs, Z) - if pole: - cR[0, :] = psi_axis_R - cZ[0, :] = psi_axis_Z + if pole: + cR[0, :] = psi_axis_R + cZ[0, :] = psi_axis_Z - return cR, cZ + return cR, cZ class TransformedPformComponent: From 70557fba908063f3c0d48d1e508764812992beb0 Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 10 Oct 2026 16:11:45 +0200 Subject: [PATCH 32/50] Pin cunumpy to the commit with copyable fake cupyx splines (cunumpy#97) Co-Authored-By: Claude Sonnet 5.5 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 3e637a235..6294f4294 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ classifiers = [ ] dependencies = [ "numpy<=2.5.0", - "cunumpy @ git+https://github.com/max-models/cunumpy.git@fb0f7c0afaac7994dd62147c1fa39e1d9618318a", + "cunumpy @ git+https://github.com/max-models/cunumpy.git@14fb37038524bd7b39ff90b10c29e04ae24ec0ad", "pyccel>=2.2.0, <=2.2.3", "maybempi>=0.1.2", "feectools>=0.8.0, <=0.8.0", From 822af6461d75fe964eb4e40bc30ca93adf0af2ad Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 10 Oct 2026 17:10:49 +0200 Subject: [PATCH 33/50] Pin cunumpy to devel with copyable fake cupyx splines (cunumpy#97 merged) Co-Authored-By: Claude Sonnet 5.5 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 6294f4294..fbfb128fa 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ classifiers = [ ] dependencies = [ "numpy<=2.5.0", - "cunumpy @ git+https://github.com/max-models/cunumpy.git@14fb37038524bd7b39ff90b10c29e04ae24ec0ad", + "cunumpy @ git+https://github.com/max-models/cunumpy.git@b62b2b7e785cd506ab296d73f97e9938c58e70f3", "pyccel>=2.2.0, <=2.2.3", "maybempi>=0.1.2", "feectools>=0.8.0, <=0.8.0", From dce63a96258bff52fd5d22ca37eca8c2b04d9e1e Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 10 Oct 2026 19:25:18 +0200 Subject: [PATCH 34/50] Pin cunumpy to devel (f996393: copyable fake cupyx splines, faster type checking) Co-Authored-By: Claude Sonnet 5.5 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index fbfb128fa..2542ffd6f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,7 +24,7 @@ classifiers = [ ] dependencies = [ "numpy<=2.5.0", - "cunumpy @ git+https://github.com/max-models/cunumpy.git@b62b2b7e785cd506ab296d73f97e9938c58e70f3", + "cunumpy @ git+https://github.com/max-models/cunumpy.git@f9963935338b323ed6c4b212020776062a029038", "pyccel>=2.2.0, <=2.2.3", "maybempi>=0.1.2", "feectools>=0.8.0, <=0.8.0", From 8fcb59f8d3794362dcc4c3abf84826dc81afde9c Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 7 Oct 2026 19:46:18 +0200 Subject: [PATCH 35/50] cupy array fixes --- src/struphy/feec/psydac_derham.py | 3 ++- src/struphy/pic/sorting.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 7394b5b10..dfdb34e4c 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -2326,7 +2326,8 @@ def __init__( if xp.get_backend() == "cupy" and np.any((degree < 1) | (degree > 8)): raise ValueError("CUDA spline degrees must be between 1 and 8.") pn = xp.asarray(degree) - knots = tuple(xp.ascontiguousarray(t, dtype=float) for t in derham.V0fem.knots) + # the knots are host arrays; cupy.ascontiguousarray does not take NumPy arrays, xp.asarray copies them + knots = tuple(xp.asarray(np.ascontiguousarray(t, dtype=float)) for t in derham.V0fem.knots) starts = (self.starts,) if isinstance(self._vector_stencil, StencilVector) else self.starts kinds = derham.spline_attributes[self.space_key].spline_types_pyccel self._args_eval = tuple( diff --git a/src/struphy/pic/sorting.py b/src/struphy/pic/sorting.py index 2d331d6c1..74db8c182 100644 --- a/src/struphy/pic/sorting.py +++ b/src/struphy/pic/sorting.py @@ -1,4 +1,5 @@ import logging +import math from typing import TYPE_CHECKING import cunumpy as xp @@ -226,7 +227,7 @@ def _set_boxes(self): n_particles = self._markers_shape[0] n_mkr = int(n_particles / n_box_in) + 1 n_cols = round( - n_mkr * (1 + 1 / xp.sqrt(n_mkr) + self._box_bufsize), + n_mkr * (1 + 1 / math.sqrt(n_mkr) + self._box_bufsize), ) # cartesian boxes (extra last row stores holes/outside particles) From 8d8603abcd38f78abe6262be44bdd37ff61ba4ae Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Wed, 7 Oct 2026 19:24:11 +0200 Subject: [PATCH 36/50] Fuse the pusher preamble into one CUDA kernel --- src/struphy/pic/pushing/prepare_push_cuda.cu | 28 ++++++++++++++ src/struphy/pic/pushing/pusher.py | 39 +++++++++++++++----- 2 files changed, 58 insertions(+), 9 deletions(-) create mode 100644 src/struphy/pic/pushing/prepare_push_cuda.cu diff --git a/src/struphy/pic/pushing/prepare_push_cuda.cu b/src/struphy/pic/pushing/prepare_push_cuda.cu new file mode 100644 index 000000000..5c1dbd28e --- /dev/null +++ b/src/struphy/pic/pushing/prepare_push_cuda.cu @@ -0,0 +1,28 @@ +/** + * Prepare the marker buffer for a push, as the three column-slice assignments at the start of Pusher._push: + * + * markers[:, first_init_idx:first_shift_idx] = markers[:, :n_init] + * markers[:, first_shift_idx:residual_idx] = 0.0 + * markers[:, residual_idx:-2] = 0.0 + * + * One thread per entry of the contiguous column range [first_init_idx, n_cols - 2) of every row, so that + * neighbouring threads write neighbouring addresses of the row-major buffer and the whole range is written + * in a single pass (instead of one strided pass over the buffer per slice). + * + * @param markers Marker buffer (n_rows x n_cols, row-major), every row including holes. + * @param n_rows Number of rows of the buffer. + * @param n_cols Number of columns of the buffer. + * @param first_init_idx First column of the saved initial phase space coordinates. + * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_init_idx + n_init). + */ +extern "C" __global__ void prepare_push(double* markers, long long n_rows, int n_cols, int first_init_idx, + int n_init) { + long long width = n_cols - 2 - first_init_idx; + long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; + if (t >= n_rows * width) return; + + long long ip = t / width; + int col = first_init_idx + (int)(t - ip * width); + double* row = markers + ip * n_cols; + row[col] = col < first_init_idx + n_init ? row[col - first_init_idx] : 0.0; +} diff --git a/src/struphy/pic/pushing/pusher.py b/src/struphy/pic/pushing/pusher.py index 7f3d64961..762dd3bd9 100644 --- a/src/struphy/pic/pushing/pusher.py +++ b/src/struphy/pic/pushing/pusher.py @@ -1,9 +1,11 @@ "Accelerated particle pushing." import logging +from functools import cache +from pathlib import Path import cunumpy as xp -from cunumpy.kernels import Kernel, PyccelKernel +from cunumpy.kernels import CudaKernel, Kernel, PyccelKernel from line_profiler import profile from maybempi import MPI from scope_profiler import ProfileManager @@ -16,6 +18,13 @@ logger = logging.getLogger("struphy") +@cache +def _prepare_push_cuda() -> CudaKernel: + """The CUDA kernel that fuses the column-slice assignments at the start of :meth:`Pusher._push` + (compiled on first call).""" + return CudaKernel.from_file(Path(__file__).with_name("prepare_push_cuda.cu"), "prepare_push") + + class Pusher: r""" Class for solving particle ODEs @@ -208,17 +217,29 @@ def _push(self, dt: float): logger.debug(f"{residual_idx =}") logger.debug(f"{self.particles.n_cols =}") - init_slice = slice(first_pusher_idx, first_shift_idx) - shift_slice = slice(first_shift_idx, residual_idx) + if self._cuda: + # the three assignments below in one pass over the buffer (each slice is a strided pass on its own) + n_rows, n_cols = markers.shape + _prepare_push_cuda()( + markers, + n_rows, + n_cols, + first_pusher_idx, + 3 + vdim, + n_threads=n_rows * (n_cols - 2 - first_pusher_idx), + ) + else: + init_slice = slice(first_pusher_idx, first_shift_idx) + shift_slice = slice(first_shift_idx, residual_idx) - # save initial phase space coordinates - markers[:, init_slice] = markers[:, : 3 + vdim] + # save initial phase space coordinates + markers[:, init_slice] = markers[:, : 3 + vdim] - # set boundary shifts to zero - markers[:, shift_slice] = 0.0 + # set boundary shifts to zero + markers[:, shift_slice] = 0.0 - # clear buffer columns starting from residual index, dont clear ID (last column) and loc_box - markers[:, residual_idx:-2] = 0.0 + # clear buffer columns starting from residual index, dont clear ID (last column) and loc_box + markers[:, residual_idx:-2] = 0.0 rank = self.particles.mpi_rank logger.debug(f"rank {rank}: starting {self.kernel} ...") From 30d81a13cb99b1b21cf1482f664d31193275a854 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 15:02:41 +0200 Subject: [PATCH 37/50] Rename first_init_idx to first_pusher_idx The pusher's scratch columns were called first_pusher_idx in Particles and first_init_idx in the kernel arguments; use first_pusher_idx everywhere. Co-Authored-By: Claude Opus 5.5 --- src/struphy/kernel_arguments/pusher_args.cuh | 2 +- src/struphy/kernel_arguments/pusher_args_cuda.py | 4 ++-- .../kernel_arguments/pusher_args_kernels.py | 2 +- .../cc_lin_mhd_5d_gradB_kernels.py | 4 ++-- .../cc_lin_mhd_5d_gradB_dg_kernels.py | 6 +++--- .../eval_gradB_ediff/eval_gradB_ediff_kernels.py | 6 +++--- .../kernels/bstar_2form/bstar_2form_kernels.py | 6 +++--- .../bstar_parallel_3form_kernels.py | 6 +++--- .../driftkinetic_hamiltonian_kernels.py | 6 +++--- .../grad_driftkinetic_hamiltonian_kernels.py | 4 ++-- ...push_deterministic_diffusion_stage_kernels.py | 4 ++-- .../push_eta_stage/push_eta_stage_cuda.cu | 8 ++++---- .../push_eta_stage/push_eta_stage_kernels.py | 6 +++--- ..._Bstar_discrete_gradient_1st_order_kernels.py | 8 ++++---- ...discrete_gradient_1st_order_newton_kernels.py | 8 ++++---- ..._Bstar_discrete_gradient_2nd_order_kernels.py | 8 ++++---- .../push_gc_Bstar_explicit_multistage_kernels.py | 10 ++++++---- ...ush_gc_bxEstar_dg_1st_order_newton_kernels.py | 6 +++--- ...xEstar_discrete_gradient_1st_order_kernels.py | 6 +++--- ...xEstar_discrete_gradient_2nd_order_kernels.py | 6 +++--- ...ush_gc_bxEstar_explicit_multistage_kernels.py | 6 +++--- .../push_gc_cc_J2_dg_Hdiv_kernels.py | 6 +++--- .../push_gc_cc_J2_dg_init_Hdiv_kernels.py | 4 ++-- .../push_gc_cc_J2_stage_H1vec_kernels.py | 6 +++--- .../push_gc_cc_J2_stage_Hdiv_kernels.py | 6 +++--- .../push_pc_eta_stage_H1vec_kernels.py | 4 ++-- .../push_pc_eta_stage_Hcurl_kernels.py | 4 ++-- .../push_pc_eta_stage_Hdiv_kernels.py | 4 ++-- .../push_vxb_analytic/push_vxb_analytic_cuda.cu | 4 ++-- .../push_vxb_analytic_kernels.py | 4 ++-- .../push_vxb_implicit/push_vxb_implicit_cuda.cu | 4 ++-- .../push_vxb_implicit_kernels.py | 4 ++-- .../kernels/unit_b_1form/unit_b_1form_kernels.py | 4 ++-- src/struphy/pic/pushing/prepare_push_cuda.cu | 16 ++++++++-------- .../pic/pushing/pusher_utilities_kernels.cuh | 6 +++--- .../pic/pushing/pusher_utilities_kernels.py | 6 +++--- .../propagators/current_coupling_5d_gradb.py | 12 ++++++------ 37 files changed, 109 insertions(+), 107 deletions(-) diff --git a/src/struphy/kernel_arguments/pusher_args.cuh b/src/struphy/kernel_arguments/pusher_args.cuh index 5c85f3d45..59815fced 100644 --- a/src/struphy/kernel_arguments/pusher_args.cuh +++ b/src/struphy/kernel_arguments/pusher_args.cuh @@ -12,7 +12,7 @@ struct MarkerArgs { int vdim; int weight_idx; int first_diagnostics_idx; - int first_init_idx; + int first_pusher_idx; int first_shift_idx; int residual_idx; int first_free_idx; diff --git a/src/struphy/kernel_arguments/pusher_args_cuda.py b/src/struphy/kernel_arguments/pusher_args_cuda.py index 0c9911d72..08fefff46 100644 --- a/src/struphy/kernel_arguments/pusher_args_cuda.py +++ b/src/struphy/kernel_arguments/pusher_args_cuda.py @@ -39,7 +39,7 @@ class CudaMarkerArguments(CudaStructArguments): ("vdim", "int"), ("weight_idx", "int"), ("first_diagnostics_idx", "int"), - ("first_init_idx", "int"), + ("first_pusher_idx", "int"), ("first_shift_idx", "int"), ("residual_idx", "int"), ("first_free_idx", "int"), @@ -71,7 +71,7 @@ def __init__( self.weight_idx = weight_idx self.n_markers = markers.shape[0] self.first_diagnostics_idx = first_diagnostics_idx - self.first_init_idx = first_pusher_idx + self.first_pusher_idx = first_pusher_idx self.first_shift_idx = first_shift_idx self.residual_idx = residual_idx self.first_free_idx = first_free_idx diff --git a/src/struphy/kernel_arguments/pusher_args_kernels.py b/src/struphy/kernel_arguments/pusher_args_kernels.py index c5b3f52d8..2e66dc58e 100644 --- a/src/struphy/kernel_arguments/pusher_args_kernels.py +++ b/src/struphy/kernel_arguments/pusher_args_kernels.py @@ -71,7 +71,7 @@ def __init__( # useful indices self.first_diagnostics_idx = first_diagnostics_idx - self.first_init_idx = first_pusher_idx + self.first_pusher_idx = first_pusher_idx self.first_shift_idx = first_shift_idx # starting idx for eta-shifts due to boundary conditions self.residual_idx = residual_idx # residual in iterative solvers self.first_free_idx = first_free_idx # index after which auxiliary saving is possible diff --git a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py index 4428c130d..e50a5e7a4 100644 --- a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py +++ b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB/cc_lin_mhd_5d_gradB_kernels.py @@ -99,7 +99,7 @@ def cc_lin_mhd_5d_gradB( markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx mu_idx = args_markers.mu_idx # allocate for magnetic field evaluation @@ -130,7 +130,7 @@ def cc_lin_mhd_5d_gradB( continue # if particle is refilled - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue # marker positions diff --git a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py index 7705768f5..015b4506a 100644 --- a/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py +++ b/src/struphy/pic/accumulation/kernels/cc_lin_mhd_5d_gradB_dg/cc_lin_mhd_5d_gradB_dg_kernels.py @@ -71,7 +71,7 @@ def cc_lin_mhd_5d_gradB_dg( markers = args_markers.markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx # allocate for magnetic field evaluation eta_diff = empty(3, dtype=float) @@ -108,10 +108,10 @@ def cc_lin_mhd_5d_gradB_dg( continue # marker positions, mid point - eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_init_idx : first_init_idx + 3]) / 2.0 + eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_pusher_idx : first_pusher_idx + 3]) / 2.0 eta_mid[:] = mod(eta_mid[:], 1.0) - eta_diff[:] = markers[ip, 0:3] - markers[ip, first_init_idx : first_init_idx + 3] + eta_diff[:] = markers[ip, 0:3] - markers[ip, first_pusher_idx : first_pusher_idx + 3] # marker weight and velocity weight = markers[ip, 5] diff --git a/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py b/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py index c21807b7c..e33183e72 100644 --- a/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py +++ b/src/struphy/pic/diagnostics/kernels/eval_gradB_ediff/eval_gradB_ediff_kernels.py @@ -42,7 +42,7 @@ def eval_gradB_ediff( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): @@ -51,10 +51,10 @@ def eval_gradB_ediff( continue # marker positions, mid point - eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_init_idx : first_init_idx + 3]) / 2.0 + eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_pusher_idx : first_pusher_idx + 3]) / 2.0 eta_mid[:] = mod(eta_mid[:], 1.0) - eta_diff = markers[ip, 0:3] - markers[ip, first_init_idx : first_init_idx + 3] + eta_diff = markers[ip, 0:3] - markers[ip, first_pusher_idx : first_pusher_idx + 3] # marker weight and velocity weight = markers[ip, 5] diff --git a/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py b/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py index ea1ba6ad6..42bed58ef 100644 --- a/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py +++ b/src/struphy/pic/pushing/kernels/bstar_2form/bstar_2form_kernels.py @@ -50,7 +50,7 @@ def bstar_2form( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -59,13 +59,13 @@ def bstar_2form( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n # spline evaluation diff --git a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py index f72014dd2..b6356a1b0 100644 --- a/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py +++ b/src/struphy/pic/pushing/kernels/bstar_parallel_3form/bstar_parallel_3form_kernels.py @@ -45,7 +45,7 @@ def bstar_parallel_3form( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -54,13 +54,13 @@ def bstar_parallel_3form( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n # evaluate Jacobian, result in dfm diff --git a/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py b/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py index 26acb7d20..d622133b3 100644 --- a/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py +++ b/src/struphy/pic/pushing/kernels/driftkinetic_hamiltonian/driftkinetic_hamiltonian_kernels.py @@ -44,7 +44,7 @@ def driftkinetic_hamiltonian( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -53,13 +53,13 @@ def driftkinetic_hamiltonian( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v = alpha[3] * v_k + (1.0 - alpha[3]) * v_n mu = markers[ip, mu_idx] diff --git a/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py b/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py index 1e89b6e36..417d384e7 100644 --- a/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py +++ b/src/struphy/pic/pushing/kernels/grad_driftkinetic_hamiltonian/grad_driftkinetic_hamiltonian_kernels.py @@ -58,7 +58,7 @@ def grad_driftkinetic_hamiltonian( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -67,7 +67,7 @@ def grad_driftkinetic_hamiltonian( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) diff --git a/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py b/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py index a12243267..3b891be49 100644 --- a/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_deterministic_diffusion_stage/push_deterministic_diffusion_stage_kernels.py @@ -49,7 +49,7 @@ def push_deterministic_diffusion_stage( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -121,7 +121,7 @@ def push_deterministic_diffusion_stage( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * a[stage] * k + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu index 978a5f1d5..0f5a1938c 100644 --- a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_cuda.cu @@ -14,7 +14,7 @@ * @param c Butcher nodes c (n_stages entries; unused, as in the pyccel kernel). * @param n_stages Number of stages; a raw device pointer carries no length. * - * Holes (first_init_idx == -1) and boundary particles (last column == -2) are skipped. + * Holes (first_pusher_idx == -1) and boundary particles (last column == -2) are skipped. * The kinetic boundary conditions are applied per marker after the update. */ extern "C" __global__ void push_eta_stage(double dt, int stage, MarkerArgs args_markers, DomainArgs args_domain, @@ -24,12 +24,12 @@ extern "C" __global__ void push_eta_stage(double dt, int stage, MarkerArgs args_ int n_markers = args_markers.n_markers; if (ip >= n_markers) return; - int first_init_idx = args_markers.first_init_idx; + int first_pusher_idx = args_markers.first_pusher_idx; int first_free_idx = args_markers.first_free_idx; double last = stage == n_stages - 1 ? 1. : 0.; // check if marker is a hole or a boundary particle - if (args_markers.markers(ip, first_init_idx) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) + if (args_markers.markers(ip, first_pusher_idx) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) return; double e1 = args_markers.markers(ip, 0); @@ -51,7 +51,7 @@ extern "C" __global__ void push_eta_stage(double dt, int stage, MarkerArgs args_ // accumulation for last stage args_markers.markers(ip, first_free_idx + j) += dt * b[stage] * k[j]; // update positions for intermediate stages or last stage - args_markers.markers(ip, j) = args_markers.markers(ip, first_init_idx + j) + dt * a[stage] * k[j] + + args_markers.markers(ip, j) = args_markers.markers(ip, first_pusher_idx + j) + dt * a[stage] * k[j] + last * args_markers.markers(ip, first_free_idx + j); } diff --git a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py index b037d4746..e5f0245f1 100644 --- a/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_eta_stage/push_eta_stage_kernels.py @@ -43,7 +43,7 @@ def push_eta_stage( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx if stage == n_stages - 1: @@ -57,7 +57,7 @@ def push_eta_stage( # fmt: on for ip in range(n_markers): # check if marker is a hole or a boundary particle - if markers[ip, first_init_idx] == -1.0 or markers[ip, -1] == -2.0: + if markers[ip, first_pusher_idx] == -1.0 or markers[ip, -1] == -2.0: continue e1 = markers[ip, 0] @@ -85,7 +85,7 @@ def push_eta_stage( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * a[stage] * k + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py index 32e0b3d64..4f0a4ac73 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order/push_gc_Bstar_discrete_gradient_1st_order_kernels.py @@ -82,25 +82,25 @@ def push_gc_Bstar_discrete_gradient_1st_order( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) eta_diff[:] = eta_k - eta_n v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v_mid = (v_k + v_n) / 2.0 v_diff = v_k - v_n diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py index fb88521fc..8f296d666 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_1st_order_newton/push_gc_Bstar_discrete_gradient_1st_order_newton_kernels.py @@ -143,23 +143,23 @@ def push_gc_Bstar_discrete_gradient_1st_order_newton( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_k_shifted[:] = eta_k + markers[ip, first_shift_idx : first_shift_idx + 3] eta_diff[:] = eta_k_shifted - eta_n v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v_diff = v_k - v_n mu = markers[ip, mu_idx] diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py index 5fe3aadac..3c6b8d82e 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_discrete_gradient_2nd_order/push_gc_Bstar_discrete_gradient_2nd_order_kernels.py @@ -108,25 +108,25 @@ def push_gc_Bstar_discrete_gradient_2nd_order( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) eta_diff[:] = eta_k - eta_n v_k = markers[ip, 3] - v_n = markers[ip, first_init_idx + 3] + v_n = markers[ip, first_pusher_idx + 3] v_mid = (v_k + v_n) / 2.0 v_diff = v_k - v_n diff --git a/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py index f47a562a9..03ac967b6 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_Bstar_explicit_multistage/push_gc_Bstar_explicit_multistage_kernels.py @@ -80,7 +80,7 @@ def push_gc_Bstar_explicit_multistage( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -93,7 +93,7 @@ def push_gc_Bstar_explicit_multistage( for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue # if stage == 0.: @@ -207,11 +207,13 @@ def push_gc_Bstar_explicit_multistage( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * a[stage] * k + last * markers[ip, first_free_idx : first_free_idx + 3] ) - markers[ip, 3] = markers[ip, first_init_idx + 3] + dt * a[stage] * k_v + last * markers[ip, first_free_idx + 3] + markers[ip, 3] = ( + markers[ip, first_pusher_idx + 3] + dt * a[stage] * k_v + last * markers[ip, first_free_idx + 3] + ) # apply kinetic boundary conditions pusher_utilities_kernels.apply_kinetic_bc_marker(ip, args_markers, args_domain, False) diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py index f383d4329..3d9c1a3e2 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_dg_1st_order_newton/push_gc_bxEstar_dg_1st_order_newton_kernels.py @@ -129,18 +129,18 @@ def push_gc_bxEstar_dg_1st_order_newton( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_k_shifted[:] = eta_k + markers[ip, first_shift_idx : first_shift_idx + 3] eta_diff[:] = eta_k_shifted - eta_n diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py index ca3507264..7c26a86dd 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_1st_order/push_gc_bxEstar_discrete_gradient_1st_order_kernels.py @@ -75,18 +75,18 @@ def push_gc_bxEstar_discrete_gradient_1st_order( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py index 619512ae4..bd0457365 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_discrete_gradient_2nd_order/push_gc_bxEstar_discrete_gradient_2nd_order_kernels.py @@ -91,18 +91,18 @@ def push_gc_bxEstar_discrete_gradient_2nd_order( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx residual_idx = args_markers.residual_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is converged or a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta_mid[:] = (eta_k + eta_n) / 2.0 eta_mid[:] = mod(eta_mid, 1.0) diff --git a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py index c68033f5d..905bfb8e0 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_bxEstar_explicit_multistage/push_gc_bxEstar_explicit_multistage_kernels.py @@ -66,7 +66,7 @@ def push_gc_bxEstar_explicit_multistage( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -79,7 +79,7 @@ def push_gc_bxEstar_explicit_multistage( for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta1 = markers[ip, 0] @@ -171,7 +171,7 @@ def push_gc_bxEstar_explicit_multistage( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * a[stage] * k + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py index 624ab42b5..c10e36e0a 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_Hdiv/push_gc_cc_J2_dg_Hdiv_kernels.py @@ -80,7 +80,7 @@ def push_gc_cc_J2_dg_Hdiv( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): @@ -90,7 +90,7 @@ def push_gc_cc_J2_dg_Hdiv( # marker positions, mid point eta_old[:] = markers[ip, 0:3] - eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_init_idx : first_init_idx + 3]) / 2.0 + eta_mid[:] = (markers[ip, 0:3] + markers[ip, first_pusher_idx : first_pusher_idx + 3]) / 2.0 eta_mid[:] = mod(eta_mid[:], 1.0) v = markers[ip, 3] @@ -204,6 +204,6 @@ def push_gc_cc_J2_dg_Hdiv( e /= abs_b_star_para e /= det_df - markers[ip, 0:3] = markers[ip, first_init_idx : first_init_idx + 3] - dt * e + markers[ip, 0:3] = markers[ip, first_pusher_idx : first_pusher_idx + 3] - dt * e markers[ip, 0:3] *= alpha markers[ip, 0:3] += eta_old * (1.0 - alpha) diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py index 8fcaafd1c..388c069cc 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_dg_init_Hdiv/push_gc_cc_J2_dg_init_Hdiv_kernels.py @@ -67,12 +67,12 @@ def push_gc_cc_J2_dg_init_Hdiv( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta1 = markers[ip, 0] diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py index a3cb0e4e5..d40c24fc7 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_H1vec/push_gc_cc_J2_stage_H1vec_kernels.py @@ -85,7 +85,7 @@ def push_gc_cc_J2_stage_H1vec( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -98,7 +98,7 @@ def push_gc_cc_J2_stage_H1vec( for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta1 = markers[ip, 0] @@ -203,7 +203,7 @@ def push_gc_cc_J2_stage_H1vec( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] - dt * a[stage] * e + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py index ab5e5f89d..c2a82a5f5 100644 --- a/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_gc_cc_J2_stage_Hdiv/push_gc_cc_J2_stage_Hdiv_kernels.py @@ -81,7 +81,7 @@ def push_gc_cc_J2_stage_Hdiv( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -98,7 +98,7 @@ def push_gc_cc_J2_stage_Hdiv( continue # if particle is refilled - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue eta1 = markers[ip, 0] @@ -204,7 +204,7 @@ def push_gc_cc_J2_stage_Hdiv( # update positions for intermediate stages or last stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] - dt * a[stage] * e + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py index 5946bdf4b..e9fe1c0c5 100644 --- a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_H1vec/push_pc_eta_stage_H1vec_kernels.py @@ -66,7 +66,7 @@ def push_pc_eta_stage_H1vec( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -131,7 +131,7 @@ def push_pc_eta_stage_H1vec( # update markers for the next stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * k * a[stage] + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py index 33004306d..55e02b6ff 100644 --- a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hcurl/push_pc_eta_stage_Hcurl_kernels.py @@ -59,7 +59,7 @@ def push_pc_eta_stage_Hcurl( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -127,7 +127,7 @@ def push_pc_eta_stage_Hcurl( # update markers for the next stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * k * a[stage] + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py index 6ced6a2f9..0cae6016c 100644 --- a/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_pc_eta_stage_Hdiv/push_pc_eta_stage_Hdiv_kernels.py @@ -59,7 +59,7 @@ def push_pc_eta_stage_Hdiv( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # get number of stages @@ -128,7 +128,7 @@ def push_pc_eta_stage_Hdiv( # update markers for the next stage markers[ip, 0:3] = ( - markers[ip, first_init_idx : first_init_idx + 3] + markers[ip, first_pusher_idx : first_pusher_idx + 3] + dt * k * a[stage] + last * markers[ip, first_free_idx : first_free_idx + 3] ) diff --git a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu index bc96bfd77..d14f81ec6 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_cuda.cu @@ -28,10 +28,10 @@ extern "C" __global__ void push_vxb_analytic(double dt, int stage, MarkerArgs ar int n_markers = args_markers.n_markers; if (ip >= n_markers) return; - int first_init_idx = args_markers.first_init_idx; + int first_pusher_idx = args_markers.first_pusher_idx; // check if marker is a hole - if (args_markers.markers(ip, first_init_idx) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) + if (args_markers.markers(ip, first_pusher_idx) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) return; double e1 = args_markers.markers(ip, 0); diff --git a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py index 3341e4794..defa85cc3 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_vxb_analytic/push_vxb_analytic_kernels.py @@ -54,7 +54,7 @@ def push_vxb_analytic( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx # fmt: off #$ omp parallel private (ip, e1, e2, e3, v, dfm, det_df, span1, span2, span3, b_form, b_cart, b_abs, b_norm, vpar, vxb_norm, vperp, b_normxvperp) @@ -62,7 +62,7 @@ def push_vxb_analytic( # fmt: on for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0 or markers[ip, -1] == -2.0: + if markers[ip, first_pusher_idx] == -1.0 or markers[ip, -1] == -2.0: continue e1 = markers[ip, 0] diff --git a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu index 1190f3e25..8fbb4a7e2 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu +++ b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_cuda.cu @@ -27,10 +27,10 @@ extern "C" __global__ void push_vxb_implicit(double dt, int stage, MarkerArgs ar int n_markers = args_markers.n_markers; if (ip >= n_markers) return; - int first_init_idx = args_markers.first_init_idx; + int first_pusher_idx = args_markers.first_pusher_idx; // check if marker is a hole - if (args_markers.markers(ip, first_init_idx) == -1.) return; + if (args_markers.markers(ip, first_pusher_idx) == -1.) return; double e1 = args_markers.markers(ip, 0); double e2 = args_markers.markers(ip, 1); diff --git a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py index 505959518..2b1c5facf 100644 --- a/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py +++ b/src/struphy/pic/pushing/kernels/push_vxb_implicit/push_vxb_implicit_kernels.py @@ -65,7 +65,7 @@ def push_vxb_implicit( # get marker arguments markers = args_markers.markers n_markers = args_markers.n_markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx # fmt: off #$ omp parallel firstprivate(b_prod) private (ip, e1, e2, e3, v, dfm, det_df, span1, span2, span3, b_form, b_cart, rhs, lhs, lhs_inv, vec, res) @@ -73,7 +73,7 @@ def push_vxb_implicit( # fmt: on for ip in range(n_markers): # check if marker is a hole - if markers[ip, first_init_idx] == -1.0: + if markers[ip, first_pusher_idx] == -1.0: continue e1 = markers[ip, 0] diff --git a/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py b/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py index 07460ad4e..77cd4aa0a 100644 --- a/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py +++ b/src/struphy/pic/pushing/kernels/unit_b_1form/unit_b_1form_kernels.py @@ -39,7 +39,7 @@ def unit_b_1form( markers = args_markers.markers n_markers = args_markers.n_markers mu_idx = args_markers.mu_idx - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx for ip in range(n_markers): @@ -48,7 +48,7 @@ def unit_b_1form( continue eta_k[:] = markers[ip, 0:3] + markers[ip, first_shift_idx : first_shift_idx + 3] - eta_n[:] = markers[ip, first_init_idx : first_init_idx + 3] + eta_n[:] = markers[ip, first_pusher_idx : first_pusher_idx + 3] eta[:] = alpha[:3] * eta_k + (1.0 - alpha[:3]) * eta_n eta[:] = mod(eta, 1.0) diff --git a/src/struphy/pic/pushing/prepare_push_cuda.cu b/src/struphy/pic/pushing/prepare_push_cuda.cu index 5c1dbd28e..71c02556d 100644 --- a/src/struphy/pic/pushing/prepare_push_cuda.cu +++ b/src/struphy/pic/pushing/prepare_push_cuda.cu @@ -1,28 +1,28 @@ /** * Prepare the marker buffer for a push, as the three column-slice assignments at the start of Pusher._push: * - * markers[:, first_init_idx:first_shift_idx] = markers[:, :n_init] + * markers[:, first_pusher_idx:first_shift_idx] = markers[:, :n_init] * markers[:, first_shift_idx:residual_idx] = 0.0 * markers[:, residual_idx:-2] = 0.0 * - * One thread per entry of the contiguous column range [first_init_idx, n_cols - 2) of every row, so that + * One thread per entry of the contiguous column range [first_pusher_idx, n_cols - 2) of every row, so that * neighbouring threads write neighbouring addresses of the row-major buffer and the whole range is written * in a single pass (instead of one strided pass over the buffer per slice). * * @param markers Marker buffer (n_rows x n_cols, row-major), every row including holes. * @param n_rows Number of rows of the buffer. * @param n_cols Number of columns of the buffer. - * @param first_init_idx First column of the saved initial phase space coordinates. - * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_init_idx + n_init). + * @param first_pusher_idx First column of the saved initial phase space coordinates. + * @param n_init Number of saved coordinates, 3 + vdim (first_shift_idx = first_pusher_idx + n_init). */ -extern "C" __global__ void prepare_push(double* markers, long long n_rows, int n_cols, int first_init_idx, +extern "C" __global__ void prepare_push(double* markers, long long n_rows, int n_cols, int first_pusher_idx, int n_init) { - long long width = n_cols - 2 - first_init_idx; + long long width = n_cols - 2 - first_pusher_idx; long long t = (long long)blockDim.x * blockIdx.x + threadIdx.x; if (t >= n_rows * width) return; long long ip = t / width; - int col = first_init_idx + (int)(t - ip * width); + int col = first_pusher_idx + (int)(t - ip * width); double* row = markers + ip * n_cols; - row[col] = col < first_init_idx + n_init ? row[col - first_init_idx] : 0.0; + row[col] = col < first_pusher_idx + n_init ? row[col - first_pusher_idx] : 0.0; } diff --git a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh index a40dc406a..48b06ef19 100644 --- a/src/struphy/pic/pushing/pusher_utilities_kernels.cuh +++ b/src/struphy/pic/pushing/pusher_utilities_kernels.cuh @@ -38,7 +38,7 @@ __device__ inline void reflect_velocity(long long ip, Array2D markers, * @param newton Accumulate periodic shifts if true; overwrite them otherwise. * * Call after updating the position. Boundary type 0 wraps periodically, - * type 1 mirrors position and logical velocity and sets first_init_idx to -1, + * type 1 mirrors position and logical velocity and sets first_pusher_idx to -1, * type 2 clears all columns except the ID to -1, and type 3 is handled on * the host. The marker view carries shape and strides; j is a CUDA-only * column index replacing Pyccel array slices. @@ -46,7 +46,7 @@ __device__ inline void reflect_velocity(long long ip, Array2D markers, __device__ inline void apply_kinetic_bc_marker(int ip, const MarkerArgs& args_markers, const DomainArgs& args_domain, bool newton) { const long long* bc_type = args_markers.bc_type; - int first_init_idx = args_markers.first_init_idx; + int first_pusher_idx = args_markers.first_pusher_idx; int first_shift_idx = args_markers.first_shift_idx; // Remove markers before applying periodic or reflecting boundaries. @@ -92,7 +92,7 @@ __device__ inline void apply_kinetic_bc_marker(int ip, const MarkerArgs& args_ma } } if (n_reflected == 0) return; - args_markers.markers(ip, first_init_idx) = -1.; + args_markers.markers(ip, first_pusher_idx) = -1.; for (int axis = 0; axis < 3; ++axis) { if (reflected[axis]) reflect_velocity(ip, args_markers.markers, args_domain, axis); } diff --git a/src/struphy/pic/pushing/pusher_utilities_kernels.py b/src/struphy/pic/pushing/pusher_utilities_kernels.py index 78b9c7d70..b21c1bfb0 100644 --- a/src/struphy/pic/pushing/pusher_utilities_kernels.py +++ b/src/struphy/pic/pushing/pusher_utilities_kernels.py @@ -27,7 +27,7 @@ def apply_kinetic_bc_marker( * 0 (periodic): wrap :math:`\eta_i` into :math:`[0, 1)` and store the shift :math:`\pm 1` at ``first_shift_idx + axis`` (set for ``newton=False``, accumulated for ``newton=True``). * 1 (reflect): mirror :math:`\eta_i` at the boundary, flip the logical velocity component - and mark the marker as done for multi-stage/iterative pushers (``first_init_idx = -1``). + and mark the marker as done for multi-stage/iterative pushers (``first_pusher_idx = -1``). * 2 (remove): turn the marker into a hole (all columns except the ID set to -1). * 3 (refill): skipped here, handled by :meth:`~struphy.pic.base.Particles.apply_kinetic_bc`. @@ -48,7 +48,7 @@ def apply_kinetic_bc_marker( markers = args_markers.markers bc_type = args_markers.bc_type - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_shift_idx = args_markers.first_shift_idx # remove @@ -94,7 +94,7 @@ def apply_kinetic_bc_marker( if n_reflected == 0: return - markers[ip, first_init_idx] = -1.0 + markers[ip, first_pusher_idx] = -1.0 # reflect velocities (Jacobian at the reflected position) dfm = empty((3, 3), dtype=float) diff --git a/src/struphy/propagators/current_coupling_5d_gradb.py b/src/struphy/propagators/current_coupling_5d_gradb.py index 9fac1ec42..a596074df 100644 --- a/src/struphy/propagators/current_coupling_5d_gradb.py +++ b/src/struphy/propagators/current_coupling_5d_gradb.py @@ -459,14 +459,14 @@ def __call__(self, dt): holes = particles.holes args_markers = particles.args_markers markers = args_markers.markers - first_init_idx = args_markers.first_init_idx + first_pusher_idx = args_markers.first_pusher_idx first_free_idx = args_markers.first_free_idx # clear buffer - markers[:, first_init_idx:-2] = 0.0 + markers[:, first_pusher_idx:-2] = 0.0 # save old marker positions - markers[:, first_init_idx : first_init_idx + 3] = markers[:, :3] + markers[:, first_pusher_idx : first_pusher_idx + 3] = markers[:, :3] # sum up total magnetic field b_full1 = b_eq + b_tilde (in-place) b_full = self._b2.copy(out=self._b_full) @@ -530,7 +530,7 @@ def __call__(self, dt): diffs = self.update_feec_variables(u=u_new) # clear the buffer - markers[:, first_init_idx:-2] = 0.0 + markers[:, first_pusher_idx:-2] = 0.0 # update_weights if self.variables.energetic_ions.species.weights_params.control_variate: @@ -661,7 +661,7 @@ def __call__(self, dt): sum_u_diff_loc = xp.sum((u_diff.toarray() ** 2)) sum_H_diff_loc = xp.sum( - (markers[~holes, :3] - markers[~holes, first_init_idx : first_init_idx + 3]) ** 2, + (markers[~holes, :3] - markers[~holes, first_pusher_idx : first_pusher_idx + 3]) ** 2, ) buffer_array = xp.array([sum_u_diff_loc]) @@ -853,7 +853,7 @@ def __call__(self, dt): diffs = self.update_feec_variables(u=u_new) # clear the buffer - markers[:, first_init_idx:-2] = 0.0 + markers[:, first_pusher_idx:-2] = 0.0 # update_weights if self.variables.energetic_ions.species.weights_params.control_variate: From 225bcae42976a7690b7c9ebc97feaaa6b34ab769 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 20:49:16 +0200 Subject: [PATCH 38/50] 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 39/50] 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 From 77a25b5613dd93b99af23651552b83fa59ca0dc7 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 20:49:16 +0200 Subject: [PATCH 40/50] Marker exchange between ranks on the device (CUDA-aware MPI) On the CuPy backend, Particles._sendrecv_markers hands the marker rows to Isend/Irecv as device buffers through cunumpy.mpi.mpi_buffer (directly after synchronize_for_mpi with CUDA-aware MPI, staged through host memory otherwise). Rows from each rank arrive in a contiguous slice of one device receive buffer, and one device scatter fills the holes. The per-rank counts are small NumPy arrays on every backend (they are host integers already), and a uniform alpha is a Python scalar, so the exchange makes no host/device copy and no implicit host read of a device array. The NumPy path is unchanged. Adds pic/tests/test_sorting_device.py (fake CuPy under mpirun, and a GPU variant) and a CI step running it on the fake CuPy. Solves #698, part of #650. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/test-PR-unit-mpi.yml | 6 + CUDA_STRATEGY.md | 20 +- src/struphy/pic/base.py | 129 ++++++++++++- src/struphy/pic/tests/test_sorting_device.py | 190 +++++++++++++++++++ 4 files changed, 338 insertions(+), 7 deletions(-) create mode 100644 src/struphy/pic/tests/test_sorting_device.py diff --git a/.github/workflows/test-PR-unit-mpi.yml b/.github/workflows/test-PR-unit-mpi.yml index e3a13dbc5..7bb53fbe2 100644 --- a/.github/workflows/test-PR-unit-mpi.yml +++ b/.github/workflows/test-PR-unit-mpi.yml @@ -76,3 +76,9 @@ jobs: source /struphy_fortran_/env_fortran_/bin/activate cd ${{ env.STRUPHY_PATH }} mpirun --oversubscribe -n ${{ matrix.n-procs }} pytest -v --with-mpi -m mpi_pic pic/tests + + - name: Run the marker exchange on the fake CuPy + run: | + source /struphy_fortran_/env_fortran_/bin/activate + cd ${{ env.STRUPHY_PATH }} + mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi pic/tests/test_sorting_device.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 6d2c60e72..68035f68c 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -365,7 +365,7 @@ matrix data). - **Kernel launch configuration.** One thread per marker with a fixed block size for now. Kernels over grid points (accumulation, FEEC) need their own launch sizes; `CudaKernel` needs a `grid`/`block` override (PR 12+), or cunumpy's. - **Accumulation strategy.** Atomics vs. sort-then-reduce; see PR 12+. Decided by measurement on the first accumulation kernel. - **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. `xp.mpi_is_cuda_aware()` detects it; the exchange in `Particles.mpi_sort_markers` has to be checked for host staging buffers. +- **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 `Derham` with `polar_splines=True` raises on CuPy (`PolarExtractionBlocksC1` builds SciPy sparse matrices, and the @@ -644,3 +644,21 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a The GPU tests have not run on an H100; the CPU emulation and the CPU regression tests are the gate. - **Parity cases.** The spline mappings are added to the geometry kernels' entries of `pic/tests/cuda_parity_cases.py` (the per-folder `_test_args.py` modules were replaced by that module in PR 16). + +## Marker exchange implementation notes (#698) + +- `Particles._sendrecv_markers` on CuPy (`_sendrecv_markers_device`): the rows to send (gathered on the device per + destination) and one device receive buffer go to `Isend`/`Irecv` through `xp.mpi.mpi_buffer`, i.e. directly after + `synchronize_for_mpi` with CUDA-aware MPI, staged through (pinned) host memory otherwise (with a warning, once). + Rows from rank `i` arrive in a contiguous slice of the receive buffer; one device scatter fills the holes. + Zero-count messages are skipped (both sides know the counts). The NumPy path is unchanged. +- Whether MPI is CUDA-aware: the answer recorded by `xp.mpi.mpi_is_cuda_aware()`/`set_mpi_cuda_aware()`, else a + collective probe of `mpi_comm` at the first exchange (`mpi_sort_markers` is collective). +- Counts: `send_info`/`recv_info` are small NumPy arrays on every backend. The send counts are sizes of `nonzero` + results, which are host integers already (CuPy synchronizes to size them), so the `Alltoall` of the counts needs no + copy. Uniform `alpha` is a Python scalar (no host-to-device copy of a 3-vector). +- Tests: `pic/tests/test_sorting_device.py` sorts the same markers on NumPy and on CuPy over all ranks and compares + the marker arrays; with CUDA-aware MPI no transfer (`count_transfers`) and no implicit host read of a device array + during the exchange, staged: one `to_host` per non-empty send and one `to_device`. Without a GPU it runs on the fake + CuPy (`mpirun -n N env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi pic/tests/test_sorting_device.py`, in the PIC MPI + workflow); the GPU variant has not run (no GPU here). diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 99db7dbfd..cae11c476 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2,6 +2,7 @@ import os import warnings from abc import ABCMeta, abstractmethod +from contextlib import ExitStack from typing import TYPE_CHECKING import cunumpy as xp @@ -1757,6 +1758,12 @@ def mpi_sort_markers( remove_ghost : bool Remove ghost particles before send. + + Notes + ----- + On the CuPy backend, the marker rows are exchanged as device buffers with CUDA-aware MPI + (staged through host memory otherwise), see :meth:`_sendrecv_markers_device`. The + per-rank counts are small host arrays on every backend (:meth:`_sendrecv_all_to_all`). """ if remove_ghost: self._remove_ghost_particles() @@ -4363,9 +4370,19 @@ def _sendrecv_determine_mtbs( """ # position that determines the sorting (including periodic shift of boundary conditions) if not isinstance(alpha, xp.ndarray): - alpha = xp.array(alpha, dtype=float) - assert alpha.size == 3 - assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) + # check a tuple/list on the host: on CuPy, asserting on a device comparison reads it back + alpha_host = np.array(alpha, dtype=float) + assert alpha_host.size == 3 + assert np.all(alpha_host >= 0.0) and np.all(alpha_host <= 1.0) + if alpha_host[0] == alpha_host[1] == alpha_host[2]: + # the common case (alpha = 1 or alpha = 0.5 in every direction): a Python scalar, + # which needs no host-to-device copy on CuPy and gives the same values on NumPy + alpha = float(alpha_host[0]) + else: + alpha = xp.asarray(alpha_host) + else: + assert alpha.size == 3 + assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) bi = self.first_pusher_idx self._sorting_etas[:] = ( alpha * (self.markers[:, :3] + self.markers[:, bi + 3 + self.vdim : bi + 3 + self.vdim + 3]) @@ -4465,8 +4482,11 @@ def _sendrecv_get_destinations(self, send_inds): Amount of particles sent to i-th process. """ - # One entry for each process - send_info = xp.zeros(self.mpi_size, dtype=int) + # One entry for each process. The counts are host integers on every backend: the + # size of a nonzero() result is known on the host (CuPy synchronizes to size it), so + # a small NumPy array holds them without any host/device copy, and the Alltoall of + # the counts (_sendrecv_all_to_all) passes host buffers to MPI. + send_info = np.zeros(self.mpi_size, dtype=int) # Gathered once and reused for every rank below, instead of re-gathering # self.markers[send_inds] and self._sorting_etas[send_inds] fresh on every @@ -4530,9 +4550,16 @@ def _sendrecv_all_to_all(self, send_info): ------- recv_info : array[int] Amount of marticles to be received from i-th process. + + Notes + ----- + ``send_info`` and ``recv_info`` are small NumPy arrays (one integer per rank) on every + backend: the counts size the receive buffers and the holes to fill, which Python needs + on the host anyway. On CuPy, only the marker rows themselves are device buffers + (:meth:`_sendrecv_markers`). """ - recv_info = xp.zeros(self.mpi_size, dtype=int) + recv_info = np.zeros(self.mpi_size, dtype=int) self.mpi_comm.Alltoall(send_info, recv_info) @@ -4551,6 +4578,10 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): Indices of empty rows in markers after send. """ + if xp.get_array_backend(self._markers) == "cupy": + self._sendrecv_markers_device(recv_info, hole_inds_after_send) + return + # i-th entry holds the number (not the index) of the first hole to be filled by data from process i first_hole = xp.cumsum(recv_info) - recv_info @@ -4594,6 +4625,92 @@ def _sendrecv_markers(self, recv_info, hole_inds_after_send): for req in send_reqs: req.Wait() + def _mpi_cuda_aware(self) -> bool: + """Whether MPI takes device (CuPy) buffers, for :meth:`_sendrecv_markers_device`. + + Uses the answer recorded by :func:`cunumpy.mpi.mpi_is_cuda_aware` or + :func:`cunumpy.mpi.set_mpi_cuda_aware` (e.g. by feectools or at startup). If nothing + was recorded, probes :attr:`mpi_comm` once; the probe is collective, which is safe + here because :meth:`mpi_sort_markers` is called by every rank of :attr:`mpi_comm`. + """ + cuda_aware = xp.mpi.get_mpi_cuda_aware() + if cuda_aware is None: + cuda_aware = xp.mpi.mpi_is_cuda_aware(self.mpi_comm) + if not cuda_aware and not Particles._warned_mpi_staging: + Particles._warned_mpi_staging = True + logger.warning( + "MPI is not CUDA-aware: marker rows exchanged between ranks are staged through host memory.", + ) + return cuda_aware + + _warned_mpi_staging = False + + def _sendrecv_markers_device(self, recv_info, hole_inds_after_send): + """:meth:`_sendrecv_markers` on the CuPy backend: marker rows go to MPI as device buffers. + + With CUDA-aware MPI (:meth:`_mpi_cuda_aware`), the rows to send (:attr:`_send_list`, + gathered on the device by :meth:`_sendrecv_get_destinations`) and one device receive + buffer are handed to ``Isend``/``Irecv`` directly, after + :func:`cunumpy.mpi.synchronize_for_mpi` (inside :func:`cunumpy.mpi.mpi_buffer`), so MPI + reads what the kernels wrote. The rows from rank ``i`` arrive in a contiguous slice of + the receive buffer, and one device scatter puts all of them into the holes. Without + CUDA-aware MPI, :func:`~cunumpy.mpi.mpi_buffer` stages the same buffers through host + memory (counted by :func:`cunumpy.profiling.count_transfers`). + + The bookkeeping stays on the host without host/device copies: the counts are host + integers (``recv_info`` from the Alltoall, the send counts are the row counts of the + send buffers), and so is the number of holes. Ranks exchange no message when the + count is zero (both sides know the count, so they skip consistently). + + Parameters + ---------- + recv_info : array[int] + Amount of markers to be received from i-th process (NumPy array). + + hole_inds_after_send : array[int] + Indices of empty rows in markers after send (device array). + """ + recv_info = np.asarray(recv_info) + n_recv = int(recv_info.sum()) + if hole_inds_after_send.size < n_recv: + warnings.warn( + f'Strong load imbalance detected: \ +number of holes ({hole_inds_after_send.size}) on rank {self.mpi_rank} \ +is smaller than number of incoming particles ({n_recv}). \ +Increasing the value of "bufsize" in the markers parameters for the next run.', + ) + self.mpi_comm.Abort() + + cuda_aware = self._mpi_cuda_aware() + + # rows from rank i land in recvbuf[first_row[i] : first_row[i] + recv_info[i]] + first_row = np.cumsum(recv_info) - recv_info + recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=float) + + with ExitStack() as buffers: + if n_recv > 0: + recv_mpi = buffers.enter_context( + xp.mpi.mpi_buffer(recvbuf, send=False, recv=True, cuda_aware=cuda_aware), + ) + reqs = [] + for i in range(self.mpi_size): + if i == self.mpi_rank: + continue + if recv_info[i] > 0: + rows = recv_mpi[first_row[i] : first_row[i] + recv_info[i]] + reqs.append(self.mpi_comm.Irecv(rows, source=i, tag=i)) + data = self._send_list[i] + if data.shape[0] > 0: + send_mpi = buffers.enter_context(xp.mpi.mpi_buffer(data, cuda_aware=cuda_aware)) + reqs.append(self.mpi_comm.Isend(send_mpi, dest=i, tag=self.mpi_rank)) + # wait inside the block: the buffers must stay valid until MPI is done, and staged + # receive rows are copied to the device when the block ends + for req in reqs: + req.Wait() + + if n_recv > 0: + self._markers[hole_inds_after_send[:n_recv]] = recvbuf + class Tesselation: """ diff --git a/src/struphy/pic/tests/test_sorting_device.py b/src/struphy/pic/tests/test_sorting_device.py new file mode 100644 index 000000000..55c5d09f5 --- /dev/null +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -0,0 +1,190 @@ +"""Marker exchange between ranks on the CuPy backend (:meth:`Particles.mpi_sort_markers`). + +The marker rows are exchanged as device buffers with CUDA-aware MPI, and staged through +host memory otherwise. Without a GPU, the tests run on cunumpy's fake CuPy, which must be +installed before Python imports cunumpy, so they are skipped unless the fake is active:: + + mpirun -n 2 env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi src/struphy/pic/tests/test_sorting_device.py + +Fake CuPy arrays live in host memory, so a non-CUDA-aware MPI library takes them, and the +CUDA-aware path (the probe of :func:`cunumpy.mpi.mpi_is_cuda_aware` succeeds) is what runs by +default. The staged path is forced with :func:`cunumpy.mpi.set_mpi_cuda_aware`. +""" + +import sys +from contextlib import contextmanager + +import cunumpy as xp +import numpy as np +import pytest +from cunumpy.kernel_testing import requires_cupy +from maybempi import MPI + +from struphy import BoundaryParameters, LoadingParameters +from struphy.pic.particles import Particles6D + + +def _fake_cupy_active() -> bool: + return bool(getattr(sys.modules.get("cupy"), "__cunumpy_fake__", False)) + + +requires_fake_cupy = pytest.mark.skipif( + not _fake_cupy_active(), + reason="needs the fake CuPy: run under `mpirun -n N env CUNUMPY_FAKE_CUPY=1 pytest ...`", +) + + +@contextmanager +def _count_host_reads(): + """Count reads of fake CuPy arrays by the host (``int()``, ``bool()``, iteration, ``.get()``, ...). + + :func:`cunumpy.profiling.count_transfers` sees the copies made through cunumpy; this also sees + the implicit ones (a device scalar in an ``if``, ``list(device_array)``), which synchronize + and copy on a GPU. + """ + import cupy + + reads = [] + names = ("__int__", "__float__", "__bool__", "__index__", "__iter__", "get") + originals = {name: getattr(cupy.ndarray, name) for name in names} + original_getattr = cupy.ndarray.__getattr__ + + def counting(name, method): + def wrapper(self, *args, **kwargs): + reads.append(name) + return method(self, *args, **kwargs) + + return wrapper + + def counting_getattr(self, name): + if name in ("item", "tolist"): + reads.append(name) + return original_getattr(self, name) + + for name, method in originals.items(): + setattr(cupy.ndarray, name, counting(name, method)) + cupy.ndarray.__getattr__ = counting_getattr + try: + yield reads + finally: + for name, method in originals.items(): + setattr(cupy.ndarray, name, method) + cupy.ndarray.__getattr__ = original_getattr + + +def _make_particles(comm, bc): + return Particles6D( + comm_world=comm, + loading_params=LoadingParameters(Np=2000, seed=1234), + boundary_params=BoundaryParameters(bc=(bc, bc, bc)), + ) + + +def _shift_positions(markers: np.ndarray, valid: np.ndarray, step: int) -> np.ndarray: + """New positions in [0, 1) for the valid rows, far enough away that most markers change rank.""" + new = markers.copy() + shift = np.array([0.37, 0.21, 0.53]) * (step + 1) + new[valid, :3] = (markers[valid, :3] + shift) % 1.0 + return new + + +def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", count_reads: bool = True): + """Sort the same markers on NumPy and CuPy across all ranks and compare the marker arrays. + + ``cuda_aware=None`` uses the probe (:func:`cunumpy.mpi.mpi_is_cuda_aware`); ``True``/``False`` + force the device-buffer or the host-staged exchange. + """ + comm = MPI.COMM_WORLD + rank, size = comm.Get_rank(), comm.Get_size() + + ref = _make_particles(comm, bc) + ref.draw_markers(sort=False) + ref.mpi_sort_markers() + + # the reference is sorted on NumPy, the device markers inside `use_backend("cupy")` blocks + on_device = lambda: xp.use_backend("cupy") # noqa: E731 + + previous = xp.mpi.get_mpi_cuda_aware() + try: + with on_device(): + if cuda_aware is not None: + xp.mpi.set_mpi_cuda_aware(cuda_aware) + else: + xp.mpi.set_mpi_cuda_aware(None) + cuda_aware = xp.mpi.mpi_is_cuda_aware(comm) + + dev = _make_particles(comm, bc) + assert xp.get_array_backend(dev.markers) == "cupy" + assert dev.markers.shape == ref.markers.shape + + for step in range(2): + new = _shift_positions(ref.markers, ref.valid_mks, step) + + # number of rows this rank has to send, from the host reference + etas = new[ref.valid_mks, :3] + left, right = ref.domain_array[rank, 0::3], ref.domain_array[rank, 1::3] + right = np.where(right == 1.0, np.nextafter(right, 2.0), right) + n_send = int(np.count_nonzero(~np.all((etas >= left) & (etas < right), axis=1))) + n_send_total = comm.allreduce(n_send) + assert size == 1 or n_send_total > 0 + + ref.markers[:] = new + ref.update_holes() + n_before = ref.n_mks_loc + ref.mpi_sort_markers(do_test=True) + n_recv = ref.n_mks_loc - (n_before - n_send) + + with on_device(): + dev.markers[:] = xp.to_cupy(new) + dev.update_holes() + with xp.profiling.count_transfers() as transfers, _count_host_reads() as reads: + dev.mpi_sort_markers() + + if cuda_aware: + assert transfers.total == 0, transfers.report() + else: + # one staged copy per non-empty send buffer, one copy back of the receive buffer + n_dest = sum(1 for i in range(size) if i != rank and dev._send_list[i].shape[0] > 0) + assert transfers.to_host == n_dest, transfers.report() + assert transfers.to_device == (1 if n_recv > 0 else 0), transfers.report() + sent_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_host") + assert sent_bytes == n_send * ref.markers.shape[1] * 8 + + if count_reads and cuda_aware: + # (the staged exchange reads the send buffers with .get(), counted above) + assert reads == [], f"host reads of device arrays during the exchange: {reads}" + + with on_device(): + assert np.array_equal(xp.to_numpy(dev.markers), ref.markers) + assert np.array_equal(xp.to_numpy(dev.valid_mks), ref.valid_mks) + n_dev = dev.n_mks_loc + assert n_dev == ref.n_mks_loc + + with on_device(): + # all markers are on the right rank now: nothing moves + dev.mpi_sort_markers(do_test=True) + assert np.array_equal(xp.to_numpy(dev.markers), ref.markers) + finally: + xp.mpi.set_mpi_cuda_aware(previous) + + +@requires_fake_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("cuda_aware", [None, True, False]) +@pytest.mark.parametrize("bc", ["periodic", "remove"]) +def test_mpi_sort_markers_fake_cupy(cuda_aware, bc): + """On the fake CuPy: the exchange gives the NumPy result, without host/device copies when MPI is CUDA-aware.""" + check_mpi_sort_markers_device(cuda_aware, bc=bc) + + +@requires_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("bc", ["periodic", "remove"]) +def test_mpi_sort_markers_gpu(bc): + """On a GPU (``mpirun -n 2``, one GPU per rank): the exchange gives the NumPy result; with + CUDA-aware MPI without host/device copies.""" + check_mpi_sort_markers_device(None, bc=bc, count_reads=False) + + +if __name__ == "__main__": + check_mpi_sort_markers_device(None) From 0ea9085afb0ee05ad4ecaa5e3da9b9222d8c7759 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 9 Oct 2026 16:21:37 +0200 Subject: [PATCH 41/50] Use xp.mpi.exchange --- pyproject.toml | 2 +- src/struphy/pic/base.py | 85 +++++++++----------- src/struphy/pic/tests/test_sorting_device.py | 7 +- 3 files changed, 42 insertions(+), 52 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c75808225..fdeda2015 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 @ git+https://github.com/max-models/cunumpy.git@26a640b1d47d5677663785e49b1c44d343540df1", "pyccel>=2.2.0, <=2.2.3", "maybempi>=0.1.2", "feectools>=0.8.0, <=0.8.0", diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index cae11c476..5e1d217b2 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -2,7 +2,6 @@ import os import warnings from abc import ABCMeta, abstractmethod -from contextlib import ExitStack from typing import TYPE_CHECKING import cunumpy as xp @@ -4646,67 +4645,55 @@ def _mpi_cuda_aware(self) -> bool: _warned_mpi_staging = False def _sendrecv_markers_device(self, recv_info, hole_inds_after_send): - """:meth:`_sendrecv_markers` on the CuPy backend: marker rows go to MPI as device buffers. - - With CUDA-aware MPI (:meth:`_mpi_cuda_aware`), the rows to send (:attr:`_send_list`, - gathered on the device by :meth:`_sendrecv_get_destinations`) and one device receive - buffer are handed to ``Isend``/``Irecv`` directly, after - :func:`cunumpy.mpi.synchronize_for_mpi` (inside :func:`cunumpy.mpi.mpi_buffer`), so MPI - reads what the kernels wrote. The rows from rank ``i`` arrive in a contiguous slice of - the receive buffer, and one device scatter puts all of them into the holes. Without - CUDA-aware MPI, :func:`~cunumpy.mpi.mpi_buffer` stages the same buffers through host - memory (counted by :func:`cunumpy.profiling.count_transfers`). - - The bookkeeping stays on the host without host/device copies: the counts are host - integers (``recv_info`` from the Alltoall, the send counts are the row counts of the - send buffers), and so is the number of holes. Ranks exchange no message when the - count is zero (both sides know the count, so they skip consistently). + """Exchange marker rows and scatter incoming rows into empty slots. + + :func:`cunumpy.mpi.exchange` handles GPU synchronization, optional host + staging, and completion of all MPI transfers. Parameters ---------- - recv_info : array[int] - Amount of markers to be received from i-th process (NumPy array). - - hole_inds_after_send : array[int] - Indices of empty rows in markers after send (device array). + recv_info : array[int] + Number of markers to receive from each process (NumPy array). + hole_inds_after_send : array[int] + Indices of empty rows in markers after send (device array). """ recv_info = np.asarray(recv_info) n_recv = int(recv_info.sum()) + if hole_inds_after_send.size < n_recv: warnings.warn( - f'Strong load imbalance detected: \ -number of holes ({hole_inds_after_send.size}) on rank {self.mpi_rank} \ -is smaller than number of incoming particles ({n_recv}). \ -Increasing the value of "bufsize" in the markers parameters for the next run.', + f"Strong load imbalance detected: " + f"number of holes ({hole_inds_after_send.size}) " + f"on rank {self.mpi_rank} is smaller than number of " + f"incoming particles ({n_recv}). " + 'Increase "bufsize" in the markers parameters for the next run.', ) self.mpi_comm.Abort() - cuda_aware = self._mpi_cuda_aware() - - # rows from rank i land in recvbuf[first_row[i] : first_row[i] + recv_info[i]] + recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=self.markers.dtype) first_row = np.cumsum(recv_info) - recv_info - recvbuf = xp.empty((n_recv, self.markers.shape[1]), dtype=float) - with ExitStack() as buffers: - if n_recv > 0: - recv_mpi = buffers.enter_context( - xp.mpi.mpi_buffer(recvbuf, send=False, recv=True, cuda_aware=cuda_aware), - ) - reqs = [] - for i in range(self.mpi_size): - if i == self.mpi_rank: - continue - if recv_info[i] > 0: - rows = recv_mpi[first_row[i] : first_row[i] + recv_info[i]] - reqs.append(self.mpi_comm.Irecv(rows, source=i, tag=i)) - data = self._send_list[i] - if data.shape[0] > 0: - send_mpi = buffers.enter_context(xp.mpi.mpi_buffer(data, cuda_aware=cuda_aware)) - reqs.append(self.mpi_comm.Isend(send_mpi, dest=i, tag=self.mpi_rank)) - # wait inside the block: the buffers must stay valid until MPI is done, and staged - # receive rows are copied to the device when the block ends - for req in reqs: - req.Wait() + sends = [] + receives = [] + for rank in range(self.mpi_size): + if rank == self.mpi_rank: + continue + + if recv_info[rank] > 0: + start = int(first_row[rank]) + stop = start + int(recv_info[rank]) + receives.append((recvbuf[start:stop], rank, rank)) + + data = self._send_list[rank] + if data.shape[0] > 0: + sends.append((data, rank, self.mpi_rank)) + + xp.mpi.exchange( + self.mpi_comm, + sends=sends, + receives=receives, + cuda_aware=self._mpi_cuda_aware(), + ) if n_recv > 0: self._markers[hole_inds_after_send[:n_recv]] = recvbuf diff --git a/src/struphy/pic/tests/test_sorting_device.py b/src/struphy/pic/tests/test_sorting_device.py index 55c5d09f5..dc8fa9462 100644 --- a/src/struphy/pic/tests/test_sorting_device.py +++ b/src/struphy/pic/tests/test_sorting_device.py @@ -143,12 +143,15 @@ def check_mpi_sort_markers_device(cuda_aware: bool | None, bc: str = "periodic", if cuda_aware: assert transfers.total == 0, transfers.report() else: - # one staged copy per non-empty send buffer, one copy back of the receive buffer + # exchange stages each non-empty send and receive buffer separately n_dest = sum(1 for i in range(size) if i != rank and dev._send_list[i].shape[0] > 0) + n_sources = sum(1 for i in range(size) if i != rank and ref._recvbufs[i].shape[0] > 0) assert transfers.to_host == n_dest, transfers.report() - assert transfers.to_device == (1 if n_recv > 0 else 0), transfers.report() + assert transfers.to_device == n_sources, transfers.report() sent_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_host") assert sent_bytes == n_send * ref.markers.shape[1] * 8 + received_bytes = sum(e.nbytes for e in transfers.events if e.kind == "to_device") + assert received_bytes == n_recv * ref.markers.shape[1] * 8 if count_reads and cuda_aware: # (the staged exchange reads the send buffers with .get(), counted above) From 569904f768229a45a6ca0c39f58cfc4fdff5984e Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 19:12:45 +0200 Subject: [PATCH 42/50] CUDA version of the linear_vlasov_ampere accumulation - linear_vlasov_ampere_cuda.cu: same arguments as pyccel, one thread per marker row, matrix blocks as Array6D views (cunumpy 0.6.1) - device fill_mat, fill_mat_vec (filler_kernels.cuh), m_v_fill_b_v1_symm (particle_to_mat_kernels.cuh) and outer (linalg_kernels.cuh), with atomic adds - parity cases (Cuboid, Colella, HollowTorus, 3d spline) and a device-helper test of m_v_fill_b_v1_symm, on a GPU and by CPU emulation - test_accum_matrix_cupy.py: the matrix path of Accumulator on CuPy vs NumPy (device data, ghost regions, transposed blocks, Schur operator and solve), without a GPU on the fake CuPy with emulated kernel launches (emulated_launches() in pic/tests/cuda_emulation.py) - CUDA_STRATEGY.md: porting order and implementation notes Co-Authored-By: Claude Opus 5.5 --- CUDA_STRATEGY.md | 59 ++++- src/struphy/linear_algebra/linalg_kernels.cuh | 11 + .../pic/accumulation/filler_kernels.cuh | 125 ++++++++++- .../linear_vlasov_ampere_cuda.cu | 90 ++++++++ .../accumulation/particle_to_mat_kernels.cuh | 69 ++++++ src/struphy/pic/tests/cuda_emulation.py | 66 +++++- src/struphy/pic/tests/cuda_parity_cases.py | 17 ++ src/struphy/pic/tests/kernel_test_args.py | 27 ++- .../pic/tests/test_accum_matrix_cupy.py | 209 ++++++++++++++++++ src/struphy/pic/tests/test_cuda_emulation.py | 12 + src/struphy/pic/tests/test_cuda_parity.py | 4 +- src/struphy/pic/tests/test_device_helpers.py | 61 +++++ 12 files changed, 730 insertions(+), 20 deletions(-) create mode 100644 src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu create mode 100644 src/struphy/pic/tests/test_accum_matrix_cupy.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 68035f68c..5f8635336 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -26,7 +26,8 @@ Progress is tracked in struphy-hub/struphy#650. - [ ] **PR 17: struphy on the feectools CUDA stack** — the `feectools` submodule points at the top of the feectools CUDA stack (`cuda-4-device-kernels`, [feectools#88](https://github.com/struphy-hub/feectools/pull/88)) instead of `devel-tiny`, so the struphy CUDA PRs run against feectools with device stencil operations, MPI with device buffers and one GPU per rank. Moves along with the stack; before the struphy stack is merged into `devel`, the feectools stack is merged into `devel-tiny` and the submodule points there again (see [feectools](#feectools)). - [x] **PR 18: Geometry evaluation on the GPU for all analytic mappings** (#681, open) — CUDA versions of the four geometry entry kernels (`kernel_evaluate_pic`, `kernel_evaluate`, `kernel_pullpush_pic`, `kernel_pullpush` in `geometry/kernels/`), built on device versions of the whole metric chain (`f`, `df`, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`/`push`/`tran`) for every analytic mapping (`kind_map` 10–12, 20–22, 30–32). Restores CuPy particle runs (weight initialization evaluates `jacobian_det`), and removes the Cuboid-only checks in `Pusher`, the accumulators and `reflect`. Parity arguments cover every analytic mapping (see [PR 18](#pr-18-geometry-evaluation-for-all-analytic-mappings) and the [PR 18 implementation notes](#pr-18-implementation-notes)). - [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines and `EQDSKequilibrium` stay host-only (see [PR 19](#pr-19-spline-mappings) and the [PR 19 implementation notes](#pr-19-implementation-notes)). -- [ ] **Next (order to be confirmed)**: 6D array views in cunumpy and the blocked matrix accumulations `vlasov_maxwell`, `linear_vlasov_ampere` (steps 1–2 of the [porting order](#porting-order)); one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, including when to compile the kernels, which needs the [feectools stack](#feectools) merged and released first; a decision on hand-written CUDA vs. code generation before the guiding-center kernels (step 3). +- [x] **`linear_vlasov_ampere` on the GPU** (part of #688) — the first matrix accumulation: CUDA version of `linear_vlasov_ampere` (`EfieldWeightsCoupling`) on cunumpy's `Array6D` views (cunumpy 0.6.1), device versions of `fill_mat`, `fill_mat_vec`, `m_v_fill_b_v1_symm` and `outer`, and the matrix path of `Accumulator` checked on the CuPy backend (device data, ghost regions, transposed blocks, Schur operator) without a GPU through cunumpy's fake CuPy with emulated kernel launches (see the [implementation notes](#linear_vlasov_ampere-implementation-notes)). +- [ ] **Next (order to be confirmed)**: the blocked matrix accumulation `vlasov_maxwell` (step 1 of the [porting order](#porting-order)); one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, including when to compile the kernels, which needs the [feectools stack](#feectools) merged and released first; a decision on hand-written CUDA vs. code generation before the guiding-center kernels (step 3). - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)). - [x] **PR 13: Move the kernel infrastructure to cunumpy** — `Kernel`, `KernelCatalog`, `CudaKernel` and `Argument` are replaced by their cunumpy counterparts; each owner has a single `args_*` object on both backends, and `pusher_args.cuh` is generated (see [Moving to cunumpy](#moving-to-cunumpy-pr-13)). Kernels and device helpers only change their includes. - [ ] **CI**: a GPU runner that runs the CUDA tests (can happen any time; until then the GPU tests are run by hand on an H100 before each PR that touches CUDA code is merged). @@ -64,7 +65,7 @@ CUDA kernels can be added one by one. If the code runs on the GPU and needs a ke | `src/struphy/utils/cuda_arguments.py` | `CUDA_STRUCTS`, `CUDA_OPTIONS`, `write_pusher_header()` and `write_local_projectors_header()` | | `src/struphy/kernel_arguments/pusher_args.cuh`, `local_projectors_args.cuh` | the C structs `MarkerArgs`, `DerhamArgs`, `DomainArgs` and `LocalProjectorsArgs`, generated from the CUDA classes, using cunumpy array views | | `src/struphy/geometry/base.py`, `src/struphy/pic/base.py`, `src/struphy/feec/psydac_derham.py` | `Domain.args_domain`, `Particles.args_markers` and `Derham.args_derham` are the pyccel class on NumPy and the CUDA class on CuPy, chosen once at construction; every kernel call goes through a `Kernel` | -| `src/struphy/*/kernels/` (`pic/pushing`, `pic/accumulation`, `pic/diagnostics`, `pic/sph`, `bsplines`, `geometry`, `feec`, `feec/local_projectors`) | every kernel called from Python with argument objects, one folder each: 44 pusher/evaluation (incl. `reflect`), 16 accumulation, 10 marker diagnostics, 4 SPH evaluation, 3 spline evaluation, 4 geometry, 1 FEEC utility and 8 local projector kernels. Each folder's `__init__.py` declares its `Kernel`; the code imports it (`from struphy.geometry.kernels.kernel_evaluate import kernel_evaluate`). Fifteen have CUDA versions, with parity cases in `pic/tests/cuda_parity_cases.py`, among them the four geometry kernels (PR 18) | +| `src/struphy/*/kernels/` (`pic/pushing`, `pic/accumulation`, `pic/diagnostics`, `pic/sph`, `bsplines`, `geometry`, `feec`, `feec/local_projectors`) | every kernel called from Python with argument objects, one folder each: 44 pusher/evaluation (incl. `reflect`), 16 accumulation, 10 marker diagnostics, 4 SPH evaluation, 3 spline evaluation, 4 geometry, 1 FEEC utility and 8 local projector kernels. Each folder's `__init__.py` declares its `Kernel`; the code imports it (`from struphy.geometry.kernels.kernel_evaluate import kernel_evaluate`). Sixteen have CUDA versions (`linear_vlasov_ampere` the latest), with parity cases in `pic/tests/cuda_parity_cases.py`, among them the four geometry kernels (PR 18) | | `src/struphy/geometry/evaluation_kernels.cuh`, `transform_kernels.cuh`, `spline_mappings_kernels.cuh`, `geometry/domains//_cuda.cuh` | device versions of the mapping helpers for every mapping: `f`/`df` per analytic domain (`kind_map` 10–12, 20–22, 30–32, PR 18) and the spline mappings `spline_3d`, `spline_2d_straight`, `spline_2d_torus` (`kind_map` 0–2, PR 19), the `kind_map` switch, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`, `push`, `tran` | | `src/struphy/pic/tests/test_cuda_parity.py`, `test_cuda_emulation.py` | parity of every CUDA kernel with pyccel on the cases of `pic/tests/cuda_parity_cases.py` (`PARITY_CASES`): on a GPU (`assert_kernels_agree`), and without one by CPU emulation | | `src/struphy/pic/tests/test_kernel_backends.py` | the generated-header tests, the check that each CUDA class mirrors its pyccel class, signature check of all kernels, owner class selection, and (on a GPU) struct layout and copy/pickle checks | @@ -277,7 +278,7 @@ from the struphy tracking issue ([#650](https://github.com/struphy-hub/struphy/i Kernels are ported in the order that completes one model after the other on the GPU. Each step lists only the kernels it adds; P = pusher, A = accumulation, E = marker evaluation. Status: -done (✓), in this PR (PR 14), blocked (⏸). +done (✓), in a PR (PR 14), blocked (⏸). **Step 0 – Vlasov** (PRs 11–12, ✓): `push_eta_stage`, `push_vxb_analytic`, `push_vxb_implicit`, `reflect`. @@ -289,14 +290,14 @@ done (✓), in this PR (PR 14), blocked (⏸). |---|---|---|---|---| | 1 | `push_v_with_efield` | P | `VlasovAmpereCoupling`, `PushVinForceField` | PR 14 | | 2 | `charge_density_0form` | A (vector) | initial Poisson solve of these models | PR 14 | -| 3 | `vlasov_maxwell` | A (matrix + vector) | `VlasovAmpereCoupling` | ⏸ 6D views | +| 3 | `vlasov_maxwell` | A (matrix + vector) | `VlasovAmpereCoupling` | next (6D views available) | **Step 2 – LinearVlasovAmpère/Maxwell (δf)** | # | Kernel | Type | Used by | Status | |---|---|---|---|---| | 4 | `push_weights_with_efield_lin_va` | P | `EfieldWeightsCoupling` | PR 14 | -| 5 | `linear_vlasov_ampere` | A (matrix + vector) | `EfieldWeightsCoupling` | ⏸ 6D views | +| 5 | `linear_vlasov_ampere` | A (matrix + vector) | `EfieldWeightsCoupling` | ✓ (#688) | **Step 3 – Guiding center: ToyDrift, then DriftKineticElectrostaticAdiabatic.** The default algorithm is `discrete_gradient_1st_order`; its chain comes first. @@ -348,9 +349,9 @@ The three spaces (H1vec/Hcurl/Hdiv) differ only in the basis; port one, then the 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, -and array views with more than 4 dimensions in cunumpy (all matrix accumulations write 6D stencil -matrix data). +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. ## Testing @@ -662,3 +663,45 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a during the exchange, staged: one `to_host` per non-empty send and one `to_device`. Without a GPU it runs on the fake CuPy (`mpirun -n N env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi pic/tests/test_sorting_device.py`, in the PIC MPI workflow); the GPU variant has not run (no GPU here). +## `linear_vlasov_ampere` implementation notes + +The first matrix accumulation on the GPU (part of #688, tracked in #650): `linear_vlasov_ampere`, the accumulation of +`EfieldWeightsCoupling` in `LinearVlasovAmpereOneSpecies` (step 2 of the [porting order](#porting-order)). + +- **6D views.** The six blocks `mat11, mat12, mat13, mat22, mat23, mat33` are `Array6D` (cunumpy 0.6.1, strided, + like the `Array3D` vectors); `cunumpy.kernels.CudaKernel` packs the `StencilMatrix._data` blocks without + copying them. `f0_values` is a `const double*`, as in `push_weights_with_efield_lin_va`. +- **Device helpers, one to one.** `filler_kernels.cuh` gets `fill_mat` and `fill_mat_vec` next to `fill_vec`, + `particle_to_mat_kernels.cuh` gets `m_v_fill_b_v1_symm` next to `vec_fill_b_v0`, and `linalg_kernels.cuh` gets + `outer`; same names, arguments and loop order as pyccel (`starts` and `pads` are `const long long*`). Every addition + into matrix or vector data is a `cunumpy_atomic_add`, as in `charge_density_0form`; the spline values pyccel keeps + in `args_derham.bn1, ..., bd3` are in the thread-local `SplineScratch`. Of the remaining `m_v_fill_*`/`mat_fill_*` + fillers, each one is ported with the first kernel that calls it. +- **Kernel.** `linear_vlasov_ampere_cuda.cu`: same arguments in the same order as pyccel, one thread per marker row, + holes and boundary particles skipped; `df`, `matrix_inv`, `matrix_vector`, `outer`, then `m_v_fill_b_v1_symm`. The + 3 x 3 `filling_m` is nine row-major entries; the CUDA-only `factor` is pyccel's `f0_values[ip] / (markers[ip, 7] * + n_markers_tot)`, computed once as pyccel does. +- **`Accumulator` on CuPy.** Nothing in `Accumulator` had to change: on the CuPy backend `StencilMatrix._data` and + `StencilVector._data` are device arrays (feectools allocates them with `xp`), the accumulator zeroes them in place and + passes the same arrays to the kernel; `exchange_assembly_data`, `update_ghost_regions`, the transposed blocks of the + symmetric matrix (`StencilMatrix.transpose`, a CUDA kernel in feectools), the scaled operator of `SchurSolver` and its + `dot` all work on them. `test_accum_matrix_cupy.py` checks this against NumPy: the kernel data are the blocks' data + (no copies), the nine arrays, the dense blocks, the accumulated vector, `BC.dot(x)` and three CG iterations of the + Schur solve agree, and nothing is copied to the host from the accumulation to `BC.dot(x)`. On a GPU it runs with + `assert_no_transfers` (not run yet); without one it runs in a subprocess on cunumpy's fake CuPy, with every + `CudaKernel` launch emulated on the CPU in place on the fake device arrays (`emulated_launches()` in + `pic/tests/cuda_emulation.py`; struphy's and feectools' kernels, struct arguments included). This would fit into + cunumpy (fake CuPy with emulated launches). +- **Found, not fixed here.** feectools' `StencilVector.inner` returns its result on the host (`xp.to_numpy`), so each + inner product of the CG solve copies a scalar from the device; the Schur solve is therefore outside the transfer + guard. `MassMatrixPreconditioner`, the default preconditioner of `EfieldWeightsCoupling`, cannot be created on CuPy + (`is_circulant` asserts a host array; the device Kronecker solve is a separate PR). `LinearVlasovAmpereOneSpecies` + end to end on the GPU (#689) also needs the feectools stencil 6D views and device Kronecker solve PRs and the first + H100 run (#687). +- **Tests.** Parity cases `PARITY_CASES["linear_vlasov_ampere"]`: 129 markers (a hole and a boundary particle) in + Cuboid, Colella, HollowTorus and a 3d spline mapping, data shaped like the real accumulation matrices of clamped + splines of degrees 2, 3, 1 on 8 cells (`v1_symm_accumulation_data()` in `pic/tests/kernel_test_args.py`); + `rtol = 1e-12`, `atol = 1e-10` (atomics change the summation order; the entries reach 1e4). The CPU emulation agrees + with pyccel to round-off and catches a transposed `DF`. `m_v_fill_b_v1_symm` is compared with pyccel through a + wrapper kernel (`test_v1_symm_filler` on a GPU, `test_emulated_v1_symm_filler` by emulation). GPU tests have not run + on an H100. diff --git a/src/struphy/linear_algebra/linalg_kernels.cuh b/src/struphy/linear_algebra/linalg_kernels.cuh index 1b72c25fe..4636f2ea3 100644 --- a/src/struphy/linear_algebra/linalg_kernels.cuh +++ b/src/struphy/linear_algebra/linalg_kernels.cuh @@ -100,4 +100,15 @@ __device__ inline void transpose(const double* a, double* b) { for (int i = 0; i < 3; ++i) for (int j = 0; j < 3; ++j) b[3 * i + j] = a[3 * j + i]; } +/** + * Outer product of two vectors of length 3, as in linalg_kernels.outer. + * + * @param a First input vector with three entries. + * @param b Second input vector with three entries. + * @param c Output matrix c_ij = a_i * b_j with nine row-major entries; must not overlap a or b. + */ +__device__ inline void outer(const double* a, const double* b, double* c) { + for (int i = 0; i < 3; ++i) + for (int j = 0; j < 3; ++j) c[3 * i + j] = a[i] * b[j]; +} } diff --git a/src/struphy/pic/accumulation/filler_kernels.cuh b/src/struphy/pic/accumulation/filler_kernels.cuh index 37f3867d9..8a33d7773 100644 --- a/src/struphy/pic/accumulation/filler_kernels.cuh +++ b/src/struphy/pic/accumulation/filler_kernels.cuh @@ -1,5 +1,5 @@ -// Device versions of the fillers in filler_kernels.py: add the contribution of one marker to a stencil vector. -// Many markers write to the same entries, so every addition is atomic. +// Device versions of the fillers in filler_kernels.py: add the contribution of one marker to a stencil vector +// or to a block of a stencil matrix. Many markers write to the same entries, so every addition is atomic. #pragma once #include "cunumpy/array_view.cuh" #include "cunumpy/atomic.cuh" @@ -40,4 +40,125 @@ __device__ inline void fill_vec(int pi1, int pi2, int pi3, const double* bi1, co } } } + +/** + * Add basis functions times filling to a block of a stencil matrix, as in filler_kernels.fill_mat. + * + * @param pi1 Spline degree of the codomain (row indices) along the first axis. + * @param pi2 Spline degree of the codomain along the second axis. + * @param pi3 Spline degree of the codomain along the third axis. + * @param pj1 Spline degree of the domain (column indices) along the first axis. + * @param pj2 Spline degree of the domain along the second axis. + * @param pj3 Spline degree of the domain along the third axis. + * @param bi1 The pi1 + 1 non-vanishing N- or D-spline values of the codomain along the first axis. + * @param bi2 The pi2 + 1 non-vanishing spline values of the codomain along the second axis. + * @param bi3 The pi3 + 1 non-vanishing spline values of the codomain along the third axis. + * @param bj1 The pj1 + 1 non-vanishing spline values of the domain along the first axis. + * @param bj2 The pj2 + 1 non-vanishing spline values of the domain along the second axis. + * @param bj3 The pj3 + 1 non-vanishing spline values of the domain along the third axis. + * @param span1 Knot span index along the first axis. + * @param span2 Knot span index along the second axis. + * @param span3 Knot span index along the third axis. + * @param starts Start indices of the codomain (three entries). + * @param pads Paddings of the codomain (three entries). + * @param mat Stencil matrix data (the _data of a StencilMatrix, shape (n1, n2, n3, 2 * pads + 1)), any strides; + * written atomically. + * @param filling Number multiplied by the basis functions. + * + * The summation order of concurrent threads (and the last bits of the result) depends on the thread schedule. + */ +__device__ inline void fill_mat(int pi1, int pi2, int pi3, int pj1, int pj2, int pj3, const double* bi1, + const double* bi2, const double* bi3, const double* bj1, const double* bj2, + const double* bj3, int span1, int span2, int span3, const long long* starts, + const long long* pads, Array6D mat, double filling) { + for (int il1 = 0; il1 <= pi1; ++il1) { + long long i1 = span1 + il1 - starts[0]; + double b1 = bi1[il1] * filling; + for (int il2 = 0; il2 <= pi2; ++il2) { + long long i2 = span2 + il2 - starts[1]; + double b2 = b1 * bi2[il2]; + for (int il3 = 0; il3 <= pi3; ++il3) { + long long i3 = span3 + il3 - starts[2]; + double b3 = b2 * bi3[il3]; + + for (int jl1 = 0; jl1 <= pj1; ++jl1) { + long long j1 = pads[0] + jl1 - il1; + double b4 = b3 * bj1[jl1]; + for (int jl2 = 0; jl2 <= pj2; ++jl2) { + long long j2 = pads[1] + jl2 - il2; + double b5 = b4 * bj2[jl2]; + for (int jl3 = 0; jl3 <= pj3; ++jl3) { + long long j3 = pads[2] + jl3 - il3; + double b6 = b5 * bj3[jl3]; + + cunumpy_atomic_add(&mat(i1, i2, i3, j1, j2, j3), b6); + } + } + } + } + } + } +} + +/** + * Add basis functions times fillings to a block of a stencil matrix and to a stencil vector, as in + * filler_kernels.fill_mat_vec. + * + * @param pi1 Spline degree of the codomain (row indices) along the first axis. + * @param pi2 Spline degree of the codomain along the second axis. + * @param pi3 Spline degree of the codomain along the third axis. + * @param pj1 Spline degree of the domain (column indices) along the first axis. + * @param pj2 Spline degree of the domain along the second axis. + * @param pj3 Spline degree of the domain along the third axis. + * @param bi1 The pi1 + 1 non-vanishing N- or D-spline values of the codomain along the first axis. + * @param bi2 The pi2 + 1 non-vanishing spline values of the codomain along the second axis. + * @param bi3 The pi3 + 1 non-vanishing spline values of the codomain along the third axis. + * @param bj1 The pj1 + 1 non-vanishing spline values of the domain along the first axis. + * @param bj2 The pj2 + 1 non-vanishing spline values of the domain along the second axis. + * @param bj3 The pj3 + 1 non-vanishing spline values of the domain along the third axis. + * @param span1 Knot span index along the first axis. + * @param span2 Knot span index along the second axis. + * @param span3 Knot span index along the third axis. + * @param starts Start indices of the codomain (three entries). + * @param pads Paddings of the codomain (three entries). + * @param mat Stencil matrix data (the _data of a StencilMatrix), any strides; written atomically. + * @param filling_mat Number multiplied by the basis functions and added to mat. + * @param vec Stencil vector data of the codomain (the _data of a StencilVector), any strides; written atomically. + * @param filling_vec Number multiplied by the basis functions of the codomain and added to vec. + */ +__device__ inline void fill_mat_vec(int pi1, int pi2, int pi3, int pj1, int pj2, int pj3, const double* bi1, + const double* bi2, const double* bi3, const double* bj1, const double* bj2, + const double* bj3, int span1, int span2, int span3, const long long* starts, + const long long* pads, Array6D mat, double filling_mat, + Array3D vec, double filling_vec) { + for (int il1 = 0; il1 <= pi1; ++il1) { + long long i1 = span1 + il1 - starts[0]; + double b1 = bi1[il1]; + for (int il2 = 0; il2 <= pi2; ++il2) { + long long i2 = span2 + il2 - starts[1]; + double b2 = b1 * bi2[il2]; + for (int il3 = 0; il3 <= pi3; ++il3) { + long long i3 = span3 + il3 - starts[2]; + double b3 = b2 * bi3[il3]; + + cunumpy_atomic_add(&vec(i1, i2, i3), b3 * filling_vec); + + for (int jl1 = 0; jl1 <= pj1; ++jl1) { + long long j1 = pads[0] + jl1 - il1; + double b4 = b3 * bj1[jl1] * filling_mat; + for (int jl2 = 0; jl2 <= pj2; ++jl2) { + long long j2 = pads[1] + jl2 - il2; + double b5 = b4 * bj2[jl2]; + for (int jl3 = 0; jl3 <= pj3; ++jl3) { + long long j3 = pads[2] + jl3 - il3; + double b6 = b5 * bj3[jl3]; + + cunumpy_atomic_add(&mat(i1, i2, i3, j1, j2, j3), b6); + } + } + } + } + } + } +} } diff --git a/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu b/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu new file mode 100644 index 000000000..387013af2 --- /dev/null +++ b/src/struphy/pic/accumulation/kernels/linear_vlasov_ampere/linear_vlasov_ampere_cuda.cu @@ -0,0 +1,90 @@ +#include "struphy/geometry/evaluation_kernels.cuh" +#include "struphy/linear_algebra/linalg_kernels.cuh" +#include "struphy/pic/accumulation/particle_to_mat_kernels.cuh" + +/** + * Accumulate the matrix and vector of the linear Vlasov-Ampere system (delta-f) into V1, as in + * linear_vlasov_ampere_kernels.linear_vlasov_ampere. + * + * Filling functions (the factors alpha^2 kappa^2 / v_th^2 and alpha^2 kappa are applied by the propagator): + * A_p^{mu, nu} = f_0(eta_p, v_p) / (N s_0) [DF^{-1}(eta_p) v_p]_mu [DF^{-1}(eta_p) v_p]_nu, + * B_p^mu = w_p [DF^{-1}(eta_p) v_p]_mu; + * one thread per marker row. + * + * @param args_markers Marker buffer (n_markers x n_cols, row-major) and the total marker number Np; the weight is in + * column 6 and s_0 in column 7, as in pyccel. + * @param args_derham Spline degrees (1 to 8), knots and start indices. + * @param args_domain Mapping arguments; every spline and analytic mapping. + * @param mat11 (1,1)-block data of the symmetric V1 -> V1 block matrix (the _data of a StencilMatrix, 6D, any + * strides); written with atomic additions. + * @param mat12 (1,2)-block data; written atomically. + * @param mat13 (1,3)-block data; written atomically. + * @param mat22 (2,2)-block data; written atomically. + * @param mat23 (2,3)-block data; written atomically. + * @param mat33 (3,3)-block data; written atomically. + * @param vec1 First component of the V1 vector data; written atomically. + * @param vec2 Second component of the V1 vector data; written atomically. + * @param vec3 Third component of the V1 vector data; written atomically. + * @param f0_values Value of f0 for each marker row (n_markers entries). + * + * Holes (markers[ip, 0] == -1) and boundary particles (markers[ip, -1] == -2) are skipped. The pyccel loop runs over + * all rows (shape(markers)[0]), which is n_markers here. Many threads add to the same entries, so the summation + * order (and the last bits of the result) differs from the serial pyccel loop. + */ +extern "C" __global__ void linear_vlasov_ampere(MarkerArgs args_markers, DerhamArgs args_derham, DomainArgs args_domain, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + Array3D vec1, Array3D vec2, Array3D vec3, + const double* f0_values) { + // CUDA-only: the marker row of this thread replaces the pyccel loop variable + int ip = blockDim.x * blockIdx.x + threadIdx.x; + + // get number of markers + int n_markers = args_markers.n_markers; + if (ip >= n_markers) return; + + // total number of markers (weights are w_p = delta f_p / (N * s_0)) + int n_markers_tot = args_markers.Np; + + // only do something if particle is a "true" particle (i.e. not a hole) + if (args_markers.markers(ip, 0) == -1. || args_markers.markers(ip, args_markers.markers.shape[1] - 1) == -2.) + return; + + // allocate for metric coeffs (3 x 3 matrices as nine row-major entries) + double dfm[9], df_inv[9]; + + // allocate for filling + double v[3], df_inv_v[3], filling_m[9], filling_v[3]; + + // marker positions + double eta1 = args_markers.markers(ip, 0); + double eta2 = args_markers.markers(ip, 1); + double eta3 = args_markers.markers(ip, 2); + + // get velocity + v[0] = args_markers.markers(ip, 3); + v[1] = args_markers.markers(ip, 4); + v[2] = args_markers.markers(ip, 5); + + // evaluate Jacobian, result in dfm + struphy_cuda::df(eta1, eta2, eta3, args_domain, dfm); + + // invert Jacobian matrix + struphy_cuda::matrix_inv(dfm, df_inv); + + // compute DF^{-1} v + struphy_cuda::matrix_vector(df_inv, v, df_inv_v); + + // filling_m = alpha^2 * kappa^2 * f0 / (N * s_0 * v_th^2) * (DF^{-1} v_p)_mu * (DF^{-1} v_p)_nu + struphy_cuda::outer(df_inv_v, df_inv_v, filling_m); + double factor = f0_values[ip] / (args_markers.markers(ip, 7) * n_markers_tot); + for (int k = 0; k < 9; ++k) filling_m[k] *= factor; + + // filling_v = alpha^2 * kappa * w_p * DL^{-1} * v_p + for (int k = 0; k < 3; ++k) filling_v[k] = args_markers.markers(ip, 6) * df_inv_v[k]; + + // call the appropriate matvec filler + struphy_cuda::m_v_fill_b_v1_symm(args_derham, eta1, eta2, eta3, mat11, mat12, mat13, mat22, mat23, mat33, + filling_m[0], filling_m[1], filling_m[2], filling_m[4], filling_m[5], + filling_m[8], vec1, vec2, vec3, filling_v[0], filling_v[1], filling_v[2]); +} diff --git a/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh b/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh index 7cabaee45..ade422b97 100644 --- a/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh +++ b/src/struphy/pic/accumulation/particle_to_mat_kernels.cuh @@ -28,4 +28,73 @@ __device__ inline void vec_fill_b_v0(const DerhamArgs& args_derham, double eta1, fill_vec(pn1, pn2, pn3, scratch.bn1, scratch.bn2, scratch.bn3, span1, span2, span3, args_derham.starts, vec, fill); } + +/** + * Add one marker's contribution to the symmetric blocks (1,1), (1,2), (1,3), (2,2), (2,3), (3,3) of a V1 -> V1 + * block matrix and to a V1 vector, as in particle_to_mat_kernels.m_v_fill_b_v1_symm. + * + * @param args_derham Spline degrees (1 to 8), knots and start indices; pn is also the padding of the blocks. + * @param eta1 Logical position along the first axis. + * @param eta2 Logical position along the second axis. + * @param eta3 Logical position along the third axis. + * @param mat11 (mu=1, nu=1)-block data (D N N x D N N), shape (n1, n2, n3, 2 * pn + 1); written atomically. + * @param mat12 (mu=1, nu=2)-block data (D N N x N D N); written atomically. + * @param mat13 (mu=1, nu=3)-block data (D N N x N N D); written atomically. + * @param mat22 (mu=2, nu=2)-block data (N D N x N D N); written atomically. + * @param mat23 (mu=2, nu=3)-block data (N D N x N N D); written atomically. + * @param mat33 (mu=3, nu=3)-block data (N N D x N N D); written atomically. + * @param fill11 Number multiplied by the basis functions and added to mat11. + * @param fill12 Number multiplied by the basis functions and added to mat12. + * @param fill13 Number multiplied by the basis functions and added to mat13. + * @param fill22 Number multiplied by the basis functions and added to mat22. + * @param fill23 Number multiplied by the basis functions and added to mat23. + * @param fill33 Number multiplied by the basis functions and added to mat33. + * @param vec1 mu=1 component of the V1 vector data (D N N); written atomically. + * @param vec2 mu=2 component of the V1 vector data (N D N); written atomically. + * @param vec3 mu=3 component of the V1 vector data (N N D); written atomically. + * @param fill1 Number multiplied by the basis functions and added to vec1. + * @param fill2 Number multiplied by the basis functions and added to vec2. + * @param fill3 Number multiplied by the basis functions and added to vec3. + * + * Pyccel keeps the spline values in args_derham.bn1, ..., bd3; here they are in a thread-local SplineScratch. + */ +__device__ inline void m_v_fill_b_v1_symm(const DerhamArgs& args_derham, double eta1, double eta2, double eta3, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + double fill11, double fill12, double fill13, double fill22, + double fill23, double fill33, Array3D vec1, Array3D vec2, + Array3D vec3, double fill1, double fill2, double fill3) { + // degrees of the basis functions : B-splines (pn) and D-splines (pd) + int pn1 = args_derham.pn[0]; + int pn2 = args_derham.pn[1]; + int pn3 = args_derham.pn[2]; + + int pd1 = pn1 - 1; + int pd2 = pn2 - 1; + int pd3 = pn3 - 1; + + // CUDA-only scratch holds the spans and the spline values pyccel keeps in args_derham + SplineScratch scratch; + get_spans(eta1, eta2, eta3, args_derham, scratch); + int span1 = scratch.span1, span2 = scratch.span2, span3 = scratch.span3; + + // fill matrix entries + fill_mat_vec(pd1, pn2, pn3, pd1, pn2, pn3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bd1, scratch.bn2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat11, fill11, vec1, fill1); + + fill_mat_vec(pn1, pd2, pn3, pn1, pd2, pn3, scratch.bn1, scratch.bd2, scratch.bn3, scratch.bn1, scratch.bd2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat22, fill22, vec2, fill2); + + fill_mat_vec(pn1, pn2, pd3, pn1, pn2, pd3, scratch.bn1, scratch.bn2, scratch.bd3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat33, fill33, vec3, fill3); + + fill_mat(pd1, pn2, pn3, pn1, pd2, pn3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bn1, scratch.bd2, + scratch.bn3, span1, span2, span3, args_derham.starts, args_derham.pn, mat12, fill12); + + fill_mat(pd1, pn2, pn3, pn1, pn2, pd3, scratch.bd1, scratch.bn2, scratch.bn3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat13, fill13); + + fill_mat(pn1, pd2, pn3, pn1, pn2, pd3, scratch.bn1, scratch.bd2, scratch.bn3, scratch.bn1, scratch.bn2, + scratch.bd3, span1, span2, span3, args_derham.starts, args_derham.pn, mat23, fill23); +} } diff --git a/src/struphy/pic/tests/cuda_emulation.py b/src/struphy/pic/tests/cuda_emulation.py index 0e8db4f12..ca3c76d01 100644 --- a/src/struphy/pic/tests/cuda_emulation.py +++ b/src/struphy/pic/tests/cuda_emulation.py @@ -6,8 +6,10 @@ the fields of each struct as separate parameters, rebuilds the structs and calls the real kernel. """ +import contextlib + import numpy as np -from cunumpy.kernel_testing import emulate_cuda_kernel +from cunumpy.kernel_testing import emulate_cuda_kernel, fake_cupy_active from cunumpy.kernels import CudaKernel from struphy.utils.cuda_arguments import CUDA_OPTIONS @@ -55,11 +57,12 @@ def _struct_field(value, name): return getattr(value, name) -def emulate_struct_kernel(kernel: CudaKernel, *args, n_threads=None): +def emulate_struct_kernel(kernel: CudaKernel, *args, n_threads=None, grid=None, block=None): """Emulate `kernel` with host arguments; struct arguments are the pyccel argument objects (NumPy arrays). - Arrays (also those inside the argument objects) are updated in place, as by a launch. Without `n_threads`, - one thread per row of the first array is launched, as in a real launch. + Arrays (also those inside the argument objects) are updated in place, as by a launch. Without `n_threads` or + `grid`, one thread per row of the first array is launched, as in a real launch; `grid` and `block` give the + launch shape explicitly. """ flat = [] for p, value in zip(kernel.signature, args): @@ -67,9 +70,58 @@ def emulate_struct_kernel(kernel: CudaKernel, *args, n_threads=None): flat.append(value) else: flat += [_struct_field(value, f.name) for f in p.struct.fields] - if n_threads is None: - n_threads = next(value.shape[0] for value in flat if isinstance(value, np.ndarray)) - grid, block = kernel.launch_shape(n_threads) + if grid is None: + if n_threads is None: + n_threads = next(value.shape[0] for value in flat if isinstance(value, np.ndarray)) + grid, block = kernel.launch_shape(n_threads, block=block) source, name = _wrapper(kernel) wrapper = CudaKernel(source, name, source_dir=kernel.source_dir, **CUDA_OPTIONS) emulate_cuda_kernel(wrapper, *flat, grid=grid, block=block, options=EMULATION_OPTIONS) + + +def _host(value): + """The host buffer behind a fake CuPy array (the same memory, no copy); other values unchanged.""" + if isinstance(value, np.ndarray) or not hasattr(value, "__cuda_array_interface__"): + return value + return object.__getattribute__(value, "_a") + + +class _HostFields: + """A CUDA argument object (``Cuda*Arguments``) whose array attributes are their fake CuPy host buffers.""" + + def __init__(self, args): + self._args = args + + def __getattr__(self, name): + return _host(getattr(self._args, name)) + + +@contextlib.contextmanager +def emulated_launches(): + """On cunumpy's fake CuPy, run every ``CudaKernel`` launch by CPU emulation, in place on the fake device arrays. + + The fake CuPy (``CUNUMPY_FAKE_CUPY=1``) keeps device arrays in host memory and rejects host/device mixing, but + cannot launch kernels. Inside this context a launch emulates the kernel (struct arguments included, through + :func:`emulate_struct_kernel`) on the host buffers of the fake arrays, so code that launches CUDA kernels (struphy's + and feectools') runs end to end on the CuPy backend without a GPU. Each launch compiles the kernel, so this is slow. + """ + if not fake_cupy_active(): + raise RuntimeError("emulated_launches() needs cunumpy's fake CuPy (CUNUMPY_FAKE_CUPY=1)") + original = CudaKernel.__call__ + + def launch(self, *args, n_threads=None, grid=None, block=None, shared_mem=0, stream=None): + grid, block = self.launch_shape(n_threads, grid=grid, block=block, args=args) + if any(p.struct is not None for p in self.signature): + host_args = [_host(a) if p.struct is None else _HostFields(a) for p, a in zip(self.signature, args)] + emulate_struct_kernel(self, *host_args, grid=grid, block=block) + else: + host_args = [_host(a) for a in args] + emulate_cuda_kernel( + self, *host_args, grid=grid, block=block, shared_mem=shared_mem, options=EMULATION_OPTIONS + ) + + CudaKernel.__call__ = launch + try: + yield + finally: + CudaKernel.__call__ = original diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index e2ce0a5d5..85380b6e2 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -29,6 +29,7 @@ marker_arguments, spline_coefficients, spline_evaluation_arguments, + v1_symm_accumulation_data, ) @@ -79,6 +80,19 @@ def charge_density_0form_args(bc): return (args_markers, derham_arguments(), args_domain, xp.zeros((18, 20, 16))) +def linear_vlasov_ampere_args(domain): + """129 markers (a hole and a boundary particle) in domain `domain` of geometry_domain(); zeroed V1 data.""" + args_markers, _ = marker_arguments((0, 0, 0)) + f0_values = xp.asarray(np.random.default_rng(14).random(N_MARKERS)) + return ( + args_markers, + derham_arguments(), + geometry_domain(domain).args_domain, + *v1_symm_accumulation_data(), + f0_values, + ) + + # ---------------------------------------------------------------- spline evaluation @@ -267,6 +281,9 @@ def size_of(index, per_thread=1): "reflect": ParityCases((0, 1, 2), reflect_args, n_threads=size_of(2)), # atomics add in another order than the serial loop "charge_density_0form": ParityCases(((0, 0, 0), (2, 0, 1)), charge_density_0form_args, rtol=1e-12, atol=1e-13), + # Cuboid, Colella and HollowTorus (non-diagonal DF) and a 3d spline mapping, by index into geometry_domain(); + # atomics add in another order than the serial loop, and the entries reach 1e4 (atol for cancelling sums) + "linear_vlasov_ampere": ParityCases((0, 2, 5, 12), linear_vlasov_ampere_args, rtol=1e-12, atol=1e-10), "eval_spline_mpi_markers": ParityCases(SPLINE_KINDS, eval_spline_mpi_markers_args), "eval_spline_mpi_matrix": ParityCases(SPLINE_KINDS, eval_spline_mpi_grid_args(sparse=False), n_threads=size_of(-1)), "eval_spline_mpi_sparse_meshgrid": ParityCases( diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index d413b5380..eaed48adf 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -49,14 +49,37 @@ def butcher_arguments(method): return xp.asarray(butcher.a_stage), xp.asarray(butcher.b), xp.asarray(butcher.c), butcher.n_stages +DERHAM_DEGREE = (2, 3, 1) +DERHAM_CELLS = 8 + + def derham_arguments(): """Splines of degrees 2, 3, 1 on 8 cells, starting at index 0.""" - degree = np.array([2, 3, 1], dtype=np.int64) - knots = [np.r_[np.zeros(p), np.linspace(0, 1, 9), np.ones(p)] for p in degree] + degree = np.array(DERHAM_DEGREE, dtype=np.int64) + knots = [np.r_[np.zeros(p), np.linspace(0, 1, DERHAM_CELLS + 1), np.ones(p)] for p in degree] args_class = CudaDerhamArguments if xp.get_backend() == "cupy" else DerhamArguments return args_class(xp.asarray(degree), *(xp.asarray(t) for t in knots), xp.zeros(3, dtype=np.int64)) +def v1_symm_accumulation_data(): + """Zeroed data of a symmetric V1 -> V1 accumulation matrix and a V1 vector, for :func:`derham_arguments`. + + Returns the six independent blocks ``mat11, mat12, mat13, mat22, mat23, mat33`` and the three vector components + ``vec1, vec2, vec3`` on the active backend, shaped like the ``_data`` of the ``StencilMatrix`` blocks and + ``StencilVector`` components that ``Accumulator(..., "Hcurl", ..., symmetry="symm", add_vector=True)`` passes on + one process for clamped splines of degrees 2, 3, 1 on 8 cells: component ``mu`` has D-splines (``n - 1`` + functions) along axis ``mu`` and N-splines (``n = 8 + p``) along the others, every axis is padded by ``p`` on both + sides, and the matrix blocks have ``2 p + 1`` diagonals per axis (the paddings are the degrees of V0). + """ + degree = np.array(DERHAM_DEGREE) + n_spline = DERHAM_CELLS + degree + rows = [tuple(n_spline - (np.arange(3) == mu) + 2 * degree) for mu in range(3)] + diagonals = tuple(2 * degree + 1) + mats = tuple(xp.zeros(rows[mu] + diagonals) for mu, nu in ((0, 0), (0, 1), (0, 2), (1, 1), (1, 2), (2, 2))) + vecs = tuple(xp.zeros(rows[mu]) for mu in range(3)) + return mats + vecs + + def spline_coefficients(n=3, seed=11): """Random coefficients, large enough for every span of :func:`derham_arguments`.""" rng = np.random.default_rng(seed) diff --git a/src/struphy/pic/tests/test_accum_matrix_cupy.py b/src/struphy/pic/tests/test_accum_matrix_cupy.py new file mode 100644 index 000000000..cf0aa62c6 --- /dev/null +++ b/src/struphy/pic/tests/test_accum_matrix_cupy.py @@ -0,0 +1,209 @@ +"""The matrix path of :class:`~struphy.pic.accumulation.particles_to_grid.Accumulator` on the CuPy backend. + +``linear_vlasov_ampere`` accumulates the symmetric V1 -> V1 block matrix and the V1 vector of ``EfieldWeightsCoupling``; +the matrix data (``StencilMatrix._data``) must be device arrays that are zeroed, passed to the CUDA kernel and turned +into the Schur operator without host copies. Checked against the NumPy backend on a GPU and, without one, on +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 + +import cunumpy as xp +import numpy as np +import pytest +from cunumpy.kernel_testing import emulation_compiler, requires_cupy + +from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env + +ALPHA, KAPPA, VTH, DT = 1.3, 0.7, 0.9, 0.05 +N_SCHUR_ITERATIONS = 3 + + +@dataclass +class MarkerOwner: + """What the accumulators read from ``Particles``: the marker arguments, no clones. + + The same markers on both backends (``Particles`` draws them with the backend's random generator). + """ + + args_markers: Any + clone_config: Any = None + + +def accumulate_linear_vlasov_ampere(guard=None): + """Accumulate ``linear_vlasov_ampere`` like ``EfieldWeightsCoupling`` on the active backend, then solve with it. + + Colella mapping (non-diagonal Jacobian), clamped splines in eta1 and periodic ones in eta2, eta3, so the three + V1 components have different sizes. `guard` (a context manager factory, e.g. one that fails on host transfers) + encloses the accumulation, the scaled operator and its application, as in a time step of the propagator; the + Schur solve is not guarded (see below). + + Returns a dict of host arrays: the nine kernel data arrays (``mat11`` ... ``vec3``) after the accumulation (ghost + regions exchanged), the assembled block matrix as dense blocks, the accumulated vector after the boundary + operators, ``BC.dot(x)`` with the operator scaled as in ``EfieldWeightsCoupling`` and the Schur solution after + ``N_SCHUR_ITERATIONS`` iterations of the conjugate gradient method. + """ + import contextlib + + from feectools.linalg.stencil import StencilMatrix + + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.psydac_derham import Derham + from struphy.geometry.domains import Colella + from struphy.io.options import DerhamOptions + from struphy.linear_algebra.schur_solver import SchurSolver + from struphy.linear_algebra.solver import SolverParameters + from struphy.pic.accumulation.kernels.linear_vlasov_ampere import linear_vlasov_ampere + from struphy.pic.accumulation.particles_to_grid import Accumulator + from struphy.pic.tests.kernel_test_args import N_MARKERS, marker_arguments + from struphy.topology.grids import TensorProductGrid + + guard = guard or contextlib.nullcontext + derham = Derham( + TensorProductGrid(num_elements=(6, 5, 4)), + DerhamOptions(degree=(2, 3, 1), bcs=(("free", "free"), None, None)), + comm=None, + ) + domain = Colella(Lx=2.0, Ly=3.0, alpha=0.05, Lz=4.0) + mass_ops = WeightedMassOperators(derham, domain) + args_markers, _ = marker_arguments((0, 0, 0)) + accum = Accumulator( + MarkerOwner(args_markers), + "Hcurl", + linear_vlasov_ampere, + mass_ops, + domain.args_domain, + add_vector=True, + symmetry="symm", + ) + matrix = accum.operators[0].matrix + + # the kernel writes into the data of the operator's blocks (and vector), not into copies + blocks = [matrix[a, b] for a in range(3) for b in range(a, 3)] + assert all(isinstance(block, StencilMatrix) for block in blocks) + assert all(data is block._data for data, block in zip(accum._args_data[:6], blocks)) + assert all(data is vec._data for data, vec in zip(accum._args_data[6:], accum._vectors[0].blocks)) + assert all(xp.is_gpu(data) == (xp.get_backend() == "cupy") for data in accum._args_data) + + rng = np.random.default_rng(16) + f0_values = xp.asarray(rng.random(N_MARKERS)) + for data in accum._args_data: + data[...] = 1.0 # left over from a previous step: the accumulator zeroes it + + x = accum._vectors[0].space.zeros() + for block in x.blocks: + block._data[...] = xp.asarray(rng.normal(size=block._data.shape)) + x.update_ghost_regions() + + schur_solver = SchurSolver( + mass_ops.M1, + ALPHA**2 * KAPPA**2 * accum.operators[0] / (4 * VTH**2), + "pcg", + # a few iterations (each one launches nine stencil dot kernels, which the emulation compiles one by one) + solver_params=SolverParameters(tol=1e-30, maxiter=N_SCHUR_ITERATIONS), + ) + + with guard(): + accum(f0_values) + # as in EfieldWeightsCoupling.__call__ + schur_solver.BC = accum.operators[0] + schur_solver.BC *= (-1) * ALPHA**2 * KAPPA**2 / (4 * VTH**2) + bc_x = schur_solver.BC.dot(x) + byn = accum.vectors[0].copy() + byn *= ALPHA**2 * KAPPA / 2.0 + # not guarded: feectools' StencilVector.inner returns its result on the host (xp.to_numpy), so every inner + # product of the conjugate gradient method copies one scalar from the device + solution, info = schur_solver(xn=x, Byn=byn, dt=DT) + + out = {name: xp.to_numpy(data) for name, data in zip(MATRIX_DATA_NAMES, accum._args_data)} + for a in range(3): + for b in range(3): + out[f"block{a}{b}"] = xp.to_numpy(matrix[a, b].toarray()) + out["vector"] = xp.to_numpy(accum.vectors[0].toarray()) + out["bc_x"] = xp.to_numpy(bc_x.toarray()) + out["solution"] = xp.to_numpy(solution.toarray()) + out["niter"] = np.array(info["niter"]) + return out + + +MATRIX_DATA_NAMES = ("mat11", "mat12", "mat13", "mat22", "mat23", "mat33", "vec1", "vec2", "vec3") + + +def compare_with_numpy(result): + """Compare a CuPy-backend `result` of :func:`accumulate_linear_vlasov_ampere` with the NumPy backend's.""" + with xp.use_backend("numpy"): + expected = accumulate_linear_vlasov_ampere() + assert expected["niter"] == result["niter"] == N_SCHUR_ITERATIONS + assert np.any(expected["mat12"] != 0.0) and np.any(expected["vec3"] != 0.0) + # the transposed blocks of the symmetric matrix are filled from the accumulated ones + np.testing.assert_allclose(expected["block10"], expected["block01"].T, rtol=1e-14, atol=1e-14) + for name, value in expected.items(): + scale = max(np.max(np.abs(value)), 1.0) + # atomics add in another order than the serial loop; the solve repeats it + np.testing.assert_allclose(result[name], value, rtol=1e-10, atol=1e-12 * scale, err_msg=name) + + +@requires_cupy +def test_accumulator_matrix_on_cupy(): + """On a GPU: the accumulated matrix and vector, the Schur operator and solve agree with NumPy; the accumulation and + the operator run without host/device transfers.""" + from cunumpy.profiling import assert_no_transfers + + with xp.use_backend("cupy"): + accumulate_linear_vlasov_ampere() # warm up (NVRTC compilation) + result = accumulate_linear_vlasov_ampere(guard=assert_no_transfers) + compare_with_numpy(result) + + +def fake_cupy_transfer_guard(): + """On cunumpy's fake CuPy: fail on every explicit host copy of a device array (``.get()``, ``cupy.asnumpy``). + + The fake CuPy already rejects implicit conversions (``numpy.asarray`` of a device array, mixing host and device + arrays); this closes the explicit ones, which are legal but are transfers inside a time step. + """ + import contextlib + + import cupy + + @contextlib.contextmanager + def guard(): + def fail(*args, **kwargs): + raise AssertionError("host copy of a device array during the accumulation") + + saved = cupy.ndarray.get, cupy.asnumpy + cupy.ndarray.get, cupy.asnumpy = fail, fail + try: + yield + finally: + cupy.ndarray.get, cupy.asnumpy = saved + + return guard + + +def check_accumulator_matrix_fake_cupy(): + """Run in a child process with ``CUNUMPY_FAKE_CUPY=1``: the CuPy backend with emulated kernel launches.""" + from struphy.pic.tests.cuda_emulation import emulated_launches + + with emulated_launches(), xp.use_backend("cupy"): + result = accumulate_linear_vlasov_ampere(guard=fake_cupy_transfer_guard()) + compare_with_numpy(result) + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +@pytest.mark.skipif(emulation_compiler() is None, reason="no C++ compiler") +def test_accumulator_matrix_fake_cupy(): + """Without a GPU: the matrix path of the Accumulator on the CuPy backend (fake CuPy, emulated launches). + + Device data are zeroed and passed to the CUDA kernel, the ghost regions, the transposed blocks, the scaled + operator and the Schur solve run on device arrays, and no array is copied to the host from the accumulation to + the operator's application. Runs in a + 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:] diff --git a/src/struphy/pic/tests/test_cuda_emulation.py b/src/struphy/pic/tests/test_cuda_emulation.py index 793bb9525..f2ab111a2 100644 --- a/src/struphy/pic/tests/test_cuda_emulation.py +++ b/src/struphy/pic/tests/test_cuda_emulation.py @@ -113,3 +113,15 @@ def test_emulated_transform_helpers(domain_index): out = np.zeros(len(expected)) emulate_struct_kernel(kernel, *inputs, args_domain, out, len(out), n_threads=len(out)) np.testing.assert_allclose(out, expected, rtol=1e-10, atol=1e-10) + + +@requires_compiler +def test_emulated_v1_symm_filler(): + """test_device_helpers.test_v1_symm_filler by CPU emulation: the device m_v_fill_b_v1_symm agrees with pyccel.""" + from struphy.pic.tests.test_device_helpers import v1_symm_filler_case + + kernel, (eta, fills), args_derham, data, expected = v1_symm_filler_case() + emulate_struct_kernel(kernel, eta, fills, args_derham, *data, len(eta), n_threads=len(eta)) + for result, reference in zip(data, expected): + assert np.any(reference != 0.0) + np.testing.assert_allclose(result, reference, rtol=1e-12, atol=1e-13) diff --git a/src/struphy/pic/tests/test_cuda_parity.py b/src/struphy/pic/tests/test_cuda_parity.py index 75d53245f..0b0a14c9e 100644 --- a/src/struphy/pic/tests/test_cuda_parity.py +++ b/src/struphy/pic/tests/test_cuda_parity.py @@ -10,6 +10,7 @@ from cunumpy.profiling import assert_no_transfers from struphy.pic.accumulation.kernels.charge_density_0form import charge_density_0form +from struphy.pic.accumulation.kernels.linear_vlasov_ampere import linear_vlasov_ampere from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage from struphy.pic.pushing.kernels.push_v_with_efield import push_v_with_efield from struphy.pic.pushing.kernels.push_vxb_analytic import push_vxb_analytic @@ -108,7 +109,7 @@ def device_only(value, *args, **kwargs): def test_vlasov_kernel_coverage(): - """Vlasov (PushEta, both PushVxB algorithms) and the ported Vlasov-Ampere kernels have device kernels.""" + """Vlasov (PushEta, both PushVxB algorithms) and the ported (linear) Vlasov-Ampere kernels have device kernels.""" for kernel in ( push_eta_stage, push_vxb_analytic, @@ -116,5 +117,6 @@ def test_vlasov_kernel_coverage(): push_v_with_efield, push_weights_with_efield_lin_va, charge_density_0form, + linear_vlasov_ampere, ): assert kernel.has_cuda, kernel.name diff --git a/src/struphy/pic/tests/test_device_helpers.py b/src/struphy/pic/tests/test_device_helpers.py index a2390ef4d..b771d8efc 100644 --- a/src/struphy/pic/tests/test_device_helpers.py +++ b/src/struphy/pic/tests/test_device_helpers.py @@ -391,3 +391,64 @@ def test_transform_helpers(domain_index): kernel, inputs, _, expected = transform_helper_case(domain_index) result = _run_on_device(kernel, inputs, domain_index, len(expected)) np.testing.assert_allclose(result, expected, rtol=1e-10, atol=1e-10) + + +# CUDA-only wrapper: one thread per point calls m_v_fill_b_v1_symm with the point's nine fillings +V1_SYMM_FILLER_SOURCE = r""" +#include "struphy/pic/accumulation/particle_to_mat_kernels.cuh" +extern "C" __global__ void fill_v1_symm(const double* eta, const double* fills, DerhamArgs args_derham, + Array6D mat11, Array6D mat12, Array6D mat13, + Array6D mat22, Array6D mat23, Array6D mat33, + Array3D vec1, Array3D vec2, Array3D vec3, int n) { + int ip = blockDim.x * blockIdx.x + threadIdx.x; + if (ip >= n) return; + const double* fill = fills + 9 * ip; + struphy_cuda::m_v_fill_b_v1_symm(args_derham, eta[3 * ip], eta[3 * ip + 1], eta[3 * ip + 2], mat11, mat12, mat13, + mat22, mat23, mat33, fill[0], fill[1], fill[2], fill[3], fill[4], fill[5], vec1, + vec2, vec3, fill[6], fill[7], fill[8]); +} +""" + + +def v1_symm_filler_case(): + """Wrapper kernel, host inputs and pyccel reference for the device m_v_fill_b_v1_symm. + + 65 points (the corners of the unit cube among them) with random fillings, accumulated into zeroed V1 data shaped + like the real accumulation matrices (:func:`~struphy.pic.tests.kernel_test_args.v1_symm_accumulation_data`). + Used by :func:`test_v1_symm_filler` (GPU) and by its CPU emulation in ``test_cuda_emulation.py``. Returns + ``kernel, (eta, fills), args_derham, data, expected``: ``args_derham`` is the pyccel class and ``data`` the nine + zeroed host arrays in kernel order. + """ + from struphy.pic.accumulation import particle_to_mat_kernels + from struphy.pic.tests.kernel_test_args import derham_arguments, v1_symm_accumulation_data + + rng = np.random.default_rng(15) + eta = np.r_[np.array([[0.0, 0.0, 0.0], [1.0, 1.0, 1.0], [0.0, 1.0, 0.5]]), rng.random((62, 3))] + fills = rng.normal(size=(len(eta), 9)) + with cunumpy.use_backend("numpy"): + args_derham = derham_arguments() + data = v1_symm_accumulation_data() + expected = tuple(a.copy() for a in data) + for point, fill in zip(eta, fills): + particle_to_mat_kernels.m_v_fill_b_v1_symm( + args_derham, *point, *expected[:6], *fill[:6], *expected[6:], *fill[6:] + ) + kernel = CudaKernel(V1_SYMM_FILLER_SOURCE, "fill_v1_symm", **CUDA_OPTIONS) + return kernel, (eta, fills), args_derham, data, expected + + +@requires_cupy +def test_v1_symm_filler(): + """m_v_fill_b_v1_symm (and the fill_mat/fill_mat_vec it calls) on the device agrees with pyccel.""" + import cupy as cp + + from struphy.pic.tests.kernel_test_args import derham_arguments + + kernel, (eta, fills), _, data, expected = v1_symm_filler_case() + with cunumpy.use_backend("cupy"): + args_derham = derham_arguments() + device_data = [cp.asarray(a) for a in data] + kernel(cp.asarray(eta), cp.asarray(fills), args_derham, *device_data, len(eta), n_threads=len(eta)) + for result, reference in zip(device_data, expected): + assert np.any(reference != 0.0) + np.testing.assert_allclose(result.get(), reference, rtol=1e-12, atol=1e-13) From 0d15e06e6efebbbb94062d13a48a35494a0845a8 Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 23:35:08 +0200 Subject: [PATCH 43/50] Mark every column of the test hole so accumulation kernels skip it The hole in marker_arguments() was only marked in column 8 (first_init_idx), which the pushers test. Accumulation kernels test column 0, so the linear_vlasov_ampere parity case had no real hole. Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/kernel_test_args.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index eaed48adf..24e597e89 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -29,7 +29,9 @@ def marker_arguments(bc): markers[:, 3:6] = rng.uniform(-2, 2, (N_MARKERS, 3)) markers[:, 8:14] = markers[:, :6] markers[:, 18:21] = 0.0 - markers[0, 8] = -1.0 + # Row 0 is a hole. Struphy marks every column of a hole with -1: pushers test column first_init_idx (8), + # accumulation kernels test column 0. + markers[0, :] = -1.0 markers[1, -1] = -2.0 valid = np.ones(N_MARKERS, dtype=bool) valid[:2] = False From 0b6fc99bc38a4d99ab4224f3364d7680d8402c80 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 15:11:12 +0200 Subject: [PATCH 44/50] Rename first_init_idx to first_pusher_idx in the code new on this branch Co-Authored-By: Claude Opus 5.5 --- src/struphy/pic/tests/kernel_test_args.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/struphy/pic/tests/kernel_test_args.py b/src/struphy/pic/tests/kernel_test_args.py index 24e597e89..abe1ef445 100644 --- a/src/struphy/pic/tests/kernel_test_args.py +++ b/src/struphy/pic/tests/kernel_test_args.py @@ -29,7 +29,7 @@ def marker_arguments(bc): markers[:, 3:6] = rng.uniform(-2, 2, (N_MARKERS, 3)) markers[:, 8:14] = markers[:, :6] markers[:, 18:21] = 0.0 - # Row 0 is a hole. Struphy marks every column of a hole with -1: pushers test column first_init_idx (8), + # Row 0 is a hole. Struphy marks every column of a hole with -1: pushers test column first_pusher_idx (8), # accumulation kernels test column 0. markers[0, :] = -1.0 markers[1, -1] = -2.0 From 85ad5e0503388df76853b3960e74df70a64b1838 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Mon, 12 Oct 2026 11:24:10 +0200 Subject: [PATCH 45/50] Formatting --- src/struphy/polar/linear_operators.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/struphy/polar/linear_operators.py b/src/struphy/polar/linear_operators.py index 623417a64..288286240 100644 --- a/src/struphy/polar/linear_operators.py +++ b/src/struphy/polar/linear_operators.py @@ -320,7 +320,9 @@ def dot(self, v, out=None): # 3. map polar coeffs to "polar rings" (added in-place) if self.blocks_ten_to_pol is not None: - dot_parts_of_polar(self._backend_blocks("blocks_ten_to_pol"), blocks_e3, v, out, map_from_tp=False, add=True) + dot_parts_of_polar( + self._backend_blocks("blocks_ten_to_pol"), blocks_e3, v, out, map_from_tp=False, add=True + ) # "standard" operator (tensor-product vector --> polar vector) else: From 5825e1f7b8c07ac847395e0674f711b6f0086cc2 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Mon, 12 Oct 2026 12:11:02 +0200 Subject: [PATCH 46/50] Fix unit test --- src/struphy/polar/kron_2dx1d.py | 23 ++++++++++++++--------- 1 file changed, 14 insertions(+), 9 deletions(-) diff --git a/src/struphy/polar/kron_2dx1d.py b/src/struphy/polar/kron_2dx1d.py index cdd9dad8c..871b904c5 100644 --- a/src/struphy/polar/kron_2dx1d.py +++ b/src/struphy/polar/kron_2dx1d.py @@ -32,16 +32,17 @@ import logging from dataclasses import dataclass, field +import cunumpy as xp import numpy as np import scipy.sparse as spa -from feectools.fem.tensor import TensorFemSpace -from feectools.linalg.basic import LinearOperator -from feectools.linalg.block import BlockVector -from feectools.linalg.stencil import StencilVector from maybempi import MPI from scipy.linalg import lu_factor, lu_solve from scipy.sparse.linalg import spilu, splu +from feectools.fem.tensor import TensorFemSpace +from feectools.linalg.basic import LinearOperator +from feectools.linalg.block import BlockVector +from feectools.linalg.stencil import StencilVector from struphy.polar.basic import PolarDerhamSpace, PolarVector logger = logging.getLogger("struphy") @@ -573,19 +574,23 @@ def _apply(self, v: PolarVector, out: PolarVector | None, ops_3: list, ops_pol: assert out.space == self._space V = self._space - out_pol = [np.zeros_like(p) for p in v.pol] + # These Kronecker solvers and MPI transposes use SciPy/NumPy on the host. + # Move backend arrays explicitly at this boundary; implicit CuPy-to-NumPy + # conversions are rejected by CuPy and by cunumpy's fake backend. + out_pol = [np.zeros_like(xp.to_numpy(p)) for p in v.pol] for g, op_3, op_pol in zip(self._groups, ops_3, ops_pol): # 1. eta_3 direction (tensor-product part redistributed, polar coeffs replicated) tmp_tp = {} tmp_pol = {} for c in g.comps: - L = _interior(tp_component(v.tp, c)) + L = xp.to_numpy(_interior(tp_component(v.tp, c))) l1, l2, l3 = L.shape T3 = self._t3[c] X = T3.forward(np.ascontiguousarray(L.reshape(l1 * l2, l3).T)) tmp_tp[c] = np.ascontiguousarray(T3.backward(op_3(X)).T).reshape(l1, l2, l3) - tmp_pol[c] = op_3(v.pol[c].T).T if V.n_polar[c] > 0 else v.pol[c].copy() + pol = xp.to_numpy(v.pol[c]) + tmp_pol[c] = op_3(pol.T).T if V.n_polar[c] > 0 else pol.copy() # 2. poloidal plane s3 = tp_space(V, g.comps[0]).starts[2] @@ -611,7 +616,7 @@ def _apply(self, v: PolarVector, out: PolarVector | None, ops_3: list, ops_pol: F[nr:] = Z[off : off + (n1 - nr) * n2].reshape(n1 - nr, n2, k1 - k0) off += (n1 - nr) * n2 u = tp_component(out.tp, c) - _interior(u)[...] = self._tpl[c].backward(F) + _interior(u)[...] = xp.asarray(self._tpl[c].backward(F)) u.ghost_regions_in_sync = False # polar coefficients are held by all processes (each column was computed on exactly one process) @@ -619,7 +624,7 @@ def _apply(self, v: PolarVector, out: PolarVector | None, ops_3: list, ops_pol: for c in range(V.n_comps): if comm is not None and comm.Get_size() > 1 and out_pol[c].size > 0: comm.Allreduce(MPI.IN_PLACE, out_pol[c], op=MPI.SUM) - out.pol[c][:] = out_pol[c] + out.pol[c][:] = xp.asarray(out_pol[c]) return out From 12487226478249544430e3de15ac4769d5735d98 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Mon, 12 Oct 2026 13:35:32 +0200 Subject: [PATCH 47/50] ruff check --fix --- src/struphy/polar/kron_2dx1d.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/struphy/polar/kron_2dx1d.py b/src/struphy/polar/kron_2dx1d.py index 871b904c5..07cb33f79 100644 --- a/src/struphy/polar/kron_2dx1d.py +++ b/src/struphy/polar/kron_2dx1d.py @@ -35,14 +35,14 @@ import cunumpy as xp import numpy as np import scipy.sparse as spa -from maybempi import MPI -from scipy.linalg import lu_factor, lu_solve -from scipy.sparse.linalg import spilu, splu - from feectools.fem.tensor import TensorFemSpace from feectools.linalg.basic import LinearOperator from feectools.linalg.block import BlockVector from feectools.linalg.stencil import StencilVector +from maybempi import MPI +from scipy.linalg import lu_factor, lu_solve +from scipy.sparse.linalg import spilu, splu + from struphy.polar.basic import PolarDerhamSpace, PolarVector logger = logging.getLogger("struphy") From 71818ca6c6da663b061290f4acbf16135088068e Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Mon, 12 Oct 2026 13:54:08 +0200 Subject: [PATCH 48/50] Complete polar spline GPU support after the Kronecker refactor Keep Kronecker operators, Schur/LU/ILU solves, and mass preconditioners on the device using cached factors. Handle device weights during setup and support synchronized or staged polar MPI collectives. Extend CPU, fake-CuPy, and GPU parity tests to cover boundary conditions, transposed solves, MPI decompositions, and transfers after warmup. Update the polar GPU documentation and CI coverage. --- .github/workflows/test-PR-unit-mpi.yml | 6 + CUDA_STRATEGY.md | 35 ++- doc/sections/subsections/numerics-geomFE.rst | 7 + src/struphy/feec/preconditioner.py | 2 +- src/struphy/feec/projectors.py | 2 +- src/struphy/polar/basic.py | 17 +- src/struphy/polar/kron_2dx1d.py | 157 ++++++++---- src/struphy/polar/linear_operators.py | 8 +- src/struphy/polar/tests/test_polar_cupy.py | 244 ++++++++++++++++--- 9 files changed, 369 insertions(+), 109 deletions(-) diff --git a/.github/workflows/test-PR-unit-mpi.yml b/.github/workflows/test-PR-unit-mpi.yml index 7bb53fbe2..6a5b1c387 100644 --- a/.github/workflows/test-PR-unit-mpi.yml +++ b/.github/workflows/test-PR-unit-mpi.yml @@ -82,3 +82,9 @@ jobs: source /struphy_fortran_/env_fortran_/bin/activate cd ${{ env.STRUPHY_PATH }} mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi pic/tests/test_sorting_device.py + + - name: Run polar GPU MPI paths on the fake CuPy + run: | + source /struphy_fortran_/env_fortran_/bin/activate + cd ${{ env.STRUPHY_PATH }} + mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi polar/tests/test_polar_cupy.py -k mpi_fake diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 84c70890b..5a094dacd 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -25,7 +25,7 @@ Progress is tracked in struphy-hub/struphy#650. - [ ] **PR 16: One CPU and one GPU argument class per argument type** — every pyccel argument class (`MarkerArguments`, `DerhamArguments`, `DomainArguments`, `LocalProjectorsArguments`) has a CUDA counterpart of the same name with a `Cuda` prefix in `kernel_arguments/*_cuda.py`, with the same constructor and attributes. Each owner creates one of the two in `__init__`, depending on the backend; `PyccelStructArguments`, `__host_args__()` and host copies are no longer used. Every kernel that takes argument objects is called through `Kernel(PyccelKernel(...))`; without a CUDA version it raises on CuPy. The remaining entry kernels called from Python (spline, geometry and SPH evaluation, marker diagnostics, `reflect`, local projectors) move into one folder per kernel, like the pushers. The parity-test inputs move out of the kernel folders into `pic/tests/cuda_parity_cases.py`. Geometry evaluations have no CUDA version yet, so CuPy particle runs fail at setup until they are ported (see [PR 16 implementation notes](#pr-16-implementation-notes)). - [ ] **PR 17: struphy on the feectools CUDA stack** — the `feectools` submodule points at the top of the feectools CUDA stack (`cuda-4-device-kernels`, [feectools#88](https://github.com/struphy-hub/feectools/pull/88)) instead of `devel-tiny`, so the struphy CUDA PRs run against feectools with device stencil operations, MPI with device buffers and one GPU per rank. Moves along with the stack; before the struphy stack is merged into `devel`, the feectools stack is merged into `devel-tiny` and the submodule points there again (see [feectools](#feectools)). - [x] **PR 18: Geometry evaluation on the GPU for all analytic mappings** (#681, open) — CUDA versions of the four geometry entry kernels (`kernel_evaluate_pic`, `kernel_evaluate`, `kernel_pullpush_pic`, `kernel_pullpush` in `geometry/kernels/`), built on device versions of the whole metric chain (`f`, `df`, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`/`push`/`tran`) for every analytic mapping (`kind_map` 10–12, 20–22, 30–32). Restores CuPy particle runs (weight initialization evaluates `jacobian_det`), and removes the Cuboid-only checks in `Pusher`, the accumulators and `reflect`. Parity arguments cover every analytic mapping (see [PR 18](#pr-18-geometry-evaluation-for-all-analytic-mappings) and the [PR 18 implementation notes](#pr-18-implementation-notes)). -- [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines and `EQDSKequilibrium` stay host-only (see [PR 19](#pr-19-spline-mappings) and the [PR 19 implementation notes](#pr-19-implementation-notes)). +- [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines were host-only at this stage; their current CuPy support is described in [Polar splines on CuPy](#polar-splines-on-cupy-695). See also the [PR 19 implementation notes](#pr-19-implementation-notes). - [x] **`linear_vlasov_ampere` on the GPU** (part of #688) — the first matrix accumulation: CUDA version of `linear_vlasov_ampere` (`EfieldWeightsCoupling`) on cunumpy's `Array6D` views (cunumpy 0.6.1), device versions of `fill_mat`, `fill_mat_vec`, `m_v_fill_b_v1_symm` and `outer`, and the matrix path of `Accumulator` checked on the CuPy backend (device data, ghost regions, transposed blocks, Schur operator) without a GPU through cunumpy's fake CuPy with emulated kernel launches (see the [implementation notes](#linear_vlasov_ampere-implementation-notes)). - [ ] **Next (order to be confirmed)**: one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, which needs the [feectools stack](#feectools) merged and released first; the guiding-center kernels (step 3), by hand (decided in #690, see [Hand-written CUDA](#hand-written-cuda-decision-for-690)). - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)). @@ -808,12 +808,29 @@ the [porting order](#porting-order); part of #688). `xp.scipy.sparse`), made once per block list and dropped by the setters. The blocks act only on the first two or three radial rings (at most `3 n2 x 3 n2` in `eta1-eta2`, `n3 x n3` in `eta3`), so they are small; blocks without non-zeros are not copied. On NumPy nothing changes (the SciPy blocks are applied as before). -- **MPI.** The ring reductions (`dot_inner_tp_rings`, `dot_parts_of_polar`, `PolarVector.toarray`) call - `cunumpy.mpi.synchronize_for_mpi` before `Allreduce`, which gets device buffers on CuPy (CUDA-aware MPI). +- **Refactored projectors and preconditioners (#746/#747).** The 2D x 1D Kronecker operators and + `PolarMassMatrixPreconditioner` keep coefficients and work arrays on the active backend. SciPy builds LU/ILU + factors on the host at setup. On CuPy, cached `cupyx.scipy.sparse.linalg.SuperLU` objects apply those factors + with device permutations and batched sparse triangular solves. Small 1d factors and the pole Schur complement + use cached dense inverses and device matrix products. The full poloidal inverse is never made dense. + Projectors support `schur`, `lu`, `ilu`, boundary conditions and transposed solves on both backends. +- **Mass setup.** Callable weights are evaluated with coordinates on the active backend. Sampled weights and + spline quadrature data are copied explicitly to the host for the 2d reduction/assembly and factorization. + These are setup transfers; repeated operator/preconditioner applications have no coefficient transfers on + a single rank or with CUDA-aware MPI. Krylov methods still synchronize convergence scalars. +- **MPI.** Kronecker `Alltoallv` redistributions, ring reductions and `PolarVector.toarray(allreduce=True)` + use `cunumpy.mpi.mpi_buffer`: synchronized device buffers with CUDA-aware MPI, host staging otherwise. + MPI capability must be established at startup through cunumpy. The pole sub-communicator and collective + ghost-region flags from #746 are preserved. feectools' own halo exchanges still require CUDA-aware MPI + on real GPUs; staging of the polar collectives does not replace that separate dependency. - **Tests.** `polar/tests/test_polar_cupy.py` compares NumPy and CuPy for `IGAPolarCylinder` and `IGAPolarTorus`: - extraction and DOF extraction operators and their transposes, `grad`/`curl`/`div` and transposes, `PolarVector` - arithmetic, the polar projectors, polar mass matrices (GPU only) and polar `SplineFunction` coefficients - (`extract_coeffs`, restart). Without a GPU it runs on the fake CuPy with two substitutes: a dense stand-in for - `cupyx.scipy.sparse` (`set_device_sparse_module`) and the host versions of the feectools kernels (the fake cannot - launch CUDA kernels). The mass matrices need struphy's CUDA geometry kernels and are compared on the GPU only. - Field evaluation (`SplineFunction.__call__`) is not covered: it is not ported to CuPy for any space yet. + extraction/DOF operators, derivatives and transposes, vector arithmetic, all three projector solvers, + transposed and boundary-condition solves, mass matrices/preconditioners, callable and sampled weights, + and spline coefficient extraction/restart. Repeated Kronecker and preconditioner applications are checked + for transfers after warmup. The fake CuPy tests use CPU-emulated CUDA launches (including geometry/mass + kernels), with dense device-array substitutes for cuSPARSE products and triangular solves. + MPI parity tests decompose independently in eta1, eta2 and eta3 and exercise both device-buffer and + staged polar collectives; CI runs them on 1, 2, 3 and 4 ranks. Real-GPU variants are included but require + GPU hardware and the feectools CUDA stack. +- **Separate follow-ups.** `SplineFunction.__call__` field evaluation is not ported to CuPy for any space yet. + FFT eta3 solves and sharing factorizations across different projector objects are additional optimizations. diff --git a/doc/sections/subsections/numerics-geomFE.rst b/doc/sections/subsections/numerics-geomFE.rst index 7fb728c51..54f9b2d97 100644 --- a/doc/sections/subsections/numerics-geomFE.rst +++ b/doc/sections/subsections/numerics-geomFE.rst @@ -289,4 +289,11 @@ For weighted mass matrices, use :class:`~struphy.feec.preconditioner.PolarMassMa (``precond="PolarMassMatrixPreconditioner"`` in the propagator options); it is exact for axisymmetric weights and mappings. +On the CuPy backend, extraction, derivatives, Kronecker operators and these solves keep their +coefficients on the device. Factorization and mass-preconditioner assembly run once on the host; +cached device factors apply subsequent solves. Polar MPI collectives support CUDA-aware device +buffers or host staging through ``cunumpy.mpi.mpi_buffer``. The feectools halo exchanges require +CUDA-aware MPI on real GPUs. General ``SplineFunction.__call__`` field evaluation is a separate +GPU port and is not supported yet. + .. image:: ../../pics/polar_derham.png \ No newline at end of file diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 87ca0bfd2..abccfb449 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -549,7 +549,7 @@ def _polar_mass_groups(mass_operator: WeightedMassOperator, reduction: str) -> l A_pol = (E @ M12 @ E.T).tocsr() # unweighted 1d mass matrix in eta_3 - A_3 = _mass_matrix_1d(mass_operator, bases[comps[0]][2], 2, _ones_1d)[0].toarray() + A_3 = xp.to_numpy(_mass_matrix_1d(mass_operator, bases[comps[0]][2], 2, _ones_1d)[0].toarray()) groups.append(Kron2Dx1DGroup(comps, A_pol, A_3)) return groups diff --git a/src/struphy/feec/projectors.py b/src/struphy/feec/projectors.py index 06e780733..1f87674ac 100644 --- a/src/struphy/feec/projectors.py +++ b/src/struphy/feec/projectors.py @@ -615,7 +615,7 @@ def solve(self, rhs, transposed=False, apply_bc=False, out=None, x0=None): krylov = self._polar_krylov[key] # relative tolerance w.r.t. the norm of the right-hand side - rhs_norm = xp.sqrt(rhs.inner(rhs)) + rhs_norm = float(xp.sqrt(rhs.inner(rhs))) krylov.set_options(x0=x0, tol=self._polar_rtol * rhs_norm if rhs_norm > 0.0 else self._polar_rtol) x = krylov.solve(rhs, out=out) diff --git a/src/struphy/polar/basic.py b/src/struphy/polar/basic.py index b3a349ea7..b6fb1fb6a 100644 --- a/src/struphy/polar/basic.py +++ b/src/struphy/polar/basic.py @@ -1,7 +1,7 @@ from typing import TYPE_CHECKING import cunumpy as xp -from cunumpy.mpi import synchronize_for_mpi +from cunumpy.mpi import mpi_buffer from feectools.linalg.basic import Vector, VectorSpace from feectools.linalg.block import BlockVector from feectools.linalg.stencil import StencilVector @@ -347,9 +347,9 @@ def toarray(self, allreduce=False): out = self.tp.toarray()[self.space.n_rings[0] * self.space.n2[0] * self.space.n3[0] :] # allreduce tensor-product part - if self.space.comm is not None and allreduce: - synchronize_for_mpi(out) - self.space.comm.Allreduce(MPI.IN_PLACE, out, op=MPI.SUM) + if self.space.comm is not None and self.space.comm.Get_size() > 1 and allreduce: + with mpi_buffer(out, recv=True) as buf: + self.space.comm.Allreduce(MPI.IN_PLACE, buf, op=MPI.SUM) out = xp.concatenate((self.pol[0].flatten(), out)) @@ -359,11 +359,10 @@ def toarray(self, allreduce=False): out3 = self.tp[2].toarray()[self.space.n_rings[2] * self.space.n2[2] * self.space.n3[2] :] # allreduce tensor-product part - if self.space.comm is not None and allreduce: - synchronize_for_mpi(out1, out2, out3) - self.space.comm.Allreduce(MPI.IN_PLACE, out1, op=MPI.SUM) - self.space.comm.Allreduce(MPI.IN_PLACE, out2, op=MPI.SUM) - self.space.comm.Allreduce(MPI.IN_PLACE, out3, op=MPI.SUM) + if self.space.comm is not None and self.space.comm.Get_size() > 1 and allreduce: + for array in (out1, out2, out3): + with mpi_buffer(array, recv=True) as buf: + self.space.comm.Allreduce(MPI.IN_PLACE, buf, op=MPI.SUM) out = xp.concatenate( ( diff --git a/src/struphy/polar/kron_2dx1d.py b/src/struphy/polar/kron_2dx1d.py index 07cb33f79..6c89c80a6 100644 --- a/src/struphy/polar/kron_2dx1d.py +++ b/src/struphy/polar/kron_2dx1d.py @@ -44,6 +44,7 @@ from scipy.sparse.linalg import spilu, splu from struphy.polar.basic import PolarDerhamSpace, PolarVector +from struphy.polar.linear_operators import DeviceSparseMatrix, device_sparse_module logger = logging.getLogger("struphy") @@ -214,6 +215,29 @@ def poloidal_extraction_matrix( # -------------------------------------------------------------------------------------- # Poloidal solvers # -------------------------------------------------------------------------------------- +class _FactorizedSolve: + """Host factorization with a cached device solve; dense inverses are only used for small 1d factors.""" + + def __init__(self, factor, dense=False): + self._factor = factor + self._dense = dense + self._device = None + + def solve(self, Y): + if Y.shape[-1] == 0: + return Y.copy() + if not xp.is_gpu(Y): + return self._factor.solve(np.asarray(Y, order="F")) + if self._device is None: + if self._dense: + self._device = xp.asarray(self._factor.solve(np.eye(self._factor.shape[0]))) + else: + # CuPy wraps the existing SciPy factors: factorization stays at + # setup, while permutations and batched triangular solves run on device. + self._device = device_sparse_module().linalg.SuperLU(self._factor) + return self._device @ Y if self._dense else self._device.solve(Y) + + class _SparseLU: """Exact (LU) or incomplete (ILU) factorization of a sparse matrix.""" @@ -224,12 +248,13 @@ def __init__(self, A: spa.spmatrix, kind: str = "lu", drop_tol=ILU_DROP_TOL, fil self._fac = splu(A) else: self._fac = spilu(A, drop_tol=drop_tol, fill_factor=fill_factor) + self._solve = _FactorizedSolve(self._fac) self.exact = kind == "lu" def solve(self, Y: np.ndarray) -> np.ndarray: if Y.shape[-1] == 0: return Y.copy() - return self._fac.solve(np.asarray(Y, order="F")) + return self._solve.solve(Y) class _KroneckerBlockDiagonal: @@ -240,18 +265,24 @@ def __init__(self, factors: list[tuple[np.ndarray, np.ndarray]]): self._lus = [] for A1, A2 in factors: self._sizes.append((A1.shape[0], A2.shape[0])) - self._lus.append((splu(spa.csc_matrix(A1)), splu(spa.csc_matrix(A2)))) + self._lus.append( + ( + _FactorizedSolve(splu(spa.csc_matrix(A1)), dense=True), + _FactorizedSolve(splu(spa.csc_matrix(A2)), dense=True), + ) + ) self.exact = True def solve(self, Y: np.ndarray) -> np.ndarray: - out = np.empty_like(Y) + arrays = xp if xp.is_gpu(Y) else np + out = arrays.empty_like(Y) m = Y.shape[-1] off = 0 for (n1, n2), (lu1, lu2) in zip(self._sizes, self._lus): X = Y[off : off + n1 * n2].reshape(n1, n2 * m) - X = lu1.solve(np.asarray(X, order="F")).reshape(n1, n2, m) - X = np.ascontiguousarray(X.transpose(1, 0, 2)).reshape(n2, n1 * m) - X = lu2.solve(np.asarray(X, order="F")).reshape(n2, n1, m) + X = lu1.solve(X).reshape(n1, n2, m) + X = arrays.ascontiguousarray(X.transpose(1, 0, 2)).reshape(n2, n1 * m) + X = lu2.solve(X).reshape(n2, n1, m) out[off : off + n1 * n2] = X.transpose(1, 0, 2).reshape(n1 * n2, m) off += n1 * n2 return out @@ -282,20 +313,37 @@ def __init__(self, A: spa.spmatrix, pole: np.ndarray, outer_solver, batch: int = W = self._outer.solve(self._Aop[:, k0:k1].toarray()) S[:, k0:k1] -= self._Apo @ W self._S = lu_factor(S) + self._device = None self._n = n logger.debug(f"Schur solver: {n = }, pole block {self._p.size}, outer block {self._o.size}") def solve(self, Y: np.ndarray) -> np.ndarray: if Y.shape[-1] == 0: return Y.copy() - Yp = Y[self._p] - Yo = Y[self._o] + if xp.is_gpu(Y): + if self._device is None: + self._device = ( + xp.asarray(self._p), + xp.asarray(self._o), + DeviceSparseMatrix(self._Apo), + DeviceSparseMatrix(self._Aop), + xp.asarray(lu_solve(self._S, np.eye(self._p.size))), + ) + p, o, Apo, Aop, Sinv = self._device + pole_solve = lambda rhs: Sinv @ rhs + arrays = xp + else: + p, o, Apo, Aop = self._p, self._o, self._Apo, self._Aop + pole_solve = lambda rhs: lu_solve(self._S, rhs) + arrays = np + Yp = Y[p] + Yo = Y[o] t = self._outer.solve(Yo) - Xp = lu_solve(self._S, Yp - self._Apo @ t) - Xo = self._outer.solve(Yo - self._Aop @ Xp) - X = np.empty_like(Y) - X[self._p] = Xp - X[self._o] = Xo + Xp = pole_solve(Yp - Apo.dot(t)) + Xo = self._outer.solve(Yo - Aop.dot(Xp)) + X = arrays.empty_like(Y) + X[p] = Xp + X[o] = Xo return X @@ -411,13 +459,14 @@ def forward(self, local: np.ndarray) -> np.ndarray: if self._comm is None: return local[..., c0:c1] - sendbuf = np.concatenate([local[..., a:b].ravel() for a, b in self._chunks]) - recvbuf = np.empty(int(self._recv_counts.sum()), dtype=float) - self._comm.Alltoallv( - [sendbuf, (self._send_counts, self._send_displs), MPI.DOUBLE], - [recvbuf, (self._recv_counts, self._recv_displs), MPI.DOUBLE], - ) - full = np.empty((*self._global_shape, c1 - c0), dtype=float) + sendbuf = xp.concatenate([local[..., a:b].ravel() for a, b in self._chunks]) + recvbuf = xp.empty(int(self._recv_counts.sum()), dtype=float) + with xp.mpi.mpi_buffer(sendbuf) as send, xp.mpi.mpi_buffer(recvbuf, send=False, recv=True) as recv: + self._comm.Alltoallv( + [send, (self._send_counts, self._send_displs), MPI.DOUBLE], + [recv, (self._recv_counts, self._recv_displs), MPI.DOUBLE], + ) + full = xp.empty((*self._global_shape, c1 - c0), dtype=float) for (s, e), d, n in zip(self._boxes, self._recv_displs, self._recv_counts): shape = tuple(ei - si + 1 for si, ei in zip(s, e)) full[tuple(slice(si, ei + 1) for si, ei in zip(s, e))] = recvbuf[d : d + n].reshape(*shape, c1 - c0) @@ -428,15 +477,16 @@ def backward(self, full: np.ndarray) -> np.ndarray: if self._comm is None: return full - sendbuf = np.concatenate( + sendbuf = xp.concatenate( [full[tuple(slice(si, ei + 1) for si, ei in zip(s, e))].ravel() for s, e in self._boxes], ) - recvbuf = np.empty(int(self._send_counts.sum()), dtype=float) - self._comm.Alltoallv( - [sendbuf, (self._recv_counts, self._recv_displs), MPI.DOUBLE], - [recvbuf, (self._send_counts, self._send_displs), MPI.DOUBLE], - ) - local = np.empty((*self._local_shape, self._m_local), dtype=float) + recvbuf = xp.empty(int(self._send_counts.sum()), dtype=float) + with xp.mpi.mpi_buffer(sendbuf) as send, xp.mpi.mpi_buffer(recvbuf, send=False, recv=True) as recv: + self._comm.Alltoallv( + [send, (self._recv_counts, self._recv_displs), MPI.DOUBLE], + [recv, (self._send_counts, self._send_displs), MPI.DOUBLE], + ) + local = xp.empty((*self._local_shape, self._m_local), dtype=float) for (a, b), d, n in zip(self._chunks, self._send_displs, self._send_counts): local[..., a:b] = recvbuf[d : d + n].reshape(*self._local_shape, b - a) return local @@ -574,22 +624,19 @@ def _apply(self, v: PolarVector, out: PolarVector | None, ops_3: list, ops_pol: assert out.space == self._space V = self._space - # These Kronecker solvers and MPI transposes use SciPy/NumPy on the host. - # Move backend arrays explicitly at this boundary; implicit CuPy-to-NumPy - # conversions are rejected by CuPy and by cunumpy's fake backend. - out_pol = [np.zeros_like(xp.to_numpy(p)) for p in v.pol] + out_pol = [xp.zeros_like(p) for p in v.pol] for g, op_3, op_pol in zip(self._groups, ops_3, ops_pol): # 1. eta_3 direction (tensor-product part redistributed, polar coeffs replicated) tmp_tp = {} tmp_pol = {} for c in g.comps: - L = xp.to_numpy(_interior(tp_component(v.tp, c))) + L = _interior(tp_component(v.tp, c)) l1, l2, l3 = L.shape T3 = self._t3[c] - X = T3.forward(np.ascontiguousarray(L.reshape(l1 * l2, l3).T)) - tmp_tp[c] = np.ascontiguousarray(T3.backward(op_3(X)).T).reshape(l1, l2, l3) - pol = xp.to_numpy(v.pol[c]) + X = T3.forward(xp.ascontiguousarray(L.reshape(l1 * l2, l3).T)) + tmp_tp[c] = xp.ascontiguousarray(T3.backward(op_3(X)).T).reshape(l1, l2, l3) + pol = v.pol[c] tmp_pol[c] = op_3(pol.T).T if V.n_polar[c] > 0 else pol.copy() # 2. poloidal plane @@ -602,7 +649,7 @@ def _apply(self, v: PolarVector, out: PolarVector | None, ops_3: list, ops_pol: nr = V.n_rings[c] blocks.append(tmp_pol[c][:, s3 + k0 : s3 + k1]) blocks.append(full[c][nr:].reshape(-1, k1 - k0)) - Z = op_pol(np.concatenate(blocks, axis=0)) if k1 > k0 else np.concatenate(blocks, axis=0) + Z = op_pol(xp.concatenate(blocks, axis=0)) if k1 > k0 else xp.concatenate(blocks, axis=0) off = 0 for c in g.comps: @@ -616,15 +663,16 @@ def _apply(self, v: PolarVector, out: PolarVector | None, ops_3: list, ops_pol: F[nr:] = Z[off : off + (n1 - nr) * n2].reshape(n1 - nr, n2, k1 - k0) off += (n1 - nr) * n2 u = tp_component(out.tp, c) - _interior(u)[...] = xp.asarray(self._tpl[c].backward(F)) + _interior(u)[...] = self._tpl[c].backward(F) u.ghost_regions_in_sync = False # polar coefficients are held by all processes (each column was computed on exactly one process) comm = V.comm for c in range(V.n_comps): if comm is not None and comm.Get_size() > 1 and out_pol[c].size > 0: - comm.Allreduce(MPI.IN_PLACE, out_pol[c], op=MPI.SUM) - out.pol[c][:] = xp.asarray(out_pol[c]) + with xp.mpi.mpi_buffer(out_pol[c], recv=True) as buf: + comm.Allreduce(MPI.IN_PLACE, buf, op=MPI.SUM) + out.pol[c][:] = out_pol[c] return out @@ -643,8 +691,14 @@ class PolarKron2Dx1DOperator(_PolarKron2Dx1DBase): """ def dot(self, v: PolarVector, out: PolarVector | None = None) -> PolarVector: - ops_3 = [lambda X, A=g.A_3: A @ X for g in self._groups] - ops_pol = [lambda Y, A=g.A_pol: A @ Y for g in self._groups] + if xp.is_gpu(v.pol[0]): + if not hasattr(self, "_device_factors"): + self._device_factors = [(xp.asarray(g.A_3), DeviceSparseMatrix(g.A_pol)) for g in self._groups] + factors = self._device_factors + else: + factors = [(g.A_3, g.A_pol) for g in self._groups] + ops_3 = [lambda X, A=A: A @ X for A, _ in factors] + ops_pol = [A.dot for _, A in factors] return self._apply(v, out, ops_3, ops_pol) def transpose(self, conjugate=False) -> "PolarKron2Dx1DOperator": @@ -678,7 +732,7 @@ def __init__( super().__init__(space, groups) self._kind = poloidal_solver self._solver_kwargs = solver_kwargs - self._lu_3 = [splu(spa.csc_matrix(g.A_3)) for g in groups] + self._lu_3 = [_FactorizedSolve(splu(spa.csc_matrix(g.A_3)), dense=True) for g in groups] self._pol_solvers = [ make_poloidal_solver( g.A_pol, @@ -700,7 +754,7 @@ def poloidal_solver(self) -> str: return self._kind def dot(self, v: PolarVector, out: PolarVector | None = None) -> PolarVector: - ops_3 = [lambda X, lu=lu: lu.solve(np.asarray(X, order="F")) for lu in self._lu_3] + ops_3 = [lu.solve for lu in self._lu_3] ops_pol = [s.solve for s in self._pol_solvers] return self._apply(v, out, ops_3, ops_pol) @@ -787,7 +841,7 @@ def _serial_attrs_1d(derham, basis: str, d: int): def _global_quad_pts(derham, d: int) -> np.ndarray: """Global quadrature points (flattened over the cells) in direction d.""" - return np.asarray(derham.H1_1d_serial_splines[d].quad_grid_pts[0][0]).ravel() + return xp.to_numpy(derham.H1_1d_serial_splines[d].quad_grid_pts[0][0]).ravel() def reduced_weight_2d(weight, derham, reduction: str = "midpoint") -> np.ndarray | None: @@ -825,9 +879,12 @@ def reduced_weight_2d(weight, derham, reduction: str = "midpoint") -> np.ndarray nodes, wts = np.polynomial.legendre.leggauss(_N_AVG_3) pts3, wts3 = (nodes + 1.0) / 2.0, wts / 2.0 E1, E2, E3 = np.meshgrid(pts1, pts2, pts3, indexing="ij") - return np.asarray(weight(E1, E2, E3)).reshape(E1.shape) @ wts3 + return xp.to_numpy(weight(xp.asarray(E1), xp.asarray(E2), xp.asarray(E3))).reshape(E1.shape) @ wts3 - assert isinstance(weight, np.ndarray), f"Weight must be callable, np.ndarray or None, but is {type(weight)}." + assert isinstance(weight, np.ndarray) or xp.is_gpu(weight), ( + f"Weight must be callable, an array or None, but is {type(weight)}." + ) + weight = xp.to_numpy(weight) dom = derham.domain_decomposition nq = derham.nquads s1, s2, s3 = (dom.starts[d] * nq[d] for d in range(3)) @@ -841,7 +898,7 @@ def reduced_weight_2d(weight, derham, reduction: str = "midpoint") -> np.ndarray sums[0, s1 : s1 + l1, s2 : s2 + l2] = weight[:, :, mid] sums[1, s1 : s1 + l1, s2 : s2 + l2] = 1.0 else: - wts3 = np.asarray(derham.spline_attributes["H1"].quad_grid_wts[0][2]).ravel() + wts3 = xp.to_numpy(derham.spline_attributes["H1"].quad_grid_wts[0][2]).ravel() assert wts3.size == l3 sums[0, s1 : s1 + l1, s2 : s2 + l2] = weight @ wts3 sums[1, s1 : s1 + l1, s2 : s2 + l2] = wts3.sum() @@ -872,9 +929,9 @@ def poloidal_mass_matrix(derham, bases_row: tuple, bases_col: tuple, weight: np. att_j = [_serial_attrs_1d(derham, bases_col[d], d) for d in range(2)] def data(att): - spans = np.asarray(att.quad_grid_spans[0][0]) - basis = np.asarray(att.quad_grid_bases[0][0])[:, :, 0, :] # (cells, p+1, quad points), no derivatives - wts = np.asarray(att.quad_grid_wts[0][0]) # (cells, quad points) + spans = xp.to_numpy(att.quad_grid_spans[0][0]) + basis = xp.to_numpy(att.quad_grid_bases[0][0])[:, :, 0, :] # (cells, p+1, quad points), no derivatives + wts = xp.to_numpy(att.quad_grid_wts[0][0]) # (cells, quad points) n = att.nbasis[0][0] p = basis.shape[1] - 1 return spans, basis, wts, n, p diff --git a/src/struphy/polar/linear_operators.py b/src/struphy/polar/linear_operators.py index 288286240..4e39f7797 100644 --- a/src/struphy/polar/linear_operators.py +++ b/src/struphy/polar/linear_operators.py @@ -1,5 +1,5 @@ import cunumpy as xp -from cunumpy.mpi import synchronize_for_mpi +from cunumpy.mpi import mpi_buffer from feectools.linalg.basic import LinearOperator from feectools.linalg.block import BlockVector, BlockVectorSpace from feectools.linalg.stencil import StencilVector, StencilVectorSpace @@ -110,9 +110,9 @@ def _drop_device_blocks(self, name): def _allreduce_sum(comm, buf): """In-place sum over the ranks of comm (buf may live on the device: CUDA-aware MPI).""" - if comm is not None: - synchronize_for_mpi(buf) - comm.Allreduce(MPI.IN_PLACE, buf, op=MPI.SUM) + if comm is not None and comm.Get_size() > 1: + with mpi_buffer(buf, recv=True) as buffer: + comm.Allreduce(MPI.IN_PLACE, buffer, op=MPI.SUM) class PolarExtractionOperator(_BlocksOnBackend, LinearOperator): diff --git a/src/struphy/polar/tests/test_polar_cupy.py b/src/struphy/polar/tests/test_polar_cupy.py index 4b26bac3a..2e7dc2c5f 100644 --- a/src/struphy/polar/tests/test_polar_cupy.py +++ b/src/struphy/polar/tests/test_polar_cupy.py @@ -6,12 +6,11 @@ compares extraction operators, derivatives, ``PolarVector`` arithmetic, mass matrices, projectors and polar ``SplineFunction`` coefficients. -Without a GPU it runs on cunumpy's fake CuPy (which rejects host/device mixing like CuPy), in a subprocess, with -two substitutes for what the fake cannot do: the device sparse module (``cupyx.scipy.sparse``) is replaced by -:class:`DenseDeviceSparse` (dense products of device arrays), and kernels are called in their host version (the -fake cannot launch CUDA kernels; the feectools stencil kernels then copy their arrays to the host and back). +Without a GPU it runs on cunumpy's fake CuPy with CPU-emulated CUDA launches and a dense device stand-in +for cupyx sparse products and triangular solves. Host/device transfers remain visible to cunumpy's counters. """ +import os import subprocess import sys @@ -39,18 +38,32 @@ def __init__(self, host): def dot(self, x): return self._array @ x + class linalg: + class SuperLU: + """Device-array triangular solves with the same permutations as CuPy's SuperLU.""" -def make_domain(mapping): + def __init__(self, factor): + self.L = cunumpy.asarray(factor.L.toarray()) + self.U = cunumpy.asarray(factor.U.toarray()) + self.rows = cunumpy.asarray(np.argsort(factor.perm_r)) + self.cols = cunumpy.asarray(factor.perm_c) + + def solve(self, rhs): + y = cunumpy.linalg.solve(self.L, rhs[self.rows]) + return cunumpy.linalg.solve(self.U, y)[self.cols] + + +def make_domain(mapping, num_elements=(6, 6)): from struphy import domains if mapping == "IGAPolarCylinder": - return domains.IGAPolarCylinder(num_elements=(6, 6), degree=(2, 3), a=1.2, Lz=3.0) + return domains.IGAPolarCylinder(num_elements=num_elements, degree=(2, 3), a=1.2, Lz=3.0) if mapping == "IGAPolarTorus": - return domains.IGAPolarTorus(num_elements=(6, 6), degree=(3, 2), a=0.8, R0=3.0, sfl=True) + return domains.IGAPolarTorus(num_elements=num_elements, degree=(3, 2), a=0.8, R0=3.0, sfl=True) raise ValueError(mapping) -def make_polar_derham(mapping, with_mass=True): +def make_polar_derham(mapping, with_mass=True, poloidal_solver="schur", mpi_dims_mask=None): """Polar Derham (with projectors), its domain and mass matrices (or None), on the active backend.""" from maybempi import MPI @@ -59,14 +72,25 @@ def make_polar_derham(mapping, with_mass=True): from struphy.io.options import DerhamOptions from struphy.topology.grids import TensorProductGrid - domain = make_domain(mapping) + size = MPI.COMM_WORLD.Get_size() + num_elements = (max(6, 4 * size), max(6, 3 * size), 4) + domain = make_domain(mapping, num_elements[:2]) options = DerhamOptions( degree=(domain.degree[0], domain.degree[1], 1), bcs=(("dirichlet", "free"), None, None), nquads_proj=(4, 4, 2), polar_splines=True, + polar_poloidal_solver=poloidal_solver, + ) + derham = Derham( + TensorProductGrid( + num_elements=num_elements, + mpi_dims_mask=mpi_dims_mask, + ), + options, + comm=MPI.COMM_WORLD, + domain=domain, ) - derham = Derham(TensorProductGrid(num_elements=(6, 6, 4)), options, comm=MPI.COMM_WORLD, domain=domain) return domain, derham, WeightedMassOperators(derham, domain) if with_mass else None @@ -78,10 +102,14 @@ def _stencils(vec): return [tp] if isinstance(tp, StencilVector) else list(tp.blocks) -def to_host(vec): +def to_host(vec, interior=False): """Data of a Stencil-/Block-/PolarVector as a list of NumPy arrays (polar coeffs first).""" pol = list(vec.pol) if hasattr(vec, "pol") else [] - return [cunumpy.to_numpy(a) for a in pol] + [cunumpy.to_numpy(s._data) for s in _stencils(vec)] + from struphy.polar.kron_2dx1d import _interior + + return [cunumpy.to_numpy(a) for a in pol] + [ + cunumpy.to_numpy(_interior(s) if interior else s._data) for s in _stencils(vec) + ] def fill(vec, seed): @@ -104,18 +132,40 @@ def assert_on_backend(vec, backend): assert cunumpy.is_gpu(a) == (backend == "cupy"), type(a) -def polar_results(mapping, backend, with_mass=True): +def polar_results(mapping, backend, with_mass=True, poloidal_solver="schur", mpi_dims_mask=None): """Results of the polar operations on `backend`, as host arrays (dict name -> list of arrays).""" from struphy.polar.basic import PolarVector res = {} with cunumpy.use_backend(backend): - domain, derham, mass = make_polar_derham(mapping, with_mass) + domain, derham, mass = make_polar_derham(mapping, with_mass, poloidal_solver, mpi_dims_mask) def record(name, vec): assert_on_backend(vec, backend) res[name] = to_host(vec) + # Host assembly accepts both backend callables and device-resident sampled weights. + from struphy.polar.kron_2dx1d import reduced_weight_2d + + pts = [cunumpy.asarray(a).ravel() for a in derham.spline_attributes["H1"].quad_grid_pts[0]] + mesh = cunumpy.meshgrid(*pts, indexing="ij") + for reduction in ("midpoint", "average"): + + def weight(e1, e2, e3): + return 1.0 + e1 + e2 + (e3 if reduction == "average" else 0.0) + + sampled = weight(*mesh) + from_array = reduced_weight_2d(sampled, derham, reduction) + from_function = reduced_weight_2d(weight, derham, reduction) + np.testing.assert_allclose(from_array, from_function, rtol=1e-12, atol=1e-12) + res["weight-" + reduction] = [from_array] + if with_mass: + from struphy.feec.preconditioner import PolarMassMatrixPreconditioner + + W = mass.create_weighted_mass("H1", "H1", weights=[[sampled]], assemble=True) + pc = PolarMassMatrixPreconditioner(W, poloidal_solver=poloidal_solver, weight_reduction=reduction) + record("array-pc-" + reduction, pc.dot(fill(W.domain.zeros(), seed=70))) + for n, form in enumerate(FORMS): E = derham.extraction_ops[form] x = fill(E.domain.zeros(), seed=10 + n) @@ -139,11 +189,21 @@ def record(name, vec): w *= 3.0 record(f"inplace{form}", w) res[f"dot{form}"] = [np.asarray(float(y.dot(z)))] - res[f"toarray{form}"] = [cunumpy.to_numpy(y.toarray())] + res[f"toarray{form}"] = [cunumpy.to_numpy(y.toarray(allreduce=True))] # mass matrices (E M E^T with the polar extraction operators) if with_mass: - record(f"M{form}", getattr(mass, "M" + form).dot(y)) + from struphy.feec.preconditioner import PolarMassMatrixPreconditioner + + M = getattr(mass, "M" + form) + record(f"M{form}", M.dot(y)) + pc = PolarMassMatrixPreconditioner(M, poloidal_solver=poloidal_solver) + record(f"pcM{form}", pc.dot(M.dot(y))) + # Factor/device caches are populated by the first application. + if backend == "cupy" and (derham.comm.Get_size() == 1 or cunumpy.mpi.get_mpi_cuda_aware()): + with cunumpy.profiling.count_transfers() as transfers: + pc.dot(y) + assert transfers.total == 0, transfers.report() for name, op in (("grad", derham.grad), ("curl", derham.curl), ("div", derham.div)): x = fill(PolarVector(op.domain), seed=40) @@ -151,7 +211,30 @@ def record(name, vec): record(name, op.dot(x)) record(name + "T", op.transpose().dot(y)) - # commuting projectors (polar case: iterative solve with the polar extraction operators) + # Direct Kronecker projectors, or ILU-preconditioned iterations. + for form in ("0", "1", "2", "3"): + P = getattr(derham, "P" + form) + x = fill(P.I.domain.zeros(), seed=60) + for transposed in (False, True): + for apply_bc in (False, True): + target = P.boundary_op.dot(x) if apply_bc else x + op = (P.I0T if apply_bc else P.IT) if transposed else (P.I0 if apply_bc else P.I) + rhs = op.dot(target) + solved = P.solve(rhs, transposed=transposed, apply_bc=apply_bc) + solved.update_ghost_regions() + target.update_ghost_regions() + record(f"solve{form}-{transposed}-{apply_bc}", solved) + for actual, expected in zip(to_host(solved, interior=True), to_host(target, interior=True)): + np.testing.assert_allclose(actual, expected, rtol=1e-9, atol=1e-10) + # Test operator and preconditioner applications without counting Krylov scalar reductions. + P.pc.dot(x) + if backend == "cupy" and (derham.comm.Get_size() == 1 or cunumpy.mpi.get_mpi_cuda_aware()): + with cunumpy.profiling.count_transfers() as transfers: + P.I.dot(x) + P.pc.dot(x) + assert transfers.total == 0, transfers.report() + + # Commuting projection of analytic functions. def f(e1, e2, e3): return cunumpy.sin(np.pi * e1) * cunumpy.cos(2 * np.pi * e2) * cunumpy.cos(2 * np.pi * e3) + e1**2 @@ -176,10 +259,10 @@ def f(e1, e2, e3): return res -def check_polar_backends(mapping, with_mass=True): +def check_polar_backends(mapping, with_mass=True, poloidal_solver="schur", mpi_dims_mask=None): """The polar operations agree on the NumPy and the CuPy backend.""" - host = polar_results(mapping, "numpy", with_mass) - device = polar_results(mapping, "cupy", with_mass) + host = polar_results(mapping, "numpy", with_mass, poloidal_solver, mpi_dims_mask) + device = polar_results(mapping, "cupy", with_mass, poloidal_solver, mpi_dims_mask) assert host.keys() == device.keys() for name in host: assert len(host[name]) == len(device[name]), name @@ -187,19 +270,17 @@ def check_polar_backends(mapping, with_mass=True): np.testing.assert_allclose(b, a, rtol=1e-9, atol=1e-11, err_msg=name) -def check_polar_backends_fake_cupy(mapping): - """:func:`check_polar_backends` on the fake CuPy (see the module docstring), without the mass matrices. - - The mass matrix weights are evaluated by struphy's CUDA geometry kernels, which take CUDA argument objects - and cannot fall back to their pyccel versions; the polar mass matrices are compared on the GPU only. - """ - from cunumpy._dispatch import Kernel - +def check_polar_backends_fake_cupy(mapping, poloidal_solver="schur"): + """Full polar parity, including mass matrices, with CPU-emulated CUDA launches.""" + from struphy.pic.tests.cuda_emulation import emulated_launches from struphy.polar import linear_operators - Kernel.get_kernel = lambda self: self._host_kernel - linear_operators.set_device_sparse_module(DenseDeviceSparse) - check_polar_backends(mapping, with_mass=False) + previous = linear_operators.set_device_sparse_module(DenseDeviceSparse) + try: + with emulated_launches(): + check_polar_backends(mapping, poloidal_solver=poloidal_solver) + finally: + linear_operators.set_device_sparse_module(previous) def test_device_blocks_numpy_backend(): @@ -233,14 +314,15 @@ def test_device_blocks_numpy_backend(): @pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") @pytest.mark.parametrize("mapping", MAPPINGS) -def test_polar_fake_cupy(mapping): +@pytest.mark.parametrize("poloidal_solver", ("schur", "lu", "ilu")) +def test_polar_fake_cupy(mapping, poloidal_solver): """Without a GPU: polar Derham on cunumpy's fake CuPy agrees with the NumPy backend. Runs in a subprocess because the fake CuPy must be installed before cunumpy is imported. """ code = ( "from struphy.polar.tests.test_polar_cupy import check_polar_backends_fake_cupy; " - f"check_polar_backends_fake_cupy({mapping!r})" + f"check_polar_backends_fake_cupy({mapping!r}, {poloidal_solver!r})" ) result = subprocess.run( [sys.executable, "-c", code], env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), capture_output=True, text=True @@ -250,6 +332,98 @@ def test_polar_fake_cupy(mapping): @requires_cupy @pytest.mark.parametrize("mapping", MAPPINGS) -def test_polar_on_cupy(mapping): +@pytest.mark.parametrize("poloidal_solver", ("schur", "lu", "ilu")) +def test_polar_on_cupy(mapping, poloidal_solver): """Polar Derham on the GPU (cupyx.scipy.sparse for the polar blocks) agrees with the NumPy backend.""" - check_polar_backends(mapping) + check_polar_backends(mapping, poloidal_solver=poloidal_solver) + + +def check_polar_mpi(cuda_aware, mpi_dims_mask, fake=True): + """Backend parity with both Kronecker redistributions and ranks away from the pole.""" + from contextlib import nullcontext + + from struphy.pic.tests.cuda_emulation import emulated_launches + from struphy.polar.linear_operators import set_device_sparse_module + + previous_mpi = cunumpy.mpi.get_mpi_cuda_aware() + previous_sparse = set_device_sparse_module(DenseDeviceSparse) if fake else None + try: + if cuda_aware is None: + # feectools halo exchanges require CUDA-aware MPI. Never probe + # an unknown MPI build with real device addresses (it may crash). + if os.environ.get("STRUPHY_TEST_CUDA_AWARE_MPI") != "1": + pytest.skip("full GPU MPI parity needs STRUPHY_TEST_CUDA_AWARE_MPI=1 on a CUDA-aware MPI build") + cuda_aware = True + cunumpy.mpi.set_mpi_cuda_aware(cuda_aware) + with emulated_launches() if fake else nullcontext(): + check_polar_backends("IGAPolarCylinder", with_mass=False, mpi_dims_mask=mpi_dims_mask) + finally: + cunumpy.mpi.set_mpi_cuda_aware(previous_mpi) + if fake: + set_device_sparse_module(previous_sparse) + + +@pytest.mark.skipif( + not getattr(sys.modules.get("cupy"), "__cunumpy_fake__", False), reason="requires CUNUMPY_FAKE_CUPY=1" +) +@pytest.mark.mpi +@pytest.mark.parametrize("cuda_aware", (True, False)) +@pytest.mark.parametrize("mpi_dims_mask", ((True, False, False), (False, True, False), (False, False, True))) +def test_polar_mpi_fake_cupy(cuda_aware, mpi_dims_mask): + check_polar_mpi(cuda_aware, mpi_dims_mask) + + +@requires_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("mpi_dims_mask", ((True, False, False), (False, True, False), (False, False, True))) +def test_polar_mpi_on_cupy(mpi_dims_mask): + check_polar_mpi(None, mpi_dims_mask, fake=False) + + +def check_polar_collectives(cuda_aware): + """Polar collectives alone, including zero-column chunks, independent of feectools halos.""" + from maybempi import MPI + + from struphy.polar.kron_2dx1d import _BoxTranspose + from struphy.polar.linear_operators import _allreduce_sum + + comm = MPI.COMM_WORLD + rank, size = comm.Get_rank(), comm.Get_size() + previous = cunumpy.mpi.get_mpi_cuda_aware() + try: + cunumpy.mpi.set_mpi_cuda_aware(cuda_aware) + with cunumpy.use_backend("cupy"): + local = cunumpy.asarray(np.arange(6.0).reshape(2, 3) + 100 * rank) + transpose = _BoxTranspose(comm, (2 * rank,), (2 * rank + 1,), (2 * size,), 3) + with cunumpy.profiling.count_transfers() as transfers: + full = transpose.forward(local) + restored = transpose.backward(full) + expected = np.concatenate([np.arange(6.0).reshape(2, 3) + 100 * r for r in range(size)]) + a, b = transpose.my_chunk + np.testing.assert_array_equal(cunumpy.to_numpy(full), expected[:, a:b]) + np.testing.assert_array_equal(cunumpy.to_numpy(restored), cunumpy.to_numpy(local)) + if cuda_aware or size == 1: + assert transfers.total == 0, transfers.report() + else: + assert transfers.to_host == transfers.to_device == 2, transfers.report() + values = cunumpy.asarray(np.arange(5.0) + rank) + _allreduce_sum(comm, values) + np.testing.assert_array_equal(cunumpy.to_numpy(values), size * np.arange(5.0) + size * (size - 1) / 2) + finally: + cunumpy.mpi.set_mpi_cuda_aware(previous) + + +@pytest.mark.skipif( + not getattr(sys.modules.get("cupy"), "__cunumpy_fake__", False), reason="requires CUNUMPY_FAKE_CUPY=1" +) +@pytest.mark.mpi +@pytest.mark.parametrize("cuda_aware", (True, False)) +def test_polar_collectives_mpi_fake(cuda_aware): + check_polar_collectives(cuda_aware) + + +@requires_cupy +@pytest.mark.mpi +def test_polar_collectives_mpi_on_cupy(): + """Real device buffers staged through a regular MPI build; no CUDA-aware MPI needed.""" + check_polar_collectives(False) From 72954b111518cbc2129ad9c759e2607a7d5b4f9f Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Mon, 12 Oct 2026 13:57:55 +0200 Subject: [PATCH 49/50] Update strategy --- CUDA_STRATEGY.md | 4 +++- src/struphy/polar/linear_operators.py | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 5a094dacd..d0c6698f3 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -831,6 +831,8 @@ the [porting order](#porting-order); part of #688). kernels), with dense device-array substitutes for cuSPARSE products and triangular solves. MPI parity tests decompose independently in eta1, eta2 and eta3 and exercise both device-buffer and staged polar collectives; CI runs them on 1, 2, 3 and 4 ranks. Real-GPU variants are included but require - GPU hardware and the feectools CUDA stack. + GPU hardware and the feectools CUDA stack. Full real-GPU MPI parity is opt-in with + `STRUPHY_TEST_CUDA_AWARE_MPI=1` on a known CUDA-aware MPI build; the separate staged-collective test + runs with regular MPI too, without involving feectools halos. - **Separate follow-ups.** `SplineFunction.__call__` field evaluation is not ported to CuPy for any space yet. FFT eta3 solves and sharing factorizations across different projector objects are additional optimizations. diff --git a/src/struphy/polar/linear_operators.py b/src/struphy/polar/linear_operators.py index 4e39f7797..1da50c261 100644 --- a/src/struphy/polar/linear_operators.py +++ b/src/struphy/polar/linear_operators.py @@ -40,7 +40,7 @@ def device_sparse_module(): class DeviceSparseMatrix: - """Device copy of a (small) host matrix of the polar operators, for products with device arrays. + """Device CSR copy of a host polar matrix, for products with device arrays. The polar blocks (:class:`~struphy.polar.extraction_operators.PolarExtractionBlocksC1`) are setup data built with SciPy on the host. On the CuPy backend the polar operators apply them to coefficients that live on the @@ -109,7 +109,7 @@ def _drop_device_blocks(self, name): def _allreduce_sum(comm, buf): - """In-place sum over the ranks of comm (buf may live on the device: CUDA-aware MPI).""" + """In-place sum with synchronized device buffers or host staging.""" if comm is not None and comm.Get_size() > 1: with mpi_buffer(buf, recv=True) as buffer: comm.Allreduce(MPI.IN_PLACE, buffer, op=MPI.SUM) From b80b20fb130c852ba9d10d65fb5721e3db929809 Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Mon, 12 Oct 2026 14:07:13 +0200 Subject: [PATCH 50/50] update docs --- CUDA_STRATEGY.md | 8 +++++--- doc/sections/subsections/numerics-geomFE.rst | 5 +++-- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index d0c6698f3..e9711804f 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -146,7 +146,7 @@ from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage # a Kerne - The CUDA argument objects hold references. If an owner reallocates an array, it must rebuild its arguments at the same place, exactly like for the pyccel arguments. - `Domain` (PR 5): arrays that already have the dtype and layout the kernels expect (`float64`/`int64`, C-contiguous) are referenced, not copied; otherwise one device copy is made when the arguments are built (e.g. `degree`, a tuple). Arguments are recreated after deepcopy or unpickling. Spline mappings (e.g. `IGAPolarCylinder`) cannot be created on CuPy yet (`interp_mapping` passes CuPy arrays to `scipy.sparse`), so CUDA `args_domain` exists only for analytic mappings for now. - `Particles` (PR 6): particle arrays, validity masks and boundary-condition codes are allocated through `xp`. Scalar MPI gathers use small NumPy buffers (the own entry filled first, see `MockComm` above) and copy the result back to the active backend. -- `Derham` (PR 7): depends on feectools running on the CuPy backend (struphy-hub/feectools#85, merged). Data that describes the spline spaces (knots, quadrature and projection grids, `spline_types_pyccel`) is host data on every backend; only coefficients and stencil matrices live on the device. `domain_array`, `index_array(_N/_D)` and `neighbours` are gathered with NumPy MPI buffers and converted with `xp.asarray`. `args_derham` is built from the host knots, degrees and starts on NumPy, and from one device copy of them on CuPy. The pyccel scratch arrays (`bn1`, ..., `bd3`) are not part of the CUDA arguments; they become per-thread local arrays in CUDA (PR 10). Polar splines run on CuPy since #695 (see [Polar splines on CuPy](#polar-splines-on-cupy-695)). Not supported on CuPy yet: local projectors (`NotImplementedError` at construction) and field evaluation (`SplineFunction.__call__` still calls pyccel kernels with device coefficients). +- `Derham` (PR 7): depends on feectools running on the CuPy backend (struphy-hub/feectools#85, merged). Data that describes the spline spaces (knots, quadrature and projection grids, `spline_types_pyccel`) is host data on every backend; only coefficients and stencil matrices live on the device. `domain_array`, `index_array(_N/_D)` and `neighbours` are gathered with NumPy MPI buffers and converted with `xp.asarray`. `args_derham` is built from the host knots, degrees and starts on NumPy, and from one device copy of them on CuPy. The pyccel scratch arrays (`bn1`, ..., `bd3`) are not part of the CUDA arguments; they become per-thread local arrays in CUDA (PR 10). Polar splines run on CuPy since #695 (see [Polar splines on CuPy](#polar-splines-on-cupy-695)). Not supported on CuPy yet: local projectors (`NotImplementedError` at construction) and fixed local-grid evaluation (`SplineFunction.eval_tp_fixed_loc` still calls a pyccel helper with device coefficients). ### PR 8: Argument structs in a shared header (complete, open) @@ -834,5 +834,7 @@ the [porting order](#porting-order); part of #688). GPU hardware and the feectools CUDA stack. Full real-GPU MPI parity is opt-in with `STRUPHY_TEST_CUDA_AWARE_MPI=1` on a known CUDA-aware MPI build; the separate staged-collective test runs with regular MPI too, without involving feectools halos. -- **Separate follow-ups.** `SplineFunction.__call__` field evaluation is not ported to CuPy for any space yet. - FFT eta3 solves and sharing factorizations across different projector objects are additional optimizations. +- **Separate follow-ups.** `SplineFunction.__call__` already dispatches CUDA kernels for marker and dense/sparse + grid evaluation. Remaining evaluation work is the host-only `eval_tp_fixed_loc` helper and distributed + field-evaluation MPI handling/coverage. FFT eta3 solves and sharing factorizations across different projector + objects are additional optimizations. diff --git a/doc/sections/subsections/numerics-geomFE.rst b/doc/sections/subsections/numerics-geomFE.rst index 54f9b2d97..59a90211f 100644 --- a/doc/sections/subsections/numerics-geomFE.rst +++ b/doc/sections/subsections/numerics-geomFE.rst @@ -293,7 +293,8 @@ On the CuPy backend, extraction, derivatives, Kronecker operators and these solv coefficients on the device. Factorization and mass-preconditioner assembly run once on the host; cached device factors apply subsequent solves. Polar MPI collectives support CUDA-aware device buffers or host staging through ``cunumpy.mpi.mpi_buffer``. The feectools halo exchanges require -CUDA-aware MPI on real GPUs. General ``SplineFunction.__call__`` field evaluation is a separate -GPU port and is not supported yet. +CUDA-aware MPI on real GPUs. ``SplineFunction.__call__`` dispatches CUDA kernels for marker +and dense/sparse grid evaluation. The fixed local-grid helper ``eval_tp_fixed_loc`` remains +host-only; distributed field-evaluation MPI handling and coverage are separate follow-ups. .. image:: ../../pics/polar_derham.png \ No newline at end of file