Cunumpy helpers

README.md

@@ -377,7 +377,7 @@ kernel(particles.kernel_args, dt, n_threads=n)  # host or CUDA kernel
 Kernels ported from pyccel index arrays like `markers[ip, j]`, which needs
 shapes and strides rather than bare pointers. The shipped header
 `cunumpy/array_view.cuh` (found by every `CudaKernel`) provides the strided
-views `Array1D<T>` to `Array3D<T>`; a parameter or struct field of that type
+views `Array1D<T>` to `Array4D<T>`; a parameter or struct field of that type
 takes a CuPy array, contiguous or not, and indexes `a(i, j)`. The struct can be
 generated from the annotations of the pyccel argument class, so the Python
 class is the one definition, and written to a header that a test keeps in sync:

docs/source/best-practices.md

@@ -37,7 +37,12 @@ A condensed checklist. Each item links to the guide with the reasoning.
   kernels by profile order, switch to `"raise"` when done. ([Porting
   kernels](kernels/overview.md))
 * Keep the CUDA kernel's argument list identical to the host kernel's; put both
-  in one folder.
+  in one folder, and test it with `catalog.check_signatures()`.
+* Compile Pyccel host kernels ahead of time with the `pyccel` command, or at run
+  time with `from_package(..., compile_host=...)`, and keep a NumPy version as
+  `host_fallback`.
+* Use `dispatch="arrays"` if host arrays reach kernels while CuPy is active.
+* Ship `.cu`/`.cuh` files as package data.
 * Declare `outputs` on host kernels so the fallback copies back only what was
   written, and never forget an argument that is written.
 * Pass device arrays to `CudaKernel`; build argument objects once with
@@ -58,6 +63,8 @@ A condensed checklist. Each item links to the guide with the reasoning.
   everywhere and use the GPU where there is one. ([Testing
   kernels](kernels/testing.md))
 * One `assert_kernels_agree` test over `catalog.parity_cases()`.
+* `emulate_cuda_kernel` tests, so CPU-only CI checks the CUDA arithmetic.
+* Seed `xp.random_streams` with `(seed, rank)` and draw only from it.
 * Debug crashes with `CUNUMPY_CUDA_DEBUG=1`, then `compute-sanitizer`.
   ([Debugging](kernels/debugging.md))
 * Time with `timed_region()`, profile with `nvtx_range()` and `nsys`.

docs/source/examples/particle_recipes.py

