diff --git a/CHANGELOG.md b/CHANGELOG.md index 25ab0ce..0baab30 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -29,6 +29,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - `KernelCatalog.from_package(..., include_dirs=None)` is now an explicit keyword; by default the source root of the top-level package (the directory containing it) is an include directory of every CUDA kernel, in addition to the kernel's own folder, so kernels can `#include "my_pkg/common.cuh"`. ### Added +- `cunumpy/morton.cuh` and `xp.morton_keys`, `morton_encode`, `morton_decode`, `morton_scales`, `MAX_MORTON_LEVELS`: Morton (Z-order) keys of 2D and 3D points (`uint64`, up to 32 and 21 bits per axis), the same in a kernel and on the host (bit for bit), for sorting particles along a space-filling curve and building quadtrees and octrees from sorted keys. +- `xp.sort_by_key(keys, *arrays)`: one stable argsort of `keys` applied to every array; returns the sorted keys, the order and the sorted arrays. - `Kernel.from_folder(package, ...)`: the kernel of one kernel folder, with the options of `KernelCatalog.from_package` (`host_suffix`, `compile_host`, `dispatch`, `include_dirs`, CUDA options...). Every version in the folder is an implementation: `.py` (`"pyccel"`, compiled with `compile_host`, and `"python"`, uncompiled), `_numba.py`, `_numpy.py` and `_cuda.cu`. A folder's own `__init__.py` can declare `kernel = xp.Kernel.from_folder(__name__, ...)`, so code imports the kernel from where it is written; `from_package` now builds each kernel with it. `Kernel.implementations` lists them, `Kernel.selected(device=False)` tells which one a call runs. - `xp.HostImplementations` and `xp.HOST_IMPLEMENTATIONS`: the host implementations of a kernel, loaded on first use; a call runs the default, the first available of pyccel, numba and NumPy (the uncompiled Python version, with a warning, if none is), or the one chosen with `xp.set_kernel_implementation(name)` / `with xp.use_kernel_implementation(name):` / `CUNUMPY_KERNEL_IMPLEMENTATION=name` (read at import), like `set_backend`/`use_backend`/`ARRAY_BACKEND`; a chosen implementation that a kernel lacks or cannot load raises instead of running another. `xp.get_kernel_implementation()` reads the setting. - `xp.as_kernel_array(value, like, dtype=None)` and `xp.kernel_output(out, like, dtype=None)`: bring the arguments of a `dispatch="arrays"` kernel to the side of the main array `like` (CuPy or NumPy, C-contiguous, `dtype`), without a copy when they already fit; `kernel_output` yields the buffer the kernel writes and copies it back into `out` if it had to be converted. diff --git a/docs/source/api.md b/docs/source/api.md index 294acbb..9fdffdf 100644 --- a/docs/source/api.md +++ b/docs/source/api.md @@ -189,6 +189,18 @@ separately); a negative key drops the value; keys must be smaller than `n_segments`. The result keeps a floating-point or complex dtype and is `float64` otherwise. +### `sort_by_key(keys, *arrays)` + +Stable argsort of the 1D `keys` (CuPy's radix sort on the device), applied to +every array along axis 0, in one call: + +```python +keys, order, positions, charges = xp.sort_by_key(keys, positions, charges) +``` + +Returns `(keys[order], order, *(a[order] for a in arrays))`, `order` as +`int64`. Equal keys keep their order, so the result is reproducible. + ## Count transfers A transfer inside a time loop is the classic performance bug of a GPU port: @@ -1892,6 +1904,42 @@ normal numbers can differ in the last bits (`log`, `sqrt`, `sin`, `cos` on the GPU are not the host's). The generator passes the Random123 known-answer tests. Use a different `counter` for every random decision of a step. +### `cunumpy/morton.cuh` and `morton_keys` + +Morton (Z-order) keys: the bits of a point's integer cell coordinates, +interleaved into one `uint64`. Sorted by key, nearby points are nearby in +memory, and the points of every node of a quadtree (2D) or octree (3D) on the +same box form a contiguous range, the starting point of tree builds on the +GPU. + +```python +keys = xp.morton_keys(positions, lower, upper, levels) # (n, 2|3) -> (n,) uint64 +keys, order, positions = xp.sort_by_key(keys, positions) +node = keys >> np.uint64(ndim * (levels - level)) # node index at `level` +cells = xp.morton_decode(node, ndim) # its integer coordinates +key = xp.morton_encode(ix, iy) # from integer cells +scales = xp.morton_scales(lower, upper, levels) # 2**levels / (upper - lower) +``` + +```c +#include + +unsigned long long cunumpy_morton_key2(x, y, lower_x, lower_y, scale_x, scale_y, levels); +unsigned long long cunumpy_morton_key3(x, y, z, lower_x, ..., scale_x, ..., levels); +unsigned long long cunumpy_morton_encode2(ix, iy); // and _encode3(ix, iy, iz) +unsigned long long cunumpy_morton_cell(x, lower, scale, levels); +unsigned long long cunumpy_morton_spread2(v); // and _compact2, _spread3, _compact3 +``` + +`levels` is the number of bits per axis, at most 32 in 2D and 21 in 3D +(`xp.MAX_MORTON_LEVELS`). Axis 0 is the lowest bit of every group of `ndim` +bits; the top group is the child of the root. The cell along an axis is +`floor((x - lower) * scale)` clipped to `[0, 2**levels - 1]`: points on a cell +boundary go to the upper cell, points outside the box to the nearest face, and +`lower > upper` reverses the axis. Given the `morton_scales` of the host, the +kernel functions return the host keys bit for bit. All host functions run on +NumPy and CuPy arrays. + ### `cunumpy/reduce.cuh` Warp- and block-level reductions for hand-written kernels: in-kernel diff --git a/docs/source/guides/particle-codes.md b/docs/source/guides/particle-codes.md index 584ecbc..660a25a 100644 --- a/docs/source/guides/particle-codes.md +++ b/docs/source/guides/particle-codes.md @@ -41,6 +41,12 @@ need: 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. +`xp.sort_by_key(keys, positions, velocities, weights)` does the argsort and +the reordering of several arrays in one call. + +For a tree code, or for better locality in 2D and 3D, sort by Morton key +(`xp.morton_keys`, see the API page) instead of by cell: the markers of every +quadtree or octree node are then a contiguous range of the sorted arrays. ## Deposit without atomics: sort, then reduce diff --git a/docs/source/kernels/cuda-kernel.md b/docs/source/kernels/cuda-kernel.md index b109f6f..7687f97 100644 --- a/docs/source/kernels/cuda-kernel.md +++ b/docs/source/kernels/cuda-kernel.md @@ -156,6 +156,7 @@ Pass extra include directories with `include_dirs=[...]` and NVRTC flags with | `` | `CUNUMPY_THREAD_1D(i, n)`, `_2D`, `_3D`, `CUNUMPY_GRID_STRIDE_1D(i, n)` | | `` | strided views `Array1D` to `Array4D` | | `` | `cunumpy_atomic_add` and indexed 2D/3D variants, see [Accumulation kernels](accumulation.md) | +| `` | Morton (Z-order) keys `cunumpy_morton_key2(x, y, ...)`, `_key3`, equal to `xp.morton_keys` on the host | | `` | counter-based random numbers `cunumpy_uniform(seed, stream, counter)`, `cunumpy_normal2(...)`, equal to `xp.philox_uniform` on the host | `xp.cuda_include_dir()` returns their directory for use with other compilers. diff --git a/src/cunumpy/LLM_GUIDE.md b/src/cunumpy/LLM_GUIDE.md index 54b8e41..cf61288 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -75,6 +75,7 @@ https://max-models.github.io/cunumpy/ and in `docs/source/` of the repository. | test a CUDA kernel's arithmetic without a GPU | `cunumpy.testing.emulate_cuda_kernel(kernel, *numpy_args, n_threads=n)` (C++ compiler; shared memory and __syncthreads ok, no warp ops; `shared_mem=` for extern shared) | | shared-memory budget of a block | `xp.max_shared_memory_per_block()` (48 KiB without a GPU) | | random numbers inside a kernel, equal on the host | `#include `: `cunumpy_uniform(seed, particle_id, step)`; host: `xp.philox_uniform(seed, ids, step)` | +| sort points along a Z-curve / quadtree or octree nodes as contiguous ranges | `keys = xp.morton_keys(pos, lower, upper, levels)`, `keys, order, pos = xp.sort_by_key(keys, pos)`; in a kernel `#include `: `cunumpy_morton_key2(x, y, x0, y0, sx, sy, levels)` with `xp.morton_scales(...)` | | one thread per marker without passing n_threads | `CudaKernel(..., n_threads_from="first_array")` | | copy device arrays to the host for output without stalling | `xp.HostStaging(shape, dtype)`: `c = staging.copy(a)` ... `c.result()` | | PIC recipes (compaction, sort by cell, MPI exchange, graphs) | docs guide "Particle codes" | @@ -138,6 +139,7 @@ with xp.mpi_buffer(a) as buf: comm.Send(buf, ...) # host array, CUDA- with xp.mpi_buffer(a, send=False, recv=True) as buf: ... # array, or pinned staging copy xp.set_mpi_cuda_aware(True | False | None), xp.get_mpi_cuda_aware() xp.segment_sum(values, keys, n_segments) # out[k] = sum(values[keys == k]); keys < 0 dropped +keys, order, a, b = xp.sort_by_key(keys, a, b) # stable argsort applied to every array xp.require_version("0.4.0") # ImportError if cunumpy is older ``` diff --git a/src/cunumpy/__init__.py b/src/cunumpy/__init__.py index 5206d75..40ac479 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -35,6 +35,13 @@ use_kernel_implementation, ) from .mirror import DeviceMirror +from .morton import ( + MAX_MORTON_LEVELS, + morton_decode, + morton_encode, + morton_keys, + morton_scales, +) from .petsc import petsc_vec from .philox import ( philox4x32_10, @@ -88,6 +95,7 @@ set_device, set_device_for_rank, set_mpi_cuda_aware, + sort_by_key, stream, synchronize, synchronize_for_mpi, @@ -135,6 +143,7 @@ def require_version(minimum: str) -> None: "DEBUG_OPTIONS", "DEFAULT_SHARED_MEMORY_PER_BLOCK", "HOST_IMPLEMENTATIONS", + "MAX_MORTON_LEVELS", "CompiledHostKernel", "CudaArguments", "CudaKernel", @@ -187,6 +196,10 @@ def require_version(minimum: str) -> None: "local_rank", "max_shared_memory_per_block", "memory_info", + "morton_decode", + "morton_encode", + "morton_keys", + "morton_scales", "mpi_buffer", "mpi_is_cuda_aware", "numpy_backend", @@ -213,6 +226,7 @@ def require_version(minimum: str) -> None: "set_device_for_rank", "set_kernel_implementation", "set_mpi_cuda_aware", + "sort_by_key", "stream", "synchronize", "synchronize_for_mpi", diff --git a/src/cunumpy/__init__.pyi b/src/cunumpy/__init__.pyi index 5de4eec..c5b0239 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -32,6 +32,11 @@ from .kernel import get_kernel_implementation as get_kernel_implementation from .kernel import set_kernel_implementation as set_kernel_implementation from .kernel import use_kernel_implementation as use_kernel_implementation from .mirror import DeviceMirror as DeviceMirror +from .morton import MAX_MORTON_LEVELS as MAX_MORTON_LEVELS +from .morton import morton_decode as morton_decode +from .morton import morton_encode as morton_encode +from .morton import morton_keys as morton_keys +from .morton import morton_scales as morton_scales from .petsc import petsc_vec as petsc_vec from .philox import philox4x32_10 as philox4x32_10 from .philox import philox_normal as philox_normal @@ -74,6 +79,7 @@ def mpi_buffer( array: Any, *, send: bool = ..., recv: bool = ..., cuda_aware: bool | None = ... ) -> Generator[Any]: ... def segment_sum(values: Any, keys: Any, n_segments: int) -> Any: ... +def sort_by_key(keys: Any, *arrays: Any) -> tuple[Any, ...]: ... def as_kernel_array(value: Any, like: Any, dtype: Any = ...) -> Any: ... @contextmanager def kernel_output(out: Any, like: Any, dtype: Any = ...) -> Generator[Any]: ... diff --git a/src/cunumpy/cuda/include/cunumpy/morton.cuh b/src/cunumpy/cuda/include/cunumpy/morton.cuh new file mode 100644 index 0000000..82feb74 --- /dev/null +++ b/src/cunumpy/cuda/include/cunumpy/morton.cuh @@ -0,0 +1,128 @@ +// cunumpy/morton.cuh: Morton (Z-order) keys for kernels. +// +// A Morton key interleaves the bits of the integer cell coordinates of a +// point. Sorting points by their keys orders them along a Z-shaped curve, and +// the points of every quadtree/octree node on the same box form a contiguous +// range of the sorted array. The keys are equal to those of the host function +// cunumpy.morton_keys when the kernel gets the same lower corner and the +// scales of cunumpy.morton_scales(lower, upper, levels): +// +// #include +// +// extern "C" __global__ void keys2d(const double* pos, unsigned long long* key, +// long long n, double x0, double y0, +// double sx, double sy, int levels) { +// long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; +// if (i >= n) return; +// key[i] = cunumpy_morton_key2(pos[2 * i], pos[2 * i + 1], +// x0, y0, sx, sy, levels); +// } +// +// With `levels` bits per axis a key has ndim * levels bits; axis 0 is the +// lowest bit of every group of ndim bits, and the top group is the child of +// the root. The cell along an axis is floor((x - lower) * scale), clipped to +// [0, 2^levels - 1]: the same operations, in the same order, as on the host, +// so the keys are bit-identical. Positions must be finite. +// +// Plain integer and double arithmetic only, so the header also compiles as +// C++ (see cunumpy.testing.emulate_cuda_kernel). + +#ifndef CUNUMPY_MORTON_CUH +#define CUNUMPY_MORTON_CUH + +// Spread the low 32 bits of v to the even bits of a 64-bit word. +__host__ __device__ __forceinline__ unsigned long long cunumpy_morton_spread2( + unsigned long long v) +{ + v &= 0xFFFFFFFFull; + v = (v | (v << 16)) & 0x0000FFFF0000FFFFull; + v = (v | (v << 8)) & 0x00FF00FF00FF00FFull; + v = (v | (v << 4)) & 0x0F0F0F0F0F0F0F0Full; + v = (v | (v << 2)) & 0x3333333333333333ull; + v = (v | (v << 1)) & 0x5555555555555555ull; + return v; +} + +// Inverse of cunumpy_morton_spread2: gather the even bits of v. +__host__ __device__ __forceinline__ unsigned long long cunumpy_morton_compact2( + unsigned long long v) +{ + v &= 0x5555555555555555ull; + v = (v ^ (v >> 1)) & 0x3333333333333333ull; + v = (v ^ (v >> 2)) & 0x0F0F0F0F0F0F0F0Full; + v = (v ^ (v >> 4)) & 0x00FF00FF00FF00FFull; + v = (v ^ (v >> 8)) & 0x0000FFFF0000FFFFull; + v = (v ^ (v >> 16)) & 0xFFFFFFFFull; + return v; +} + +// Spread the low 21 bits of v to every third bit of a 64-bit word. +__host__ __device__ __forceinline__ unsigned long long cunumpy_morton_spread3( + unsigned long long v) +{ + v &= 0x1FFFFFull; + v = (v | (v << 32)) & 0x001F00000000FFFFull; + v = (v | (v << 16)) & 0x001F0000FF0000FFull; + v = (v | (v << 8)) & 0x100F00F00F00F00Full; + v = (v | (v << 4)) & 0x10C30C30C30C30C3ull; + v = (v | (v << 2)) & 0x1249249249249249ull; + return v; +} + +// Inverse of cunumpy_morton_spread3. +__host__ __device__ __forceinline__ unsigned long long cunumpy_morton_compact3( + unsigned long long v) +{ + v &= 0x1249249249249249ull; + v = (v ^ (v >> 2)) & 0x10C30C30C30C30C3ull; + v = (v ^ (v >> 4)) & 0x100F00F00F00F00Full; + v = (v ^ (v >> 8)) & 0x001F0000FF0000FFull; + v = (v ^ (v >> 16)) & 0x001F00000000FFFFull; + v = (v ^ (v >> 32)) & 0x1FFFFFull; + return v; +} + +// Key of integer cells (ix, iy) / (ix, iy, iz), as cunumpy.morton_encode. +__host__ __device__ __forceinline__ unsigned long long cunumpy_morton_encode2( + unsigned long long ix, unsigned long long iy) +{ + return cunumpy_morton_spread2(ix) | (cunumpy_morton_spread2(iy) << 1); +} + +__host__ __device__ __forceinline__ unsigned long long cunumpy_morton_encode3( + unsigned long long ix, unsigned long long iy, unsigned long long iz) +{ + return cunumpy_morton_spread3(ix) | (cunumpy_morton_spread3(iy) << 1) + | (cunumpy_morton_spread3(iz) << 2); +} + +// Cell of x along one axis: floor((x - lower) * scale) in [0, 2^levels - 1]. +__host__ __device__ __forceinline__ unsigned long long cunumpy_morton_cell( + double x, double lower, double scale, int levels) +{ + const double top = (double)((1ull << levels) - 1ull); + double c = floor((x - lower) * scale); + c = c < 0.0 ? 0.0 : c; + c = c > top ? top : c; + return (unsigned long long)c; +} + +// Key of a point, as cunumpy.morton_keys with scales = morton_scales(...). +__host__ __device__ __forceinline__ unsigned long long cunumpy_morton_key2( + double x, double y, double lower_x, double lower_y, + double scale_x, double scale_y, int levels) +{ + return cunumpy_morton_encode2(cunumpy_morton_cell(x, lower_x, scale_x, levels), + cunumpy_morton_cell(y, lower_y, scale_y, levels)); +} + +__host__ __device__ __forceinline__ unsigned long long cunumpy_morton_key3( + double x, double y, double z, double lower_x, double lower_y, double lower_z, + double scale_x, double scale_y, double scale_z, int levels) +{ + return cunumpy_morton_encode3(cunumpy_morton_cell(x, lower_x, scale_x, levels), + cunumpy_morton_cell(y, lower_y, scale_y, levels), + cunumpy_morton_cell(z, lower_z, scale_z, levels)); +} + +#endif // CUNUMPY_MORTON_CUH diff --git a/src/cunumpy/morton.py b/src/cunumpy/morton.py new file mode 100644 index 0000000..3d5e0be --- /dev/null +++ b/src/cunumpy/morton.py @@ -0,0 +1,189 @@ +"""Morton (Z-order) keys, the same on host and device. + +The host side of ``cunumpy/morton.cuh``. A Morton key interleaves the bits of +the integer cell coordinates of a point, so sorting points by their keys +orders them along a Z-shaped space-filling curve: points close in space end up +close in memory (better locality for gathers and neighbour loops), and the +points of every node of a quadtree (2D) or octree (3D) on the same box are a +contiguous range of the sorted array:: + + scales = xp.morton_scales(lower, upper, levels) + keys = xp.morton_keys(positions, lower, upper, levels) # uint64, one per point + keys, order, positions, charges = xp.sort_by_key(keys, positions, charges) + node = keys >> np.uint64(ndim * (levels - level)) # node index at `level` + +Bit layout: with ``levels`` bits per axis the key has ``ndim * levels`` bits; +axis 0 is the lowest bit of every group of ``ndim`` bits. The top group is the +child of the root a point lies in, the next group the child of that child, and +so on. In a kernel, ``cunumpy_morton_key2`` / ``cunumpy_morton_key3`` with the +:func:`morton_scales` of the host return exactly the keys of +:func:`morton_keys` (the cell index is ``floor((x - lower) * scale)`` in both). +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import numpy as np + +from .philox import _module + +__all__ = [ + "MAX_MORTON_LEVELS", + "morton_decode", + "morton_encode", + "morton_keys", + "morton_scales", +] + +#: Largest number of bits per axis that fits a uint64 key, by dimension. +MAX_MORTON_LEVELS = {2: 32, 3: 21} + +_SPREAD = { + 2: ( + (16, 0x0000FFFF0000FFFF), + (8, 0x00FF00FF00FF00FF), + (4, 0x0F0F0F0F0F0F0F0F), + (2, 0x3333333333333333), + (1, 0x5555555555555555), + ), + 3: ( + (32, 0x001F00000000FFFF), + (16, 0x001F0000FF0000FF), + (8, 0x100F00F00F00F00F), + (4, 0x10C30C30C30C30C3), + (2, 0x1249249249249249), + ), +} +_LOW_BITS = {2: 0xFFFFFFFF, 3: 0x1FFFFF} + + +def _check_levels(ndim: int, levels: int) -> None: + if ndim not in MAX_MORTON_LEVELS: + raise ValueError(f"Morton keys support 2 or 3 dimensions, got {ndim}") + if not 1 <= levels <= MAX_MORTON_LEVELS[ndim]: + raise ValueError( + f"levels must be between 1 and {MAX_MORTON_LEVELS[ndim]} in {ndim}D, " + f"got {levels}" + ) + + +def _spread(xpm: Any, value: Any, ndim: int) -> Any: + value = value & xpm.uint64(_LOW_BITS[ndim]) + for shift, mask in _SPREAD[ndim]: + value = (value | (value << xpm.uint64(shift))) & xpm.uint64(mask) + return value + + +def _compact(xpm: Any, value: Any, ndim: int) -> Any: + steps = _SPREAD[ndim] + value = value & xpm.uint64(steps[-1][1]) + masks = [mask for _, mask in steps[:-1]][::-1] + [_LOW_BITS[ndim]] + shifts = [shift for shift, _ in steps][::-1] + for shift, mask in zip(shifts, masks): + value = (value ^ (value >> xpm.uint64(shift))) & xpm.uint64(mask) + return value + + +def morton_encode(*cells: Any) -> Any: + """Interleave the integer cell coordinates of every point into a uint64 key. + + Parameters + ---------- + *cells : arrays of non-negative integers + One array per axis (2 or 3 of them), broadcast against each other; + values must fit in 32 bits (2D) or 21 bits (3D). + + Returns + ------- + array of uint64 + The keys, bit ``ndim * b + a`` holding bit ``b`` of axis ``a``. + """ + ndim = len(cells) + _check_levels(ndim, 1) + xpm = _module(*cells) + cells = xpm.broadcast_arrays(*(xpm.asarray(c).astype(xpm.uint64) for c in cells)) + key = xpm.zeros(cells[0].shape, dtype=xpm.uint64) + for axis, cell in enumerate(cells): + key |= _spread(xpm, cell, ndim) << xpm.uint64(axis) + return key + + +def morton_decode(keys: Any, ndim: int) -> tuple[Any, ...]: + """The integer cell coordinates of Morton keys, the inverse of :func:`morton_encode`. + + Returns + ------- + tuple of arrays of uint64 + One array per axis, shaped like `keys`. + """ + _check_levels(ndim, 1) + xpm = _module(keys) + keys = xpm.asarray(keys).astype(xpm.uint64) + return tuple(_compact(xpm, keys >> xpm.uint64(axis), ndim) for axis in range(ndim)) + + +def morton_scales(lower: Sequence[float], upper: Sequence[float], levels: int) -> Any: + """Cells per unit length of every axis, ``2**levels / (upper - lower)``. + + The numbers to pass to ``cunumpy_morton_key2`` / ``_key3`` in a kernel so + that it computes the keys of :func:`morton_keys`. A float64 NumPy array. + """ + lower = np.asarray(lower, dtype=np.float64) + upper = np.asarray(upper, dtype=np.float64) + if lower.shape != upper.shape or lower.ndim != 1: + raise ValueError( + f"lower and upper must be sequences of equal length, got shapes " + f"{lower.shape} and {upper.shape}" + ) + _check_levels(lower.shape[0], levels) + if np.any(upper == lower): + raise ValueError("upper and lower must differ on every axis") + return float(2**levels) / (upper - lower) + + +def morton_keys( + positions: Any, + lower: Sequence[float], + upper: Sequence[float], + levels: int, +) -> Any: + """Morton keys of points in the box ``[lower, upper]``, ``levels`` bits per axis. + + The cell of a point along an axis is ``floor((x - lower) * scale)`` with the + :func:`morton_scales`, clipped to ``[0, 2**levels - 1]``, so points on or + outside the box get the cell at the nearest face. A point exactly on a cell + boundary belongs to the upper cell. ``lower > upper`` on an axis reverses + that axis (cell 0 at ``lower``). + + Parameters + ---------- + positions : array of float, shape (n, ndim) + The points, ``ndim`` 2 or 3; must be finite. + lower, upper : sequence of float + Corners of the box, one value per axis. + levels : int + Bits per axis: 1 to 32 in 2D, 1 to 21 in 3D. + + Returns + ------- + array of uint64, shape (n,) + On the backend of `positions`. + """ + xpm = _module(positions) + positions = xpm.asarray(positions) + if positions.ndim != 2: + raise ValueError(f"positions must have shape (n, ndim), got {positions.shape}") + ndim = positions.shape[1] + scales = morton_scales(lower, upper, levels) + if scales.shape[0] != ndim: + raise ValueError(f"lower and upper need {ndim} values, got {scales.shape[0]}") + top = float(2**levels - 1) + cells = [] + for axis in range(ndim): + cell = xpm.floor( + (positions[:, axis] - float(lower[axis])) * float(scales[axis]) + ) + cells.append(xpm.clip(cell, 0.0, top).astype(xpm.uint64)) + return morton_encode(*cells) diff --git a/src/cunumpy/xp.py b/src/cunumpy/xp.py index 3ed2762..155f3f5 100644 --- a/src/cunumpy/xp.py +++ b/src/cunumpy/xp.py @@ -1039,6 +1039,47 @@ def segment_sum(values: Any, keys: Any, n_segments: int) -> Any: return out +def sort_by_key(keys: Any, *arrays: Any) -> tuple[Any, ...]: + """Sort `keys` and reorder every array the same way, in one stable argsort. + + The usual first step of a particle code on the GPU: sort the particles by + cell index or Morton key (:func:`cunumpy.morton_keys`), then work on + contiguous ranges. The sort is stable, so equal keys keep their order and + the result is reproducible:: + + keys, order, positions, charges = xp.sort_by_key(keys, positions, charges) + + Parameters + ---------- + keys : array, shape (n,) + The sort keys. + *arrays : arrays + Arrays with ``n`` rows, on the backend of `keys`, reordered along + axis 0. + + Returns + ------- + tuple + ``(sorted_keys, order, *sorted_arrays)``: ``order`` (int64) is the + permutation, ``sorted_keys = keys[order]``, and each sorted array is + ``array[order]`` (a new array). + """ + if get_array_backend(keys) == "cupy": + import cupy as xpm # its argsort is a stable radix sort + else: + xpm = np + keys = xpm.asarray(keys) + if keys.ndim != 1: + raise ValueError(f"keys must be 1D, got shape {keys.shape}") + for array in arrays: + if array.shape[:1] != keys.shape: + raise ValueError( + f"every array needs {keys.shape[0]} rows, got shape {array.shape}" + ) + order = xpm.argsort(keys, kind="stable").astype(xpm.int64, copy=False) + return (keys[order], order, *(array[order] for array in arrays)) + + def to_cunumpy(array: Any) -> Any: """Convert an array to the currently active backend. diff --git a/tests/unit/test_morton.py b/tests/unit/test_morton.py new file mode 100644 index 0000000..de15868 --- /dev/null +++ b/tests/unit/test_morton.py @@ -0,0 +1,222 @@ +"""Tests for Morton keys (host functions and cunumpy/morton.cuh) and sort_by_key.""" + +from pathlib import Path + +import numpy as np +import pytest + +import cunumpy as xp +from cunumpy import CudaKernel, cuda_include_dir +from cunumpy.testing import emulate_cuda_kernel, emulation_compiler + + +def interleave(cells, levels): + """Bit-by-bit reference for morton_encode.""" + ndim = len(cells) + key = 0 + for bit in range(levels): + for axis, cell in enumerate(cells): + key |= ((int(cell) >> bit) & 1) << (ndim * bit + axis) + return key + + +@pytest.mark.parametrize("ndim", [2, 3]) +def test_encode_matches_bitwise_reference(ndim): + levels = xp.MAX_MORTON_LEVELS[ndim] + rng = np.random.default_rng(0) + cells = rng.integers(0, 2**levels, size=(ndim, 500), dtype=np.uint64) + cells[:, 0] = 2**levels - 1 # all bits set + cells[:, 1] = 0 + keys = xp.morton_encode(*cells) + assert keys.dtype == np.uint64 + expected = [interleave(cells[:, i], levels) for i in range(cells.shape[1])] + assert keys.tolist() == expected + + +@pytest.mark.parametrize("ndim", [2, 3]) +def test_decode_inverts_encode(ndim): + levels = xp.MAX_MORTON_LEVELS[ndim] + rng = np.random.default_rng(1) + cells = rng.integers(0, 2**levels, size=(ndim, 1000), dtype=np.uint64) + decoded = xp.morton_decode(xp.morton_encode(*cells), ndim) + for axis in range(ndim): + np.testing.assert_array_equal(decoded[axis], cells[axis]) + + +def test_encode_broadcasts_and_rejects_bad_dimensions(): + keys = xp.morton_encode(np.arange(4)[:, None], np.arange(3)) + assert keys.shape == (4, 3) + assert keys[2, 1] == interleave((2, 1), 2) + with pytest.raises(ValueError, match="2 or 3 dimensions"): + xp.morton_encode(np.arange(3)) + + +def test_keys_cells_and_clipping(): + levels = 3 # 8 cells per axis on [0, 1] + positions = np.array( + [ + [0.0, 0.0], + [0.125, 0.0], # on a cell boundary: the upper cell + [0.999, 0.5], + [1.0, 1.0], # upper face: the last cell + [-5.0, 7.0], # outside: the nearest face + ] + ) + keys = xp.morton_keys(positions, [0.0, 0.0], [1.0, 1.0], levels) + cells = [(0, 0), (1, 0), (7, 4), (7, 7), (0, 7)] + assert keys.tolist() == [interleave(c, levels) for c in cells] + + +def test_reversed_axis(): + # lower > upper on y: cell 0 at the top, like a quadtree with y < mid as + # its second quadrant bit + keys = xp.morton_keys(np.array([[0.2, 0.9], [0.2, 0.1]]), [0, 1], [1, 0], 1) + assert keys.tolist() == [0, 2] + + +def test_keys_validate_arguments(): + with pytest.raises(ValueError, match="levels"): + xp.morton_keys(np.zeros((3, 2)), [0, 0], [1, 1], 33) + with pytest.raises(ValueError, match="levels"): + xp.morton_keys(np.zeros((3, 3)), [0, 0, 0], [1, 1, 1], 22) + with pytest.raises(ValueError, match="differ"): + xp.morton_keys(np.zeros((3, 2)), [0, 0], [1, 0], 4) + with pytest.raises(ValueError, match="need 2 values"): + xp.morton_keys(np.zeros((3, 2)), [0, 0, 0], [1, 1, 1], 4) + with pytest.raises(ValueError, match=r"\(n, ndim\)"): + xp.morton_keys(np.zeros(3), [0, 0], [1, 1], 4) + + +def test_sorted_keys_make_tree_nodes_contiguous(): + rng = np.random.default_rng(2) + positions = rng.random((2000, 2)) + levels = 10 + keys, _, sorted_positions = xp.sort_by_key( + xp.morton_keys(positions, [0, 0], [1, 1], levels), positions + ) + for level in (1, 2, 3): + node = keys >> np.uint64(2 * (levels - level)) + assert np.all(node[1:] >= node[:-1]) + # every point of a node lies in that node's square + size = 0.5**level + cx, cy = xp.morton_decode(node, 2) + assert np.all(np.floor(sorted_positions[:, 0] / size) == cx) + assert np.all(np.floor(sorted_positions[:, 1] / size) == cy) + + +def test_sort_by_key_is_stable_and_reorders_all_arrays(): + keys = np.array([2, 0, 1, 0, 2], dtype=np.uint64) + ids = np.arange(5) + rows = np.arange(10.0).reshape(5, 2) + sorted_keys, order, sorted_ids, sorted_rows = xp.sort_by_key(keys, ids, rows) + assert order.dtype == np.int64 + assert order.tolist() == [1, 3, 2, 0, 4] + assert sorted_keys.tolist() == [0, 0, 1, 2, 2] + np.testing.assert_array_equal(sorted_ids, ids[order]) + np.testing.assert_array_equal(sorted_rows, rows[order]) + assert len(xp.sort_by_key(keys)) == 2 + + +def test_sort_by_key_validates_shapes(): + with pytest.raises(ValueError, match="1D"): + xp.sort_by_key(np.zeros((2, 2))) + with pytest.raises(ValueError, match="3 rows"): + xp.sort_by_key(np.zeros(3), np.zeros(4)) + + +def test_header_is_shipped(): + header = (Path(cuda_include_dir()) / "cunumpy" / "morton.cuh").read_text() + for name in ( + "cunumpy_morton_key2", + "cunumpy_morton_key3", + "cunumpy_morton_encode3", + ): + assert name in header + + +KEYS = r""" +#include "cunumpy/morton.cuh" +extern "C" __global__ +void keys2(const double* pos, unsigned long long* key, long long n, + double x0, double y0, double sx, double sy, int levels) { + long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; + if (i >= n) return; + key[i] = cunumpy_morton_key2(pos[2 * i], pos[2 * i + 1], x0, y0, sx, sy, levels); +} +extern "C" __global__ +void keys3(const double* pos, unsigned long long* key, long long n, + double x0, double y0, double z0, double sx, double sy, double sz, + int levels) { + long long i = blockIdx.x * (long long)blockDim.x + threadIdx.x; + if (i >= n) return; + key[i] = cunumpy_morton_key3(pos[3 * i], pos[3 * i + 1], pos[3 * i + 2], + x0, y0, z0, sx, sy, sz, levels); +} +""" + + +def _cases(ndim): + rng = np.random.default_rng(3) + lower = [-1.5, 0.25, 2.0][:ndim] + upper = [2.5, -0.75, 3.0][:ndim] # the second axis is reversed + positions = rng.uniform(-2.0, 3.5, size=(997, ndim)) + levels = xp.MAX_MORTON_LEVELS[ndim] + # points on cell boundaries, where rounding would show + edges = np.array(lower) + np.arange(5)[:, None] / xp.morton_scales( + lower, upper, levels + ) + positions = np.ascontiguousarray(np.vstack([positions, edges])) + return positions, lower, upper, levels + + +def _device_keys(run, ndim): + positions, lower, upper, levels = _cases(ndim) + n = positions.shape[0] + keys = np.zeros(n, dtype=np.uint64) + scales = xp.morton_scales(lower, upper, levels) + args = (*lower, *scales.tolist(), levels) + run(CudaKernel(KEYS, f"keys{ndim}"), positions, keys, n, *args, n_threads=n) + return keys, xp.morton_keys(positions, lower, upper, levels) + + +@pytest.mark.skipif(emulation_compiler() is None, reason="no C++ compiler") +@pytest.mark.parametrize("ndim", [2, 3]) +def test_header_matches_the_host_keys_in_emulation(ndim): + device, host = _device_keys(emulate_cuda_kernel, ndim) + np.testing.assert_array_equal(device, host) + + +def _run_on_gpu(kernel, *args, n_threads): + import cupy as cp + + device = [cp.asarray(a) if isinstance(a, np.ndarray) else a for a in args] + kernel(*device, n_threads=n_threads) + for host, dev in zip(args, device): + if isinstance(host, np.ndarray): + host[...] = cp.asnumpy(dev) + + +@pytest.mark.parametrize("ndim", [2, 3]) +def test_header_matches_the_host_keys_on_gpu(ndim): + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + device, host = _device_keys(_run_on_gpu, ndim) + np.testing.assert_array_equal(device, host) + + +def test_cupy_arrays_stay_on_the_device(): + if not xp.cupy_available(): + pytest.skip("CuPy not installed or not functional") + import cupy as cp + + positions, lower, upper, levels = _cases(2) + keys = xp.morton_keys(cp.asarray(positions), lower, upper, levels) + assert isinstance(keys, cp.ndarray) + np.testing.assert_array_equal( + cp.asnumpy(keys), xp.morton_keys(positions, lower, upper, levels) + ) + sorted_keys, order, _ = xp.sort_by_key(keys, cp.asarray(positions)) + assert isinstance(order, cp.ndarray) + assert bool((sorted_keys[1:] >= sorted_keys[:-1]).all()) + cx, _ = xp.morton_decode(sorted_keys, 2) + assert isinstance(cx, cp.ndarray)