Add strided kernel arrays, last_axis launch size and compact_by_mask axis

CHANGELOG.md

@@ -26,8 +26,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   on the fake CuPy, scalar reads of device arrays. They are in `counter.syncs`
   and the report but not in `total`; `assert_no_transfers(syncs=True)` rejects them.
   The real CuPy's own `float(a)` cannot be observed from Python and is not counted.
-- `xp.algorithms.compact_by_mask(mask, *arrays)` moves the masked rows of arrays to
-  the front, in place and in order, and returns their number.
+- `xp.algorithms.compact_by_mask(mask, *arrays, axis=0)` moves the masked entries
+  of arrays to the front along `axis`, in place and in order, and returns their
+  number; `axis=-1` compacts component-major `(ncomp, N)` and `(N,)` arrays together.
+- `as_kernel_array(..., strided=True)` and `kernel_output(..., strided=True)` take a
+  NumPy array in C order with gaps (positive strides, each at least the extent of
+  the next axis, e.g. `storage[:, :n]`) unchanged for host kernels, instead of
+  copying it to a C-contiguous array; Pyccel's wrappers take such arrays.
+- `CudaKernel(n_threads_from="last_axis")`: one thread per entry of the last axis of
+  the first array argument, for component-major `(ncomp, N)` marker arrays.
 - Kernel outputs: a name in `PyccelKernel(outputs=...)` also finds a positional
   argument and an index a keyword argument, using the parameter names of the
   function (or the new `parameters=`); a `Kernel` supplies those of its host

docs/source/api.md

