From 20b2fe53f8f0d8f386f00d2a79cc65b3e7a6918b Mon Sep 17 00:00:00 2001 From: Tobias Buck Date: Sun, 12 Jul 2026 21:49:56 +0200 Subject: [PATCH 1/4] ci(docs): publish the Sphinx site to GitHub Pages Adds a Docs workflow that builds the Sphinx HTML and deploys it to GitHub Pages via actions/upload-pages-artifact + actions/deploy-pages. It builds on every push and PR (catching doc breakage) but only deploys on pushes to main or a manual workflow_dispatch, so PRs never touch the live site. Links the hosted docs from the README. Co-Authored-By: Claude Opus 4.8 --- .github/workflows/docs.yml | 62 ++++++++++++++++++++++++++++++++++++++ README.md | 15 +++++---- 2 files changed, 71 insertions(+), 6 deletions(-) create mode 100644 .github/workflows/docs.yml diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 0000000..13f7f51 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,62 @@ +name: Docs + +on: + push: + branches: [main] + pull_request: + workflow_dispatch: + +# Least-privilege token, plus the Pages deploy permissions. +permissions: + contents: read + pages: write + id-token: write + +# Allow only one concurrent deployment; let an in-progress run finish. +concurrency: + group: pages + cancel-in-progress: false + +jobs: + build: + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + + - name: Install package + docs toolchain + # jax is pinned to match the version the docs are validated against + # (autodoc imports the package); sphinx/furo/myst come from ".[docs]". + run: | + python -m pip install --upgrade pip + pip install -e ".[docs]" "jax==0.8.3" "jaxlib==0.8.3" + + - name: Build HTML + run: | + python -m sphinx -b html docs docs/_build/html + touch docs/_build/html/.nojekyll + + - name: Upload Pages artifact + uses: actions/upload-pages-artifact@v3 + with: + path: docs/_build/html + + # Publish only for main (post-merge) or a manual run; PRs build without + # deploying so doc breakage is caught without touching the live site. + deploy: + needs: build + if: github.event_name == 'push' || github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + steps: + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v4 diff --git a/README.md b/README.md index a171405..617575c 100644 --- a/README.md +++ b/README.md @@ -188,8 +188,15 @@ pytest --cov=yggdrax --cov-report=term-missing ## Documentation -The reference documentation is built with Sphinx (autodoc + napoleon, numpy -docstring style): +Hosted documentation: **https://tobias-buck.de/yggdrax/** + +It includes a concepts guide (trees/backends, Morton ordering, MAC variants, +buffer glossary), the backend contract, and the full API reference generated +from docstrings. The site is built with Sphinx (autodoc + napoleon, numpy +docstring style) and published to GitHub Pages on every push to `main` by the +`Docs` workflow. + +To build it locally: ```bash pip install -e ".[docs]" @@ -197,10 +204,6 @@ python -m sphinx -b html docs docs/_build/html # open docs/_build/html/index.html ``` -It includes a concepts guide (trees/backends, Morton ordering, MAC variants, -buffer glossary), the backend contract, and the full API reference generated -from docstrings. - ## Project Structure - `yggdrax/tree.py`, `yggdrax/_tree_impl.py`: tree building and radix internals From f4caf9a5506ea54a4081cf32bacd01a7e8d6683a Mon Sep 17 00:00:00 2001 From: TobiBu Date: Thu, 16 Jul 2026 12:57:17 +0200 Subject: [PATCH 2/4] perf(walk): batch the dual-tree walk's per-key-prefix argsorts (4 -> 2 sorts) The static-radix fused fast-lane per step is launch/host-bound (GPU idle ~58-66% on A100 @200k): ~4200 host-launched micro-kernels, of which the dual-tree walk's `_per_key_prefix` contributes FOUR 262144-element argsorts inside the while-body (each lowering to ~150 bitonic device kernels). Within `_dual_tree_walk_impl`, the far (fwd/bwd) and near (fwd/bwd) per-key prefixes are INDEPENDENT of one another -- only their downstream slot-index use is sequential, not the argsort itself. Batch each pair into a single `jax.vmap(_per_key_prefix)` over a stacked (2, B) array, so XLA emits one batched bitonic sort per pair instead of two. `vmap` is semantically identical to the two sequential calls, so forces are bit-identical. Measured @200k / leaf256 / A100 (fused device-only lane, min-of-3 x 80 steps): - fused per-step HLO sort ops: 5 -> 3 - ms/step: 121.4 -> 112.5 (~7% faster) - energy conservation over 60 steps: max|dE/E0| = 2.48e-5 (unchanged) Not a null result (unlike the earlier off-critical-path M2L-count reduction): the walk argsorts are on the serial critical path, so cutting their launch count moves the wall. Co-Authored-By: Claude Opus 4.8 --- yggdrax/_interactions_impl.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/yggdrax/_interactions_impl.py b/yggdrax/_interactions_impl.py index 2e66716..3d00bbc 100644 --- a/yggdrax/_interactions_impl.py +++ b/yggdrax/_interactions_impl.py @@ -1480,7 +1480,14 @@ def _far_update(carry): accept_mask, pair_tags_rev, as_index(-1) ) - fwd_prefix = _per_key_prefix(accept_targets_a, accept_mask, total_nodes) + # fwd/bwd per-key prefixes are INDEPENDENT (only their downstream slot + # use is sequential), so batch the two 256k-element argsorts into ONE + # vmapped sort -> halves the walk's dominant sort-kernel launch count. + _far_prefixes = jax.vmap(_per_key_prefix)( + jnp.stack([accept_targets_a, accept_sources_a], axis=0), + jnp.stack([accept_mask, accept_mask], axis=0), + ) + fwd_prefix = _far_prefixes[0] fwd_slot = far_counts_c[safe_targets] + fwd_prefix fwd_ok = accept_mask & (fwd_slot < as_index(max_interactions_per_node)) far_overflow_c = far_overflow_c | jnp.any( @@ -1505,7 +1512,7 @@ def _far_update(carry): ) far_counts_after_fwd = far_counts_c + fwd_incr - bwd_prefix = _per_key_prefix(accept_sources_a, accept_mask, total_nodes) + bwd_prefix = _far_prefixes[1] bwd_slot = far_counts_after_fwd[safe_sources] + bwd_prefix bwd_ok = accept_mask & (bwd_slot < as_index(max_interactions_per_node)) far_overflow_c = far_overflow_c | jnp.any( @@ -1576,7 +1583,13 @@ def _near_update(carry): safe_pos_t = jnp.where(near_mask, near_pos_t, as_index(0)) safe_pos_s = jnp.where(near_mask, near_pos_s, as_index(0)) - fwd_near_prefix = _per_key_prefix(safe_pos_t, near_mask, num_leaves) + # Batch the two independent near per-key prefix argsorts into one + # vmapped sort (same rationale as the far pair above). + _near_prefixes = jax.vmap(_per_key_prefix)( + jnp.stack([safe_pos_t, safe_pos_s], axis=0), + jnp.stack([near_mask, near_mask], axis=0), + ) + fwd_near_prefix = _near_prefixes[0] fwd_near_slot = near_counts_c[safe_pos_t] + fwd_near_prefix fwd_near_ok = near_mask & ( fwd_near_slot < as_index(max_neighbors_per_leaf) @@ -1606,7 +1619,7 @@ def _near_update(carry): ) near_counts_after_fwd = near_counts_c + fwd_near_incr - bwd_near_prefix = _per_key_prefix(safe_pos_s, near_mask, num_leaves) + bwd_near_prefix = _near_prefixes[1] bwd_near_slot = near_counts_after_fwd[safe_pos_s] + bwd_near_prefix bwd_near_ok = near_mask & ( bwd_near_slot < as_index(max_neighbors_per_leaf) From 3b1138fa296fdd26ea39347e06bcdd593261c645 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Thu, 16 Jul 2026 14:00:36 +0200 Subject: [PATCH 3/4] feat(distributed): backend-selectable tree build + leaf-only KD-tree traversal Multi-GPU tree build and cross-walk now support all three backends. Distributed backend selector (radix + octree): - local_tree: DISTRIBUTED_TREE_TYPES + _validate_distributed_tree_type + shard_map-safe _build_local_tree dispatch; distributed_tree_moments / build_local_moments gain tree_type (root via argmin(parent)). - let: tree_type threaded to the local build in build_distributed_coarse_tree / classify_against_remote / distributed_let_import. Coarse tree stays radix (internal remote-COM representation; robust at coarse leaf_size=1). Leaf-only (bucket) median-split KD-tree (build_leaf_kdtree, LeafKDTree): - All points in leaves; internal nodes are pure split planes; compact contiguous node_ranges tiling [0,N) + nodes_by_level -> satisfies the same topology contract as radix/octree, giving complete pair coverage. The heap KD-tree (build_kdtree) used by the KNN kernels is untouched; the generic pipeline / tree_type="kdtree" / distributed now use the leaf-only build. Validated: 46 kd unit tests (heap KNN preserved + leaf-kd interactions parity vs radix); distributed CPU 9 passed incl. kd cross-walk complete disjoint source partition (the invariant the heap KD-tree violated). GPU validation in progress. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/distributed/test_local_tree_backends.py | 199 +++++++++++++ yggdrax/distributed/let.py | 46 ++- yggdrax/distributed/local_tree.py | 107 ++++++- yggdrax/kdtree.py | 279 +++++++++++++++++- yggdrax/tree.py | 6 +- 5 files changed, 614 insertions(+), 23 deletions(-) create mode 100644 tests/distributed/test_local_tree_backends.py diff --git a/tests/distributed/test_local_tree_backends.py b/tests/distributed/test_local_tree_backends.py new file mode 100644 index 0000000..b5f46a8 --- /dev/null +++ b/tests/distributed/test_local_tree_backends.py @@ -0,0 +1,199 @@ +"""Multi-GPU distributed tree build + cross-walk across backends. + +`distributed_tree_moments` / the LET drivers gained a `tree_type` selector +(``radix``, ``octree``, ``kdtree``). The build correctness anchor is +backend-independent: the union of the per-domain coarse moments must reproduce +the *global* total mass and centre of mass a single-device build would give. + +For ``kdtree`` the distributed path uses the leaf-only bucket build +(:func:`yggdrax.kdtree.build_leaf_kdtree`), which stores all points in leaves and +so satisfies the same topology contract as radix/octree -- unlike the heap +KD-tree used by the KNN kernels. The cross-walk test asserts the actual thing +that was previously broken for KD: a complete, disjoint source partition. + + XLA_FLAGS=--xla_force_host_platform_device_count=4 JAX_PLATFORMS=cpu \ + pytest tests/distributed/test_local_tree_backends.py -o addopts="" -q +""" + +import jax.numpy as jnp +import numpy as np +import pytest + +from yggdrax.bounds import infer_bounds +from yggdrax.distributed import ( + build_distributed_coarse_tree, + device_count, + distributed_tree_moments, + dual_tree_walk_cross, + make_mesh, +) +from yggdrax.distributed.local_tree import ( + DISTRIBUTED_TREE_TYPES, + _build_local_tree, + _validate_distributed_tree_type, +) +from yggdrax.geometry import compute_tree_geometry + +_NDEV = min(4, device_count()) +_LEAF = 8 +_PER_DEV = 16 +_N = _PER_DEV * _NDEV +_CAP = 4 * _PER_DEV + +needs_devices = pytest.mark.skipif( + device_count() < 2, reason="distributed tree build needs >= 2 devices" +) + + +# ---- validation (no devices needed) ---- +def test_distributed_tree_types_advertised(): + assert DISTRIBUTED_TREE_TYPES == ("radix", "octree", "kdtree") + + +def test_unknown_tree_type_rejected(): + with pytest.raises(ValueError, match="Unsupported tree_type"): + _validate_distributed_tree_type("bogus") + + +# ---- distributed build reconstruction: radix baseline + octree + kdtree ---- +@pytest.fixture(scope="module") +def built(): + """Run distributed_tree_moments once per backend; reuse across assertions.""" + mesh = make_mesh(_NDEV) + rng = np.random.default_rng(10) + pts = rng.uniform(-1.0, 1.0, size=(_N, 3)).astype(np.float32) + mass = rng.uniform(0.5, 2.0, size=(_N,)).astype(np.float32) + out = {} + for tt in ("octree", "kdtree"): + out[tt] = distributed_tree_moments( + mesh, + jnp.asarray(pts), + jnp.asarray(mass), + leaf_size=_LEAF, + output_capacity=_CAP, + equalize=True, + tree_type=tt, + ) + return pts, mass, out + + +@needs_devices +@pytest.mark.parametrize("tree_type", ["octree", "kdtree"]) +def test_domain_moments_reconstruct_global(built, tree_type): + pts, mass, out = built + res = out[tree_type] + domain_mass = np.asarray(res.domain_mass) + domain_com = np.asarray(res.domain_com) + counts = np.asarray(res.counts) + + assert int(counts.sum()) == _N + np.testing.assert_allclose(domain_mass.sum(), mass.sum(), rtol=1e-4) + global_com = (mass[:, None] * pts).sum(0) / mass.sum() + recon_com = (domain_mass[:, None] * domain_com).sum(0) / domain_mass.sum() + np.testing.assert_allclose(recon_com, global_com, rtol=1e-3, atol=1e-4) + + +@needs_devices +@pytest.mark.parametrize("tree_type", ["octree", "kdtree"]) +def test_shared_top_tree_replicated(built, tree_type): + _, _, out = built + res = out[tree_type] + top_mass = np.asarray(res.top_mass).reshape(_NDEV, _NDEV) + domain_mass = np.asarray(res.domain_mass) + for g in range(_NDEV): + np.testing.assert_allclose(top_mass[g], top_mass[0], rtol=1e-5) + np.testing.assert_allclose(top_mass[0], domain_mass, rtol=1e-5) + + +@needs_devices +@pytest.mark.parametrize("tree_type", ["octree", "kdtree"]) +def test_coarse_tree_root_is_global(built, tree_type): + """build_distributed_coarse_tree: coarse root mass == global (any local backend).""" + pts, mass, _ = built + mesh = make_mesh(_NDEV) + metrics = build_distributed_coarse_tree( + mesh, + jnp.asarray(pts), + jnp.asarray(mass), + leaf_size=_LEAF, + output_capacity=_CAP, + equalize=True, + tree_type=tree_type, + ) + np.testing.assert_allclose(np.asarray(metrics.coarse_root_mass), mass.sum(), rtol=1e-3) + + +# ---- cross-tree walk over leaf-only KD trees: complete disjoint source partition ---- +def _kd_build(points, leaf_size=_LEAF): + pts = jnp.asarray(points) + tree, ps, _ms = _build_local_tree( + pts, jnp.ones(pts.shape[0], jnp.float32), infer_bounds(pts), + tree_type="kdtree", leaf_size=leaf_size, + ) + geom = compute_tree_geometry(tree, ps, max_leaf_size=leaf_size) + return tree, geom + + +def test_kdtree_cross_walk_partitions_sources(): + """Every target leaf's far ancestors + near leaves tile ALL source particles once. + + This is the invariant the heap KD-tree violated (internal-node pivots absent + from the near-list); the leaf-only build must satisfy it exactly. + """ + rng = np.random.default_rng(3) + tgt = rng.uniform(-1.0, 0.2, size=(80, 3)).astype(np.float32) + src = rng.uniform(-0.2, 1.0, size=(80, 3)).astype(np.float32) + # shared frame + allpts = jnp.asarray(np.concatenate([tgt, src], 0)) + bounds = infer_bounds(allpts) + + def build(points): + pts = jnp.asarray(points) + tree, ps, _ = _build_local_tree( + pts, jnp.ones(pts.shape[0], jnp.float32), bounds, + tree_type="kdtree", leaf_size=_LEAF, + ) + return tree, compute_tree_geometry(tree, ps, max_leaf_size=_LEAF) + + t_tree, t_geom = build(tgt) + s_tree, s_geom = build(src) + res = dual_tree_walk_cross( + t_tree, t_geom, s_tree, s_geom, 0.5, mac_type="bh", + max_interactions_per_node=512, max_neighbors_per_leaf=512, max_pair_queue=16384, + ) + assert not bool(res.queue_overflow | res.far_overflow | res.near_overflow) + + parent = np.asarray(t_tree.parent) + s_ranges = np.asarray(s_tree.node_ranges) + tt = np.asarray(res.interaction_targets) + ss = np.asarray(res.interaction_sources) + m = tt >= 0 + far_map = {} + for t_node, s_node in zip(tt[m].tolist(), ss[m].tolist()): + far_map.setdefault(t_node, []).append(s_node) + near_off = np.asarray(res.neighbor_offsets) + near_idx = np.asarray(res.neighbor_indices) + near_cnt = np.asarray(res.neighbor_counts) + leaf_indices = np.asarray(res.leaf_indices) + n_source = int(src.shape[0]) + root = int(np.argmin(parent)) + + def src_particles(node): + lo, hi = int(s_ranges[node, 0]), int(s_ranges[node, 1]) + return list(range(lo, hi + 1)) + + for row, leaf_node in enumerate(leaf_indices.tolist()): + far_sources, node = [], leaf_node + while True: + far_sources.extend(far_map.get(node, [])) + if node == root: + break + node = int(parent[node]) + o, c = int(near_off[row]), int(near_cnt[row]) + covered = [] + for sn in far_sources + near_idx[o:o + c].tolist(): + covered.extend(src_particles(sn)) + assert set(covered) == set(range(n_source)), ( + f"leaf_row {row}: covered {len(set(covered))}/{n_source}" + ) + assert len(covered) == n_source, f"leaf_row {row}: overlap {len(covered)}" diff --git a/yggdrax/distributed/let.py b/yggdrax/distributed/let.py index 5c62ec2..e51a748 100644 --- a/yggdrax/distributed/let.py +++ b/yggdrax/distributed/let.py @@ -27,12 +27,12 @@ import jax.numpy as jnp from jaxtyping import Array -from .._tree_impl import build_tree from ..dtypes import INDEX_DTYPE, as_index from ..geometry import compute_tree_geometry from ..tree import Tree from ..tree_moments import compute_tree_mass_moments from .comm import _COUNT_DTYPE, ragged_all_to_all_exchange +from .local_tree import _build_local_tree, _validate_distributed_tree_type from .sharding import AXIS_NAME # Global particle id = source_domain * _GID_STRIDE + local_sorted_index. Lets an @@ -106,7 +106,9 @@ def gather_global_coarse_tree( Each frontier node becomes a "coarse particle" (its COM, weighted by node mass). Building over identical gathered input yields the same tree on every - device -- no agreement round needed. + device -- no agreement round needed. The coarse tree is always radix (an + internal representation over remote COMs; the backend is immaterial and + radix is robust at ``coarse_leaf_size=1``). """ n_top = frontier.mass.shape[0] @@ -122,8 +124,8 @@ def gather_global_coarse_tree( frontier.node_range, axis_name, tiled=True ) # [ncoarse,2] - tree, pos_sorted, mass_sorted, _inv = build_tree( - coms, masses, bounds, return_reordered=True, leaf_size=int(coarse_leaf_size) + tree, pos_sorted, mass_sorted = _build_local_tree( + coms, masses, bounds, tree_type="radix", leaf_size=int(coarse_leaf_size) ) geometry = compute_tree_geometry( tree, pos_sorted, max_leaf_size=int(coarse_leaf_size) @@ -164,6 +166,7 @@ def build_distributed_coarse_tree( output_capacity: int, num_samples: int = 8, equalize: bool = True, + tree_type: str = "radix", axis_name: str = AXIS_NAME, ) -> CoarseTreeMetrics: """Decompose, build per-GPU trees, and gather the global coarse tree. @@ -171,8 +174,12 @@ def build_distributed_coarse_tree( Returns global-view :class:`CoarseTreeMetrics`. The full per-device :class:`GlobalCoarseTree` is built inside the ``shard_map`` and consumed by later LET stages; this driver surfaces the invariants worth testing. + ``tree_type`` selects the local per-device backend (``"radix"``, + ``"octree"``, or ``"kdtree"``); the coarse tree is always radix. """ + _validate_distributed_tree_type(tree_type) + try: # stable across recent JAX versions from jax import shard_map except ImportError: # pragma: no cover @@ -200,8 +207,8 @@ def fn(pos, mass): p, m, c, cnt, ndev, output_capacity=output_capacity, axis_name=axis_name ) p, m = sanitize_padding(p, m, cnt) - tree, pos_sorted, mass_sorted, _ = build_tree( - p, m, bounds, return_reordered=True, leaf_size=leaf_size + tree, pos_sorted, mass_sorted = _build_local_tree( + p, m, bounds, tree_type=tree_type, leaf_size=leaf_size ) moments = compute_tree_mass_moments(tree, pos_sorted, mass_sorted) node_mass = moments.mass @@ -209,6 +216,9 @@ def fn(pos, mass): root = as_index(jnp.argmin(jnp.asarray(tree.parent))) fr = build_coarse_frontier(tree, node_mass, node_com) + # Coarse tree is always radix: it is an internal representation over + # remote frontier COMs (backend-immaterial), and radix is robust at the + # coarse leaf_size=1 that KD/octree bucket builds cannot honour exactly. gct = gather_global_coarse_tree(fr, bounds=bounds, axis_name=axis_name) croot = as_index(jnp.argmin(jnp.asarray(gct.tree.parent))) @@ -325,6 +335,7 @@ def classify_against_remote( max_pair_queue: int = 8192, num_samples: int = 8, equalize: bool = True, + tree_type: str = "radix", axis_name: str = AXIS_NAME, ) -> ClassifyMetrics: """Cross-walk each GPU's local tree against the remote-only coarse tree. @@ -332,8 +343,12 @@ def classify_against_remote( Produces the far (remote M2L) and near (remote import) classifications that Phase-3 stage 3 turns into an actual particle import. This driver surfaces the invariants; the far/near *lists* are consumed by later stages. + ``tree_type`` selects the local per-device backend (``"radix"``, + ``"octree"``, or ``"kdtree"``); the coarse tree is always radix. """ + _validate_distributed_tree_type(tree_type) + try: # stable across recent JAX versions from jax import shard_map except ImportError: # pragma: no cover @@ -362,8 +377,8 @@ def fn(pos, mass): p, m, c, cnt, ndev, output_capacity=output_capacity, axis_name=axis_name ) p, m = sanitize_padding(p, m, cnt) - tree, pos_sorted, mass_sorted, _ = build_tree( - p, m, bounds, return_reordered=True, leaf_size=leaf_size + tree, pos_sorted, mass_sorted = _build_local_tree( + p, m, bounds, tree_type=tree_type, leaf_size=leaf_size ) geom = compute_tree_geometry(tree, pos_sorted, max_leaf_size=leaf_size) moments = compute_tree_mass_moments(tree, pos_sorted, mass_sorted) @@ -371,6 +386,7 @@ def fn(pos, mass): own_mass = moments.mass[root] fr = build_coarse_frontier(tree, moments.mass, moments.center_of_mass) + # Coarse tree is always radix (internal remote-COM representation). rct = build_remote_coarse_tree(fr, ndev, bounds=bounds, axis_name=axis_name) rroot = as_index(jnp.argmin(jnp.asarray(rct.tree.parent))) @@ -574,9 +590,16 @@ def distributed_let_import( max_pair_queue: int = 8192, num_samples: int = 8, equalize: bool = True, + tree_type: str = "radix", axis_name: str = AXIS_NAME, ) -> ImportMetrics: - """Full LET pipeline through the halo import (decompose -> classify -> import).""" + """Full LET pipeline through the halo import (decompose -> classify -> import). + + ``tree_type`` selects the local per-device backend (``"radix"``, + ``"octree"``, or ``"kdtree"``); the coarse tree is always radix. + """ + + _validate_distributed_tree_type(tree_type) try: # stable across recent JAX versions from jax import shard_map @@ -609,12 +632,13 @@ def fn(pos, mass): p, m, c, cnt, ndev, output_capacity=output_capacity, axis_name=axis_name ) p, m = sanitize_padding(p, m, cnt) - tree, pos_sorted, mass_sorted, _ = build_tree( - p, m, bounds, return_reordered=True, leaf_size=leaf_size + tree, pos_sorted, mass_sorted = _build_local_tree( + p, m, bounds, tree_type=tree_type, leaf_size=leaf_size ) geom = compute_tree_geometry(tree, pos_sorted, max_leaf_size=leaf_size) moments = compute_tree_mass_moments(tree, pos_sorted, mass_sorted) fr = build_coarse_frontier(tree, moments.mass, moments.center_of_mass) + # Coarse tree is always radix (internal remote-COM representation). rct = build_remote_coarse_tree(fr, ndev, bounds=bounds, axis_name=axis_name) res = dual_tree_walk_cross_impl( tree, diff --git a/yggdrax/distributed/local_tree.py b/yggdrax/distributed/local_tree.py index b4ccedc..a695063 100644 --- a/yggdrax/distributed/local_tree.py +++ b/yggdrax/distributed/local_tree.py @@ -29,6 +29,79 @@ from .partition import equalize_domain, global_bounds, sfc_partition from .sharding import AXIS_NAME +# Backends whose topology satisfies the distributed contract (binary parent/ +# child links, compact contiguous ``node_ranges`` tiling, and ``nodes_by_level``) +# that :mod:`~yggdrax.distributed.cross_walk` and the moment stages consume. +# ``octree`` rides the radix binary topology (augmented with explicit cell +# buffers the cross-walk ignores); ``kdtree`` uses the leaf-only bucket build +# (:func:`yggdrax.kdtree.build_leaf_kdtree`), which stores all points in leaves +# and exposes the same contract (the heap KD-tree used by the KNN kernels does +# not, and is not used here). +DISTRIBUTED_TREE_TYPES: tuple[str, ...] = ("radix", "octree", "kdtree") + + +def _validate_distributed_tree_type(tree_type: str) -> None: + """Raise a clear error for tree types the distributed path cannot handle.""" + + if tree_type not in DISTRIBUTED_TREE_TYPES: + supported = ", ".join(f"'{t}'" for t in DISTRIBUTED_TREE_TYPES) + raise ValueError( + f"Unsupported tree_type '{tree_type}'. Supported: ({supported})." + ) + + +def _build_local_tree( + positions: Array, + masses: Array, + bounds: tuple[Array, Array], + *, + tree_type: str, + leaf_size: int, +): + """Backend-dispatched local tree build, safe to call inside ``shard_map``. + + Returns ``(topology, positions_sorted, masses_sorted)``. Only *un-jitted* + impls are used so the build traces inside a manual ``shard_map`` body (a + nested top-level ``@jax.jit`` would re-enter with a mesh mismatch). All + supported backends expose the parent/child/``node_ranges``/``nodes_by_level`` + contract downstream stages rely on: ``radix`` directly, ``octree`` by + augmenting the radix topology with explicit cell buffers, and ``kdtree`` via + the leaf-only bucket build (all points in leaves). + """ + + _validate_distributed_tree_type(tree_type) + if tree_type == "octree": + # Local imports: avoid importing the public tree module (which pulls in + # the kdtree/octree backends) at package import time. + from ..octree import augment_radix_topology_with_octree + from ..tree import _ADAPTIVE_OCTREE_REFINEMENT_DEFAULTS, _build_octree_result + + topo, pos_sorted, mass_sorted, _inv = _build_octree_result( + positions, + masses, + build_mode="adaptive", + bounds=bounds, + return_reordered=True, + workspace=None, + return_workspace=False, + leaf_size=leaf_size, + **_ADAPTIVE_OCTREE_REFINEMENT_DEFAULTS, + ) + return augment_radix_topology_with_octree(topo), pos_sorted, mass_sorted + + if tree_type == "kdtree": + from ..dtypes import INDEX_DTYPE + from ..kdtree import build_leaf_kdtree + + topo = build_leaf_kdtree(positions, leaf_size=leaf_size) + pidx = jnp.asarray(topo.particle_indices, dtype=INDEX_DTYPE) + return topo, positions[pidx], masses[pidx] + + tree, pos_sorted, mass_sorted, _inv = build_tree( + positions, masses, bounds, return_reordered=True, leaf_size=leaf_size + ) + return tree, pos_sorted, mass_sorted + @dataclass class DistributedTreeMoments: @@ -74,18 +147,21 @@ def build_local_moments( *, bounds: tuple[Array, Array], leaf_size: int, + tree_type: str = "radix", axis_name: str = AXIS_NAME, ): """Build this device's local tree and gather the shared top tree. Returns ``(node_mass, node_com, domain_mass, domain_com, top_mass, top_com)`` for this device. ``bounds`` must be the *global* box (shared - across devices) so geometry/Morton codes are consistent. + across devices) so geometry/Morton codes are consistent. ``tree_type`` + selects the local backend (``"radix"`` or ``"octree"``; see + :data:`DISTRIBUTED_TREE_TYPES`). """ positions, masses = sanitize_padding(positions, masses, count) - tree, pos_sorted, mass_sorted, _inv = build_tree( - positions, masses, bounds, return_reordered=True, leaf_size=leaf_size + tree, pos_sorted, mass_sorted = _build_local_tree( + positions, masses, bounds, tree_type=tree_type, leaf_size=leaf_size ) # Geometry is computed for completeness / downstream MAC use; keep the leaf # cap bounded so no num_particles-sized staging buffer is allocated. @@ -94,10 +170,12 @@ def build_local_moments( node_mass = moments.mass # [total_nodes] node_com = moments.center_of_mass # [total_nodes, 3] - # Root (node 0) of the local LBVH spans the whole domain -> its moment is - # the domain's aggregate (coarsest) multipole. - domain_mass = node_mass[0] - domain_com = node_com[0] + # The tree root spans the whole domain -> its moment is the domain's + # aggregate (coarsest) multipole. Locate it by parent link (== node 0 for + # the radix LBVH; robust for the octree-augmented topology too). + root = jnp.argmin(jnp.asarray(tree.parent)) + domain_mass = node_mass[root] + domain_com = node_com[root] # Shared top tree: every device learns every domain's coarse moment. top_mass = jax.lax.all_gather(domain_mass[None], axis_name, tiled=True) # [ndev] @@ -114,6 +192,7 @@ def distributed_tree_moments( output_capacity: int, num_samples: int = 8, equalize: bool = True, + tree_type: str = "radix", axis_name: str = AXIS_NAME, ) -> DistributedTreeMoments: """Decompose, build per-GPU local trees, and gather the shared top tree. @@ -122,8 +201,12 @@ def distributed_tree_moments( the local tree build into one ``shard_map``. ``positions`` (``[N, 3]``) and ``masses`` (``[N]``) are sharded evenly over the mesh (``N`` divisible by ``mesh.size``). Uses the *global* bounding box for consistent geometry. + ``tree_type`` selects the per-device backend (``"radix"`` or ``"octree"``; + see :data:`DISTRIBUTED_TREE_TYPES`). """ + _validate_distributed_tree_type(tree_type) + try: # stable across recent JAX versions from jax import shard_map except ImportError: # pragma: no cover @@ -148,7 +231,14 @@ def fn(pos, mass): p, m, c, cnt, ndev, output_capacity=output_capacity, axis_name=axis_name ) nm, nc, dm, dc, tm, tc = build_local_moments( - p, m, cnt, ndev, bounds=bounds, leaf_size=leaf_size, axis_name=axis_name + p, + m, + cnt, + ndev, + bounds=bounds, + leaf_size=leaf_size, + tree_type=tree_type, + axis_name=axis_name, ) return nm, nc, dm[None], dc[None], tm, tc, cnt[None] @@ -176,6 +266,7 @@ def fn(pos, mass): __all__ = [ + "DISTRIBUTED_TREE_TYPES", "DistributedTreeMoments", "build_local_moments", "distributed_tree_moments", diff --git a/yggdrax/kdtree.py b/yggdrax/kdtree.py index e43edc1..f6b5c6b 100644 --- a/yggdrax/kdtree.py +++ b/yggdrax/kdtree.py @@ -3,13 +3,16 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Literal +from typing import Literal, Optional import jax import jax.numpy as jnp +import numpy as np from beartype import beartype from jaxtyping import Array, jaxtyped +from .dtypes import INDEX_DTYPE + @dataclass(frozen=True) class KDTree: @@ -628,6 +631,280 @@ def build_kdtree(points: Array, *, leaf_size: int = 32) -> KDTree: ) +@dataclass(frozen=True) +class LeafKDTree: + """Leaf-only (bucket) median-split KD-tree satisfying the radix contract. + + Unlike :class:`KDTree` (a left-balanced heap with one pivot point per node, + used by the KNN query kernels), this container stores **all** points in + leaves and exposes the binary-tree topology contract shared with the radix + and octree backends: ``parent``/``left_child``/``right_child``, a compact + contiguous ``node_ranges`` (inclusive) whose leaf portion tiles ``[0, N)``, + and ``nodes_by_level``. That lets it flow through the generic geometry / + dual-tree-walk / moments pipeline (and the multi-GPU distributed path) + exactly like radix/octree, giving complete pair coverage -- the heap + KD-tree cannot, because its internal-node pivots are absent from the leaves. + + ``split_dim``/``split_value`` (per node; ``-1``/``nan`` on leaves) record the + median cut so KD-tree detection in the interaction backends still fires and + the tree remains a faithful median-split KD-tree. + """ + + points: Array + particle_indices: Array + node_ranges: Array + parent: Array + left_child: Array + right_child: Array + leaf_nodes: Array + nodes_by_level: Array + split_dim: Array + split_value: Array + num_particles: Array + use_morton_geometry: Array + num_internal_nodes: int + leaf_size: int + + @property + def num_points(self) -> int: + """Return the number of reference points in the tree.""" + + return int(self.points.shape[0]) + + @property + def dimension(self) -> int: + """Return spatial dimensionality of reference points.""" + + return int(self.points.shape[1]) + + +def _register_leaf_kdtree_pytree() -> None: + if getattr(LeafKDTree, "_yggdrax_pytree_registered", False): + return + + def flatten(tree: LeafKDTree): + children = ( + tree.points, + tree.particle_indices, + tree.node_ranges, + tree.parent, + tree.left_child, + tree.right_child, + tree.leaf_nodes, + tree.nodes_by_level, + tree.split_dim, + tree.split_value, + tree.num_particles, + tree.use_morton_geometry, + ) + aux = (tree.leaf_size, tree.num_internal_nodes) + return children, aux + + def unflatten(aux, children): + leaf_size, num_internal_nodes = aux + ( + points, + particle_indices, + node_ranges, + parent, + left_child, + right_child, + leaf_nodes, + nodes_by_level, + split_dim, + split_value, + num_particles, + use_morton_geometry, + ) = children + return LeafKDTree( + points=points, + particle_indices=particle_indices, + node_ranges=node_ranges, + parent=parent, + left_child=left_child, + right_child=right_child, + leaf_nodes=leaf_nodes, + nodes_by_level=nodes_by_level, + split_dim=split_dim, + split_value=split_value, + num_particles=num_particles, + use_morton_geometry=use_morton_geometry, + num_internal_nodes=num_internal_nodes, + leaf_size=leaf_size, + ) + + jax.tree_util.register_pytree_node(LeafKDTree, flatten, unflatten) + setattr(LeafKDTree, "_yggdrax_pytree_registered", True) + + +_register_leaf_kdtree_pytree() + + +def _leaf_kdtree_depth(n: int, leaf_size: int) -> int: + """Resolve the number of median-split levels for a leaf-only KD-tree. + + Chosen so every leaf bucket holds ``<= leaf_size`` points, but never so deep + that a leaf would be empty (``2**D <= n``). With balanced median splits this + guarantees non-empty leaves of size ``floor(n/2**D)..ceil(n/2**D)``. + """ + + if n <= leaf_size: + return 0 + need = -(-n // leaf_size) # ceil(n / leaf_size) + depth = max(1, (need - 1).bit_length()) # ceil(log2(need)) + return min(depth, (n).bit_length() - 1) # floor(log2(n)) -> 2**D <= n + + +def _build_leaf_kdtree_topology(points: Array, leaf_size: int) -> dict: + """Build a leaf-only median-split KD-tree via level-parallel bucket splits. + + All points end up in leaves; internal nodes are pure split planes. Returns a + dict of the radix-contract arrays. Tree *structure* (parent/child links, + leaf ids, level order) depends only on the static depth and is built as + NumPy constants; only the data-dependent arrays (permutation, ranges, split + metadata) are traced. + """ + + points = jnp.asarray(points) + n = int(points.shape[0]) + idx = INDEX_DTYPE + + depth = _leaf_kdtree_depth(n, int(leaf_size)) + num_leaves = 1 << depth + num_internal = num_leaves - 1 + total = 2 * num_leaves - 1 + + order = jnp.arange(n, dtype=idx) + bucket = jnp.zeros(n, dtype=idx) + split_dim_nodes = -jnp.ones(total, dtype=idx) + split_value_nodes = jnp.full(total, jnp.nan, dtype=points.dtype) + row = jnp.arange(n, dtype=idx) + + for level in range(depth): + nb = 1 << level + pts = points[order] + seg_max = jax.ops.segment_max(pts, bucket, num_segments=nb) + seg_min = jax.ops.segment_min(pts, bucket, num_segments=nb) + wdim = jnp.argmax(seg_max - seg_min, axis=-1).astype(idx) # [nb] + coord = jnp.take_along_axis( + pts, wdim[bucket][:, None], axis=-1 + ).squeeze(-1) + # Sort points by (bucket, coord) so each bucket's points are contiguous + # and ascending along the split dimension. + bucket, coord_sorted, order = jax.lax.sort( + (bucket, coord, order), dimension=0, num_keys=2 + ) + counts = jax.ops.segment_sum( + jnp.ones(n, dtype=idx), bucket, num_segments=nb + ) + starts = jnp.cumsum(counts) - counts + within = row - starts[bucket] + half = (counts + 1) // 2 # ceil -> left child gets the larger half + go_right = (within >= half[bucket]).astype(idx) + node_ids = (nb - 1) + jnp.arange(nb, dtype=idx) + split_dim_nodes = split_dim_nodes.at[node_ids].set(wdim) + boundary = jnp.clip(starts + half, 0, n - 1) + split_value_nodes = split_value_nodes.at[node_ids].set(coord_sorted[boundary]) + bucket = bucket * 2 + go_right + + # Final stable grouping by leaf bucket -> contiguous leaf ranges. + bucket, order = jax.lax.sort((bucket, order), dimension=0, num_keys=1) + particle_indices = order + leaf_counts = jax.ops.segment_sum( + jnp.ones(n, dtype=idx), bucket, num_segments=num_leaves + ) + leaf_start = jnp.cumsum(leaf_counts) - leaf_counts + leaf_end_incl = leaf_start + leaf_counts - 1 + + # ---- structural arrays (depend only on depth -> NumPy constants) ---- + ids_np = np.arange(total) + parent_np = np.where(ids_np > 0, (ids_np - 1) // 2, -1) + left_np = 2 * np.arange(num_internal) + 1 + right_np = 2 * np.arange(num_internal) + 2 + leaf_nodes_np = num_internal + np.arange(num_leaves) + levels_np = np.floor(np.log2(ids_np + 1)).astype(np.int64) + nodes_by_level_np = np.lexsort((ids_np, levels_np)) # by level, then id + + # ---- node_ranges: leaves from bucket extents, internals bottom-up union ---- + node_ranges = jnp.zeros((total, 2), dtype=idx) + leaf_node_ids = jnp.asarray(leaf_nodes_np, dtype=idx) + node_ranges = node_ranges.at[leaf_node_ids, 0].set(leaf_start) + node_ranges = node_ranges.at[leaf_node_ids, 1].set(leaf_end_incl) + for level in range(depth - 1, -1, -1): + nb = 1 << level + node_ids_np = np.arange(nb) + (nb - 1) + lc = jnp.asarray(2 * node_ids_np + 1, dtype=idx) + rc = jnp.asarray(2 * node_ids_np + 2, dtype=idx) + node_ids = jnp.asarray(node_ids_np, dtype=idx) + node_ranges = node_ranges.at[node_ids, 0].set(node_ranges[lc, 0]) + node_ranges = node_ranges.at[node_ids, 1].set(node_ranges[rc, 1]) + + return { + "particle_indices": particle_indices, + "node_ranges": node_ranges, + "parent": jnp.asarray(parent_np, dtype=idx), + "left_child": jnp.asarray(left_np, dtype=idx), + "right_child": jnp.asarray(right_np, dtype=idx), + "leaf_nodes": leaf_node_ids, + "nodes_by_level": jnp.asarray(nodes_by_level_np, dtype=idx), + "split_dim": split_dim_nodes, + "split_value": split_value_nodes, + "num_internal": int(num_internal), + } + + +def build_leaf_kdtree( + points: Array, + *, + leaf_size: int = 32, + bounds: Optional[tuple[Array, Array]] = None, +) -> LeafKDTree: + """Build a leaf-only (bucket) median-split KD-tree. + + Every point is stored in a leaf; internal nodes are pure split planes. The + result satisfies the same topology contract as the radix/octree backends, + so it feeds the generic geometry, dual-tree-walk, moments, and distributed + pipelines and gives complete pair coverage. ``bounds`` is accepted for a + uniform builder signature but ignored (median splits are data-driven). + + Parameters + ---------- + points + Reference points of shape ``(n_points, dim)``. + leaf_size + Maximum points per leaf bucket. + + Returns + ------- + LeafKDTree + Leaf-only KD-tree container. + """ + + del bounds # median splits are data-driven; no fixed box needed + points_arr = _validate_points(points) + if leaf_size < 1: + raise ValueError(f"leaf_size must be >= 1, received {leaf_size}") + + topo = _build_leaf_kdtree_topology(points_arr, int(leaf_size)) + n = int(points_arr.shape[0]) + return LeafKDTree( + points=points_arr, + particle_indices=topo["particle_indices"], + node_ranges=topo["node_ranges"], + parent=topo["parent"], + left_child=topo["left_child"], + right_child=topo["right_child"], + leaf_nodes=topo["leaf_nodes"], + nodes_by_level=topo["nodes_by_level"], + split_dim=topo["split_dim"], + split_value=topo["split_value"], + num_particles=jnp.asarray(n, dtype=INDEX_DTYPE), + use_morton_geometry=jnp.asarray(False), + num_internal_nodes=topo["num_internal"], + leaf_size=int(leaf_size), + ) + + def _query_neighbors_dense( tree: KDTree, queries_arr: Array, diff --git a/yggdrax/tree.py b/yggdrax/tree.py index cc162b8..85d961b 100644 --- a/yggdrax/tree.py +++ b/yggdrax/tree.py @@ -14,8 +14,8 @@ from . import _tree_impl from .bounds import infer_bounds from .dtypes import INDEX_DTYPE, as_index -from .kdtree import KDTree as KDTreeTopology -from .kdtree import build_kdtree +from .kdtree import LeafKDTree as KDTreeTopology +from .kdtree import build_leaf_kdtree from .octree import OctreeTopology, augment_radix_topology_with_octree from .protocols import MortonLeafBoundsProtocol @@ -884,7 +884,7 @@ def from_particles( f"'{build_mode}' for kdtree. Supported: ('adaptive',)" ) - topology = build_kdtree(positions, leaf_size=leaf_size) + topology = build_leaf_kdtree(positions, leaf_size=leaf_size) if return_reordered: idx = jnp.asarray(topology.particle_indices, dtype=INDEX_DTYPE) pos_sorted = positions[idx] From 731a924a71723b34391f47c01e41718f8afd13f6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:02:37 +0000 Subject: [PATCH 4/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/distributed/test_local_tree_backends.py | 37 +++++++++++++------ yggdrax/kdtree.py | 8 +--- 2 files changed, 28 insertions(+), 17 deletions(-) diff --git a/tests/distributed/test_local_tree_backends.py b/tests/distributed/test_local_tree_backends.py index b5f46a8..d5a1354 100644 --- a/tests/distributed/test_local_tree_backends.py +++ b/tests/distributed/test_local_tree_backends.py @@ -120,15 +120,20 @@ def test_coarse_tree_root_is_global(built, tree_type): equalize=True, tree_type=tree_type, ) - np.testing.assert_allclose(np.asarray(metrics.coarse_root_mass), mass.sum(), rtol=1e-3) + np.testing.assert_allclose( + np.asarray(metrics.coarse_root_mass), mass.sum(), rtol=1e-3 + ) # ---- cross-tree walk over leaf-only KD trees: complete disjoint source partition ---- def _kd_build(points, leaf_size=_LEAF): pts = jnp.asarray(points) tree, ps, _ms = _build_local_tree( - pts, jnp.ones(pts.shape[0], jnp.float32), infer_bounds(pts), - tree_type="kdtree", leaf_size=leaf_size, + pts, + jnp.ones(pts.shape[0], jnp.float32), + infer_bounds(pts), + tree_type="kdtree", + leaf_size=leaf_size, ) geom = compute_tree_geometry(tree, ps, max_leaf_size=leaf_size) return tree, geom @@ -150,16 +155,26 @@ def test_kdtree_cross_walk_partitions_sources(): def build(points): pts = jnp.asarray(points) tree, ps, _ = _build_local_tree( - pts, jnp.ones(pts.shape[0], jnp.float32), bounds, - tree_type="kdtree", leaf_size=_LEAF, + pts, + jnp.ones(pts.shape[0], jnp.float32), + bounds, + tree_type="kdtree", + leaf_size=_LEAF, ) return tree, compute_tree_geometry(tree, ps, max_leaf_size=_LEAF) t_tree, t_geom = build(tgt) s_tree, s_geom = build(src) res = dual_tree_walk_cross( - t_tree, t_geom, s_tree, s_geom, 0.5, mac_type="bh", - max_interactions_per_node=512, max_neighbors_per_leaf=512, max_pair_queue=16384, + t_tree, + t_geom, + s_tree, + s_geom, + 0.5, + mac_type="bh", + max_interactions_per_node=512, + max_neighbors_per_leaf=512, + max_pair_queue=16384, ) assert not bool(res.queue_overflow | res.far_overflow | res.near_overflow) @@ -191,9 +206,9 @@ def src_particles(node): node = int(parent[node]) o, c = int(near_off[row]), int(near_cnt[row]) covered = [] - for sn in far_sources + near_idx[o:o + c].tolist(): + for sn in far_sources + near_idx[o : o + c].tolist(): covered.extend(src_particles(sn)) - assert set(covered) == set(range(n_source)), ( - f"leaf_row {row}: covered {len(set(covered))}/{n_source}" - ) + assert set(covered) == set( + range(n_source) + ), f"leaf_row {row}: covered {len(set(covered))}/{n_source}" assert len(covered) == n_source, f"leaf_row {row}: overlap {len(covered)}" diff --git a/yggdrax/kdtree.py b/yggdrax/kdtree.py index f6b5c6b..73cf20b 100644 --- a/yggdrax/kdtree.py +++ b/yggdrax/kdtree.py @@ -786,17 +786,13 @@ def _build_leaf_kdtree_topology(points: Array, leaf_size: int) -> dict: seg_max = jax.ops.segment_max(pts, bucket, num_segments=nb) seg_min = jax.ops.segment_min(pts, bucket, num_segments=nb) wdim = jnp.argmax(seg_max - seg_min, axis=-1).astype(idx) # [nb] - coord = jnp.take_along_axis( - pts, wdim[bucket][:, None], axis=-1 - ).squeeze(-1) + coord = jnp.take_along_axis(pts, wdim[bucket][:, None], axis=-1).squeeze(-1) # Sort points by (bucket, coord) so each bucket's points are contiguous # and ascending along the split dimension. bucket, coord_sorted, order = jax.lax.sort( (bucket, coord, order), dimension=0, num_keys=2 ) - counts = jax.ops.segment_sum( - jnp.ones(n, dtype=idx), bucket, num_segments=nb - ) + counts = jax.ops.segment_sum(jnp.ones(n, dtype=idx), bucket, num_segments=nb) starts = jnp.cumsum(counts) - counts within = row - starts[bucket] half = (counts + 1) // 2 # ceil -> left child gets the larger half