Set up sorting boxes on the CuPy backend

[max-models]

Closes #726

What changed

• Particles setup with boxes (pic/base.py):
  • n_boxes uses math.prod instead of xp.prod(xp.array(...)). On CuPy the old version gave a 0-d device array, and SortingBoxes._set_boxes failed with type ndarray doesn't define __round__ method.
  • domain_array has one host copy, _domain_array_host, made once in __init__. The host bookkeeping uses it: n_cells, the per-rank cell counts, is_domain_boundary (Python bools), _get_neighbouring_proc (the pyccel distance used to get 0-d device arrays, so ParticlesSPH setup failed on CuPy) and the neighbour-rank check.
• SortingBoxes._set_boxes (pic/sorting.py): _neighbours is built on the host with the pyccel initialize_neighbours and moved to the active backend once with xp.to_cunumpy.
• Box sorting kernels (pic/kernels/<name>/, one folder per kernel as in CUDA_STRATEGY.md): box_particles, fill_boxes, sort_and_fill_boxes and box_and_sort_particles, declared with Kernel.from_folder(..., missing_cuda="fallback"). On NumPy they run the pyccel kernel. On CuPy they are an explicit host fallback (Port the diffusion and SPH kernels to CUDA #694 adds <name>_cuda.cu without changing the call sites). The helpers stay in sorting_kernels.py.
  • do_sort() without SPH ghosts calls only box_and_sort_particles: assign the boxes, count per box, sort, refill boxes. On CuPy the markers cross once in each direction. The per-box count is done on the host (no device bincount, no implicit sync).
  • put_particles_in_boxes calls box_particles, then (SPH) the ghost exchange and fill_boxes.
  • boxes is rebuilt after the sort, so it points to the sorted rows (before, it still pointed to the rows before the sort).
  • All kernels use SortingBoxes.box_index for the box column.

Time loop of non-SPH models

For a non-SPH model (e.g. VlasovAmpereOneSpecies), setting boxes_per_dim alone does not sort anything during the time loop. Pushers only call put_particles_in_boxes when sorting_boxes.communicate is true, which is SPH only. Box sorting runs only:

• during setup, via draw_markers(sort=True), when SortingParameters(do_sort=True);
• every sort_step steps in sim.py, via do_sort(), when EnvironmentOptions.sort_step > 0.

Both paths now work on CuPy through the host fallback above. The fallback still has to be ported (#694).

Remaining for #694

• CUDA versions of the four box kernels (<name>_cuda.cu in their folders).
• SPH evaluation and ghost-box communication on the device. I checked that ParticlesSPH setup works on CuPy, also on several ranks. Drawing markers with a domain needs CUDA launches, which the fake CuPy cannot run, so nothing after setup was exercised for SPH.

Tests

• test_sorting.py:
  • check_sorting_boxes_on_cupy (serial): creates Particles6D(boxes_per_dim=(4, 3, 2)) on NumPy and on CuPy and compares neighbours. It then runs do_sort() on the same markers on both backends and compares the markers and box arrays. It also checks that the markers cross host↔device exactly once in each direction (xp.profiling.count_transfers), and checks the ParticlesSPH setup. It runs as test_sorting_boxes_on_cupy (real CuPy, mpi_skip) and test_sorting_boxes_on_cupy_fake_cupy (fake CuPy via run_fake_cupy_child).
  • assert_boxes_match_markers: every row of boxes lists exactly the marker rows of that box. It is used in test_sorting, in the new test_boxes_valid_after_do_sort (both sort paths) and in the CuPy checks.
• test_sorting_device.py: check_sorting_boxes_device runs on every rank (fake CuPy in CI on 1–4 ranks; test_sorting_boxes_gpu on a GPU). It does the same do_sort() comparison and transfer count with a real decomposition, checks that the box kernel used this rank's domain_array row, and compares the ParticlesSPH neighbouring procs (26 directions), boundary flags and boundary boxes with NumPy.
• Locally: mpirun -n 1..4 env CUNUMPY_FAKE_CUPY=1 pytest --with-mpi pic/tests/test_sorting_device.py and mpirun -n 2/3 pytest --with-mpi pic/tests/test_sorting.py pass.

Merge order

Based on devel. It does not depend on the unmerged CUDA PRs.

🤖 Generated with Claude Code