@@ -271,7 +271,7 @@ keys, order, positions, charges = xp.algorithms.sort_by_key(keys, positions, cha
 Returns `(keys[order], order, *(a[order] for a in arrays))`, `order` as
 `int64`. Equal keys keep their order, so the result is reproducible.
 
-### `algorithms.compact_by_mask(mask, *arrays)`
+### `algorithms.compact_by_mask(mask, *arrays, axis=0)`
 
 Moves the rows where the boolean `mask` is True to the front of every array, in
 place and in their original order, and returns how many there are. Typical use:
@@ -284,7 +284,14 @@ markers, weights = markers[:n], weights[:n]
 
 The rows after the first `n` are unspecified. The count is needed on the host,
 so on CuPy each call synchronizes once. The mask and the arrays must be on the
-same backend.
+same backend. `axis` is the axis of every array that the mask indexes; `-1`
+compacts component-major arrays, `(ncomp, N)` next to `(N,)`, along their
+marker axis:
+
+```python
+n = xp.algorithms.compact_by_mask(alive, positions, weights, axis=-1)
+positions, weights = positions[:, :n], weights[:n]
+```
 
 ## Count transfers
 
@@ -400,7 +407,7 @@ class DeviceParticles(xp.arguments.CudaArguments):
         super().__init__(self.markers, self.degree, self.markers.shape[0])
 ```
 
-### `kernels.as_kernel_array(value, like, dtype=None)`, `kernels.kernel_output(out, like, dtype=None)`
+### `kernels.as_kernel_array(value, like, dtype=None, *, strided=False)`, `kernels.kernel_output(out, like, dtype=None, *, strided=False)`
 
 For the arguments of a `Kernel` with `dispatch="arrays"`, whose choice follows
 the arrays. `as_kernel_array` returns `value` on the side of `like` (a CuPy
@@ -412,6 +419,13 @@ itself if `as_kernel_array` takes it unchanged, else a converted copy whose
 contents are written into `out` (on its own side) when the block ends without
 an error.
 
+With `strided=True` a NumPy array for a host kernel is also taken unchanged
+when it is in C order with gaps (positive strides, each at least the extent of
+the next axis), e.g. the first `n` columns `storage[:, :n]` of a marker buffer.
+Pyccel's wrappers take such arrays without a copy; a host kernel that needs
+contiguous memory must not use it. F-ordered, transposed and negative-stride
+arrays are still copied, and CuPy arrays are always made C-contiguous.
+
 ```python
 with xp.kernels.kernel_output(result, like=field, dtype=float) as buffer:
     gather(
@@ -1096,7 +1110,8 @@ Explicit `n_threads` or `grid` always takes precedence. Set `n_threads_from`
 (constructor argument or settable property) to a callable such as
 `lambda args: args[0].size` for flattened element kernels, or
 `lambda args: args[0].shape[::-1]` for kernels whose x index follows columns.
-`"first_array"` always uses the first axis; None disables inference and requires
+`"first_array"` always uses the first axis and `"last_axis"` the last one (one
+thread per entry of a component-major `(ncomp, N)` or `(N,)` array); None disables inference and requires
 explicit launch sizes. The same defaults apply through `kernels.Kernel` and
 `kernel_testing.assert_kernels_agree`, and in CPU emulation.
 

docs/source/guides/array-ordering.md

@@ -84,6 +84,12 @@ markers[...] = markers.copy()  # C-ordered temporary, F-ordered destination
 assert markers.flags.f_contiguous
 ```
 
+Component-major arrays, `(ncomp, N)` with the markers along the last axis, keep
+a C-ordered buffer and a contiguous `positions[d]` per component without any
+order flag; `compact_by_mask(alive, positions, weights, axis=-1)` compacts them
+in place, and the prefix `positions[:, :n]` is C order with gaps, which Pyccel
+kernels take as it is (`as_kernel_array(..., strided=True)`).
+
 Likewise, `xp.algorithms.compact_by_mask(alive, markers)` gathers selected rows
 and writes them into the front of the original buffer, preserving that
 buffer's layout. Only the returned count of leading rows is valid. However,
@@ -120,8 +126,14 @@ Some convenience helpers deliberately require C order:
 | Helper | Contract |
 | --- | --- |
 | `xp.as_device_array` | Returns a C-contiguous device array, copying if needed. |
-| `xp.kernels.as_kernel_array` | Returns a C-contiguous array on the side of `like`. |
-| `xp.kernels.kernel_output` | Yields a C-contiguous working buffer and copies updates back into the original output when needed. |
+| `xp.kernels.as_kernel_array` | Returns a C-contiguous array on the side of `like`; with `strided=True`, also a NumPy array in C order with gaps, unchanged. |
+| `xp.kernels.kernel_output` | Yields a C-contiguous working buffer (with `strided=True`, also a NumPy output in C order with gaps) and copies updates back into the original output when needed. |
+
+"C order with gaps" means positive strides, each at least the extent of the next
+axis: `a[:, :n]`, `a[::2]` or `a[:, ::2]` of a C-contiguous `a`. Pyccel's
+wrappers take such arrays without a copy. They refuse F-ordered arrays with a
+`TypeError` and abort the process on negative strides, so `strided=True` still
+copies those.
 
 For an F-order kernel, allocate or normalize its arguments explicitly instead
 of using these helpers to prepare F-ordered inputs.

docs/source/guides/execution-helpers.md

@@ -157,7 +157,8 @@ markers, weights = markers[:n], weights[:n]
 Moves the rows where the boolean mask is True to the front of every array, in
 place and in order, and returns their number. The rows after the first `n` are
 unspecified. The count is needed on the host, so on CuPy each call
-synchronizes once.
+synchronizes once. For component-major arrays, `(ncomp, N)` next to `(N,)`,
+pass `axis=-1` and continue with `positions[:, :n]`.
 
 ## Prepare cell ranges and segment reductions
 

docs/source/kernels/cuda-kernel.md

@@ -99,7 +99,8 @@ axpy(2.0, x, y, x.size)  # infer n_threads = x.shape[0]
 Explicit sizes override inference. Set `n_threads_from=lambda args: args[0].size`
 for a flattened element kernel, or another callback for a different axis order
 or logical work count. `n_threads_from=None` requires explicit sizes;
-`"first_array"` always selects only the first axis. `launch_shape(args=(...))`
+`"first_array"` always selects only the first axis, `"last_axis"` only the last
+one (one thread per marker of component-major `(ncomp, N)` arrays). `launch_shape(args=(...))`
 inspects the inferred grid/block without compilation or a launch. CPU emulation
 and host/CUDA parity tests use the same defaults.
 

docs/source/kernels/dispatch.md

@@ -279,6 +279,11 @@ with xp.kernels.kernel_output(result, like=field, dtype=float) as buffer:
     catalog["gather"](convert(positions), convert(field), buffer, n_threads=n)
 ```
 
+Pyccel host kernels also take arrays in C order with gaps, such as the stored
+markers `storage[:, :n]` of a component-major buffer with spare capacity. Pass
+`strided=True` to both helpers to hand those over without a copy on the host
+(see [Array ordering and strides](../guides/array-ordering.md)).
+
 ## Same parameters on both sides
 
 The call site is the same for both kernels only if they take the same

src/cunumpy/LLM_GUIDE.md

@@ -74,15 +74,15 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository.
 | launch a hand-written CUDA C kernel | `xp.kernels.CudaKernel(source, "name")` / `CudaKernel.from_file(path)` |
 | run a Metal (MSL) kernel on an Apple silicon GPU, NumPy float32 in and out (float64 raises; `float64="cast"` computes in float32); not part of `Kernel` dispatch | `xp.kernels.MetalKernel(body, inputs=[...], outputs=[...])(*args, out=arrays, n_threads=n)`; check `xp.kernels.metal_available()` |
 | call host-only code (SciPy, file readers) with arguments of either backend | `xp.host_call(fun, *args)`; `@xp.evaluate_on_host` on a method; `@xp.setup_on_host` on `__init__` |
-| keep the live rows of particle arrays at the front | `n = xp.algorithms.compact_by_mask(alive, markers, weights)` |
+| keep the live rows of particle arrays at the front | `n = xp.algorithms.compact_by_mask(alive, markers, weights)`; component-major `(ncomp, N)`: `axis=-1` |
 | run CUDA kernel launches on the CPU in a test (fake CuPy) | `with kernel_testing.emulated_launches(): ...`; `kernel_testing.host_buffer(a)` reads a fake array; struct arguments are read through their fields |
 | find the arguments a host kernel writes (copy only those back) | `PyccelKernel(fn, outputs=("out",))` (names work positionally); `xp.kernels.outputs_from_annotations(fn)`; `Kernel.from_folder(..., outputs="annotations")` |
 | host kernel + CUDA port, chosen by backend | `xp.kernels.Kernel(host_fn, cuda_kernel_or_None)` |
 | many kernels in a package, ported incrementally | `xp.kernels.KernelCatalog.from_package(__name__, missing_cuda="fallback")` |
 | host kernels compiled at first call (your compile function), NumPy fallback | `from_package(..., host_suffix="_pyccel", compile_host=my_compile, host_fallback={...})` -> `xp.kernels.CompiledHostKernel` |
 | host arrays reach kernels while CuPy is active | `Kernel(..., dispatch="arrays")` / `from_package(..., dispatch="arrays")`: CUDA only for device arguments |
 | one kernel folder declares its kernel in its own `__init__.py` | `kernel = xp.kernels.Kernel.from_folder(__name__, host_suffix="_pyccel", compile_host=..., dispatch="arrays")`; `<name>_numba.py`, `<name>_numpy.py` in the folder are further host implementations |
-| bring a `dispatch="arrays"` kernel's arguments to the side of the main array | `xp.kernels.as_kernel_array(a, like=grid, dtype=float)`; outputs: `with xp.kernels.kernel_output(out, like=grid, dtype=float) as buf:` |
+| bring a `dispatch="arrays"` kernel's arguments to the side of the main array | `xp.kernels.as_kernel_array(a, like=grid, dtype=float)`; outputs: `with xp.kernels.kernel_output(out, like=grid, dtype=float) as buf:`; `strided=True` passes NumPy views in C order with gaps (`a[:, :n]`) to Pyccel hosts without a copy |
 | choose the host implementation (pyccel/numba/numpy/python) | `xp.kernels.set_host_kernel_implementation("numpy")`, `with xp.kernels.use_host_kernel_implementation(...)`, `CUNUMPY_HOST_KERNEL_IMPLEMENTATION=numpy`; default: first available of pyccel, numba, numpy; `kernel.selected()` |
 | require CUDA for device kernel dispatch | `xp.kernels.set_device_kernel_implementation("cuda")`, `get_device_kernel_implementation()`, `with xp.kernels.use_device_kernel_implementation(...)`, `CUNUMPY_DEVICE_KERNEL_IMPLEMENTATION=cuda`; default `None` preserves `missing_cuda` policy; explicit CUDA rejects host fallback |
 | check host and CUDA kernels take the same parameters | `catalog.check_signatures()` (in a unit test) |
@@ -305,8 +305,10 @@ xp.cuda.cuda_include_dir()
 * Default `n_threads_from="auto"` infers from the first array, including arrays
   in supported argument objects: 1D block -> shape[0], 2D -> shape[:2], 3D ->
   shape[:3] (x, y, z). Per-call block overrides apply. Explicit n_threads/grid
-  wins; use a callback for flattened `.size`, reversed axes or another work
-  count, or None to require explicit sizes. Missing arrays/axes raise.
+  wins; `"last_axis"` uses the first array's last axis (component-major
+  `(ncomp, N)` marker arrays); use a callback for flattened `.size`, reversed
+  axes or another work count, or None to require explicit sizes. Missing
+  arrays/axes raise.
 
 Shipped CUDA headers (always on the include path):
 

src/cunumpy/_algorithms.py

@@ -253,30 +253,40 @@ def sort_by_key(keys: Any, *arrays: Any) -> tuple[Any, ...]:
     return (keys[order], order, *(array[order] for array in arrays))
 
 
-def compact_by_mask(mask: Any, *arrays: Any) -> int:
-    """Move the rows where `mask` is True to the front of every array, in place.
+def compact_by_mask(mask: Any, *arrays: Any, axis: int = 0) -> int:
+    """Move the entries where `mask` is True to the front of every array, in place.
 
     The usual step after particles left the domain or were absorbed: keep the
     live ones at the front of the marker array (and of the arrays that go with
-    it) and continue with ``markers[:n]``. The order of the kept rows is
+    it) and continue with ``markers[:n]``. The order of the kept entries is
     preserved, so the result is reproducible::
 
         n = xp.algorithms.compact_by_mask(alive, markers, weights)
         markers, weights = markers[:n], weights[:n]
 
