diff --git a/.github/workflows/reusable-scheduled.yml b/.github/workflows/reusable-scheduled.yml index c24a03184..e4652bd23 100644 --- a/.github/workflows/reusable-scheduled.yml +++ b/.github/workflows/reusable-scheduled.yml @@ -107,7 +107,7 @@ jobs: TEST_FILES="bsplines/tests/ console/tests/ fields_background/tests/ geometry/tests/ initial/tests/ kinetic_background/tests/ linear_algebra/tests/" ;; shard-2) - TEST_FILES="propagators/tests/ ode/tests/ polar/tests/ post_processing/tests/ dispersion_relations/tests/ io/tests/ simulation/tests/" + TEST_FILES="propagators/tests/ ode/tests/ polar/tests/ post_processing/tests/ dispersion_relations/tests/ io/tests/ simulation/tests/ models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py models/tests/test_gpu_VlasovAmpereOneSpecies.py" ;; shard-3) TEST_FILES="feec/tests/" diff --git a/.github/workflows/reusable-unit-testing.yml b/.github/workflows/reusable-unit-testing.yml index 2e8f59e7b..89944e67b 100644 --- a/.github/workflows/reusable-unit-testing.yml +++ b/.github/workflows/reusable-unit-testing.yml @@ -39,7 +39,7 @@ jobs: - shard: shard-1 test_files: bsplines/tests/ console/tests/ fields_background/tests/ geometry/tests/ initial/tests/ kinetic_background/tests/ linear_algebra/tests/ - shard: shard-2 - test_files: propagators/tests/ ode/tests/ polar/tests/ post_processing/tests/ dispersion_relations/tests/ io/tests/ simulation/tests/ + test_files: propagators/tests/ ode/tests/ polar/tests/ post_processing/tests/ dispersion_relations/tests/ io/tests/ simulation/tests/ models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py models/tests/test_gpu_VlasovAmpereOneSpecies.py - shard: shard-3 test_files: feec/tests/ - shard: shard-4 diff --git a/.github/workflows/test-PR-unit-mpi.yml b/.github/workflows/test-PR-unit-mpi.yml index 6a5b1c387..1743e82a9 100644 --- a/.github/workflows/test-PR-unit-mpi.yml +++ b/.github/workflows/test-PR-unit-mpi.yml @@ -88,3 +88,11 @@ jobs: source /struphy_fortran_/env_fortran_/bin/activate cd ${{ env.STRUPHY_PATH }} mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi polar/tests/test_polar_cupy.py -k mpi_fake + + - name: Run the end-to-end GPU model tests on the fake CuPy + # 8 elements of degree 3: at most 2 ranks + if: matrix.n-procs <= 2 + run: | + source /struphy_fortran_/env_fortran_/bin/activate + cd ${{ env.STRUPHY_PATH }} + mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py models/tests/test_gpu_VlasovAmpereOneSpecies.py diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index a7724a594..69f8866b5 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -31,6 +31,7 @@ Progress is tracked in struphy-hub/struphy#650. - [x] **FEEC assembly kernels in kernel folders** (#411) — the 16 assembly kernels of `feec/mass.py`, `feec/boundary_mass.py` and `feec/basis_projection_ops.py` move from `feec/mass_kernels.py` and `feec/basis_projection_kernels.py` into one folder each under `feec/kernels/`; pure restructuring, ready for CUDA ports (see [FEEC assembly kernels](#feec-assembly-kernels-411)). - [ ] **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)). +- [ ] **End-to-end tests of `LinearVlasovAmpereOneSpecies` and `VlasovAmpereOneSpecies`** (#689) — `models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py` and `test_gpu_VlasovAmpereOneSpecies.py` (shared helpers in `models/tests/gpu_e2e.py`) run a small Cuboid case (320 Sobol markers, 8 elements, 3 steps; for `VlasovAmpereOneSpecies` Landau damping with full-f weights and an `e1_v1` binning plot) on NumPy and CuPy, compare energies, field coefficients and markers, and count the transfers of the time loop: in `model.integrate` only the 8-byte residual norm of each Krylov iteration (feectools' convergence test), scalar-sized downloads in the diagnostics, one download per saved dataset in the output. They run on a GPU, or on the fake CuPy with every CUDA launch emulated on the CPU (`CUNUMPY_FAKE_CUPY=1`, also under `mpirun -n 2`; `CudaKernel.compile` is skipped there); without either, `test_fake_cupy_matches_numpy` runs the comparison in a serial child process. Both pass on the fake CuPy on 1 and 2 ranks (not yet on a GPU); on 2 ranks only with the markers MPI-sorted after loading (#764). CI runs them in unit shard-2 (serial and on 2 ranks on NumPy) and on the fake CuPy on 1 and 2 ranks (`test-PR-unit-mpi.yml`). The setup still assembles the mass matrices and vectors with host kernels (`kernel_1d_mat`, `kernel_3d_mat`, `kernel_3d_vec`, whose arrays are copied to the host and back at every call; outside the time loop, so not in the transfer budget); CUDA versions of these are open. Needs the device Kronecker solve, the 6D stencil views, the `linear_vlasov_ampere`/`vlasov_maxwell` CUDA kernels, `MassMatrixPreconditioner` on CuPy and device inner products, and a feectools release with them. - [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). @@ -274,6 +275,16 @@ the `pyproject.toml` pin (today `feectools>=0.3.0, <=0.3.0`) point at that relea deprecated (removed in 0.6): import them from `cunumpy.cuda` and `cunumpy.mpi`, and require `cunumpy>=0.5.0, <0.6` like struphy. +**Mass-matrix preconditioners on CuPy** (`feec/preconditioner.py`, for #689). The Kronecker preconditioners +(`MassMatrixPreconditioner`) build their 1d mass matrices and 1d solvers +(`FFTSolver` for circulant, `SparseSolver` otherwise) on the host on every backend; the process-local 1d stencil +matrices (factors of the `KroneckerStencilMatrix`) and the diagonal scaling live on the device. Applying them makes +no host copies: `KroneckerLinearSolver` solves device data with dense inverses built once from the host 1d solvers +([feectools#96](https://github.com/struphy-hub/feectools/pull/96)). Tested with fake CuPy, emulated launches and +`assert_no_transfers` (`feec/tests/test_preconditioner_cupy.py`). `StiffnessPreconditioner` builds its 1d data the +same way, but `KroneckerSumSolver` (feectools 0.7.0) does not yet run on the device: its dense "solvers" hold device +matrices, from which `KroneckerLinearSolver` tries to build dense inverses on the host. + New feectools work for the GPU follows the same pattern: a `cuda--` PR, linked in the table above and from the struphy tracking issue ([#650](https://github.com/struphy-hub/struphy/issues/650)). diff --git a/src/struphy/console/test.py b/src/struphy/console/test.py index 454ee1050..0af89df1f 100644 --- a/src/struphy/console/test.py +++ b/src/struphy/console/test.py @@ -60,6 +60,9 @@ def struphy_test( f"{LIBPATH}/post_processing/tests/", f"{LIBPATH}/propagators/tests/", f"{LIBPATH}/simulation/tests/", + # end-to-end GPU tests (NumPy reference, and the CuPy comparison on a GPU or the fake CuPy) + f"{LIBPATH}/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py", + f"{LIBPATH}/models/tests/test_gpu_VlasovAmpereOneSpecies.py", ] if mpi > 1: diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index abccfb449..97ca621f3 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -332,7 +332,10 @@ def weight_1d(c: int, d: int) -> Callable | xp.ndarray: return _kronecker_approximation(mass_operator, bc, weight_1d) def _core_diagonal(self): - return self._mass_operator.matrix.diagonal() + # diagonal of the mass matrix, including the band entries that wrap onto the main diagonal + # (StencilMatrix.diagonal() keeps only the entry at offset 0) + M = self._mass_operator.matrix + return _diagonal_operator(M.domain, _local_diagonal(M)) @property def dim_reduce(self) -> int | None: @@ -870,13 +873,11 @@ def _build_approximation(self, bc): stiffness.append(S_d) mass.append(M_d) - local_mass.append(_process_local_matrix_1d(xp.asarray(M_d), comp.coeff_space, d, domain_decomposition)) + local_mass.append(_process_local_matrix_1d(M_d, comp.coeff_space, d, domain_decomposition)) local_stiffness.append( None if S_d is None - else _process_local_matrix_1d( - xp.asarray(coeffs[d] * S_d), comp.coeff_space, d, domain_decomposition - ) + else _process_local_matrix_1d(coeffs[d] * S_d, comp.coeff_space, d, domain_decomposition) ) # stabilization: sigma_c * (M_1 x M_2 x M_3 - F_1 x F_2 x F_3), or the identity @@ -1156,7 +1157,7 @@ def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]: if isinstance(op, StencilDiagonalMatrix): return [op._data] if isinstance(op, StencilMatrix): - return [op.diagonal()._data] + return [_stencil_local_diagonal(op)] if isinstance(op, KroneckerStencilMatrix): # outer product of the diagonals of the (process-local) factors on the local rows diags = [] @@ -1165,9 +1166,10 @@ def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]: # the factor may own more rows than the codomain on this process off = A.codomain.pads[0] * A.codomain.shifts[0] + s - A.codomain.starts[0] diags.append(A._data[off : off + e - s + 1, A.pads[0]]) + # outer product by broadcasting (works on every backend, unlike ufunc.outer) out = diags[0] for d in diags[1:]: - out = xp.multiply.outer(out, d) + out = out[..., None] * d return [out] if isinstance(op, _DenseKroneckerProduct): return [op.local_diagonal()] @@ -1179,6 +1181,31 @@ def _local_diagonal(op: LinearOperator) -> list[xp.ndarray]: raise NotImplementedError(f"Diagonal of {type(op).__name__} is not supported.") +def _stencil_local_diagonal(M: StencilMatrix) -> xp.ndarray: + """ + Local diagonal of a StencilMatrix, including the band entries that wrap onto the main diagonal. + + In a periodic direction with fewer points than the band width (``2p + 1 > n``, e.g. a single + element), the band entries at the offsets ``k = m n`` (``m != 0``) are on the main diagonal too; + :meth:`StencilMatrix.diagonal` returns only the entry at ``k = 0``. + """ + V = M.codomain + assert all(m == 1 for m in V.shifts) and all(m == 1 for m in M.domain.shifts), "Only shifts 1 are supported." + offsets = [ + [k for k in range(-p, p + 1) if k % n == 0] if periodic else [0] + for n, p, periodic in zip(V.npts, M.pads, V.periods) + ] + diag = M.diagonal()._data + if all(o == [0] for o in offsets): + return diag + idx = M._get_diagonal_indices() + rows, cols = idx[: V.ndim], idx[V.ndim :] + for ks in itertools.product(*offsets): + if any(ks): + diag = diag + M._data[rows + tuple(j + k for j, k in zip(cols, ks))] + return diag + + def _diagonal_operator(space: VectorSpace, diags: list[xp.ndarray]) -> LinearOperator: """Diagonal operator on ``space`` (Stencil- or BlockVectorSpace) with the given local diagonals.""" if isinstance(space, BlockVectorSpace): @@ -1450,27 +1477,34 @@ def _apply_bc_1d(M: StencilMatrix, bc_d: tuple[bool, bool]) -> None: apply_essential_bc_stencil(M, axis=0, ext=ext, order=0, identity=True) -def _solver_1d(M: StencilMatrix, M_arr: xp.ndarray) -> FFTSolver | SparseSolver: - """Direct solver for the 1d matrix M (dense copy M_arr): FFT if circulant, else sparse LU.""" +def _solver_1d(M: StencilMatrix, M_arr: np.ndarray) -> FFTSolver | SparseSolver: + """ + Direct solver for the 1d matrix M (dense host copy M_arr): FFT if circulant, else sparse LU. + + The solver is a host solver on every backend, as expected by KroneckerLinearSolver: on the + CuPy backend, the latter solves device data with the dense inverse it builds from this solver + once, so that applying the preconditioner makes no host copies. + """ if is_circulant(M_arr): return FFTSolver(M_arr) return SparseSolver(M.tosparse()) def _process_local_matrix_1d( - M_arr: xp.ndarray, coeff_space: StencilVectorSpace, d: int, domain_decomposition: DomainDecomposition + M_arr: np.ndarray, coeff_space: StencilVectorSpace, d: int, domain_decomposition: DomainDecomposition ) -> StencilMatrix: """ Process-local 1d factor of a KroneckerStencilMatrix on ``coeff_space``. The factor lives on a 1d space without communicator that owns the same rows (``starts[d]`` to ``ends[d]``) as ``coeff_space`` in direction ``d`` on this process, - as required by KroneckerStencilMatrix. + as required by KroneckerStencilMatrix. The band is filled on the host and copied to the + data of the factor (active backend) in one step. Parameters ---------- - M_arr : xp.ndarray - The global 1d matrix as dense array. + M_arr : np.ndarray + The global 1d matrix as dense array (a device array is copied to the host). coeff_space : StencilVectorSpace The (distributed) 3d coefficient space. @@ -1489,21 +1523,24 @@ def _process_local_matrix_1d( cart_1d = CartDecomposition(domain_decomposition, [n], [[s]], [[e]], [p], [1]) V_local = StencilVectorSpace(cart_1d) M_local = StencilMatrix(V_local, V_local) + M_arr = xp.to_numpy(M_arr) # copy the rows owned by this process: entry (i, j) is stored at row i - s + p and diagonal # index k + p with the offset k = j - i; in periodic directions k is taken nearest to 0 # (in [-n/2, n/2)), so that the main diagonal is at k = 0 also if 2p + 1 > n - rows, cols = xp.nonzero(M_arr) + rows, cols = np.nonzero(M_arr) on_process = (rows >= s) & (rows <= e) rows, cols = rows[on_process], cols[on_process] k = cols - rows if coeff_space.periods[d]: k = (k + n // 2) % n - n // 2 - assert xp.all(xp.abs(k) <= p), "Entries outside the band of the stencil matrix." - M_local._data[rows - s + p, k + p] = M_arr[rows, cols] + assert np.all(np.abs(k) <= p), "Entries outside the band of the stencil matrix." + data = np.zeros(M_local._data.shape, dtype=M_local._data.dtype) + data[rows - s + p, k + p] = M_arr[rows, cols] + M_local._data[...] = xp.asarray(data) # check if stencil matrix was built correctly - assert xp.allclose(M_local.toarray()[s : e + 1], M_arr[s : e + 1]) + assert np.allclose(xp.to_numpy(M_local.toarray())[s : e + 1], M_arr[s : e + 1]) return M_local @@ -1556,7 +1593,8 @@ def _kronecker_approximation( if bc is not None and _has_essential_bc(mass_operator, basis, c, d): _apply_bc_1d(M, bc[d]) - M_arr = M.toarray() + # the dense 1d matrix and the 1d solver are host setup data on every backend + M_arr = xp.to_numpy(M.toarray()) solvercells.append(_solver_1d(M, M_arr)) matrixcells.append(_process_local_matrix_1d(M_arr, femspace_c.coeff_space, d, domain_decomposition)) @@ -1885,26 +1923,31 @@ class FFTSolver(BandedSolver): Solve the equation Ax = b for x, assuming A is a circulant matrix. b can contain multiple right-hand sides (RHS) and is of shape (#RHS, N). + A host (NumPy) solver on every backend: the 1d matrices of Kronecker solvers are small + setup data. On the CuPy backend, KroneckerLinearSolver solves device data with the dense + inverse it builds from this solver (by solving for the identity on the host), so ``solve`` + is only called with host arrays. + Parameters ---------- - circmat : xp.ndarray - Generic circulant matrix. + circmat : np.ndarray + Generic circulant matrix (a device array is copied to the host). """ def __init__(self, circmat): - assert isinstance(circmat, xp.ndarray) + circmat = xp.to_numpy(circmat) assert is_circulant(circmat) from scipy.linalg import solve_circulant # deferred: scipy.linalg is slow to import - self._space = xp.ndarray + self._space = np.ndarray # copy: circmat is still used by the caller (e.g. for the process-local stencil matrix) - self._column = xp.array(circmat[:, 0], copy=True) + self._column = np.array(circmat[:, 0], copy=True) # stabilize a singular matrix once, here, so that all solves use the same matrix try: - solve_circulant(self._column, xp.ones_like(self._column)) - except xp.linalg.LinAlgError: + solve_circulant(self._column, np.ones_like(self._column)) + except np.linalg.LinAlgError: eps = 1e-4 logger.info(f"Stabilizing singular preconditioning FFTSolver with {eps =}:") self._column[0] *= 1.0 + eps @@ -1923,13 +1966,13 @@ def solve(self, rhs, out=None, transposed=False): Parameters ---------- - rhs : xp.ndarray - The right-hand sides to solve for. The vectors are assumed to be given in C-contiguous order, + rhs : np.ndarray + The right-hand sides to solve for (host array). The vectors are assumed to be given in C-contiguous order, i.e. if multiple right-hand sides are given, then rhs is a two-dimensional array with the 0-th index denoting the number of the right-hand side, and the 1-st index denoting the element inside a right-hand side. - out : xp.ndarray, optional + out : np.ndarray, optional Output vector. If given, it has to have the same shape and datatype as rhs. transposed : bool @@ -1938,6 +1981,7 @@ def solve(self, rhs, out=None, transposed=False): from scipy.linalg import solve_circulant + assert isinstance(rhs, np.ndarray), f"FFTSolver solves on the host, got {type(rhs).__name__}." assert rhs.T.shape[0] == self._column.size if out is None: @@ -1956,7 +2000,7 @@ def is_circulant(mat): Parameters ---------- mat : array[float] - The matrix that is checked to be circulant. + The matrix that is checked to be circulant (host or device array; checked on the host). Returns ------- @@ -1964,11 +2008,11 @@ def is_circulant(mat): Whether the matrix is circulant (=True) or not (=False). """ - assert isinstance(mat, xp.ndarray) + mat = xp.to_numpy(mat) assert len(mat.shape) == 2 assert mat.shape[0] == mat.shape[1] # circulant: every row is the previous one shifted by one, i.e. mat[i, j] = mat[0, (j - i) % n] n = mat.shape[0] - idx = (xp.arange(n)[None, :] - xp.arange(n)[:, None]) % n - return bool(xp.allclose(mat, mat[0][idx])) + idx = (np.arange(n)[None, :] - np.arange(n)[:, None]) % n + return bool(np.allclose(mat, mat[0][idx])) diff --git a/src/struphy/feec/tests/test_mass_matrices.py b/src/struphy/feec/tests/test_mass_matrices.py index 23af76b82..799646edf 100644 --- a/src/struphy/feec/tests/test_mass_matrices.py +++ b/src/struphy/feec/tests/test_mass_matrices.py @@ -1742,6 +1742,40 @@ def test_transpose_and_copy(): assert xp.allclose(A.T.toarray(), -A_arr, atol=1e-14) +@pytest.mark.parametrize( + "num_elements, degree", [((16, 1, 1), (3, 1, 1)), ((16, 2, 1), (2, 2, 1)), ((8, 6, 4), (2, 2, 2))] +) +@pytest.mark.parametrize("space", ["M0", "M1", "M2", "M3"]) +def test_mass_preconditioner_diagonal_scaling_periodic_exact(num_elements, degree, space): + """On the periodic logical cube the mass matrices are Kronecker products, so MassMatrixPreconditioner with + diagonal scaling inverts them exactly; also with fewer points than the band width in a direction (one element), + where band entries wrap onto the main diagonal.""" + + import numpy as np + from maybempi import MPI + + from struphy import domains + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.preconditioner import MassMatrixPreconditioner + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + domain = domains.Cuboid(r1=12.56) + derham = Derham(TensorProductGrid(num_elements=num_elements), DerhamOptions(degree=degree), comm=MPI.COMM_WORLD) + M = getattr(WeightedMassOperators(derham, domain), space) + + rng = np.random.default_rng(0) + v = M.domain.zeros() + for block in v.blocks if hasattr(v, "blocks") else [v]: + block._data[:] = rng.random(block._data.shape) + v.update_ghost_regions() + + w = MassMatrixPreconditioner(M, diagonal_scaling=True).dot(M.dot(v)) + err = ((w - v).inner(w - v) / v.inner(v)) ** 0.5 + assert err < 1e-12, f"{space}, {num_elements=}, {degree=}: |P M v - v| / |v| = {err}" + + if __name__ == "__main__": # test_mass( # num_elements=(32, 32, 32), diff --git a/src/struphy/feec/tests/test_preconditioner_cupy.py b/src/struphy/feec/tests/test_preconditioner_cupy.py new file mode 100644 index 000000000..9139e592b --- /dev/null +++ b/src/struphy/feec/tests/test_preconditioner_cupy.py @@ -0,0 +1,150 @@ +"""Mass-matrix preconditioners on the CuPy backend (struphy-hub/struphy#689, part of #650). + +The 1d mass matrices and solvers of the Kronecker preconditioners are host (NumPy) setup data on every +backend; the Kronecker stencil matrices, the diagonal scaling and the vectors they are applied to live on +the device, and feectools' KroneckerLinearSolver solves device data with dense inverses built from the host +1d solvers. These tests build the preconditioners on the CuPy backend, check that applying them makes no +host/device transfers, and compare their action with the NumPy backend: + +* without a GPU, with cunumpy's fake CuPy (host memory, CuPy semantics) in a subprocess, the CUDA kernels + launched on the way (feectools' stencil kernels, struphy's geometry kernels) run by CPU emulation; +* on a GPU (skipped otherwise), directly. +""" + +import cunumpy +import numpy as np +import pytest + +from struphy.geometry.tests.test_domain import _cupy_installed, run_fake_cupy_child + +requires_cupy = pytest.mark.skipif(not cunumpy.cupy_available(), reason="CuPy/GPU not available") + +SPACES = ("M0", "M1", "M2") +# (class name, keyword arguments) +PRECONDITIONERS = ( + ("MassMatrixPreconditioner", {"diagonal_scaling": False}), + ("MassMatrixPreconditioner", {}), + ("MassMatrixPreconditioner", {"weight_reduction": "average"}), + ("MassMatrixPreconditioner", {"dim_reduce": None}), +) +BCS = { + "periodic": (None, None, None), + "clamped": (("dirichlet", "dirichlet"), ("free", "free"), ("free", "dirichlet")), +} + + +def _build(backend, bcs): + """Derham (6 x 5 x 4 elements, degrees (2, 3, 1)) and mass operators on a Colella mapping, on `backend`.""" + from maybempi import MPI + + from struphy import domains + from struphy.feec.mass import WeightedMassOperators + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.topology.grids import TensorProductGrid + + with cunumpy.use_backend(backend): + options = DerhamOptions(degree=(2, 3, 1), bcs=bcs) + derham = Derham(TensorProductGrid(num_elements=(6, 5, 4)), options, comm=MPI.COMM_WORLD) + return WeightedMassOperators(derham, domains.Colella()) + + +def _blocks(v): + """The stencil vectors of a stencil or block vector.""" + return list(v.blocks) if hasattr(v, "blocks") else [v] + + +def _random_vector(space, rng): + v = space.zeros() + for block in _blocks(v): + block._data[...] = cunumpy.asarray(rng.random(block._data.shape)) + v.update_ghost_regions() + return v + + +def _to_device(v_host, space): + """A vector of `space` (active backend) with the data of the host vector `v_host`.""" + v = space.zeros() + for block, block_host in zip(_blocks(v), _blocks(v_host)): + block._data[...] = cunumpy.asarray(block_host._data) + v.update_ghost_regions() + return v + + +def _to_host(v): + return np.concatenate([cunumpy.to_numpy(block._data).ravel() for block in _blocks(v)]) + + +def check_preconditioners_on_cupy(bcs_name): + """Every preconditioner of :data:`PRECONDITIONERS` for M0, M1 and M2 gives the NumPy result on the CuPy backend. + + Applying the preconditioners (``dot``, also in place) makes no host/device transfers. + """ + from struphy.feec import preconditioner + + bcs = BCS[bcs_name] + host, device = _build("numpy", bcs), _build("cupy", bcs) + rng = np.random.default_rng(1) + + for space in SPACES: + with cunumpy.use_backend("numpy"): + M_host = getattr(host, space) + v_host = _random_vector(M_host.domain, rng) + with cunumpy.use_backend("cupy"): + M_device = getattr(device, space) + v_device = _to_device(v_host, M_device.domain) + + for name, kwargs in PRECONDITIONERS: + cls = getattr(preconditioner, name) + case = (bcs_name, space, name, kwargs) + with cunumpy.use_backend("numpy"): + expected = _to_host(cls(M_host, **kwargs).dot(v_host)) + with cunumpy.use_backend("cupy"): + pc = cls(M_device, **kwargs) + out_inplace = pc.codomain.zeros() + with cunumpy.profiling.assert_no_transfers(): + out = pc.dot(v_device) + # in-place application + pc.dot(v_device, out=out_inplace) + for block in _blocks(out): + assert cunumpy.is_gpu(block._data), case + assert np.linalg.norm(expected) > 0, case + np.testing.assert_allclose(_to_host(out), expected, rtol=1e-12, atol=1e-12 * np.abs(expected).max()) + np.testing.assert_allclose(_to_host(out_inplace), expected, rtol=1e-12, atol=1e-12 * np.abs(expected).max()) + + +def check_preconditioners_on_fake_cupy(bcs_name): + """:func:`check_preconditioners_on_cupy` with the CUDA kernel launches emulated on the CPU.""" + from cunumpy.kernel_testing import emulated_launches, fake_cupy_active + + from struphy.pic.tests.cuda_emulation import EMULATION_OPTIONS + + assert fake_cupy_active() + with emulated_launches(options=EMULATION_OPTIONS): + check_preconditioners_on_cupy(bcs_name) + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +@pytest.mark.parametrize("bcs_name", list(BCS)) +def test_mass_preconditioners_fake_cupy(bcs_name): + """Without a GPU: the preconditioners build and apply on cunumpy's fake CuPy, which rejects host/device mixing. + + Runs in a serial subprocess (on MPI rank 0 only) because the fake CuPy must be installed before cunumpy is imported. + """ + run_fake_cupy_child( + "from struphy.feec.tests.test_preconditioner_cupy import check_preconditioners_on_fake_cupy; " + f"check_preconditioners_on_fake_cupy({bcs_name!r})" + ) + + +@requires_cupy +@pytest.mark.parametrize("bcs_name", list(BCS)) +def test_mass_preconditioners_on_cupy(bcs_name): + """On a GPU, or in this process on the fake CuPy (``CUNUMPY_FAKE_CUPY=1``, launches emulated on the CPU): the + preconditioners give the NumPy result on the CuPy backend.""" + from cunumpy.kernel_testing import fake_cupy_active + + if fake_cupy_active(): + check_preconditioners_on_fake_cupy(bcs_name) + else: + check_preconditioners_on_cupy(bcs_name) diff --git a/src/struphy/models/tests/gpu_e2e.py b/src/struphy/models/tests/gpu_e2e.py new file mode 100644 index 000000000..332a48492 --- /dev/null +++ b/src/struphy/models/tests/gpu_e2e.py @@ -0,0 +1,255 @@ +"""Shared helpers of the end-to-end GPU tests of the models (issue #689). + +A test runs a small simulation once on the NumPy backend and once on the CuPy backend, compares the two, and checks +that the time loop of the CuPy run makes no host/device transfers through cunumpy, apart from an explicit budget for +the diagnostics and the output (:class:`TransferBudget`). + +The CuPy run needs a GPU, or cunumpy's fake CuPy (``CUNUMPY_FAKE_CUPY=1``), on which every CUDA launch is emulated +on the CPU (:func:`device_context`). Without either, a test can still run the comparison in a serial child process +on the fake CuPy (:func:`run_fake_cupy_child`). + +Only transfers made through cunumpy are counted (``to_numpy``, ``to_cupy``, host fallbacks of kernels, ...). +Implicit synchronizations such as ``float(device_scalar)`` are not seen; use ``nsys`` for those. +""" + +import contextlib +import os +from collections.abc import Callable +from dataclasses import dataclass, field + +import cunumpy as xp +import h5py +import numpy as np +import pytest +from cunumpy.kernel_testing import fake_cupy_active +from cunumpy.profiling import TransferCounter, count_transfers +from maybempi import MPI + +from struphy.io.output_handling import DataContainer + +# transfer kinds that move data between host and device (device_copy is device-only and allowed) +HOST_DEVICE_KINDS = ("to_host", "to_device", "kernel_conversion", "fallback") + + +def cupy_backend_available() -> bool: + """Whether the CuPy backend can run here: a GPU, or the fake CuPy (with emulated launches).""" + return fake_cupy_active() or xp.cupy_available() + + +requires_cupy_backend = pytest.mark.skipif( + not cupy_backend_available(), reason="needs a GPU or the fake CuPy (CUNUMPY_FAKE_CUPY=1)" +) + + +@contextlib.contextmanager +def device_context(): + """Context for the CuPy run: on the fake CuPy every CUDA launch is emulated on the CPU, on a GPU nothing. + + On the fake CuPy, :meth:`CudaKernel.compile` (called by ``Simulation.compile_cuda_kernels`` before the time loop) + does nothing: the fake CuPy cannot compile CUDA, and the emulated launches compile their own C++ version. + ``compile_cuda_kernels`` still checks that every kernel of the time loop has a CUDA version. + """ + if not fake_cupy_active(): + yield + return + from unittest import mock + + from cunumpy.kernel_testing import emulated_launches + from cunumpy.kernels import CudaKernel + + from struphy.pic.tests.cuda_emulation import EMULATION_OPTIONS + + with ( + emulated_launches(options=EMULATION_OPTIONS), + mock.patch.object(CudaKernel, "compile", lambda self, **kw: None), + ): + yield + + +@dataclass +class TransferBudget: + """Transfers recorded per phase of the time loop (the setup before it is not counted). + + ``integrate`` (the propagators) may only copy scalars to the host: the convergence test of the Krylov solvers + reads the 8-byte residual norm once per iteration (feectools keeps everything else of the solve on the device). + No arrays, no host-to-device copies, no host kernels. ``diagnostics`` (scalars, saved markers, binned + distribution functions) may copy small results to the host: at most one scalar per tracked scalar and call, + never a host-to-device copy or a host kernel. ``output`` (``DataContainer.save_data``) may copy each saved + dataset to the host once, nothing more. + """ + + n_steps: int + integrate: TransferCounter = field(default_factory=TransferCounter) + diagnostics: TransferCounter = field(default_factory=TransferCounter) + output: TransferCounter = field(default_factory=TransferCounter) + # allowed device-to-host copies by `output`: one per saved dataset, and their bytes + output_allowed_copies: int = 0 + output_allowed_bytes: int = 0 + n_integrate_calls: int = 0 + n_diagnostics_calls: int = 0 + + def report(self) -> str: + return ( + f"integrate:\n{self.integrate.report()}\n" + f"diagnostics:\n{self.diagnostics.report()}\n" + f"output:\n{self.output.report()}" + ) + + def check(self, n_scalars: int): + report = self.report() + assert self.n_integrate_calls == self.n_steps, report + + # 1. the propagators: only scalars to the host (solver convergence tests), nothing else + for e in self.integrate.events: + if e.kind == "to_host": + assert e.nbytes is not None and e.nbytes <= 8, report + else: + assert e.kind not in HOST_DEVICE_KINDS, report + + # 2. diagnostics: no host-to-device copies, no host kernels, only scalar-sized downloads + for kind in ("to_device", "kernel_conversion", "fallback"): + assert self.diagnostics.count(kind) == 0, report + assert self.diagnostics.to_host <= n_scalars * self.n_diagnostics_calls, report + assert self.diagnostics.bytes_to_host <= 8 * n_scalars * self.n_diagnostics_calls, report + + # 3. output: each saved dataset once to the host + for kind in ("to_device", "kernel_conversion", "fallback"): + assert self.output.count(kind) == 0, report + assert self.output.to_host <= self.output_allowed_copies, report + assert self.output.bytes_to_host <= self.output_allowed_bytes, report + + +class _Patcher: + """Minimal ``monkeypatch``: set attributes, and restore them in reverse order.""" + + def __init__(self): + self._saved = [] + + def setattr(self, obj, name, value): + self._saved.append((obj, name, getattr(obj, name))) + setattr(obj, name, value) + + def undo(self): + while self._saved: + obj, name, value = self._saved.pop() + setattr(obj, name, value) + + +def instrument(sim, patcher, n_steps: int) -> TransferBudget: + """Count the transfers of the time loop of ``sim.run()`` per phase. + + Counting starts once the HDF5 datasets are created, the last step of the setup in :meth:`Simulation.run`. The + model methods are patched on the instance, ``save_data`` on the class (the ``DataContainer`` is created inside + ``run``). + """ + budget = TransferBudget(n_steps=n_steps) + in_loop = [False] + model = sim.model + + def counted(fn, phase: str, calls: str | None = None): + def wrapper(*args, **kwargs): + if not in_loop[0]: + return fn(*args, **kwargs) + with count_transfers() as counter: + out = fn(*args, **kwargs) + getattr(budget, phase).events.extend(counter.events) + if calls is not None: + setattr(budget, calls, getattr(budget, calls) + 1) + return out + + return wrapper + + original_init_datasets = sim._initialize_hdf5_datasets + + def init_datasets(*args, **kwargs): + out = original_init_datasets(*args, **kwargs) + in_loop[0] = True + return out + + patcher.setattr(sim, "_initialize_hdf5_datasets", init_datasets) + patcher.setattr(model, "integrate", counted(model.integrate, "integrate", "n_integrate_calls")) + patcher.setattr( + model, + "update_scalar_quantities", + counted(model.update_scalar_quantities, "diagnostics", "n_diagnostics_calls"), + ) + for name in ("update_markers_to_be_saved", "update_distr_functions"): + patcher.setattr(model, name, counted(getattr(model, name), "diagnostics")) + + original_save_data = DataContainer.save_data + + def save_data(self, keys=None): + if not in_loop[0]: + return original_save_data(self, keys=keys) + for key in self._dset_dict if keys is None else keys: + val = self._dset_dict[key] + budget.output_allowed_copies += 1 + budget.output_allowed_bytes += int(getattr(val, "nbytes", 8)) + with count_transfers() as counter: + original_save_data(self, keys=keys) + budget.output.events.extend(counter.events) + + patcher.setattr(DataContainer, "save_data", save_data) + return budget + + +@dataclass +class RunResult: + """Host copies of what is compared between the backends.""" + + scalars: dict # name -> time series (rank 0 only, else empty) + fields: dict # FEEC variable name -> local coefficients at the end + markers: np.ndarray # all valid markers (gathered from all ranks), sorted by marker ID + n_scalars: int + budget: TransferBudget + + +def gather_markers(particles) -> np.ndarray: + """The valid markers of all ranks on the host, sorted by marker ID (the last column).""" + comm = MPI.COMM_WORLD + local = np.asarray(xp.to_numpy(particles.markers_wo_holes_and_ghost)) + parts = comm.allgather(local) if comm.Get_size() > 1 else [local] + markers = np.concatenate(parts, axis=0) + return markers[np.argsort(markers[:, -1])] + + +def run( + backend: str, + make_simulation: Callable[[str, str], object], + out_folders: str, + n_steps: int, + fields: Callable[[object], dict], + particles: Callable[[object], object], +) -> RunResult: + """Run the simulation made by ``make_simulation(out_folders, sim_folder)`` on ``backend``; copy results to host. + + ``fields(model)`` returns the FEEC variables to compare (name -> ``FEECVariable``), ``particles(model)`` the + ``Particles`` object whose markers are compared. + """ + with xp.use_backend(backend), device_context() if backend == "cupy" else contextlib.nullcontext(): + sim = make_simulation(out_folders, backend) + patcher = _Patcher() + budget = instrument(sim, patcher, n_steps) + try: + sim.run() + finally: + patcher.undo() + + model = sim.model + coeffs = {name: np.asarray(xp.to_numpy(var.spline.vector.toarray())) for name, var in fields(model).items()} + markers = gather_markers(particles(model)) + + scalars = {} + if MPI.COMM_WORLD.Get_rank() == 0: + with h5py.File(os.path.join(sim.env.path_out, "data", "data_proc0.hdf5"), "r") as f: + for key in model.scalars.dct: + scalars[key] = f["scalar"][key][()].ravel() + + return RunResult(scalars=scalars, fields=coeffs, markers=markers, n_scalars=len(model.scalars.dct), budget=budget) + + +def shared_out_folder(tmp_path_factory, name: str) -> str: + """One output folder for all ranks (the one of rank 0; ``tmp_path`` differs per rank).""" + comm = MPI.COMM_WORLD + path = str(tmp_path_factory.mktemp(name)) if comm.Get_rank() == 0 else None + return comm.bcast(path, root=0) diff --git a/src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py b/src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py new file mode 100644 index 000000000..b6a9a6f51 --- /dev/null +++ b/src/struphy/models/tests/test_gpu_LinearVlasovAmpereOneSpecies.py @@ -0,0 +1,194 @@ +"""End-to-end GPU test of :class:`~struphy.models.LinearVlasovAmpereOneSpecies` (issue #689). + +A small run of the model (the energy-conservation setup of +``verification/test_verif_LinearVlasovAmpereOneSpecies.py`` with fewer markers) is done once on +the NumPy backend and once on the CuPy backend. The test checks that + +1. both runs give the same energies, electric-field coefficients and marker weights/positions, + within tolerances explained in :func:`compare_runs`, and +2. the time loop of the CuPy run does no host/device transfers through cunumpy, apart from an + explicit budget for the diagnostics and the output (see :class:`~struphy.models.tests.gpu_e2e.TransferBudget`). + +The CuPy part needs a GPU, or cunumpy's fake CuPy (``CUNUMPY_FAKE_CUPY=1``), on which the CUDA launches are emulated +on the CPU, and is skipped otherwise. It is MPI-compatible: under +``mpirun -n 2 pytest --with-mpi`` both runs use the same domain decomposition, and the markers are +gathered and compared by marker ID. + +The shared parts (transfer budget, instrumentation, gathering the markers) are in +:mod:`struphy.models.tests.gpu_e2e`. +""" + +import logging + +import numpy as np +import pytest +from maybempi import MPI + +from struphy import ( + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + SavingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + grids, + maxwellians, + perturbations, +) +from struphy.linear_algebra.solver import SolverParameters +from struphy.models import LinearVlasovAmpereOneSpecies +from struphy.models.tests.gpu_e2e import RunResult, requires_cupy_backend, run, shared_out_folder + +logger = logging.getLogger("struphy") + +AMPLITUDE = 0.1 +R1 = 12.56 # box of length 4*pi, k = 0.5 +NUM_ELEMENTS = (8, 1, 1) +PPC = 40 # 320 markers +DT = 0.1 +N_STEPS = 3 +SOLVER_TOL = 1e-12 + + +def make_simulation(out_folders: str, sim_folder: str) -> Simulation: + """The model on a periodic 1D Cuboid, with the active cunumpy backend. + + Deterministic marker loading (``sobol_standard``) so that both backends start from the same + markers: ``pseudo_random`` draws from ``xp.random``, whose streams differ between NumPy and CuPy. + No B0 and no E0, so the time step is ``PushEta`` followed by ``EfieldWeightsCoupling``. No sorting + boxes (the verification test uses them): the model does not need them, and the box sorting + kernels have no CUDA version yet. + """ + model = LinearVlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False, with_E0=False) + + env = EnvironmentOptions(out_folders=out_folders, sim_folder=sim_folder) + time_opts = Time(dt=DT, Tend=round(N_STEPS * DT, 14)) # 3 * 0.1 > 0.3 would add a step + domain = domains.Cuboid(r1=R1) + grid = grids.TensorProductGrid(num_elements=NUM_ELEMENTS) + derham_opts = DerhamOptions(degree=(3, 1, 1)) + + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(ppc=PPC, loading="sobol_standard"), + weights_params=WeightsParameters(), + boundary_params=BoundaryParameters(), + # no sorting boxes: optional for this model, and box sorting has no device version yet + sorting_params=SortingParameters(), + saving_params=SavingParameters(), + bufsize=0.4, + ) + + model.propagators.push_eta.options = model.propagators.push_eta.Options() + # a tight solver tolerance, so that the backends differ by round-off and not by where pcg stops + model.propagators.coupling_Eweights.options = model.propagators.coupling_Eweights.Options( + solver_params=SolverParameters(tol=SOLVER_TOL), + ) + model.initial_poisson.options = model.initial_poisson.Options( + stab_mat="M0", + stab_eps=1e-6, + solver_params=SolverParameters(tol=SOLVER_TOL), + ) + + model.kinetic_ions.var.add_background(maxwellians.Maxwellian3D(n=(1.0, None))) + perturbation = perturbations.ModesCos(ls=(1,), amps=(AMPLITUDE,)) + model.kinetic_ions.var.add_initial_condition(maxwellians.Maxwellian3D(n=(1.0, perturbation))) + + return Simulation( + model=model, + env=env, + time_opts=time_opts, + domain=domain, + grid=grid, + derham_opts=derham_opts, + ) + + +def run_model(backend: str, out_folders: str) -> RunResult: + """Run the simulation on ``backend`` and copy the results to the host.""" + return run( + backend, + make_simulation, + out_folders, + n_steps=N_STEPS, + fields=lambda model: {"e_field": model.em_fields.e_field}, + particles=lambda model: model.kinetic_ions.var.particles, + ) + + +def check_reference(ref: RunResult): + """Physics sanity of the NumPy run: the comparison is only meaningful if the reference is right.""" + assert ref.markers.shape[0] == PPC * int(np.prod(NUM_ELEMENTS)) + if MPI.COMM_WORLD.Get_rank() == 0: + en_E, en_tot = ref.scalars["en_E"], ref.scalars["en_tot"] + assert en_tot.size == N_STEPS + 1 + # energy moves between field and particles ... + assert abs(en_E[-1] - en_E[0]) > 1e-3 * en_tot[0], en_E + # ... and the Crank-Nicolson step conserves the total + drift = np.max(np.abs(en_tot - en_tot[0])) / en_tot[0] + assert drift < 1e-6, f"{en_tot = }, {drift = }" + + +def compare_runs(ref: RunResult, gpu: RunResult): + """Compare the CuPy run against the NumPy run. + + Tolerances: + + * Velocities do not change (no E0, no B0) and positions are pushed explicitly + (``eta += dt * v / r1``, mod 1): the same floating-point operations on both backends, up to + fused multiply-adds on the GPU, so ``rtol = atol = 1e-13``. + * Weights, field coefficients and energies come out of ``EfieldWeightsCoupling``: a + Crank-Nicolson step whose Schur complement is solved by pcg to a relative residual of + ``SOLVER_TOL = 1e-12``, and whose matrix and vector are accumulated over markers (on the GPU + with atomic adds, i.e. in a different order). On this grid of 8 elements the preconditioned + system has a condition number of order 10-100, so the two solutions differ by at most + ~1e-10 relative to their norm after 3 steps; the tolerance ``1e-8`` (relative to the largest + entry) leaves a factor 100 against that and is still far below the change per step (~1e-2). + """ + rtol = 1e-8 + np.testing.assert_array_equal(gpu.markers[:, -1], ref.markers[:, -1]) # same markers + np.testing.assert_allclose(gpu.markers[:, 3:6], ref.markers[:, 3:6], rtol=1e-13, atol=1e-13) + np.testing.assert_allclose(gpu.markers[:, 0:3], ref.markers[:, 0:3], rtol=1e-13, atol=1e-13) + + w_ref, w_gpu = ref.markers[:, 6], gpu.markers[:, 6] + np.testing.assert_allclose(w_gpu, w_ref, rtol=0, atol=rtol * np.max(np.abs(w_ref))) + + e_ref, e_gpu = ref.fields["e_field"], gpu.fields["e_field"] + assert e_gpu.shape == e_ref.shape + np.testing.assert_allclose(e_gpu, e_ref, rtol=0, atol=rtol * np.max(np.abs(e_ref))) + + if MPI.COMM_WORLD.Get_rank() == 0: + for key, series in ref.scalars.items(): + np.testing.assert_allclose(gpu.scalars[key], series, rtol=rtol, err_msg=key) + + +@pytest.fixture +def out_folders(tmp_path_factory): + return shared_out_folder(tmp_path_factory, "gpu_e2e_LinearVlasovAmpereOneSpecies") + + +def test_numpy_reference(out_folders): + """The NumPy run that the GPU run is compared against; on NumPy nothing can be transferred.""" + ref = run_model("numpy", out_folders) + check_reference(ref) + ref.budget.check(n_scalars=ref.n_scalars) + assert ref.budget.integrate.total == ref.budget.diagnostics.total == ref.budget.output.total == 0 + + +@requires_cupy_backend +def test_cupy_matches_numpy_without_transfers(out_folders): + """The CuPy run agrees with the NumPy run, and its time loop stays on the device. + + On a GPU, or on the fake CuPy (``CUNUMPY_FAKE_CUPY=1``) with every CUDA launch emulated on the CPU. + """ + ref = run_model("numpy", out_folders) + check_reference(ref) + gpu = run_model("cupy", out_folders) + compare_runs(ref, gpu) + gpu.budget.check(n_scalars=gpu.n_scalars) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-ra"]) diff --git a/src/struphy/models/tests/test_gpu_VlasovAmpereOneSpecies.py b/src/struphy/models/tests/test_gpu_VlasovAmpereOneSpecies.py new file mode 100644 index 000000000..c647847e1 --- /dev/null +++ b/src/struphy/models/tests/test_gpu_VlasovAmpereOneSpecies.py @@ -0,0 +1,217 @@ +"""End-to-end GPU test of :class:`~struphy.models.VlasovAmpereOneSpecies` (issue #689). + +A small Landau-damping run (the setup of ``verification/test_verif_VlasovAmpereOneSpecies.py`` with fewer cells and +markers, full-f weights, and a larger perturbation so that energy moves within three steps) is done once on the NumPy +backend and once on the CuPy backend. The test checks that + +1. both runs give the same energies, electric-field coefficients and markers, within the tolerances explained in + :func:`compare_runs`, and +2. the time loop of the CuPy run does no host/device transfers through cunumpy, apart from an explicit budget for + the diagnostics and the output (:class:`~struphy.models.tests.gpu_e2e.TransferBudget`). + +A time step is ``PushEta`` followed by ``VlasovAmpereCoupling`` (no B0). The setup also runs the initial Poisson +solve (``charge_density_0form`` and the ``L2Projector`` with its ``MassMatrixPreconditioner``) on the device. + +The CuPy part needs a GPU, or cunumpy's fake CuPy (``CUNUMPY_FAKE_CUPY=1``), on which every CUDA launch is emulated on +the CPU. Without either, :func:`test_fake_cupy_matches_numpy` runs the comparison in a serial child process on the +fake CuPy. The tests are MPI-compatible (``mpirun -n 2 pytest --with-mpi``): both runs use the same domain +decomposition, and the markers are gathered and compared by marker ID. +""" + +import logging + +import numpy as np +import pytest +from maybempi import MPI + +from struphy import ( + BinningPlot, + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + SavingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + grids, + maxwellians, + perturbations, +) +from struphy.geometry.tests.test_domain import run_fake_cupy_child +from struphy.linear_algebra.solver import SolverParameters +from struphy.models import VlasovAmpereOneSpecies +from struphy.models.tests.gpu_e2e import ( + RunResult, + cupy_backend_available, + requires_cupy_backend, + run, + shared_out_folder, +) + +logger = logging.getLogger("struphy") + +AMPLITUDE = 0.1 +R1 = 12.56 # box of length 4*pi, k = 0.5 +NUM_ELEMENTS = (8, 1, 1) +PPC = 40 # 320 markers +DT = 0.1 +N_STEPS = 3 +SOLVER_TOL = 1e-12 + + +def make_simulation(out_folders: str, sim_folder: str) -> Simulation: + """The model on a periodic 1D Cuboid, with the active cunumpy backend. + + Deterministic marker loading (``sobol_standard``), so that both backends start from the same markers: + ``pseudo_random`` draws from ``xp.random``, whose streams differ between NumPy and CuPy. No sorting boxes, since + ``SortingBoxes`` does not work on the CuPy backend yet (#726); the model does not need them. Full-f weights: with + the control variate, ``kinetic_energy`` at t=0 is computed with the full-f weights and afterwards with the delta-f + weights, so ``total_energy`` cannot be checked. + """ + model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False) + + env = EnvironmentOptions(out_folders=out_folders, sim_folder=sim_folder, save_restart=False) + time_opts = Time(dt=DT, Tend=round(N_STEPS * DT, 14)) + + model.kinetic_ions.set_markers( + loading_params=LoadingParameters(ppc=PPC, loading="sobol_standard"), + weights_params=WeightsParameters(), + boundary_params=BoundaryParameters(), + sorting_params=SortingParameters(), + # the binned distribution function of the verification test, so that the diagnostics run on the device too + saving_params=SavingParameters( + binning_plots=(BinningPlot(slice="e1_v1", n_bins=(16, 16), ranges=((0.0, 1.0), (-5.0, 5.0))),), + ), + bufsize=0.4, + ) + + model.propagators.push_eta.options = model.propagators.push_eta.Options() + # tight solver tolerances, so that the backends differ by round-off and not by where pcg stops + model.propagators.coupling_va.options = model.propagators.coupling_va.Options( + solver_params=SolverParameters(tol=SOLVER_TOL), + ) + # stab_eps=1e-6, not the default 1e-14: the Monte-Carlo net charge puts (net charge) / stab_eps into the constant + # mode of phi, so with 1e-14 phi is ~1e14 and E = -grad(phi) loses ~14 digits to cancellation (5% of E here, on + # either backend) + model.initial_poisson.options = model.initial_poisson.Options( + stab_mat="M0", + stab_eps=1e-6, + solver_params=SolverParameters(tol=SOLVER_TOL), + ) + + model.kinetic_ions.var.add_background(maxwellians.Maxwellian3D(n=(1.0, None))) + perturbation = perturbations.ModesCos(ls=(1,), amps=(AMPLITUDE,)) + model.kinetic_ions.var.add_initial_condition(maxwellians.Maxwellian3D(n=(1.0, perturbation))) + + return Simulation( + model=model, + env=env, + time_opts=time_opts, + domain=domains.Cuboid(r1=R1), + grid=grids.TensorProductGrid(num_elements=NUM_ELEMENTS), + derham_opts=DerhamOptions(degree=(3, 1, 1)), + ) + + +def run_model(backend: str, out_folders: str) -> RunResult: + """Run the simulation on ``backend`` and copy the results to the host.""" + return run( + backend, + make_simulation, + out_folders, + n_steps=N_STEPS, + fields=lambda model: {"e_field": model.em_fields.e_field, "phi": model.em_fields.phi}, + particles=lambda model: model.kinetic_ions.var.particles, + ) + + +def check_reference(ref: RunResult): + """Physics sanity of the NumPy run: the comparison is only meaningful if the reference is right.""" + assert ref.markers.shape[0] == PPC * int(np.prod(NUM_ELEMENTS)) + if MPI.COMM_WORLD.Get_rank() == 0: + en_E = ref.scalars["electric_energy"] + en_tot = ref.scalars["total_energy"] + assert en_tot.size == N_STEPS + 1 + # energy moves between field and particles ... + assert abs(en_E[-1] - en_E[0]) > 1e-2 * abs(en_E[0]), en_E + # ... and the Crank-Nicolson coupling conserves the total from step to step, to the solver tolerance. The value + # at t=0 differs by ~2e-4 relative (the initial diagnostic, not the time stepping), so it is checked loosely. + drift = np.max(np.abs(en_tot[1:] - en_tot[1])) / abs(en_tot[1]) + assert drift < 1e-9, f"{en_tot = }, {drift = }" + assert abs(en_tot[1] - en_tot[0]) < 1e-3 * abs(en_tot[0]), en_tot + + +def compare_runs(ref: RunResult, gpu: RunResult): + """Compare the CuPy run against the NumPy run. + + Velocities, the electric field and the energies come out of ``VlasovAmpereCoupling``, a Crank-Nicolson step whose + Schur complement is solved by pcg to a relative residual of ``SOLVER_TOL = 1e-12``, with the matrix and vector + accumulated over markers (on the GPU with atomic adds, i.e. in another order). On this grid of 8 elements the + preconditioned system has a condition number of order 10-100, so the two solutions differ by ~1e-10 relative to + their norm after three steps. Positions are pushed explicitly with these velocities, so they inherit + ``DT * 1e-10``. The tolerance ``1e-8`` (relative to the largest entry) leaves a factor 100 and is still far below + the change per step (~5e-2 in the electric energy). Full-f weights do not change, so they must agree to round-off. + """ + rtol = 1e-8 + np.testing.assert_array_equal(gpu.markers[:, -1], ref.markers[:, -1]) # same markers + + def close(a, b, name): + np.testing.assert_allclose(a, b, rtol=0, atol=rtol * max(np.max(np.abs(b)), 1e-300), err_msg=name) + + close(gpu.markers[:, 0:3], ref.markers[:, 0:3], "positions") + close(gpu.markers[:, 3:6], ref.markers[:, 3:6], "velocities") + np.testing.assert_allclose(gpu.markers[:, 6], ref.markers[:, 6], rtol=1e-13, atol=0, err_msg="weights") + for name, coeffs in ref.fields.items(): + assert gpu.fields[name].shape == coeffs.shape, name + close(gpu.fields[name], coeffs, name) + + if MPI.COMM_WORLD.Get_rank() == 0: + for key, series in ref.scalars.items(): + np.testing.assert_allclose(gpu.scalars[key], series, rtol=rtol, err_msg=key) + + +def check_on_cupy(out_folders: str): + """The whole comparison, for the fake-CuPy child process of :func:`test_fake_cupy_matches_numpy`.""" + ref = run_model("numpy", out_folders) + check_reference(ref) + gpu = run_model("cupy", out_folders) + compare_runs(ref, gpu) + gpu.budget.check(n_scalars=gpu.n_scalars) + + +@pytest.fixture +def out_folders(tmp_path_factory): + return shared_out_folder(tmp_path_factory, "gpu_e2e_VlasovAmpereOneSpecies") + + +def test_numpy_reference(out_folders): + """The NumPy run that the GPU run is compared against; on NumPy nothing can be transferred.""" + ref = run_model("numpy", out_folders) + check_reference(ref) + ref.budget.check(n_scalars=ref.n_scalars) + assert ref.budget.integrate.total == ref.budget.diagnostics.total == ref.budget.output.total == 0 + + +@requires_cupy_backend +def test_cupy_matches_numpy_without_transfers(out_folders): + """The CuPy run agrees with the NumPy run, and its time loop stays on the device. + + On a GPU, or on the fake CuPy (``CUNUMPY_FAKE_CUPY=1``) with every CUDA launch emulated on the CPU; under + ``mpirun -n 2`` with two ranks. + """ + check_on_cupy(out_folders) + + +@pytest.mark.skipif(cupy_backend_available(), reason="runs in this process (test_cupy_matches_numpy_without_transfers)") +def test_fake_cupy_matches_numpy(out_folders): + """Without a GPU and without ``CUNUMPY_FAKE_CUPY``: the comparison in a serial child process on the fake CuPy.""" + run_fake_cupy_child( + f"from struphy.models.tests.test_gpu_VlasovAmpereOneSpecies import check_on_cupy; check_on_cupy({out_folders!r})" + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-ra"]) diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index ffd2e9fea..037cdc5af 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -1146,12 +1146,14 @@ def draw_markers( All needed parameters can be set in the parameter file, see :ref:`params_yml`. - An initial sorting will be performed if sort is given as True (default) and sorting_params were given to the init. + On several processes, the markers are always sent to the process of their domain after drawing + (:meth:`mpi_sort_markers`). They are also sorted into boxes if sort is True (default) and sorting + boxes were set up. Parameters ---------- sort : Bool - Wether to sort the particules in boxes after initial drawing (only if sorting params were passed) + Whether to sort the markers into boxes after initial drawing (only if sorting boxes were set up). """ import scipy.special as sp @@ -1258,10 +1260,13 @@ def draw_markers( # 2. plain sobol numbers with skip of first 1000 numbers elif self.loading == "sobol_standard": - self.phasespace_coords = sobol_seq.i4_sobol_generate( - 3 + self.vdim, - n_mks_load_loc, - 1000 + (n_mks_load_cum_sum - self.n_mks_load)[self._mpi_rank], + # generated on the host, then copied to the device (setup only) + self.phasespace_coords = xp.to_cunumpy( + sobol_seq.i4_sobol_generate( + 3 + self.vdim, + int(n_mks_load_loc), + 1000 + int((n_mks_load_cum_sum - self.n_mks_load)[self._mpi_rank]), + ) ) # 3. symmetric sobol numbers in all 6 dimensions with skip of first 1000 numbers @@ -1272,10 +1277,12 @@ def draw_markers( # each sobol point yields 64 symmetric markers; round up and truncate the last group n_sobol = -(-self.n_mks_load // 64) - temp_markers = sobol_seq.i4_sobol_generate( - 3 + self.vdim, - n_sobol[self._mpi_rank], - 1000 + (xp.cumsum(n_sobol) - n_sobol)[self._mpi_rank], + temp_markers = xp.to_cunumpy( + sobol_seq.i4_sobol_generate( + 3 + self.vdim, + int(n_sobol[self._mpi_rank]), + 1000 + int((xp.cumsum(n_sobol) - n_sobol)[self._mpi_rank]), + ) ) temp_symmetric = xp.zeros((64 * n_sobol[self._mpi_rank], 3 + self.vdim), dtype=float) @@ -1405,10 +1412,15 @@ def draw_markers( assert xp.all(~self.holes[:n_mks_load_loc]) assert xp.all(self.holes[n_mks_load_loc:]) - if self._initialized_sorting and sort: - logger.info("\nSorting the markers after initial draw") - if self.mpi_comm is not None: - self.mpi_sort_markers() + # send the markers to the process of their domain: needed on several processes whether or not there are + # sorting boxes, since accumulations before the first push (e.g. an initial Poisson solve) assume it (#764) + sort_boxes = self._initialized_sorting and sort + if self.mpi_comm is not None and (self.mpi_size > 1 or sort_boxes): + logger.info("\nSending the markers to their processes after initial draw") + self.mpi_sort_markers() + + if sort_boxes: + logger.info("\nSorting the markers into boxes after initial draw") self.do_sort() logger.info("Done.") @@ -2587,6 +2599,12 @@ def _allocate_marker_array(self, dry_run: bool = False): self._bc_type, ) + # on the CuPy backend with several ranks, find out now (collectively, in the setup) whether MPI takes device + # buffers; otherwise the probe runs on the first marker exchange, inside the first time step. + # The probe is collective: all ranks of mpi_comm must allocate their marker arrays together, never one alone. + if xp.get_backend() == "cupy" and self.mpi_size > 1: + self._mpi_cuda_aware() + def _initialize_sorting_boxes(self): """Initializes the sorting boxes. diff --git a/src/struphy/pic/sobol_seq.py b/src/struphy/pic/sobol_seq.py index 41a5e1f9c..15e92ff51 100644 --- a/src/struphy/pic/sobol_seq.py +++ b/src/struphy/pic/sobol_seq.py @@ -13,13 +13,16 @@ PYTHON versions by Corrado Chisari Original code is available from http://people.sc.fsu.edu/~jburkardt/py_src/sobol/sobol.html + + Runs on the host (NumPy) on every backend: it is a loop over scalars, and the caller copies the result + to the device. """ from __future__ import division import logging -import cunumpy as xp +import numpy as np logger = logging.getLogger("struphy") @@ -62,7 +65,7 @@ def i4_bit_hi1(n): Output, integer BIT, the number of bits base 2. """ - i = xp.floor(n) + i = np.floor(n) bit = 0 while i > 0: bit += 1 @@ -107,7 +110,7 @@ def i4_bit_lo0(n): Output, integer BIT, the position of the low 1 bit. """ bit = 1 - i = xp.floor(n) + i = np.floor(n) while i != 2 * (i // 2): bit += 1 i //= 2 @@ -125,7 +128,7 @@ def i4_sobol_generate(dim_num, n, skip=1): Output, real R(M,N), the points. """ - r = xp.full((n, dim_num), xp.nan) + r = np.full((n, dim_num), np.nan) for j in range(n): seed = j + skip r[j, 0:dim_num], next_seed = i4_sobol(dim_num, seed) @@ -226,8 +229,8 @@ def i4_sobol(dim_num, seed): seed_save = -1 # Initialize (part of) V. - v = xp.zeros((dim_max, log_max)) - v[0:40, 0] = xp.transpose( + v = np.zeros((dim_max, log_max)) + v[0:40, 0] = np.transpose( [ 1, 1, @@ -272,7 +275,7 @@ def i4_sobol(dim_num, seed): ], ) - v[2:40, 1] = xp.transpose( + v[2:40, 1] = np.transpose( [ 1, 3, @@ -315,7 +318,7 @@ def i4_sobol(dim_num, seed): ], ) - v[3:40, 2] = xp.transpose( + v[3:40, 2] = np.transpose( [ 7, 5, @@ -357,7 +360,7 @@ def i4_sobol(dim_num, seed): ], ) - v[5:40, 3] = xp.transpose( + v[5:40, 3] = np.transpose( [ 1, 7, @@ -397,7 +400,7 @@ def i4_sobol(dim_num, seed): ], ) - v[7:40, 4] = xp.transpose( + v[7:40, 4] = np.transpose( [ 9, 3, @@ -435,15 +438,15 @@ def i4_sobol(dim_num, seed): ], ) - v[13:40, 5] = xp.transpose( + v[13:40, 5] = np.transpose( [37, 33, 7, 5, 11, 39, 63, 27, 17, 15, 23, 29, 3, 21, 13, 31, 25, 9, 49, 33, 19, 29, 11, 19, 27, 15, 25], ) - v[19:40, 6] = xp.transpose( + v[19:40, 6] = np.transpose( [13, 33, 115, 41, 79, 17, 29, 119, 75, 73, 105, 7, 59, 65, 21, 3, 113, 61, 89, 45, 107], ) - v[37:40, 7] = xp.transpose([7, 23, 39]) + v[37:40, 7] = np.transpose([7, 23, 39]) # Set POLY. poly = [ @@ -522,7 +525,7 @@ def i4_sobol(dim_num, seed): # Expand this bit pattern to separate components of the logical array INCLUD. j = poly[i - 1] - includ = xp.zeros(m) + includ = np.zeros(m) for k in range(m, 0, -1): j2 = j // 2 includ[k - 1] = j != 2 * j2 @@ -536,7 +539,7 @@ def i4_sobol(dim_num, seed): for k in range(1, m + 1): l *= 2 if includ[k - 1]: - newv = xp.bitwise_xor(int(newv), int(l * v[i - 1, j - k - 1])) + newv = np.bitwise_xor(int(newv), int(l * v[i - 1, j - k - 1])) v[i - 1, j - 1] = newv # Multiply columns of V by appropriate power of 2. @@ -547,16 +550,16 @@ def i4_sobol(dim_num, seed): # RECIPD is 1/(common denominator of the elements in V). recipd = 1.0 / (2 * l) - lastq = xp.zeros(dim_num) + lastq = np.zeros(dim_num) - seed = int(xp.floor(seed)) + seed = int(np.floor(seed)) if seed < 0: seed = 0 l = 1 if seed == 0: - lastq = xp.zeros(dim_num) + lastq = np.zeros(dim_num) elif seed == seed_save + 1: # Find the position of the right-hand zero in SEED. @@ -564,12 +567,12 @@ def i4_sobol(dim_num, seed): elif seed <= seed_save: seed_save = 0 - lastq = xp.zeros(dim_num) + lastq = np.zeros(dim_num) for seed_temp in range(int(seed_save), int(seed)): l = i4_bit_lo0(seed_temp) for i in range(1, dim_num + 1): - lastq[i - 1] = xp.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) + lastq[i - 1] = np.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) l = i4_bit_lo0(seed) @@ -577,7 +580,7 @@ def i4_sobol(dim_num, seed): for seed_temp in range(int(seed_save + 1), int(seed)): l = i4_bit_lo0(seed_temp) for i in range(1, dim_num + 1): - lastq[i - 1] = xp.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) + lastq[i - 1] = np.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) l = i4_bit_lo0(seed) @@ -590,10 +593,10 @@ def i4_sobol(dim_num, seed): return # Calculate the new components of QUASI. - quasi = xp.zeros(dim_num) + quasi = np.zeros(dim_num) for i in range(1, dim_num + 1): quasi[i - 1] = lastq[i - 1] * recipd - lastq[i - 1] = xp.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) + lastq[i - 1] = np.bitwise_xor(int(lastq[i - 1]), int(v[i - 1, l - 1])) seed_save = seed seed += 1 @@ -643,11 +646,11 @@ def i4_uniform(a, b, seed): logger.info("I4_UNIFORM - Fatal error!") logger.info(" Input SEED = 0!") - seed = xp.floor(seed) + seed = np.floor(seed) a = round(a) b = round(b) - seed = xp.mod(seed, 2147483647) + seed = np.mod(seed, 2147483647) if seed < 0: seed += 2147483647 @@ -701,7 +704,7 @@ def prime_ge(n): Output, integer P, the smallest prime number that is greater than or equal to N. """ - p = max(xp.ceil(n), 2) + p = max(np.ceil(n), 2) while not is_prime(p): p += 1 @@ -725,7 +728,7 @@ def is_prime(n): return False # All primes >3 are of the form 6n+1 or 6n+5 (6n, 6n+2, 6n+4 are 2-divisible, 6n+3 is 3-divisible) p = 5 - root = int(xp.ceil(xp.sqrt(n))) + root = int(np.ceil(np.sqrt(n))) while p <= root: if n % p == 0 or n % (p + 2) == 0: return False diff --git a/src/struphy/pic/tests/test_draw_parallel.py b/src/struphy/pic/tests/test_draw_parallel.py index ff18a3321..09020c588 100644 --- a/src/struphy/pic/tests/test_draw_parallel.py +++ b/src/struphy/pic/tests/test_draw_parallel.py @@ -177,6 +177,56 @@ def test_marker_ids(loading, Np=1000): assert n_at_origin == 0 +@pytest.mark.parametrize("loading", ["pseudo_random", "sobol_standard", "sobol_antithetic"]) +@pytest.mark.parametrize("sort", [True, False]) +def test_draw_without_boxes_on_rank_domain(loading, sort, Np=1000): + """Without sorting boxes, every marker is on the domain of its process right after drawing (#764), + for ``sort=True`` and ``sort=False`` (which only concerns the boxes), and no marker is lost.""" + + import cunumpy as xp + from maybempi import MPI + + from struphy import LoadingParameters, domains + from struphy.feec.psydac_derham import Derham + from struphy.io.options import DerhamOptions + from struphy.pic.particles import Particles6D + from struphy.topology.grids import TensorProductGrid + + comm = MPI.COMM_WORLD + + grid = TensorProductGrid(num_elements=[8, 5, 3]) + derham = Derham(grid, DerhamOptions(degree=[1, 1, 1]), comm=comm) + domain_decomp = (derham.domain_array, derham.domain_decomposition.nprocs) + + loading_params = LoadingParameters( + Np=Np, + seed=1234, + loading=loading, + moments=(0.0, 0.0, 0.0, 1.0, 1.0, 1.0), + ) + + # no sorting_params: no sorting boxes + particles = Particles6D( + comm_world=comm, + domain_decomp=domain_decomp, + loading_params=loading_params, + domain=domains.Cuboid(), + ) + assert particles.sorting_boxes is None + + particles.draw_markers(sort=sort) + + positions = particles.markers_wo_holes[:, :3] + on_rank = xp.all(particles._is_on_domain_of_rank(positions, particles.mpi_rank), axis=1) + n_off_rank = int(xp.sum(~on_rank)) + n_markers = positions.shape[0] + if particles.mpi_size > 1: + n_markers = particles.mpi_comm.allreduce(n_markers) + + assert n_off_rank == 0, f"rank {particles.mpi_rank}: {n_off_rank} markers off the process domain after drawing" + assert n_markers == Np + + if __name__ == "__main__": # test_draw([8, 9, 10], [2, 3, 4], [False, False, True], ['Cuboid', { # 'l1': 1., 'r1': 2., 'l2': 10., 'r2': 20., 'l3': 100., 'r3': 200.}]) diff --git a/src/struphy/propagators/vlasov_ampere_coupling.py b/src/struphy/propagators/vlasov_ampere_coupling.py index 0c8c76414..77a129ae0 100644 --- a/src/struphy/propagators/vlasov_ampere_coupling.py +++ b/src/struphy/propagators/vlasov_ampere_coupling.py @@ -214,6 +214,8 @@ def __call__(self, dt): self._e_sum += self.variables.e.spline.vector self._e_sum += self._e_tmp self._e_sum *= 0.5 + # the pusher kernel reads the ghost regions directly (the solver output may leave them stale) + self._e_sum.update_ghost_regions() # Update velocities self._pusher(dt)