Collection - Run Struphy kernels on NVIDIA GPUs (CUDA) #650

Tracking issue for running Struphy's compute kernels on NVIDIA GPUs, next to the existing pyccel (CPU) kernels.

The design, principles, per-PR details and open questions are in CUDA_STRATEGY.md, which is kept up to date in each PR. This issue tracks progress. Each step is a small PR that keeps the CPU code path working and tested. The PRs are stacked: each branch contains the previous one, so they are reviewed and merged in order.

Goal: a developer adding a kernel creates a folder <package>/kernels/<name>/ with:

• the pyccel kernel <name>_kernels.py;
• an __init__.py that declares <name> = Kernel.from_folder(__name__, structs=CUDA_STRUCTS);
• optionally, <name>_cuda.cu, plus its parity cases in pic/tests/cuda_parity_cases.py.

Loading, dispatch and the pyccel/CUDA parity tests are done by the infrastructure. The parity tests run on a GPU, and without one through CPU emulation. On the GPU, a kernel without a CUDA version raises a clear error instead of silently falling back to the CPU.

Where things live:

• Infrastructure: cunumpy 0.6. Kernel, KernelCatalog, PyccelKernel and CudaKernel are in cunumpy.kernels; CudaStructArguments, CudaStruct and write_cuda_header are in cunumpy.arguments; the array views come from cunumpy/array_view.cuh; the testing helpers are in cunumpy.kernel_testing.
• Argument classes, in pairs in kernel_arguments/: a pyccel class (MarkerArguments, DerhamArguments, DomainArguments, LocalProjectorsArguments) for the NumPy backend, and a CUDA class (Cuda<Name>) with the same constructor for the CuPy backend. The owner (Particles, Derham, Domain, the local projectors) creates the one for its backend.
• CUDA kernels see only the C structs generated from those classes (pusher_args.cuh, local_projectors_args.cuh).
• FEEC data on the GPU (stencil vectors and matrices, MPI exchange, one GPU per rank) is in feectools.

Done

