From 1ad4cd6fba1146bbf0ebbd1364bd8218b935723a Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 19:12:45 +0200 Subject: [PATCH 01/23] 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/23] 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/23] 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/23] 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/23] 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/23] 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 1979ea91fa3ddb3a47448f6c17f10d674cdb464f Mon Sep 17 00:00:00 2001 From: Max Date: Wed, 7 Oct 2026 23:35:08 +0200 Subject: [PATCH 07/23] 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 08/23] 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 09/23] 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 10/23] 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 11/23] 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 0e74e15bb8cfc0abc28741280c740ce50b047359 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 9 Oct 2026 15:02:41 +0200 Subject: [PATCH 12/23] 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 13/23] 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 05bc8d95e6d55b82224157915b37ee39728e74fc Mon Sep 17 00:00:00 2001 From: Max Lindqvist Date: Fri, 9 Oct 2026 16:21:37 +0200 Subject: [PATCH 14/23] 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 15/23] 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 16/23] 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 17/23] 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 18/23] 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 19/23] 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 20/23] 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 21/23] 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 22/23] 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 23/23] 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