From ed5dfc69f73edc8b7a898f8355c825186a04f1fa Mon Sep 17 00:00:00 2001 From: Max Date: Sat, 3 Oct 2026 14:25:06 +0200 Subject: [PATCH] Support morton keys --- CHANGELOG.md | 2 + docs/source/api.md | 48 +++++ docs/source/guides/particle-codes.md | 6 + docs/source/kernels/cuda-kernel.md | 1 + src/cunumpy/LLM_GUIDE.md | 2 + src/cunumpy/__init__.py | 14 ++ src/cunumpy/__init__.pyi | 6 + src/cunumpy/cuda/include/cunumpy/morton.cuh | 128 +++++++++++ src/cunumpy/morton.py | 189 +++++++++++++++++ src/cunumpy/xp.py | 41 ++++ tests/unit/test_morton.py | 222 ++++++++++++++++++++ 11 files changed, 659 insertions(+) create mode 100644 src/cunumpy/cuda/include/cunumpy/morton.cuh create mode 100644 src/cunumpy/morton.py create mode 100644 tests/unit/test_morton.py diff --git a/CHANGELOG.md b/CHANGELOG.md index f41ad41..548e698 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,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. - `cunumpy/random.cuh` and `xp.philox_uniform`, `philox_uniform2`, `philox_normal`, `philox_normal2`, `philox4x32_10`: counter-based random numbers (Philox4x32-10, passing the Random123 known-answer tests) as a pure function of `(seed, stream, counter)`, the same in a kernel and on the host (uniform numbers bit for bit, normal numbers up to the last bits of the math functions), so kernels that draw random numbers can be compared with their host versions. - `CudaKernel(..., n_threads_from="first_array")`: one thread per row of the first array argument when a launch gives neither `n_threads` nor `grid`. - `CudaKernel` launches with `shared_mem` above 48 KiB set the kernel's `max_dynamic_shared_size_bytes` once, up to the device's opt-in limit, and raise `ValueError` beyond it. diff --git a/docs/source/api.md b/docs/source/api.md index 06dee24..dacbc5b 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: @@ -1823,6 +1835,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 72a80cf..6cdc8b2 100644 --- a/src/cunumpy/LLM_GUIDE.md +++ b/src/cunumpy/LLM_GUIDE.md @@ -72,6 +72,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" | @@ -135,6 +136,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 4aa6692..269c8ce 100644 --- a/src/cunumpy/__init__.py +++ b/src/cunumpy/__init__.py @@ -25,6 +25,13 @@ from .fusion import fuse from .kernel import CompiledHostKernel, KernelArguments, PyccelKernel, resolve_host_args 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, @@ -76,6 +83,7 @@ set_device, set_device_for_rank, set_mpi_cuda_aware, + sort_by_key, stream, synchronize, synchronize_for_mpi, @@ -122,6 +130,7 @@ def require_version(minimum: str) -> None: __all__ = [ "DEBUG_OPTIONS", "DEFAULT_SHARED_MEMORY_PER_BLOCK", + "MAX_MORTON_LEVELS", "CompiledHostKernel", "CudaArguments", "CudaKernel", @@ -170,6 +179,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", @@ -195,6 +208,7 @@ def require_version(minimum: str) -> None: "set_device", "set_device_for_rank", "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 f54f7ad..f3f7574 100644 --- a/src/cunumpy/__init__.pyi +++ b/src/cunumpy/__init__.pyi @@ -27,6 +27,11 @@ from .fusion import fuse as fuse from .kernel import CompiledHostKernel as CompiledHostKernel from .kernel import PyccelKernel as PyccelKernel 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 @@ -69,6 +74,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 require_version(minimum: str) -> None: ... def device_count() -> int: ... def memory_info() -> tuple[int, int] | None: ... 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 a039a35..42e9678 100644 --- a/src/cunumpy/xp.py +++ b/src/cunumpy/xp.py @@ -1038,6 +1038,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)