Morton keys

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: `<name><host_suffix>.py` (`"pyccel"`, compiled with `compile_host`, and `"python"`, uncompiled), `<name>_numba.py`, `<name>_numpy.py` and `<name>_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.

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 <cunumpy/morton.cuh>
+
+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

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
 

docs/source/kernels/cuda-kernel.md

@@ -156,6 +156,7 @@ Pass extra include directories with `include_dirs=[...]` and NVRTC flags with
 | `<cunumpy/index.cuh>` | `CUNUMPY_THREAD_1D(i, n)`, `_2D`, `_3D`, `CUNUMPY_GRID_STRIDE_1D(i, n)` |
 | `<cunumpy/array_view.cuh>` | strided views `Array1D<T>` to `Array4D<T>` |
 | `<cunumpy/atomic.cuh>` | `cunumpy_atomic_add` and indexed 2D/3D variants, see [Accumulation kernels](accumulation.md) |
+| `<cunumpy/morton.cuh>` | Morton (Z-order) keys `cunumpy_morton_key2(x, y, ...)`, `_key3`, equal to `xp.morton_keys` on the host |
 | `<cunumpy/random.cuh>` | 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.

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/random.cuh>`: `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.cuh>`: `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
 ```
 

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",

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]: ...

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 <cunumpy/morton.cuh>
+//
+//     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