diff --git a/.github/workflows/test-PR-unit-mpi.yml b/.github/workflows/test-PR-unit-mpi.yml index 7bb53fbe2..6a5b1c387 100644 --- a/.github/workflows/test-PR-unit-mpi.yml +++ b/.github/workflows/test-PR-unit-mpi.yml @@ -82,3 +82,9 @@ jobs: source /struphy_fortran_/env_fortran_/bin/activate cd ${{ env.STRUPHY_PATH }} mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi pic/tests/test_sorting_device.py + + - name: Run polar GPU MPI paths on the fake CuPy + run: | + source /struphy_fortran_/env_fortran_/bin/activate + cd ${{ env.STRUPHY_PATH }} + mpirun --oversubscribe -n ${{ matrix.n-procs }} env CUNUMPY_FAKE_CUPY=1 pytest -v --with-mpi polar/tests/test_polar_cupy.py -k mpi_fake diff --git a/CUDA_STRATEGY.md b/CUDA_STRATEGY.md index 6880a52ef..e9711804f 100644 --- a/CUDA_STRATEGY.md +++ b/CUDA_STRATEGY.md @@ -25,7 +25,7 @@ Progress is tracked in struphy-hub/struphy#650. - [ ] **PR 16: One CPU and one GPU argument class per argument type** — every pyccel argument class (`MarkerArguments`, `DerhamArguments`, `DomainArguments`, `LocalProjectorsArguments`) has a CUDA counterpart of the same name with a `Cuda` prefix in `kernel_arguments/*_cuda.py`, with the same constructor and attributes. Each owner creates one of the two in `__init__`, depending on the backend; `PyccelStructArguments`, `__host_args__()` and host copies are no longer used. Every kernel that takes argument objects is called through `Kernel(PyccelKernel(...))`; without a CUDA version it raises on CuPy. The remaining entry kernels called from Python (spline, geometry and SPH evaluation, marker diagnostics, `reflect`, local projectors) move into one folder per kernel, like the pushers. The parity-test inputs move out of the kernel folders into `pic/tests/cuda_parity_cases.py`. Geometry evaluations have no CUDA version yet, so CuPy particle runs fail at setup until they are ported (see [PR 16 implementation notes](#pr-16-implementation-notes)). - [ ] **PR 17: struphy on the feectools CUDA stack** — the `feectools` submodule points at the top of the feectools CUDA stack (`cuda-4-device-kernels`, [feectools#88](https://github.com/struphy-hub/feectools/pull/88)) instead of `devel-tiny`, so the struphy CUDA PRs run against feectools with device stencil operations, MPI with device buffers and one GPU per rank. Moves along with the stack; before the struphy stack is merged into `devel`, the feectools stack is merged into `devel-tiny` and the submodule points there again (see [feectools](#feectools)). - [x] **PR 18: Geometry evaluation on the GPU for all analytic mappings** (#681, open) — CUDA versions of the four geometry entry kernels (`kernel_evaluate_pic`, `kernel_evaluate`, `kernel_pullpush_pic`, `kernel_pullpush` in `geometry/kernels/`), built on device versions of the whole metric chain (`f`, `df`, `det_df`, `df_inv`, `g`, `g_inv`, `select_metric_coeff`, `pull`/`push`/`tran`) for every analytic mapping (`kind_map` 10–12, 20–22, 30–32). Restores CuPy particle runs (weight initialization evaluates `jacobian_det`), and removes the Cuboid-only checks in `Pusher`, the accumulators and `reflect`. Parity arguments cover every analytic mapping (see [PR 18](#pr-18-geometry-evaluation-for-all-analytic-mappings) and the [PR 18 implementation notes](#pr-18-implementation-notes)). -- [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines and `EQDSKequilibrium` stay host-only (see [PR 19](#pr-19-spline-mappings) and the [PR 19 implementation notes](#pr-19-implementation-notes)). +- [x] **PR 19: Spline mappings on the GPU** ([#683](https://github.com/struphy-hub/struphy/pull/683)) — `kind_map` 0–2 (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC): `DomainArgs` gets array views for `t1..3`, `ind1..3` and `cx/cy/cz` (shapes needed on the device), spline-mapped `Domain`s can be created on the CuPy backend (control points fitted on the host), and `spline_3d`, `spline_2d_straight`, `spline_2d_torus` get device versions in the `kind_map` switch, so every mapping runs on CuPy and `check_mapping_on_device` is gone. Parity arguments and device-helper tests add four spline mappings. Polar splines were host-only at this stage; their current CuPy support is described in [Polar splines on CuPy](#polar-splines-on-cupy-695). See also the [PR 19 implementation notes](#pr-19-implementation-notes). - [x] **`linear_vlasov_ampere` on the GPU** (part of #688) — the first matrix accumulation: CUDA version of `linear_vlasov_ampere` (`EfieldWeightsCoupling`) on cunumpy's `Array6D` views (cunumpy 0.6.1), device versions of `fill_mat`, `fill_mat_vec`, `m_v_fill_b_v1_symm` and `outer`, and the matrix path of `Accumulator` checked on the CuPy backend (device data, ghost regions, transposed blocks, Schur operator) without a GPU through cunumpy's fake CuPy with emulated kernel launches (see the [implementation notes](#linear_vlasov_ampere-implementation-notes)). - [ ] **Next (order to be confirmed)**: one complete model (`VlasovAmpereOneSpecies`) on the GPU end to end, which needs the [feectools stack](#feectools) merged and released first; the guiding-center kernels (step 3), by hand (decided in #690, see [Hand-written CUDA](#hand-written-cuda-decision-for-690)). - [ ] **feectools**: the FEEC side (stencil vectors and matrices, MPI exchange, GPU binding) in feectools, stacked PRs [#85](https://github.com/struphy-hub/feectools/pull/85) (merged into `cuda-development`), [#86](https://github.com/struphy-hub/feectools/pull/86), [#87](https://github.com/struphy-hub/feectools/pull/87), [#88](https://github.com/struphy-hub/feectools/pull/88), integrated by [#90](https://github.com/struphy-hub/feectools/pull/90); needed before the end-to-end model run, not for PR 18/19 (see [feectools](#feectools)). @@ -146,7 +146,7 @@ from struphy.pic.pushing.kernels.push_eta_stage import push_eta_stage # a Kerne - The CUDA argument objects hold references. If an owner reallocates an array, it must rebuild its arguments at the same place, exactly like for the pyccel arguments. - `Domain` (PR 5): arrays that already have the dtype and layout the kernels expect (`float64`/`int64`, C-contiguous) are referenced, not copied; otherwise one device copy is made when the arguments are built (e.g. `degree`, a tuple). Arguments are recreated after deepcopy or unpickling. Spline mappings (e.g. `IGAPolarCylinder`) cannot be created on CuPy yet (`interp_mapping` passes CuPy arrays to `scipy.sparse`), so CUDA `args_domain` exists only for analytic mappings for now. - `Particles` (PR 6): particle arrays, validity masks and boundary-condition codes are allocated through `xp`. Scalar MPI gathers use small NumPy buffers (the own entry filled first, see `MockComm` above) and copy the result back to the active backend. -- `Derham` (PR 7): depends on feectools running on the CuPy backend (struphy-hub/feectools#85, merged). Data that describes the spline spaces (knots, quadrature and projection grids, `spline_types_pyccel`) is host data on every backend; only coefficients and stencil matrices live on the device. `domain_array`, `index_array(_N/_D)` and `neighbours` are gathered with NumPy MPI buffers and converted with `xp.asarray`. `args_derham` is built from the host knots, degrees and starts on NumPy, and from one device copy of them on CuPy. The pyccel scratch arrays (`bn1`, ..., `bd3`) are not part of the CUDA arguments; they become per-thread local arrays in CUDA (PR 10). Not supported on CuPy yet: local projectors (`NotImplementedError` at construction), polar splines (`NotImplementedError` at construction since PR 19, see the [PR 19 implementation notes](#pr-19-implementation-notes)), and field evaluation (`SplineFunction.__call__` still calls pyccel kernels with device coefficients). +- `Derham` (PR 7): depends on feectools running on the CuPy backend (struphy-hub/feectools#85, merged). Data that describes the spline spaces (knots, quadrature and projection grids, `spline_types_pyccel`) is host data on every backend; only coefficients and stencil matrices live on the device. `domain_array`, `index_array(_N/_D)` and `neighbours` are gathered with NumPy MPI buffers and converted with `xp.asarray`. `args_derham` is built from the host knots, degrees and starts on NumPy, and from one device copy of them on CuPy. The pyccel scratch arrays (`bn1`, ..., `bd3`) are not part of the CUDA arguments; they become per-thread local arrays in CUDA (PR 10). Polar splines run on CuPy since #695 (see [Polar splines on CuPy](#polar-splines-on-cupy-695)). Not supported on CuPy yet: local projectors (`NotImplementedError` at construction) and fixed local-grid evaluation (`SplineFunction.eval_tp_fixed_loc` still calls a pyccel helper with device coefficients). ### PR 8: Argument structs in a shared header (complete, open) @@ -208,7 +208,7 @@ For each kernel: add `_cuda.cu`; the parity test from PR 11 picks it up; ` ### PR 19: Spline mappings - `DomainArgs`: `t1..t3` become `Array1D` (`find_span` needs the number of knots), `ind1..3` `Array2D` and `cx/cy/cz` `Array3D` (the device needs their shapes); `CudaDomainArguments.fields` and the generated `pusher_args.cuh` change, the pyccel `DomainArguments` does not. -- Spline-mapped domains (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC) can be created on the CuPy backend: the control-point fitting (SciPy) runs on the host and the control points are copied to the device once. Polar splines (`Derham` with `polar_splines=True`) stay host-only and raise on CuPy. +- Spline-mapped domains (`IGAPolarCylinder`, `IGAPolarTorus`, `Tokamak`, GVEC, DESC) can be created on the CuPy backend: the control-point fitting (SciPy) runs on the host and the control points are copied to the device once. Polar splines (`Derham` with `polar_splines=True`) stayed host-only and raised on CuPy (on CuPy since #695). - Device versions of `spline_3d(_df)`, `spline_2d_straight(_df)`, `spline_2d_torus(_df)` in `geometry/spline_mappings_kernels.cuh`, using the B-spline helpers of PR 10; the parity arguments of the geometry kernels add the spline mappings. ## Moving to cunumpy (PR 13) @@ -348,10 +348,9 @@ The three spaces (H1vec/Hcurl/Hdiv) differ only in the basis; port one, then the `push_v_sph_pressure`, `push_v_sph_pressure_ideal_gas`, `push_v_viscosity`, `div_u_weak_1form`. Infrastructure that gates the steps, independent of the kernels: mappings are no gate any more (every CUDA kernel -accepts every mapping since PR 19), except for what is still host-only around them: polar splines in `Derham` cannot be -created on CuPy (see [Open questions](#open-questions)). Multi-rank marker sorting stays on the device since #712. Array -views with more than 4 dimensions (all matrix accumulations write 6D stencil matrix data) are in cunumpy since 0.6.1 -(`Array6D`); `linear_vlasov_ampere` is the first kernel that uses them. +accepts every mapping since PR 19); polar splines in `Derham` (#695) and MHD equilibria (#696) run on CuPy too. Multi-rank +marker sorting stays on the device since #712. Array views with more than 4 dimensions (all matrix accumulations write 6D +stencil matrix data) are in cunumpy since 0.6.1 (`Array6D`); `linear_vlasov_ampere` is the first kernel that uses them. ## Testing @@ -368,11 +367,9 @@ views with more than 4 dimensions (all matrix accumulations write 6D stencil mat - **Marker layout.** The markers array is row-major (`n_markers × n_cols`). With one thread per marker, the memory accesses are strided. This is fine for now (each thread reads a few neighbouring columns), but a column-major or struct-of-arrays layout may be faster later. This would affect the CPU code too, so it is out of scope here. The array view represents strides explicitly on the CUDA side. - **MPI + GPUs.** One GPU per MPI rank (`xp.bind_local_device()` before `MPI_Init`, with feectools#86/#87), and GPU-aware MPI for the marker exchange, so markers do not go through the host. The marker exchange in `Particles.mpi_sort_markers` uses device buffers since #698 (see [Marker exchange implementation notes](#marker-exchange-implementation-notes-698)); the SPH ghost-box exchange (`_sendrecv_markers_boxes`) does not yet. - **Single-source alternatives.** Decided (#690): no code generation; every CUDA kernel is written by hand, see [Hand-written CUDA](#hand-written-cuda-decision-for-690). -- **Polar splines on the GPU** (left after PR 19): spline mappings run on the device, but - `Derham` with `polar_splines=True` raises on CuPy (`PolarExtractionBlocksC1` builds SciPy sparse matrices, and the - polar extraction operators would apply them to device stencil data; needs `cupyx.scipy.sparse` or kernels). Needed - once a model with a polar domain runs on the GPU. (MHD equilibria run on CuPy since #696, see - [MHD equilibria](#mhd-equilibria-on-cupy-696).) +- **Polar splines and MHD equilibria on the GPU** (left after PR 19): both run on CuPy now, polar splines since #695 + (see [Polar splines on CuPy](#polar-splines-on-cupy-695)) and MHD equilibria since #696 (see + [MHD equilibria](#mhd-equilibria-on-cupy-696)). - **GVEC/DESC fields on the device.** `GVECequilibrium` and `DESCequilibrium` still evaluate B, J, p and n through gvec/DESC on the host, one copy per call (#696). Plan: evaluate them once at setup on a grid, interpolate onto 3d tensor-product splines (`make_interp_spline` per axis + `NdBSpline`) and evaluate those on the device, with native @@ -690,8 +687,8 @@ Every mapping runs on the GPU: the spline mappings (`kind_map` 0–2) join the a - **Still host-only.** `EQDSKequilibrium` cannot be created on CuPy (a `Tokamak` on CuPy takes an equilibrium created on the NumPy backend, or builds its default one there; fixed in #696). Polar splines: `PolarExtractionBlocksC1` builds SciPy sparse matrices from the control points, and the polar extraction operators apply them to the stencil data, which lives on - the device on CuPy; `Derham` raises `NotImplementedError` for `polar_splines=True` on CuPy (see - [Open questions](#open-questions)). + the device on CuPy; `Derham` raised `NotImplementedError` for `polar_splines=True` on CuPy (resolved in #695, see + [Polar splines on CuPy](#polar-splines-on-cupy-695)). - **Device helpers.** `b_splines_slim`, `b_der_splines_slim` (`bsplines/bsplines_kernels.cuh`), `evaluation_kernel_2d` (`bsplines/evaluation_kernels_2d.cuh`), `evaluation_kernel_3d` (`bsplines/evaluation_kernels_3d.cuh`) and the six spline mappings in `geometry/spline_mappings_kernels.cuh`, with the pyccel names and arguments (`const DomainArgs&` @@ -799,3 +796,45 @@ the [porting order](#porting-order); part of #688). `VlasovAmpereCoupling`, `MaxwellWeakAmpere` and the initial `PoissonSolve`) cannot be created on CuPy, CG inner products return host scalars, the FEEC operators (`curl`, `grad`, mass matrices, Schur solves) need the feectools CUDA stack, and the first H100 run (#687) is pending. + +## Polar splines on CuPy (#695) + +- **Blocks on the host.** `PolarExtractionBlocksC1` builds every block (extraction, DOF extraction, polar + derivatives) with NumPy/SciPy from a host copy of the control points (`xp.to_numpy(domain.cx)`), on every backend: + they are setup data, and `.toarray()`, `.T` and the restart left inverse (`SplineFunction._restart_extraction_op`) + keep working on them. +- **Products on the device.** `PolarExtractionOperator` and `PolarLinearOperator` keep their `blocks_*` as host + matrices; on CuPy, `dot` uses device copies (`DeviceSparseMatrix`, a `cupyx.scipy.sparse.csr_matrix` via + `xp.scipy.sparse`), made once per block list and dropped by the setters. The blocks act only on the first two or + three radial rings (at most `3 n2 x 3 n2` in `eta1-eta2`, `n3 x n3` in `eta3`), so they are small; blocks without + non-zeros are not copied. On NumPy nothing changes (the SciPy blocks are applied as before). +- **Refactored projectors and preconditioners (#746/#747).** The 2D x 1D Kronecker operators and + `PolarMassMatrixPreconditioner` keep coefficients and work arrays on the active backend. SciPy builds LU/ILU + factors on the host at setup. On CuPy, cached `cupyx.scipy.sparse.linalg.SuperLU` objects apply those factors + with device permutations and batched sparse triangular solves. Small 1d factors and the pole Schur complement + use cached dense inverses and device matrix products. The full poloidal inverse is never made dense. + Projectors support `schur`, `lu`, `ilu`, boundary conditions and transposed solves on both backends. +- **Mass setup.** Callable weights are evaluated with coordinates on the active backend. Sampled weights and + spline quadrature data are copied explicitly to the host for the 2d reduction/assembly and factorization. + These are setup transfers; repeated operator/preconditioner applications have no coefficient transfers on + a single rank or with CUDA-aware MPI. Krylov methods still synchronize convergence scalars. +- **MPI.** Kronecker `Alltoallv` redistributions, ring reductions and `PolarVector.toarray(allreduce=True)` + use `cunumpy.mpi.mpi_buffer`: synchronized device buffers with CUDA-aware MPI, host staging otherwise. + MPI capability must be established at startup through cunumpy. The pole sub-communicator and collective + ghost-region flags from #746 are preserved. feectools' own halo exchanges still require CUDA-aware MPI + on real GPUs; staging of the polar collectives does not replace that separate dependency. +- **Tests.** `polar/tests/test_polar_cupy.py` compares NumPy and CuPy for `IGAPolarCylinder` and `IGAPolarTorus`: + extraction/DOF operators, derivatives and transposes, vector arithmetic, all three projector solvers, + transposed and boundary-condition solves, mass matrices/preconditioners, callable and sampled weights, + and spline coefficient extraction/restart. Repeated Kronecker and preconditioner applications are checked + for transfers after warmup. The fake CuPy tests use CPU-emulated CUDA launches (including geometry/mass + kernels), with dense device-array substitutes for cuSPARSE products and triangular solves. + MPI parity tests decompose independently in eta1, eta2 and eta3 and exercise both device-buffer and + staged polar collectives; CI runs them on 1, 2, 3 and 4 ranks. Real-GPU variants are included but require + GPU hardware and the feectools CUDA stack. Full real-GPU MPI parity is opt-in with + `STRUPHY_TEST_CUDA_AWARE_MPI=1` on a known CUDA-aware MPI build; the separate staged-collective test + runs with regular MPI too, without involving feectools halos. +- **Separate follow-ups.** `SplineFunction.__call__` already dispatches CUDA kernels for marker and dense/sparse + grid evaluation. Remaining evaluation work is the host-only `eval_tp_fixed_loc` helper and distributed + field-evaluation MPI handling/coverage. FFT eta3 solves and sharing factorizations across different projector + objects are additional optimizations. diff --git a/doc/sections/subsections/numerics-geomFE.rst b/doc/sections/subsections/numerics-geomFE.rst index 7fb728c51..59a90211f 100644 --- a/doc/sections/subsections/numerics-geomFE.rst +++ b/doc/sections/subsections/numerics-geomFE.rst @@ -289,4 +289,12 @@ For weighted mass matrices, use :class:`~struphy.feec.preconditioner.PolarMassMa (``precond="PolarMassMatrixPreconditioner"`` in the propagator options); it is exact for axisymmetric weights and mappings. +On the CuPy backend, extraction, derivatives, Kronecker operators and these solves keep their +coefficients on the device. Factorization and mass-preconditioner assembly run once on the host; +cached device factors apply subsequent solves. Polar MPI collectives support CUDA-aware device +buffers or host staging through ``cunumpy.mpi.mpi_buffer``. The feectools halo exchanges require +CUDA-aware MPI on real GPUs. ``SplineFunction.__call__`` dispatches CUDA kernels for marker +and dense/sparse grid evaluation. The fixed local-grid helper ``eval_tp_fixed_loc`` remains +host-only; distributed field-evaluation MPI handling and coverage are separate follow-ups. + .. image:: ../../pics/polar_derham.png \ No newline at end of file diff --git a/src/struphy/feec/preconditioner.py b/src/struphy/feec/preconditioner.py index 87ca0bfd2..abccfb449 100644 --- a/src/struphy/feec/preconditioner.py +++ b/src/struphy/feec/preconditioner.py @@ -549,7 +549,7 @@ def _polar_mass_groups(mass_operator: WeightedMassOperator, reduction: str) -> l A_pol = (E @ M12 @ E.T).tocsr() # unweighted 1d mass matrix in eta_3 - A_3 = _mass_matrix_1d(mass_operator, bases[comps[0]][2], 2, _ones_1d)[0].toarray() + A_3 = xp.to_numpy(_mass_matrix_1d(mass_operator, bases[comps[0]][2], 2, _ones_1d)[0].toarray()) groups.append(Kron2Dx1DGroup(comps, A_pol, A_3)) return groups diff --git a/src/struphy/feec/projectors.py b/src/struphy/feec/projectors.py index 06e780733..1f87674ac 100644 --- a/src/struphy/feec/projectors.py +++ b/src/struphy/feec/projectors.py @@ -615,7 +615,7 @@ def solve(self, rhs, transposed=False, apply_bc=False, out=None, x0=None): krylov = self._polar_krylov[key] # relative tolerance w.r.t. the norm of the right-hand side - rhs_norm = xp.sqrt(rhs.inner(rhs)) + rhs_norm = float(xp.sqrt(rhs.inner(rhs))) krylov.set_options(x0=x0, tol=self._polar_rtol * rhs_norm if rhs_norm > 0.0 else self._polar_rtol) x = krylov.solve(rhs, out=out) diff --git a/src/struphy/feec/psydac_derham.py b/src/struphy/feec/psydac_derham.py index 2eec857e2..a870d58b0 100644 --- a/src/struphy/feec/psydac_derham.py +++ b/src/struphy/feec/psydac_derham.py @@ -619,12 +619,6 @@ def __init__( raise NotImplementedError( "Local projectors (DerhamOptions.local_projectors=True) are not supported on the CuPy backend yet." ) - if polar_splines and (xp.get_backend() == "cupy"): - # PolarExtractionBlocksC1 builds SciPy sparse matrices from the control points, and the polar - # extraction operators apply them to the stencil data, which lives on the device on CuPy. - raise NotImplementedError( - "Polar splines (DerhamOptions.polar_splines=True) are not supported on the CuPy backend yet." - ) # number of elements and spline degrees in each direction assert len(num_elements) == 3 @@ -2782,16 +2776,17 @@ def _restart_extraction_op(self): n_comps = W.n_comps # stack blocks of E (polar coeffs x polar rings); incompatible blocks (None) are zero - rows = xp.cumsum([0] + list(W.n_polar)) - cols = xp.cumsum([0] + [n_r * n_2 for n_r, n_2 in zip(W.n_rings, W.n2)]) - E_full = xp.zeros((rows[-1], cols[-1]), dtype=float) + # (host setup data on every backend, like the blocks of E) + rows = np.cumsum([0] + list(W.n_polar)) + cols = np.cumsum([0] + [n_r * n_2 for n_r, n_2 in zip(W.n_rings, W.n2)]) + E_full = np.zeros((rows[-1], cols[-1]), dtype=float) for m in range(n_comps): for n in range(n_comps): if E.blocks_ten_to_pol[m][n] is not None: E_full[rows[m] : rows[m + 1], cols[n] : cols[n + 1]] = E.blocks_ten_to_pol[m][n].toarray() # E^T has full column rank, hence pinv(E^T) @ E^T = identity on polar coeffs - L_full = xp.linalg.pinv(E_full.T) + L_full = np.linalg.pinv(E_full.T) blocks = [ [ diff --git a/src/struphy/feec/tests/test_derham_gpu.py b/src/struphy/feec/tests/test_derham_gpu.py index 64c55a0d8..d8f3df46a 100644 --- a/src/struphy/feec/tests/test_derham_gpu.py +++ b/src/struphy/feec/tests/test_derham_gpu.py @@ -67,20 +67,3 @@ def test_local_projectors_not_supported_on_cupy(): """Local projectors have no device implementation yet; the Derham fails when it is created.""" with cunumpy.use_backend("cupy"), pytest.raises(NotImplementedError, match="Local projectors"): make_derham(local_projectors=True) - - -@requires_cupy -def test_polar_splines_not_supported_on_cupy(): - """Polar splines have no device implementation yet; the Derham fails when it is created.""" - from maybempi import MPI - - from struphy import domains - from struphy.feec.psydac_derham import Derham - from struphy.io.options import DerhamOptions - from struphy.topology.grids import TensorProductGrid - - with cunumpy.use_backend("cupy"): - domain = domains.IGAPolarCylinder(num_elements=(8, 6), degree=(2, 3)) - options = DerhamOptions(degree=(2, 3, 1), bcs=(("dirichlet", "free"), None, None), polar_splines=True) - with pytest.raises(NotImplementedError, match="Polar splines"): - Derham(TensorProductGrid(num_elements=(8, 6, 4)), options, comm=MPI.COMM_WORLD, domain=domain) diff --git a/src/struphy/polar/basic.py b/src/struphy/polar/basic.py index 645f6a45d..b6fb1fb6a 100644 --- a/src/struphy/polar/basic.py +++ b/src/struphy/polar/basic.py @@ -1,6 +1,7 @@ from typing import TYPE_CHECKING import cunumpy as xp +from cunumpy.mpi import mpi_buffer from feectools.linalg.basic import Vector, VectorSpace from feectools.linalg.block import BlockVector from feectools.linalg.stencil import StencilVector @@ -346,8 +347,9 @@ def toarray(self, allreduce=False): out = self.tp.toarray()[self.space.n_rings[0] * self.space.n2[0] * self.space.n3[0] :] # allreduce tensor-product part - if self.space.comm is not None and allreduce: - self.space.comm.Allreduce(MPI.IN_PLACE, out, op=MPI.SUM) + if self.space.comm is not None and self.space.comm.Get_size() > 1 and allreduce: + with mpi_buffer(out, recv=True) as buf: + self.space.comm.Allreduce(MPI.IN_PLACE, buf, op=MPI.SUM) out = xp.concatenate((self.pol[0].flatten(), out)) @@ -357,10 +359,10 @@ def toarray(self, allreduce=False): out3 = self.tp[2].toarray()[self.space.n_rings[2] * self.space.n2[2] * self.space.n3[2] :] # allreduce tensor-product part - if self.space.comm is not None and allreduce: - self.space.comm.Allreduce(MPI.IN_PLACE, out1, op=MPI.SUM) - self.space.comm.Allreduce(MPI.IN_PLACE, out2, op=MPI.SUM) - self.space.comm.Allreduce(MPI.IN_PLACE, out3, op=MPI.SUM) + if self.space.comm is not None and self.space.comm.Get_size() > 1 and allreduce: + for array in (out1, out2, out3): + with mpi_buffer(array, recv=True) as buf: + self.space.comm.Allreduce(MPI.IN_PLACE, buf, op=MPI.SUM) out = xp.concatenate( ( diff --git a/src/struphy/polar/extraction_operators.py b/src/struphy/polar/extraction_operators.py index 0de114d15..a6caa9e14 100644 --- a/src/struphy/polar/extraction_operators.py +++ b/src/struphy/polar/extraction_operators.py @@ -1,4 +1,5 @@ import cunumpy as xp +import numpy as np from struphy.eigenvalue_solvers.derivatives import grad_1d_matrix @@ -26,29 +27,33 @@ class PolarExtractionBlocksC1: A linear combination of n_rings rings is used to construct n_polar polar basis functions/DOFs and "first tp ring" basis functions/DOFs. + The blocks are setup data: they are built on the host (NumPy, SciPy sparse) on every array backend, + from a host copy of the control points. On the CuPy backend the polar operators + (:mod:`struphy.polar.linear_operators`) copy them to the device once and apply them there. + Parameters ---------- - cx : array-like - Control points defining the x-component of a 2D B-spline mapping. + domain : struphy.geometry.base.Domain + Spline mapping (IGAPolarCylinder, IGAPolarTorus, ...) whose control points define the polar splines. - cy : array-like - Control points defining the y-component of a 2D B-spline mapping. + derham : struphy.feec.psydac_derham.Derham + Discrete Derham complex. """ def __init__(self, domain, derham): from scipy.sparse import csr_matrix as csr - # get control points - cx = domain.cx[:, :, 0] - cy = domain.cy[:, :, 0] + # get control points (host copies; they live on the device on the CuPy backend) + cx = xp.to_numpy(domain.cx[:, :, 0]) + cy = xp.to_numpy(domain.cy[:, :, 0]) self._cx = cx self._cy = cy self._pole = (cx[0, 0], cy[0, 0]) - assert xp.all(cx[0] == self.pole[0]) - assert xp.all(cy[0] == self.pole[1]) + assert np.all(cx[0] == self.pole[0]) + assert np.all(cy[0] == self.pole[1]) self._n0 = cx.shape[0] self._n1 = cx.shape[1] @@ -70,14 +75,14 @@ def __init__(self, domain, derham): self._tau = max( [ ((self.cx[1] - self.pole[0]) * (-2)).max(), - ((self.cx[1] - self.pole[0]) - xp.sqrt(3) * (self.cy[1] - self.pole[1])).max(), - ((self.cx[1] - self.pole[0]) + xp.sqrt(3) * (self.cy[1] - self.pole[1])).max(), + ((self.cx[1] - self.pole[0]) - np.sqrt(3) * (self.cy[1] - self.pole[1])).max(), + ((self.cx[1] - self.pole[0]) + np.sqrt(3) * (self.cy[1] - self.pole[1])).max(), ], ) # barycentric coordinates - self._xi_0 = xp.zeros((3, self.n1), dtype=float) - self._xi_1 = xp.zeros((3, self.n1), dtype=float) + self._xi_0 = np.zeros((3, self.n1), dtype=float) + self._xi_1 = np.zeros((3, self.n1), dtype=float) self._xi_0[:, :] = 1 / 3 @@ -85,12 +90,12 @@ def __init__(self, domain, derham): self._xi_1[1, :] = ( 1 / 3 - 1 / (3 * self.tau) * (self.cx[1] - self.pole[0]) - + xp.sqrt(3) / (3 * self.tau) * (self.cy[1] - self.pole[1]) + + np.sqrt(3) / (3 * self.tau) * (self.cy[1] - self.pole[1]) ) self._xi_1[2, :] = ( 1 / 3 - 1 / (3 * self.tau) * (self.cx[1] - self.pole[0]) - - xp.sqrt(3) / (3 * self.tau) * (self.cy[1] - self.pole[1]) + - np.sqrt(3) / (3 * self.tau) * (self.cy[1] - self.pole[1]) ) # remove small values @@ -102,17 +107,17 @@ def __init__(self, domain, derham): # ============= basis extraction operator for discrete 0-forms ================ # first n_rings tp rings --> "polar coeffs" - e0_blocks_ten_to_pol = xp.block([self.xi_0, self.xi_1]) + e0_blocks_ten_to_pol = np.block([self.xi_0, self.xi_1]) self._e_ten_to_pol["0"] = [[csr(e0_blocks_ten_to_pol)]] # ============ basis extraction operator for discrete 1-forms (Hcurl) ========= # first n_rings tp rings --> "polar coeffs" - e1_11_blocks_ten_to_pol = xp.zeros((self.n_polar[1][0], self.n_rings[1][0] * self.n1), dtype=float) - e1_12_blocks_ten_to_pol = xp.zeros((self.n_polar[1][0], self.n_rings[1][1] * self.d1), dtype=float) + e1_11_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][0] * self.n1), dtype=float) + e1_12_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][1] * self.d1), dtype=float) - e1_21_blocks_ten_to_pol = xp.zeros((self.n_polar[1][1], self.n_rings[1][0] * self.n1), dtype=float) - e1_22_blocks_ten_to_pol = xp.zeros((self.n_polar[1][1], self.n_rings[1][1] * self.d1), dtype=float) + e1_21_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][0] * self.n1), dtype=float) + e1_22_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][1] * self.d1), dtype=float) # 1st component for l in range(2): @@ -135,7 +140,7 @@ def __init__(self, domain, derham): # =============== basis extraction operator for discrete 1-forms (Hdiv) ========= # first n_rings tp rings --> "polar coeffs" - e3_blocks_ten_to_pol = xp.zeros((self.n_polar[3][0], self.n_rings[3][0] * self.d1), dtype=float) + e3_blocks_ten_to_pol = np.zeros((self.n_polar[3][0], self.n_rings[3][0] * self.d1), dtype=float) self._e_ten_to_pol["2"] = [ [csr(e1_22_blocks_ten_to_pol), csr(-e1_21_blocks_ten_to_pol), None], @@ -161,7 +166,7 @@ def __init__(self, domain, derham): self._p_ten_to_ten = {} # first n_rings tp rings --> "polar coeffs" - p0_blocks_ten_to_pol = xp.zeros((self.n_polar[0][0], self.n_rings[0][0] * self.n1), dtype=float) + p0_blocks_ten_to_pol = np.zeros((self.n_polar[0][0], self.n_rings[0][0] * self.n1), dtype=float) # !! NOTE: for odd spline degrees and periodic splines the first Greville point sometimes does NOT start at zero!! if domain.degree[1] % 2 != 0 and not (abs(derham.V0fem.spaces[1].interpolation_grid[0]) < 1e-14): @@ -176,15 +181,15 @@ def __init__(self, domain, derham): self._p_ten_to_pol["0"] = [[csr(p0_blocks_ten_to_pol)]] # first n_rings + 1 tp rings --> "first tp ring" - p0_blocks_ten_to_ten = xp.block([0 * xp.identity(self.n1)] * self.n_rings[0][0] + [xp.identity(self.n1)]) + p0_blocks_ten_to_ten = np.block([0 * np.identity(self.n1)] * self.n_rings[0][0] + [np.identity(self.n1)]) self._p_ten_to_ten["0"] = [[csr(p0_blocks_ten_to_ten)]] # =========== projection extraction operator for discrete 1-forms (Hcurl) ======== # first n_rings tp rings --> "polar coeffs" - p1_11_blocks_ten_to_pol = xp.zeros((self.n_polar[1][0], self.n_rings[1][0] * self.n1), dtype=float) - p1_22_blocks_ten_to_pol = xp.zeros((self.n_polar[1][1], self.n_rings[1][1] * self.d1), dtype=float) + p1_11_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][0] * self.n1), dtype=float) + p1_22_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][1] * self.d1), dtype=float) # !! NOTE: PSYDAC's first integration interval sometimes start at < 0 !! if derham.V3fem.spaces[1].histopolation_grid[0] < -1e-14: @@ -196,8 +201,8 @@ def __init__(self, domain, derham): p1_22_blocks_ten_to_pol[1, (self.d1 + 0 * self.d1 // 3) : (self.d1 + 1 * self.d1 // 3)] = 1.0 p1_22_blocks_ten_to_pol[1, (self.d1 + 1 * self.d1 // 3) : (self.d1 + 2 * self.d1 // 3)] = 1.0 - p1_12_blocks_ten_to_pol = xp.zeros((self.n_polar[1][0], self.n_rings[1][1] * self.d1), dtype=float) - p1_21_blocks_ten_to_pol = xp.zeros((self.n_polar[1][1], self.n_rings[1][0] * self.d1), dtype=float) + p1_12_blocks_ten_to_pol = np.zeros((self.n_polar[1][0], self.n_rings[1][1] * self.d1), dtype=float) + p1_21_blocks_ten_to_pol = np.zeros((self.n_polar[1][1], self.n_rings[1][0] * self.d1), dtype=float) self._p_ten_to_pol["1"] = [ [csr(p1_11_blocks_ten_to_pol), csr(p1_12_blocks_ten_to_pol), None], @@ -206,26 +211,26 @@ def __init__(self, domain, derham): ] # first n_rings + 1 tp rings --> "first tp ring" - p1_11_blocks_ten_to_ten = xp.zeros((self.n1, self.n1), dtype=float) + p1_11_blocks_ten_to_ten = np.zeros((self.n1, self.n1), dtype=float) # !! NOTE: for odd spline degrees and periodic splines the first Greville point sometimes does NOT start at zero!! if domain.degree[1] % 2 != 0 and not (abs(derham.V0fem.spaces[1].interpolation_grid[0]) < 1e-14): - p1_11_blocks_ten_to_ten[:, 3 * self.n1 // 3 - 1] = -xp.roll(self.xi_1[0], -1) - p1_11_blocks_ten_to_ten[:, 1 * self.n1 // 3 - 1] = -xp.roll(self.xi_1[1], -1) - p1_11_blocks_ten_to_ten[:, 2 * self.n1 // 3 - 1] = -xp.roll(self.xi_1[2], -1) + p1_11_blocks_ten_to_ten[:, 3 * self.n1 // 3 - 1] = -np.roll(self.xi_1[0], -1) + p1_11_blocks_ten_to_ten[:, 1 * self.n1 // 3 - 1] = -np.roll(self.xi_1[1], -1) + p1_11_blocks_ten_to_ten[:, 2 * self.n1 // 3 - 1] = -np.roll(self.xi_1[2], -1) else: p1_11_blocks_ten_to_ten[:, 0 * self.n1 // 3] = -self.xi_1[0] p1_11_blocks_ten_to_ten[:, 1 * self.n1 // 3] = -self.xi_1[1] p1_11_blocks_ten_to_ten[:, 2 * self.n1 // 3] = -self.xi_1[2] - p1_11_blocks_ten_to_ten += xp.identity(self.n1) + p1_11_blocks_ten_to_ten += np.identity(self.n1) - p1_11_blocks_ten_to_ten = xp.block([p1_11_blocks_ten_to_ten, xp.identity(self.n1)]) + p1_11_blocks_ten_to_ten = np.block([p1_11_blocks_ten_to_ten, np.identity(self.n1)]) - p1_22_blocks_ten_to_ten = xp.block([0 * xp.identity(self.d1)] * self.n_rings[1][1] + [xp.identity(self.d1)]) + p1_22_blocks_ten_to_ten = np.block([0 * np.identity(self.d1)] * self.n_rings[1][1] + [np.identity(self.d1)]) - p1_12_blocks_ten_to_ten = xp.zeros((self.d1, (self.n_rings[1][1] + 1) * self.d1), dtype=float) - p1_21_blocks_ten_to_ten = xp.zeros((self.n1, (self.n_rings[1][0] + 1) * self.n1), dtype=float) + p1_12_blocks_ten_to_ten = np.zeros((self.d1, (self.n_rings[1][1] + 1) * self.d1), dtype=float) + p1_21_blocks_ten_to_ten = np.zeros((self.n1, (self.n_rings[1][0] + 1) * self.n1), dtype=float) self._p_ten_to_ten["1"] = [ [csr(p1_11_blocks_ten_to_ten), csr(p1_12_blocks_ten_to_ten), None], @@ -236,7 +241,7 @@ def __init__(self, domain, derham): # ========== projection extraction operator for discrete 1-forms (Hdiv) ========== # first n_rings tp rings --> "polar coeffs" - p3_blocks_ten_to_pol = xp.zeros((self.n_polar[3][0], self.n_rings[3][0] * self.d1), dtype=float) + p3_blocks_ten_to_pol = np.zeros((self.n_polar[3][0], self.n_rings[3][0] * self.d1), dtype=float) self._p_ten_to_pol["2"] = [ [csr(p1_22_blocks_ten_to_pol), csr(p1_21_blocks_ten_to_pol), None], @@ -245,24 +250,24 @@ def __init__(self, domain, derham): ] # first n_rings + 1 tp rings --> "first tp ring" - p3_blocks_ten_to_ten = xp.zeros((self.d1, self.d1), dtype=float) + p3_blocks_ten_to_ten = np.zeros((self.d1, self.d1), dtype=float) - a0 = xp.diff(self.xi_1[1], append=self.xi_1[1, 0]) - a1 = xp.diff(self.xi_1[2], append=self.xi_1[2, 0]) + a0 = np.diff(self.xi_1[1], append=self.xi_1[1, 0]) + a1 = np.diff(self.xi_1[2], append=self.xi_1[2, 0]) # !! NOTE: PSYDAC's first integration interval sometimes start at < 0 !! if derham.V3fem.spaces[1].histopolation_grid[0] < -1e-14: p3_blocks_ten_to_ten[:, (0 * self.n1 // 3 + 1) : (1 * self.n1 // 3 + 1)] = ( - -xp.roll(a0, +1)[:, None] - xp.roll(a1, +1)[:, None] + -np.roll(a0, +1)[:, None] - np.roll(a1, +1)[:, None] ) - p3_blocks_ten_to_ten[:, (1 * self.n1 // 3 + 1) : (2 * self.n1 // 3 + 1)] = -xp.roll(a1, +1)[:, None] + p3_blocks_ten_to_ten[:, (1 * self.n1 // 3 + 1) : (2 * self.n1 // 3 + 1)] = -np.roll(a1, +1)[:, None] else: p3_blocks_ten_to_ten[:, 0 * self.n1 // 3 : 1 * self.n1 // 3] = -a0[:, None] - a1[:, None] p3_blocks_ten_to_ten[:, 1 * self.n1 // 3 : 2 * self.n1 // 3] = -a1[:, None] - p3_blocks_ten_to_ten += xp.identity(self.d1) + p3_blocks_ten_to_ten += np.identity(self.d1) - p3_blocks_ten_to_ten = xp.block([p3_blocks_ten_to_ten, xp.identity(self.d1)]) + p3_blocks_ten_to_ten = np.block([p3_blocks_ten_to_ten, np.identity(self.d1)]) self._p_ten_to_ten["2"] = [ [csr(p1_22_blocks_ten_to_ten), csr(p1_21_blocks_ten_to_ten), None], @@ -295,39 +300,39 @@ def __init__(self, domain, derham): # ======================= discrete gradient ====================================== # "polar coeffs" to "polar coeffs" - grad_pol_to_pol_1 = xp.zeros((self.n_polar[1][0], self.n_polar[0][0]), dtype=float) - grad_pol_to_pol_2 = xp.array([[-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]]) - grad_pol_to_pol_3 = xp.identity(self.n_polar[0][0], dtype=float) + grad_pol_to_pol_1 = np.zeros((self.n_polar[1][0], self.n_polar[0][0]), dtype=float) + grad_pol_to_pol_2 = np.array([[-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]]) + grad_pol_to_pol_3 = np.identity(self.n_polar[0][0], dtype=float) self._grad_pol_to_pol = [[csr(grad_pol_to_pol_1)], [csr(grad_pol_to_pol_2)], [csr(grad_pol_to_pol_3)]] # "polar coeffs" to "first tp ring" - grad_pol_to_ten_1 = xp.zeros(((self.n_rings[1][0] + 1) * self.n1, self.n_polar[0][0])) - grad_pol_to_ten_2 = xp.zeros(((self.n_rings[1][1] + 1) * self.d1, self.n_polar[0][0])) - grad_pol_to_ten_3 = xp.zeros(((self.n_rings[0][0] + 1) * self.n1, self.n_polar[0][0])) + grad_pol_to_ten_1 = np.zeros(((self.n_rings[1][0] + 1) * self.n1, self.n_polar[0][0])) + grad_pol_to_ten_2 = np.zeros(((self.n_rings[1][1] + 1) * self.d1, self.n_polar[0][0])) + grad_pol_to_ten_3 = np.zeros(((self.n_rings[0][0] + 1) * self.n1, self.n_polar[0][0])) grad_pol_to_ten_1[-self.n1 :, :] = -self.xi_1.T self._grad_pol_to_ten = [[csr(grad_pol_to_ten_1)], [csr(grad_pol_to_ten_2)], [csr(grad_pol_to_ten_3)]] # eta_3 direction - grad_e3_1 = xp.identity(self.n2, dtype=float) - grad_e3_2 = xp.identity(self.n2, dtype=float) - grad_e3_3 = grad_1d_matrix(derham.spl_kind[2], self.n2) + grad_e3_1 = np.identity(self.n2, dtype=float) + grad_e3_2 = np.identity(self.n2, dtype=float) + grad_e3_3 = xp.to_numpy(grad_1d_matrix(derham.spl_kind[2], self.n2)) self._grad_e3 = [[csr(grad_e3_1)], [csr(grad_e3_2)], [csr(grad_e3_3)]] # =========================== discrete curl ====================================== # "polar coeffs" to "polar coeffs" - curl_pol_to_pol_12 = xp.identity(self.n_polar[1][1], dtype=float) - curl_pol_to_pol_13 = xp.array([[-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]]) + curl_pol_to_pol_12 = np.identity(self.n_polar[1][1], dtype=float) + curl_pol_to_pol_13 = np.array([[-1.0, 1.0, 0.0], [-1.0, 0.0, 1.0]]) - curl_pol_to_pol_21 = xp.identity(self.n_polar[1][0], dtype=float) - curl_pol_to_pol_23 = xp.zeros((self.n_polar[2][1], self.n_polar[0][0]), dtype=float) + curl_pol_to_pol_21 = np.identity(self.n_polar[1][0], dtype=float) + curl_pol_to_pol_23 = np.zeros((self.n_polar[2][1], self.n_polar[0][0]), dtype=float) - curl_pol_to_pol_31 = xp.zeros((self.n_polar[3][0], self.n_polar[1][0]), dtype=float) - curl_pol_to_pol_32 = xp.zeros((self.n_polar[3][0], self.n_polar[1][1]), dtype=float) + curl_pol_to_pol_31 = np.zeros((self.n_polar[3][0], self.n_polar[1][0]), dtype=float) + curl_pol_to_pol_32 = np.zeros((self.n_polar[3][0], self.n_polar[1][1]), dtype=float) self._curl_pol_to_pol = [ [None, csr(-curl_pol_to_pol_12), csr(curl_pol_to_pol_13)], @@ -336,14 +341,14 @@ def __init__(self, domain, derham): ] # "polar coeffs" to "first tp ring" - curl_pol_to_ten_12 = xp.zeros(((self.n_rings[2][0] + 1) * self.d1, self.n_polar[1][1])) - curl_pol_to_ten_13 = xp.zeros(((self.n_rings[2][0] + 1) * self.d1, self.n_polar[0][0])) + curl_pol_to_ten_12 = np.zeros(((self.n_rings[2][0] + 1) * self.d1, self.n_polar[1][1])) + curl_pol_to_ten_13 = np.zeros(((self.n_rings[2][0] + 1) * self.d1, self.n_polar[0][0])) - curl_pol_to_ten_21 = xp.zeros(((self.n_rings[2][1] + 1) * self.n1, self.n_polar[1][0])) - curl_pol_to_ten_23 = xp.zeros(((self.n_rings[2][1] + 1) * self.n1, self.n_polar[0][0])) + curl_pol_to_ten_21 = np.zeros(((self.n_rings[2][1] + 1) * self.n1, self.n_polar[1][0])) + curl_pol_to_ten_23 = np.zeros(((self.n_rings[2][1] + 1) * self.n1, self.n_polar[0][0])) - curl_pol_to_ten_31 = xp.zeros(((self.n_rings[3][0] + 1) * self.n1, self.n_polar[1][0])) - curl_pol_to_ten_32 = xp.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[1][1])) + curl_pol_to_ten_31 = np.zeros(((self.n_rings[3][0] + 1) * self.n1, self.n_polar[1][0])) + curl_pol_to_ten_32 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[1][1])) curl_pol_to_ten_23[-self.n1 :, :] = -self.xi_1.T @@ -360,14 +365,14 @@ def __init__(self, domain, derham): ] # eta_3 direction - curl_e3_12 = grad_1d_matrix(derham.spl_kind[2], self.n2) - curl_e3_13 = xp.identity(self.d2) + curl_e3_12 = xp.to_numpy(grad_1d_matrix(derham.spl_kind[2], self.n2)) + curl_e3_13 = np.identity(self.d2) - curl_e3_21 = grad_1d_matrix(derham.spl_kind[2], self.n2) - curl_e3_23 = xp.identity(self.d2) + curl_e3_21 = xp.to_numpy(grad_1d_matrix(derham.spl_kind[2], self.n2)) + curl_e3_23 = np.identity(self.d2) - curl_e3_31 = xp.identity(self.n2) - curl_e3_32 = xp.identity(self.n2) + curl_e3_31 = np.identity(self.n2) + curl_e3_32 = np.identity(self.n2) self._curl_e3 = [ [None, csr(curl_e3_12), csr(curl_e3_13)], @@ -378,16 +383,16 @@ def __init__(self, domain, derham): # =========================== discrete div ====================================== # "polar coeffs" to "polar coeffs" - div_pol_to_pol_1 = xp.zeros((self.n_polar[3][0], self.n_polar[2][0]), dtype=float) - div_pol_to_pol_2 = xp.zeros((self.n_polar[3][0], self.n_polar[2][1]), dtype=float) - div_pol_to_pol_3 = xp.identity(self.n_polar[3][0], dtype=float) + div_pol_to_pol_1 = np.zeros((self.n_polar[3][0], self.n_polar[2][0]), dtype=float) + div_pol_to_pol_2 = np.zeros((self.n_polar[3][0], self.n_polar[2][1]), dtype=float) + div_pol_to_pol_3 = np.identity(self.n_polar[3][0], dtype=float) self._div_pol_to_pol = [[csr(div_pol_to_pol_1), csr(div_pol_to_pol_2), csr(div_pol_to_pol_3)]] # "polar coeffs" to "first tp ring" - div_pol_to_ten_1 = xp.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[2][0])) - div_pol_to_ten_2 = xp.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[2][1])) - div_pol_to_ten_3 = xp.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[3][0])) + div_pol_to_ten_1 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[2][0])) + div_pol_to_ten_2 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[2][1])) + div_pol_to_ten_3 = np.zeros(((self.n_rings[3][0] + 1) * self.d1, self.n_polar[3][0])) for l in range(2): for j in range(self.d1, 2 * self.d1): @@ -398,9 +403,9 @@ def __init__(self, domain, derham): self._div_pol_to_ten = [[csr(div_pol_to_ten_1), csr(div_pol_to_ten_2), csr(div_pol_to_ten_3)]] # eta_3 direction - div_e3_1 = xp.identity(self.d2, dtype=float) - div_e3_2 = xp.identity(self.d2, dtype=float) - div_e3_3 = grad_1d_matrix(derham.spl_kind[2], self.n2) + div_e3_1 = np.identity(self.d2, dtype=float) + div_e3_2 = np.identity(self.d2, dtype=float) + div_e3_3 = xp.to_numpy(grad_1d_matrix(derham.spl_kind[2], self.n2)) self._div_e3 = [[csr(div_e3_1), csr(div_e3_2), csr(div_e3_3)]] diff --git a/src/struphy/polar/kron_2dx1d.py b/src/struphy/polar/kron_2dx1d.py index cdd9dad8c..6c89c80a6 100644 --- a/src/struphy/polar/kron_2dx1d.py +++ b/src/struphy/polar/kron_2dx1d.py @@ -32,6 +32,7 @@ import logging from dataclasses import dataclass, field +import cunumpy as xp import numpy as np import scipy.sparse as spa from feectools.fem.tensor import TensorFemSpace @@ -43,6 +44,7 @@ from scipy.sparse.linalg import spilu, splu from struphy.polar.basic import PolarDerhamSpace, PolarVector +from struphy.polar.linear_operators import DeviceSparseMatrix, device_sparse_module logger = logging.getLogger("struphy") @@ -213,6 +215,29 @@ def poloidal_extraction_matrix( # -------------------------------------------------------------------------------------- # Poloidal solvers # -------------------------------------------------------------------------------------- +class _FactorizedSolve: + """Host factorization with a cached device solve; dense inverses are only used for small 1d factors.""" + + def __init__(self, factor, dense=False): + self._factor = factor + self._dense = dense + self._device = None + + def solve(self, Y): + if Y.shape[-1] == 0: + return Y.copy() + if not xp.is_gpu(Y): + return self._factor.solve(np.asarray(Y, order="F")) + if self._device is None: + if self._dense: + self._device = xp.asarray(self._factor.solve(np.eye(self._factor.shape[0]))) + else: + # CuPy wraps the existing SciPy factors: factorization stays at + # setup, while permutations and batched triangular solves run on device. + self._device = device_sparse_module().linalg.SuperLU(self._factor) + return self._device @ Y if self._dense else self._device.solve(Y) + + class _SparseLU: """Exact (LU) or incomplete (ILU) factorization of a sparse matrix.""" @@ -223,12 +248,13 @@ def __init__(self, A: spa.spmatrix, kind: str = "lu", drop_tol=ILU_DROP_TOL, fil self._fac = splu(A) else: self._fac = spilu(A, drop_tol=drop_tol, fill_factor=fill_factor) + self._solve = _FactorizedSolve(self._fac) self.exact = kind == "lu" def solve(self, Y: np.ndarray) -> np.ndarray: if Y.shape[-1] == 0: return Y.copy() - return self._fac.solve(np.asarray(Y, order="F")) + return self._solve.solve(Y) class _KroneckerBlockDiagonal: @@ -239,18 +265,24 @@ def __init__(self, factors: list[tuple[np.ndarray, np.ndarray]]): self._lus = [] for A1, A2 in factors: self._sizes.append((A1.shape[0], A2.shape[0])) - self._lus.append((splu(spa.csc_matrix(A1)), splu(spa.csc_matrix(A2)))) + self._lus.append( + ( + _FactorizedSolve(splu(spa.csc_matrix(A1)), dense=True), + _FactorizedSolve(splu(spa.csc_matrix(A2)), dense=True), + ) + ) self.exact = True def solve(self, Y: np.ndarray) -> np.ndarray: - out = np.empty_like(Y) + arrays = xp if xp.is_gpu(Y) else np + out = arrays.empty_like(Y) m = Y.shape[-1] off = 0 for (n1, n2), (lu1, lu2) in zip(self._sizes, self._lus): X = Y[off : off + n1 * n2].reshape(n1, n2 * m) - X = lu1.solve(np.asarray(X, order="F")).reshape(n1, n2, m) - X = np.ascontiguousarray(X.transpose(1, 0, 2)).reshape(n2, n1 * m) - X = lu2.solve(np.asarray(X, order="F")).reshape(n2, n1, m) + X = lu1.solve(X).reshape(n1, n2, m) + X = arrays.ascontiguousarray(X.transpose(1, 0, 2)).reshape(n2, n1 * m) + X = lu2.solve(X).reshape(n2, n1, m) out[off : off + n1 * n2] = X.transpose(1, 0, 2).reshape(n1 * n2, m) off += n1 * n2 return out @@ -281,20 +313,37 @@ def __init__(self, A: spa.spmatrix, pole: np.ndarray, outer_solver, batch: int = W = self._outer.solve(self._Aop[:, k0:k1].toarray()) S[:, k0:k1] -= self._Apo @ W self._S = lu_factor(S) + self._device = None self._n = n logger.debug(f"Schur solver: {n = }, pole block {self._p.size}, outer block {self._o.size}") def solve(self, Y: np.ndarray) -> np.ndarray: if Y.shape[-1] == 0: return Y.copy() - Yp = Y[self._p] - Yo = Y[self._o] + if xp.is_gpu(Y): + if self._device is None: + self._device = ( + xp.asarray(self._p), + xp.asarray(self._o), + DeviceSparseMatrix(self._Apo), + DeviceSparseMatrix(self._Aop), + xp.asarray(lu_solve(self._S, np.eye(self._p.size))), + ) + p, o, Apo, Aop, Sinv = self._device + pole_solve = lambda rhs: Sinv @ rhs + arrays = xp + else: + p, o, Apo, Aop = self._p, self._o, self._Apo, self._Aop + pole_solve = lambda rhs: lu_solve(self._S, rhs) + arrays = np + Yp = Y[p] + Yo = Y[o] t = self._outer.solve(Yo) - Xp = lu_solve(self._S, Yp - self._Apo @ t) - Xo = self._outer.solve(Yo - self._Aop @ Xp) - X = np.empty_like(Y) - X[self._p] = Xp - X[self._o] = Xo + Xp = pole_solve(Yp - Apo.dot(t)) + Xo = self._outer.solve(Yo - Aop.dot(Xp)) + X = arrays.empty_like(Y) + X[p] = Xp + X[o] = Xo return X @@ -410,13 +459,14 @@ def forward(self, local: np.ndarray) -> np.ndarray: if self._comm is None: return local[..., c0:c1] - sendbuf = np.concatenate([local[..., a:b].ravel() for a, b in self._chunks]) - recvbuf = np.empty(int(self._recv_counts.sum()), dtype=float) - self._comm.Alltoallv( - [sendbuf, (self._send_counts, self._send_displs), MPI.DOUBLE], - [recvbuf, (self._recv_counts, self._recv_displs), MPI.DOUBLE], - ) - full = np.empty((*self._global_shape, c1 - c0), dtype=float) + sendbuf = xp.concatenate([local[..., a:b].ravel() for a, b in self._chunks]) + recvbuf = xp.empty(int(self._recv_counts.sum()), dtype=float) + with xp.mpi.mpi_buffer(sendbuf) as send, xp.mpi.mpi_buffer(recvbuf, send=False, recv=True) as recv: + self._comm.Alltoallv( + [send, (self._send_counts, self._send_displs), MPI.DOUBLE], + [recv, (self._recv_counts, self._recv_displs), MPI.DOUBLE], + ) + full = xp.empty((*self._global_shape, c1 - c0), dtype=float) for (s, e), d, n in zip(self._boxes, self._recv_displs, self._recv_counts): shape = tuple(ei - si + 1 for si, ei in zip(s, e)) full[tuple(slice(si, ei + 1) for si, ei in zip(s, e))] = recvbuf[d : d + n].reshape(*shape, c1 - c0) @@ -427,15 +477,16 @@ def backward(self, full: np.ndarray) -> np.ndarray: if self._comm is None: return full - sendbuf = np.concatenate( + sendbuf = xp.concatenate( [full[tuple(slice(si, ei + 1) for si, ei in zip(s, e))].ravel() for s, e in self._boxes], ) - recvbuf = np.empty(int(self._send_counts.sum()), dtype=float) - self._comm.Alltoallv( - [sendbuf, (self._recv_counts, self._recv_displs), MPI.DOUBLE], - [recvbuf, (self._send_counts, self._send_displs), MPI.DOUBLE], - ) - local = np.empty((*self._local_shape, self._m_local), dtype=float) + recvbuf = xp.empty(int(self._send_counts.sum()), dtype=float) + with xp.mpi.mpi_buffer(sendbuf) as send, xp.mpi.mpi_buffer(recvbuf, send=False, recv=True) as recv: + self._comm.Alltoallv( + [send, (self._recv_counts, self._recv_displs), MPI.DOUBLE], + [recv, (self._send_counts, self._send_displs), MPI.DOUBLE], + ) + local = xp.empty((*self._local_shape, self._m_local), dtype=float) for (a, b), d, n in zip(self._chunks, self._send_displs, self._send_counts): local[..., a:b] = recvbuf[d : d + n].reshape(*self._local_shape, b - a) return local @@ -573,7 +624,7 @@ def _apply(self, v: PolarVector, out: PolarVector | None, ops_3: list, ops_pol: assert out.space == self._space V = self._space - out_pol = [np.zeros_like(p) for p in v.pol] + out_pol = [xp.zeros_like(p) for p in v.pol] for g, op_3, op_pol in zip(self._groups, ops_3, ops_pol): # 1. eta_3 direction (tensor-product part redistributed, polar coeffs replicated) @@ -583,9 +634,10 @@ def _apply(self, v: PolarVector, out: PolarVector | None, ops_3: list, ops_pol: L = _interior(tp_component(v.tp, c)) l1, l2, l3 = L.shape T3 = self._t3[c] - X = T3.forward(np.ascontiguousarray(L.reshape(l1 * l2, l3).T)) - tmp_tp[c] = np.ascontiguousarray(T3.backward(op_3(X)).T).reshape(l1, l2, l3) - tmp_pol[c] = op_3(v.pol[c].T).T if V.n_polar[c] > 0 else v.pol[c].copy() + X = T3.forward(xp.ascontiguousarray(L.reshape(l1 * l2, l3).T)) + tmp_tp[c] = xp.ascontiguousarray(T3.backward(op_3(X)).T).reshape(l1, l2, l3) + pol = v.pol[c] + tmp_pol[c] = op_3(pol.T).T if V.n_polar[c] > 0 else pol.copy() # 2. poloidal plane s3 = tp_space(V, g.comps[0]).starts[2] @@ -597,7 +649,7 @@ def _apply(self, v: PolarVector, out: PolarVector | None, ops_3: list, ops_pol: nr = V.n_rings[c] blocks.append(tmp_pol[c][:, s3 + k0 : s3 + k1]) blocks.append(full[c][nr:].reshape(-1, k1 - k0)) - Z = op_pol(np.concatenate(blocks, axis=0)) if k1 > k0 else np.concatenate(blocks, axis=0) + Z = op_pol(xp.concatenate(blocks, axis=0)) if k1 > k0 else xp.concatenate(blocks, axis=0) off = 0 for c in g.comps: @@ -618,7 +670,8 @@ def _apply(self, v: PolarVector, out: PolarVector | None, ops_3: list, ops_pol: comm = V.comm for c in range(V.n_comps): if comm is not None and comm.Get_size() > 1 and out_pol[c].size > 0: - comm.Allreduce(MPI.IN_PLACE, out_pol[c], op=MPI.SUM) + with xp.mpi.mpi_buffer(out_pol[c], recv=True) as buf: + comm.Allreduce(MPI.IN_PLACE, buf, op=MPI.SUM) out.pol[c][:] = out_pol[c] return out @@ -638,8 +691,14 @@ class PolarKron2Dx1DOperator(_PolarKron2Dx1DBase): """ def dot(self, v: PolarVector, out: PolarVector | None = None) -> PolarVector: - ops_3 = [lambda X, A=g.A_3: A @ X for g in self._groups] - ops_pol = [lambda Y, A=g.A_pol: A @ Y for g in self._groups] + if xp.is_gpu(v.pol[0]): + if not hasattr(self, "_device_factors"): + self._device_factors = [(xp.asarray(g.A_3), DeviceSparseMatrix(g.A_pol)) for g in self._groups] + factors = self._device_factors + else: + factors = [(g.A_3, g.A_pol) for g in self._groups] + ops_3 = [lambda X, A=A: A @ X for A, _ in factors] + ops_pol = [A.dot for _, A in factors] return self._apply(v, out, ops_3, ops_pol) def transpose(self, conjugate=False) -> "PolarKron2Dx1DOperator": @@ -673,7 +732,7 @@ def __init__( super().__init__(space, groups) self._kind = poloidal_solver self._solver_kwargs = solver_kwargs - self._lu_3 = [splu(spa.csc_matrix(g.A_3)) for g in groups] + self._lu_3 = [_FactorizedSolve(splu(spa.csc_matrix(g.A_3)), dense=True) for g in groups] self._pol_solvers = [ make_poloidal_solver( g.A_pol, @@ -695,7 +754,7 @@ def poloidal_solver(self) -> str: return self._kind def dot(self, v: PolarVector, out: PolarVector | None = None) -> PolarVector: - ops_3 = [lambda X, lu=lu: lu.solve(np.asarray(X, order="F")) for lu in self._lu_3] + ops_3 = [lu.solve for lu in self._lu_3] ops_pol = [s.solve for s in self._pol_solvers] return self._apply(v, out, ops_3, ops_pol) @@ -782,7 +841,7 @@ def _serial_attrs_1d(derham, basis: str, d: int): def _global_quad_pts(derham, d: int) -> np.ndarray: """Global quadrature points (flattened over the cells) in direction d.""" - return np.asarray(derham.H1_1d_serial_splines[d].quad_grid_pts[0][0]).ravel() + return xp.to_numpy(derham.H1_1d_serial_splines[d].quad_grid_pts[0][0]).ravel() def reduced_weight_2d(weight, derham, reduction: str = "midpoint") -> np.ndarray | None: @@ -820,9 +879,12 @@ def reduced_weight_2d(weight, derham, reduction: str = "midpoint") -> np.ndarray nodes, wts = np.polynomial.legendre.leggauss(_N_AVG_3) pts3, wts3 = (nodes + 1.0) / 2.0, wts / 2.0 E1, E2, E3 = np.meshgrid(pts1, pts2, pts3, indexing="ij") - return np.asarray(weight(E1, E2, E3)).reshape(E1.shape) @ wts3 + return xp.to_numpy(weight(xp.asarray(E1), xp.asarray(E2), xp.asarray(E3))).reshape(E1.shape) @ wts3 - assert isinstance(weight, np.ndarray), f"Weight must be callable, np.ndarray or None, but is {type(weight)}." + assert isinstance(weight, np.ndarray) or xp.is_gpu(weight), ( + f"Weight must be callable, an array or None, but is {type(weight)}." + ) + weight = xp.to_numpy(weight) dom = derham.domain_decomposition nq = derham.nquads s1, s2, s3 = (dom.starts[d] * nq[d] for d in range(3)) @@ -836,7 +898,7 @@ def reduced_weight_2d(weight, derham, reduction: str = "midpoint") -> np.ndarray sums[0, s1 : s1 + l1, s2 : s2 + l2] = weight[:, :, mid] sums[1, s1 : s1 + l1, s2 : s2 + l2] = 1.0 else: - wts3 = np.asarray(derham.spline_attributes["H1"].quad_grid_wts[0][2]).ravel() + wts3 = xp.to_numpy(derham.spline_attributes["H1"].quad_grid_wts[0][2]).ravel() assert wts3.size == l3 sums[0, s1 : s1 + l1, s2 : s2 + l2] = weight @ wts3 sums[1, s1 : s1 + l1, s2 : s2 + l2] = wts3.sum() @@ -867,9 +929,9 @@ def poloidal_mass_matrix(derham, bases_row: tuple, bases_col: tuple, weight: np. att_j = [_serial_attrs_1d(derham, bases_col[d], d) for d in range(2)] def data(att): - spans = np.asarray(att.quad_grid_spans[0][0]) - basis = np.asarray(att.quad_grid_bases[0][0])[:, :, 0, :] # (cells, p+1, quad points), no derivatives - wts = np.asarray(att.quad_grid_wts[0][0]) # (cells, quad points) + spans = xp.to_numpy(att.quad_grid_spans[0][0]) + basis = xp.to_numpy(att.quad_grid_bases[0][0])[:, :, 0, :] # (cells, p+1, quad points), no derivatives + wts = xp.to_numpy(att.quad_grid_wts[0][0]) # (cells, quad points) n = att.nbasis[0][0] p = basis.shape[1] - 1 return spans, basis, wts, n, p diff --git a/src/struphy/polar/linear_operators.py b/src/struphy/polar/linear_operators.py index 3cba9d5dc..1da50c261 100644 --- a/src/struphy/polar/linear_operators.py +++ b/src/struphy/polar/linear_operators.py @@ -1,4 +1,5 @@ import cunumpy as xp +from cunumpy.mpi import mpi_buffer from feectools.linalg.basic import LinearOperator from feectools.linalg.block import BlockVector, BlockVectorSpace from feectools.linalg.stencil import StencilVector, StencilVectorSpace @@ -8,8 +9,113 @@ from struphy.linear_algebra.linalg_kron import kron_matvec_2d from struphy.polar.basic import PolarDerhamSpace, PolarVector +# Sparse module used to put the polar blocks on the device on the CuPy backend; ``None`` means ``xp.scipy.sparse`` +# (``cupyx.scipy.sparse``). Tests without a GPU replace it with :func:`set_device_sparse_module`. +_DEVICE_SPARSE_MODULE = None -class PolarExtractionOperator(LinearOperator): + +def set_device_sparse_module(module): + """Set the sparse module that builds the device copies of the polar blocks on the CuPy backend. + + Parameters + ---------- + module : module | None + Must provide ``csr_matrix(host_csr_matrix)`` returning an object with ``.dot(device_array)``. + ``None`` restores the default, ``xp.scipy.sparse`` (``cupyx.scipy.sparse`` on the CuPy backend). + + Returns + ------- + previous : module | None + The module that was set before. + """ + global _DEVICE_SPARSE_MODULE + previous = _DEVICE_SPARSE_MODULE + _DEVICE_SPARSE_MODULE = module + return previous + + +def device_sparse_module(): + """The sparse module for the device copies of the polar blocks (``cupyx.scipy.sparse`` by default).""" + return xp.scipy.sparse if _DEVICE_SPARSE_MODULE is None else _DEVICE_SPARSE_MODULE + + +class DeviceSparseMatrix: + """Device CSR copy of a host polar matrix, for products with device arrays. + + The polar blocks (:class:`~struphy.polar.extraction_operators.PolarExtractionBlocksC1`) are setup data built + with SciPy on the host. On the CuPy backend the polar operators apply them to coefficients that live on the + device; this class holds a CSR copy of a block on the device (``cupyx.scipy.sparse``) and computes ``A @ x`` + there (SpMV/SpMM on the GPU). Blocks without non-zero entries (or with an empty dimension) are not copied, + their products are zeros. + + Parameters + ---------- + host : scipy.sparse matrix | numpy.ndarray + The host matrix. + """ + + def __init__(self, host): + host = csr_matrix(host) + self._shape = host.shape + self._dtype = host.dtype + if host.nnz == 0 or 0 in host.shape: + self._device = None + else: + self._device = device_sparse_module().csr_matrix(host) + + @property + def shape(self): + return self._shape + + @property + def dtype(self): + return self._dtype + + def dot(self, x): + """Product with a device array x of shape (shape[1],) or (shape[1], k).""" + assert x.shape[0] == self.shape[1], f"{x.shape =} does not match {self.shape =}" + if self._device is None: + return xp.zeros((self.shape[0],) + tuple(x.shape[1:]), dtype=float) + # kron_matvec_2d passes transposed views; cuSPARSE takes a contiguous dense operand + return self._device.dot(xp.ascontiguousarray(x)) + + +def device_blocks(blocks): + """2D nested list of host blocks (``None`` allowed) -> same list with :class:`DeviceSparseMatrix` copies.""" + if blocks is None: + return None + return [[None if blk is None else DeviceSparseMatrix(blk) for blk in row] for row in blocks] + + +class _BlocksOnBackend: + """Mixin: the blocks of a polar operator for the active array backend. + + The blocks themselves (properties ``blocks_*``) stay host matrices on every backend. On the NumPy backend + :meth:`_backend_blocks` returns them unchanged; on the CuPy backend it returns device copies + (:func:`device_blocks`), made once per block list and dropped when the list is replaced. + """ + + def _backend_blocks(self, name): + blocks = getattr(self, name) + if blocks is None or xp.get_backend() == "numpy": + return blocks + cache = self.__dict__.setdefault("_device_blocks_cache", {}) + if name not in cache: + cache[name] = device_blocks(blocks) + return cache[name] + + def _drop_device_blocks(self, name): + self.__dict__.get("_device_blocks_cache", {}).pop(name, None) + + +def _allreduce_sum(comm, buf): + """In-place sum with synchronized device buffers or host staging.""" + if comm is not None and comm.Get_size() > 1: + with mpi_buffer(buf, recv=True) as buffer: + comm.Allreduce(MPI.IN_PLACE, buffer, op=MPI.SUM) + + +class PolarExtractionOperator(_BlocksOnBackend, LinearOperator): """ Linear operator mapping from Stencil-/BlockVectorSpace (V) to PolarDerhamSpace (W). @@ -149,6 +255,7 @@ def blocks_ten_to_pol(self, blocks): check_blocks(blocks, self.blocks_ten_to_pol_shapes) self._blocks_ten_to_pol = blocks + self._drop_device_blocks("blocks_ten_to_pol") @property def blocks_ten_to_ten_shapes(self): @@ -168,6 +275,7 @@ def blocks_ten_to_ten(self, blocks): check_blocks(blocks, self.blocks_ten_to_ten_shapes) self._blocks_ten_to_ten = blocks + self._drop_device_blocks("blocks_ten_to_ten") @property def blocks_e3(self): @@ -204,13 +312,17 @@ def dot(self, v, out=None): assert out.space == self._codomain v.tp.copy(out=out) + blocks_e3 = self._backend_blocks("blocks_e3") + # 2. map "first tp ring" to "polar rings" + "first tp ring" if self.blocks_ten_to_ten is not None: - dot_parts_of_polar(self.blocks_ten_to_ten, self.blocks_e3, v, out, map_from_tp=True) + dot_parts_of_polar(self._backend_blocks("blocks_ten_to_ten"), blocks_e3, v, out, map_from_tp=True) # 3. map polar coeffs to "polar rings" (added in-place) if self.blocks_ten_to_pol is not None: - dot_parts_of_polar(self.blocks_ten_to_pol, self.blocks_e3, v, out, map_from_tp=False, add=True) + dot_parts_of_polar( + self._backend_blocks("blocks_ten_to_pol"), blocks_e3, v, out, map_from_tp=False, add=True + ) # "standard" operator (tensor-product vector --> polar vector) else: @@ -226,13 +338,15 @@ def dot(self, v, out=None): # 1. identity operation on outer tp zone out.tp = v + blocks_e3 = self._backend_blocks("blocks_e3") + # 2. map from "polar rings" to polar coeffs if self.blocks_ten_to_pol is not None: - dot_inner_tp_rings(self.blocks_ten_to_pol, self.blocks_e3, v, out, map_to_tp=False) + dot_inner_tp_rings(self._backend_blocks("blocks_ten_to_pol"), blocks_e3, v, out, map_to_tp=False) # 3. map to "polar rings" + "first tp ring" to "first tp ring" if self.blocks_ten_to_ten is not None: - dot_inner_tp_rings(self.blocks_ten_to_ten, self.blocks_e3, v, out, map_to_tp=True) + dot_inner_tp_rings(self._backend_blocks("blocks_ten_to_ten"), blocks_e3, v, out, map_to_tp=True) return out @@ -267,7 +381,7 @@ def transpose(self, conjugate=False): ) -class PolarLinearOperator(LinearOperator): +class PolarLinearOperator(_BlocksOnBackend, LinearOperator): """ Linear operator mapping from PolarDerhamSpace (V) to PolarDerhamSpace (W). @@ -427,6 +541,8 @@ def blocks_pol_to_ten(self, blocks): check_blocks(blocks, self.blocks_pol_to_ten_shapes) self._blocks_pol_to_ten = blocks + self._drop_device_blocks("blocks_pol_to_ten") + @property def blocks_pol_to_pol_shapes(self): return self._blocks_pol_to_pol_shapes @@ -447,6 +563,8 @@ def blocks_pol_to_pol(self, blocks): check_blocks(blocks, self.blocks_pol_to_pol_shapes) self._blocks_pol_to_pol = blocks + self._drop_device_blocks("blocks_pol_to_pol") + @property def blocks_e3_shapes(self): return self._blocks_e3_shapes @@ -467,6 +585,8 @@ def blocks_e3(self, blocks): check_blocks(blocks, self.blocks_e3_shapes) self._blocks_e3 = blocks + self._drop_device_blocks("blocks_e3") + def dot(self, v, out=None): """ Dot product mapping from PolarVector to PolarVector @@ -502,16 +622,19 @@ def dot(self, v, out=None): self.tp_operator.dot(v.tp, out=out.tp) out.set_tp_coeffs_to_zero() + blocks_e3 = self._backend_blocks("blocks_e3") + blocks_pol_to_ten = self._backend_blocks("blocks_pol_to_ten") + # 2. polar to polar (dense) - dot_pol_pol(self.blocks_pol_to_pol, self.blocks_e3, v, out) + dot_pol_pol(self._backend_blocks("blocks_pol_to_pol"), blocks_e3, v, out) # 3. contributions are added in-place if self.transposed: # "first tp ring" to polar - dot_inner_tp_rings(self.blocks_pol_to_ten, self.blocks_e3, v.tp, out, map_to_tp=False, add=True) + dot_inner_tp_rings(blocks_pol_to_ten, blocks_e3, v.tp, out, map_to_tp=False, add=True) else: # polar to "first tp ring" - dot_parts_of_polar(self.blocks_pol_to_ten, self.blocks_e3, v, out.tp, map_from_tp=False, add=True) + dot_parts_of_polar(blocks_pol_to_ten, blocks_e3, v, out.tp, map_from_tp=False, add=True) return out @@ -661,8 +784,7 @@ def dot_inner_tp_rings(blocks_e1_e2, blocks_e3, v, out, map_to_tp, add=False): # sum up local dot products (polar coeffs are held by all processes, the "first tp ring" only at the pole) comm = polar_space.pole_comm if map_to_tp else polar_space.comm - if comm is not None and comm.Get_size() > 1: - comm.Allreduce(MPI.IN_PLACE, res, op=MPI.SUM) + _allreduce_sum(comm, res) # write result to output polar vector (in-place) # (only the rank owning the first eta_1 coeffs holds the "first tp ring", see PolarDerhamSpace) @@ -799,9 +921,7 @@ def dot_parts_of_polar(blocks_e1_e2, blocks_e3, v, out, map_from_tp, add=False): res += kron_matvec_2d([block_e1_e2, block_e3], v.pol[n]).reshape(n_rings_out[m], n2, n3_out[m]) if map_from_tp: - comm = polar_space.pole_comm - if comm is not None and comm.Get_size() > 1: - comm.Allreduce(MPI.IN_PLACE, res, op=MPI.SUM) + _allreduce_sum(polar_space.pole_comm, res) s1, s2, s3 = out_starts[m] e1, e2, e3 = out_ends[m] diff --git a/src/struphy/polar/tests/test_polar_cupy.py b/src/struphy/polar/tests/test_polar_cupy.py new file mode 100644 index 000000000..2e7dc2c5f --- /dev/null +++ b/src/struphy/polar/tests/test_polar_cupy.py @@ -0,0 +1,429 @@ +"""Polar splines on the CuPy backend (struphy-hub/struphy#695, CUDA_STRATEGY.md). + +The polar blocks (:class:`~struphy.polar.extraction_operators.PolarExtractionBlocksC1`) are built on the host on +every backend; on CuPy the polar operators apply device copies of them (``cupyx.scipy.sparse``) to the device +coefficients. :func:`check_polar_backends` builds the same polar Derham on the NumPy and the CuPy backend and +compares extraction operators, derivatives, ``PolarVector`` arithmetic, mass matrices, projectors and polar +``SplineFunction`` coefficients. + +Without a GPU it runs on cunumpy's fake CuPy with CPU-emulated CUDA launches and a dense device stand-in +for cupyx sparse products and triangular solves. Host/device transfers remain visible to cunumpy's counters. +""" + +import os +import subprocess +import sys + +import cunumpy +import numpy as np +import pytest + +from struphy.geometry.tests.test_domain import _cupy_installed, serial_child_env + +requires_cupy = pytest.mark.skipif(not cunumpy.cupy_available(), reason="CuPy/GPU not available") + +MAPPINGS = ("IGAPolarCylinder", "IGAPolarTorus") +FORMS = ("0", "1", "2", "3", "v") +SPACE_IDS = {"0": "H1", "1": "Hcurl", "2": "Hdiv", "3": "L2", "v": "H1vec"} + + +class DenseDeviceSparse: + """Stand-in for ``cupyx.scipy.sparse`` on the fake CuPy: the block as a dense array of the active backend.""" + + class csr_matrix: + def __init__(self, host): + self._array = cunumpy.asarray(host.toarray()) + self.shape = host.shape + + def dot(self, x): + return self._array @ x + + class linalg: + class SuperLU: + """Device-array triangular solves with the same permutations as CuPy's SuperLU.""" + + def __init__(self, factor): + self.L = cunumpy.asarray(factor.L.toarray()) + self.U = cunumpy.asarray(factor.U.toarray()) + self.rows = cunumpy.asarray(np.argsort(factor.perm_r)) + self.cols = cunumpy.asarray(factor.perm_c) + + def solve(self, rhs): + y = cunumpy.linalg.solve(self.L, rhs[self.rows]) + return cunumpy.linalg.solve(self.U, y)[self.cols] + + +def make_domain(mapping, num_elements=(6, 6)): + from struphy import domains + + if mapping == "IGAPolarCylinder": + return domains.IGAPolarCylinder(num_elements=num_elements, degree=(2, 3), a=1.2, Lz=3.0) + if mapping == "IGAPolarTorus": + return domains.IGAPolarTorus(num_elements=num_elements, degree=(3, 2), a=0.8, R0=3.0, sfl=True) + raise ValueError(mapping) + + +def make_polar_derham(mapping, with_mass=True, poloidal_solver="schur", mpi_dims_mask=None): + """Polar Derham (with projectors), its domain and mass matrices (or None), on the active backend.""" + from maybempi import MPI + + 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 + + size = MPI.COMM_WORLD.Get_size() + num_elements = (max(6, 4 * size), max(6, 3 * size), 4) + domain = make_domain(mapping, num_elements[:2]) + options = DerhamOptions( + degree=(domain.degree[0], domain.degree[1], 1), + bcs=(("dirichlet", "free"), None, None), + nquads_proj=(4, 4, 2), + polar_splines=True, + polar_poloidal_solver=poloidal_solver, + ) + derham = Derham( + TensorProductGrid( + num_elements=num_elements, + mpi_dims_mask=mpi_dims_mask, + ), + options, + comm=MPI.COMM_WORLD, + domain=domain, + ) + return domain, derham, WeightedMassOperators(derham, domain) if with_mass else None + + +def _stencils(vec): + """The Stencil- (or Block-) vector part of vec as a list of StencilVectors.""" + from feectools.linalg.stencil import StencilVector + + tp = vec.tp if hasattr(vec, "tp") else vec + return [tp] if isinstance(tp, StencilVector) else list(tp.blocks) + + +def to_host(vec, interior=False): + """Data of a Stencil-/Block-/PolarVector as a list of NumPy arrays (polar coeffs first).""" + pol = list(vec.pol) if hasattr(vec, "pol") else [] + from struphy.polar.kron_2dx1d import _interior + + return [cunumpy.to_numpy(a) for a in pol] + [ + cunumpy.to_numpy(_interior(s) if interior else s._data) for s in _stencils(vec) + ] + + +def fill(vec, seed): + """Fill vec (on the active backend) with reproducible random values (same on every backend).""" + rng = np.random.default_rng(seed) + if hasattr(vec, "pol"): + vec.pol = [cunumpy.asarray(rng.random(p.shape)) for p in vec.pol] + for s in _stencils(vec): + s._data[...] = cunumpy.asarray(rng.random(s._data.shape)) + if hasattr(vec, "pol"): + vec.set_tp_coeffs_to_zero() + vec.update_ghost_regions() + return vec + + +def assert_on_backend(vec, backend): + """All data of vec lives on the given backend.""" + pol = list(vec.pol) if hasattr(vec, "pol") else [] + for a in pol + [s._data for s in _stencils(vec)]: + assert cunumpy.is_gpu(a) == (backend == "cupy"), type(a) + + +def polar_results(mapping, backend, with_mass=True, poloidal_solver="schur", mpi_dims_mask=None): + """Results of the polar operations on `backend`, as host arrays (dict name -> list of arrays).""" + from struphy.polar.basic import PolarVector + + res = {} + with cunumpy.use_backend(backend): + domain, derham, mass = make_polar_derham(mapping, with_mass, poloidal_solver, mpi_dims_mask) + + def record(name, vec): + assert_on_backend(vec, backend) + res[name] = to_host(vec) + + # Host assembly accepts both backend callables and device-resident sampled weights. + from struphy.polar.kron_2dx1d import reduced_weight_2d + + pts = [cunumpy.asarray(a).ravel() for a in derham.spline_attributes["H1"].quad_grid_pts[0]] + mesh = cunumpy.meshgrid(*pts, indexing="ij") + for reduction in ("midpoint", "average"): + + def weight(e1, e2, e3): + return 1.0 + e1 + e2 + (e3 if reduction == "average" else 0.0) + + sampled = weight(*mesh) + from_array = reduced_weight_2d(sampled, derham, reduction) + from_function = reduced_weight_2d(weight, derham, reduction) + np.testing.assert_allclose(from_array, from_function, rtol=1e-12, atol=1e-12) + res["weight-" + reduction] = [from_array] + if with_mass: + from struphy.feec.preconditioner import PolarMassMatrixPreconditioner + + W = mass.create_weighted_mass("H1", "H1", weights=[[sampled]], assemble=True) + pc = PolarMassMatrixPreconditioner(W, poloidal_solver=poloidal_solver, weight_reduction=reduction) + record("array-pc-" + reduction, pc.dot(fill(W.domain.zeros(), seed=70))) + + for n, form in enumerate(FORMS): + E = derham.extraction_ops[form] + x = fill(E.domain.zeros(), seed=10 + n) + y = fill(PolarVector(E.codomain), seed=20 + n) + record(f"E{form}", E.dot(x)) + record(f"E{form}T", E.transpose().dot(y)) + P = derham.dofs_extraction_ops[form] + record(f"P{form}", P.dot(x)) + record(f"P{form}T", P.transpose().dot(y)) + + # PolarVector arithmetic + z = fill(PolarVector(E.codomain), seed=30 + n) + record(f"add{form}", y + z) + record(f"sub{form}", y - z) + record(f"mul{form}", y * 1.5) + record(f"rmul{form}", 0.5 * y) + record(f"neg{form}", -y) + w = y.copy() + w += z + w -= 2.0 * z + w *= 3.0 + record(f"inplace{form}", w) + res[f"dot{form}"] = [np.asarray(float(y.dot(z)))] + res[f"toarray{form}"] = [cunumpy.to_numpy(y.toarray(allreduce=True))] + + # mass matrices (E M E^T with the polar extraction operators) + if with_mass: + from struphy.feec.preconditioner import PolarMassMatrixPreconditioner + + M = getattr(mass, "M" + form) + record(f"M{form}", M.dot(y)) + pc = PolarMassMatrixPreconditioner(M, poloidal_solver=poloidal_solver) + record(f"pcM{form}", pc.dot(M.dot(y))) + # Factor/device caches are populated by the first application. + if backend == "cupy" and (derham.comm.Get_size() == 1 or cunumpy.mpi.get_mpi_cuda_aware()): + with cunumpy.profiling.count_transfers() as transfers: + pc.dot(y) + assert transfers.total == 0, transfers.report() + + for name, op in (("grad", derham.grad), ("curl", derham.curl), ("div", derham.div)): + x = fill(PolarVector(op.domain), seed=40) + y = fill(PolarVector(op.codomain), seed=41) + record(name, op.dot(x)) + record(name + "T", op.transpose().dot(y)) + + # Direct Kronecker projectors, or ILU-preconditioned iterations. + for form in ("0", "1", "2", "3"): + P = getattr(derham, "P" + form) + x = fill(P.I.domain.zeros(), seed=60) + for transposed in (False, True): + for apply_bc in (False, True): + target = P.boundary_op.dot(x) if apply_bc else x + op = (P.I0T if apply_bc else P.IT) if transposed else (P.I0 if apply_bc else P.I) + rhs = op.dot(target) + solved = P.solve(rhs, transposed=transposed, apply_bc=apply_bc) + solved.update_ghost_regions() + target.update_ghost_regions() + record(f"solve{form}-{transposed}-{apply_bc}", solved) + for actual, expected in zip(to_host(solved, interior=True), to_host(target, interior=True)): + np.testing.assert_allclose(actual, expected, rtol=1e-9, atol=1e-10) + # Test operator and preconditioner applications without counting Krylov scalar reductions. + P.pc.dot(x) + if backend == "cupy" and (derham.comm.Get_size() == 1 or cunumpy.mpi.get_mpi_cuda_aware()): + with cunumpy.profiling.count_transfers() as transfers: + P.I.dot(x) + P.pc.dot(x) + assert transfers.total == 0, transfers.report() + + # Commuting projection of analytic functions. + def f(e1, e2, e3): + return cunumpy.sin(np.pi * e1) * cunumpy.cos(2 * np.pi * e2) * cunumpy.cos(2 * np.pi * e3) + e1**2 + + record("Pi0", derham.P0(f)) + record("Pi1", derham.P1([f, f, f])) + record("Pi2", derham.P2([f, f, f])) + record("Pi3", derham.P3(f)) + + # polar SplineFunction: polar coefficients, extracted tensor-product coefficients + for form in ("0", "2"): + field = derham.create_spline_function("f" + form, SPACE_IDS[form]) + assert isinstance(field.vector, PolarVector) + fill(field.vector, seed=50) + field.extract_coeffs() + record("field" + form, field.vector) + record("field_stencil" + form, field.vector_stencil) + # restart: polar coefficients recovered from the extracted tensor-product coefficients + restored = field._restart_extraction_op().dot(field.vector_stencil) + record("restart" + form, restored) + np.testing.assert_allclose(to_host(restored)[0], to_host(field.vector)[0], rtol=1e-10, atol=1e-12) + + return res + + +def check_polar_backends(mapping, with_mass=True, poloidal_solver="schur", mpi_dims_mask=None): + """The polar operations agree on the NumPy and the CuPy backend.""" + host = polar_results(mapping, "numpy", with_mass, poloidal_solver, mpi_dims_mask) + device = polar_results(mapping, "cupy", with_mass, poloidal_solver, mpi_dims_mask) + assert host.keys() == device.keys() + for name in host: + assert len(host[name]) == len(device[name]), name + for a, b in zip(host[name], device[name]): + np.testing.assert_allclose(b, a, rtol=1e-9, atol=1e-11, err_msg=name) + + +def check_polar_backends_fake_cupy(mapping, poloidal_solver="schur"): + """Full polar parity, including mass matrices, with CPU-emulated CUDA launches.""" + from struphy.pic.tests.cuda_emulation import emulated_launches + from struphy.polar import linear_operators + + previous = linear_operators.set_device_sparse_module(DenseDeviceSparse) + try: + with emulated_launches(): + check_polar_backends(mapping, poloidal_solver=poloidal_solver) + finally: + linear_operators.set_device_sparse_module(previous) + + +def test_device_blocks_numpy_backend(): + """On the NumPy backend the polar operators apply their host (SciPy) blocks, as before.""" + import scipy.sparse as sp + + from struphy.polar.linear_operators import DeviceSparseMatrix, set_device_sparse_module + + with cunumpy.use_backend("numpy"): + domain, derham, _ = make_polar_derham("IGAPolarCylinder", with_mass=False) + E = derham.extraction_ops["1"] + assert E._backend_blocks("blocks_ten_to_pol") is E.blocks_ten_to_pol + assert E._backend_blocks("blocks_e3") is E.blocks_e3 + assert all(sp.issparse(b) for row in E.blocks_ten_to_pol for b in row if b is not None) + + # the device copy (here with the dense stand-in on the NumPy backend) computes the same products + previous = set_device_sparse_module(DenseDeviceSparse) + try: + rng = np.random.default_rng(0) + for row in derham.dofs_extraction_ops["1"].blocks_ten_to_ten + E.blocks_ten_to_pol: + for blk in row: + if blk is None: + continue + x = rng.random((blk.shape[1], 3)) + np.testing.assert_allclose(DeviceSparseMatrix(blk).dot(x), blk.dot(x), rtol=1e-14) + empty = DeviceSparseMatrix(sp.csr_matrix((0, 4))) + assert empty.dot(np.ones((4, 2))).shape == (0, 2) + finally: + set_device_sparse_module(previous) + + +@pytest.mark.skipif(_cupy_installed(), reason="the fake CuPy cannot replace an installed CuPy") +@pytest.mark.parametrize("mapping", MAPPINGS) +@pytest.mark.parametrize("poloidal_solver", ("schur", "lu", "ilu")) +def test_polar_fake_cupy(mapping, poloidal_solver): + """Without a GPU: polar Derham on cunumpy's fake CuPy agrees with the NumPy backend. + + Runs in a subprocess because the fake CuPy must be installed before cunumpy is imported. + """ + code = ( + "from struphy.polar.tests.test_polar_cupy import check_polar_backends_fake_cupy; " + f"check_polar_backends_fake_cupy({mapping!r}, {poloidal_solver!r})" + ) + result = subprocess.run( + [sys.executable, "-c", code], env=serial_child_env(CUNUMPY_FAKE_CUPY="1"), capture_output=True, text=True + ) + assert result.returncode == 0, result.stderr[-4000:] + + +@requires_cupy +@pytest.mark.parametrize("mapping", MAPPINGS) +@pytest.mark.parametrize("poloidal_solver", ("schur", "lu", "ilu")) +def test_polar_on_cupy(mapping, poloidal_solver): + """Polar Derham on the GPU (cupyx.scipy.sparse for the polar blocks) agrees with the NumPy backend.""" + check_polar_backends(mapping, poloidal_solver=poloidal_solver) + + +def check_polar_mpi(cuda_aware, mpi_dims_mask, fake=True): + """Backend parity with both Kronecker redistributions and ranks away from the pole.""" + from contextlib import nullcontext + + from struphy.pic.tests.cuda_emulation import emulated_launches + from struphy.polar.linear_operators import set_device_sparse_module + + previous_mpi = cunumpy.mpi.get_mpi_cuda_aware() + previous_sparse = set_device_sparse_module(DenseDeviceSparse) if fake else None + try: + if cuda_aware is None: + # feectools halo exchanges require CUDA-aware MPI. Never probe + # an unknown MPI build with real device addresses (it may crash). + if os.environ.get("STRUPHY_TEST_CUDA_AWARE_MPI") != "1": + pytest.skip("full GPU MPI parity needs STRUPHY_TEST_CUDA_AWARE_MPI=1 on a CUDA-aware MPI build") + cuda_aware = True + cunumpy.mpi.set_mpi_cuda_aware(cuda_aware) + with emulated_launches() if fake else nullcontext(): + check_polar_backends("IGAPolarCylinder", with_mass=False, mpi_dims_mask=mpi_dims_mask) + finally: + cunumpy.mpi.set_mpi_cuda_aware(previous_mpi) + if fake: + set_device_sparse_module(previous_sparse) + + +@pytest.mark.skipif( + not getattr(sys.modules.get("cupy"), "__cunumpy_fake__", False), reason="requires CUNUMPY_FAKE_CUPY=1" +) +@pytest.mark.mpi +@pytest.mark.parametrize("cuda_aware", (True, False)) +@pytest.mark.parametrize("mpi_dims_mask", ((True, False, False), (False, True, False), (False, False, True))) +def test_polar_mpi_fake_cupy(cuda_aware, mpi_dims_mask): + check_polar_mpi(cuda_aware, mpi_dims_mask) + + +@requires_cupy +@pytest.mark.mpi +@pytest.mark.parametrize("mpi_dims_mask", ((True, False, False), (False, True, False), (False, False, True))) +def test_polar_mpi_on_cupy(mpi_dims_mask): + check_polar_mpi(None, mpi_dims_mask, fake=False) + + +def check_polar_collectives(cuda_aware): + """Polar collectives alone, including zero-column chunks, independent of feectools halos.""" + from maybempi import MPI + + from struphy.polar.kron_2dx1d import _BoxTranspose + from struphy.polar.linear_operators import _allreduce_sum + + comm = MPI.COMM_WORLD + rank, size = comm.Get_rank(), comm.Get_size() + previous = cunumpy.mpi.get_mpi_cuda_aware() + try: + cunumpy.mpi.set_mpi_cuda_aware(cuda_aware) + with cunumpy.use_backend("cupy"): + local = cunumpy.asarray(np.arange(6.0).reshape(2, 3) + 100 * rank) + transpose = _BoxTranspose(comm, (2 * rank,), (2 * rank + 1,), (2 * size,), 3) + with cunumpy.profiling.count_transfers() as transfers: + full = transpose.forward(local) + restored = transpose.backward(full) + expected = np.concatenate([np.arange(6.0).reshape(2, 3) + 100 * r for r in range(size)]) + a, b = transpose.my_chunk + np.testing.assert_array_equal(cunumpy.to_numpy(full), expected[:, a:b]) + np.testing.assert_array_equal(cunumpy.to_numpy(restored), cunumpy.to_numpy(local)) + if cuda_aware or size == 1: + assert transfers.total == 0, transfers.report() + else: + assert transfers.to_host == transfers.to_device == 2, transfers.report() + values = cunumpy.asarray(np.arange(5.0) + rank) + _allreduce_sum(comm, values) + np.testing.assert_array_equal(cunumpy.to_numpy(values), size * np.arange(5.0) + size * (size - 1) / 2) + finally: + cunumpy.mpi.set_mpi_cuda_aware(previous) + + +@pytest.mark.skipif( + not getattr(sys.modules.get("cupy"), "__cunumpy_fake__", False), reason="requires CUNUMPY_FAKE_CUPY=1" +) +@pytest.mark.mpi +@pytest.mark.parametrize("cuda_aware", (True, False)) +def test_polar_collectives_mpi_fake(cuda_aware): + check_polar_collectives(cuda_aware) + + +@requires_cupy +@pytest.mark.mpi +def test_polar_collectives_mpi_on_cupy(): + """Real device buffers staged through a regular MPI build; no CUDA-aware MPI needed.""" + check_polar_collectives(False)