@@ -0,0 +1,100 @@
+"""Recipes for particle codes with cunumpy: the code of the "Particle codes" page.
+
+Every function runs on NumPy and on CuPy arrays (``xp`` follows the active
+backend); ``tests/unit/test_particle_recipes.py`` runs them.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+import cunumpy as xp
+
+
+def remove_dead(markers, alive):
+    """Drop the markers whose ``alive`` flag is False (a new, compact array)."""
+    return markers[alive]
+
+
+def compact_in_place(markers, alive):
+    """Move the live markers to the front of a preallocated buffer; return their count.
+
+    ``markers[:n]`` are the live markers afterwards, in their original order,
+    and the rows behind them are free for injection. The right-hand side is a
+    copy (boolean indexing), so source and destination may overlap.
+    """
+    n_alive = int(alive.sum())
+    markers[:n_alive] = markers[alive]
+    return n_alive
+
+
+def sort_by_cell(positions, lower, cell_size, n_cells):
+    """Order the markers by cell; return the order, the cells and the cell offsets.
+
+    After ``markers = markers[order]`` the markers of cell ``c`` are
+    ``markers[offsets[c]:offsets[c + 1]]``. The sort is stable, so markers
+    keep their relative order within a cell (reproducible results).
+    """
+    cell = xp.floor((positions - lower) / cell_size).astype(xp.int64)
+    cell = xp.clip(cell, 0, n_cells - 1)
+    order = xp.argsort(cell, kind="stable")
+    counts = xp.bincount(cell, minlength=n_cells)
+    offsets = xp.zeros(n_cells + 1, dtype=xp.int64)
+    offsets[1:] = xp.cumsum(counts)
+    return order, cell[order], offsets
+
+
+def deposit_nearest_cell(cell, weights, n_cells):
+    """Sum the weights per cell without atomics: the sort-then-reduce deposit."""
+    return xp.segment_sum(weights, cell, n_cells)
+
+
+def pack_for_ranks(markers, destination, n_ranks):
+    """Group the markers by destination rank; return the send buffer and the counts.
+
+    ``destination[i]`` is the rank marker ``i`` moves to (its own rank to stay).
+    The markers for rank ``r`` are the rows ``displacements[r]`` to
+    ``displacements[r] + counts[r]`` of the buffer.
+    """
+    order = xp.argsort(destination, kind="stable")
+    counts = xp.bincount(destination, minlength=n_ranks)
+    return xp.ascontiguousarray(markers[order]), counts
+
+
+def exchange(comm, markers, destination):
+    """Send every marker to its destination rank (``MPI_Alltoallv``); return the received ones.
+
+    Works for host and device arrays, with or without CUDA-aware MPI
+    (``xp.mpi_buffer`` stages device buffers through the host when needed).
+    """
+    n_ranks = comm.Get_size()
+    width = markers.shape[1]
+    sendbuf, counts = pack_for_ranks(markers, destination, n_ranks)
+    send_counts = np.asarray(xp.to_numpy(counts), dtype=np.int64)
+    recv_counts = np.empty(n_ranks, dtype=np.int64)
+    comm.Alltoall(send_counts, recv_counts)  # how many rows come from each rank
+    received = xp.empty((int(recv_counts.sum()), width), dtype=markers.dtype)
+
+    def displacements(counts):
+        return np.concatenate([[0], np.cumsum(counts)[:-1]])
+
+    with (
+        xp.mpi_buffer(sendbuf) as send,
+        xp.mpi_buffer(received, send=False, recv=True) as recv,
+    ):
+        comm.Alltoallv(
+            [send, send_counts * width, displacements(send_counts) * width, None],
+            [recv, recv_counts * width, displacements(recv_counts) * width, None],
+        )
+    return received
+
+
+def thermal_velocities(seed, particle_ids, step, v_th):
+    """Maxwellian velocities from (seed, particle id, step): no generator state.
+
+    The same numbers as ``cunumpy_normal2(seed, id, step, ...)`` in a kernel
+    (up to the last bits of the math functions), whatever the order of the
+    particles or the number of ranks.
+    """
+    z0, z1 = xp.philox_normal2(seed, particle_ids, step)
+    return v_th * z0, v_th * z1

docs/source/guides/mpi.md

@@ -95,6 +95,47 @@ columns are not (copy them with `xp.ascontiguousarray()` first). No
 synchronization is needed after MPI returns; kernels launched afterwards see
 the received data.
 
+### 5. One call site for both MPI builds: `mpi_buffer`
+
+Code that must also run with an MPI library that is not CUDA-aware (or on
+the NumPy backend) stages device buffers through the host. `mpi_buffer()`
+does the right thing for each case, so the MPI call is written once:
+
+```python
+xp.mpi_is_cuda_aware(comm)  # once at startup; the answer is remembered
+
+with xp.mpi_buffer(send_r) as sendbuf, xp.mpi_buffer(recv_l, send=False, recv=True) as recvbuf:
+    comm.Sendrecv(sendbuf, dest=right, recvbuf=recvbuf, source=left)
+```
+
+A host array is yielded as it is. A device array is yielded as it is (after
+`synchronize_for_mpi`) when MPI is CUDA-aware, and otherwise replaced by a
+pinned host copy: filled from the device before the block when `send=True`,
+copied back into the device array after the block when `recv=True`. The
+copies are counted by `count_transfers()`, so a GPU run with a plain MPI
+build is visible in the transfer report. Without a recorded answer (no
+`mpi_is_cuda_aware()` call and no `set_mpi_cuda_aware()`), a device array
+raises instead of guessing.
+
+## Reproducible random numbers
+
+Each rank needs its own random stream, and a run is reproducible only if every
+draw comes from a seeded generator. Seed `xp.random_streams` once, after MPI is
+initialized, and draw from it everywhere:
+
+```python
+xp.random_streams.seed(config.seed, rank=comm.Get_rank())
+
+positions = xp.random_streams.random((n, 3))
+velocities = xp.random_streams.normal(0.0, v_th, (n, 3))
+rng = xp.random_streams.generator()  # for other distributions
+```
+
+Rank `r` draws the stream `(seed, r)`; the same seed and number of ranks give
+the same results, and different ranks never share numbers. Avoid unseeded
+generators (`np.random.default_rng()` without a seed) anywhere in the time loop:
+one of them makes every run different.
+
 ## Launching
 
 ```bash

docs/source/guides/particle-codes.md

@@ -0,0 +1,186 @@
+# Particle codes
+
+Recipes for the parts of a particle-in-cell (or any particle) code around the
+kernels: removing and sorting markers, depositing without atomics, moving
+markers between MPI ranks, random numbers per particle, output, and running a
+time step as a CUDA graph. The NumPy recipes are the functions of
+`docs/source/examples/particle_recipes.py`, which the test suite runs; each
+works on NumPy and CuPy arrays alike.
+
+## Remove dead markers
+
+Markers leave the domain, hit a wall or are absorbed. Keep a boolean `alive`
+array and compact:
+
+```{literalinclude} ../examples/particle_recipes.py
+:pyobject: remove_dead
+```
+
+To avoid reallocating, keep the markers in a preallocated buffer with a count
+of live rows, and move the live ones to the front; the free rows behind them
+take injected markers:
+
+```{literalinclude} ../examples/particle_recipes.py
+:pyobject: compact_in_place
+```
+
+`alive.sum()` on the GPU is a reduction plus one small transfer (the count);
+do it once per step, not per species and kernel.
+
+## Sort markers by cell
+
+Sorting by cell makes deposits and gathers read and write memory in order,
+which is often worth more than the sort costs, and gives each cell a
+contiguous range of markers, which binary collisions and per-cell diagnostics
+need:
+
+```{literalinclude} ../examples/particle_recipes.py
+:pyobject: sort_by_cell
+```
+
+Apply `order` to every per-marker array (positions, velocities, weights, ids),
+e.g. by keeping them as columns of one `(n, k)` array. A stable sort keeps the
+result independent of how the markers were ordered before.
+
+## Deposit without atomics: sort, then reduce
+
+With markers sorted by cell, a nearest-cell deposit is a segmented sum,
+deterministic on both backends:
+
+```{literalinclude} ../examples/particle_recipes.py
+:pyobject: deposit_nearest_cell
+```
+
+For linear (cloud-in-cell) weights, deposit each of the two (2D: four, 3D:
+eight) neighbours with its weight, one `segment_sum` each.
+
+The two other GPU strategies, as kernels:
+
+* **Global atomics** (`cunumpy_atomic_add` from `<cunumpy/atomic.cuh>`): the
+  simplest; slow when many threads hit the same cells, and the summation order
+  varies between runs (results differ in the last bits).
+* **Per-block shared memory**: each block deposits into a copy of the grid in
+  shared memory, then adds it to the global grid once per cell. Fast for small
+  grids. Check the size with `xp.max_shared_memory_per_block()` and pass
+  `shared_mem=` at the launch; above 48 KiB the kernel is set up for the larger
+  limit automatically.
+
+```c
+extern "C" __global__
+void deposit(Array1D<double> x, Array1D<double> w, Array1D<double> rho,
+             double lower, double dx, int nx) {
+    extern __shared__ double block_rho[];
+    for (int k = threadIdx.x; k < nx; k += blockDim.x) block_rho[k] = 0.0;
+    __syncthreads();
+    long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x;
+    if (i < x.shape[0]) { /* cunumpy_atomic_add(&block_rho[cell], ...) */ }
+    __syncthreads();
+    for (int k = threadIdx.x; k < nx; k += blockDim.x)
+        cunumpy_atomic_add(&rho(k), block_rho[k]);
+}
+```
+
+`cunumpy.testing.emulate_cuda_kernel(..., shared_mem=8 * nx)` runs such a
+kernel on the CPU, barriers included, so it can be checked against the host
+version without a GPU.
+
+## Move markers between MPI ranks
+
+With a spatial decomposition, markers that leave a rank's subdomain move to the
+rank that owns their new position: compute each marker's destination rank,
+group the markers by destination, exchange the counts, then the markers:
+
+```{literalinclude} ../examples/particle_recipes.py
+:pyobject: pack_for_ranks
+```
+
+```{literalinclude} ../examples/particle_recipes.py
+:pyobject: exchange
+```
+
+`xp.mpi_buffer` hands device arrays to MPI directly when it is CUDA-aware and
+stages them through host memory otherwise (see [MPI](mpi.md)). Only the counts
+are host arrays. Remove the markers that left with the compaction above, and
+append the received ones.
+
+## Random numbers per particle
+
+A time loop is reproducible only if every random number comes from a seed.
+Counter-based random numbers make that independent of the order of the
+markers, the launch shape and the number of ranks: particle `id` at step
+`step` always draws the same numbers. On the host:
+
+```{literalinclude} ../examples/particle_recipes.py
+:pyobject: thermal_velocities
+```
+
+and in a kernel, with `<cunumpy/random.cuh>`, the same numbers (uniforms
+bit for bit):
+
+```c
+#include <cunumpy/random.cuh>
+double z0, z1;
+cunumpy_normal2(seed, particle_id, step, &z0, &z1);
+v[i] = v_th * z0;
+```
+
+Use a different counter for every random decision of a step (e.g.
+`4 * step + 0` for injection, `4 * step + 1` for collisions) so that they are
+independent. For draws that need not be per particle (e.g. a collision
+operator's own sampling), `xp.random_streams` gives one seeded generator per
+rank.
+
+## Write output without stalling the GPU
+
+```python
+staging = xp.HostStaging(rho.shape, rho.dtype)  # once
+pending = []
+for step in range(n_steps):
+    advance()
+    if step % output_every == 0:
+        pending.append((step, staging.copy(rho)))  # returns at once
+    while pending and pending[0][1].ready():
+        s, copy = pending.pop(0)
+        h5file[f"rho/{s}"] = copy.result()
+```
+
+`copy()` snapshots the array on the device and copies the snapshot to pinned
+host memory on its own stream, so the next steps overwrite `rho` while the copy
+runs. See `HostStaging` in the [API reference](../api.md).
+
+## Run a time step as a CUDA graph
+
+A step of a small simulation is dozens of short kernels, and the launch
+overhead (a few microseconds each) can dominate. CuPy can capture the launches
+of a step on a stream once and replay them:
+
+```python
+stream = cp.cuda.Stream(non_blocking=True)
+catalog.compile_all()  # compile before capturing: no compilation in a graph
+with stream:
+    stream.begin_capture()
+    step_kernels()  # the kernel launches of one step
+    graph = stream.end_capture()
+for _ in range(n_steps):
+    graph.launch(stream)
+stream.synchronize()
+```
+
+A graph replays exactly the captured launches: the same arrays (by address),
+the same scalar arguments and launch shapes. So allocate every array before
+capturing, keep scalars that change per step in device arrays, and keep host
+synchronization (`.get()`, `float(x)`, `alive.sum()` read on the host, MPI)
+outside the captured part. Debug mode skips its synchronization while a stream
+is capturing.
+
+## Choose the block size from the compiled kernel
+
+```python
+raw = kernel.compile()  # the cupy.RawKernel
+raw.num_regs, raw.max_threads_per_block, raw.shared_size_bytes
+```
+
+A kernel that uses many registers per thread cannot run 1024 threads per block;
+`max_threads_per_block` is the limit for this kernel on this device. Start
+with 128 or 256 threads per block, then time a few sizes on the target GPU
+(`xp.timed_region`) for the kernels that dominate a step.

docs/source/guides/portable-code.md

@@ -130,8 +130,9 @@ of the installed CuPy version. The common differences:
   device array on CuPy. Keep it as an array, or call `float()` and accept the
   synchronization.
 * SciPy functions accept only NumPy arrays; CuPy has its own `cupyx.scipy`.
-  Convert with `to_numpy()` at that boundary, or branch on
-  `xp.get_array_backend(a)`.
+  Use `xp.scipy`, which is the one for the active backend (see
+  [Solvers and fluid updates](solvers.md)); for functions `cupyx.scipy` lacks,
+  convert with `to_numpy()` at that boundary.
 * Plotting, HDF5 and most I/O libraries need host arrays: `to_numpy()` first.
 
 When a function genuinely needs a different implementation per backend, branch