From 0e6d1c119b720be325f6e58deee7ffc0ea700d69 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Sun, 6 Sep 2026 09:58:26 +0200 Subject: [PATCH 1/2] fix(strict): carry the eager walk's validated capacities into the traced refresh The fused strict_run_v2 lane rebuilt its interaction lists inside the compiled velocity-Verlet scan with the preset traversal capacities (at N=200k: pair queue 65536, 256 neighbours per leaf). Eagerly those caps are irrelevant -- yggdrax's retry ladder grows them until the walk fits -- but under jit the overflow flags are tracers, the ladder has one attempt, and yggdrax returns the truncated result. Neighbour rows of up to 781 were cut to 256 (theta 1.0) or 128 (theta 0.6, the queue overflowed too), 85 % of the near-field entries vanished, and every step after the first carried a wrong force: 5.8 % relative L2 at theta 1.0, ~60 % at theta 0.6. fallback_count stayed 0 and no diagnostic fired. Fix: the strict streamed builder now passes a retry_logger to catch the ladder's success queue and reports the capacities it ran with plus what it observed (longest row, far-pair count); the eager prepare stores that on the engine, and the traced refresh raises its caps to cover it with headroom (neighbours and far pairs to pow2(1.5x observed), queue to pow2(2x ladder)). The scan's capacity_ok carry additionally fails when a refreshed row reaches the traced neighbour cap or the far list its buffer, so a future overflow is an error rather than a wrong force. The queue itself stays unverifiable under tracing -- hence the 2x margin. Regression test recovers the applied force from the trajectory and compares it with an eager prepare+evaluate at the same positions (comparing against x0 is wrong here: close pairs move O(0.1) per step and dominate the L2 norm). Co-Authored-By: Claude Fable 5.1 --- jaccpot/runtime/_interaction_cache.py | 105 +++++++++++- jaccpot/runtime/fmm_prepare.py | 117 +++++++++++++ jaccpot/runtime/fmm_strict_run.py | 52 ++++-- .../test_strict_run_v2_refresh_capacity.py | 160 ++++++++++++++++++ 4 files changed, 420 insertions(+), 14 deletions(-) create mode 100644 tests/integration/test_strict_run_v2_refresh_capacity.py diff --git a/jaccpot/runtime/_interaction_cache.py b/jaccpot/runtime/_interaction_cache.py index bcff1dab..a2780157 100644 --- a/jaccpot/runtime/_interaction_cache.py +++ b/jaccpot/runtime/_interaction_cache.py @@ -952,9 +952,24 @@ def _build_dual_tree_artifacts_split_strict_streamed( traversal_config: Optional[DualTreeTraversalConfig], pair_policy: Optional[PairPolicy], policy_state: Optional[AdaptivePolicyState], + capacity_report: Optional[Callable[[dict], None]] = None, + max_neighbors_per_leaf_override: Optional[int] = None, + compact_far_pair_capacity_override: Optional[int] = None, ) -> _DualTreeArtifacts: """Strict static fast-lane: single compact shared far+near build call. + ``capacity_report`` receives, after a successful build, the capacities the + walk actually ran with plus (eager only) what it observed: the queue the + yggdrax ladder settled on, the compact far-pair cap, the far-pair count, the + longest neighbour row and the total edge count. The fused traced refresh + cannot grow capacities -- under ``jit`` the overflow flags are tracers and + yggdrax returns the truncated result -- so it must be told what the eager + prepare needed. ``max_neighbors_per_leaf_override`` and + ``compact_far_pair_capacity_override`` are how it is told: each only ever + RAISES the corresponding capacity. (Found 2026-09-06: at N=200k / leaf 256 + the preset cap of 256 neighbours per leaf against rows of 781 cut 85 % of the + near field out of every step after the first, with no diagnostic firing.) + This path intentionally avoids generic split-builder host branching and callback plumbing. It is valid only for streamed compact far-pairs with no dense/grouped/interactions payload requests. @@ -987,6 +1002,14 @@ def _build_dual_tree_artifacts_split_strict_streamed( MAC alone. policy_state : Optional[AdaptivePolicyState] State the pair policy reads. Meaningless without ``pair_policy``. + capacity_report : Optional[Callable[[dict], None]] + Called once after a successful build with the capacities the walk ran + with (queue, far-pair cap, neighbour cap) and, eagerly, what it observed + (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. + compact_far_pair_capacity_override : Optional[int] + Floor for the compact far-pair cap; only ever raises it. Returns ------- @@ -1021,6 +1044,10 @@ def _build_dual_tree_artifacts_split_strict_streamed( process_block_resolved = ( None if pair_process_block is None else int(pair_process_block) ) + if max_neighbors_per_leaf_override is not None: + max_neighbors_per_leaf = max( + int(max_neighbors_per_leaf), int(max_neighbors_per_leaf_override) + ) flat_compact_enabled = os.environ.get( "JACCPOT_STATIC_STRICT_FUSED_FLAT_COMPACT_FAR_PAIRS", "1" @@ -1034,6 +1061,10 @@ def _build_dual_tree_artifacts_split_strict_streamed( raise ValueError( "JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP must be positive" ) + if compact_far_pair_capacity_override is not None: + compact_far_pair_capacity = max( + int(compact_far_pair_capacity), int(compact_far_pair_capacity_override) + ) # Opt-in: build far/near from the device-resident per-leaf treecode walk # instead of the host-iterated yggdrax dual-tree walk (kills the walk launch @@ -1081,6 +1112,20 @@ def _build_dual_tree_artifacts_split_strict_streamed( attempt_queue = max_pair_queue_resolved attempt_far_cap = compact_far_pair_capacity grew: list[str] = [] + # The eager yggdrax ladder may settle on a LARGER queue than requested; the + # "success" event is the only place that number is reported, and the traced + # refresh needs it (see ``capacity_report``). + ladder_success: dict = {} + + def _capture_ladder(event: DualTreeRetryEvent) -> None: + if str(event.status) == "success": + ladder_success.update( + queue_capacity=int(event.queue_capacity), + interaction_capacity=int(event.interaction_capacity), + far_pair_count=int(event.far_pair_count), + near_pair_count=int(event.near_pair_count), + ) + for attempt in range(_STRICT_STREAMED_RETRY_ATTEMPTS): try: ( @@ -1097,7 +1142,7 @@ def _build_dual_tree_artifacts_split_strict_streamed( max_pair_queue=attempt_queue, process_block=process_block_resolved, traversal_config=None, - retry_logger=None, + retry_logger=_capture_ladder, timing_callback=None, compact_far_pair_capacity=attempt_far_cap, pair_policy=pair_policy, @@ -1145,6 +1190,48 @@ def _build_dual_tree_artifacts_split_strict_streamed( attempt_far_cap = grown else: # pragma: no cover - the loop always breaks or raises raise RuntimeError("strict streamed dual-tree walk did not run") + if capacity_report is not None: + counts_arr = getattr(neighbor_list, "counts", None) + offsets_arr = getattr(neighbor_list, "offsets", None) + fp_count = getattr(compact_far_pairs, "far_pair_count", None) + traced = isinstance(counts_arr, Tracer) or isinstance(fp_count, Tracer) + report = dict( + traced=bool(traced), + max_pair_queue_requested=( + None if attempt_queue is None else int(attempt_queue) + ), + queue_capacity=int( + ladder_success.get( + "queue_capacity", + ( + _STRICT_STREAMED_QUEUE_FLOOR + if attempt_queue is None + else int(attempt_queue) + ), + ) + ), + compact_far_pair_capacity=( + None if attempt_far_cap is None else int(attempt_far_cap) + ), + max_neighbors_per_leaf_used=int(max_neighbors_per_leaf), + grew=list(grew), + ) + if not traced: + try: + report["far_pair_count"] = ( + int(fp_count) + if fp_count is not None + else int(compact_far_pairs.sources.shape[0]) + ) + report["max_neighbors_observed"] = ( + int(jnp.max(counts_arr)) if int(counts_arr.shape[0]) else 0 + ) + report["total_neighbors"] = ( + int(offsets_arr[-1]) if int(offsets_arr.shape[0]) else 0 + ) + except Exception: # pragma: no cover - diagnostics only + pass + capacity_report(report) return _DualTreeArtifacts( interactions=None, neighbor_list=neighbor_list, @@ -2169,6 +2256,9 @@ def _build_dual_tree_artifacts( jit_traversal: bool = True, timing_callback: Optional[Callable[[str, float], None]] = None, planner_hint: Optional[_RefreshDualPlannerHint] = None, + strict_capacity_report: Optional[Callable[[dict], None]] = None, + strict_max_neighbors_per_leaf_override: Optional[int] = None, + strict_compact_far_pair_capacity_override: Optional[int] = None, ) -> tuple[_DualTreeArtifacts, Optional[_InteractionCacheEntry]]: """Construct or reuse dual-tree traversal products for a tree. @@ -2240,6 +2330,12 @@ def _build_dual_tree_artifacts( planner_hint : Optional[_RefreshDualPlannerHint] Hint from a previous refresh, letting the planner skip work whose answer is already known. + strict_capacity_report : Optional[Callable[[dict], None]] + 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_compact_far_pair_capacity_override : Optional[int] + Forwarded to the strict streamed builder; only ever raises the cap. Returns ------- @@ -2316,6 +2412,13 @@ def _build_dual_tree_artifacts( traversal_config=traversal_config, pair_policy=pair_policy, policy_state=policy_state, + capacity_report=strict_capacity_report, + max_neighbors_per_leaf_override=( + strict_max_neighbors_per_leaf_override + ), + compact_far_pair_capacity_override=( + strict_compact_far_pair_capacity_override + ), ) 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 46a9e8ac..d9a81429 100644 --- a/jaccpot/runtime/fmm_prepare.py +++ b/jaccpot/runtime/fmm_prepare.py @@ -1409,10 +1409,22 @@ def _record_dual_artifact_substage(name: str, elapsed: float) -> None: max_leaf_size=int(tree_artifacts.leaf_cap), ) ) + ( + runtime_traversal_config, + strict_nbr_override, + strict_far_override, + _strict_capacity_report, + ) = self._strict_fused_capacity_handoff( + runtime_traversal_config=runtime_traversal_config, + suppress_host_side_effects=suppress_host_side_effects, + ) dual_artifacts, cache_entry = _build_dual_tree_artifacts( tree_artifacts.tree, tree_artifacts.upward.geometry, geometry_factory=geometry_factory, + strict_capacity_report=_strict_capacity_report, + strict_max_neighbors_per_leaf_override=strict_nbr_override, + strict_compact_far_pair_capacity_override=strict_far_override, theta=theta_val, mac_type=mac_type_val, dehnen_radius_scale=dehnen_radius_scale, @@ -3262,6 +3274,99 @@ def _prepare_downward_with_artifacts( p_gears=p_gears, ) + def _strict_fused_capacity_handoff( + self, + *, + runtime_traversal_config: Optional[DualTreeTraversalConfig], + suppress_host_side_effects: bool, + ) -> tuple[ + Optional[DualTreeTraversalConfig], + Optional[int], + Optional[int], + Callable[[dict], None], + ]: + """Carry the eager walk's validated capacities into the traced refresh. + + Eagerly, yggdrax's retry ladder grows the pair queue and the per-leaf + neighbour cap until the walk fits, so the eager state is exact whatever + the preset caps say. Under the fused velocity-Verlet scan the SAME walk + runs traced: the overflow flags are tracers, the ladder has one attempt, + and yggdrax returns the truncated result. Measured 2026-09-06 at + N=200k, leaf 256, theta 0.6: the preset caps (queue 65536, 256 + neighbours per leaf) against rows of up to 781 cut the near field to + 15 % from step 2 on -- relative force error 3 -- with no diagnostic. + + So the eager pass records what it needed (``_strict_fused_validated_caps``, + via the returned report callback) and the traced pass raises its caps to + cover that with headroom: neighbours and far pairs to 1.5x the observed + maximum (next power of two), the queue to 2x the ladder's answer. The + queue cannot be verified after the fact -- its overflow flag never leaves + the trace -- hence the larger margin; the neighbour and far-pair caps ARE + re-checked inside the scan by ``strict_run_v2`` from + ``_strict_fused_traced_caps``, which the same callback records on the + traced build. + + Parameters + ---------- + runtime_traversal_config : Optional[DualTreeTraversalConfig] + Capacities resolved for this build; widened (never shrunk) on the + traced path when the eager pass recorded a larger validated queue. + suppress_host_side_effects : bool + ``True`` on the traced hot path (the refresh inside the compiled + scan); selects which of the two attributes the report lands in. + + Returns + ------- + Optional[DualTreeTraversalConfig] + The traversal config to build with. + Optional[int] + ``max_neighbors_per_leaf`` floor for the strict streamed builder. + Optional[int] + Compact far-pair cap floor for the strict streamed builder. + Callable[[dict], None] + Report callback that stores the builder's capacities on the engine. + """ + + def _pow2_ceil(value: int) -> int: + value = max(1, int(value)) + return 1 << (value - 1).bit_length() + + nbr_override: Optional[int] = None + far_override: Optional[int] = None + validated = getattr(self, "_strict_fused_validated_caps", None) + if bool(suppress_host_side_effects) and isinstance(validated, dict): + observed_rows = validated.get("max_neighbors_observed") + if observed_rows is not None: + nbr_override = _pow2_ceil(int(1.5 * int(observed_rows)) + 1) + observed_far = validated.get("far_pair_count") + if observed_far is not None: + far_override = _pow2_ceil(int(1.5 * int(observed_far)) + 1) + validated_queue = validated.get("queue_capacity") + if validated_queue is not None and runtime_traversal_config is not None: + widened_queue = max( + int(runtime_traversal_config.max_pair_queue), + _pow2_ceil(2 * int(validated_queue)), + ) + if widened_queue != int(runtime_traversal_config.max_pair_queue): + runtime_traversal_config = DualTreeTraversalConfig( + max_pair_queue=int(widened_queue), + process_block=int(runtime_traversal_config.process_block), + max_interactions_per_node=int( + runtime_traversal_config.max_interactions_per_node + ), + max_neighbors_per_leaf=int( + runtime_traversal_config.max_neighbors_per_leaf + ), + ) + + def _report(report: dict) -> None: + if bool(report.get("traced")): + self._strict_fused_traced_caps = dict(report) + else: + self._strict_fused_validated_caps = dict(report) + + return runtime_traversal_config, nbr_override, far_override, _report + def _prepare_state_dual_and_downward_strict_streamed_fast( self, *, @@ -3324,10 +3429,22 @@ def _prepare_state_dual_and_downward_strict_streamed_fast( max_leaf_size=int(tree_artifacts.leaf_cap), ) ) + ( + runtime_traversal_config, + strict_nbr_override, + strict_far_override, + _strict_capacity_report, + ) = self._strict_fused_capacity_handoff( + runtime_traversal_config=runtime_traversal_config, + suppress_host_side_effects=suppress_host_side_effects, + ) dual_artifacts, cache_entry = _build_dual_tree_artifacts( tree_artifacts.tree, tree_artifacts.upward.geometry, geometry_factory=geometry_factory, + strict_capacity_report=_strict_capacity_report, + strict_max_neighbors_per_leaf_override=strict_nbr_override, + strict_compact_far_pair_capacity_override=strict_far_override, 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 7582b970..16bd0a57 100644 --- a/jaccpot/runtime/fmm_strict_run.py +++ b/jaccpot/runtime/fmm_strict_run.py @@ -819,20 +819,40 @@ def _evaluate_self(prepared_in: PreparedStateLike, state_in: Array) -> Array: def _static_target_block_capacity_ok( prepared_in: PreparedStateLike, ) -> Array: + offsets = jnp.asarray(prepared_in.neighbor_list.offsets) + counts = offsets[1:] - offsets[:-1] + ok = jnp.asarray(True) padded = getattr( prepared_in, "nearfield_target_block_source_leaf_ids_padded", None, ) - if padded is None: - return jnp.asarray(True) - padded_arr = jnp.asarray(padded) - if padded_arr.ndim != 3 or int(padded_arr.shape[1]) == 0: - return jnp.asarray(True) - offsets = jnp.asarray(prepared_in.neighbor_list.offsets) - counts = offsets[1:] - offsets[:-1] - capacity = int(padded_arr.shape[1]) * int(padded_arr.shape[2]) - return jnp.all(counts <= jnp.asarray(capacity, dtype=counts.dtype)) + if padded is not None: + padded_arr = jnp.asarray(padded) + if padded_arr.ndim == 3 and int(padded_arr.shape[1]) > 0: + capacity = int(padded_arr.shape[1]) * int(padded_arr.shape[2]) + ok = ok & jnp.all( + counts <= jnp.asarray(capacity, dtype=counts.dtype) + ) + # Traversal-capacity saturation guard. The traced refresh walk runs + # with fixed caps and yggdrax cannot raise on overflow under jit; a + # neighbour row that fills its cap, or a far-pair list that fills + # its buffer, means entries were dropped and the force is wrong. + # The caps are host constants recorded while the refresh traced + # (``_strict_fused_traced_caps``), so this is a static comparison. + traced_caps = getattr(self, "_strict_fused_traced_caps", None) + if isinstance(traced_caps, dict): + nbr_cap = traced_caps.get("max_neighbors_per_leaf_used") + if nbr_cap is not None and int(counts.shape[0]) > 0: + ok = ok & ( + jnp.max(counts) < jnp.asarray(int(nbr_cap), counts.dtype) + ) + far_cap = traced_caps.get("compact_far_pair_capacity") + far_pairs = getattr(prepared_in, "compact_far_pairs", None) + far_count = getattr(far_pairs, "far_pair_count", None) + if far_cap is not None and far_count is not None: + ok = ok & (jnp.asarray(far_count) < jnp.asarray(int(far_cap))) + return ok def _refresh_and_evaluate_endpoint( prepared_in: PreparedStateLike, @@ -1017,11 +1037,17 @@ def _skip(_): "JACCPOT_LARGE_N_STATIC_TARGET_BLOCKS_MAX_PER_LEAF", "32", ) + traced_caps = getattr(self, "_strict_fused_traced_caps", None) or {} raise RuntimeError( - "fused payload static target-block cap exceeded during " - "compiled velocity-Verlet scan: max_blocks_per_leaf=" - f"{max_blocks}. Increase " - "JACCPOT_LARGE_N_STATIC_TARGET_BLOCKS_MAX_PER_LEAF." + "a fixed capacity saturated inside the compiled velocity-Verlet " + "scan, so the refreshed interaction lists are truncated and the " + "forces from that step on are wrong. Checked: static target-block " + f"cap (max_blocks_per_leaf={max_blocks}), traced neighbour cap " + f"({traced_caps.get('max_neighbors_per_leaf_used')} per leaf) and " + f"compact far-pair cap ({traced_caps.get('compact_far_pair_capacity')}). " + "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." ) except Exception as exc: if bool( diff --git a/tests/integration/test_strict_run_v2_refresh_capacity.py b/tests/integration/test_strict_run_v2_refresh_capacity.py new file mode 100644 index 00000000..4ad1d152 --- /dev/null +++ b/tests/integration/test_strict_run_v2_refresh_capacity.py @@ -0,0 +1,160 @@ +"""The fused ``strict_run_v2`` refresh must not truncate its interaction lists. + +Regression test for the 2026-09-06 finding: on ``large_n_gpu``/static_radix at +N=200k, leaf 256, theta 0.6 the eager prepare built exact lists (rows up to +781, via yggdrax's retry ladder) while the traced refresh inside the compiled +velocity-Verlet scan ran the same walk with the preset capacities (256 +neighbours per leaf, queue 65536), could not read the overflow flags under +``jit``, and silently kept 15 % of the near field. Every step after the first +then carried a wrong force -- ~60 % at theta 0.6, 6 % at theta 1.0 -- while +``fallback_count`` stayed 0 and no overflow diagnostic fired. An energy check +between two lanes with the same bug, or any single-step probe, cannot see it. + +The check recovers the force the scan actually applied from the trajectory -- +from rest ``a0 = 2 (x1 - x0) / dt^2`` and ``a_k = (x_{k+1} - 2 x_k + x_{k-1}) +/ dt^2`` -- and compares it with an eager prepare+evaluate AT THE SAME +POSITIONS ``x_k``. Comparing against the force at ``x0`` instead is wrong: +this sample has unresolved close pairs whose accelerations dominate the L2 norm +and move far in one step, so the true field changes by order unity between +steps (that false alarm cost half a day). +""" + +from __future__ import annotations + +import os + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +pytestmark = [pytest.mark.slow] + +_FUSED_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": "64", + "JACCPOT_LARGE_N_NEIGHBOR_EDGE_PROFILE_FIXED_CAP": "2097152", + "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": "131072", + "JACCPOT_LARGE_N_COMPILED_STATE_MODE": "on", + "JACCPOT_LARGE_N_RADIX_FAST_PAYLOAD_IN_FUSED": "1", +} + + +def _plummer(n: int, seed: int = 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 _rel_l2(a, b): + a = np.asarray(a, np.float64) + b = np.asarray(b, np.float64) + return float(np.linalg.norm(a - b) / np.linalg.norm(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): + 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)) + + from jaccpot import ( + FarFieldConfig, + FastMultipoleMethod, + FMMAdvancedConfig, + NearFieldConfig, + TreeConfig, + ) + + pos, mass = _plummer(n) + solver = FastMultipoleMethod( + preset="large_n_gpu", + runtime_path="large_n", + basis="real", + theta=theta, + G=1.0, + softening=1e-7, + 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=order, + ) + P = jnp.asarray(pos) + M = jnp.asarray(mass) + + # eager reference: the eval-only seam builds the exact lists + prepared_eager, eval_fn = solver.strict_fused_prepared_eval_fn( + positions=P, masses=M, leaf_size=leaf, max_order=order, theta=theta + ) + a_eager = np.asarray(jax.block_until_ready(eval_fn(prepared_eager))) + eager_rows = np.asarray(prepared_eager.neighbor_list.counts) + assert ( + int(eager_rows.max()) > 256 + ), "this test needs rows longer than the preset cap to be meaningful" + + def eager_force_at(x): + p, ev = solver.strict_fused_prepared_eval_fn( + positions=jnp.asarray(x, jnp.float32), + masses=M, + leaf_size=leaf, + max_order=order, + theta=theta, + ) + return np.asarray(jax.block_until_ready(ev(p)), np.float64) + + dt = 1e-2 + state0 = jnp.stack([P, jnp.zeros((n, 3), jnp.float32)], axis=1) + final, prepared_out, history = solver.strict_run_v2( + state=state0, + masses=M, + dt=dt, + num_steps=2, + refresh_every=1, + leaf_size=leaf, + max_order=order, + theta=theta, + prepared_state=None, + return_prepared_state=True, + return_history=True, + ) + hist = np.asarray(jax.block_until_ready(history), np.float64) + xs = [pos.astype(np.float64)] + [hist[k, :, 0, :] for k in range(hist.shape[0])] + a0_scan = 2.0 * (xs[1] - xs[0]) / dt**2 # force at x0 (fresh prepare) + a1_scan = (xs[2] - 2.0 * xs[1] + xs[0]) / dt**2 # force at x1 (traced refresh) + + err0 = _rel_l2(a0_scan, a_eager) + err1 = _rel_l2(a1_scan, eager_force_at(xs[1])) + # float32 positions limit the recovery to a few 1e-3; the truncated refresh + # gave ~0.6 here (rows capped at 128 against 781). + assert err0 < 1e-2, err0 + assert err1 < 1e-2, (err0, err1) + + # and the refreshed state's lists are full-length, not a capped subset + refreshed_rows = np.asarray(prepared_out.neighbor_list.counts) + assert int(refreshed_rows.max()) > 256, int(refreshed_rows.max()) + assert abs(int(refreshed_rows.max()) - int(eager_rows.max())) <= 8, ( + int(refreshed_rows.max()), + int(eager_rows.max()), + ) + caps = solver._impl._strict_fused_traced_caps + assert caps["max_neighbors_per_leaf_used"] > int(refreshed_rows.max()) From ca57749eeb964f74c5c1adceddb0e530740e7865 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Mon, 7 Sep 2026 09:43:37 +0200 Subject: [PATCH 2/2] fix(strict): do not widen the compact far-pair cap the caller named The capacity hand-off carried three caps into the traced refresh. Two of them belong there; the compact far-pair cap does not, and widening it broke tests/integration/test_fmm.py::test_strict_fused_compact_far_pair_cap_fails, which pins that a too-small cap fails loudly instead of truncating. The far-pair cap is already checked under tracing -- _raw_to_compact_far_pairs raises through a debug callback -- so it never truncated silently, and it is set explicitly by JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP: raising a cap the caller named would turn a deliberate memory bound into a no-op, the same reasoning _resolve_tracing_traversal_config already applies to explicit capacities. The neighbour cap, the one whose overflow flag is a tracer nobody reads and the one that produced the wrong forces, is still widened. The scan's saturation guard keeps both arms and now says which is which: the neighbour arm is the only defence for a silent failure, the far-pair arm a second, cheap line behind yggdrax's own raise. Verified: the pinned test raises again (CPU), and test_strict_run_v2_refresh_keeps_full_neighbor_lists still passes on an A100 with the override gone -- rows 781, force error 6.5e-4 at step 2. Co-Authored-By: Claude Fable 5.1 --- jaccpot/runtime/_interaction_cache.py | 29 ++++++++----------- jaccpot/runtime/fmm_prepare.py | 40 ++++++++++++--------------- jaccpot/runtime/fmm_strict_run.py | 13 ++++++--- 3 files changed, 38 insertions(+), 44 deletions(-) diff --git a/jaccpot/runtime/_interaction_cache.py b/jaccpot/runtime/_interaction_cache.py index a2780157..6fdd721d 100644 --- a/jaccpot/runtime/_interaction_cache.py +++ b/jaccpot/runtime/_interaction_cache.py @@ -954,7 +954,6 @@ 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, - compact_far_pair_capacity_override: Optional[int] = None, ) -> _DualTreeArtifacts: """Strict static fast-lane: single compact shared far+near build call. @@ -964,11 +963,17 @@ def _build_dual_tree_artifacts_split_strict_streamed( longest neighbour row and the total edge count. The fused traced refresh cannot grow capacities -- under ``jit`` the overflow flags are tracers and yggdrax returns the truncated result -- so it must be told what the eager - prepare needed. ``max_neighbors_per_leaf_override`` and - ``compact_far_pair_capacity_override`` are how it is told: each only ever - RAISES the corresponding capacity. (Found 2026-09-06: at N=200k / leaf 256 - the preset cap of 256 neighbours per leaf against rows of 781 cut 85 % of the - near field out of every step after the first, with no diagnostic firing.) + prepare needed. ``max_neighbors_per_leaf_override`` is how it is told, and it + only ever RAISES the cap. (Found 2026-09-06: at N=200k / leaf 256 the preset + cap of 256 neighbours per leaf against rows of 781 cut 85 % of the near field + out of every step after the first, with no diagnostic firing.) + + Only the neighbour cap is carried over, because it is the only one that fails + SILENTLY: the compact far-pair cap is checked under tracing too + (``_raw_to_compact_far_pairs`` raises through a debug callback), so a + too-small one is already loud, and raising a cap the caller named in + ``JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP`` would turn a deliberate + memory bound into a no-op. This path intentionally avoids generic split-builder host branching and callback plumbing. It is valid only for streamed compact far-pairs with no @@ -1008,8 +1013,6 @@ 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. - compact_far_pair_capacity_override : Optional[int] - Floor for the compact far-pair cap; only ever raises it. Returns ------- @@ -1061,10 +1064,6 @@ def _build_dual_tree_artifacts_split_strict_streamed( raise ValueError( "JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP must be positive" ) - if compact_far_pair_capacity_override is not None: - compact_far_pair_capacity = max( - int(compact_far_pair_capacity), int(compact_far_pair_capacity_override) - ) # Opt-in: build far/near from the device-resident per-leaf treecode walk # instead of the host-iterated yggdrax dual-tree walk (kills the walk launch @@ -2258,7 +2257,6 @@ 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_compact_far_pair_capacity_override: Optional[int] = None, ) -> tuple[_DualTreeArtifacts, Optional[_InteractionCacheEntry]]: """Construct or reuse dual-tree traversal products for a tree. @@ -2334,8 +2332,6 @@ 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_compact_far_pair_capacity_override : Optional[int] - Forwarded to the strict streamed builder; only ever raises the cap. Returns ------- @@ -2416,9 +2412,6 @@ def _build_dual_tree_artifacts( max_neighbors_per_leaf_override=( strict_max_neighbors_per_leaf_override ), - compact_far_pair_capacity_override=( - strict_compact_far_pair_capacity_override - ), ) 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 6d8e8190..b1d8c7ab 100644 --- a/jaccpot/runtime/fmm_prepare.py +++ b/jaccpot/runtime/fmm_prepare.py @@ -1467,7 +1467,6 @@ def _record_dual_artifact_substage(name: str, elapsed: float) -> None: ( runtime_traversal_config, strict_nbr_override, - strict_far_override, _strict_capacity_report, ) = self._strict_fused_capacity_handoff( runtime_traversal_config=runtime_traversal_config, @@ -1479,7 +1478,6 @@ 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_compact_far_pair_capacity_override=strict_far_override, theta=theta_val, mac_type=mac_type_val, dehnen_radius_scale=dehnen_radius_scale, @@ -3343,7 +3341,6 @@ def _strict_fused_capacity_handoff( ) -> tuple[ Optional[DualTreeTraversalConfig], Optional[int], - Optional[int], Callable[[dict], None], ]: """Carry the eager walk's validated capacities into the traced refresh. @@ -3353,19 +3350,26 @@ def _strict_fused_capacity_handoff( the preset caps say. Under the fused velocity-Verlet scan the SAME walk runs traced: the overflow flags are tracers, the ladder has one attempt, and yggdrax returns the truncated result. Measured 2026-09-06 at - N=200k, leaf 256, theta 0.6: the preset caps (queue 65536, 256 - neighbours per leaf) against rows of up to 781 cut the near field to - 15 % from step 2 on -- relative force error 3 -- with no diagnostic. + N=200k, leaf 256: the preset caps (queue 65536, 256 neighbours per leaf) + against rows of up to 781 cut the near field to 15 % from step 2 on -- + relative force error ~60 % at theta 0.6, 5.8 % at theta 1.0 -- with no + diagnostic. So the eager pass records what it needed (``_strict_fused_validated_caps``, via the returned report callback) and the traced pass raises its caps to - cover that with headroom: neighbours and far pairs to 1.5x the observed - maximum (next power of two), the queue to 2x the ladder's answer. The - queue cannot be verified after the fact -- its overflow flag never leaves - the trace -- hence the larger margin; the neighbour and far-pair caps ARE - re-checked inside the scan by ``strict_run_v2`` from - ``_strict_fused_traced_caps``, which the same callback records on the - traced build. + cover that with headroom: the neighbour cap to 1.5x the observed longest + row (next power of two), the queue to 2x the ladder's answer. The queue + cannot be verified after the fact -- its overflow flag never leaves the + trace -- hence the larger margin; the neighbour cap IS re-checked inside + the scan by ``strict_run_v2`` from ``_strict_fused_traced_caps``, which + the same callback records on the traced build. + + The compact far-pair cap is deliberately NOT carried over. It is checked + under tracing already (``_raw_to_compact_far_pairs`` raises through a + debug callback), so a too-small one fails loudly instead of truncating + silently, and it is set explicitly by + ``JACCPOT_STATIC_STRICT_FUSED_COMPACT_FAR_PAIR_CAP`` -- widening a cap the + caller named would make a deliberate memory bound a no-op. Parameters ---------- @@ -3382,8 +3386,6 @@ def _strict_fused_capacity_handoff( The traversal config to build with. Optional[int] ``max_neighbors_per_leaf`` floor for the strict streamed builder. - Optional[int] - Compact far-pair cap floor for the strict streamed builder. Callable[[dict], None] Report callback that stores the builder's capacities on the engine. """ @@ -3393,15 +3395,11 @@ def _pow2_ceil(value: int) -> int: return 1 << (value - 1).bit_length() nbr_override: Optional[int] = None - far_override: Optional[int] = None validated = getattr(self, "_strict_fused_validated_caps", None) if bool(suppress_host_side_effects) and isinstance(validated, dict): observed_rows = validated.get("max_neighbors_observed") if observed_rows is not None: nbr_override = _pow2_ceil(int(1.5 * int(observed_rows)) + 1) - observed_far = validated.get("far_pair_count") - if observed_far is not None: - far_override = _pow2_ceil(int(1.5 * int(observed_far)) + 1) validated_queue = validated.get("queue_capacity") if validated_queue is not None and runtime_traversal_config is not None: widened_queue = max( @@ -3426,7 +3424,7 @@ def _report(report: dict) -> None: else: self._strict_fused_validated_caps = dict(report) - return runtime_traversal_config, nbr_override, far_override, _report + return runtime_traversal_config, nbr_override, _report def _prepare_state_dual_and_downward_strict_streamed_fast( self, @@ -3493,7 +3491,6 @@ def _prepare_state_dual_and_downward_strict_streamed_fast( ( runtime_traversal_config, strict_nbr_override, - strict_far_override, _strict_capacity_report, ) = self._strict_fused_capacity_handoff( runtime_traversal_config=runtime_traversal_config, @@ -3505,7 +3502,6 @@ 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_compact_far_pair_capacity_override=strict_far_override, 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 16bd0a57..83f24359 100644 --- a/jaccpot/runtime/fmm_strict_run.py +++ b/jaccpot/runtime/fmm_strict_run.py @@ -835,10 +835,15 @@ def _static_target_block_capacity_ok( counts <= jnp.asarray(capacity, dtype=counts.dtype) ) # Traversal-capacity saturation guard. The traced refresh walk runs - # with fixed caps and yggdrax cannot raise on overflow under jit; a - # neighbour row that fills its cap, or a far-pair list that fills - # its buffer, means entries were dropped and the force is wrong. - # The caps are host constants recorded while the refresh traced + # with fixed caps, and a neighbour row that fills its cap -- or a + # far-pair list that fills its buffer -- means entries were dropped + # and the force is wrong. The neighbour cap is the one that fails + # SILENTLY: yggdrax's near-overflow flag is a tracer under jit and + # nothing reads it, which is the defect this guard exists for. The + # far-pair cap already raises through a debug callback in + # ``_raw_to_compact_far_pairs``, so its arm here is a second, cheap + # line of defence rather than the only one. Both caps are host + # constants recorded while the refresh traced # (``_strict_fused_traced_caps``), so this is a static comparison. traced_caps = getattr(self, "_strict_fused_traced_caps", None) if isinstance(traced_caps, dict):