+    Component-major marker arrays, ``(ncomp, N)`` next to ``(N,)`` scalars, keep
+    the markers along the last axis of each, so compact that one::
+
+        n = xp.algorithms.compact_by_mask(alive, positions, weights, axis=-1)
+        positions, weights = positions[:, :n], weights[:n]
+
     Parameters
     ----------
     mask : array of bool, shape (n,)
-        True for the rows to keep.
+        True for the entries to keep.
     *arrays : arrays
-        Arrays with ``n`` rows (any further axes), on the backend of `mask`.
-        Rows ``[:count]`` hold the kept rows afterwards; the rows after them
-        are unspecified, so ignore them (or overwrite them).
+        Arrays with ``n`` entries along `axis` (any other axes), on the backend
+        of `mask`. Entries ``[:count]`` along `axis` hold the kept ones
+        afterwards; the entries after them are unspecified, so ignore them (or
+        overwrite them).
+    axis : int
+        The axis of every array that `mask` indexes, the first by default;
+        ``-1`` is the last axis of each array, whatever its number of axes.
 
     Returns
     -------
     int
-        The number of kept rows. Its value is needed on the host, so on CuPy
+        The number of kept entries. Its value is needed on the host, so on CuPy
         the call synchronizes once per call (counted by
         :func:`~cunumpy.profiling.count_transfers` where it can be seen).
     """
@@ -287,13 +297,23 @@ def compact_by_mask(mask: Any, *arrays: Any) -> int:
         raise TypeError(
             f"mask must be a 1D boolean array, got dtype {mask.dtype}, {mask.ndim}D",
         )
+    axes = []
     for i, array in enumerate(arrays):
-        if array.ndim < 1 or array.shape[0] != mask.shape[0]:
+        if not -array.ndim <= axis < array.ndim:
             raise ValueError(
-                f"array {i} has shape {array.shape}, expected {mask.shape[0]} rows",
+                f"array {i} has {array.ndim} axes, so it has no axis {axis}",
             )
+        array_axis = axis % array.ndim
+        if array.shape[array_axis] != mask.shape[0]:
+            raise ValueError(
+                f"array {i} has shape {array.shape}, expected {mask.shape[0]} "
+                f"entries along axis {axis}",
+            )
+        axes.append(array_axis)
     rows = xpm.nonzero(mask)[0]
     n_kept = int(rows.size)
-    for array in arrays:
-        array[:n_kept] = array[rows]  # the right side is a copy: no overlap problem
+    for array, array_axis in zip(arrays, axes, strict=True):
+        front = (slice(None),) * array_axis + (slice(0, n_kept),)
+        # the right side is a copy: no overlap problem
+        array[front] = xpm.take(array, rows, axis=array_axis)
     return n_kept

src/cunumpy/_cuda_kernel.py

@@ -1797,6 +1797,20 @@ def _first_array_length(args: tuple[Any, ...]) -> int:
     return typing.cast(int, _first_array_shape(args, 1))
 
 
+def _last_axis_length(args: tuple[Any, ...]) -> int:
+    """``n_threads_from="last_axis"``: the last axis of the first array argument.
+
+    One thread per entry of a component-major array, ``(ncomp, N)`` or ``(N,)``.
+    """
+    shape = next(_array_shapes_in(args), None)
+    if shape is None:
+        raise TypeError(
+            "inferring n_threads needs an array argument; pass n_threads or grid, "
+            "or set n_threads_from to a callable",
+        )
+    return int(shape[-1])
+
+
 def _as_shape(value: int | Sequence[int], what: str) -> tuple[int, ...]:
     shape = (value,) if isinstance(value, (int, np.integer)) else tuple(value)
     if not 1 <= len(shape) <= 3:
@@ -1876,11 +1890,12 @@ class CudaKernel:
         (:func:`cunumpy.cuda.set_cuda_debug`, ``CUNUMPY_CUDA_DEBUG``) at every
         launch; True or False fix it for this kernel. The compile options are
         fixed when the kernel is compiled.
-    n_threads_from : {"auto", "first_array"} | callable | None
+    n_threads_from : {"auto", "first_array", "last_axis"} | callable | None
         Default "auto" infers thread counts from the first array's leading
         shape axes, matching the block dimensionality (1D: one thread per row).
         Arrays in supported argument objects are included. "first_array"
-        always uses the first axis. A callable receives the positional argument
+        always uses the first axis, "last_axis" the last one (one thread per
+        entry of a component-major ``(ncomp, N)`` or ``(N,)`` array). A callable receives the positional argument
         tuple; None requires an explicit launch size. Explicit `n_threads` or
         `grid` overrides inference.
 
@@ -2160,7 +2175,9 @@ def n_threads_from(self) -> Callable[[tuple[Any, ...]], Any] | None:
         Default ``"auto"`` uses the first array's leading axes, matching the
         launch block dimensions; 1D launches use its first axis, one thread per
         row. Supported argument objects are searched in field/argument order.
-        ``"first_array"`` always uses its first axis. None disables inference.
+        ``"first_array"`` always uses its first axis, ``"last_axis"`` its last
+        one (one thread per entry of a component-major ``(ncomp, N)`` or
+        ``(N,)`` array). None disables inference.
         Settable, also on the ``cuda_kernel`` of a :class:`~cunumpy.kernels.Kernel`.
         """
         return self._n_threads_from
@@ -2175,9 +2192,12 @@ def n_threads_from(
             value = self._default_n_threads
         elif isinstance(value, str) and value == "first_array":
             value = _first_array_length
+        elif isinstance(value, str) and value == "last_axis":
+            value = _last_axis_length
         if value is not None and not callable(value):
             raise TypeError(
-                "n_threads_from must be callable, 'auto', 'first_array' or None"
+                "n_threads_from must be callable, 'auto', 'first_array', "
+                "'last_axis' or None"
             )
         self._automatic_threads = automatic
         self._n_threads_from = value