• PRs 1–12 (Cuda pr 1 proof of concept #643, Cuda pr 2 kernel files #644, Cuda pr 3 kernel catalog #645, Cuda pr 4 pusher #646, Cuda pr 5 domain #649, Cuda pr 6 particles on gpu #651, Cuda pr 7 derham #652, Cuda pr 8 cuda headers #653, Cuda pr 9 kernel folders #663, Cuda pr 10 device helpers #665, Cuda pr 11 push eta #666, Cuda pr 12 particle kernels #667): proof of concept, CUDA source files, kernel catalog, Pusher accepting a Kernel, Domain/Particles/Derham on the GPU, shared CUDA header, one folder per pusher kernel, device helper functions, and the first CUDA kernels (push_eta_stage, push_vxb_analytic, push_vxb_implicit, reflect, spline evaluation).
• PR 13: Move the kernel infrastructure to cunumpy (Cuda pr 13 cunumpy #668): struphy's own dispatch code is deleted, and pusher_args.cuh is generated.
• PR 14: Accumulation folders and the first Vlasov–Ampère kernels (Cuda pr 14 vlasov ampere #670): push_v_with_efield, push_weights_with_efield_lin_va, charge_density_0form (atomics), and CPU emulation of every CUDA kernel.
• PR 15: Use cunumpy as intended (Cuda PR 15 - Use cunumpy as intended: kernels as imports, no catalog lookups #671): every folder declares its kernel, and the code imports it.
• PR 16: One CPU and one GPU argument class per argument type (CUDA PR 16 - One pyccel and one CUDA argument class per argument type #679):
  • plain pyccel/CUDA class pairs replace cunumpy's PyccelStructArguments;
  • every kernel that takes argument objects lives in its own folder (31 more kernels: diagnostics, SPH, spline, geometry, FEEC and local projector kernels, and reflect);
  • prepare_kernel is removed, so CUDA kernels compile on first use;
  • the parity-test inputs moved from <name>_test_args.py files in the kernel folders into pic/tests/cuda_parity_cases.py;
  • resolved Replace MARKER with a 2D view #673, Remove the _evaluate_cuda helper in psydac_derham.py #674 and Remove the _cuda_reflect caller, we should always call reflect, both on GPU and CPU #675.
• feectools CUDA stack, merged into devel-tiny and released as feectools 0.3.1:
  • CUDA 1 - Run feectools with the CuPy backend feectools#90: feectools on the CuPy backend.
  • CUDA 2 - MPI with device buffers feectools#86: MPI with device buffers.
  • CUDA 3 - One GPU per MPI rank feectools#87: one GPU per MPI rank.
  • CUDA 4 - Stencil operations on the device, one folder per kernel feectools#88: stencil dot, transpose, inner and axpy on the device, one folder per kernel, with parity and CPU-emulation tests.
• cunumpy 0.6.0 on PyPI. Driven by this work:
  • PyccelStructArguments and __host_args__ were removed in favour of explicit pyccel/CUDA class pairs;
  • the C-contiguous views CArray1D–CArray4D were added;
  • CudaKernel moved to cunumpy.kernels, and the argument classes to cunumpy.arguments;
  • MPI detection moved to maybempi.

Open PRs (merge in this order)

• PR 17: feectools with CUDA (CUDA PR 17 - struphy on the feectools CUDA stack #680, base devel): the feectools submodule points at devel-tiny, which contains the feectools CUDA stack, and pyproject.toml pins feectools>=0.3.1, <=0.3.1.
• PR 18: Geometry evaluation on the GPU for all analytic mappings (CUDA PR 18 - Geometry evaluation on the GPU for all analytic mappings #681):
  • CUDA versions of kernel_evaluate_pic, kernel_evaluate, kernel_pullpush_pic and kernel_pullpush, with the full metric chain (f, df, det_df, df_inv, g, g_inv, pull/push) for kind_map 10–12, 20–22 and 30–32;
  • the Cuboid-only restriction of all CUDA kernels is lifted, and CuPy particle runs work again.
• PR 19: Spline mappings on the GPU (CUDA PR 19 - Spline mappings on the GPU #683, kind_map 0–2):
  • DomainArgs gets array views for t1..3, ind1..3 and cx/cy/cz;
  • IGAPolarCylinder, IGAPolarTorus, Tokamak, GVEC and DESC can be created on CuPy (their setup runs on the host);
  • device versions of spline_3d, spline_2d_straight and spline_2d_torus.
• PR 20: Use cunumpy 0.6.0 (Use cunumpy 0.6.1 #685): new import locations (cunumpy.kernels, cunumpy.arguments) and regenerated headers. It needs Update to cunumpy 0.6.1 feectools#94 (feectools on cunumpy 0.6), which it points its submodule at.

None of these PRs has run on a GPU. The CPU emulation (every CUDA kernel compared with pyccel without a GPU) and the CPU regression tests are the gate so far.

Next steps

Each step is a sub-issue, labelled GPU and CUDA:

1. Merge the open CUDA PRs and release feectools on cunumpy 0.6 #686 Merge the open CUDA PRs and release feectools on cunumpy 0.6
2. First GPU run of the CUDA tests (H100) #687 First GPU run of the CUDA tests (H100)
3. 6D array views in cunumpy for the matrix accumulations #688 6D array views in cunumpy for the matrix accumulations (vlasov_maxwell, linear_vlasov_ampere)
4. Run VlasovAmpereOneSpecies end to end on the GPU #689 Run VlasovAmpereOneSpecies end to end on the GPU (including when kernels are compiled)
5. Decide between hand-written CUDA and code generation before the guiding-centre kernels #690 Decide between hand-written CUDA and code generation before the guiding-centre kernels
6. Port the remaining models, in the porting order of CUDA_STRATEGY.md:
   • Port the guiding-centre kernels to CUDA (ToyDrift, DriftKineticElectrostaticAdiabatic) #691 guiding centre (ToyDrift, DriftKineticElectrostaticAdiabatic)
   • Port the hybrid MHD-kinetic 6D kernels to CUDA (current and pressure coupling) #692 hybrid MHD-kinetic 6D (current and pressure coupling)
   • Port the hybrid 5D kernels to CUDA #693 hybrid 5D
   • Port the diffusion and SPH kernels to CUDA #694 diffusion and SPH
7. What still runs only on the host:
   • Polar splines on the CuPy backend #695 polar splines on CuPy
   • MHD equilibria on the CuPy backend #696 MHD equilibria on CuPy
   • Field evaluation on the device (SplineFunction and diagnostics) #697 field evaluation on the device
   • Put each kernel into its own folder and file #411 FEEC assembly kernels into folders
8. Multi-rank marker exchange with CUDA-aware MPI #698 Multi-rank marker exchange with CUDA-aware MPI
9. GPU runner in CI for the CUDA tests #699 GPU runner in CI for the CUDA tests
10. Smaller items:
    • Add more struphy cuda namespaces #672 CUDA namespaces
    • Move kernel creation #678 spline-evaluation argument tuples
    • Declare outputs for the PyccelKernels #342 declare outputs= for the PyccelKernels
    • cunumpy features that would replace struphy helpers #700 cunumpy features that would replace struphy helpers

Earlier sub-issues: #411, #672, #678 and #342 (open), and #673, #674, #675 and #134 (closed).

Unrelated bugs found along the way go into their own issues and PRs.

🤖 Generated with Claude Code