GPU: row-major marker buffer and hole rows make column operations bandwidth-bound #707

Summary

On the CuPy backend, a large share of the time in model.integrate is not spent in the pusher kernels but in
column operations on the marker buffer (Particles.markers). Two properties of the buffer cause this:

1. Row-major (array-of-structs) layout. markers has shape (n_rows, n_cols) in C order, so a column
   (or a narrow range of columns) is strided by n_cols * 8 bytes. Any operation on a few columns touches the
   cache lines of every row, i.e. it costs about as much memory traffic as sweeping the whole buffer.
2. Hole rows are processed like valid ones. The buffer has room for n_rows ≈ (1 + bufsize) * Np rows
   (in the benchmark below: 8,002,000 rows for 4,000,000 markers). Holes can sit anywhere in it and nothing
   tracks a high-water mark, so column operations and kernels always run over all n_rows rows.

Proposal: store the markers column-wise (struct-of-arrays, e.g. shape (n_cols, n_rows) in C order or
(n_rows, n_cols) in F order) on the GPU, and keep valid markers compact (or track the range of rows that
can hold valid markers) so that launches can be limited to it.

The preamble in Pusher._push is

markers[:, first_pusher_idx:first_shift_idx] = markers[:, : 3 + vdim]
markers[:, first_shift_idx:residual_idx] = 0.0
markers[:, residual_idx:-2] = 0.0

Each of these is one strided pass over the full 1.66 GB buffer, this takes longer than the kernel call itself.

The preamble can be fused into one CUDA kernel (src/struphy/pic/pushing/prepare_push_cuda.cu), with
one thread per entry of the contiguous column range so that writes are coalesced. See #708

Skipping holes inside the kernel gains only ~10%, since every thread still has to read column 0 of its row
to know whether it is a hole, which touches the same cache lines.

Why holes cannot simply be skipped today

• The CUDA pushers detect a hole by markers[ip, first_init_idx] == -1; that -1 is written by the preamble
  copying column 0. A marker that became a hole in the previous push still has its old position in
  first_init_idx until the next preamble.
• Iterative pushers take xp.max(markers[:, residual_idx]) and count non-converged markers over all rows,
  relying on hole rows having a zero residual.
• There is no high-water mark of valid rows; holes are created anywhere by boundary conditions and MPI sorting.

Suggested direction

1. Column-wise marker storage on the GPU. Every column operation (preamble, markers[:, 0] == -1 in
   update_holes, valid_mks masks, diagnostics) becomes a contiguous, fully coalesced sweep over only the
   columns involved. For the preamble that is 15 of 26 columns written and 6 read, i.e. well below the
   whole-buffer floor above. The kernels access args_markers.markers(ip, j) through Array2D, so the
   accessor could absorb the layout change; the Python side (markers[:, slice], markers[valid_mks],
   MPI send buffers, HDF5 output) needs auditing.
2. Keep valid markers compact, or track a high-water mark (n_rows_used) updated wherever markers are
   created, removed or moved (loading, boundary conditions, mpi_sort_markers, refilling), so that launches
   cover [0, n_rows_used) instead of [0, n_rows). In the benchmark this halves the work of every
   full-buffer operation. Hole rows inside that range would still need the -1 marker and a zero residual.
3. Outside model.integrate, the same trace shows ~6 ms per step of boolean-mask work over all rows in the
   diagnostics (cupy_scan_naive, cupy_getitem_mask, cupy_copy__bool_bool), which would benefit from 1. and 2. as well.

Plan for part 1: struct-of-arrays via order="F"

Allocating the buffer as xp.zeros((n_rows, n_cols), order="F") gives exactly the struct-of-arrays memory layout, and every markers[ip, j] and markers[:, a:b] keeps working unchanged. The alternative, storing the transpose with shape (n_cols, n_rows), would mean rewriting every index expression. So F order is the way to do it, but it can't just be a flag. Findings below are on cuda-namespaces (6c2a772, the top of the CUDA stack), with struphy's pyccel 2.2.1.

What enforces the layout explicitly

1. pyccel checks the memory order at every call. An F-ordered array passed to a float[:,:] argument raises TypeError: Wrong argument type for argument markers : argument does not have the expected ordering, and the reverse fails too. Declaring float[:,:](order=F) fixes it, and that also works on a pyccel class attribute (tested). Places to change:

   • MarkerArguments.markers (kernel_arguments/pusher_args_kernels.py), which ~110 kernel files receive;
   • 20 kernels that take markers directly: sorting, sampling, utilities, reflect, the diagnostics, eval_spline_mpi_markers, kernel_evaluate_pic, kernel_pullpush_pic and the orbits post-processing.

   Every caller must then pass F-ordered arrays. That includes copies such as markers_wo_holes (markers[~holes]), which come out in C order.

