diff --git a/docs/small_leaves_2026-09.md b/docs/small_leaves_2026-09.md index 03052846..20c6a025 100644 --- a/docs/small_leaves_2026-09.md +++ b/docs/small_leaves_2026-09.md @@ -38,8 +38,9 @@ both serialised by the atomic-add lowering) = 169 ms of the 410 ms leaf-64 step, 2. **Contention-free chunk reduction** (`_m2l.py::_chunk_segment_scatter_add`): segmented `associative_scan` + out-of-bounds sink (`mode="drop"`, unique in-bounds indices). Leaf 64: 410 -> 237 ms/step; leaf 128: 188 -> 145; leaf 32: 1168 -> 649. -3. **Target-tiled CSR M2L Pallas kernel** (`jaccpot/pallas/m2l_real_csr.py`, flag - `JACCPOT_STATIC_STRICT_FUSED_M2L_CSR=1`): one program per target owns its local row; rotations +3. **Target-tiled CSR M2L Pallas kernel** (`jaccpot/pallas/m2l_real_csr.py`; the DEFAULT on Ampere+ + GPUs since 2026-09-10, `JACCPOT_STATIC_STRICT_FUSED_M2L_CSR=0` restores the chunked pure-JAX lanes, + `JACCPOT_M2L_CSR_INTERPRET=1` runs it interpreted on CPU): one program per target owns its local row; rotations assembled on chip from the two alignment angles in the centred (degree, m) layout, degree-only z-core, no per-pair operand in HBM. Microbench on an idle A100 (`m2l_csr_microbench.json`): diff --git a/docs/tree_walk_2026-09.md b/docs/tree_walk_2026-09.md new file mode 100644 index 00000000..c2e8893f --- /dev/null +++ b/docs/tree_walk_2026-09.md @@ -0,0 +1,83 @@ +# The per-step tree walk on the fused single-GPU lane (2026-09-10) + +Plan `~/.claude/plans/ok-then-please-make-replicated-rain.md` (Odisseo box). Follows +`small_leaves_2026-09.md`, which ended with the far field solved (CSR M2L kernel) and the traced dual-tree +walk as the remaining wall: ~100 ms of a 177 ms step at leaf 64. yggdrax PR #74 carries the walk side +(`docs/traversal_walk_cost.md` there has the isolated-walk table); this note is the jaccpot side. + +## What the walk cost and why + +yggdrax's `_dual_tree_walk_impl`, one call per step at the full traced queue, spent its time on its OUTPUT +layout, not on the traversal: dense per-node rows (`total_nodes x max_interactions_per_node`, 820 MB at +leaf 64) carried through the `while_loop` and copied in and out by `lax.cond` identity branches every +round, a full-queue argsort per round to place each pair in its rows, four `segment_sum`s per round, and a +post-loop flatten over `nodes x K` slots. Isolated on one tree (A100, N=200k, theta 0.6) it took 503 ms per +walk at leaf 64 and 1291 ms at leaf 32. yggdrax's flat-emission `dual_tree_walk_mutual` -- same wavefront, +each unordered pair appended once to a flat list with a cumsum -- produces the same far and near pairs as +SETS in 13-40 ms. The scatter-lowering promise (`unique_indices`) is not a lever (0.10 ms per 1M scatter +either way); the dense rows, the conditionals and the sorts are. + +## The lane + +The flat walk is the DEFAULT of the strict fused lane (since 2026-09-10; `JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK=0` +restores the traced dual walk). `_build_dual_tree_artifacts_split_strict_streamed` routes to +`_build_flat_walk_artifacts_strict_streamed` (`runtime/_interaction_cache.py`), which calls +`dual_tree_walk_mutual` with the dual walk's own `mac_extents` (`_build_mac_extents(...)[0]`) and +`mac_type`, then: + +* un-mutualises the far pairs INTERLEAVED (`[b->a, a->b]` per canonical pair) so the live pairs stay a + prefix -- every M2L consumer masks `idx < far_pair_count`, and a concatenation would silently drop the + second direction; capacity is the compact far-pair cap; +* builds the leaf neighbour CSR from the directed near pairs with one stable argsort by target leaf and + `searchsorted` offsets; width is the neighbour-edge cap, so eager and traced carries match without a pad; +* eager: a queue ladder on `queue_overflow`; a far or near overflow doubles the capacity when the caller did + not name it (`JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP` / `JACCPOT_LARGE_N_NEIGHBOR_EDGE_PROFILE_FIXED_CAP` + absent from the environment; floors 131072 and 2^21 directed pairs, ceilings 2^26 and 2^28) and raises + naming the cap when it did (a named cap is never widened). The widths the eager pass settled on are handed + to the traced refresh -- and to any later eager prepare -- as floors by `_strict_fused_capacity_handoff`, + so the `lax.scan` carry shapes match and never shrink; + traced: any overflow saturates `far_pair_count` to the capacity, which trips `strict_run_v2`'s existing + saturation guard -- a truncated refresh is fatal, not silent; +* always emits the capacity report (the treecode graft's early return left the guard dark) with + `peak_wavefront`; `_strict_fused_capacity_handoff` sizes the traced queue as pow2(1.5 x peak). In the real + fused geometry the leaf-64 peak is 466,626 pairs over 24 rounds, so the traced queue is 2^20. + +Supported: `mac_type` bh/dehnen, `pair_policy=None`, the flat compact far-pair layout. A configuration outside +that (the treecode walk requested, `mac_type='engblom'`, a solver-owned pair policy such as `dehnen_error`, +`JACCPOT_STATIC_STRICT_FUSED_FLAT_COMPACT_FAR_PAIRS=0`) falls back to the dual walk QUIETLY while the flag is +merely defaulted, and RAISES when the flag was set to `1` explicitly -- then the caller asked for a walk it +cannot have, and silence would hand it the wrong one (`tests/unit/runtime/test_flat_walk_default_dispatch.py`). +Indices: the harness sets +`YGGDRAX_INDEX_PRECISION=int32` and `JACCPOT_INDEX_PRECISION=int32` for the fused lane (read at import; set +both, yggdrax falls back to jaccpot's variable but not the reverse). + +## Per step (`strict_run_v2`, N=200k Plummer, p=4, idle A100, no foreign process; flat walk + CSR M2L + int32) + +| leaf | theta | small-leaves start | + CSR M2L | **+ flat walk** | eval-only | near kernel | M2L kernel | upward | downward (M2L + walk + L2L) | aggL2 | +|---|---|---|---|---|---|---|---|---|---|---| +| 256 | 0.6 | 131.8 | 120.3 | **96.2** | 68.6 | 79.9 | 1.7 | 5.4 | 8.4 | 8.27e-4 | +| 128 | 0.6 | 187.5 | 122.3 | **71.3** | 42.7 | 48.5 | 5.1 | 8.5 | 12.1 | 1.15e-3 | +| 64 | 0.6 | 409.6 | 176.7 | **63.3** | 23.6 | 26.5 | 13.0 | 10.5 | 25.4 | 1.20e-3 | +| 32 | 0.6 | 1168 | 509.4 | **82.5** | 16.2 | -- | -- | 12.9 | 53.7 | 1.30e-3 | +| 64 | 0.8 | 192.0 | 113.5 | **41.6** | 21.3 | 13.0 | 7.1 | 10.8 | 16.9 | 5.33e-3 | +| 32 | 0.8 | -- | 364.1 | **53.6** | 12.8 | -- | -- | 11.0 | 33.5 | 1.44e-2 | + +Launches per step 3.8-4.6k (from 7-32k). Forces identical to the dual-walk lane to 4 digits at every leaf +(the lists are the same sets; only the fp32 order changes); the #333 truncation check passes with the flag +(`tests/integration/test_strict_run_v2_refresh_capacity.py` is parametrised over it). + +**The per-step optimum moved to leaf 64: 63.3 ms against 96.2 at leaf 256 with the same lanes and 120 before +them -- 1.9x per step at 200k, 6.5x against the leaf-64 step the small-leaves work started from.** At leaf 32 +the walk is again the largest item (downward 53.7 with M2L ~25): the wavefront's long thin tail (46-61 rounds +at ~0.15 ms each) and a 2^21 queue; a narrow-width branch for the tail rounds (plan Tier 3) is the next lever +there, not at leaf 64. + +## Traps recorded + +* `tests/conftest.py` used to put the sibling `/export/home/tbuck/yggdrax` checkout (a paper branch) at the + front of `sys.path` unconditionally: every local jaccpot pytest run exercised THAT yggdrax. It now honours + `YGGDRAX_WORKTREE`; the bench harness's `sitecustomize` repoints the yggdrax editable finder the same way. +* `autocvd -l -o -q` can return an empty device list when every card is busy; the job then runs on CPU and + GPU-only tests skip silently. Check the device line in the log. +* The strict fused prepared-eval seam exists only above the large-N threshold on a GPU, so the lane's wiring + test runs in the GPU suite at N=70k. diff --git a/jaccpot/runtime/_interaction_cache.py b/jaccpot/runtime/_interaction_cache.py index 6fdd721d..cea2d741 100644 --- a/jaccpot/runtime/_interaction_cache.py +++ b/jaccpot/runtime/_interaction_cache.py @@ -912,6 +912,11 @@ def _record(name: str, start: Optional[float]) -> None: _STRICT_STREAMED_FAR_PAIR_FLOOR = 131_072 _STRICT_STREAMED_RETRY_LIMIT = 1 << 25 _STRICT_STREAMED_RETRY_ATTEMPTS = 12 +# Flat-walk lane: eager floors and ceilings for the two capacities the caller did +# NOT name (a named cap is never widened -- #333's rule). Directed pair counts. +_FLAT_WALK_NEAR_EDGE_FLOOR = 1 << 21 +_FLAT_WALK_NEAR_EDGE_LIMIT = 1 << 28 +_FLAT_WALK_FAR_PAIR_LIMIT = 1 << 26 def _strict_streamed_retry_diag(grew: list[str]) -> None: @@ -954,6 +959,7 @@ def _build_dual_tree_artifacts_split_strict_streamed( policy_state: Optional[AdaptivePolicyState], capacity_report: Optional[Callable[[dict], None]] = None, max_neighbors_per_leaf_override: Optional[int] = None, + flat_walk_capacity_floor: Optional[dict] = None, ) -> _DualTreeArtifacts: """Strict static fast-lane: single compact shared far+near build call. @@ -1013,6 +1019,11 @@ def _build_dual_tree_artifacts_split_strict_streamed( (far-pair count, longest neighbour row, total edges). ``None`` skips it. max_neighbors_per_leaf_override : Optional[int] Floor for the per-leaf neighbour cap; only ever raises it. + flat_walk_capacity_floor : Optional[dict] + Floors for the flat-walk lane's directed far and near capacities + (``compact_far_pair_capacity`` / ``near_edge_capacity``), carried over from + the eager pass so the traced refresh builds the same widths. Applies only + to a capacity the caller did not name in the environment. Returns ------- @@ -1058,7 +1069,10 @@ def _build_dual_tree_artifacts_split_strict_streamed( compact_far_pair_capacity = None if flat_compact_enabled: compact_far_pair_capacity = int( - os.environ.get("JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP", "131072") + os.environ.get( + "JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP", + str(_STRICT_STREAMED_FAR_PAIR_FLOOR), + ) ) if compact_far_pair_capacity <= 0: raise ValueError( @@ -1071,6 +1085,89 @@ def _build_dual_tree_artifacts_split_strict_streamed( treecode_enabled = os.environ.get( "JACCPOT_STATIC_STRICT_FUSED_TREECODE_WALK", "0" ) not in ("0", "false", "False", "off", "OFF") + # DEFAULT since 2026-09-10: yggdrax's flat-emission wavefront walk with the + # dual walk's own MAC extents -- the same lists as sets at a fraction of the + # per-step cost (leaf-64 step at N=200k: 177 -> 63 ms). Set + # JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK=0 for the traced dual walk. A + # configuration the flat walk cannot carry (the treecode walk requested, a + # solver-owned pair policy, a MAC other than bh/dehnen, the non-flat far-pair + # layout) falls back to the dual walk quietly while the flag is merely + # defaulted; an EXPLICIT "1" against such a configuration raises, because then + # the caller asked for a walk it cannot have and silence would hand it the + # wrong one. See _build_flat_walk_artifacts_strict_streamed. + flat_walk_raw = os.environ.get("JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK") + flat_walk_explicit = flat_walk_raw is not None + flat_walk_enabled = (flat_walk_raw if flat_walk_explicit else "1") not in ( + "0", + "false", + "False", + "off", + "OFF", + ) + flat_walk_blocker: Optional[str] = None + if flat_walk_enabled: + if treecode_enabled: + flat_walk_blocker = ( + "JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK and " + "JACCPOT_STATIC_STRICT_FUSED_TREECODE_WALK are both set; pick one walk." + ) + elif pair_policy is not None or policy_state is not None: + flat_walk_blocker = ( + "JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK cannot carry a solver-owned " + "pair policy (mac_type='dehnen_error' / adaptive_error_model=" + "'dehnen_paper'): the flat walk's acceptance test is the geometric " + "MAC on per-node extents with no policy seam. Unset the env flag, " + "or use mac_type='dehnen'." + ) + elif str(mac_type) not in ("bh", "dehnen"): + flat_walk_blocker = ( + "JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK supports mac_type 'bh' and " + f"'dehnen' only, got {mac_type!r}." + ) + elif compact_far_pair_capacity is None: + flat_walk_blocker = ( + "JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK needs the flat compact " + "far-pair layout (JACCPOT_STATIC_STRICT_FUSED_FLAT_COMPACT_FAR_PAIRS=1)." + ) + if flat_walk_blocker is not None: + if flat_walk_explicit: + raise RuntimeError(flat_walk_blocker) + flat_walk_enabled = False + if flat_walk_enabled: + assert compact_far_pair_capacity is not None + floor = dict(flat_walk_capacity_floor or {}) + # A cap the caller NAMED is a deliberate memory bound: honoured exactly, + # never widened (#333). An unnamed one starts at its floor -- raised to + # what an earlier eager pass needed -- and the eager ladder grows it. + near_edge_env = os.environ.get( + "JACCPOT_LARGE_N_NEIGHBOR_EDGE_PROFILE_FIXED_CAP" + ) + near_edge_named = near_edge_env is not None + if near_edge_named: + near_edge_capacity = int(near_edge_env) + else: + near_edge_capacity = max( + _FLAT_WALK_NEAR_EDGE_FLOOR, int(floor.get("near_edge_capacity") or 0) + ) + far_named = "JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP" in os.environ + far_capacity = int(compact_far_pair_capacity) + if not far_named: + far_capacity = max( + far_capacity, int(floor.get("compact_far_pair_capacity") or 0) + ) + return _build_flat_walk_artifacts_strict_streamed( + tree=tree, + geometry=geometry, + theta=theta, + mac_type=mac_type, + dehnen_radius_scale=dehnen_radius_scale, + compact_far_pair_capacity=far_capacity, + near_edge_capacity=near_edge_capacity, + max_pair_queue=max_pair_queue_resolved, + capacity_report=capacity_report, + far_named=far_named, + near_edge_named=near_edge_named, + ) if treecode_enabled: if pair_policy is not None or policy_state is not None: # The treecode walk evaluates its own device-resident `_mac_ok` from @@ -1626,6 +1723,316 @@ def _env_int(name, default): ) +def _build_flat_walk_artifacts_strict_streamed( + *, + tree: Tree, + geometry: TreeGeometry, + theta: float, + mac_type: MACType, + dehnen_radius_scale: float, + compact_far_pair_capacity: int, + near_edge_capacity: int, + max_pair_queue: Optional[int], + capacity_report: Optional[Callable[[dict], None]] = None, + far_named: bool = True, + near_edge_named: bool = True, +) -> _DualTreeArtifacts: + """Far pairs and leaf neighbours from yggdrax's flat-emission wavefront walk. + + The strict fused lane's alternative to ``build_compact_far_pairs_and_leaf_ + neighbor_lists`` (plan "tree walk", 2026-09-10). Measured in isolation on an + A100 at N=200k / leaf 64 / queue 2^20 the traced dual-tree walk costs 503 ms + per call and ``dual_tree_walk_mutual`` 40 ms (26 ms with int32 indices) for + IDENTICAL far and near pair counts: the dual walk carries dense per-node + output rows (``total_nodes x max_interactions_per_node`` -- 820 MB at leaf 64 + -- plus ``num_leaves x max_neighbors_per_leaf``) through the loop, sorts every + round to place each pair in its rows, and flattens ``total_nodes x K`` slots + afterwards; the mutual walk appends each unordered pair once to a flat list + with a cumsum. Fed the dual walk's own ``mac_extents`` and ``mac_type`` it + produces the same lists as SETS (``tests/unit/test_dual_tree_walk_mutual.py`` + in yggdrax pins that), so only the fp32 summation order of a force changes. + + What this builder adds on top of the walk: + + * far pairs un-mutualised and INTERLEAVED -- ``[b->a, a->b]`` per canonical + pair -- so the live pairs are a prefix of the buffer. Every M2L consumer + masks ``idx < far_pair_count`` (``kernels/_m2l.py``, the CSR Pallas lane, + ``_large_n_grad``); concatenating the two directions would put the second + half beyond that prefix and silently drop it. Capacity is + ``compact_far_pair_capacity`` (``2 x far_cap``), eager and traced alike, as + ``_compact_prefix_with_fixed_capacity`` does for the dual walk. + * the leaf neighbour CSR from the directed near pairs with one stable argsort + by target leaf and ``searchsorted`` offsets -- no per-node row buffer, no + duplicate-index scatter. Width is ``near_edge_capacity`` (``2 x near_cap``), + which is exactly the fixed edge cap ``_trim_radix_fast_lane_neighbor_list`` + pads to, so eager and traced carries match without a pad. + * overflow: eager, the queue is doubled on ``queue_overflow``; a far or near + overflow doubles the capacity when the caller did not name it + (``far_named`` / ``near_edge_named`` False) and raises naming the cap when + it did (a named cap is never widened -- #333's rule). Traced, + the three flags saturate ``far_pair_count`` to the capacity, which trips + the strict runner's existing ``far_pair_count < capacity`` arm of the + saturation guard (``fmm_strict_run.py``), so a truncated refresh is fatal + rather than silent. The capacity report is ALWAYS emitted (the treecode + graft's early return left that guard dark) and carries ``peak_wavefront`` + so the traced queue can be sized from data. + + Parameters + ---------- + tree : Tree + Built static-radix tree (leaves are the last ``num_leaves`` nodes). + geometry : TreeGeometry + Node centres and extents the MAC is evaluated against. + theta : float + Opening angle. + mac_type : MACType + ``bh`` or ``dehnen``; ``engblom`` is supported by the walk but not by the + strict lane's callers and is refused at the seam. + dehnen_radius_scale : float + Radius inflation for the Dehnen MAC. + compact_far_pair_capacity : int + Directed far-pair capacity (even). + near_edge_capacity : int + Directed near-pair capacity, i.e. the neighbour-edge cap. + max_pair_queue : Optional[int] + Wavefront capacity; ``None`` starts the eager ladder at the floor. + capacity_report : Optional[Callable[[dict], None]] + Receives the capacities used (and, eager only, what was observed). + far_named : bool + Whether ``compact_far_pair_capacity`` was named by the caller + (``JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP``); a named cap raises + on overflow, an unnamed one is doubled eagerly up to + ``_FLAT_WALK_FAR_PAIR_LIMIT``. + near_edge_named : bool + Same for ``near_edge_capacity`` + (``JACCPOT_LARGE_N_NEIGHBOR_EDGE_PROFILE_FIXED_CAP``; ceiling + ``_FLAT_WALK_NEAR_EDGE_LIMIT``). + + Returns + ------- + _DualTreeArtifacts + Compact far pairs and the leaf neighbour list; no interactions payload. + + Raises + ------ + ValueError + If a capacity is not even / positive. + RuntimeError + Eager far or near overflow of a named cap (or of an unnamed one past its + ceiling), or a queue that will not fit within the retry ceiling. + """ + from yggdrax._interactions_impl import _build_mac_extents + from yggdrax.interactions import dual_tree_walk_mutual + + if int(compact_far_pair_capacity) <= 0 or int(compact_far_pair_capacity) % 2: + raise ValueError( + "JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP must be a positive " + f"even number on the flat-walk lane, got {compact_far_pair_capacity}" + ) + if int(near_edge_capacity) <= 0 or int(near_edge_capacity) % 2: + raise ValueError( + "JACCPOT_LARGE_N_NEIGHBOR_EDGE_PROFILE_FIXED_CAP must be a positive even " + f"number on the flat-walk lane, got {near_edge_capacity}" + ) + far_cap = int(compact_far_pair_capacity) // 2 + near_cap = int(near_edge_capacity) // 2 + + topo = tree.topology + num_internal = int(topo.left_child.shape[0]) + total_nodes = int(topo.parent.shape[0]) + num_leaves = total_nodes - num_internal + idx = topo.parent.dtype + left_full = jnp.concatenate( + [jnp.asarray(topo.left_child, idx), jnp.full((num_leaves,), -1, idx)] + ) + right_full = jnp.concatenate( + [jnp.asarray(topo.right_child, idx), jnp.full((num_leaves,), -1, idx)] + ) + root_idx = jnp.argmin(topo.parent).astype(idx) + centers = jnp.asarray(geometry.center) + mac_extents, _leaf_extents = _build_mac_extents( + topo.parent, geometry, num_internal, str(mac_type), float(dehnen_radius_scale) + ) + mac_extents = jnp.asarray(mac_extents, dtype=centers.dtype) + + queue = ( + _STRICT_STREAMED_QUEUE_FLOOR if max_pair_queue is None else int(max_pair_queue) + ) + grew: list[str] = [] + walk = None + traced = False + # Cap growth (unnamed caps only) is bounded by the ceilings; queue growth by + # the retry budget. Only queue doublings count against that budget, so a tiny + # unnamed cap cannot exhaust it. + queue_attempts = 0 + while True: + walk = dual_tree_walk_mutual( + left_full, + right_full, + centers, + mac_extents, + float(theta), + root_idx, + max_pair_queue=int(queue), + far_cap=far_cap, + near_cap=near_cap, + mac_type=str(mac_type), + ) + traced = isinstance(walk.queue_overflow, Tracer) + if traced: + break + far_ovf = bool(walk.far_overflow) + near_ovf = bool(walk.near_overflow) + queue_ovf = bool(walk.queue_overflow) + if far_ovf and (far_named or 2 * far_cap >= _FLAT_WALK_FAR_PAIR_LIMIT): + raise RuntimeError( + "flat-walk far pairs overflowed: capacity " + f"{2 * far_cap} directed pairs" + + ( + " (the caller named it, so it is not widened here). Raise " + if far_named + else f" (the eager ceiling is {_FLAT_WALK_FAR_PAIR_LIMIT}). Set " + ) + + "JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP." + ) + if near_ovf and (near_edge_named or 2 * near_cap >= _FLAT_WALK_NEAR_EDGE_LIMIT): + raise RuntimeError( + "flat-walk near pairs overflowed: capacity " + f"{2 * near_cap} directed pairs" + + ( + " (the caller named it, so it is not widened here). Raise " + if near_edge_named + else f" (the eager ceiling is {_FLAT_WALK_NEAR_EDGE_LIMIT}). Set " + ) + + "JACCPOT_LARGE_N_NEIGHBOR_EDGE_PROFILE_FIXED_CAP." + ) + if far_ovf: + grew.append(f"compact_far_pair_capacity {2 * far_cap}->{4 * far_cap}") + far_cap *= 2 + if near_ovf: + grew.append(f"near_edge_capacity {2 * near_cap}->{4 * near_cap}") + near_cap *= 2 + if queue_ovf: + grown = int(queue) * 2 + queue_attempts += 1 + if ( + queue_attempts >= _STRICT_STREAMED_RETRY_ATTEMPTS + or grown > _STRICT_STREAMED_RETRY_LIMIT + ): + raise RuntimeError( + "max_pair_queue overflowed on the flat wavefront walk and " + f"re-planning did not fit it: grew to {queue} (ceiling " + f"{_STRICT_STREAMED_RETRY_LIMIT}) over {queue_attempts} attempts; " + "the walk needed a peak wavefront of " + f"{int(walk.peak_wavefront)}. Pass jaccpot.TraversalOverrides(" + "max_pair_queue=...) explicitly." + ) + grew.append(f"max_pair_queue {queue}->{grown}") + queue = grown + if not (far_ovf or near_ovf or queue_ovf): + if grew: + _strict_streamed_retry_diag(grew) + break + assert walk is not None + compact_far_pair_capacity = 2 * far_cap + near_edge_capacity = 2 * near_cap + + # --- far pairs: directed, interleaved, prefix-live, capacity-width --- + far_live = jnp.arange(far_cap, dtype=idx) < walk.far_count + fa = jnp.where(far_live, walk.far_a, -1).astype(idx) + fb = jnp.where(far_live, walk.far_b, -1).astype(idx) + far_sources = jnp.stack([fb, fa], axis=1).reshape((2 * far_cap,)) + far_targets = jnp.stack([fa, fb], axis=1).reshape((2 * far_cap,)) + far_tags = jnp.full((2 * far_cap,), -1, dtype=idx) + any_overflow = walk.far_overflow | walk.near_overflow | walk.queue_overflow + # Saturate on ANY overflow: the strict runner's guard tests + # ``far_pair_count < compact_far_pair_capacity`` and this is how the near and + # queue flags reach it under trace. Eager overflow raised above, so this only + # bites inside the compiled scan. + far_pair_count = jnp.where( + any_overflow, + jnp.asarray(2 * far_cap, idx), + (2 * walk.far_count).astype(idx), + ) + compact_far_pairs = CompactTaggedFarPairs( + sources=far_sources, + targets=far_targets, + tags=far_tags, + far_pair_count=far_pair_count, + ) + + # --- near pairs: directed, one stable sort by target leaf, CSR --- + near_live = jnp.arange(near_cap, dtype=idx) < walk.near_count + na = jnp.where(near_live, walk.near_a, 0).astype(idx) + nb = jnp.where(near_live, walk.near_b, 0).astype(idx) + tgt = jnp.concatenate([na, nb]) + src = jnp.concatenate([nb, na]) + valid = jnp.concatenate([near_live, near_live]) + # static radix: leaves are the last ``num_leaves`` nodes + tgt_leaf = tgt - jnp.asarray(num_internal, idx) + key = jnp.where(valid, tgt_leaf, jnp.asarray(num_leaves, idx)) + perm = jnp.argsort(key, stable=True) + sorted_key = key[perm] + neighbors = jnp.where(valid[perm], src[perm], jnp.asarray(0, idx)) + offsets = jnp.searchsorted( + sorted_key, jnp.arange(num_leaves + 1, dtype=idx), side="left" + ).astype(idx) + counts = offsets[1:] - offsets[:-1] + leaf_nodes = jnp.arange(num_internal, total_nodes, dtype=idx) + neighbor_list = NodeNeighborList( + offsets=offsets, + neighbors=neighbors, + leaf_indices=leaf_nodes, + counts=counts, + particle_order_leaf_indices=leaf_nodes, + particle_order_to_native_leaf=jnp.arange(num_leaves, dtype=idx), + neighbor_leaf_positions=jnp.zeros((num_leaves, 0), dtype=idx), + target_block_leaf_ids=jnp.zeros((0,), dtype=idx), + target_block_source_leaf_ids=jnp.zeros((0, 0), dtype=idx), + target_block_valid_mask=jnp.zeros((0, 0), dtype=bool), + target_block_offsets=jnp.zeros((num_leaves + 1,), dtype=idx), + target_block_size=0, + ) + + if capacity_report is not None: + report = dict( + traced=bool(traced), + flat_walk=True, + max_pair_queue_requested=int(queue), + queue_capacity=int(queue), + compact_far_pair_capacity=int(compact_far_pair_capacity), + near_edge_capacity=int(near_edge_capacity), + far_named=bool(far_named), + near_edge_named=bool(near_edge_named), + # no per-leaf row cap on this lane; None switches the guard's row arm off + max_neighbors_per_leaf_used=None, + grew=list(grew), + ) + if not traced: + report["far_pair_count"] = 2 * int(walk.far_count) + report["total_neighbors"] = 2 * int(walk.near_count) + report["max_neighbors_observed"] = int(jnp.max(counts)) if num_leaves else 0 + report["peak_wavefront"] = int(walk.peak_wavefront) + report["rounds"] = int(walk.rounds) + capacity_report(report) + + return _DualTreeArtifacts( + interactions=None, + neighbor_list=neighbor_list, + traversal_result=None, + compact_far_pairs=compact_far_pairs, + dense_buffers=None, + grouped_buffers=None, + grouped_segment_starts=None, + grouped_segment_lengths=None, + grouped_segment_class_ids=None, + grouped_segment_sort_permutation=None, + grouped_segment_group_ids=None, + grouped_segment_unique_targets=None, + grouped_chunk_size=None, + ) + + def _dual_tree_build_grouped_buffers( *, tree: Tree, @@ -2257,6 +2664,7 @@ def _build_dual_tree_artifacts( planner_hint: Optional[_RefreshDualPlannerHint] = None, strict_capacity_report: Optional[Callable[[dict], None]] = None, strict_max_neighbors_per_leaf_override: Optional[int] = None, + strict_flat_walk_capacity_floor: Optional[dict] = None, ) -> tuple[_DualTreeArtifacts, Optional[_InteractionCacheEntry]]: """Construct or reuse dual-tree traversal products for a tree. @@ -2332,6 +2740,8 @@ def _build_dual_tree_artifacts( Forwarded to the strict streamed builder as ``capacity_report``. strict_max_neighbors_per_leaf_override : Optional[int] Forwarded to the strict streamed builder; only ever raises the cap. + strict_flat_walk_capacity_floor : Optional[dict] + Forwarded to the strict streamed builder as ``flat_walk_capacity_floor``. Returns ------- @@ -2412,6 +2822,7 @@ def _build_dual_tree_artifacts( max_neighbors_per_leaf_override=( strict_max_neighbors_per_leaf_override ), + flat_walk_capacity_floor=strict_flat_walk_capacity_floor, ) if strict_streamed_split else _build_dual_tree_artifacts_split( diff --git a/jaccpot/runtime/fmm_prepare.py b/jaccpot/runtime/fmm_prepare.py index b1d8c7ab..6c98b6a5 100644 --- a/jaccpot/runtime/fmm_prepare.py +++ b/jaccpot/runtime/fmm_prepare.py @@ -1467,6 +1467,7 @@ def _record_dual_artifact_substage(name: str, elapsed: float) -> None: ( runtime_traversal_config, strict_nbr_override, + strict_flat_floor, _strict_capacity_report, ) = self._strict_fused_capacity_handoff( runtime_traversal_config=runtime_traversal_config, @@ -1478,6 +1479,7 @@ def _record_dual_artifact_substage(name: str, elapsed: float) -> None: geometry_factory=geometry_factory, strict_capacity_report=_strict_capacity_report, strict_max_neighbors_per_leaf_override=strict_nbr_override, + strict_flat_walk_capacity_floor=strict_flat_floor, theta=theta_val, mac_type=mac_type_val, dehnen_radius_scale=dehnen_radius_scale, @@ -3341,6 +3343,7 @@ def _strict_fused_capacity_handoff( ) -> tuple[ Optional[DualTreeTraversalConfig], Optional[int], + Optional[dict], Callable[[dict], None], ]: """Carry the eager walk's validated capacities into the traced refresh. @@ -3386,6 +3389,10 @@ def _strict_fused_capacity_handoff( The traversal config to build with. Optional[int] ``max_neighbors_per_leaf`` floor for the strict streamed builder. + Optional[dict] + Floors for the flat-walk lane's directed far / near capacities (the + widths the eager pass settled on, so the traced carry matches; an + eager re-prepare starts from them too, so widths never shrink). Callable[[dict], None] Report callback that stores the builder's capacities on the engine. """ @@ -3395,16 +3402,33 @@ def _pow2_ceil(value: int) -> int: return 1 << (value - 1).bit_length() nbr_override: Optional[int] = None + flat_floor: Optional[dict] = None validated = getattr(self, "_strict_fused_validated_caps", None) + if isinstance(validated, dict) and bool(validated.get("flat_walk")): + flat_floor = { + k: int(validated[k]) + for k in ("compact_far_pair_capacity", "near_edge_capacity") + if validated.get(k) is not None + } if bool(suppress_host_side_effects) and isinstance(validated, dict): + flat_walk = bool(validated.get("flat_walk")) observed_rows = validated.get("max_neighbors_observed") - if observed_rows is not None: + if observed_rows is not None and not flat_walk: + # The flat walk has no per-leaf row buffer, so no row cap to widen. nbr_override = _pow2_ceil(int(1.5 * int(observed_rows)) + 1) validated_queue = validated.get("queue_capacity") + peak_wavefront = validated.get("peak_wavefront") if validated_queue is not None and runtime_traversal_config is not None: + if flat_walk and peak_wavefront is not None: + # Sized from DATA: the eager flat walk reports the largest + # wavefront any round needed, so 1.5x that (pow2) replaces the + # 2x-the-rung rule -- the per-round cost of the traced walk is + # linear in this capacity (3-6x saved at leaf 64, measured). + target_queue = _pow2_ceil(int(1.5 * int(peak_wavefront))) + else: + target_queue = _pow2_ceil(2 * int(validated_queue)) widened_queue = max( - int(runtime_traversal_config.max_pair_queue), - _pow2_ceil(2 * int(validated_queue)), + int(runtime_traversal_config.max_pair_queue), target_queue ) if widened_queue != int(runtime_traversal_config.max_pair_queue): runtime_traversal_config = DualTreeTraversalConfig( @@ -3424,7 +3448,7 @@ def _report(report: dict) -> None: else: self._strict_fused_validated_caps = dict(report) - return runtime_traversal_config, nbr_override, _report + return runtime_traversal_config, nbr_override, flat_floor, _report def _prepare_state_dual_and_downward_strict_streamed_fast( self, @@ -3491,6 +3515,7 @@ def _prepare_state_dual_and_downward_strict_streamed_fast( ( runtime_traversal_config, strict_nbr_override, + strict_flat_floor, _strict_capacity_report, ) = self._strict_fused_capacity_handoff( runtime_traversal_config=runtime_traversal_config, @@ -3502,6 +3527,7 @@ def _prepare_state_dual_and_downward_strict_streamed_fast( geometry_factory=geometry_factory, strict_capacity_report=_strict_capacity_report, strict_max_neighbors_per_leaf_override=strict_nbr_override, + strict_flat_walk_capacity_floor=strict_flat_floor, theta=theta_val, mac_type=mac_type_val, dehnen_radius_scale=dehnen_radius_scale, diff --git a/jaccpot/runtime/fmm_strict_run.py b/jaccpot/runtime/fmm_strict_run.py index 83f24359..56ed04de 100644 --- a/jaccpot/runtime/fmm_strict_run.py +++ b/jaccpot/runtime/fmm_strict_run.py @@ -1053,6 +1053,15 @@ def _skip(_): "Raise JACCPOT_LARGE_N_STATIC_TARGET_BLOCKS_MAX_PER_LEAF, pass " "jaccpot.TraversalOverrides(max_neighbors_per_leaf=...), or raise " "JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP." + + ( + " On the flat-walk lane (the default; " + "JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK=0 for the dual walk) " + "the far-pair count saturates on ANY overflow -- far, near " + "(JACCPOT_LARGE_N_NEIGHBOR_EDGE_PROFILE_FIXED_CAP) or queue " + "(TraversalOverrides(max_pair_queue=...)) -- so check all three." + if traced_caps.get("flat_walk") + else "" + ) ) except Exception as exc: if bool( diff --git a/jaccpot/runtime/kernels/_downward_prep.py b/jaccpot/runtime/kernels/_downward_prep.py index 6aa3cd1a..1b243d23 100644 --- a/jaccpot/runtime/kernels/_downward_prep.py +++ b/jaccpot/runtime/kernels/_downward_prep.py @@ -59,9 +59,10 @@ def _m2l_csr_pallas_active() -> bool: """Whether the flat real-basis M2L runs the target-tiled CSR Pallas kernel. - Opt-in (``JACCPOT_STATIC_STRICT_FUSED_M2L_CSR=1``) and only where it can - lower: an Ampere+ GPU, or ``JACCPOT_M2L_CSR_INTERPRET=1`` for CPU parity - tests. Read at trace time through :mod:`jaccpot._env`, never at import. + Default on (since 2026-09-10; ``JACCPOT_STATIC_STRICT_FUSED_M2L_CSR=0`` + restores the chunked pure-JAX lanes) and only where it can lower: an + Ampere+ GPU, or ``JACCPOT_M2L_CSR_INTERPRET=1`` for CPU parity tests. Read at + trace time through :mod:`jaccpot._env`, never at import. Returns ------- @@ -71,7 +72,7 @@ def _m2l_csr_pallas_active() -> bool: """ from jaccpot._env import env_flag - if not env_flag("JACCPOT_STATIC_STRICT_FUSED_M2L_CSR", False): + if not env_flag("JACCPOT_STATIC_STRICT_FUSED_M2L_CSR", True): return False if env_flag("JACCPOT_M2L_CSR_INTERPRET", False): return True @@ -544,8 +545,9 @@ def _solidfmm_downward_accumulate_from_multipoles( picks grouped over flat, ``farfield_mode`` picks class-major over pair-grouped within grouped, and on the flat path ``pair_count <= chunk_size`` picks full-batch over a chunked scan. All four compute the same operator; they - differ in how the pair list is blocked. A fifth, opt-in flat lane for the - real basis (:func:`_m2l_csr_pallas_active`) hands the whole pair list to the + differ in how the pair list is blocked. A fifth flat lane for the real + basis, the default on Ampere+ GPUs + (:func:`_m2l_csr_pallas_active`), hands the whole pair list to the target-tiled CSR Pallas kernel, which owns one local row per program and so needs neither the chunked scan nor its per-chunk scatter. diff --git a/tests/conftest.py b/tests/conftest.py index 42f4bca5..ec5ee897 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -10,7 +10,12 @@ if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) -YGGDRAX_ROOT = REPO_ROOT.parent / "yggdrax" +# The sibling checkout is the default; ``YGGDRAX_WORKTREE`` names another (a +# git worktree on a branch under test) -- the shared checkout is often on a +# paper branch, and a test run must be able to pin the yggdrax it exercises. +YGGDRAX_ROOT = pathlib.Path( + os.environ.get("YGGDRAX_WORKTREE") or (REPO_ROOT.parent / "yggdrax") +) if YGGDRAX_ROOT.exists() and str(YGGDRAX_ROOT) not in sys.path: sys.path.insert(0, str(YGGDRAX_ROOT)) diff --git a/tests/integration/test_strict_run_v2_refresh_capacity.py b/tests/integration/test_strict_run_v2_refresh_capacity.py index 4ad1d152..78b3aaf0 100644 --- a/tests/integration/test_strict_run_v2_refresh_capacity.py +++ b/tests/integration/test_strict_run_v2_refresh_capacity.py @@ -67,12 +67,18 @@ def _rel_l2(a, b): @pytest.mark.skipif( jax.default_backend() != "gpu", reason="fused strict lane is GPU-only" ) -def test_strict_run_v2_refresh_keeps_full_neighbor_lists(monkeypatch): +@pytest.mark.parametrize("flat_walk", ["0", "1"], ids=["dual_walk", "flat_walk"]) +def test_strict_run_v2_refresh_keeps_full_neighbor_lists(monkeypatch, flat_walk): n = 200_000 leaf, order, theta = 256, 3, 0.6 for key, val in _FUSED_ENV.items(): monkeypatch.setenv(key, val) monkeypatch.setenv("JACCPOT_STATIC_STRICT_FUSED_PROFILE_SET", str(n)) + # The flat-emission walk (plan "tree walk", 2026-09-10) has no per-leaf row + # cap: its neighbour list is a CSR of width JACCPOT_LARGE_N_NEIGHBOR_EDGE_ + # PROFILE_FIXED_CAP, and every overflow saturates the far-pair count into + # the same guard. Same force, same lists as sets, different fp32 order. + monkeypatch.setenv("JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK", flat_walk) from jaccpot import ( FarFieldConfig, @@ -157,4 +163,11 @@ def eager_force_at(x): int(eager_rows.max()), ) caps = solver._impl._strict_fused_traced_caps - assert caps["max_neighbors_per_leaf_used"] > int(refreshed_rows.max()) + if flat_walk == "1": + assert caps["flat_walk"] is True + assert caps["near_edge_capacity"] > int(refreshed_rows.sum()) + validated = solver._impl._strict_fused_validated_caps + assert validated["peak_wavefront"] > 0 + assert caps["queue_capacity"] >= validated["peak_wavefront"] + else: + assert caps["max_neighbors_per_leaf_used"] > int(refreshed_rows.max()) diff --git a/tests/unit/runtime/test_flat_walk_default_dispatch.py b/tests/unit/runtime/test_flat_walk_default_dispatch.py new file mode 100644 index 00000000..0b388571 --- /dev/null +++ b/tests/unit/runtime/test_flat_walk_default_dispatch.py @@ -0,0 +1,195 @@ +"""The strict streamed seam takes the flat walk BY DEFAULT and falls back sanely. + +``_build_dual_tree_artifacts_split_strict_streamed`` (``runtime/_interaction_cache.py``) +routes to the flat-emission walk unless ``JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK=0``. +A configuration the flat walk cannot carry -- the treecode walk requested, a MAC +other than bh/dehnen, a solver-owned pair policy, the non-flat far-pair layout -- +falls back to the dual walk QUIETLY while the flag is merely defaulted, and RAISES +when the flag was set to ``1`` explicitly (the caller asked for a walk it cannot +have). Capacities the caller did not name grow eagerly from their floor and the +floor handed over from an earlier eager pass is honoured; a named cap is exact. +Pinned on CPU on a real tree, every commit. +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +import pytest + +pytest.importorskip("yggdrax") +from yggdrax import DualTreeTraversalConfig +from yggdrax._geometry_impl import compute_tree_geometry +from yggdrax.tree import Tree + +import jaccpot.runtime._interaction_cache as ic + +_LEAF = 8 +_FLAG = "JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK" +_NEAR_CAP = "JACCPOT_LARGE_N_NEIGHBOR_EDGE_PROFILE_FIXED_CAP" +_FAR_CAP = "JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP" + + +@pytest.fixture(scope="module") +def tree_and_geometry(): + points = jax.random.uniform(jax.random.PRNGKey(5), (512, 3), dtype=jnp.float64) + masses = jnp.ones((512,), dtype=jnp.float64) + tree = Tree.from_particles(points, masses, leaf_size=_LEAF, tree_type="radix") + geometry = compute_tree_geometry( + tree.topology, tree.positions_sorted, max_leaf_size=_LEAF + ) + return tree, geometry + + +@pytest.fixture(autouse=True) +def _clean_env(monkeypatch): + for k in ( + _FLAG, + _NEAR_CAP, + _FAR_CAP, + "JACCPOT_STATIC_STRICT_FUSED_TREECODE_WALK", + "JACCPOT_STATIC_STRICT_FUSED_FLAT_COMPACT_FAR_PAIRS", + ): + monkeypatch.delenv(k, raising=False) + + +def _seam(tree, geometry, *, mac_type="dehnen", report=None, floor=None): + return ic._build_dual_tree_artifacts_split_strict_streamed( + tree=tree, + geometry=geometry, + theta=0.5, + mac_type=mac_type, + dehnen_radius_scale=1.0, + max_pair_queue=None, + pair_process_block=None, + traversal_config=DualTreeTraversalConfig( + max_pair_queue=1 << 14, + process_block=256, + max_interactions_per_node=1024, + max_neighbors_per_leaf=512, + ), + pair_policy=None, + policy_state=None, + capacity_report=report, + flat_walk_capacity_floor=floor, + ) + + +def _count(monkeypatch, name): + calls = {"n": 0} + real = getattr(ic, name) + + def counting(*a, **k): + calls["n"] += 1 + return real(*a, **k) + + monkeypatch.setattr(ic, name, counting) + return calls + + +def test_flat_walk_is_the_default(monkeypatch, tree_and_geometry): + tree, geometry = tree_and_geometry + flat = _count(monkeypatch, "_build_flat_walk_artifacts_strict_streamed") + reports = [] + art = _seam(tree, geometry, report=reports.append) + assert flat["n"] == 1 + (report,) = reports + assert report["flat_walk"] is True and report["peak_wavefront"] > 0 + assert report["far_named"] is False and report["near_edge_named"] is False + # unnamed caps sit at their floors (the 512-particle tree fits them) + assert report["near_edge_capacity"] == ic._FLAT_WALK_NEAR_EDGE_FLOOR + assert report["compact_far_pair_capacity"] == ic._STRICT_STREAMED_FAR_PAIR_FLOOR + assert int(art.compact_far_pairs.far_pair_count) > 0 + + +def test_flag_zero_takes_the_dual_walk(monkeypatch, tree_and_geometry): + tree, geometry = tree_and_geometry + flat = _count(monkeypatch, "_build_flat_walk_artifacts_strict_streamed") + monkeypatch.setenv(_FLAG, "0") + reports = [] + art = _seam(tree, geometry, report=reports.append) + assert flat["n"] == 0 + (report,) = reports + assert not report.get("flat_walk") + assert int(art.compact_far_pairs.far_pair_count) > 0 + + +@pytest.mark.parametrize( + "blocker", + ["treecode", "mac", "layout"], +) +def test_defaulted_flag_falls_back_and_explicit_flag_raises( + monkeypatch, tree_and_geometry, blocker +): + tree, geometry = tree_and_geometry + mac_type = "dehnen" + if blocker == "treecode": + # the treecode graft is exercised elsewhere; here only the routing + # matters, so its builder is replaced by one that returns a known + # (correctly typed) artifacts object built by the dual walk. + monkeypatch.setenv(_FLAG, "0") + sentinel = _seam(tree, geometry) + monkeypatch.delenv(_FLAG) + monkeypatch.setenv("JACCPOT_STATIC_STRICT_FUSED_TREECODE_WALK", "1") + monkeypatch.setattr( + ic, "_build_treecode_artifacts_strict_streamed", lambda **k: sentinel + ) + elif blocker == "mac": + mac_type = "engblom" + else: + monkeypatch.setenv("JACCPOT_STATIC_STRICT_FUSED_FLAT_COMPACT_FAR_PAIRS", "0") + flat = _count(monkeypatch, "_build_flat_walk_artifacts_strict_streamed") + + # defaulted: quiet fallback, the flat builder is never entered + out = _seam(tree, geometry, mac_type=mac_type) + assert flat["n"] == 0 + if blocker == "treecode": + assert out is sentinel + else: + assert int(out.compact_far_pairs.far_pair_count) > 0 or blocker == "layout" + + # explicit: the caller asked for a walk it cannot have + monkeypatch.setenv(_FLAG, "1") + with pytest.raises(RuntimeError, match="JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK"): + _seam(tree, geometry, mac_type=mac_type) + assert flat["n"] == 0 + + +def test_named_near_cap_is_exact_and_unnamed_grows(monkeypatch, tree_and_geometry): + tree, geometry = tree_and_geometry + monkeypatch.setenv(_NEAR_CAP, "16") + monkeypatch.setenv(_FAR_CAP, "16") + with pytest.raises(RuntimeError, match="named it"): + _seam(tree, geometry) + monkeypatch.delenv(_NEAR_CAP) + monkeypatch.delenv(_FAR_CAP) + monkeypatch.setattr(ic, "_FLAT_WALK_NEAR_EDGE_FLOOR", 16) + monkeypatch.setattr(ic, "_STRICT_STREAMED_FAR_PAIR_FLOOR", 16) + reports = [] + _seam(tree, geometry, report=reports.append) + (report,) = reports + assert ( + report["near_edge_capacity"] > 16 and report["compact_far_pair_capacity"] > 16 + ) + assert report["grew"], report + + +def test_capacity_floor_from_an_earlier_pass_is_honoured( + monkeypatch, tree_and_geometry +): + tree, geometry = tree_and_geometry + reports = [] + _seam( + tree, + geometry, + report=reports.append, + floor={"near_edge_capacity": 1 << 23, "compact_far_pair_capacity": 1 << 19}, + ) + (report,) = reports + assert report["near_edge_capacity"] == 1 << 23 + assert report["compact_far_pair_capacity"] == 1 << 19 + # a NAMED cap ignores the floor + monkeypatch.setenv(_NEAR_CAP, str(1 << 15)) + reports.clear() + _seam(tree, geometry, report=reports.append, floor={"near_edge_capacity": 1 << 23}) + assert reports[0]["near_edge_capacity"] == 1 << 15 diff --git a/tests/unit/runtime/test_flat_walk_lane_wiring.py b/tests/unit/runtime/test_flat_walk_lane_wiring.py new file mode 100644 index 00000000..43115078 --- /dev/null +++ b/tests/unit/runtime/test_flat_walk_lane_wiring.py @@ -0,0 +1,149 @@ +"""The flat-walk lane is the solver's default walk, ``=0`` restores the dual walk, +and the two are force-neutral against each other. + +Pattern of ``test_m2l_csr_lane_wiring.py``: a lane that is silently not reached +would pass any parity check, so the builder entry is counted. The strict fused +lane that owns this seam exists only above the large-N threshold on a GPU +(``LargeNPreparedState``), so this runs in the ordinary GPU suite at N=70k with +the env of ``tests/integration/test_strict_run_v2_refresh_capacity.py``. +""" + +from __future__ import annotations + +import jax +import numpy as np +import pytest + +import jaccpot.runtime._interaction_cache as ic + +pytestmark = pytest.mark.skipif( + jax.default_backend() != "gpu", reason="the strict fused large-N lane is GPU-only" +) +_N = 70_000 + +_ENV = { + "JACCPOT_STATIC_STRICT_GPU_MODE": "on", + "JACCPOT_STATIC_STRICT_FUSED_MODE": "on", + "JACCPOT_LARGE_N_STATIC_TARGET_BLOCKS": "1", + "JACCPOT_LARGE_N_TARGET_BLOCK_SIZE": "4", + "JACCPOT_LARGE_N_STATIC_TARGET_BLOCKS_MAX_PER_LEAF": "auto", + "JACCPOT_LARGE_N_NEIGHBOR_EDGE_PROFILE_FIXED_CAP": "4194304", + "JACCPOT_STATIC_STRICT_REQUIRE_EXACT_CAP_PROFILE_MATCH": "0", + "JACCPOT_STATIC_STRICT_FUSED_DEVICE_ONLY": "1", + "JACCPOT_STATIC_STRICT_FUSED_DISALLOW_HOST_SEGMENT_FALLBACK": "1", + "JACCPOT_STATIC_STRICT_FUSED_FLAT_COMPACT_FAR_PAIRS": "1", + "JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP": "1048576", + "JACCPOT_LARGE_N_COMPILED_STATE_MODE": "on", + "JACCPOT_LARGE_N_RADIX_FAST_PAYLOAD_IN_FUSED": "1", + "JACCPOT_LARGE_N_RADIX_FAST_PAYLOAD_MAX_MB": "0", +} + + +def _plummer(n, seed=0): + rng = np.random.default_rng(seed) + x = rng.uniform(0.0, 1.0, size=n) + r = 1.0 / np.sqrt(x ** (-2.0 / 3.0) - 1.0) + mu = rng.uniform(-1.0, 1.0, size=n) + phi = rng.uniform(0.0, 2.0 * np.pi, size=n) + st = np.sqrt(1.0 - mu * mu) + pos = np.stack([r * st * np.cos(phi), r * st * np.sin(phi), r * mu], 1) + return pos.astype(np.float32), np.full(n, 1.0 / n, np.float32) + + +def _force(monkeypatch, n, leaf, flat): + import jax + import jax.numpy as jnp + + from jaccpot import ( + FarFieldConfig, + FastMultipoleMethod, + FMMAdvancedConfig, + NearFieldConfig, + TreeConfig, + ) + + for k, v in _ENV.items(): + monkeypatch.setenv(k, v) + monkeypatch.setenv("JACCPOT_STATIC_STRICT_FUSED_PROFILE_SET", str(n)) + if flat is None: + monkeypatch.delenv("JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK", raising=False) + else: + monkeypatch.setenv( + "JACCPOT_STATIC_STRICT_FUSED_FLAT_WALK", "1" if flat else "0" + ) + pos, mass = _plummer(n) + solver = FastMultipoleMethod( + preset="large_n_gpu", + runtime_path="large_n", + basis="real", + theta=0.6, + G=1.0, + softening=1e-3, + working_dtype=jnp.float32, + advanced=FMMAdvancedConfig( + tree=TreeConfig(mode="static_radix", leaf_target=leaf), + farfield=FarFieldConfig(mode="auto"), + nearfield=NearFieldConfig(mode="auto"), + mac_type="dehnen", + ), + fixed_order=3, + ) + prepared, ev = solver.strict_fused_prepared_eval_fn( + positions=jnp.asarray(pos), + masses=jnp.asarray(mass), + leaf_size=leaf, + max_order=3, + theta=0.6, + ) + acc = np.asarray(jax.block_until_ready(ev(prepared)), np.float64) + caps = dict(getattr(solver._impl, "_strict_fused_validated_caps", None) or {}) + return acc, caps + + +def test_default_on_and_flag_zero_takes_the_dual_walk(monkeypatch): + calls = {"n": 0} + real = ic._build_flat_walk_artifacts_strict_streamed + + def counting(*a, **k): + calls["n"] += 1 + return real(*a, **k) + + monkeypatch.setattr(ic, "_build_flat_walk_artifacts_strict_streamed", counting) + a_dual, caps_dual = _force(monkeypatch, _N, 64, flat=False) + assert calls["n"] == 0 + assert not caps_dual.get("flat_walk") + a_flat, caps_flat = _force(monkeypatch, _N, 64, flat=None) + assert calls["n"] >= 1, "the flat-walk lane was never entered by default" + assert caps_flat.get("flat_walk") is True + assert caps_flat.get("peak_wavefront", 0) > 0 + assert np.all(np.isfinite(a_flat)) + rel = np.linalg.norm(a_flat - a_dual) / np.linalg.norm(a_dual) + # same lists as sets; only the fp32 summation order differs + assert rel < 2e-5, rel + + +def test_both_walk_flags_set_is_refused(monkeypatch): + monkeypatch.setenv("JACCPOT_STATIC_STRICT_FUSED_TREECODE_WALK", "1") + with pytest.raises(RuntimeError, match="pick one walk"): + _force(monkeypatch, _N, 64, flat=True) + + +def test_treecode_flag_alone_wins_over_the_defaulted_flat_walk(monkeypatch): + calls = {"flat": 0, "treecode": 0} + real_flat = ic._build_flat_walk_artifacts_strict_streamed + real_tree = ic._build_treecode_artifacts_strict_streamed + + def count_flat(*a, **k): + calls["flat"] += 1 + return real_flat(*a, **k) + + def count_tree(*a, **k): + calls["treecode"] += 1 + return real_tree(*a, **k) + + monkeypatch.setattr(ic, "_build_flat_walk_artifacts_strict_streamed", count_flat) + monkeypatch.setattr(ic, "_build_treecode_artifacts_strict_streamed", count_tree) + monkeypatch.setenv("JACCPOT_STATIC_STRICT_FUSED_TREECODE_WALK", "1") + acc, _ = _force(monkeypatch, _N, 64, flat=None) + assert calls["flat"] == 0 and calls["treecode"] >= 1, calls + assert np.all(np.isfinite(acc)) diff --git a/tests/unit/runtime/test_flat_walk_production_seam.py b/tests/unit/runtime/test_flat_walk_production_seam.py new file mode 100644 index 00000000..1f617089 --- /dev/null +++ b/tests/unit/runtime/test_flat_walk_production_seam.py @@ -0,0 +1,259 @@ +"""The flat-walk seam of the strict fused lane: same lists as the dual walk, as sets. + +``_build_flat_walk_artifacts_strict_streamed`` (plan "tree walk", 2026-09-10) +replaces yggdrax's traced dual-tree walk with its flat-emission +``dual_tree_walk_mutual`` fed the dual walk's own ``mac_extents`` and +``mac_type``. The contract pinned here on a real tree, on CPU, every commit: + +* far pairs are DIRECTED and PREFIX-LIVE (every consumer masks + ``idx < far_pair_count``), each canonical pair present in both directions, the + tail ``-1``; +* the far set equals the dual walk's far set, the per-leaf neighbour sets equal + the dual walk's, no self neighbour, no duplicate ``(leaf, neighbour)``; +* the neighbour CSR is valid (offsets monotone, counts consistent, width equal to + the edge cap); +* the capacity report is emitted with ``peak_wavefront`` and marks the lane; +* a too-small far or near capacity raises eagerly, naming the knob, when the + caller named it; an unnamed one is doubled eagerly and the report says so. +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +pytest.importorskip("yggdrax") +from yggdrax import DualTreeTraversalConfig +from yggdrax._geometry_impl import compute_tree_geometry +from yggdrax._interactions_impl import build_interactions_and_neighbors +from yggdrax.tree import Tree + +from jaccpot.runtime._interaction_cache import ( + _build_flat_walk_artifacts_strict_streamed, +) + +_LEAF = 8 + + +@pytest.fixture(scope="module") +def tree_and_geometry(): + points = jax.random.uniform(jax.random.PRNGKey(3), (512, 3), dtype=jnp.float64) + masses = jnp.ones((512,), dtype=jnp.float64) + tree = Tree.from_particles(points, masses, leaf_size=_LEAF, tree_type="radix") + geometry = compute_tree_geometry( + tree.topology, tree.positions_sorted, max_leaf_size=_LEAF + ) + num_internal = int(tree.topology.left_child.shape[0]) + total_nodes = int(tree.topology.parent.shape[0]) + return tree, geometry, num_internal, total_nodes + + +def _flat( + tree, + geometry, + *, + theta=0.5, + mac_type="dehnen", + scale=1.0, + far_cap=1 << 15, + near_cap=1 << 15, + queue=1 << 14, + report=None, + **kw, +): + return _build_flat_walk_artifacts_strict_streamed( + tree=tree, + geometry=geometry, + theta=theta, + mac_type=mac_type, + dehnen_radius_scale=scale, + compact_far_pair_capacity=far_cap, + near_edge_capacity=near_cap, + max_pair_queue=queue, + capacity_report=report, + **kw, + ) + + +def _dual_sets(tree, geometry, *, theta, mac_type, scale=1.0): + config = DualTreeTraversalConfig( + max_pair_queue=1 << 15, + process_block=64, + max_interactions_per_node=2048, + max_neighbors_per_leaf=2048, + ) + _i, neighbors, result = build_interactions_and_neighbors( + tree.topology, + geometry, + theta=theta, + traversal_config=config, + mac_type=mac_type, + dehnen_radius_scale=scale, + return_result=True, + ) + assert not ( + bool(result.queue_overflow) + or bool(result.far_overflow) + or bool(result.near_overflow) + ) + src = np.asarray(result.interaction_sources) + tgt = np.asarray(result.interaction_targets) + live = (src >= 0) & (tgt >= 0) + far = set(zip(tgt[live].tolist(), src[live].tolist())) + offsets, counts = np.asarray(neighbors.offsets), np.asarray(neighbors.counts) + nbrs, leaves = np.asarray(neighbors.neighbors), np.asarray(neighbors.leaf_indices) + near = {} + for row, leaf in enumerate(leaves.tolist()): + near[leaf] = {int(nbrs[int(offsets[row]) + k]) for k in range(int(counts[row]))} + return far, near + + +def _flat_far_set(cfp): + n = int(cfp.far_pair_count) + src, tgt = np.asarray(cfp.sources), np.asarray(cfp.targets) + assert np.all(src[:n] >= 0) and np.all(tgt[:n] >= 0), "live prefix has no padding" + assert np.all(src[n:] == -1) and np.all(tgt[n:] == -1), "-1 tail after the prefix" + pairs = list(zip(tgt[:n].tolist(), src[:n].tolist())) + assert len(set(pairs)) == len(pairs), "duplicate directed far pair" + return set(pairs) + + +def _flat_near_sets(nl, num_internal, total_nodes): + offsets, counts = np.asarray(nl.offsets), np.asarray(nl.counts) + nbrs, leaves = np.asarray(nl.neighbors), np.asarray(nl.leaf_indices) + assert np.array_equal(leaves, np.arange(num_internal, total_nodes)) + assert np.all(np.diff(offsets) >= 0) and np.array_equal( + offsets[1:] - offsets[:-1], counts + ) + near = {} + for row, leaf in enumerate(leaves.tolist()): + block = nbrs[int(offsets[row]) : int(offsets[row]) + int(counts[row])].tolist() + assert leaf not in block, "self neighbour" + assert len(set(block)) == len(block), "duplicate neighbour" + assert all( + num_internal <= b < total_nodes for b in block + ), "neighbour is not a leaf" + near[leaf] = set(block) + return near + + +@pytest.mark.parametrize("mac_type", ["bh", "dehnen"]) +@pytest.mark.parametrize("theta", [0.3, 0.5, 0.9]) +def test_flat_walk_lists_equal_the_dual_walk_as_sets( + tree_and_geometry, mac_type, theta +): + tree, geometry, num_internal, total_nodes = tree_and_geometry + far_ref, near_ref = _dual_sets(tree, geometry, theta=theta, mac_type=mac_type) + reports = [] + art = _flat(tree, geometry, theta=theta, mac_type=mac_type, report=reports.append) + assert art.interactions is None and art.traversal_result is None + far = _flat_far_set(art.compact_far_pairs) + assert far == far_ref + assert all((b, a) in far for a, b in far), "every far pair in both directions" + near = _flat_near_sets(art.neighbor_list, num_internal, total_nodes) + assert near == near_ref + assert int(art.neighbor_list.neighbors.shape[0]) == 1 << 15 # edge-cap width + (report,) = reports + assert report["flat_walk"] is True and report["traced"] is False + assert report["peak_wavefront"] > 0 and report["rounds"] > 0 + assert report["far_pair_count"] == len(far) + assert report["total_neighbors"] == sum(len(v) for v in near.values()) + assert report["max_neighbors_per_leaf_used"] is None + + +def test_dehnen_radius_scale_is_honoured(tree_and_geometry): + tree, geometry, num_internal, total_nodes = tree_and_geometry + far_ref, near_ref = _dual_sets( + tree, geometry, theta=0.5, mac_type="dehnen", scale=1.4 + ) + art = _flat(tree, geometry, theta=0.5, mac_type="dehnen", scale=1.4) + assert _flat_far_set(art.compact_far_pairs) == far_ref + assert _flat_near_sets(art.neighbor_list, num_internal, total_nodes) == near_ref + far_unscaled, _ = _dual_sets( + tree, geometry, theta=0.5, mac_type="dehnen", scale=1.0 + ) + assert far_unscaled != far_ref + + +def test_traced_call_keeps_capacity_width_and_reports_traced(tree_and_geometry): + tree, geometry, num_internal, total_nodes = tree_and_geometry + reports = [] + eager = _flat(tree, geometry, report=reports.append) + + def run(positions_sorted): + geom = compute_tree_geometry( + tree.topology, positions_sorted, max_leaf_size=_LEAF + ) + art = _flat(tree, geom, report=reports.append) + return ( + art.compact_far_pairs.sources, + art.compact_far_pairs.targets, + art.compact_far_pairs.far_pair_count, + art.neighbor_list.neighbors, + art.neighbor_list.counts, + ) + + src, tgt, n, nbrs, counts = jax.jit(run)(tree.positions_sorted) + assert int(n) == int(eager.compact_far_pairs.far_pair_count) + assert src.shape == eager.compact_far_pairs.sources.shape + assert nbrs.shape == eager.neighbor_list.neighbors.shape + assert np.array_equal(np.asarray(counts), np.asarray(eager.neighbor_list.counts)) + assert set( + zip(np.asarray(tgt)[: int(n)].tolist(), np.asarray(src)[: int(n)].tolist()) + ) == _flat_far_set(eager.compact_far_pairs) + assert [r["traced"] for r in reports] == [False, True] + assert "peak_wavefront" not in reports[1] + + +def test_far_and_near_overflow_raise_eagerly_naming_the_knob(tree_and_geometry): + tree, geometry, _, _ = tree_and_geometry + with pytest.raises(RuntimeError, match="COMPACT_FAR_PAIR_CAP"): + _flat(tree, geometry, far_cap=8) + with pytest.raises(RuntimeError, match="NEIGHBOR_EDGE_PROFILE_FIXED_CAP"): + _flat(tree, geometry, near_cap=8) + + +def test_unnamed_far_and_near_caps_grow_eagerly(tree_and_geometry): + tree, geometry, num_internal, total_nodes = tree_and_geometry + reports = [] + ref = _flat(tree, geometry) + art = _flat( + tree, + geometry, + far_cap=8, + near_cap=8, + far_named=False, + near_edge_named=False, + report=reports.append, + ) + (report,) = reports + assert report["compact_far_pair_capacity"] > 8 and report["near_edge_capacity"] > 8 + assert report["far_named"] is False and report["near_edge_named"] is False + assert any(g.startswith("compact_far_pair_capacity") for g in report["grew"]) + assert any(g.startswith("near_edge_capacity") for g in report["grew"]) + # widths follow the grown caps, and the lists are the same sets as before + assert art.compact_far_pairs.sources.shape[0] == report["compact_far_pair_capacity"] + assert art.neighbor_list.neighbors.shape[0] == report["near_edge_capacity"] + assert _flat_far_set(art.compact_far_pairs) == _flat_far_set(ref.compact_far_pairs) + assert _flat_near_sets(art.neighbor_list, num_internal, total_nodes) == ( + _flat_near_sets(ref.neighbor_list, num_internal, total_nodes) + ) + + +def test_odd_capacities_are_rejected(tree_and_geometry): + tree, geometry, _, _ = tree_and_geometry + with pytest.raises(ValueError, match="even"): + _flat(tree, geometry, far_cap=(1 << 15) + 1) + with pytest.raises(ValueError, match="even"): + _flat(tree, geometry, near_cap=(1 << 15) + 1) + + +def test_queue_ladder_grows_from_a_small_queue(tree_and_geometry): + tree, geometry, _, _ = tree_and_geometry + reports = [] + art = _flat(tree, geometry, queue=8, report=reports.append) + (report,) = reports + assert report["queue_capacity"] > 8 and report["grew"], report + assert int(art.compact_far_pairs.far_pair_count) > 0 diff --git a/tests/unit/runtime/test_m2l_csr_lane_wiring.py b/tests/unit/runtime/test_m2l_csr_lane_wiring.py index 0e9fdf28..676453b0 100644 --- a/tests/unit/runtime/test_m2l_csr_lane_wiring.py +++ b/tests/unit/runtime/test_m2l_csr_lane_wiring.py @@ -1,6 +1,7 @@ -"""The opt-in CSR M2L lane is (a) actually taken and (b) force-neutral. +"""The CSR M2L lane (default on where it lowers) is (a) taken and (b) force-neutral. -``JACCPOT_STATIC_STRICT_FUSED_M2L_CSR=1`` routes the flat real-basis M2L of +``JACCPOT_STATIC_STRICT_FUSED_M2L_CSR`` (default ``1`` since 2026-09-10; ``0`` +restores the chunked pure-JAX lanes) routes the flat real-basis M2L of ``_solidfmm_downward_accumulate_from_multipoles`` to :func:`jaccpot.pallas.m2l_real_csr.m2l_real_csr_pallas`. On CPU the kernel runs in interpret mode (``JACCPOT_M2L_CSR_INTERPRET=1``). A lane that is silently not @@ -16,8 +17,17 @@ from jaccpot.runtime.kernels._downward_prep import _m2l_csr_pallas_active -def test_flag_off_by_default(monkeypatch): +def test_default_on_exactly_where_the_kernel_lowers(monkeypatch): monkeypatch.delenv("JACCPOT_STATIC_STRICT_FUSED_M2L_CSR", raising=False) + monkeypatch.delenv("JACCPOT_M2L_CSR_INTERPRET", raising=False) + assert _m2l_csr_pallas_active() is csr_mod.pallas_m2l_real_csr_supported() + monkeypatch.setenv("JACCPOT_M2L_CSR_INTERPRET", "1") + assert _m2l_csr_pallas_active() is True + + +def test_flag_zero_switches_the_lane_off(monkeypatch): + monkeypatch.setenv("JACCPOT_STATIC_STRICT_FUSED_M2L_CSR", "0") + monkeypatch.setenv("JACCPOT_M2L_CSR_INTERPRET", "1") assert _m2l_csr_pallas_active() is False @@ -77,7 +87,7 @@ def solve(): ) return np.asarray(acc, np.float64) - monkeypatch.delenv("JACCPOT_STATIC_STRICT_FUSED_M2L_CSR", raising=False) + monkeypatch.setenv("JACCPOT_STATIC_STRICT_FUSED_M2L_CSR", "0") a_ref = solve() calls = {"n": 0}