[max-models]

Part 7 of 8 of the split of #772 (cuda-aware-mpi-fixes) into stacked PRs. Requires struphy-hub/feectools#104 to be merged into feectools devel-tiny first (PR 1 bumps the submodule).

• Pseudo-random markers are drawn in chunks of 2**20 (was 10000): 16 instead of 1600 iterations for 16M markers, each a few small kernels and a device sync on the GPU. The NumPy draws do not depend on the chunk size.

• update_weights finds the indices of the valid markers once and reads only the quantities it needs (xp.take of one row of the buffer each), instead of a boolean selection per property and a count of the valid markers.

• mpi_sort_markers: the ghost bookkeeping is skipped without sorting boxes, where no ghosts can exist. There are no barriers around the device exchange, which waits for every send and receive (cunumpy.mpi.exchange). The host path keeps them. The sorting positions for alpha = 1 are one add instead of five kernels.

• PushVxB without a magnetic field variable computes the total field once, so there is no ghost exchange (an MPI synchronization point) before every push.

• Device MPI sort (Particles._mpi_sort_device, struphy.pic.mpi_sort_cuda): on the GPU one kernel (sort_destinations) finds the destination of every marker and counts the markers per destination, instead of a pass over the buffer per direction and per rank. The leaving markers are listed once, stably sorted by destination and gathered once per destination rank. The result is that of the host path (the same markers in the same columns; lost markers become holes), at the cost of one 8-byte read per rank for the counts. STRUPHY_DEVICE_SORT=1 or 0 forces it on (also on the fake CuPy) or off. The original branch imported mpi_sort_cuda, but it was never committed (and called an undefined _to_numpy_for_kernel), so it is written new here for the (n_cols, n_markers) buffer of devel. test_mpi_sort_markers_device_kernel_fake_cupy compares it with the NumPy sort, CUDA-aware and staged, periodic and removing boundaries, on 2 and 4 ranks.

Tests

Local, macOS without a GPU, in a fresh environment with feectools 1d8d47c and the pinned cunumpy, after merging devel (with #774 and #775). The PR passes test_cuda_parity, test_cuda_emulation, test_kernel_backends, test_sorting, test_accum_matrix_cupy, test_preconditioner_cupy, test_poisson and test_output on its own. The top of the stack also passes the mass-matrix, preconditioner, Derham-GPU, fused-gradient, Kronecker-block, Poisson, shear Alfvén, post-processing, polar fake-CuPy and all pic tests (1477 passed), the examples test for the two profiling cases, the sorting, Poisson and fused-gradient tests under mpirun -n 2, and the device MPI sort on the fake CuPy on 2 and 4 ranks. The Kronecker products match the stencil mass matrices to 5e-16 for M0, M1, M2, M3, Mv, M1n and M2n on a Cuboid, periodic and Dirichlet, on 1 and 2 ranks.
Not tested on a GPU here. The CUDA kernels and the device paths need a run on the cluster.

Stack (merge in order)

0. Device ghost-exchange plans, fused block exchanges, single-reduction PCG (CUDA-aware MPI) feectools#104 (feectools, into devel-tiny)

#	PR	base
1	#782 update feectools (device ghost exchange, single-reduction PCG); pcg_sr solver option	devel
2	#783 host/device conversions in post-processing; math.sqrt in Species	#782
3	#784 CUDA mass-matrix assembly kernel kernel_3d_mat	#783
4	#785 apply separable mass matrices as Kronecker products of 1D banded matrices	#784
5	#786 Poisson/diffusion solve with fewer passes over the vectors	#785
6	#787 marker rows prepared inside the push kernels	#786
7	#788 chunked marker loading, row-based weight update, leaner MPI sorting; static field in PushVxB (this PR)	#787
8	#789 profiling cases Poisson cube and Vlasov slab (CPU node vs GPU, GPU scaling)	#788

🤖 Generated with Claude Code