2. CudaMarkerArguments rejects it. _device_array raises "must be a C-contiguous array". The CUDA kernels themselves already handle any layout, because Array2D carries strides, so only this check needs relaxing.

3. Hand-written row-major CUDA. prepare_push_cuda.cu (Fuse the pusher preamble into one CUDA kernel #708) indexes markers + ip * n_cols.

What silently creates C-ordered arrays

• a.copy(): NumPy and CuPy default to order="C", unlike np.copy, which defaults to "K".
• Boolean and fancy row indexing (markers[valid_mks]), xp.concatenate (the send lists), MPI receive buffers and h5py reads.

These are harmless as long as the result is only written back into the F buffer with markers[idx] = .... They break when the result is passed to a kernel. The MPI exchange (mpi_sort_markers) already builds temporary buffers and scatters them back, so it doesn't depend on the layout. Only its row gathers become strided.

Quick CPU benchmark

One thread, Fortran backend, 2M markers × 26 columns. The kernel is pusher-like: it reads 9 columns and writes 6. The preamble is the three column assignments from Pusher._push, run as NumPy operations.

So the kernel is about 8% slower and the column work is 8× faster, even on the CPU. This argues for one layout on both backends rather than SoA on the GPU only. Splitting by backend would also break the pyccel fallback on CuPy, because converting a device array to NumPy keeps its F order, and the pyccel signature would reject it. This was a toy kernel; real kernels with spline evaluation and OpenMP still need measuring (step 5).

Steps

• 0. One place that creates the buffer.
  • Particles._markers is only assigned in pic/base.py (xp.zeros((self.n_rows, self.n_cols))). Add a layout constant used there and for _lost_markers, the receive buffers and the send lists.
  • Add a _check_layout() assert that runs in debug mode and in the tests.
• 1. pyccel signatures.
  • Change MarkerArguments.markers and the 20 direct markers: "float[:, :]" arguments to float[:, :](order=F).
  • Convert other 2D per-marker arrays that reach kernels (callers of lost_markers and markers_wo_holes), or use xp.asfortranarray(...) at those call sites.
  • Recompile, then run the PIC tests. The ordering TypeError points to every caller that was missed.
• 2. CUDA.
  • Relax _device_array for markers to accept F-contiguous arrays.
  • Rewrite prepare_push_cuda.cu (Fuse the pusher preamble into one CUDA kernel #708) for column-major, or drop it. With F order the preamble is three contiguous column-block copies (markers[:, a:b] = markers[:, :k]), which CuPy runs as coalesced copies. Measure before deciding.
  • Rerun the parity tests (pic/tests/test_cuda_parity.py) with F-ordered marker inputs: in CPU emulation now, and on the H100 in First GPU run of the CUDA tests (H100) #687.
• 3. Python-side audit (pic/base.py, pic/sorting.py, sampling, diagnostics, post_processing, marker evaluation in geometry/base.py).
  • .copy(), .reshape() and .ravel() on marker arrays: use order="K", or make sure the result never reaches a kernel.
  • Restart and HDF5 output: datasets stay (n_markers, n_cols) on disk, so files don't change. Writing a strided row subset gets slower; measure it for large outputs.
  • Box sorting and compact_by_mask row moves: each moved row now means n_cols strided accesses. Measure; this happens once per sort, not once per push.
• 4. Layout test. Allocate, load, push, sort, exchange (2 ranks, fake CuPy), save and restart. Check that the buffer stays F-ordered and that results match the C-ordered baseline bit-for-bit.
• 5. Benchmarks before merging.
  • CPU: a few real models (VlasovAmpere, guiding centre, SPH) with OpenMP, C vs F.
  • GPU: the trace above (preamble, update_holes, valid_mks masks, diagnostics).
  • If a CPU kernel regresses badly, fix that kernel (e.g. read a marker's columns into locals once) rather than splitting the layout by backend.
• 6. Part 2, separate PR: track the range of rows that can hold valid markers (n_rows_used) so that launches skip trailing holes. This is independent of the layout and roughly halves the work again in the benchmark above.

Steps 0–4 are mostly mechanical: ~20 annotations, one pyccel class, one CUDA check, one CUDA kernel, plus the audit and tests. The risk is in the audit (step 3) and in CPU performance (step 5).