perf(tree): node depths by pointer doubling — the KD-tree traversal was dispatch-bound

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)

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:

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