From fb3db6f2191877312b42312844637c71dd7c0b3e Mon Sep 17 00:00:00 2001 From: TobiBu Date: Sun, 6 Sep 2026 00:37:44 +0200 Subject: [PATCH] perf(tree): node depths by pointer doubling -- the KD-tree traversal was dispatch-bound `get_node_levels` fell back to a Python loop of `num_nodes - 1` max-relaxation rounds whenever a topology carried no `node_level` field. The radix and octree topologies carry one; the KD-tree's does not, so it was the only backend taking that fallback -- and `_result_to_interactions` calls it on every traversal. Measured on an A100, N = 1e5, leaf 64, theta 0.6, mac dehnen: the fallback is 6385 ms of the KD-tree's 6665 ms traversal, against 0.01 ms for radix's, and it costs a flat 1.56 ms per node at every size (46.9 ms / 30 nodes, 799 ms / 510, 6385 ms / 4094). cProfile attributes 32 896 eagerly dispatched primitives to the KD-tree walk and 134 to radix's -- the same disease as the octree build. The walk itself was never the problem: it runs the same number of wavefront generations as radix's (16 vs 16 at N = 1e4, 22 vs 24 at 1e5) at the same mean occupancy, and the raw walk is *faster* for the KD-tree at N = 1e5 (61.6 ms against 64.6 ms). Depths now come from a device-side pointer-doubling `lax.while_loop`: one dispatched computation, O(log depth) rounds. `_compute_node_depths`, which already did this inside the jitted walk, becomes a thin alias so the walk's depths and the interaction list's levels have one implementation. Verified equal to the relaxation element for element on heaps, paths and forests with padded nodes at 18 sizes, and on all three backends -- where the radix and octree topologies also agree with their own `node_level` field. Co-Authored-By: Claude Opus 5 (1M context) --- tests/unit/test_node_levels.py | 128 +++++++++++++++++++++++++++++++++ yggdrax/_interactions_impl.py | 36 ++-------- yggdrax/tree.py | 66 ++++++++++++++--- 3 files changed, 191 insertions(+), 39 deletions(-) create mode 100644 tests/unit/test_node_levels.py diff --git a/tests/unit/test_node_levels.py b/tests/unit/test_node_levels.py new file mode 100644 index 0000000..98b7143 --- /dev/null +++ b/tests/unit/test_node_levels.py @@ -0,0 +1,128 @@ +"""Per-node depths, against the relaxation they replaced. + +``get_node_levels`` used to fall back to a Python loop of ``num_nodes - 1`` +rounds whenever a topology carried no ``node_level`` field -- ``O(num_nodes)`` +eagerly dispatched primitives, 1.56 ms per node. Only the KD-tree takes that +fallback, and it is why its dual-tree traversal was 50x the radix backend's at +N = 1e5 while doing strictly less device work. It is now a pointer-doubling +``lax.while_loop``: one dispatched computation, ``O(log depth)`` rounds. + +These tests pin the depths to the definition -- distance to the root along the +parent chain -- rather than to either implementation. +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from yggdrax import Tree +from yggdrax.dtypes import INDEX_DTYPE, as_index +from yggdrax.tree import get_node_levels, node_levels_from_parent + + +def _levels_by_relaxation(parent: np.ndarray) -> np.ndarray: + """The pre-change fallback, verbatim: one max-relaxation round per node.""" + parent_j = jnp.asarray(parent, dtype=INDEX_DTYPE) + num_nodes = int(parent_j.shape[0]) + if num_nodes == 0: + return np.zeros((0,), dtype=np.int64) + levels = jnp.zeros((num_nodes,), dtype=INDEX_DTYPE) + parent_safe = jnp.where(parent_j >= 0, parent_j, as_index(0)) + for _ in range(max(num_nodes - 1, 0)): + candidate = jnp.where( + parent_j >= 0, levels[parent_safe] + as_index(1), as_index(0) + ) + levels = jnp.maximum(levels, candidate) + return np.asarray(levels) + + +def _heap_parents(n: int) -> np.ndarray: + """Parent links of a heap-ordered binary tree of ``n`` nodes.""" + return np.array([-1] + [(i - 1) // 2 for i in range(1, n)], dtype=np.int32) + + +def _path_parents(n: int) -> np.ndarray: + """Parent links of a single chain -- the worst case for depth.""" + return np.array([-1] + list(range(n - 1)), dtype=np.int32) + + +_SIZES = [1, 2, 3, 4, 5, 7, 8, 15, 16, 17, 31, 32, 63, 64, 100, 255, 256, 257] + + +@pytest.mark.parametrize("n", _SIZES) +def test_heap_depths_match_the_relaxation(n: int) -> None: + parent = _heap_parents(n) + got = np.asarray(node_levels_from_parent(jnp.asarray(parent, dtype=INDEX_DTYPE))) + np.testing.assert_array_equal(got, _levels_by_relaxation(parent)) + # A heap's depth is the bit length of the 1-based index. + expected = np.array([int(i + 1).bit_length() - 1 for i in range(n)]) + np.testing.assert_array_equal(got, expected) + + +@pytest.mark.parametrize("n", _SIZES) +def test_path_depths_match_the_relaxation(n: int) -> None: + parent = _path_parents(n) + got = np.asarray(node_levels_from_parent(jnp.asarray(parent, dtype=INDEX_DTYPE))) + np.testing.assert_array_equal(got, _levels_by_relaxation(parent)) + np.testing.assert_array_equal(got, np.arange(n)) + + +@pytest.mark.parametrize("seed", [0, 1, 2]) +def test_forests_with_padded_nodes_match_the_relaxation(seed: int) -> None: + """Several roots and unattached nodes: every ``parent < 0`` is depth 0.""" + rng = np.random.default_rng(seed) + n = 257 + parent = np.array( + [-1] + + [ + int(rng.integers(0, max(1, i))) if rng.random() > 0.1 else -1 + for i in range(1, n) + ], + dtype=np.int32, + ) + got = np.asarray(node_levels_from_parent(jnp.asarray(parent, dtype=INDEX_DTYPE))) + np.testing.assert_array_equal(got, _levels_by_relaxation(parent)) + np.testing.assert_array_equal(got[parent < 0], 0) + + +def test_empty_topology_returns_empty_levels() -> None: + class _Empty: + parent = jnp.zeros((0,), dtype=INDEX_DTYPE) + + assert get_node_levels(_Empty()).shape == (0,) + + +@pytest.mark.parametrize("backend", ["radix", "octree", "kdtree"]) +def test_backend_levels_agree_with_the_relaxation(backend: str) -> None: + """Every backend's depths, derived and (where carried) declared. + + The KD-tree is the one that matters: it carries no ``node_level``, so it is + the only backend whose depths ``get_node_levels`` has to derive. + """ + key = jax.random.PRNGKey(0) + kp, km = jax.random.split(key) + positions = jax.random.uniform( + kp, (2000, 3), minval=-1.0, maxval=1.0, dtype=jnp.float32 + ) + masses = jax.random.uniform(km, (2000,), minval=0.5, maxval=1.5, dtype=jnp.float32) + tree = Tree.from_particles( + positions, + masses, + tree_type=backend, + build_mode="adaptive", + leaf_size=64, + return_reordered=True, + ) + topo = tree.topology if hasattr(tree, "topology") else tree + parent = np.asarray(jnp.asarray(topo.parent, dtype=INDEX_DTYPE)) + derived = np.asarray( + node_levels_from_parent(jnp.asarray(parent, dtype=INDEX_DTYPE)) + ) + np.testing.assert_array_equal(derived, _levels_by_relaxation(parent)) + if hasattr(topo, "node_level"): + declared = np.asarray(jnp.asarray(topo.node_level, dtype=INDEX_DTYPE)) + np.testing.assert_array_equal(derived, declared) + np.testing.assert_array_equal(np.asarray(get_node_levels(topo)), derived) diff --git a/yggdrax/_interactions_impl.py b/yggdrax/_interactions_impl.py index bf46c85..dfa2a1d 100644 --- a/yggdrax/_interactions_impl.py +++ b/yggdrax/_interactions_impl.py @@ -38,6 +38,7 @@ get_level_offsets, get_node_levels, get_nodes_by_level, + node_levels_from_parent, ) # Each node only needs to interact with a bounded number of well-separated @@ -978,37 +979,14 @@ def body_fn(state): def _compute_node_depths(parent: Array) -> Array: """Return the depth of every node (root depth = 0). - Uses pointer doubling for O(log depth) convergence. Each node - keeps a *depth-to-root* counter and a shortcut pointer. On each - round the shortcut doubles its reach and accumulated depth - contributions are propagated. + Thin alias for :func:`yggdrax.tree.node_levels_from_parent`, which is the + single implementation of the pointer-doubling depth pass -- the same one + ``get_node_levels`` falls back to when a topology carries no ``node_level`` + field, so the walk's depths and the interaction list's levels cannot drift + apart. """ - total_nodes = parent.shape[0] - is_root = parent < 0 - # dist[i] = accumulated distance along the shortcut chain. - # Initially 1 for non-root nodes (edge to parent), 0 for root. - dist = jnp.where(is_root, as_index(0), as_index(1)) - # shortcut[i] = parent[i] for non-root, i for root. - shortcut = jnp.where( - is_root, - jnp.arange(total_nodes, dtype=parent.dtype), - parent, - ) - - def cond_fn(state): - _sc, _d, changed = state - return changed - - def body_fn(state): - sc, d, _changed = state - # Pointer doubling: add distance of shortcut target. - new_d = d + d[sc] - new_sc = sc[sc] - changed = jnp.any(new_sc != sc) - return new_sc, new_d, changed - _, depth, _ = lax.while_loop(cond_fn, body_fn, (shortcut, dist, jnp.bool_(True))) - return depth + return node_levels_from_parent(parent) def _compute_effective_extents(parent: Array, extents: Array) -> Array: diff --git a/yggdrax/tree.py b/yggdrax/tree.py index df5d07e..e879733 100644 --- a/yggdrax/tree.py +++ b/yggdrax/tree.py @@ -1256,6 +1256,60 @@ def get_leaf_nodes(tree: object) -> Array: return jnp.arange(num_internal, total_nodes, dtype=INDEX_DTYPE) +@jax.jit +def node_levels_from_parent(parent: Array) -> Array: + """Return per-node depth from parent links, by pointer doubling. + + Each node keeps a shortcut pointer into its ancestor chain and the distance + it has already covered; a round doubles the reach of every shortcut, so the + walk to the root converges in ``O(log depth)`` rounds rather than one round + per level. Nodes whose ``parent`` is negative are roots and get depth 0, + which also makes padded nodes harmless. + + This is a device-side ``lax.while_loop``: one dispatched computation + whatever the tree's size. The relaxation it replaces was a Python loop of + ``num_nodes - 1`` rounds -- ``O(num_nodes)`` *eagerly dispatched* + primitives, which cost 1.56 ms per node and made the KD-tree's traversal + 50x the radix backend's at N = 10^5 while doing strictly less device work. + + Parameters + ---------- + parent + Per-node parent index, negative at a root, length ``n_nodes``. + + Returns + ------- + Array + Per-node depth, root depth 0, length ``n_nodes``. + """ + + total_nodes = parent.shape[0] + is_root = parent < 0 + # dist[i] is how far i has already walked along its shortcut chain. + dist = jnp.where(is_root, as_index(0), as_index(1)) + # A root shortcuts to itself, so it is the chain's fixpoint. + shortcut = jnp.where( + is_root, + jnp.arange(total_nodes, dtype=parent.dtype), + parent, + ) + + def cond_fn(state): + _shortcut, _dist, changed = state + return changed + + def body_fn(state): + sc, d, _changed = state + new_dist = d + d[sc] + new_shortcut = sc[sc] + return new_shortcut, new_dist, jnp.any(new_shortcut != sc) + + _, depth, _ = jax.lax.while_loop( + cond_fn, body_fn, (shortcut, dist, jnp.bool_(True)) + ) + return depth.astype(INDEX_DTYPE) + + def get_node_levels(tree: object) -> Array: """Return per-node depth levels, deriving from parent links when missing.""" @@ -1267,16 +1321,7 @@ def get_node_levels(tree: object) -> Array: if num_nodes == 0: return jnp.zeros((0,), dtype=INDEX_DTYPE) - levels = jnp.zeros((num_nodes,), dtype=INDEX_DTYPE) - parent_safe = jnp.where(parent >= 0, parent, as_index(0)) - for _ in range(max(num_nodes - 1, 0)): - candidate = jnp.where( - parent >= 0, - levels[parent_safe] + as_index(1), - as_index(0), - ) - levels = jnp.maximum(levels, candidate) - return levels + return node_levels_from_parent(parent) def get_num_levels(tree: object, *, node_levels: Optional[Array] = None) -> int: @@ -2067,6 +2112,7 @@ def build_fixed_depth_octree_jit( "missing_fmm_topology_fields", "missing_leaf_topology_fields", "missing_morton_topology_fields", + "node_levels_from_parent", "resolve_tree_topology", "require_fmm_core_topology", "require_fmm_topology",