diff --git a/CHANGELOG.md b/CHANGELOG.md index 1b26e2b..9df4299 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/source/api.md b/docs/source/api.md index 38dab05..3def6df 100644 --- a/docs/source/api.md +++ b/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. diff --git a/docs/source/guides/array-ordering.md b/docs/source/guides/array-ordering.md index be572f1..83b3a73 100644 --- a/docs/source/guides/array-ordering.md +++ b/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. diff --git a/docs/source/guides/execution-helpers.md b/docs/source/guides/execution-helpers.md index 8c6083b..07df405 100644 --- a/docs/source/guides/execution-helpers.md +++ b/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 diff --git a/docs/source/kernels/cuda-kernel.md b/docs/source/kernels/cuda-kernel.md index e0b0949..a501a4e 100644 --- a/docs/source/kernels/cuda-kernel.md +++ b/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. diff --git a/docs/source/kernels/dispatch.md b/docs/source/kernels/dispatch.md index 9cd99dd..61d7ce7 100644 --- a/docs/source/kernels/dispatch.md +++ b/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 diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index b0db197..1652609 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -74,7 +74,7 @@ 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)` | @@ -82,7 +82,7 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. | 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")`; `_numba.py`, `_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): diff --git a/src/cunumpy/_algorithms.py b/src/cunumpy/_algorithms.py index ab79486..974e16f 100644 --- a/src/cunumpy/_algorithms.py +++ b/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 diff --git a/src/cunumpy/_cuda_kernel.py b/src/cunumpy/_cuda_kernel.py index 7c3bd8b..80b203d 100644 --- a/src/cunumpy/_cuda_kernel.py +++ b/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 diff --git a/src/cunumpy/_kernel.py b/src/cunumpy/_kernel.py index 735d600..386d772 100644 --- a/src/cunumpy/_kernel.py +++ b/src/cunumpy/_kernel.py @@ -789,7 +789,26 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any: return kernel(*args, **kwargs) -def as_kernel_array(value: Any, like: Any, dtype: Any = None) -> Any: +def _c_ordered_with_gaps(array: np.ndarray) -> bool: + """Whether `array` is laid out in C order, possibly with gaps between entries. + + Every stride is positive and at least the extent of the next axis, e.g. the + first ``n`` columns ``a[:, :n]`` of a C-contiguous ``a`` or every other + entry ``a[::2]``. Axes of length 1 are ignored, as their stride is never used. + """ + extent = array.itemsize + for length, stride in reversed(list(zip(array.shape, array.strides, strict=True))): + if length == 1: + continue + if stride < extent: + return False + extent = stride * length + return True + + +def as_kernel_array( + value: Any, like: Any, dtype: Any = None, *, strided: bool = False +) -> Any: """`value` as an array the kernel chosen for `like` takes. On the device of `like` (a CuPy array if `like` is one, a NumPy array @@ -805,6 +824,15 @@ def as_kernel_array(value: Any, like: Any, dtype: Any = None) -> Any: For an array the kernel writes, use :func:`kernel_output`, which copies a converted array back. + + 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, such as the first ``n`` columns + ``storage[:, :n]`` of a component-major marker buffer, which C-contiguity + would copy on every call. Pyccel's wrappers take such arrays without a copy + (they refuse F order and abort on negative strides, which are still copied); + a host kernel that needs contiguous memory, e.g. one using raw pointers, + must not use it. CuPy arrays are made C-contiguous either way. """ if is_gpu(like): import cupy @@ -819,11 +847,21 @@ def as_kernel_array(value: Any, like: Any, dtype: Any = None) -> Any: nbytes=_nbytes(result), ) return result - return np.ascontiguousarray(to_numpy(value), dtype=dtype) + value = to_numpy(value) + if ( + strided + and value.ndim > 0 + and (dtype is None or value.dtype == np.dtype(dtype)) + and _c_ordered_with_gaps(value) + ): + return value + return np.ascontiguousarray(value, dtype=dtype) @contextmanager -def kernel_output(out: Any, like: Any, dtype: Any = None) -> Generator[Any]: +def kernel_output( + out: Any, like: Any, dtype: Any = None, *, strided: bool = False +) -> Generator[Any]: """The buffer a kernel chosen for `like` writes, copied back into `out` after it. Yields `out` itself if :func:`as_kernel_array` takes it unchanged (then the @@ -833,8 +871,11 @@ def kernel_output(out: Any, like: Any, dtype: Any = None) -> Generator[Any]: with xp.kernels.kernel_output(result, like=grid, dtype=float) as buffer: gather(convert(positions), grid, buffer, ...) + + `strided` is passed on to :func:`as_kernel_array`: with True, a host kernel + writes directly into a NumPy `out` in C order with gaps. """ - buffer = as_kernel_array(out, like, dtype) + buffer = as_kernel_array(out, like, dtype, strided=strided) yield buffer if buffer is not out: if is_gpu(out) and not is_gpu(buffer): diff --git a/tests/unit/test_automatic_launch.py b/tests/unit/test_automatic_launch.py index 4668a18..4ff15b4 100644 --- a/tests/unit/test_automatic_launch.py +++ b/tests/unit/test_automatic_launch.py @@ -60,6 +60,17 @@ def test_callbacks_and_explicit_opt_out(): assert kernel.launch_shape(args=(np.empty((300, 7)),)) == ((3,), (128,)) +def test_last_axis_launches_one_thread_per_marker_of_component_major_arrays(): + kernel = CudaKernel(SOURCE, "work", n_threads_from="last_axis") + # (ncomp, N) positions first: N threads, not ncomp + assert kernel.n_threads_from((2.0, np.empty((3, 1000)), np.empty(5))) == 1000 + assert kernel.launch_shape(args=(np.empty((3, 1000)),)) == ((8,), (128,)) + # a 1D per-marker array, and a 0D array before it is skipped + assert kernel.n_threads_from((np.array(1.0), np.empty(300))) == 300 + with pytest.raises(TypeError, match="needs an array argument"): + kernel.launch_shape(args=(2.0,)) + + def test_invalid_setting_preserves_auto_inference(): kernel = CudaKernel(SOURCE, "work") with pytest.raises(TypeError, match="n_threads_from"): diff --git a/tests/unit/test_compact_by_mask.py b/tests/unit/test_compact_by_mask.py index 39014b9..7796fee 100644 --- a/tests/unit/test_compact_by_mask.py +++ b/tests/unit/test_compact_by_mask.py @@ -39,5 +39,40 @@ def test_matches_boolean_indexing_on_random_masks(): def test_argument_errors(): with pytest.raises(TypeError, match="boolean"): xp.algorithms.compact_by_mask(np.array([1, 0]), np.zeros(2)) - with pytest.raises(ValueError, match="2 rows"): + with pytest.raises(ValueError, match="2 entries along axis 0"): xp.algorithms.compact_by_mask(np.array([True, False]), np.zeros(3)) + + +def test_last_axis_compacts_component_major_and_scalar_arrays_together(): + positions = np.arange(12.0).reshape(3, 4) # (ncomp, N) + weights = np.array([1.0, 2.0, 3.0, 4.0]) # (N,) + alive = np.array([False, True, True, False]) + n = xp.algorithms.compact_by_mask(alive, positions, weights, axis=-1) + assert n == 2 + np.testing.assert_array_equal(positions[:, :n], [[1, 2], [5, 6], [9, 10]]) + np.testing.assert_array_equal(weights[:n], [2.0, 3.0]) + + +def test_axis_matches_boolean_indexing_on_random_masks(): + rng = np.random.default_rng(2) + for axis in (1, -1): + a = rng.random((3, 40)) + mask = rng.random(40) < 0.5 + expected = a[:, mask] + n = xp.algorithms.compact_by_mask(mask, a, axis=axis) + np.testing.assert_array_equal(a[:, :n], expected) + + +def test_compacting_a_column_prefix_view_keeps_the_storage(): + storage = np.arange(20.0).reshape(2, 10) + view = storage[:, :6] # rows further apart than the view is wide + n = xp.algorithms.compact_by_mask(np.arange(6) % 2 == 0, view, axis=-1) + np.testing.assert_array_equal(storage[:, :n], [[0, 2, 4], [10, 12, 14]]) + np.testing.assert_array_equal(storage[:, 6:], [[6, 7, 8, 9], [16, 17, 18, 19]]) + + +def test_axis_errors(): + with pytest.raises(ValueError, match="3 entries along axis -1"): + xp.algorithms.compact_by_mask(np.ones(3, bool), np.zeros((3, 2)), axis=-1) + with pytest.raises(ValueError, match="no axis 1"): + xp.algorithms.compact_by_mask(np.ones(3, bool), np.zeros(3), axis=1) diff --git a/tests/unit/test_kernel_dispatch_arrays.py b/tests/unit/test_kernel_dispatch_arrays.py index a04f14b..0b4bd6a 100644 --- a/tests/unit/test_kernel_dispatch_arrays.py +++ b/tests/unit/test_kernel_dispatch_arrays.py @@ -594,6 +594,53 @@ def test_as_kernel_array_on_the_host(): assert ints.dtype == np.float64 and isinstance(ints, np.ndarray) +@pytest.mark.parametrize( + ("name", "make"), + [ + ("column prefix", lambda: np.zeros((3, 10))[:, :6]), + ("every other entry", lambda: np.zeros(12)[::2]), + ("every other row", lambda: np.zeros((6, 4))[::2]), + ("every other column", lambda: np.zeros((3, 8))[:, ::2]), + ("one row of a prefix", lambda: np.zeros((3, 10))[1:2, :6]), + ], +) +def test_strided_takes_c_ordered_arrays_with_gaps(name, make): + grid = np.zeros(3) + view = make() + assert not view.flags.c_contiguous or name == "one row of a prefix" + assert ( + xp.kernels.as_kernel_array(view, like=grid, dtype=float, strided=True) is view + ) + converted = xp.kernels.as_kernel_array(view, like=grid, dtype=float) + assert converted.flags.c_contiguous # the default still copies + with xp.kernels.kernel_output(view, like=grid, dtype=float, strided=True) as out: + assert out is view + + +@pytest.mark.parametrize( + ("name", "make"), + [ + ("F order", lambda: np.asfortranarray(np.zeros((3, 4)))), + ("transposed", lambda: np.zeros((4, 3)).T), + ("negative stride", lambda: np.zeros(5)[::-1]), + ("negative row stride", lambda: np.zeros((3, 4))[::-1]), + ("other dtype", lambda: np.zeros((3, 10), dtype=np.float32)[:, :6]), + ], +) +def test_strided_still_copies_other_layouts(name, make): + grid = np.zeros(3) + value = make() + converted = xp.kernels.as_kernel_array(value, like=grid, dtype=float, strided=True) + assert converted is not value + assert converted.flags.c_contiguous and converted.dtype == np.float64 + np.testing.assert_array_equal(converted, value) + + +def test_strided_without_dtype_keeps_any_dtype(): + view = np.zeros((2, 8), dtype=np.int32)[:, :5] + assert xp.kernels.as_kernel_array(view, like=np.zeros(1), strided=True) is view + + def test_kernel_output_writes_into_its_target(): grid = np.zeros(3) out = np.zeros(4)