diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index de132864f..e81695fc7 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -27,6 +27,7 @@ Progress is tracked in struphy-hub/struphy#650. - [x] **PR 18: Geometry evaluation on the GPU for all analytic mappings** (#681, open) — CUDA versions of the four geometry entry kernels (`kernel_evaluate_pic`, `kernel_evaluate`, `kernel_pullpush_pic`, `kernel_pullpush` in `geometry/kernels/`), built on device versions of the whole metric chain (`f`, `df`, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`/`push`/`tran`) for every analytic mapping (`kind_map` 10–12, 20–22, 30–32). Restores CuPy particle runs (weight initialization evaluates `jacobian_det`), and removes the Cuboid-only checks in `Pusher`, the accumulators and `reflect`. Parity arguments cover every analytic mapping (see [PR 18](#pr-18-geometry-evaluation-for-all-analytic-mappings) and the [PR 18 implementation notes](#pr-18-implementation-notes)). - [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines were host-only at this stage; their current CuPy support is described in [Polar splines on CuPy](#polar-splines-on-cupy-695). See also the [PR 19 implementation notes](#pr-19-implementation-notes). - [x] **`linear_vlasov_ampere` on the GPU** (part of #688) — the first matrix accumulation: CUDA version of `linear_vlasov_ampere` (`EfieldWeightsCoupling`) on cunumpy's `Array6D` views (cunumpy 0.6.1), device versions of `fill_mat`, `fill_mat_vec`, `m_v_fill_b_v1_symm` and `outer`, and the matrix path of `Accumulator` checked on the CuPy backend (device data, ghost regions, transposed blocks, Schur operator) without a GPU through cunumpy's fake CuPy with emulated kernel launches (see the [implementation notes](#linear_vlasov_ampere-implementation-notes)). +- [ ] **Diagnostics on the device** (struphy-hub/struphy#697) — scalar quantities, marker binning, saved markers and spline evaluation for output run on the device; the scalars reach the host in one copy per update (`Scalars.to_host`), every other saved dataset in one counted copy per output step (`DataContainer.save_data`), and the time state is host bookkeeping. CUDA version of `eval_spline_mpi_tensor_product_fixed` (`SplineFunction.eval_tp_fixed_loc`), moved into its own folder. Marker diagnostics kernels (`pic/diagnostics/kernels`) and SPH kernel density plots are follow-ups (see [Diagnostics on the device](#diagnostics-on-the-device-697)). - [ ] **Next (order to be confirmed)**: one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, which needs the [feectools stack](#feectools) merged and released first; the guiding-center kernels (step 3), by hand (decided in #690, see [Hand-written CUDA](#hand-written-cuda-decision-for-690)). - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)). - [x] **PR 13: Move the kernel infrastructure to cunumpy** — `Kernel`, `KernelCatalog`, `CudaKernel` and `Argument` are replaced by their cunumpy counterparts; each owner has a single `args_*` object on both backends, and `pusher_args.cuh` is generated (see [Moving to cunumpy](#moving-to-cunumpy-pr-13)). Kernels and device helpers only change their includes. @@ -829,6 +830,41 @@ the [porting order](#porting-order); part of #688). ghost-region flags from #746 are preserved. feectools' own halo exchanges still require CUDA-aware MPI on real GPUs; staging of the polar collectives does not replace that separate dependency. - **Tests.** `polar/tests/test_polar_cupy.py` compares NumPy and CuPy for `IGAPolarCylinder` and `IGAPolarTorus`: + extraction and DOF extraction operators and their transposes, `grad`/`curl`/`div` and transposes, `PolarVector` + arithmetic, the polar projectors, polar mass matrices (GPU only) and polar `SplineFunction` coefficients + (`extract_coeffs`, restart). Without a GPU it runs on the fake CuPy with two substitutes: a dense stand-in for + `cupyx.scipy.sparse` (`set_device_sparse_module`) and the host versions of the feectools kernels (the fake cannot + launch CUDA kernels). The mass matrices need struphy's CUDA geometry kernels and are compared on the GPU only. + Field evaluation (`SplineFunction.__call__`) is not covered: it is not ported to CuPy for any space yet. + +## Diagnostics on the device (#697) + +What a CuPy run copies to the host per output step, and where: + +- **Scalars** (`models/scalars.py`). The `value` of every scalar in a `Scalars` container is a view into one device + array of the container; `Scalars.update()` ends with `to_host()`, one `xp.to_numpy` of that array (8 bytes per + scalar). `FunctionScalar*` store the result of their function without `float()` (a device sync per scalar before), + so a device reduction stays on the device. The output (`Simulation._initialize_hdf5_datasets`) and + `print_scalar_quantities` read `Scalars.host_value(name)`. `BilinearEnergyFEEC` and `VolumeFormEnergyFEEC` call + feectools' inner products, which copy their result to the host until struphy-hub/feectools#98 (#716) is in. +- **Output** (`io/output_handling.py`). `DataContainer._as_numpy_array` is `xp.to_numpy`: the one place where device + data is copied for saving, one counted copy per saved device dataset (FEEC coefficients, saved markers, binned + distribution functions). Before, it used `.get()`, which `count_transfers` does not see. +- **Time state** (`simulation/sim.py`). `time_state` is NumPy on both backends: the time loop reads it with `float()` + every step, which was a device sync per read on CuPy. +- **Binning** (`Particles.binning`). Failed on CuPy (`xp.count_nonzero` of a list); now gathers the valid markers once + and runs `xp.histogramdd` on the device. Saved markers (`update_markers_to_be_saved`) gather once, without counting + the mask first. +- **Spline evaluation.** `SplineFunction.__call__` compared the process bounds as device scalars (six syncs per + meshgrid evaluation); the bounds are cached on the host once, and the flagging writes are masked copies. + `eval_tp_fixed_loc` (thermal energies of the variational models) called the pyccel kernel with device arrays; its + kernel `eval_spline_mpi_tensor_product_fixed` is now in `bsplines/kernels/` with a CUDA version and parity cases. + `Derham.prepare_eval_tp_fixed` (setup) computes on the host and copies the spans and basis values to the device once. +- **Not counted, still there**: boolean-mask gathers (`markers[valid_mks]`, also in `KineticEnergyPIC`) wait for the + device to know the result size; one per gather. +- **Follow-ups**: the marker diagnostics kernels in `pic/diagnostics/kernels/` (5D energies, magnetic moments, + canonical momenta, `thermal_energy`; used by the guiding-center and hybrid models) have no CUDA version; SPH kernel + density plots need the SPH kernels (step 6). extraction/DOF operators, derivatives and transposes, vector arithmetic, all three projector solvers, transposed and boundary-condition solves, mass matrices/preconditioners, callable and sampled weights, and spline coefficient extraction/restart. Repeated Kronecker and preconditioner applications are checked diff --git a/src/struphy/bsplines/evaluation_kernels_3d.py b/src/struphy/bsplines/evaluation_kernels_3d.py index d67f69781..5409b0b4d 100644 --- a/src/struphy/bsplines/evaluation_kernels_3d.py +++ b/src/struphy/bsplines/evaluation_kernels_3d.py @@ -1334,82 +1334,6 @@ def eval_spline_mpi_tensor_product_fast( ) -@stack_array("b1", "b2", "b3") -def eval_spline_mpi_tensor_product_fixed( - span1s: "int[:]", - span2s: "int[:]", - span3s: "int[:]", - b1s: "float[:,:]", - b2s: "float[:,:]", - b3s: "float[:,:]", - _data: "float[:,:,:]", - kind: "int[:]", - pn: "int[:]", - starts: "int[:]", - values: "float[:,:,:]", -): - """ - Tensor-product evaluation of a tensor-product spline, distributed, - and optimized for a fixed grid (spans and spline values have been pre-evaluated). - - Parameters - ---------- - span1s, span2s, span3s : array[int] - Knot span indices. - - b1s, b2s, b3s : array[float] - Values of p+1 non-zero basis functions at evaluation points. - - _data : array[float] - The spline coefficients c_ijk. - - kind : array[int] - Kind of 1d basis in each direction: 0 = N-spline, 1 = D-spline. - - pn : array[int] - Spline degrees of V0 in each direction. - - starts : array[float] - Start indices of splines on current process. - - values : array[float] - Return array for spline values S_ijk where ijk are the flattened indices corresponding to spans. - """ - - # allocate spline values - b1 = empty(pn[0] - kind[0] + 1, dtype=float) - b2 = empty(pn[1] - kind[1] + 1, dtype=float) - b3 = empty(pn[2] - kind[2] + 1, dtype=float) - - ni = span1s.size - nj = span2s.size - nk = span3s.size - - for i in range(ni): - span1 = span1s[i] - b1[:] = b1s[i, :] - for j in range(nj): - span2 = span2s[j] - b2[:] = b2s[j, :] - for k in range(nk): - span3 = span3s[k] - b3[:] = b3s[k, :] - - values[i, j, k] = eval_spline_mpi_kernel( - pn[0] - kind[0], - pn[1] - kind[1], - pn[2] - kind[2], - b1, - b2, - b3, - span1, - span2, - span3, - _data, - starts, - ) - - ########################################### ### Distributed Derham field evaluation ### ########################################### diff --git a/src/struphy/bsplines/kernels/eval_spline_mpi_tensor_product_fixed/__init__.py b/src/struphy/bsplines/kernels/eval_spline_mpi_tensor_product_fixed/__init__.py new file mode 100644 index 000000000..a627f430d --- /dev/null +++ b/src/struphy/bsplines/kernels/eval_spline_mpi_tensor_product_fixed/__init__.py @@ -0,0 +1,7 @@ +"""The kernel of this folder: pyccel in ``_kernels.py``, CUDA in ``_cuda.cu`` once ported.""" + +from cunumpy.kernels import Kernel + +from struphy.utils.cuda_arguments import CUDA_STRUCTS + +eval_spline_mpi_tensor_product_fixed = Kernel.from_folder(__name__, structs=CUDA_STRUCTS) diff --git a/src/struphy/bsplines/kernels/eval_spline_mpi_tensor_product_fixed/eval_spline_mpi_tensor_product_fixed_cuda.cu b/src/struphy/bsplines/kernels/eval_spline_mpi_tensor_product_fixed/eval_spline_mpi_tensor_product_fixed_cuda.cu new file mode 100644 index 000000000..fa7716720 --- /dev/null +++ b/src/struphy/bsplines/kernels/eval_spline_mpi_tensor_product_fixed/eval_spline_mpi_tensor_product_fixed_cuda.cu @@ -0,0 +1,53 @@ +#include "struphy/bsplines/evaluation_kernels_3d.cuh" + +// Same arguments, in the same order, as the pyccel kernel in eval_spline_mpi_tensor_product_fixed_kernels.py. Array +// views carry dimensions and strides, so no CUDA-only lengths are needed. + +/** + * Tensor-product evaluation of a distributed spline on a fixed grid, as in + * eval_spline_mpi_tensor_product_fixed_kernels.eval_spline_mpi_tensor_product_fixed. + * + * One thread per entry (i, j, k) of values. The knot spans and the non-zero basis values at the grid points are + * pre-evaluated (SplineFunction.eval_tp_fixed_loc, Derham.prepare_eval_tp_fixed), so the thread only sums the + * contributions of the local coefficients. + * + * @param span1s Knot span indices of the grid points along the first axis (ni entries). + * @param span2s Knot span indices along the second axis (nj entries). + * @param span3s Knot span indices along the third axis (nk entries). + * @param b1s Values of the non-zero basis functions along the first axis, shape (ni, pn[0] - kind[0] + 1). + * @param b2s Values of the non-zero basis functions along the second axis, shape (nj, pn[1] - kind[1] + 1). + * @param b3s Values of the non-zero basis functions along the third axis, shape (nk, pn[2] - kind[2] + 1). + * @param _data Spline coefficients of the current process, any strides. + * @param kind Kind of 1d basis in each direction (three entries): 0 = N-spline, 1 = D-spline. + * @param pn Spline degrees of V0 in each direction (three entries, 1 to 8). + * @param starts Start indices of the splines on the current process (three entries). + * @param values Output spline values, shape (ni, nj, nk). + * + * Pyccel copies the rows b1s[i, :], b2s[j, :], b3s[k, :] into b1, b2, b3 before calling eval_spline_mpi_kernel; here + * b1, b2, b3 are thread-local arrays of MAX_SPLINE_DEGREE + 1 entries, since the rows of a view may be strided. + */ +extern "C" __global__ void eval_spline_mpi_tensor_product_fixed(Array1D span1s, Array1D span2s, + Array1D span3s, Array2D b1s, + Array2D b2s, Array2D b3s, + Array3D _data, const long long* kind, + const long long* pn, const long long* starts, + Array3D values) { + // CUDA-only: flat index of this thread, split into the pyccel loop variables i, j, k + long long ijk = (long long)blockDim.x * blockIdx.x + threadIdx.x; + long long ni = span1s.shape[0]; + long long nj = span2s.shape[0]; + long long nk = span3s.shape[0]; + if (ijk >= ni * nj * nk) return; + long long k = ijk % nk; + long long j = (ijk / nk) % nj; + long long i = ijk / (nj * nk); + + double b1[struphy_cuda::MAX_SPLINE_DEGREE + 1], b2[struphy_cuda::MAX_SPLINE_DEGREE + 1], + b3[struphy_cuda::MAX_SPLINE_DEGREE + 1]; + for (int il = 0; il <= pn[0] - kind[0]; ++il) b1[il] = b1s(i, il); + for (int il = 0; il <= pn[1] - kind[1]; ++il) b2[il] = b2s(j, il); + for (int il = 0; il <= pn[2] - kind[2]; ++il) b3[il] = b3s(k, il); + + values(i, j, k) = struphy_cuda::eval_spline_mpi_kernel(pn[0] - kind[0], pn[1] - kind[1], pn[2] - kind[2], b1, b2, + b3, span1s(i), span2s(j), span3s(k), _data, starts); +} diff --git a/src/struphy/bsplines/kernels/eval_spline_mpi_tensor_product_fixed/eval_spline_mpi_tensor_product_fixed_kernels.py b/src/struphy/bsplines/kernels/eval_spline_mpi_tensor_product_fixed/eval_spline_mpi_tensor_product_fixed_kernels.py new file mode 100644 index 000000000..29e7157c5 --- /dev/null +++ b/src/struphy/bsplines/kernels/eval_spline_mpi_tensor_product_fixed/eval_spline_mpi_tensor_product_fixed_kernels.py @@ -0,0 +1,82 @@ +"""Tensor-product evaluation of a distributed spline on a fixed grid (spans and basis values pre-evaluated).""" + +from numpy import empty +from pyccel.decorators import stack_array + +from struphy.bsplines.evaluation_kernels_3d import eval_spline_mpi_kernel + + +@stack_array("b1", "b2", "b3") +def eval_spline_mpi_tensor_product_fixed( + span1s: "int[:]", + span2s: "int[:]", + span3s: "int[:]", + b1s: "float[:,:]", + b2s: "float[:,:]", + b3s: "float[:,:]", + _data: "float[:,:,:]", + kind: "int[:]", + pn: "int[:]", + starts: "int[:]", + values: "float[:,:,:]", +): + """ + Tensor-product evaluation of a tensor-product spline, distributed, + and optimized for a fixed grid (spans and spline values have been pre-evaluated). + + Parameters + ---------- + span1s, span2s, span3s : array[int] + Knot span indices. + + b1s, b2s, b3s : array[float] + Values of p+1 non-zero basis functions at evaluation points. + + _data : array[float] + The spline coefficients c_ijk. + + kind : array[int] + Kind of 1d basis in each direction: 0 = N-spline, 1 = D-spline. + + pn : array[int] + Spline degrees of V0 in each direction. + + starts : array[float] + Start indices of splines on current process. + + values : array[float] + Return array for spline values S_ijk where ijk are the flattened indices corresponding to spans. + """ + + # allocate spline values + b1 = empty(pn[0] - kind[0] + 1, dtype=float) + b2 = empty(pn[1] - kind[1] + 1, dtype=float) + b3 = empty(pn[2] - kind[2] + 1, dtype=float) + + ni = span1s.size + nj = span2s.size + nk = span3s.size + + for i in range(ni): + span1 = span1s[i] + b1[:] = b1s[i, :] + for j in range(nj): + span2 = span2s[j] + b2[:] = b2s[j, :] + for k in range(nk): + span3 = span3s[k] + b3[:] = b3s[k, :] + + values[i, j, k] = eval_spline_mpi_kernel( + pn[0] - kind[0], + pn[1] - kind[1], + pn[2] - kind[2], + b1, + b2, + b3, + span1, + span2, + span3, + _data, + starts, + ) diff --git a/src/struphy/bsplines/tests/test_eval_spline_mpi.py b/src/struphy/bsplines/tests/test_eval_spline_mpi.py index 99f4c850c..e0007dd24 100644 --- a/src/struphy/bsplines/tests/test_eval_spline_mpi.py +++ b/src/struphy/bsplines/tests/test_eval_spline_mpi.py @@ -699,7 +699,8 @@ def test_eval_tensor_product_grid(num_elements, degree, bcs, n_markers=10): import time - from struphy.bsplines.evaluation_kernels_3d import eval_spline_mpi_tensor_product_fixed, evaluate_tensor_product + from struphy.bsplines.evaluation_kernels_3d import evaluate_tensor_product + from struphy.bsplines.kernels.eval_spline_mpi_tensor_product_fixed import eval_spline_mpi_tensor_product_fixed from struphy.feec.basis_projection_ops import prepare_projection_of_basis from struphy.feec.psydac_derham import Derham from struphy.feec.utilities import create_equal_random_arrays as cera diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index cef3335b9..d11f4aa5d 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -29,10 +29,10 @@ from feectools.linalg.stencil import StencilVector, StencilVectorSpace from maybempi import MPI, SerialComm -from struphy.bsplines.evaluation_kernels_3d import eval_spline_mpi_tensor_product_fixed from struphy.bsplines.kernels.eval_spline_mpi_markers import eval_spline_mpi_markers from struphy.bsplines.kernels.eval_spline_mpi_matrix import eval_spline_mpi_matrix from struphy.bsplines.kernels.eval_spline_mpi_sparse_meshgrid import eval_spline_mpi_sparse_meshgrid +from struphy.bsplines.kernels.eval_spline_mpi_tensor_product_fixed import eval_spline_mpi_tensor_product_fixed from struphy.feec.linear_operators import BoundaryOperator from struphy.feec.local_projectors_kernels import get_local_problem_size, select_quasi_points from struphy.feec.projectors import CommutingProjector, CommutingProjectorLocal @@ -2225,14 +2225,16 @@ def _get_span_and_basis_for_eval_mpi(self, etas, Nspace, end): from struphy.bsplines import bsplines_kernels # Extract knot vectors, degree and kind of basis - Tn = Nspace.knots + Tn = xp.to_numpy(Nspace.knots) pn = Nspace.degree - spans = xp.zeros(etas.size, dtype=int) - bns = xp.zeros((etas.size, pn + 1), dtype=float) - bds = xp.zeros((etas.size, pn), dtype=float) - bn = xp.zeros(pn + 1, dtype=float) - bd = xp.zeros(pn, dtype=float) + # setup: computed point by point with the pyccel helpers on the host, then copied to the active backend once + etas = xp.to_numpy(etas) + spans = np.zeros(etas.size, dtype=int) + bns = np.zeros((etas.size, pn + 1), dtype=float) + bds = np.zeros((etas.size, pn), dtype=float) + bn = np.zeros(pn + 1, dtype=float) + bd = np.zeros(pn, dtype=float) for n in range(etas.size): # avoid 1. --> 0. for clamped interpolation @@ -2246,7 +2248,7 @@ def _get_span_and_basis_for_eval_mpi(self, etas, Nspace, end): bns[n] = bn bds[n] = bd - return spans, bns, bds + return xp.to_cunumpy(spans), xp.to_cunumpy(bns), xp.to_cunumpy(bds) class SplineFunction: @@ -2835,15 +2837,9 @@ def eval_tp_fixed_loc(self, spans, bases, out=None): else: assert out.shape == tuple([span.size for span in spans]) - eval_spline_mpi_tensor_product_fixed( - *spans, - *bases, - vec._data, - self.derham.spline_attributes[self.space_key].spline_types_pyccel[0], - xp.array(self.derham.degree), - xp.array(self.starts), - out, - ) + args = self._args_spline[0] + kind, pn, starts = args.kind, args.pn, args.starts + eval_spline_mpi_tensor_product_fixed(*spans, *bases, vec._data, kind, pn, starts, out, n_threads=out.size) else: out_is_none = False @@ -2866,18 +2862,10 @@ def eval_tp_fixed_loc(self, spans, bases, out=None): [span.size for span in spans], ) + args = self._args_spline[i] + kind, pn, starts = args.kind, args.pn, args.starts eval_spline_mpi_tensor_product_fixed( - *spans, - *bases[i], - vec[i]._data, - self.derham.spline_attributes[self.space_key].spline_types_pyccel[i], - xp.array( - self.derham.degree, - ), - xp.array( - self.starts[i], - ), - out[i], + *spans, *bases[i], vec[i]._data, kind, pn, starts, out[i], n_threads=out[i].size ) return out @@ -3119,28 +3107,32 @@ def _flag_pts_not_on_proc(self, *etas): ) on_proc = xp.all(is_on_proc_domain, axis=1) - markers[~on_proc, :] = -1.0 + xp.copyto(markers, -1.0, where=~on_proc[:, None]) # 3D meshgrid evaluation else: assert len(etas) == 3 E1, E2, E3 = etas - # check if eval points are "interior points" in domain_array; if so, add small offset + # check if eval points are "interior points" in domain_array; if so, add small offset. + # The bounds are compared as host floats: with the device array each `if` would wait for the device. + if not hasattr(self, "_domain_array_host"): + self._domain_array_host = xp.to_numpy(dom_arr) + dom_arr = self._domain_array_host if dom_arr[rank, 0] != 0.0: - E1[E1 == dom_arr[rank, 0]] += 1e-8 + E1 += 1e-8 * (E1 == dom_arr[rank, 0]) if dom_arr[rank, 1] != 1.0: - E1[E1 == dom_arr[rank, 1]] += 1e-8 + E1 += 1e-8 * (E1 == dom_arr[rank, 1]) if dom_arr[rank, 3] != 0.0: - E2[E2 == dom_arr[rank, 3]] += 1e-8 + E2 += 1e-8 * (E2 == dom_arr[rank, 3]) if dom_arr[rank, 4] != 1.0: - E2[E2 == dom_arr[rank, 4]] += 1e-8 + E2 += 1e-8 * (E2 == dom_arr[rank, 4]) if dom_arr[rank, 6] != 0.0: - E3[E3 == dom_arr[rank, 6]] += 1e-8 + E3 += 1e-8 * (E3 == dom_arr[rank, 6]) if dom_arr[rank, 7] != 1.0: - E3[E3 == dom_arr[rank, 7]] += 1e-8 + E3 += 1e-8 * (E3 == dom_arr[rank, 7]) # True for eval points on current process E1_on_proc = xp.logical_and( @@ -3156,10 +3148,10 @@ def _flag_pts_not_on_proc(self, *etas): E3 <= dom_arr[rank, 7], ) - # flag eval points not on current process - E1[~E1_on_proc] = -1.0 - E2[~E2_on_proc] = -1.0 - E3[~E3_on_proc] = -1.0 + # flag eval points not on current process (masked writes without a device sync) + xp.copyto(E1, -1.0, where=~E1_on_proc) + xp.copyto(E2, -1.0, where=~E2_on_proc) + xp.copyto(E3, -1.0, where=~E3_on_proc) def _add_noise( self, diff --git a/src/struphy/io/output_handling.py b/src/struphy/io/output_handling.py index d0b535cbf..9d7620563 100644 --- a/src/struphy/io/output_handling.py +++ b/src/struphy/io/output_handling.py @@ -1,6 +1,7 @@ import logging import os +import cunumpy as xp import h5py import numpy as np @@ -82,15 +83,13 @@ def dset_dict(self): @staticmethod def _as_numpy_array(val): - """Return a NumPy view/copy suitable for h5py writes.""" - if isinstance(val, np.ndarray): - return val + """Return a NumPy array suitable for h5py writes. - get = getattr(val, "get", None) - if callable(get) and "cupy" in val.__class__.__module__: - return get() - - return np.asarray(val) + Host data is returned as it is. Device data (CuPy backend) is copied to the host here, through cunumpy + (``xp.to_numpy``), so the copy is explicit and seen by ``cunumpy.profiling.count_transfers``. This is the + only place where the output copies device data: one copy per saved dataset and output step. + """ + return xp.to_numpy(val) def add_data(self, data_dict): """ diff --git a/src/struphy/models/base.py b/src/struphy/models/base.py index 94f3444b2..4b0b6219c 100644 --- a/src/struphy/models/base.py +++ b/src/struphy/models/base.py @@ -5,6 +5,7 @@ from textwrap import indent import cunumpy as xp +import numpy as np from line_profiler import profile from maybempi import MPI, SerialMPI from scope_profiler import ProfileManager @@ -454,9 +455,10 @@ def print_scalar_quantities(self): Check if scalars are not "nan" and print to screen. """ sq_str = "" - for key, scalar in self.scalars.dct.items(): - val = scalar.value[0] - assert not xp.isnan(val), f"Scalar {key} is {val}." + for key in self.scalars.dct: + # the host copy made by the last update: no device synchronization per scalar + val = self.scalars.host_value(key)[0] + assert not np.isnan(val), f"Scalar {key} is {val}." sq_str += f"{key}:".ljust(25) + "{:4.2e}\n".format(val).rjust(26) print(sq_str) @@ -561,9 +563,11 @@ def update_markers_to_be_saved(self): obj.markers[:, -1] >= 0.0, obj.markers[:, -1] < var.n_to_save, ) - n_markers_on_proc = xp.count_nonzero(markers_on_proc) + # the number of rows is known on the host after the gather; counting the mask first + # (xp.count_nonzero) would wait for the device a second time + rows = obj.markers[markers_on_proc] var.saved_markers[:] = -1.0 - var.saved_markers[:n_markers_on_proc] = obj.markers[markers_on_proc] + var.saved_markers[: rows.shape[0]] = rows @profile def update_distr_functions(self): @@ -608,7 +612,7 @@ def update_distr_functions(self): h2 = 1 / obj.boxes_per_dim[1] h3 = 1 / obj.boxes_per_dim[2] - ndim = xp.count_nonzero(xp.array([d > 1 for d in obj.boxes_per_dim])) + ndim = sum(int(d) > 1 for d in obj.boxes_per_dim) if ndim == 0: kernel_type = "gaussian_3d" else: diff --git a/src/struphy/models/scalars.py b/src/struphy/models/scalars.py index 065709a12..00faeb7b7 100644 --- a/src/struphy/models/scalars.py +++ b/src/struphy/models/scalars.py @@ -2,6 +2,7 @@ from typing import Callable, Union import cunumpy as xp +import numpy as np from maybempi import MPI from struphy.feec.mass import WeightedMassOperator @@ -15,6 +16,18 @@ _DUMMY_VARIABLE = object() +def _scalar_value(value): + """`value` (of size 1) as something that can be written into a scalar's buffer without leaving the device. + + A device result (e.g. a reduction on the CuPy backend) stays a 0-d device array; ``float()`` would wait for the + device and copy it to the host at every update. Host values (Python or NumPy scalars, such as the result of an + inner product that already went through MPI on the host) are plain floats. + """ + if isinstance(value, xp.ndarray): + return value.reshape(()) + return float(np.asarray(value).reshape(())) + + class Scalar(metaclass=ABCMeta): """Abstract base class for scalar quantities in MPI parallel simulations. @@ -140,7 +153,12 @@ def _mpi_sum(self): class Scalars: """Container for multiple Scalar objects. - Calling .update() on this container will update all contained scalars.""" + Calling .update() on this container will update all contained scalars. + + The scalars are computed where the data lives (on the device on the CuPy backend). The ``value`` of each scalar + in the container is a view into one array of the container, and :meth:`update` ends with :meth:`to_host`, which + copies that array to the host in one transfer. The output (:meth:`host_value`, printing) reads the host copy. + """ def __init__(self, **scalars: dict[str, Scalar]): for name, scalar in scalars.items(): @@ -150,6 +168,14 @@ def __init__(self, **scalars: dict[str, Scalar]): else: self._dct = {} + # one buffer for all values, so that they reach the host in a single copy + self._device_values = xp.zeros(len(self._dct), dtype=float) + self._host_values = np.zeros(len(self._dct), dtype=float) + self._index = {} + for i, (name, scalar) in enumerate(self._dct.items()): + scalar.value = self._device_values[i : i + 1] + self._index[name] = i + @property def dct(self) -> dict[str, Scalar]: return self._dct @@ -161,6 +187,19 @@ def update(self): # SumOfScalars(SumOfScalars(a, b), c), and an inner sum left up to date would keep its first value for scalar in self.dct.values(): _mark_outdated(scalar) + self.to_host() + + def to_host(self): + """Copy the values of all scalars to the host, in one transfer (the only one of the scalar diagnostics).""" + self._host_values[:] = xp.to_numpy(self._device_values) + + def host_value(self, name: str) -> np.ndarray: + """The value of scalar `name` after the last :meth:`update`, as a NumPy array of size 1. + + A view into the host buffer of the container: it follows later updates (the output registers it once). + """ + i = self._index[name] + return self._host_values[i : i + 1] def _mark_outdated(scalar: Scalar): @@ -275,7 +314,7 @@ def __init__( Scalar.__init__(self, _DUMMY_VARIABLE) def _local_update(self): - self.local_value[0] = float(self.function()) + self.local_value[0] = _scalar_value(self.function()) def _mpi_sum(self): """Communication has been handled by psydac, so no additional MPI operations are needed.""" @@ -334,7 +373,7 @@ def __init__( super().__init__(pic_variable) def _local_update(self): - self.local_value[0] = float(self.function()) + self.local_value[0] = _scalar_value(self.function()) class KineticEnergySPH(SPHScalar): @@ -372,4 +411,4 @@ def __init__( super().__init__(sph_variable) def _local_update(self): - self.local_value[0] = float(self.function()) + self.local_value[0] = _scalar_value(self.function()) diff --git a/src/struphy/models/tests/test_diagnostics_on_device.py b/src/struphy/models/tests/test_diagnostics_on_device.py new file mode 100644 index 000000000..b29ce81b0 --- /dev/null +++ b/src/struphy/models/tests/test_diagnostics_on_device.py @@ -0,0 +1,345 @@ +"""Diagnostics and output on the device (struphy-hub/struphy#697). + +During a CuPy run the per-step diagnostics (scalar quantities, binned distribution functions, saved markers) are +computed on the device, and only the reduced results reach the host: the scalars in one copy per update +(``Scalars.to_host``), every other saved dataset in one copy per output step (``DataContainer.save_data``). + +Each ``check_*`` function runs on the active backend and compares with the NumPy backend. Without a GPU they run in a +subprocess on cunumpy's fake CuPy (``CUNUMPY_FAKE_CUPY=1``), which rejects host/device mixing; there, implicit +device-to-host conversions (``float()``, ``int()``, ``bool()`` and ``__index__`` of a device array, each of which +waits for the device) raise as well, and the copies made through cunumpy are counted with +``cunumpy.profiling.count_transfers``. With a GPU the same checks run on the real CuPy (without the implicit-sync +guard, which needs the fake). The checks use mock particles and need no CUDA kernel. +""" + +import contextlib +import os +import subprocess +import sys +import tempfile +from types import SimpleNamespace + +import cunumpy as xp +import numpy as np +import pytest +from cunumpy.kernel_testing import requires_cupy + +from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env + +N_ROWS = 12 # marker rows: 8 valid markers, 2 ghosts, 2 holes +N_COLS = 13 # Particles6D columns: eta (3), v (3), w, s0, w0, 2 auxiliary, box, ID + + +def _markers(): + """Particles6D-like markers on the host: 8 valid markers, 2 ghosts (ID -2) and 2 holes (all -1).""" + rng = np.random.default_rng(697) + markers = rng.random((N_ROWS, N_COLS)) + markers[:, 3:6] = rng.normal(size=(N_ROWS, 3)) + markers[:, -1] = np.arange(N_ROWS) + markers[8:10, -1] = -2.0 # ghosts + markers[10:] = -1.0 # holes + return markers + + +class _MockParticles: + """The parts of :class:`~struphy.pic.base.Particles` used by the diagnostics, on the active backend.""" + + vdim = 3 + Np = 8 + n_lost_markers = 3 + clone_config = None + index = {"pos": slice(0, 3), "vel": slice(3, 6), "coords": slice(0, 6), "weights": 6, "s0": 7, "w0": 8, "ids": -1} + + def __init__(self): + self._markers = xp.asarray(_markers()) + self._valid_mks = ~xp.logical_or(self._markers[:, 0] == -1.0, self._markers[:, -1] == -2.0) + + @property + def markers(self): + return self._markers + + @property + def valid_mks(self): + return self._valid_mks + + @property + def markers_wo_holes_and_ghost(self): + return self.markers[self.valid_mks] + + @property + def velocities(self): + return self.markers[self.valid_mks, self.index["vel"]] + + @property + def weights(self): + return self.markers[self.valid_mks, self.index["weights"]] + + +def _fake_cupy_active(): + return bool(getattr(sys.modules.get("cupy"), "__cunumpy_fake__", False)) + + +@contextlib.contextmanager +def no_implicit_syncs(): + """On the fake CuPy, raise if a device array is converted to a host scalar implicitly (a device sync).""" + if not _fake_cupy_active(): + yield + return + import cupy + + def forbidden(name): + def raise_(self, *args, **kwargs): + raise AssertionError(f"implicit device-to-host conversion: ndarray.{name}") + + return raise_ + + names = ("__float__", "__int__", "__bool__", "__index__") + saved = {name: getattr(cupy.ndarray, name) for name in names} + for name in names: + setattr(cupy.ndarray, name, forbidden(name)) + try: + yield + finally: + for name, method in saved.items(): + setattr(cupy.ndarray, name, method) + + +@contextlib.contextmanager +def _serial_propagator_base(): + """``Propagator.derham`` without a communicator, as the PIC scalars need it for their MPI sum.""" + from struphy.propagators.base import Propagator + + saved = Propagator.__dict__["derham"] + Propagator.derham = SimpleNamespace(comm=None) + try: + yield + finally: + Propagator.derham = saved + + +def _scalars(particles): + from struphy.models.scalars import FunctionScalarPIC, KineticEnergyPIC, LostMarkersPIC, Scalars + from struphy.models.variables import PICVariable + + class _Variable(PICVariable): + def __init__(self): + self._particles = particles + + var = _Variable() + kinetic = KineticEnergyPIC(var, normalization=2.0) + # a device reduction, as LinearVlasovAmpereOneSpecies._compute_en_w returns + weights = FunctionScalarPIC(lambda: xp.sum(particles.weights**2), var) + lost = LostMarkersPIC(var) + return Scalars(en_kin=kinetic, en_w=weights, lost=lost, en_tot=kinetic + weights) + + +def _scalar_values(backend): + with xp.use_backend(backend), _serial_propagator_base(): + particles = _MockParticles() + scalars = _scalars(particles) + with xp.profiling.count_transfers() as counter, no_implicit_syncs(): + scalars.update() + values = {key: scalars.host_value(key).copy() for key in scalars.dct} + + # the host values follow the next update, and are views of one host buffer + particles.markers[:, 6] *= 2.0 + scalars.update() + doubled = scalars.host_value("en_kin")[0] + return values, counter, doubled + + +def check_scalars(): + """The scalars stay on the device; one copy of all values per update, no implicit syncs.""" + values, counter, doubled = _scalar_values(xp.get_backend()) + reference, reference_counter, reference_doubled = _scalar_values("numpy") + + assert reference_counter.events == [] + assert values.keys() == reference.keys() + for key in values: + np.testing.assert_allclose(values[key], reference[key], rtol=1e-14, err_msg=key) + np.testing.assert_allclose(doubled, reference_doubled, rtol=1e-14) + np.testing.assert_allclose(doubled, 2.0 * values["en_kin"][0], rtol=1e-14) + np.testing.assert_allclose(values["en_tot"], values["en_kin"] + values["en_w"], rtol=1e-14) + assert values["lost"][0] == 3.0 + + if xp.get_backend() == "cupy": + kinds = [event.kind for event in counter.events] + assert kinds == ["to_host"], counter.report() + assert counter.events[0].nbytes == 8 * len(values), counter.report() + + +def _binned(backend, output_quantity): + from struphy.particles.parameters import BinningPlot + from struphy.pic.base import Particles + + with xp.use_backend(backend): + particles = _MockParticles() + bin_plot = BinningPlot(slice="e1_v1", n_bins=(4, 3), ranges=((0.0, 1.0), (-2.0, 2.0))) + components = [True, False, False, True, False, False] + with xp.profiling.count_transfers() as counter, no_implicit_syncs(): + f_slice, df_slice = Particles.binning( + particles, components, bin_plot.bin_edges, output_quantity=output_quantity, divide_by_jac=False + ) + bin_plot.f[:] = f_slice + bin_plot.df[:] = df_slice + assert type(bin_plot.f) is type(particles.markers) + return xp.to_numpy(bin_plot.f), xp.to_numpy(bin_plot.df), counter + + +def check_binning(): + """Marker binning (``xp.histogramdd``) runs on the device of the markers, without transfers.""" + for output_quantity in ("density", "current_1", "energy_tensor_12", "heat_flux_2"): + f_slice, df_slice, counter = _binned(xp.get_backend(), output_quantity) + f_reference, df_reference, _ = _binned("numpy", output_quantity) + assert counter.events == [], counter.report() + assert np.any(f_reference != 0.0) and np.any(df_reference != 0.0) + np.testing.assert_allclose(f_slice, f_reference, rtol=1e-13, atol=1e-15, err_msg=output_quantity) + np.testing.assert_allclose(df_slice, df_reference, rtol=1e-13, atol=1e-15, err_msg=output_quantity) + + +def _saved_markers(backend): + from struphy.models.base import StruphyModel + from struphy.models.species import ParticleSpecies + from struphy.models.variables import PICVariable + from struphy.pic.base import Particles + + # update_markers_to_be_saved checks for a Particles instance + mock_class = type("_MockParticlesInstance", (_MockParticles, Particles), {}) + mock_class.__abstractmethods__ = frozenset() + + with xp.use_backend(backend): + particles = mock_class() + + class _Variable(PICVariable): + def __init__(self): + self._particles = particles + self._n_to_save = 5 + self._saved_markers = xp.zeros((5, N_COLS)) + + class _Species(ParticleSpecies): + def __init__(self): + self._variables = {"var": _Variable()} + + species = _Species() + model = SimpleNamespace(particle_species={"ions": species}) + with xp.profiling.count_transfers() as counter, no_implicit_syncs(): + StruphyModel.update_markers_to_be_saved(model) + return xp.to_numpy(species.variables["var"].saved_markers), counter + + +def check_saved_markers(): + """The markers to be saved are gathered on the device, without transfers.""" + saved, counter = _saved_markers(xp.get_backend()) + reference, _ = _saved_markers("numpy") + assert counter.events == [], counter.report() + # IDs 0-4 are valid markers here; the ghosts and holes have no ID below n_to_save + np.testing.assert_array_equal(reference[:, -1], [0.0, 1.0, 2.0, 3.0, 4.0]) + np.testing.assert_array_equal(saved, reference) + + +def check_output(): + """``DataContainer.save_data`` copies each device dataset once, through cunumpy, and host datasets not at all.""" + import h5py + + from struphy.io.output_handling import DataContainer + + with tempfile.TemporaryDirectory() as path: + os.makedirs(os.path.join(path, "data")) + data = DataContainer(path) + host = {"time/value": np.full(1, 0.5), "scalar/en": np.full(1, 2.0)} + device = {"feec/e": xp.arange(24.0).reshape(2, 3, 4), "kinetic/f": xp.ones((4, 3))} + data.add_data(host | device) + with xp.profiling.count_transfers() as counter, no_implicit_syncs(): + data.save_data() + with h5py.File(data.file_path) as file: + np.testing.assert_array_equal(file["feec/e"][-1], np.arange(24.0).reshape(2, 3, 4)) + np.testing.assert_array_equal(file["scalar/en"][:], [2.0, 2.0]) + + if xp.get_backend() == "cupy": + assert [event.kind for event in counter.events] == ["to_host"] * len(device), counter.report() + assert sum(event.nbytes for event in counter.events) == sum(a.nbytes for a in device.values()) + else: + assert counter.events == [] + + +def _flagged(backend): + from struphy.feec.psydac_derham import SplineFunction + + # two processes split along eta1 at 0.5; this is rank 0 + domain_array = np.array([[0.0, 0.5, 4, 0.0, 1.0, 6, 0.0, 1.0, 4], [0.5, 1.0, 4, 0.0, 1.0, 6, 0.0, 1.0, 4]]) + with xp.use_backend(backend): + spline = SplineFunction.__new__(SplineFunction) + spline._derham = SimpleNamespace(domain_array=xp.asarray(domain_array), comm=None) + E1, E2, E3 = (xp.asarray(a) for a in np.meshgrid(*[np.linspace(0.0, 1.0, 5)] * 3, indexing="ij")) + markers = xp.asarray(np.random.default_rng(1).random((10, 3))) + spline._flag_pts_not_on_proc(E1, E2, E3) # first call: caches the host bounds + with xp.profiling.count_transfers() as counter, no_implicit_syncs(): + spline._flag_pts_not_on_proc(E1, E2, E3) + spline._flag_pts_not_on_proc(markers) + return [xp.to_numpy(a) for a in (E1, E2, E3, markers)], counter + + +def check_flag_points(): + """SplineFunction flags the points of other processes on the device; the bounds are cached on the host once.""" + flagged, counter = _flagged(xp.get_backend()) + reference, _ = _flagged("numpy") + assert counter.events == [], counter.report() + assert np.any(reference[0] == -1.0) and np.any(reference[3] == -1.0) + for a, b in zip(flagged, reference): + np.testing.assert_array_equal(a, b) + + +CHECKS = ("check_scalars", "check_binning", "check_saved_markers", "check_output", "check_flag_points") + + +@pytest.mark.mpi_skip +@pytest.mark.parametrize("check", CHECKS) +def test_on_numpy(check): + """The checks on the NumPy backend: no transfers, the same results as before.""" + with xp.use_backend("numpy"): + globals()[check]() + + +@pytest.mark.mpi_skip +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +@pytest.mark.parametrize("check", CHECKS) +def test_on_fake_cupy(check): + """Without a GPU: the checks on cunumpy's fake CuPy, in a subprocess (the fake must be installed first).""" + code = f"import cunumpy as xp; xp.set_backend('cupy'); from {__name__} import {check}; {check}()" + 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:] + + +@pytest.mark.mpi_skip +@requires_cupy +@pytest.mark.parametrize("check", CHECKS) +def test_on_cupy(check): + """With a GPU: the same checks on the real CuPy.""" + with xp.use_backend("cupy"): + globals()[check]() + + +@pytest.mark.mpi_skip +@requires_cupy +def test_eval_tp_fixed_loc_on_cupy(): + """``SplineFunction.eval_tp_fixed_loc`` (the CUDA ``eval_spline_mpi_tensor_product_fixed``) agrees with NumPy.""" + from struphy.feec.tests.test_derham_gpu import make_derham + + pts = [np.linspace(0.05, 0.95, n) for n in (5, 4, 3)] + results = {} + for backend in ("numpy", "cupy"): + with xp.use_backend(backend): + derham = make_derham() + spans, bns, bds = derham.prepare_eval_tp_fixed([xp.asarray(p) for p in pts]) + for space, bases in (("H1", bns), ("L2", bds)): + field = derham.create_spline_function(space, space) + rng = np.random.default_rng(3) + field.vector._data[:] = xp.asarray(rng.random(field.vector._data.shape)) + with xp.profiling.count_transfers() as counter: + out = field.eval_tp_fixed_loc(spans, bases) + assert counter.events == [], counter.report() + results[backend, space] = xp.to_numpy(out) + for space in ("H1", "L2"): + np.testing.assert_allclose(results["cupy", space], results["numpy", space], rtol=1e-13, atol=1e-14) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index 5e1d217b2..ffd2e9fea 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1542,7 +1542,7 @@ def binning( divide_by_jac: bool = True, ): r"""Computes full-f and delta-f distribution functions via marker binning in logical space. - Numpy's histogramdd is used, following the algorithm outlined in :ref:`binning`. + ``xp.histogramdd`` (NumPy or CuPy, on the device of the markers) is used, following the algorithm outlined in :ref:`binning`. Parameters ---------- @@ -1567,7 +1567,8 @@ def binning( The reconstructed delta-f distribution function. """ - assert xp.count_nonzero(components) == len(bin_edges) + # components is a host list: counted with Python, not with xp (CuPy rejects lists) + assert sum(bool(c) for c in components) == len(bin_edges) # volume of a bin bin_vol = 1.0 @@ -1576,43 +1577,51 @@ def binning( # extend components list to number of columns of markers array _n = len(components) - slicing = components + [False] * (self.markers.shape[1] - _n) + slicing = list(components) + [False] * (self.markers.shape[1] - _n) # determine type of output quantity # Note: "density" Literal does not have "_" quantity, *v_axis = output_quantity.rsplit(sep="_", maxsplit=1) v_axis = [int(char) - 1 for char in "".join(v_axis)] # convert dimension axis to index + # the valid markers, gathered once; everything below stays on the device of the markers + markers = self.markers_wo_holes_and_ghost + velocities = markers[:, self.index["vel"]] + # determine histogram weights multiplier if quantity == "density": multiplier = 1 elif quantity == "current": - multiplier = self.velocities[:, v_axis[0]] + multiplier = velocities[:, v_axis[0]] elif quantity == "energy_tensor": - multiplier = self.velocities[:, v_axis[0]] * self.velocities[:, v_axis[1]] + multiplier = velocities[:, v_axis[0]] * velocities[:, v_axis[1]] elif quantity == "heat_flux": - velocity_norm2 = xp.linalg.norm(self.velocities, axis=1) ** 2 - multiplier = velocity_norm2 * self.velocities[:, v_axis[0]] + velocity_norm2 = xp.linalg.norm(velocities, axis=1) ** 2 + multiplier = velocity_norm2 * velocities[:, v_axis[0]] # compute weights of histogram: - _weights0 = self.weights0 * self.Np * multiplier - _weights = self.weights * self.Np * multiplier + _weights0 = markers[:, self.index["w0"]] * self.Np * multiplier + _weights = markers[:, self.index["weights"]] * self.Np * multiplier if divide_by_jac: - _weights /= self.domain.jacobian_det(self.positions, remove_outside=False) + jacobian_det = self.domain.jacobian_det(markers[:, self.index["pos"]], remove_outside=False) + _weights /= jacobian_det # _weights /= self.velocity_jacobian_det(*self.phasespace_coords.T) - _weights0 /= self.domain.jacobian_det(self.positions, remove_outside=False) + _weights0 /= jacobian_det # _weights0 /= self.velocity_jacobian_det(*self.phasespace_coords.T) + # the binned columns, one by one (an index list would have to be uploaded to the device first) + sample = xp.stack([markers[:, i] for i, binned in enumerate(slicing) if binned], axis=1) + f_slice = xp.histogramdd( - self.markers_wo_holes_and_ghost[:, slicing], + sample, bins=bin_edges, weights=_weights0, )[0] df_slice = xp.histogramdd( - self.markers_wo_holes_and_ghost[:, slicing], + sample, bins=bin_edges, weights=_weights, )[0] diff --git a/src/struphy/pic/tests/cuda_parity_cases.py b/src/struphy/pic/tests/cuda_parity_cases.py index a262ae815..1fcac72ac 100644 --- a/src/struphy/pic/tests/cuda_parity_cases.py +++ b/src/struphy/pic/tests/cuda_parity_cases.py @@ -141,6 +141,18 @@ def eval_spline_mpi_grid_args(sparse): return lambda kind: (*evaluation_grid(sparse=sparse), *spline_evaluation_arguments(kind), xp.zeros((7, 5, 4))) +def eval_spline_mpi_tensor_product_fixed_args(kind): + """Spline values on a 5 x 4 x 3 grid from pre-evaluated spans and (random) basis values.""" + _data, args_spline = spline_evaluation_arguments(kind) + kind, pn, starts = args_spline.kind, args_spline.pn, args_spline.starts + rng = np.random.default_rng(7) + degree, kinds = (2, 3, 1), xp.to_numpy(kind) + # spans of the 8 cells: p <= span <= p + 7 + spans = [xp.asarray(rng.integers(p, p + 8, size=n)) for p, n in zip(degree, (5, 4, 3))] + bases = [xp.asarray(rng.random((n, p - k + 1))) for p, k, n in zip(degree, kinds, (5, 4, 3))] + return (*spans, *bases, _data, kind, pn, starts, xp.zeros((5, 4, 3))) + + # ---------------------------------------------------------------- geometry # metric coefficients at 129 markers: every mapping for F, det(DF), DF^(-1) and G^(-1); the identity, DF and G for @@ -333,6 +345,9 @@ def size_of(index, per_thread=1): "eval_spline_mpi_sparse_meshgrid": ParityCases( SPLINE_KINDS, eval_spline_mpi_grid_args(sparse=True), n_threads=size_of(-1) ), + "eval_spline_mpi_tensor_product_fixed": ParityCases( + SPLINE_KINDS, eval_spline_mpi_tensor_product_fixed_args, n_threads=size_of(-1) + ), "kernel_evaluate_pic": ParityCases(KERNEL_EVALUATE_PIC_CASES, kernel_evaluate_pic_args, **GEOMETRY_TOLERANCES), # one thread per grid point: the first three axes of mat_f, shape (n1, n2, n3, 3, 3) "kernel_evaluate": ParityCases( diff --git a/src/struphy/pic/tests/test_kernel_backends.py b/src/struphy/pic/tests/test_kernel_backends.py index e373ddea4..f1364f368 100644 --- a/src/struphy/pic/tests/test_kernel_backends.py +++ b/src/struphy/pic/tests/test_kernel_backends.py @@ -196,7 +196,7 @@ def test_catalog_signatures(): "pic.accumulation.kernels": 16, "pic.diagnostics.kernels": 10, "pic.sph.kernels": 4, - "bsplines.kernels": 3, + "bsplines.kernels": 4, "geometry.kernels": 4, "feec.kernels": 1, "feec.local_projectors.kernels": 8, diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 52347634f..d13fd9b13 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -18,6 +18,7 @@ import cunumpy as xp import h5py +import numpy as np import yaml from cunumpy.kernels import Kernel from feectools.linalg.memory import stencil_matrix_memory @@ -664,11 +665,12 @@ def initialize_data_storage(self): self.data = DataContainer(self.env.path_out, comm=self.comm) - # time quantities (current time value, value in seconds and index) + # time quantities (current time value, value in seconds and index); host bookkeeping of the time loop, + # NumPy on both backends, so that reading them for the loop control and saving them never waits for the device self.time_state = {} - self.time_state["value"] = xp.zeros(1, dtype=float) - self.time_state["value_sec"] = xp.zeros(1, dtype=float) - self.time_state["index"] = xp.zeros(1, dtype=int) + self.time_state["value"] = np.zeros(1, dtype=float) + self.time_state["value_sec"] = np.zeros(1, dtype=float) + self.time_state["index"] = np.zeros(1, dtype=int) # add time quantities to data object for saving for key, val in self.time_state.items(): @@ -1425,10 +1427,10 @@ def _initialize_hdf5_datasets(self, data: DataContainer, size: int): """ # save scalar quantities in group 'scalar/' - for key, scalar in self.model.scalars.dct.items(): - val = scalar.value + # the host copies of the scalars (Scalars.to_host): saving them copies nothing from the device + for key in self.model.scalars.dct: key_scalar = "scalar/" + key - data.add_data({key_scalar: val}) + data.add_data({key_scalar: self.model.scalars.host_value(key)}) with h5py.File(data.file_path, "a") as file: # store grid_info only for runs with 512 ranks or smaller