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..abe1ef445 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_pusher_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 @@ -49,14 +51,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)