From 774cb6b47073c77bb9bf1d8f99bea707cf9e22b9 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Mon, 7 Sep 2026 10:26:44 +0200 Subject: [PATCH 01/12] perf(nearfield): fold the self-leaf block into the leaf-pair Pallas kernel One extra unconditional pass per program over the program's own leaf, diagonal masked before the rsqrt, gated to chunk 0 on the source-chunk grid, replaces the lax.scan over leaves (_self_contributions) on the forward prepacked lane. JACCPOT_NEARFIELD_LEAFPAIR_FOLD_SELF=0 restores the scan; the cvjp lane keeps self_acc + cvjp(...) so the gradient path is byte-identical. Kernel name gains _self so the two variants never alias in the compilation cache. Tests: fold == pairs + _self_contributions to fp32 summation order on both grids, subtile lane indexing, softening 0 (no NaN, potential lane), wide accumulator, default-off byte parity, and a Triton (sm_80) run registered in the GPU gate. Co-Authored-By: Claude Fable 5.1 --- bench/gpu_gate.py | 4 + jaccpot/nearfield/_fast_lane.py | 62 +++++-- jaccpot/pallas/nearfield_fused_leaf.py | 162 +++++++++++++----- .../operators/test_pallas_nearfield_fused.py | 138 +++++++++++++++ 4 files changed, 317 insertions(+), 49 deletions(-) diff --git a/bench/gpu_gate.py b/bench/gpu_gate.py index 0aab64f6..c101fc53 100644 --- a/bench/gpu_gate.py +++ b/bench/gpu_gate.py @@ -91,6 +91,10 @@ "test_fused_pallas_complex_m2l_matches_the_pure_jax_lane_in_gradient[False]", "test_the_production_real_fused_m2l_kernel_carries_the_axis_derivative", "test_the_production_complex_fused_m2l_kernel_carries_the_axis_derivative", + # The near-field leaf-pair kernel had NO gate entry before the self-fold + # (plan "small leaves", Phase 1); its Triton lowering is what this checks. + "test_leafpair_include_self_gpu_matches_reference[None]", + "test_leafpair_include_self_gpu_matches_reference[3]", ) # Measured on an A100 sm_80 / jax 0.10.2 and documented in ARCHITECTURE.md §9, diff --git a/jaccpot/nearfield/_fast_lane.py b/jaccpot/nearfield/_fast_lane.py index 61d846b8..43e0d243 100644 --- a/jaccpot/nearfield/_fast_lane.py +++ b/jaccpot/nearfield/_fast_lane.py @@ -605,6 +605,7 @@ def _radix_fast_lane_prepacked_pallas( num_stages: int = ..., target_subtile: Optional[int] = ..., interpret: bool = ..., + include_self: bool = ..., ) -> Array: ... @@ -625,6 +626,7 @@ def _radix_fast_lane_prepacked_pallas( num_stages: int = ..., target_subtile: Optional[int] = ..., interpret: bool = ..., + include_self: bool = ..., ) -> Tuple[Array, Array]: ... @@ -645,6 +647,7 @@ def _radix_fast_lane_prepacked_pallas( num_stages: int = 1, target_subtile: Optional[int] = None, interpret: bool = False, + include_self: bool = False, ) -> Union[Array, Tuple[Array, Array]]: """Fused Pallas leaf-pair path over the compact prepacked source-leaf layout. @@ -652,7 +655,8 @@ def _radix_fast_lane_prepacked_pallas( used by the production fused near-field lane. Source leaves are gathered by id inside the kernel (no dense per-particle source materialization), then the leaf-major result is scattered to particle order. The intra-leaf self term - is handled separately by the caller, matching the pure-JAX path. + is handled separately by the caller, matching the pure-JAX path, unless + ``include_self`` folds it into the kernel. This is the forward half of the production differentiable near field: it is what :func:`_radix_fast_lane_prepacked_accel_cvjp` calls as its primal. Needs @@ -691,6 +695,12 @@ def _radix_fast_lane_prepacked_pallas( interpret : bool Run through Pallas' reference interpreter rather than Triton; the shipped callers hardcode ``False``. + include_self : bool + Fold the intra-leaf self term into the kernel (one extra leaf pass per + program, diagonal masked) instead of leaving it to the caller's + ``_self_contributions`` scan. See + :func:`jaccpot.pallas.nearfield_fused_leaf.nearfield_leafpair_pallas`. + Default False. Returns ------- @@ -737,6 +747,7 @@ def _radix_fast_lane_prepacked_pallas( target_subtile=target_subtile, interpret=interpret, source_chunk=(None if source_chunk <= 0 else int(source_chunk)), + include_self=bool(include_self), ) pair_acc = _scatter_contributions( @@ -1422,6 +1433,17 @@ def compute_leaf_p2p_accelerations_radix_fast_lane( ) pallas_pairs = pallas_available and has_materialized_sources pallas_prepacked = pallas_available and has_prepacked_sources + # Fold the self-leaf term into the prepacked Pallas kernel (plan "small + # leaves", Phase 1). The scan it replaces launches ~5 fusions per leaf batch, + # which is what made leaf 64 slower than leaf 256 with 3x fewer pair + # evaluations. Forward lane only: the differentiable prepacked lane keeps + # ``self_acc + cvjp(...)`` so its gradient path is byte-identical + # (JACCPOT_NEARFIELD_LEAFPAIR_FOLD_SELF=0 restores the scan everywhere). + fold_self = ( + pallas_prepacked + and not (differentiable and not want_potential) + and _env_flag("JACCPOT_NEARFIELD_LEAFPAIR_FOLD_SELF", True) + ) # Potential is only implemented on the fused Pallas paths; otherwise the # caller falls back to the generic W x W path (preserving prior behavior). @@ -1454,7 +1476,32 @@ def _zeros_result(): softening_sq = jnp.asarray(float(softening) ** 2, dtype=positions.dtype) self_acc = jnp.zeros_like(positions) self_pot = jnp.zeros(positions.shape[:1], dtype=dtype) - if diag_mode != "pairs_only": + pallas_num_warps = _env_int("JACCPOT_NEARFIELD_PALLAS_NUM_WARPS", 0) + pallas_num_stages = max(1, _env_int("JACCPOT_NEARFIELD_PALLAS_NUM_STAGES", 1)) + pallas_subtile = _env_int("JACCPOT_NEARFIELD_PALLAS_TARGET_SUBTILE", 0) + if fold_self and diag_mode == "self_only": + # The folded lane's self term alone: the kernel with every source slot + # invalid and the self pass on, so the diag mode measures the pass that + # replaced the scan rather than the scan itself. + source_leaf_ids_padded = jnp.asarray(payload.source_leaf_ids, dtype=INDEX_DTYPE) + return _radix_fast_lane_prepacked_pallas( + source_leaf_ids_padded, + jnp.zeros(source_leaf_ids_padded.shape, dtype=bool), + leaf_positions, + leaf_masses, + leaf_mask, + leaf_particle_idx, + positions, + G=G, + softening_sq=softening_sq, + compute_potential=want_potential, + num_warps=(pallas_num_warps if pallas_num_warps > 0 else None), + num_stages=pallas_num_stages, + target_subtile=(pallas_subtile if pallas_subtile > 0 else None), + interpret=pallas_interpret, + include_self=True, + ) + if diag_mode != "pairs_only" and not fold_self: if want_potential: ( self_acc, @@ -1484,9 +1531,6 @@ def _zeros_result(): return self_acc if pallas_pairs: - pallas_num_warps = _env_int("JACCPOT_NEARFIELD_PALLAS_NUM_WARPS", 0) - pallas_num_stages = max(1, _env_int("JACCPOT_NEARFIELD_PALLAS_NUM_STAGES", 1)) - pallas_subtile = _env_int("JACCPOT_NEARFIELD_PALLAS_TARGET_SUBTILE", 0) pairs_result = _radix_fast_lane_pairs_pallas( positions, masses, @@ -1520,11 +1564,6 @@ def _zeros_result(): ) if pallas_prepacked: - pallas_num_warps = _env_int("JACCPOT_NEARFIELD_PALLAS_NUM_WARPS", 0) - pallas_num_stages = max( - 1, _env_int("JACCPOT_NEARFIELD_PALLAS_NUM_STAGES", 1) - ) - pallas_subtile = _env_int("JACCPOT_NEARFIELD_PALLAS_TARGET_SUBTILE", 0) if differentiable and not want_potential: # Differentiable prepacked lane: the SAME Pallas forward wrapped in # a custom_vjp whose reverse is autodiff of this lane's own tiled @@ -1585,12 +1624,15 @@ def _zeros_result(): num_stages=pallas_num_stages, target_subtile=(pallas_subtile if pallas_subtile > 0 else None), interpret=pallas_interpret, + include_self=bool(fold_self), ) # Branch on the value, not on `want_potential`. Both say the same thing -- # the callee returns a pair exactly when the flag is set -- but the flag is # a runtime bool, so it selects the fallback overload and leaves the result # a union that cannot be added to an Array. `isinstance` narrows it, and a # JAX array is never a tuple, so the discriminator is exact. + # With ``fold_self`` the self term is already inside ``prepacked_result`` + # and ``self_acc`` / ``self_pot`` are the zeros initialised above. if isinstance(prepacked_result, tuple): pair_acc, pair_pot = prepacked_result return self_acc + pair_acc, self_pot + pair_pot diff --git a/jaccpot/pallas/nearfield_fused_leaf.py b/jaccpot/pallas/nearfield_fused_leaf.py index 0b409102..cee12dd6 100644 --- a/jaccpot/pallas/nearfield_fused_leaf.py +++ b/jaccpot/pallas/nearfield_fused_leaf.py @@ -577,7 +577,7 @@ def nearfield_fused_leaf_backend(*, prefer_pallas: bool = True) -> str: # --------------------------------------------------------------------------- -@jax.jit +@partial(jax.jit, static_argnames=("include_self",)) @jaxtyped(typechecker=beartype) def nearfield_leafpair_jax( leaf_positions: Array, @@ -588,6 +588,7 @@ def nearfield_leafpair_jax( *, softening_sq: Array, G: Array, + include_self: bool = False, ) -> Array: """Reference leaf-pair near-field update in pure JAX (dense; test-scale only). @@ -615,6 +616,11 @@ def nearfield_leafpair_jax( Scalar *squared* Plummer softening, added to every squared separation. G : Array Scalar gravitational constant, applied as a plain multiplier. + include_self : bool + Also add each leaf's intra-leaf term with the diagonal removed (the + dense form of ``_self_contributions``), matching + :func:`nearfield_leafpair_pallas` with ``include_self=True``. Default + False keeps the historical cross-leaf-only reference. Returns ------- @@ -654,8 +660,18 @@ def nearfield_leafpair_jax( inv_dist3 = inv_r * inv_r * inv_r weighted = inv_dist3 * src_mass[:, None, :, :] accels = -G * jnp.sum(weighted[..., None] * diff, axis=(2, 3)) # (L, W_t, 3) - accels = jnp.where(leaf_mask[..., None], accels, 0.0) potentials = -G * jnp.sum(inv_r * src_mass[:, None, :, :], axis=(2, 3)) + if include_self: + width = int(leaf_positions.shape[1]) + identity = jnp.eye(width, dtype=bool) + diff_s = leaf_positions[:, :, None, :] - leaf_positions[:, None, :, :] + dist_sq_s = jnp.sum(diff_s * diff_s, axis=-1) + softening_sq # (L, W, W) + mask_s = leaf_mask[:, :, None] & leaf_mask[:, None, :] & (~identity) + inv_r_s = jnp.where(mask_s, lax.rsqrt(jnp.where(mask_s, dist_sq_s, 1.0)), 0.0) + weighted_s = inv_r_s * inv_r_s * inv_r_s * leaf_masses[:, None, :] + accels = accels - G * jnp.sum(weighted_s[..., None] * diff_s, axis=2) + potentials = potentials - G * jnp.sum(inv_r_s * leaf_masses[:, None, :], axis=2) + accels = jnp.where(leaf_mask[..., None], accels, 0.0) potentials = jnp.where(leaf_mask, potentials, 0.0) return jnp.concatenate([accels, potentials[..., None]], axis=-1) @@ -709,6 +725,8 @@ def _nearfield_leafpair_kernel( leaf_width: int, accum_dtype: Any = None, out_dtype: Any = None, + include_self: bool = False, + self_on_first_chunk_only: bool = False, ) -> None: """Leaf-pair near-field update for one target subtile (vector of Bt targets). @@ -777,6 +795,22 @@ def _nearfield_leafpair_kernel( Dtype of ``out_ref``. ``None`` means the input dtype (one final downcast of the wide accumulator). The chunked grid passes ``accum_dtype`` so each chunk's partial keeps the wide precision until the caller's reduce. + include_self : bool + Also sum the program's OWN leaf (``pl.program_id(0)``) against this + target subtile, diagonal removed -- the intra-leaf term that + ``jaccpot.nearfield._kernels._self_contributions`` otherwise computes as + a ``lax.scan`` over leaves. Folding it here costs one more unconditional + leaf pass per program and removes that scan's per-leaf launches, which + is what makes small leaves affordable (plan "small leaves", Phase 1). + The diagonal is masked BEFORE ``rsqrt``: at ``softening_sq == 0`` the + self pair is ``rsqrt(0) * 0 = NaN``, and even softened it would add a + spurious ``-G m_i / eps`` to the potential lane. The cross-leaf slots + are untouched, so with this off the kernel is byte-identical to before. + Static. + self_on_first_chunk_only : bool + On the chunked grid (a third grid axis over source chunks) run the self + pass on chunk 0 only; otherwise it would be counted ``n_chunks`` times. + Static; the single-pass grid leaves it False. Returns ------- @@ -799,45 +833,75 @@ def _nearfield_leafpair_kernel( else (zero, zero, zero, zero) ) + def _leaf_pass(sid, acc, exclude_lane=None): + # One source leaf ``sid`` summed into ``acc``; ``exclude_lane`` is the + # ``(Bt,)`` vector of the targets' own slot indices for the self pass + # (diagonal out), ``None`` for a cross-leaf slot -- that path is the + # historical loop verbatim, op for op. + def _lane_body(j, acc): + acc_x, acc_y, acc_z, acc_p = acc + sx = src_table_pos_ref[sid, j, 0] + sy = src_table_pos_ref[sid, j, 1] + sz = src_table_pos_ref[sid, j, 2] + sm = src_table_mass_ref[sid, j] + lane_valid = src_table_mask_ref[sid, j] + dx = tx - sx + dy = ty - sy + dz = tz - sz + dist_sq = dx * dx + dy * dy + dz * dz + soft + active = tvalid & lane_valid + if exclude_lane is not None: + active = active & (exclude_lane != j) + safe_dist_sq = jnp.where(active, dist_sq, 1.0) + inv_r = lax.rsqrt(safe_dist_sq) + inv_r = jnp.where(active, inv_r, 0.0) + inv_dist3 = inv_r * inv_r * inv_r + scale = -g_value * inv_dist3 * sm + acc_x = acc_x + scale * dx + acc_y = acc_y + scale * dy + acc_z = acc_z + scale * dz + acc_p = acc_p - g_value * inv_r * sm + return (acc_x, acc_y, acc_z, acc_p) + + if not wide: + return lax.fori_loop(0, leaf_width, _lane_body, acc) + # Two-level: this leaf's contribution accumulates in the narrow input + # dtype (leaf_width terms, so its own round-off is negligible), and only + # the per-leaf total is added into the wide running accumulator. + part = lax.fori_loop(0, leaf_width, _lane_body, (zero, zero, zero, zero)) + return tuple(a + q.astype(accum_dtype) for a, q in zip(acc, part)) + def _slot_body(s, acc): sid = source_leaf_ids_ref[0, s] slot_valid = source_valid_ref[0, s] + return lax.cond( + slot_valid, lambda acc: _leaf_pass(sid, acc), lambda acc: acc, acc + ) - def _apply(acc): - def _lane_body(j, acc): - acc_x, acc_y, acc_z, acc_p = acc - sx = src_table_pos_ref[sid, j, 0] - sy = src_table_pos_ref[sid, j, 1] - sz = src_table_pos_ref[sid, j, 2] - sm = src_table_mass_ref[sid, j] - lane_valid = src_table_mask_ref[sid, j] - dx = tx - sx - dy = ty - sy - dz = tz - sz - dist_sq = dx * dx + dy * dy + dz * dz + soft - active = tvalid & lane_valid - safe_dist_sq = jnp.where(active, dist_sq, 1.0) - inv_r = lax.rsqrt(safe_dist_sq) - inv_r = jnp.where(active, inv_r, 0.0) - inv_dist3 = inv_r * inv_r * inv_r - scale = -g_value * inv_dist3 * sm - acc_x = acc_x + scale * dx - acc_y = acc_y + scale * dy - acc_z = acc_z + scale * dz - acc_p = acc_p - g_value * inv_r * sm - return (acc_x, acc_y, acc_z, acc_p) - - if not wide: - return lax.fori_loop(0, leaf_width, _lane_body, acc) - # Two-level: this leaf's contribution accumulates in the narrow input - # dtype (leaf_width terms, so its own round-off is negligible), and only - # the per-leaf total is added into the wide running accumulator. - part = lax.fori_loop(0, leaf_width, _lane_body, (zero, zero, zero, zero)) - return tuple(a + q.astype(accum_dtype) for a, q in zip(acc, part)) - - return lax.cond(slot_valid, _apply, lambda acc: acc, acc) - - acc_x, acc_y, acc_z, acc_p = lax.fori_loop(0, num_source_slots, _slot_body, acc0) + acc = lax.fori_loop(0, num_source_slots, _slot_body, acc0) + + if include_self: + # The self-leaf block, one extra unconditional pass rather than an extra + # source slot: the payload builders and their capacity checks stay + # untouched, and a leaf is never in its own source row (a precondition + # of the twin, asserted by the audit probe), so nothing is double + # counted. The target lane's slot index within the leaf is what the + # diagonal mask compares against the source lane ``j``. + bt = int(tx.shape[0]) + lane_idx = pl.program_id(1) * bt + lax.broadcasted_iota( + jnp.int32, (bt,), 0 + ) + own_leaf = pl.program_id(0) + + def _self_pass(acc): + return _leaf_pass(own_leaf, acc, exclude_lane=lane_idx) + + if self_on_first_chunk_only: + acc = lax.cond(pl.program_id(2) == 0, _self_pass, lambda acc: acc, acc) + else: + acc = _self_pass(acc) + + acc_x, acc_y, acc_z, acc_p = acc target_out_dtype = zero.dtype if out_dtype is None else out_dtype if wide and target_out_dtype != accum_dtype: # One downcast, on the FINAL value rather than on the sum being accumulated: @@ -871,6 +935,7 @@ def nearfield_leafpair_pallas( interpret: bool = False, accum: str = "input", source_chunk: int | None = None, + include_self: bool = False, ) -> Array: """Leaf-pair near-field update with Pallas. @@ -933,6 +998,16 @@ def nearfield_leafpair_pallas( source_chunk : int | None Source slots per program on the chunked grid (see above). ``None`` or a value of at least ``S`` keeps the single-pass grid. + include_self : bool + Also add each leaf's intra-leaf (self) term, diagonal removed, inside the + kernel -- see :func:`_nearfield_leafpair_kernel`. The result then equals + this kernel with it off plus + ``jaccpot.nearfield._kernels._self_contributions`` to float32 summation + order. Requires that no leaf appears in its own ``source_leaf_ids`` row + (the same precondition the twin documents), or the self term is counted + twice. Default False: byte-identical to the historical kernel; the + kernel name gains ``_self`` when set so the two never alias in the + compilation cache. Returns ------- @@ -959,6 +1034,8 @@ def nearfield_leafpair_pallas( source_valid = jnp.asarray(source_valid, dtype=bool) softening_sq_arr = jnp.asarray([softening_sq], dtype=dtype) g_arr = jnp.asarray([G], dtype=dtype) + include_self = bool(include_self) + self_tag = "_self" if include_self else "" if leaf_positions.ndim != 3 or leaf_positions.shape[-1] != 3: raise ValueError("leaf_positions must have shape (num_leaves, W, 3)") @@ -1002,6 +1079,7 @@ def _kernel(*refs): num_source_slots=num_source_slots, leaf_width=leaf_width, accum_dtype=accum_dtype, + include_self=include_self, ) kernel = pl.pallas_call( @@ -1029,7 +1107,10 @@ def _kernel(*refs): num_warps=int(num_warps), num_stages=int(num_stages) ), interpret=bool(interpret), - name=f"nearfield_leafpair_t{bt}_s{num_source_slots}_w{leaf_width}_a{accum}", + name=( + f"nearfield_leafpair_t{bt}_s{num_source_slots}_w{leaf_width}" + f"_a{accum}{self_tag}" + ), ) out = kernel( target_positions_padded, @@ -1065,6 +1146,9 @@ def _kernel_chunk(*refs): leaf_width=leaf_width, accum_dtype=accum_dtype, out_dtype=partial_dtype, + include_self=include_self, + # counted once, on chunk 0, not n_chunks times + self_on_first_chunk_only=True, ) kernel = pl.pallas_call( @@ -1095,7 +1179,7 @@ def _kernel_chunk(*refs): interpret=bool(interpret), name=( f"nearfield_leafpair_t{bt}_s{num_source_slots}_c{chunk}" - f"_w{leaf_width}_a{accum}" + f"_w{leaf_width}_a{accum}{self_tag}" ), ) partials = kernel( diff --git a/tests/unit/operators/test_pallas_nearfield_fused.py b/tests/unit/operators/test_pallas_nearfield_fused.py index fdb63716..d7924977 100644 --- a/tests/unit/operators/test_pallas_nearfield_fused.py +++ b/tests/unit/operators/test_pallas_nearfield_fused.py @@ -475,3 +475,141 @@ def test_decoupled_equal_widths_are_untouched_and_finite(): assert bool(jnp.all(jnp.isfinite(out))) # Non-vacuity: a kernel returning zeros would satisfy "finite" trivially. assert float(jnp.max(jnp.abs(out))) > 0.0 + + +# --------------------------------------------------------------------------- +# include_self: the self-leaf block folded into the leaf-pair kernel (plan +# "small leaves", Phase 1). Reference = the kernel with the fold OFF plus the +# scan it replaces (``_self_contributions``), so the test pins the fold to the +# code it retires rather than to a re-derivation. +# --------------------------------------------------------------------------- + + +def _self_scan_reference(lp, lm, lmask, *, soft, G): + from jaccpot.nearfield._kernels import _self_contributions + + acc, pot = _self_contributions( + lp, lm, lmask, softening_sq=soft, G=G, compute_potential=True + ) + return np.concatenate([np.asarray(acc), np.asarray(pot)[..., None]], axis=-1) + + +@pytest.mark.parametrize("source_chunk", [None, 2]) +def test_leafpair_include_self_interpret_equals_pairs_plus_self_scan(source_chunk): + """Fold on == fold off + ``_self_contributions``, to float32 summation order. + + ``source_chunk=2`` exercises the chunked grid, where the self pass must run + on chunk 0 only -- counted ``n_chunks`` times it would be off by a factor + ``ceil(S / chunk)`` on the self term, which this catches. + """ + lp, lm, lmask, sids, svalid = _leafpair_inputs(seed=11, L=6, W=8, S=5) + soft = jnp.float32(0.03**2) + G = jnp.float32(1.2) + common = dict(softening_sq=soft, G=G, interpret=True, source_chunk=source_chunk) + pairs_only = nearfield_leafpair_pallas( + lp, lm, lmask, sids, svalid, include_self=False, **common + ) + folded = nearfield_leafpair_pallas( + lp, lm, lmask, sids, svalid, include_self=True, **common + ) + expected = np.asarray(pairs_only) + _self_scan_reference(lp, lm, lmask, soft=soft, G=G) + # non-vacuity: the self term must actually be present + assert not np.allclose(np.asarray(folded), np.asarray(pairs_only), atol=1e-6) + assert np.allclose(np.asarray(folded), expected, rtol=1e-5, atol=1e-6) + # and the dense twin agrees with both + twin = nearfield_leafpair_jax( + lp, lm, lmask, sids, svalid, softening_sq=soft, G=G, include_self=True + ) + assert np.allclose(np.asarray(twin), expected, rtol=1e-5, atol=1e-6) + + +def test_leafpair_include_self_interpret_subtile_lane_index(): + """With several target subtiles per leaf the diagonal mask must use the + lane's index WITHIN THE LEAF (``program_id(1) * Bt + iota``), not within + the subtile -- a wrong offset excludes the wrong source lane and keeps the + true self pair, which at softening 0 is a NaN.""" + lp, lm, lmask, sids, svalid = _leafpair_inputs(seed=12, L=5, W=8, S=4) + soft = jnp.float32(0.0) + G = jnp.float32(1.0) + folded = nearfield_leafpair_pallas( + lp, lm, lmask, sids, svalid, softening_sq=soft, G=G, + interpret=True, include_self=True, target_subtile=4, + ) + pairs_only = nearfield_leafpair_pallas( + lp, lm, lmask, sids, svalid, softening_sq=soft, G=G, + interpret=True, include_self=False, target_subtile=4, + ) + expected = np.asarray(pairs_only) + _self_scan_reference(lp, lm, lmask, soft=soft, G=G) + assert np.all(np.isfinite(np.asarray(folded))) + assert np.allclose(np.asarray(folded), expected, rtol=1e-5, atol=1e-6) + + +def test_leafpair_include_self_softening_zero_has_no_nan_and_matches_potential(): + """At ``softening_sq == 0`` the self pair is ``rsqrt(0) * 0``; the diagonal + must be masked BEFORE the rsqrt (acceleration finite) and contribute nothing + to the potential lane.""" + lp, lm, lmask, sids, svalid = _leafpair_inputs(seed=13, L=4, W=8, S=3) + soft = jnp.float32(0.0) + G = jnp.float32(0.7) + folded = nearfield_leafpair_pallas( + lp, lm, lmask, sids, svalid, softening_sq=soft, G=G, interpret=True, include_self=True + ) + got = np.asarray(folded) + assert np.all(np.isfinite(got)) + pairs_only = nearfield_leafpair_pallas( + lp, lm, lmask, sids, svalid, softening_sq=soft, G=G, interpret=True, include_self=False + ) + expected = np.asarray(pairs_only) + _self_scan_reference(lp, lm, lmask, soft=soft, G=G) + assert np.allclose(got[..., 3], expected[..., 3], rtol=1e-5, atol=1e-6) + assert np.allclose(got, expected, rtol=1e-5, atol=1e-6) + + +def test_leafpair_include_self_wide_accumulator_interpret(): + """The wide (two-level float64) accumulator takes the self block through the + same per-leaf partial.""" + lp, lm, lmask, sids, svalid = _leafpair_inputs(seed=14, L=5, W=8, S=4) + soft = jnp.float32(0.02**2) + G = jnp.float32(1.0) + folded = nearfield_leafpair_pallas( + lp, lm, lmask, sids, svalid, softening_sq=soft, G=G, + interpret=True, include_self=True, accum="wide", source_chunk=2, + ) + ref = nearfield_leafpair_jax( + lp, lm, lmask, sids, svalid, softening_sq=soft, G=G, include_self=True + ) + assert np.allclose(np.asarray(folded), np.asarray(ref), rtol=1e-5, atol=1e-6) + + +def test_leafpair_include_self_off_is_the_historical_kernel(): + """Default off: byte-identical to the cross-leaf-only kernel (the twin's + reference), so nothing outside the fold moved.""" + lp, lm, lmask, sids, svalid = _leafpair_inputs(seed=15) + soft = jnp.float32(0.05**2) + G = jnp.float32(1.3) + ref = nearfield_leafpair_jax(lp, lm, lmask, sids, svalid, softening_sq=soft, G=G) + got = nearfield_leafpair_pallas( + lp, lm, lmask, sids, svalid, softening_sq=soft, G=G, interpret=True + ) + assert np.allclose(np.asarray(got), np.asarray(ref), rtol=1e-5, atol=1e-6) + + +@pytest.mark.skipif( + not pallas_nearfield_fused_supported(), + reason="leaf-pair near-field Pallas kernel requires an Ampere+ (sm_80+) GPU", +) +@pytest.mark.parametrize("source_chunk", [None, 3]) +def test_leafpair_include_self_gpu_matches_reference(source_chunk): + """The Triton lowering of the self pass (program_id-indexed gather, iota mask, + chunk-0 gate) against the interpret-validated reference.""" + lp, lm, lmask, sids, svalid = _leafpair_inputs(seed=16, L=10, W=16, S=6) + soft = jnp.float32(0.02**2) + G = jnp.float32(1.1) + ref = nearfield_leafpair_jax( + lp, lm, lmask, sids, svalid, softening_sq=soft, G=G, include_self=True + ) + got = nearfield_leafpair_pallas( + lp, lm, lmask, sids, svalid, softening_sq=soft, G=G, + target_subtile=8, interpret=False, include_self=True, source_chunk=source_chunk, + ) + assert np.all(np.isfinite(np.asarray(got))) + assert np.allclose(np.asarray(got), np.asarray(ref), rtol=1e-5, atol=1e-5) From 8ea81ada5f218c19fb3a2844224c3a42527ba7db Mon Sep 17 00:00:00 2001 From: TobiBu Date: Mon, 7 Sep 2026 13:14:10 +0200 Subject: [PATCH 02/12] perf(m2l): contention-free per-chunk reduction into the locals _chunk_segment_scatter_add ran once per 4096-pair chunk and its two scatters were pathological for XLA's atomic-add lowering: segment_sum sent every pair of a target to the same group row, and the final .at[].add sent every non-head slot (~3800 of 4096) to node 0 with a zero value. At leaf 64 (992k far pairs, 243 chunks/step) that was 169 ms of a 410 ms step, 686 us per launch for a 4096 x 25 reduction. Now a segmented associative_scan reduces within the sorted segments with no scatter, and non-head slots carry an out-of-bounds sink index dropped by mode='drop', so the one remaining scatter has unique in-bounds indices and no atomics. Pinned to np.add.at on random chunks (padded tail, all-same-target, nothing valid, degenerate width), node-0 untouched unless targeted, deterministic. Co-Authored-By: Claude Fable 5.1 --- jaccpot/runtime/kernels/_m2l.py | 61 ++++++++----- .../test_m2l_chunk_segment_scatter.py | 90 +++++++++++++++++++ 2 files changed, 129 insertions(+), 22 deletions(-) create mode 100644 tests/unit/operators/test_m2l_chunk_segment_scatter.py diff --git a/jaccpot/runtime/kernels/_m2l.py b/jaccpot/runtime/kernels/_m2l.py index 462a4163..79c05c87 100644 --- a/jaccpot/runtime/kernels/_m2l.py +++ b/jaccpot/runtime/kernels/_m2l.py @@ -402,14 +402,30 @@ def _chunk_segment_scatter_add( """Reduce one fixed-width chunk by target index and scatter-add into locals. Sorts the chunk by target so that contributions to the same target become a - contiguous segment, reduces within segments, then scatters once. Invalid - slots are given the maximum index so they sort to the end and fall outside - the scatter. + contiguous segment, reduces within segments with a segmented prefix scan, + then scatters the one total per segment. Invalid slots are given the + maximum index so they sort to the end and fall outside the scatter. The sort makes the summation order a deterministic function of the target indices rather than of the pair order, which is what keeps the four accumulators agreeing to reassociation. + Why a segmented ``associative_scan`` and an out-of-bounds sink, not + ``segment_sum`` and index 0 (measured 2026-09-07, N=200k Plummer, A100, + ``strict_run_v2``): this function ran once per 4096-pair chunk, and at leaf + 64 (992k far pairs, 243 chunks per step) its scatter fusions cost 169 ms of + a 410 ms step -- 686 us per launch for a 4096 x 25 reduction. Both scatters + in the old body were pathological for XLA's atomic-add lowering: the + ``segment_sum`` sent every pair of a target to the SAME group row (up to a + few hundred duplicates per address, serialised), and the final + ``.at[safe_targets].add`` sent every slot that was not a segment head -- + ~3800 of 4096 -- to node 0 with a zero value, another serialised address. + The segmented scan reduces within segments with no scatter at all, and the + non-head slots now carry an index one past the end of ``local_accum``, + which ``mode="drop"`` discards without a write. Each in-bounds index is + then unique within the chunk (one segment head per target), so the final + scatter needs no atomics either. + Parameters ---------- local_accum : Array @@ -434,28 +450,29 @@ def _chunk_segment_scatter_add( tgt_sorted = tgt_chunk[sort_idx] contribs_sorted = contribs[sort_idx] valid_sorted = valid[sort_idx] - contribs_sorted = jnp.where(valid_sorted[:, None], contribs_sorted, 0) - new_group = jnp.concatenate( - ( - jnp.asarray([True], dtype=bool), - sorted_keys[1:] != sorted_keys[:-1], - ), - axis=0, + + boundary = sorted_keys[1:] != sorted_keys[:-1] + new_group = jnp.concatenate((jnp.ones((1,), dtype=bool), boundary), axis=0) + is_last = jnp.concatenate((boundary, jnp.ones((1,), dtype=bool)), axis=0) + + def _segmented_add(a: tuple[Array, Array], b: tuple[Array, Array]): + # Prefix sums that restart at every segment head: the right operand + # replaces the running sum when it starts a segment, else it adds. + va, fa = a + vb, fb = b + return jnp.where(fb[..., None], vb, va + vb), fa | fb + + segment_prefix, _ = jax.lax.associative_scan( + _segmented_add, (contribs_sorted, new_group), axis=0 ) - group_ids = jnp.cumsum(new_group.astype(INDEX_DTYPE)) - jnp.asarray( - 1, - dtype=INDEX_DTYPE, + take = is_last & valid_sorted + sink = jnp.asarray(local_accum.shape[0], dtype=INDEX_DTYPE) # out of bounds + rows_tgt = jnp.where(take, tgt_sorted, sink) + rows_val = jnp.where(take[:, None], segment_prefix, 0) + return local_accum.at[rows_tgt].add( + rows_val, mode="drop", indices_are_sorted=True, unique_indices=True ) - reduced = jax.ops.segment_sum(contribs_sorted, group_ids, chunk_size) - - unique_targets = jnp.zeros((chunk_size,), dtype=INDEX_DTYPE) - unique_targets = unique_targets.at[group_ids].set(tgt_sorted) - unique_valid = jnp.zeros((chunk_size,), dtype=bool) - unique_valid = unique_valid.at[group_ids].set(valid_sorted) - safe_targets = jnp.where(unique_valid, unique_targets, 0) - reduced = jnp.where(unique_valid[:, None], reduced, 0) - return local_accum.at[safe_targets].add(reduced) @jaxtyped(typechecker=beartype) diff --git a/tests/unit/operators/test_m2l_chunk_segment_scatter.py b/tests/unit/operators/test_m2l_chunk_segment_scatter.py new file mode 100644 index 00000000..210dbc4b --- /dev/null +++ b/tests/unit/operators/test_m2l_chunk_segment_scatter.py @@ -0,0 +1,90 @@ +"""``_chunk_segment_scatter_add``: the per-chunk M2L reduction into the locals. + +Pinned to a numpy ``np.add.at`` reference because the 2026-09-07 rewrite +(segmented ``associative_scan`` + out-of-bounds sink instead of ``segment_sum`` ++ index-0 dummies) changed every line of the body: the function must still add +exactly the valid rows to exactly their targets, count a target that appears in +many rows once per row, leave every other node untouched, and be deterministic. +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from jaccpot.runtime.kernels._m2l import _chunk_segment_scatter_add + + +def _case(seed, *, chunk_size, total_nodes, ncoef, n_valid, max_target): + rng = np.random.default_rng(seed) + local = rng.standard_normal((total_nodes, ncoef)).astype(np.float32) + contribs = rng.standard_normal((chunk_size, ncoef)).astype(np.float32) + tgt = rng.integers(0, max_target, size=chunk_size).astype(np.int32) + valid = np.zeros(chunk_size, bool) + valid[:n_valid] = True + rng.shuffle(valid) + # invalid slots may hold anything, including out-of-range garbage + tgt = np.where(valid, tgt, rng.integers(-5, total_nodes + 7, size=chunk_size)).astype(np.int32) + expected = local.astype(np.float64).copy() + np.add.at(expected, tgt[valid], contribs[valid].astype(np.float64)) + return local, contribs, tgt, valid, expected + + +@pytest.mark.parametrize( + "chunk_size,total_nodes,n_valid,max_target", + [ + (64, 40, 64, 40), # full chunk, every target hit many times + (64, 40, 37, 40), # padded tail + (64, 400, 64, 400), # mostly unique targets + (64, 40, 64, 1), # every row the SAME target (the serialised case) + (64, 40, 0, 40), # nothing valid: the accumulator is untouched + (1, 3, 1, 3), # degenerate width + ], +) +def test_chunk_segment_scatter_add_matches_np_add_at( + chunk_size, total_nodes, n_valid, max_target +): + local, contribs, tgt, valid, expected = _case( + 7, chunk_size=chunk_size, total_nodes=total_nodes, ncoef=9, + n_valid=n_valid, max_target=max_target, + ) + got = _chunk_segment_scatter_add( + jnp.asarray(local), jnp.asarray(contribs), jnp.asarray(tgt), jnp.asarray(valid), + chunk_size=chunk_size, + ) + got = np.asarray(got, np.float64) + assert got.shape == expected.shape + scale = np.max(np.abs(expected)) + 1.0 + assert np.allclose(got, expected, rtol=0, atol=2e-5 * scale) + if n_valid == 0: + np.testing.assert_array_equal(got, local.astype(np.float64)) + + +def test_chunk_segment_scatter_add_hits_node_zero_only_when_targeted(): + """Node 0 used to receive every non-head slot as a zero add; now it must be + touched only by rows that target it -- checked with a chunk that never does.""" + local, contribs, tgt, valid, expected = _case( + 3, chunk_size=32, total_nodes=16, ncoef=4, n_valid=32, max_target=16 + ) + tgt = np.where(tgt == 0, 5, tgt).astype(np.int32) + expected = local.astype(np.float64).copy() + np.add.at(expected, tgt[valid], contribs[valid].astype(np.float64)) + got = np.asarray(_chunk_segment_scatter_add( + jnp.asarray(local), jnp.asarray(contribs), jnp.asarray(tgt), jnp.asarray(valid), + chunk_size=32, + ), np.float64) + np.testing.assert_array_equal(got[0], local[0].astype(np.float64)) + assert np.allclose(got, expected, rtol=0, atol=1e-4) + + +def test_chunk_segment_scatter_add_is_deterministic_under_jit(): + local, contribs, tgt, valid, _ = _case( + 11, chunk_size=256, total_nodes=50, ncoef=25, n_valid=250, max_target=50 + ) + fn = jax.jit(lambda a, c, t, v: _chunk_segment_scatter_add(a, c, t, v, chunk_size=256)) + args = (jnp.asarray(local), jnp.asarray(contribs), jnp.asarray(tgt), jnp.asarray(valid)) + a = np.asarray(fn(*args)) + b = np.asarray(fn(*args)) + np.testing.assert_array_equal(a, b) From 20d61f1206e046c455c0bf8c054e39e3b1a27ae8 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Wed, 9 Sep 2026 17:49:47 +0200 Subject: [PATCH 03/12] feat(pallas): target-tiled real M2L kernel with on-chip rotations (CSR by target) One program per target node loops over its far-pair segment of a target-sorted source list and builds rotate -> z-translate -> rotate-back on chip from the two alignment angles: Dz as diag(cos |m| t) + A sin(|m| t) with a constant antisymmetric pattern, the constant B stack per degree, and a separable z-core Z = Zsf * outer(rinv^(n+1), rinv^k). Nothing per pair touches HBM except one (Cp,) multipole row; the program owns its output row, so the chunked lane's per-chunk argsort + segment_sum scatter disappear. Forward only. Interpret-mode parity vs m2l_rot_scale_real_batch: < 1e-10 (fp64) for orders 2..6, < 3e-4 (fp32); random CSR with empty targets and a padded -1 tail, active_pair_count truncation, on-axis deltas, and the wrapper under jit. Co-Authored-By: Claude Fable 5.1 --- jaccpot/pallas/m2l_real_csr.py | 563 ++++++++++++++++++ .../operators/test_m2l_real_csr_pallas.py | 190 ++++++ 2 files changed, 753 insertions(+) create mode 100644 jaccpot/pallas/m2l_real_csr.py create mode 100644 tests/unit/operators/test_m2l_real_csr_pallas.py diff --git a/jaccpot/pallas/m2l_real_csr.py b/jaccpot/pallas/m2l_real_csr.py new file mode 100644 index 00000000..f107291a --- /dev/null +++ b/jaccpot/pallas/m2l_real_csr.py @@ -0,0 +1,563 @@ +"""Target-tiled real-basis M2L Pallas kernel with on-chip rotations (CSR by target). + +One program per TARGET node. The program loops over that target's far-pair +segment of a source list sorted by target (CSR), and for every source builds the +whole rotate -> z-translate -> rotate-back M2L on chip from two angles, summing +into one register-resident local row. It exists because the chunked pure-JAX +lane (``runtime/kernels/_m2l.py``) is launch-bound: at N=200k, leaf 64 the far +field is a 21k-launch/step storm (~130 ns per directed pair with nothing above +27 ms in the kernel table), and the two shipped fused Pallas M2L shapes lose +because they take the world<->z rotation blocks as ``(pairs, p+1, 2p+1, 2p+1)`` +HBM operands (32 KB per pair). Here nothing per pair is materialised: the +program owns its output row, so there is no ``segment_sum`` scatter and no +argsort per chunk, and the coefficient traffic is one ``(Cp,)`` row load per pair. + +Arithmetic, per pair, in the CENTRED padded layout of +:func:`jaccpot.operators.m2l_real_rot_scale._centred_degree_maps` (degree ``l`` +occupies columns ``p-l .. p+l`` of a width ``2p+1`` row, so ``m`` sits at column +``p+m`` for every degree at once): + +* world -> z multipole block, degree ``l``: ``B_l Dz(-ax) B_l Dz(az)`` with + ``az = atan2(x, y)``, ``ax = atan2(rho, z)`` (the conventions of + :func:`jaccpot.operators.real_rotations._multipole_align_to_z_block`). + ``B_l`` is a compile-time constant stack; ``Dz(t) v = cos(|m| t) * v + + A (sin(|m| t) * v)`` with ``A`` the constant antisymmetric pattern + ``A[p+m, p-m] = -1, A[p-m, p+m] = +1`` -- read off + :func:`jaccpot.operators.real_rotations.real_Dz_diagonal`. +* z-core: ``F_n^m = sum_k sign(m) (n+k)! r^-(n+k+1) M_k^m`` from + :func:`jaccpot.operators.real_harmonics.z_m2l_translation_tables`, the single + source of truth, as the SEPARABLE dense form + ``Z = Zsf * outer(rinv^(n+1), rinv^k)`` so the radius enters through two + ``(Cp,)`` power vectors, not ``Cp^2`` transcendental calls. +* z -> world local block = transpose of the multipole block + (:func:`jaccpot.operators.real_rotations.real_rotation_from_z_axis_local`): + ``Dz(-az) B_l^T Dz(ax) B_l^T``. + +Every contraction is a broadcast-multiply + ``jnp.sum`` (no ``dot``, no TF32), +as in :mod:`jaccpot.pallas.m2l_real_fused`. Padded lanes of every constant are +exactly zero, so they are inert in every reduction. + +Forward only (the fused strict lane is forward-only); the pure-JAX lane stays the +differentiable path, and the transverse-degeneracy JVP treatment of +``m2l_rot_scale_real_batch`` is not needed here. + +HARDWARE: real Pallas GPU execution needs Ampere (sm_80+); ``interpret=True`` +runs the same arithmetic on CPU. +""" + +from __future__ import annotations + +import functools +import math +from typing import Any, Optional + +import jax +import jax.numpy as jnp +import numpy as np +from jax import lax +from jax.experimental import pallas as pl +from jaxtyping import Array + +from jaccpot.operators.real_dehnen_q import compute_real_B_matrix_multipole +from jaccpot.operators.real_harmonics import ( + sh_offset, + sh_size, + z_m2l_translation_tables, +) +from jaccpot.pallas._compat import KernelRef, pallas_backend_kwargs +from jaccpot.pallas.m2l_real_fused import pallas_m2l_real_fused_supported + +__all__ = [ + "pallas_m2l_real_csr_supported", + "m2l_real_csr_tables", + "m2l_real_csr_pair_jax", + "m2l_real_csr_jax", + "m2l_real_csr_pallas", + "csr_by_target", +] + +_TABLE_KEYS = ( + "Ppack", + "Uunpack", + "Bstack", + "BstackT", + "Apat", + "mabs", + "Zsf", + "PowOut", + "PowSrc", +) + + +def pallas_m2l_real_csr_supported() -> bool: + """True on a GPU with compute capability >= 8.0 (same predicate as the fused kernel). + + Returns + ------- + bool + Whether the Triton lowering of this kernel can run here. + """ + return pallas_m2l_real_fused_supported() + + +def _next_pow2(n: int) -> int: + n = max(1, int(n)) + return 1 << (n - 1).bit_length() + + +@functools.lru_cache(maxsize=None) +def m2l_real_csr_tables(order: int) -> dict: + """Compile-time constants of the kernel for one expansion order. + + Parameters + ---------- + order : int + Expansion order ``p``. + + Returns + ------- + dict + NumPy float64 arrays (cast to the working dtype by the caller) plus the + shape scalars: ``C = (p+1)^2``, ``Cp`` (pow2 >= C), ``W = 2p+1``, + ``Wp`` (pow2 >= W), ``Bp`` (pow2 >= p+1), ``K`` (pow2 >= p+2) radius powers. + + ``Ppack [Bp*Wp, Cp]`` / ``Uunpack [Cp, Bp*Wp]``: one-hot pack/unpack + between the packed coefficient vector and the centred ``(Bp, Wp)`` rows. + ``Bstack [Bp, Wp, Wp]``: ``B_U(l)`` centred per degree; ``BstackT`` its + per-degree transpose. ``Apat [Wp, Wp]``: the Dz sine pattern. + ``mabs [Wp]``: ``|m|`` per column (0 on padded columns). + ``Zsf [Cp, Cp]``: ``sign(m) (n+k)!`` on the valid (out, src) entries. + ``PowOut [Cp, K]`` / ``PowSrc [Cp, K]``: one-hot selectors of the radius + power ``rinv^(n+1)`` per output lane and ``rinv^k`` per source lane. + """ + p = int(order) + if p < 0: + raise ValueError("order must be >= 0") + C = sh_size(p) + W = 2 * p + 1 + Cp = _next_pow2(C) + Wp = _next_pow2(W) + Bp = _next_pow2(p + 1) + K = _next_pow2(p + 2) # radius powers 0..p+1, padded to a pow2 width for Triton + + Ppack = np.zeros((Bp * Wp, Cp), dtype=np.float64) + for ell in range(p + 1): + for m in range(-ell, ell + 1): + Ppack[ell * Wp + (p + m), sh_offset(ell) + ell + m] = 1.0 + Uunpack = Ppack.T.copy() + + Bstack = np.zeros((Bp, Wp, Wp), dtype=np.float64) + # `compute_real_B_matrix_multipole` is jitted: evaluated eagerly here so the + # table is a literal even when this cache is first filled inside a trace. + with jax.ensure_compile_time_eval(): + for ell in range(p + 1): + b = np.asarray( + compute_real_B_matrix_multipole(ell, dtype=jnp.float64), dtype=np.float64 + ) + Bstack[ell, p - ell : p + ell + 1, p - ell : p + ell + 1] = b + BstackT = np.swapaxes(Bstack, -1, -2).copy() + + Apat = np.zeros((Wp, Wp), dtype=np.float64) + mabs = np.zeros((Wp,), dtype=np.float64) + for m in range(1, p + 1): + Apat[p + m, p - m] = -1.0 + Apat[p - m, p + m] = 1.0 + mabs[p + m] = float(m) + mabs[p - m] = float(m) + + src_index, valid, fact_index, r_exponent, sign = z_m2l_translation_tables(p) + fact = np.asarray([math.factorial(i) for i in range(2 * p + 1)], dtype=np.float64) + Zsf = np.zeros((Cp, Cp), dtype=np.float64) + deg_of = np.zeros((C,), dtype=np.int64) + for n in range(p + 1): + deg_of[sh_offset(n) : sh_offset(n + 1)] = n + for out in range(C): + for k in range(p + 1): + if bool(valid[out, k]): + s = int(src_index[out, k]) + Zsf[out, s] = float(sign[out]) * float(fact[int(fact_index[out, k])]) + # r^-(n+k+1) = rinv^(n+1) * rinv^k with n = deg(out), k = deg(s) + assert int(r_exponent[out, k]) == deg_of[out] + 1 + deg_of[s] + PowOut = np.zeros((Cp, K), dtype=np.float64) + PowSrc = np.zeros((Cp, K), dtype=np.float64) + for i in range(C): + PowOut[i, deg_of[i] + 1] = 1.0 + PowSrc[i, deg_of[i]] = 1.0 + return dict( + p=p, C=C, Cp=Cp, W=W, Wp=Wp, Bp=Bp, K=K, + Ppack=Ppack, Uunpack=Uunpack, Bstack=Bstack, BstackT=BstackT, + Apat=Apat, mabs=mabs, Zsf=Zsf, PowOut=PowOut, PowSrc=PowSrc, + ) + + +def _tables_to_jnp(order: int, dtype: Any) -> dict[str, Array]: + t = m2l_real_csr_tables(order) + return {k: jnp.asarray(t[k], dtype=dtype) for k in _TABLE_KEYS} + + +# --------------------------------------------------------------------------- math +# Every helper below is written for BOTH the Pallas kernel (values loaded from +# refs) and the pure-jnp twin: plain broadcast-multiply + sum, no dot, no gather. + + +def _matvec(mat: Array, vec: Array) -> Array: + return jnp.sum(mat * vec[None, :], axis=1) + + +def _bapply(bstack: Array, rows: Array) -> Array: + """``out[l, i] = sum_j bstack[l, i, j] rows[l, j]`` (block-diagonal by degree).""" + return jnp.sum(bstack * rows[:, None, :], axis=-1) + + +def _dz(rows: Array, cosv: Array, sinv: Array, apat: Array) -> Array: + """``Dz(t)`` on every degree row at once: ``cos(|m|t) v + A (sin(|m|t) v)``.""" + sv = rows * sinv[None, :] + return rows * cosv[None, :] + jnp.sum(apat[None, :, :] * sv[:, None, :], axis=-1) + + +def _radius_powers(rinv: Array, k: int) -> Array: + """``[1, rinv, rinv^2, ..., rinv^(k-1)]`` by repeated multiplication (exact integer powers).""" + pw = [jnp.ones_like(rinv)] + for _ in range(1, k): + pw.append(pw[-1] * rinv) + return jnp.stack(pw) + + +def _m2l_pair(mult: Array, delta3: tuple, t: dict[str, Array], *, bp: int, wp: int) -> Array: + """Full real M2L for one pair from the packed row ``mult`` and ``delta = c_t - c_s``. + + Parameters + ---------- + mult : Array + Padded source multipole row ``(Cp,)``. + delta3 : tuple + ``(x, y, z)`` scalars, target centre minus source centre. + t : dict[str, Array] + Tables from :func:`_tables_to_jnp` at the working dtype. + bp : int + Padded degree count ``Bp``. Static. + wp : int + Padded row width ``Wp``. Static. + + Returns + ------- + Array + Padded local contribution ``(Cp,)``. + """ + x, y, z = delta3 + dtype = mult.dtype + rho2 = x * x + y * y + rho = jnp.sqrt(rho2) + az = jnp.arctan2(x, y) + ax = jnp.arctan2(rho, z) + r = jnp.sqrt(rho2 + z * z) + r = jnp.maximum(r, jnp.asarray(1.0e-30, dtype=dtype)) + rinv = 1.0 / r + mabs = t["mabs"] + cos_az = jnp.cos(mabs * az) + sin_az = jnp.sin(mabs * az) + cos_ax = jnp.cos(mabs * ax) + sin_ax = jnp.sin(mabs * ax) + apat = t["Apat"] + + # world -> z (multipole): B Dz(-ax) B Dz(az), applied right to left + rows = _matvec(t["Ppack"], mult).reshape(bp, wp) + rows = _dz(rows, cos_az, sin_az, apat) + rows = _bapply(t["Bstack"], rows) + rows = _dz(rows, cos_ax, -sin_ax, apat) + rows = _bapply(t["Bstack"], rows) + mrf = _matvec(t["Uunpack"], rows.reshape(bp * wp)) + + # z-core, separable radius powers + pw = _radius_powers(rinv, int(t["PowOut"].shape[1])) + rinv_out = jnp.sum(t["PowOut"] * pw[None, :], axis=1) + rinv_src = jnp.sum(t["PowSrc"] * pw[None, :], axis=1) + lz = rinv_out * _matvec(t["Zsf"], rinv_src * mrf) + + # z -> world (local) = transpose of the multipole block: Dz(-az) B^T Dz(ax) B^T + rows = _matvec(t["Ppack"], lz).reshape(bp, wp) + rows = _bapply(t["BstackT"], rows) + rows = _dz(rows, cos_ax, sin_ax, apat) + rows = _bapply(t["BstackT"], rows) + rows = _dz(rows, cos_az, -sin_az, apat) + return _matvec(t["Uunpack"], rows.reshape(bp * wp)) + + +# ----------------------------------------------------------------- pure-jnp twin + + +def m2l_real_csr_pair_jax(multipoles: Array, deltas: Array, *, order: int) -> Array: + """Per-pair local contributions with this kernel's arithmetic (the twin). + + Parameters + ---------- + multipoles : Array + ``[N, C]`` source multipoles. + deltas : Array + ``[N, 3]`` target minus source centres. + order : int + Expansion order ``p``. Static. + + Returns + ------- + Array + ``[N, C]`` local contributions, one per pair (NOT reduced by target). + """ + tb = m2l_real_csr_tables(int(order)) + C, Cp, Bp, Wp = tb["C"], tb["Cp"], tb["Bp"], tb["Wp"] + mult = jnp.asarray(multipoles) + dtype = mult.dtype + t = _tables_to_jnp(int(order), dtype) + mult_p = jnp.pad(mult, ((0, 0), (0, Cp - C))) + d = jnp.asarray(deltas, dtype=dtype) + + def one(m, dd): + return _m2l_pair(m, (dd[0], dd[1], dd[2]), t, bp=Bp, wp=Wp) + + return jax.vmap(one)(mult_p, d)[:, :C] + + +def csr_by_target( + sources: Array, + targets: Array, + *, + total_nodes: int, + active_pair_count: Optional[Array] = None, +) -> tuple[Array, Array, Array]: + """Sort a (padded) flat far-pair list by target into CSR form. + + Parameters + ---------- + sources : Array + ``[P]`` source node ids; ``-1`` (or anything negative) marks padding. + targets : Array + ``[P]`` target node ids, aligned; negative marks padding. + total_nodes : int + Number of target rows. Static. + active_pair_count : Optional[Array] + Number of leading live entries; ``None`` means every non-negative entry + is live. + + Returns + ------- + tuple[Array, Array, Array] + ``(sources_sorted [P], offsets [total_nodes], counts [total_nodes])``: + target ``t``'s sources are ``sources_sorted[offsets[t] : offsets[t] + + counts[t]]``. Padding sorts to the end and is never addressed. + """ + src = jnp.asarray(sources, dtype=jnp.int32) + tgt = jnp.asarray(targets, dtype=jnp.int32) + P = int(src.shape[0]) + valid = (src >= 0) & (tgt >= 0) + if active_pair_count is not None: + valid = valid & (jnp.arange(P, dtype=jnp.int32) < jnp.asarray(active_pair_count, jnp.int32)) + key = jnp.where(valid, tgt, jnp.asarray(total_nodes, jnp.int32)) + perm = jnp.argsort(key, stable=True) + src_sorted = jnp.where(valid[perm], src[perm], 0) + counts = jax.ops.segment_sum( + valid.astype(jnp.int32), jnp.where(valid, tgt, 0), num_segments=int(total_nodes) + ) + offsets = jnp.cumsum(counts) - counts + return src_sorted, offsets.astype(jnp.int32), counts.astype(jnp.int32) + + +def m2l_real_csr_jax( + multipoles: Array, + centers: Array, + sources: Array, + targets: Array, + *, + order: int, + active_pair_count: Optional[Array] = None, +) -> Array: + """Reference: per-pair twin reduced by target with ``segment_sum``. + + Parameters + ---------- + multipoles : Array + ``[n, C]`` node multipoles. + centers : Array + ``[n, 3]`` node centres. + sources : Array + ``[P]`` source ids (negative = padding). + targets : Array + ``[P]`` target ids (negative = padding). + order : int + Expansion order. Static. + active_pair_count : Optional[Array] + Live prefix length, as in :func:`csr_by_target`. + + Returns + ------- + Array + ``[n, C]`` local coefficient increments. + """ + n = int(multipoles.shape[0]) + src = jnp.asarray(sources, jnp.int32) + tgt = jnp.asarray(targets, jnp.int32) + valid = (src >= 0) & (tgt >= 0) + if active_pair_count is not None: + valid = valid & (jnp.arange(src.shape[0], dtype=jnp.int32) < jnp.asarray(active_pair_count, jnp.int32)) + s = jnp.where(valid, src, 0) + tt = jnp.where(valid, tgt, 0) + deltas = centers[tt] - centers[s] + contrib = m2l_real_csr_pair_jax(multipoles[s], deltas, order=order) + contrib = jnp.where(valid[:, None], contrib, 0) + return jax.ops.segment_sum(contrib, tt, num_segments=n) + + +# -------------------------------------------------------------------- the kernel + + +def _m2l_real_csr_kernel( + mult_ref: KernelRef, + cent_ref: KernelRef, + src_ref: KernelRef, + off_ref: KernelRef, + cnt_ref: KernelRef, + *table_and_out_refs: KernelRef, + bp: int, + wp: int, +) -> None: + """One program per target: loop over its CSR segment, accumulate one local row. + + Parameters + ---------- + mult_ref : KernelRef + Whole padded multipole table ``[n, Cp]`` (gathered by source id). + cent_ref : KernelRef + Whole padded centre table ``[n, 4]``. + src_ref : KernelRef + Whole target-sorted source list ``[P]``. + off_ref : KernelRef + Segment start per target ``[n]``. + cnt_ref : KernelRef + Segment length per target ``[n]``. + *table_and_out_refs : KernelRef + The ``_TABLE_KEYS`` constants (whole arrays) followed by the output ref + ``[1, Cp]``. + bp : int + ``Bp``. Static. + wp : int + ``Wp``. Static. + + Returns + ------- + None + Writes the target's local row. + """ + table_refs = table_and_out_refs[: len(_TABLE_KEYS)] + (out_ref,) = table_and_out_refs[len(_TABLE_KEYS) :] + t = {k: ref[...] for k, ref in zip(_TABLE_KEYS, table_refs)} + tgt = pl.program_id(0) + start = off_ref[tgt] + cnt = cnt_ref[tgt] + ctx = cent_ref[tgt, 0] + cty = cent_ref[tgt, 1] + ctz = cent_ref[tgt, 2] + cp = int(out_ref.shape[1]) + acc0 = jnp.zeros((cp,), dtype=out_ref.dtype) + + def body(k, acc): + sid = src_ref[start + k] + m = mult_ref[sid, :] + dx = ctx - cent_ref[sid, 0] + dy = cty - cent_ref[sid, 1] + dz_ = ctz - cent_ref[sid, 2] + return acc + _m2l_pair(m, (dx, dy, dz_), t, bp=bp, wp=wp) + + acc = lax.fori_loop(0, cnt, body, acc0) + out_ref[0, :] = acc + + +def m2l_real_csr_pallas( + multipoles: Array, + centers: Array, + sources: Array, + targets: Array, + *, + order: int, + active_pair_count: Optional[Array] = None, + interpret: bool = False, + backend: str = "triton", + num_warps: int = 1, +) -> Array: + """Local coefficient increments from a flat far-pair list, one Pallas program per target. + + Parameters + ---------- + multipoles : Array + ``[n, C]`` node multipoles (real basis). + centers : Array + ``[n, 3]`` node centres. + sources : Array + ``[P]`` source node ids; negative = padding. + targets : Array + ``[P]`` target node ids; negative = padding. Need NOT be sorted -- the + list is sorted by target here (one argsort of ``P`` keys). + order : int + Expansion order ``p``. Static. + active_pair_count : Optional[Array] + Live prefix length of the padded list; ``None`` = all non-negative. + interpret : bool + Pallas interpret mode (CPU semantics). + backend : str + Pallas GPU lowering, ``"triton"`` by default. + num_warps : int + Warps per program; the row width is ``Cp`` lanes, one warp suffices. + + Returns + ------- + Array + ``[n, C]`` local increments, same dtype as ``multipoles``. + """ + tb = m2l_real_csr_tables(int(order)) + C, Cp, Bp, Wp = tb["C"], tb["Cp"], tb["Bp"], tb["Wp"] + mult = jnp.asarray(multipoles) + dtype = mult.dtype + n = int(mult.shape[0]) + if int(mult.shape[1]) != C: + raise ValueError(f"multipoles must have {C} coefficients for order {order}") + cent = jnp.asarray(centers, dtype=dtype) + if cent.ndim != 2 or int(cent.shape[1]) != 3 or int(cent.shape[0]) != n: + raise ValueError("centers must have shape (n, 3) aligned with multipoles") + mult_p = jnp.pad(mult, ((0, 0), (0, Cp - C))) + cent_p = jnp.pad(cent, ((0, 0), (0, 1))) + src_sorted, offsets, counts = csr_by_target( + sources, targets, total_nodes=n, active_pair_count=active_pair_count + ) + P = int(src_sorted.shape[0]) + if n == 0 or P == 0: + return jnp.zeros((n, C), dtype=dtype) + tables = _tables_to_jnp(int(order), dtype) + table_arrays = [tables[k] for k in _TABLE_KEYS] + + def bs_full(arr: Array) -> pl.BlockSpec: + shp = tuple(arr.shape) + return pl.BlockSpec(shp, (lambda *_: (0,) * len(shp))) + + kernel = functools.partial(_m2l_real_csr_kernel, bp=Bp, wp=Wp) + backend_kwargs = pallas_backend_kwargs(backend, interpret) + if "compiler_params" in backend_kwargs: + # one Cp-lane row per program: a single warp is the right launch shape + backend_kwargs["compiler_params"] = type(backend_kwargs["compiler_params"])( + num_warps=int(num_warps) + ) + out = pl.pallas_call( + kernel, + grid=(n,), + in_specs=[ + bs_full(mult_p), + bs_full(cent_p), + bs_full(src_sorted), + bs_full(offsets), + bs_full(counts), + *[bs_full(a) for a in table_arrays], + ], + out_specs=pl.BlockSpec((1, Cp), lambda i: (i, 0)), + out_shape=jax.ShapeDtypeStruct((n, Cp), dtype), + interpret=bool(interpret), + **backend_kwargs, + name=f"m2l_real_csr_p{int(order)}", + )(mult_p, cent_p, src_sorted, offsets, counts, *table_arrays) + return out[:, :C] diff --git a/tests/unit/operators/test_m2l_real_csr_pallas.py b/tests/unit/operators/test_m2l_real_csr_pallas.py new file mode 100644 index 00000000..ac749857 --- /dev/null +++ b/tests/unit/operators/test_m2l_real_csr_pallas.py @@ -0,0 +1,190 @@ +"""Parity tests for the target-tiled (CSR) real-basis M2L Pallas kernel. + +``jaccpot.pallas.m2l_real_csr`` builds the rotations on chip from two angles and +sums each target's far pairs inside one program. Pinned here, in interpret mode +so it runs on CPU CI, against + +* ``m2l_rot_scale_real_batch`` -- the pure-JAX rotate/scale M2L that is THE + reference for every real M2L lane (rel err < 1e-10 at fp64, < 3e-4 at fp32, + the tolerances of ``test_m2l_real_fused_pallas.py``), reduced per target with + ``segment_sum``; +* the kernel's own pure-jnp twin (``m2l_real_csr_jax``), a literal port. + +Cases: random CSR with empty targets and a padded ``-1`` tail, an +``active_pair_count`` shorter than the live prefix, on-axis deltas (``rho = 0``, +where the alignment azimuth is undefined and the forward must still be exact), +and orders 2..6 (``Cp`` switches 16 -> 32 -> 64 across that range). +""" + +from __future__ import annotations + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from jaccpot.operators.m2l_real_rot_scale import m2l_rot_scale_real_batch +from jaccpot.operators.real_harmonics import sh_size +from jaccpot.pallas.m2l_real_csr import ( + csr_by_target, + m2l_real_csr_jax, + m2l_real_csr_pair_jax, + m2l_real_csr_pallas, + pallas_m2l_real_csr_supported, +) + + +def _case(order, dtype, *, n=12, pairs=40, seed=0, on_axis=False, pad=7): + rng = np.random.default_rng(seed) + c = sh_size(order) + mult = rng.standard_normal((n, c)).astype(dtype) + centers = (rng.standard_normal((n, 3)) * 3.0).astype(dtype) + if on_axis: + centers[:, :2] = 0.0 # every delta along z: rho == 0 + centers[:, 2] = np.arange(n) * 2.5 + # random pairs, targets 0..n-3 so the last two targets are EMPTY + tgt = rng.integers(0, n - 2, size=pairs).astype(np.int32) + src = rng.integers(0, n, size=pairs).astype(np.int32) + src = np.where(src == tgt, (src + 1) % n, src).astype(np.int32) + if on_axis: + src = np.where(src == tgt, (src + 2) % n, src).astype(np.int32) + tgt_p = np.concatenate([tgt, -np.ones(pad, np.int32)]) + src_p = np.concatenate([src, -np.ones(pad, np.int32)]) + return mult, centers, src_p, tgt_p + + +def _reference(mult, centers, src, tgt, order, active=None): + """pure-JAX rot-scale M2L per pair + segment_sum, on the live pairs only.""" + n = mult.shape[0] + valid = (src >= 0) & (tgt >= 0) + if active is not None: + valid &= np.arange(src.shape[0]) < active + s, t = src[valid], tgt[valid] + deltas = centers[t] - centers[s] + contrib = np.asarray(m2l_rot_scale_real_batch(jnp.asarray(mult[s]), jnp.asarray(deltas), order=order)) + out = np.zeros_like(mult, dtype=np.float64) + np.add.at(out, t, contrib.astype(np.float64)) + return out + + +def _relerr(a, ref): + return float(np.linalg.norm(np.asarray(a, np.float64) - ref) / (np.linalg.norm(ref) + 1e-30)) + + +def test_csr_by_target_partitions_the_live_pairs(): + _, _, src, tgt = _case(3, np.float32, n=10, pairs=33, seed=2) + n = 10 + ss, off, cnt = csr_by_target(jnp.asarray(src), jnp.asarray(tgt), total_nodes=n) + ss, off, cnt = np.asarray(ss), np.asarray(off), np.asarray(cnt) + live = tgt >= 0 + assert int(cnt.sum()) == int(live.sum()) + assert cnt[-2:].sum() == 0 # the two empty targets + np.testing.assert_array_equal(off, np.cumsum(cnt) - cnt) + for t in range(n): + got = sorted(ss[off[t] : off[t] + cnt[t]].tolist()) + assert got == sorted(src[live & (tgt == t)].tolist()) + + +@pytest.mark.parametrize("order", [2, 3, 4, 5, 6]) +def test_csr_pair_twin_matches_rot_scale_f64(order): + """The on-chip rotation assembly equals the pure-JAX rotate/scale per pair.""" + if not jax.config.jax_enable_x64: + pytest.skip("float64 disabled in this JAX runtime") + rng = np.random.default_rng(order) + c = sh_size(order) + mult = jnp.asarray(rng.standard_normal((25, c))) + deltas = rng.standard_normal((25, 3)) * 2.0 + deltas[:, 2] += 3.0 + deltas = jnp.asarray(deltas) + ref = np.asarray(m2l_rot_scale_real_batch(mult, deltas, order=order)) + got = m2l_real_csr_pair_jax(mult, deltas, order=order) + assert _relerr(got, ref) < 1e-10 + + +@pytest.mark.parametrize("order", [2, 3, 4, 5, 6]) +def test_csr_pallas_interpret_matches_rot_scale_f64(order): + if not jax.config.jax_enable_x64: + pytest.skip("float64 disabled in this JAX runtime") + mult, centers, src, tgt = _case(order, np.float64, seed=order) + ref = _reference(mult, centers, src, tgt, order) + got = m2l_real_csr_pallas( + jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), jnp.asarray(tgt), + order=order, interpret=True, + ) + assert got.shape == mult.shape + assert _relerr(got, ref) < 1e-10 + # empty targets stay exactly zero + assert np.all(np.asarray(got)[-2:] == 0.0) + + +@pytest.mark.parametrize("order", [2, 4]) +def test_csr_pallas_interpret_matches_twin_and_rot_scale_f32(order): + mult, centers, src, tgt = _case(order, np.float32, seed=10 + order) + ref = _reference(mult, centers, src, tgt, order) + got = m2l_real_csr_pallas( + jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), jnp.asarray(tgt), + order=order, interpret=True, + ) + twin = m2l_real_csr_jax( + jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), jnp.asarray(tgt), order=order + ) + assert _relerr(got, ref) < 3e-4 + assert _relerr(got, np.asarray(twin, np.float64)) < 1e-5 + + +def test_csr_pallas_interpret_active_pair_count_truncates(): + if not jax.config.jax_enable_x64: + pytest.skip("float64 disabled in this JAX runtime") + order = 3 + mult, centers, src, tgt = _case(order, np.float64, pairs=40, seed=5) + active = 23 + ref = _reference(mult, centers, src, tgt, order, active=active) + got = m2l_real_csr_pallas( + jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), jnp.asarray(tgt), + order=order, active_pair_count=jnp.asarray(active, jnp.int32), interpret=True, + ) + assert _relerr(got, ref) < 1e-10 + + +def test_csr_pallas_interpret_on_axis_deltas_are_exact(): + """rho == 0: atan2(0, 0) = 0 must give the exact forward, no NaN.""" + if not jax.config.jax_enable_x64: + pytest.skip("float64 disabled in this JAX runtime") + order = 4 + mult, centers, src, tgt = _case(order, np.float64, seed=8, on_axis=True) + ref = _reference(mult, centers, src, tgt, order) + got = np.asarray(m2l_real_csr_pallas( + jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), jnp.asarray(tgt), + order=order, interpret=True, + )) + assert np.all(np.isfinite(got)) + assert _relerr(got, ref) < 1e-10 + + +def test_csr_pallas_under_jit_with_traced_active_count(): + """The wrapper (sort + CSR + pallas_call) must trace: caps are static, counts traced.""" + order = 3 + mult, centers, src, tgt = _case(order, np.float32, seed=11) + fn = jax.jit(lambda m, c, s, t, a: m2l_real_csr_pallas(m, c, s, t, order=order, + active_pair_count=a, interpret=True)) + got = fn(jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), jnp.asarray(tgt), + jnp.asarray(40, jnp.int32)) + ref = _reference(mult, centers, src, tgt, order) + assert _relerr(got, ref) < 3e-4 + + +@pytest.mark.skipif( + not pallas_m2l_real_csr_supported(), + reason="CSR M2L Pallas kernel needs an Ampere+ (sm_80) GPU", +) +@pytest.mark.parametrize("order", [2, 4, 6]) +def test_csr_pallas_gpu_matches_rot_scale(order): + """The Triton lowering: dynamic-trip loop, row gathers by id, atan2, pow2 tiles.""" + mult, centers, src, tgt = _case(order, np.float32, n=40, pairs=400, seed=20 + order) + ref = _reference(mult, centers, src, tgt, order) + got = m2l_real_csr_pallas( + jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), jnp.asarray(tgt), + order=order, interpret=False, + ) + assert np.all(np.isfinite(np.asarray(got))) + assert _relerr(got, ref) < 3e-4 From 51b7725a2565bb9d49cc750e619e7de468cdc0d5 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Wed, 9 Sep 2026 17:54:27 +0200 Subject: [PATCH 04/12] fix(pallas): CSR M2L radius powers as exp(deg * log rinv) -- Triton has no n-ary stack Co-Authored-By: Claude Fable 5.1 --- jaccpot/pallas/m2l_real_csr.py | 44 ++++++++++++++-------------------- 1 file changed, 18 insertions(+), 26 deletions(-) diff --git a/jaccpot/pallas/m2l_real_csr.py b/jaccpot/pallas/m2l_real_csr.py index f107291a..e2c20182 100644 --- a/jaccpot/pallas/m2l_real_csr.py +++ b/jaccpot/pallas/m2l_real_csr.py @@ -28,7 +28,8 @@ :func:`jaccpot.operators.real_harmonics.z_m2l_translation_tables`, the single source of truth, as the SEPARABLE dense form ``Z = Zsf * outer(rinv^(n+1), rinv^k)`` so the radius enters through two - ``(Cp,)`` power vectors, not ``Cp^2`` transcendental calls. + ``(Cp,)`` power vectors (``exp(deg * log rinv)``, the form the fused kernel + lowers), not ``Cp^2`` transcendental calls. * z -> world local block = transpose of the multipole block (:func:`jaccpot.operators.real_rotations.real_rotation_from_z_axis_local`): ``Dz(-az) B_l^T Dz(ax) B_l^T``. @@ -84,8 +85,8 @@ "Apat", "mabs", "Zsf", - "PowOut", - "PowSrc", + "DegOut", + "DegSrc", ) @@ -119,7 +120,7 @@ def m2l_real_csr_tables(order: int) -> dict: dict NumPy float64 arrays (cast to the working dtype by the caller) plus the shape scalars: ``C = (p+1)^2``, ``Cp`` (pow2 >= C), ``W = 2p+1``, - ``Wp`` (pow2 >= W), ``Bp`` (pow2 >= p+1), ``K`` (pow2 >= p+2) radius powers. + ``Wp`` (pow2 >= W), ``Bp`` (pow2 >= p+1). ``Ppack [Bp*Wp, Cp]`` / ``Uunpack [Cp, Bp*Wp]``: one-hot pack/unpack between the packed coefficient vector and the centred ``(Bp, Wp)`` rows. @@ -127,8 +128,9 @@ def m2l_real_csr_tables(order: int) -> dict: per-degree transpose. ``Apat [Wp, Wp]``: the Dz sine pattern. ``mabs [Wp]``: ``|m|`` per column (0 on padded columns). ``Zsf [Cp, Cp]``: ``sign(m) (n+k)!`` on the valid (out, src) entries. - ``PowOut [Cp, K]`` / ``PowSrc [Cp, K]``: one-hot selectors of the radius - power ``rinv^(n+1)`` per output lane and ``rinv^k`` per source lane. + ``DegOut [Cp]`` / ``DegSrc [Cp]``: the radius exponents ``n+1`` per + output lane and ``k`` per source lane (0 on padded lanes, where ``Zsf`` + is zero anyway). """ p = int(order) if p < 0: @@ -138,7 +140,6 @@ def m2l_real_csr_tables(order: int) -> dict: Cp = _next_pow2(C) Wp = _next_pow2(W) Bp = _next_pow2(p + 1) - K = _next_pow2(p + 2) # radius powers 0..p+1, padded to a pow2 width for Triton Ppack = np.zeros((Bp * Wp, Cp), dtype=np.float64) for ell in range(p + 1): @@ -178,15 +179,14 @@ def m2l_real_csr_tables(order: int) -> dict: Zsf[out, s] = float(sign[out]) * float(fact[int(fact_index[out, k])]) # r^-(n+k+1) = rinv^(n+1) * rinv^k with n = deg(out), k = deg(s) assert int(r_exponent[out, k]) == deg_of[out] + 1 + deg_of[s] - PowOut = np.zeros((Cp, K), dtype=np.float64) - PowSrc = np.zeros((Cp, K), dtype=np.float64) - for i in range(C): - PowOut[i, deg_of[i] + 1] = 1.0 - PowSrc[i, deg_of[i]] = 1.0 + DegOut = np.zeros((Cp,), dtype=np.float64) + DegSrc = np.zeros((Cp,), dtype=np.float64) + DegOut[:C] = deg_of + 1 + DegSrc[:C] = deg_of return dict( - p=p, C=C, Cp=Cp, W=W, Wp=Wp, Bp=Bp, K=K, + p=p, C=C, Cp=Cp, W=W, Wp=Wp, Bp=Bp, Ppack=Ppack, Uunpack=Uunpack, Bstack=Bstack, BstackT=BstackT, - Apat=Apat, mabs=mabs, Zsf=Zsf, PowOut=PowOut, PowSrc=PowSrc, + Apat=Apat, mabs=mabs, Zsf=Zsf, DegOut=DegOut, DegSrc=DegSrc, ) @@ -215,14 +215,6 @@ def _dz(rows: Array, cosv: Array, sinv: Array, apat: Array) -> Array: return rows * cosv[None, :] + jnp.sum(apat[None, :, :] * sv[:, None, :], axis=-1) -def _radius_powers(rinv: Array, k: int) -> Array: - """``[1, rinv, rinv^2, ..., rinv^(k-1)]`` by repeated multiplication (exact integer powers).""" - pw = [jnp.ones_like(rinv)] - for _ in range(1, k): - pw.append(pw[-1] * rinv) - return jnp.stack(pw) - - def _m2l_pair(mult: Array, delta3: tuple, t: dict[str, Array], *, bp: int, wp: int) -> Array: """Full real M2L for one pair from the packed row ``mult`` and ``delta = c_t - c_s``. @@ -268,10 +260,10 @@ def _m2l_pair(mult: Array, delta3: tuple, t: dict[str, Array], *, bp: int, wp: i rows = _bapply(t["Bstack"], rows) mrf = _matvec(t["Uunpack"], rows.reshape(bp * wp)) - # z-core, separable radius powers - pw = _radius_powers(rinv, int(t["PowOut"].shape[1])) - rinv_out = jnp.sum(t["PowOut"] * pw[None, :], axis=1) - rinv_src = jnp.sum(t["PowSrc"] * pw[None, :], axis=1) + # z-core, separable radius powers: r^-(n+k+1) = rinv^(n+1) * rinv^k + log_rinv = jnp.log(rinv) + rinv_out = jnp.exp(t["DegOut"] * log_rinv) + rinv_src = jnp.exp(t["DegSrc"] * log_rinv) lz = rinv_out * _matvec(t["Zsf"], rinv_src * mrf) # z -> world (local) = transpose of the multipole block: Dz(-az) B^T Dz(ax) B^T From 567bde3c1ff8f0803bd56a18fbe92e99ab8c7e33 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Wed, 9 Sep 2026 17:58:24 +0200 Subject: [PATCH 05/12] feat(m2l): opt-in CSR Pallas lane for the flat real-basis M2L + G2a microbench JACCPOT_STATIC_STRICT_FUSED_M2L_CSR=1 (Ampere+, or JACCPOT_M2L_CSR_INTERPRET=1 for CPU parity) routes _solidfmm_downward_accumulate_from_multipoles' flat real lanes to m2l_real_csr_pallas: one program per target owns its local row, so the chunked scan and its per-chunk scatter are gone. Wiring test counts the kernel entry and pins the force to the chunked lane (rel < 2e-5 at N=3000, p=3). bench/m2l_csr_microbench.py: ns/pair vs the pure-JAX chunked lane (degree-batched off/on) at 1M and 6M pairs, orders 4 and 6 -- the plan's G2a gate. Co-Authored-By: Claude Fable 5.1 --- bench/m2l_csr_microbench.py | 124 ++++++++++++++++++ jaccpot/runtime/kernels/_downward_prep.py | 44 ++++++- .../unit/runtime/test_m2l_csr_lane_wiring.py | 91 +++++++++++++ 3 files changed, 257 insertions(+), 2 deletions(-) create mode 100644 bench/m2l_csr_microbench.py create mode 100644 tests/unit/runtime/test_m2l_csr_lane_wiring.py diff --git a/bench/m2l_csr_microbench.py b/bench/m2l_csr_microbench.py new file mode 100644 index 00000000..29503711 --- /dev/null +++ b/bench/m2l_csr_microbench.py @@ -0,0 +1,124 @@ +"""Go/no-go microbench for the CSR M2L Pallas kernel (plan "small leaves", gate G2a). + +Synthetic far-pair lists of 1M and 6M directed pairs over ``n`` nodes (the leaf +64 / leaf 32 counts at N=200k), orders 4 and 6, fp32: + +* ``m2l_real_csr_pallas`` (one program per target, rotations on chip), jitted; +* the pure-JAX chunked lane as production runs it -- ``m2l_rot_scale_real_batch`` + over 4096-pair chunks in a ``lax.scan`` with ``_chunk_segment_scatter_add`` -- + with ``JACCPOT_M2L_DEGREE_BATCHED`` off and on (the plan's 2.0 candidate, and + the honest pure-JAX baseline). + +Reports ns per directed pair (min of the timed calls) and the fp32 rel-L2 of +the kernel against the pure-JAX lane. Gate: <= 15 ns/pair at p=4 and parity +<= 1e-5 rel is "go" for wiring; the pure-JAX lane sits at ~130-375 ns/pair. + + CUDA_VISIBLE_DEVICES= python bench/m2l_csr_microbench.py [--pairs 1000000 6000000] [--orders 4 6] +""" + +from __future__ import annotations + +import argparse +import json +import os +import time + +import numpy as np + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--pairs", type=int, nargs="+", default=[1_000_000, 6_000_000]) + ap.add_argument("--orders", type=int, nargs="+", default=[4, 6]) + ap.add_argument("--nodes", type=int, default=12_500) + ap.add_argument("--repeats", type=int, default=5) + ap.add_argument("--chunk", type=int, default=4096) + ap.add_argument("--out", default=None) + args = ap.parse_args() + os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") + + import jax + import jax.numpy as jnp + from jax import lax + + from jaccpot.operators.m2l_real_rot_scale import m2l_rot_scale_real_batch + from jaccpot.operators.real_harmonics import sh_size + from jaccpot.pallas.m2l_real_csr import m2l_real_csr_pallas + from jaccpot.runtime.kernels._m2l import _chunk_segment_scatter_add + + dev = jax.devices()[0] + print(f"device {dev} cc={getattr(dev, 'compute_capability', '?')}", flush=True) + rng = np.random.default_rng(0) + n = int(args.nodes) + centers = jnp.asarray(rng.uniform(-1, 1, (n, 3)).astype(np.float32)) + results = [] + + def timed(fn, *a): + out = jax.block_until_ready(fn(*a)) + ts = [] + for _ in range(args.repeats): + t0 = time.perf_counter() + out = jax.block_until_ready(fn(*a)) + ts.append(time.perf_counter() - t0) + return out, min(ts), float(np.median(ts)) + + for order in args.orders: + C = sh_size(order) + mult = jnp.asarray(rng.standard_normal((n, C)).astype(np.float32)) + for P in args.pairs: + # random well-separated pairs: reject |delta| < 0.3 (the MAC would too) + src = rng.integers(0, n, P).astype(np.int32) + tgt = rng.integers(0, n, P).astype(np.int32) + d = np.asarray(centers)[tgt] - np.asarray(centers)[src] + bad = np.linalg.norm(d, axis=1) < 0.3 + src = np.where(bad, (src + 1) % n, src).astype(np.int32) + src_j, tgt_j = jnp.asarray(src), jnp.asarray(tgt) + counts = np.bincount(tgt, minlength=n) + + csr = jax.jit(lambda m, c, s, t: m2l_real_csr_pallas(m, c, s, t, order=order)) + out_k, t_k, med_k = timed(csr, mult, centers, src_j, tgt_j) + + chunk = int(args.chunk) + n_chunks = -(-P // chunk) + pad = n_chunks * chunk - P + src_p = jnp.pad(src_j, (0, pad), constant_values=0) + tgt_p = jnp.pad(tgt_j, (0, pad), constant_values=0) + + def pure(m, c, s, t, P=P, chunk=chunk, n_chunks=n_chunks): + def body(acc, i): + idx = i * chunk + jnp.arange(chunk, dtype=jnp.int32) + valid = idx < P + sc = s[idx] + tc = t[idx] + contrib = m2l_rot_scale_real_batch(m[sc], c[tc] - c[sc], order=order) + return _chunk_segment_scatter_add(acc, contrib, tc, valid, chunk_size=chunk), None + acc, _ = lax.scan(body, jnp.zeros((n, C), jnp.float32), jnp.arange(n_chunks, dtype=jnp.int32)) + return acc + + rows = {} + for db in ("0", "1"): + os.environ["JACCPOT_M2L_DEGREE_BATCHED"] = db + fn = jax.jit(pure) + out_p, t_p, med_p = timed(fn, mult, centers, src_p, tgt_p) + rows[db] = (out_p, t_p, med_p) + del fn + ref = np.asarray(rows["0"][0], np.float64) + rel = float(np.linalg.norm(np.asarray(out_k, np.float64) - ref) / np.linalg.norm(ref)) + rel_db = float(np.linalg.norm(np.asarray(rows["1"][0], np.float64) - ref) / np.linalg.norm(ref)) + row = dict(order=order, pairs=P, nodes=n, longest_row=int(counts.max()), + csr_ns_per_pair=1e9 * t_k / P, csr_ms=1e3 * t_k, csr_median_ms=1e3 * med_k, + pure_ns_per_pair=1e9 * rows["0"][1] / P, pure_ms=1e3 * rows["0"][1], + pure_db_ns_per_pair=1e9 * rows["1"][1] / P, pure_db_ms=1e3 * rows["1"][1], + rel_l2_csr_vs_pure=rel, rel_l2_db_vs_pure=rel_db, chunk=chunk) + results.append(row) + print(f"p={order} pairs={P/1e6:.0f}M longest_row={counts.max()}: CSR {row['csr_ns_per_pair']:.1f} ns/pair " + f"({row['csr_ms']:.1f} ms) | pure-JAX {row['pure_ns_per_pair']:.1f} ns/pair | degree-batched " + f"{row['pure_db_ns_per_pair']:.1f} ns/pair | rel-L2 csr {rel:.2e}, db {rel_db:.2e}", flush=True) + if args.out: + with open(args.out, "w") as fh: + json.dump(results, fh, indent=2) + print("wrote", args.out) + + +if __name__ == "__main__": + main() diff --git a/jaccpot/runtime/kernels/_downward_prep.py b/jaccpot/runtime/kernels/_downward_prep.py index 2e27e743..6aa3cd1a 100644 --- a/jaccpot/runtime/kernels/_downward_prep.py +++ b/jaccpot/runtime/kernels/_downward_prep.py @@ -56,6 +56,30 @@ __all__: list[str] = [] +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. + + Returns + ------- + bool + True when :func:`jaccpot.pallas.m2l_real_csr.m2l_real_csr_pallas` + replaces the full-batch / chunked flat lanes for the real basis. + """ + from jaccpot._env import env_flag + + if not env_flag("JACCPOT_STATIC_STRICT_FUSED_M2L_CSR", False): + return False + if env_flag("JACCPOT_M2L_CSR_INTERPRET", False): + return True + from jaccpot.pallas.m2l_real_csr import pallas_m2l_real_csr_supported + + return pallas_m2l_real_csr_supported() + + class _FarPairCOO(NamedTuple): """Compact COO-style far-pair representation for streamed M2L execution. @@ -520,7 +544,10 @@ 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. + 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 + target-tiled CSR Pallas kernel, which owns one local row per program and so + needs neither the chunked scan nor its per-chunk scatter. Parameters ---------- @@ -631,7 +658,20 @@ def _solidfmm_downward_accumulate_from_multipoles( basis_mode=basis_mode, ) else: - if pair_count <= chunk_size: + if real_basis and _m2l_csr_pallas_active(): + from jaccpot._env import env_flag + from jaccpot.pallas.m2l_real_csr import m2l_real_csr_pallas + + locals_updated = initial_locals_coeffs + m2l_real_csr_pallas( + multipoles_coeffs, + centers, + src, + tgt, + order=order, + active_pair_count=active_pair_count, + interpret=env_flag("JACCPOT_M2L_CSR_INTERPRET", False), + ) + elif pair_count <= chunk_size: locals_updated = _accumulate_m2l_fullbatch( initial_locals_coeffs, multipoles_coeffs, diff --git a/tests/unit/runtime/test_m2l_csr_lane_wiring.py b/tests/unit/runtime/test_m2l_csr_lane_wiring.py new file mode 100644 index 00000000..6f180eb0 --- /dev/null +++ b/tests/unit/runtime/test_m2l_csr_lane_wiring.py @@ -0,0 +1,91 @@ +"""The opt-in CSR M2L lane is (a) actually taken and (b) force-neutral. + +``JACCPOT_STATIC_STRICT_FUSED_M2L_CSR=1`` 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 +reached would pass any parity check, so the kernel entry is counted. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +import jaccpot.pallas.m2l_real_csr as csr_mod +from jaccpot.runtime.kernels._downward_prep import _m2l_csr_pallas_active + + +def test_flag_off_by_default(monkeypatch): + monkeypatch.delenv("JACCPOT_STATIC_STRICT_FUSED_M2L_CSR", raising=False) + assert _m2l_csr_pallas_active() is False + + +def test_flag_on_cpu_needs_interpret(monkeypatch): + monkeypatch.setenv("JACCPOT_STATIC_STRICT_FUSED_M2L_CSR", "1") + monkeypatch.delenv("JACCPOT_M2L_CSR_INTERPRET", raising=False) + import jax + + if jax.default_backend() != "gpu": + assert _m2l_csr_pallas_active() is False + monkeypatch.setenv("JACCPOT_M2L_CSR_INTERPRET", "1") + assert _m2l_csr_pallas_active() is True + + +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 test_csr_lane_is_taken_and_matches_the_chunked_lane(monkeypatch): + import jax.numpy as jnp + + from jaccpot import ( + FarFieldConfig, + FastMultipoleMethod, + FMMAdvancedConfig, + NearFieldConfig, + TreeConfig, + ) + + pos, mass = _plummer(3000) + + def solve(): + # the default (non-strict) runtime: CPU-friendly, real basis, same downward + solver = FastMultipoleMethod( + basis="real", theta=0.6, + G=1.0, softening=1e-3, working_dtype=jnp.float32, + advanced=FMMAdvancedConfig( + tree=TreeConfig(mode="static_radix", leaf_target=32), + farfield=FarFieldConfig(mode="auto"), nearfield=NearFieldConfig(mode="auto"), + mac_type="dehnen"), + fixed_order=3) + acc = solver.compute_accelerations( + jnp.asarray(pos), jnp.asarray(mass), leaf_size=32, max_order=3, theta=0.6 + ) + return np.asarray(acc, np.float64) + + monkeypatch.delenv("JACCPOT_STATIC_STRICT_FUSED_M2L_CSR", raising=False) + a_ref = solve() + + calls = {"n": 0} + real_kernel = csr_mod.m2l_real_csr_pallas + + def counting(*a, **k): + calls["n"] += 1 + return real_kernel(*a, **k) + + monkeypatch.setattr(csr_mod, "m2l_real_csr_pallas", counting) + monkeypatch.setenv("JACCPOT_STATIC_STRICT_FUSED_M2L_CSR", "1") + monkeypatch.setenv("JACCPOT_M2L_CSR_INTERPRET", "1") + a_csr = solve() + assert calls["n"] >= 1, "the CSR lane was never entered" + rel = np.linalg.norm(a_csr - a_ref) / np.linalg.norm(a_ref) + assert np.all(np.isfinite(a_csr)) + assert rel < 2e-5, rel From 0f6aeeb40f5c17b9131f07c31043c052f24d8e55 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Wed, 9 Sep 2026 18:05:24 +0200 Subject: [PATCH 06/12] perf(pallas): CSR M2L in the centred layout -- no per-pair pack/unpack, degree-only z-core The first cut held two 128x32 one-hot pack/unpack matrices and a Cp x Cp z-core in registers per program and ran at ~93 ns/pair on a shared A100 (non-record), no faster than the pure-JAX lane. The multipole table is now packed into the centred (Bp, Wp) layout once in the wrapper and the output unpacked once; the z-translation preserves m, so its operator is a Bp x Bp degree matrix per column. Half the per-pair MACs, a third of the constants; 4 warps per program. Microbench: iterate the well-separation rejection (a coincident pair made the reference NaN), assert finiteness. CSR sm_80 tests registered in the GPU gate. Co-Authored-By: Claude Fable 5.1 --- bench/gpu_gate.py | 5 + bench/m2l_csr_microbench.py | 12 +- jaccpot/pallas/m2l_real_csr.py | 243 ++++++++++++++++++++------------- 3 files changed, 162 insertions(+), 98 deletions(-) diff --git a/bench/gpu_gate.py b/bench/gpu_gate.py index c101fc53..5f23aa9b 100644 --- a/bench/gpu_gate.py +++ b/bench/gpu_gate.py @@ -95,6 +95,11 @@ # (plan "small leaves", Phase 1); its Triton lowering is what this checks. "test_leafpair_include_self_gpu_matches_reference[None]", "test_leafpair_include_self_gpu_matches_reference[3]", + # The target-tiled CSR M2L kernel (plan "small leaves", Phase 2): dynamic-trip + # loop, row gathers by id, atan2 and pow2 tiles all lower only on sm_80. + "test_csr_pallas_gpu_matches_rot_scale[2]", + "test_csr_pallas_gpu_matches_rot_scale[4]", + "test_csr_pallas_gpu_matches_rot_scale[6]", ) # Measured on an A100 sm_80 / jax 0.10.2 and documented in ARCHITECTURE.md §9, diff --git a/bench/m2l_csr_microbench.py b/bench/m2l_csr_microbench.py index 29503711..1b256b51 100644 --- a/bench/m2l_csr_microbench.py +++ b/bench/m2l_csr_microbench.py @@ -69,9 +69,13 @@ def timed(fn, *a): # random well-separated pairs: reject |delta| < 0.3 (the MAC would too) src = rng.integers(0, n, P).astype(np.int32) tgt = rng.integers(0, n, P).astype(np.int32) - d = np.asarray(centers)[tgt] - np.asarray(centers)[src] - bad = np.linalg.norm(d, axis=1) < 0.3 - src = np.where(bad, (src + 1) % n, src).astype(np.int32) + cn = np.asarray(centers) + for _ in range(50): # resample until every pair is well separated + bad = np.linalg.norm(cn[tgt] - cn[src], axis=1) < 0.3 + if not bad.any(): + break + src = np.where(bad, rng.integers(0, n, P), src).astype(np.int32) + assert not bad.any() src_j, tgt_j = jnp.asarray(src), jnp.asarray(tgt) counts = np.bincount(tgt, minlength=n) @@ -103,6 +107,8 @@ def body(acc, i): rows[db] = (out_p, t_p, med_p) del fn ref = np.asarray(rows["0"][0], np.float64) + assert np.all(np.isfinite(ref)), "pure-JAX reference has non-finite rows" + assert np.all(np.isfinite(np.asarray(out_k))), "CSR kernel produced non-finite rows" rel = float(np.linalg.norm(np.asarray(out_k, np.float64) - ref) / np.linalg.norm(ref)) rel_db = float(np.linalg.norm(np.asarray(rows["1"][0], np.float64) - ref) / np.linalg.norm(ref)) row = dict(order=order, pairs=P, nodes=n, longest_row=int(counts.max()), diff --git a/jaccpot/pallas/m2l_real_csr.py b/jaccpot/pallas/m2l_real_csr.py index e2c20182..ea9247ca 100644 --- a/jaccpot/pallas/m2l_real_csr.py +++ b/jaccpot/pallas/m2l_real_csr.py @@ -34,9 +34,13 @@ (:func:`jaccpot.operators.real_rotations.real_rotation_from_z_axis_local`): ``Dz(-az) B_l^T Dz(ax) B_l^T``. -Every contraction is a broadcast-multiply + ``jnp.sum`` (no ``dot``, no TF32), -as in :mod:`jaccpot.pallas.m2l_real_fused`. Padded lanes of every constant are -exactly zero, so they are inert in every reduction. +The kernel works entirely in the centred ``(Bp, Wp)`` layout: the wrapper packs +the multipole table into it once (an XLA gather over ``n`` rows) and unpacks the +output once, so no per-pair pack/unpack matvec and no ``Cp x Cp`` table exist in +the kernel -- the z-core preserves ``m`` and is a ``Bp x Bp`` degree operator +per column. Every contraction is a broadcast-multiply + ``jnp.sum`` (no ``dot``, +no TF32), as in :mod:`jaccpot.pallas.m2l_real_fused`. Padded lanes of every +constant are exactly zero, so they are inert in every reduction. Forward only (the fused strict lane is forward-only); the pure-JAX lane stays the differentiable path, and the transverse-degeneracy JVP treatment of @@ -75,18 +79,19 @@ "m2l_real_csr_jax", "m2l_real_csr_pallas", "csr_by_target", + "pack_centred", + "unpack_centred", ] _TABLE_KEYS = ( - "Ppack", - "Uunpack", "Bstack", "BstackT", "Apat", "mabs", - "Zsf", - "DegOut", - "DegSrc", + "signm", + "Zf", + "degn", + "degk", ) @@ -110,6 +115,10 @@ def _next_pow2(n: int) -> int: def m2l_real_csr_tables(order: int) -> dict: """Compile-time constants of the kernel for one expansion order. + Everything lives in the CENTRED ``(Bp, Wp)`` layout: row = degree ``l``, + column ``p + m``; degrees ``> p`` and columns with ``|m| > p`` are padding + and every constant is exactly zero there. + Parameters ---------- order : int @@ -119,33 +128,33 @@ def m2l_real_csr_tables(order: int) -> dict: ------- dict NumPy float64 arrays (cast to the working dtype by the caller) plus the - shape scalars: ``C = (p+1)^2``, ``Cp`` (pow2 >= C), ``W = 2p+1``, - ``Wp`` (pow2 >= W), ``Bp`` (pow2 >= p+1). + shape scalars ``C = (p+1)^2``, ``W = 2p+1``, ``Wp`` (pow2 >= W), ``Bp`` + (pow2 >= p+1), and the pack/unpack maps ``idx [Bp, Wp]`` (packed + coefficient index of each centred slot) and ``mask [Bp, Wp]``. - ``Ppack [Bp*Wp, Cp]`` / ``Uunpack [Cp, Bp*Wp]``: one-hot pack/unpack - between the packed coefficient vector and the centred ``(Bp, Wp)`` rows. ``Bstack [Bp, Wp, Wp]``: ``B_U(l)`` centred per degree; ``BstackT`` its per-degree transpose. ``Apat [Wp, Wp]``: the Dz sine pattern. - ``mabs [Wp]``: ``|m|`` per column (0 on padded columns). - ``Zsf [Cp, Cp]``: ``sign(m) (n+k)!`` on the valid (out, src) entries. - ``DegOut [Cp]`` / ``DegSrc [Cp]``: the radius exponents ``n+1`` per - output lane and ``k`` per source lane (0 on padded lanes, where ``Zsf`` - is zero anyway). + ``mabs [Wp]``: ``|m|`` per column. ``signm [Wp]``: the z-core's + ``sign(m) = (-1)^m (2 if m != 0 else 1)`` per column. + ``Zf [Bp, Bp]``: ``(n+k)!`` where ``k <= p - n``, else 0 -- the z-core + preserves ``m``, so it is a degree x degree operator per column. + ``degn [Bp]``: ``n + 1`` (radius exponent of the output degree); + ``degk [Bp]``: ``k`` (radius exponent of the source degree). """ p = int(order) if p < 0: raise ValueError("order must be >= 0") C = sh_size(p) W = 2 * p + 1 - Cp = _next_pow2(C) Wp = _next_pow2(W) Bp = _next_pow2(p + 1) - Ppack = np.zeros((Bp * Wp, Cp), dtype=np.float64) + idx = np.zeros((Bp, Wp), dtype=np.int32) + mask = np.zeros((Bp, Wp), dtype=bool) for ell in range(p + 1): for m in range(-ell, ell + 1): - Ppack[ell * Wp + (p + m), sh_offset(ell) + ell + m] = 1.0 - Uunpack = Ppack.T.copy() + idx[ell, p + m] = sh_offset(ell) + ell + m + mask[ell, p + m] = True Bstack = np.zeros((Bp, Wp, Wp), dtype=np.float64) # `compute_real_B_matrix_multipole` is jitted: evaluated eagerly here so the @@ -160,33 +169,41 @@ def m2l_real_csr_tables(order: int) -> dict: Apat = np.zeros((Wp, Wp), dtype=np.float64) mabs = np.zeros((Wp,), dtype=np.float64) + signm = np.zeros((Wp,), dtype=np.float64) + signm[p] = 1.0 for m in range(1, p + 1): Apat[p + m, p - m] = -1.0 Apat[p - m, p + m] = 1.0 mabs[p + m] = float(m) mabs[p - m] = float(m) + signm[p + m] = signm[p - m] = (-1.0 if (m % 2) else 1.0) * 2.0 + # z-core in the centred layout, cross-checked against the single source of truth src_index, valid, fact_index, r_exponent, sign = z_m2l_translation_tables(p) fact = np.asarray([math.factorial(i) for i in range(2 * p + 1)], dtype=np.float64) - Zsf = np.zeros((Cp, Cp), dtype=np.float64) - deg_of = np.zeros((C,), dtype=np.int64) + Zf = np.zeros((Bp, Bp), dtype=np.float64) + for n in range(p + 1): + for k in range(p - n + 1): + Zf[n, k] = fact[n + k] for n in range(p + 1): - deg_of[sh_offset(n) : sh_offset(n + 1)] = n - for out in range(C): - for k in range(p + 1): - if bool(valid[out, k]): - s = int(src_index[out, k]) - Zsf[out, s] = float(sign[out]) * float(fact[int(fact_index[out, k])]) - # r^-(n+k+1) = rinv^(n+1) * rinv^k with n = deg(out), k = deg(s) - assert int(r_exponent[out, k]) == deg_of[out] + 1 + deg_of[s] - DegOut = np.zeros((Cp,), dtype=np.float64) - DegSrc = np.zeros((Cp,), dtype=np.float64) - DegOut[:C] = deg_of + 1 - DegSrc[:C] = deg_of + for m in range(-n, n + 1): + out = sh_offset(n) + n + m + assert abs(sign[out] - signm[p + m]) < 1e-12 + for k in range(p + 1): + if bool(valid[out, k]): + assert int(src_index[out, k]) == sh_offset(k) + k + m # same m + assert int(r_exponent[out, k]) == n + k + 1 + assert fact[int(fact_index[out, k])] == Zf[n, k] + else: + assert k < abs(m) or k > p - n + degn = np.zeros((Bp,), dtype=np.float64) + degk = np.zeros((Bp,), dtype=np.float64) + degn[: p + 1] = np.arange(p + 1) + 1 + degk[: p + 1] = np.arange(p + 1) return dict( - p=p, C=C, Cp=Cp, W=W, Wp=Wp, Bp=Bp, - Ppack=Ppack, Uunpack=Uunpack, Bstack=Bstack, BstackT=BstackT, - Apat=Apat, mabs=mabs, Zsf=Zsf, DegOut=DegOut, DegSrc=DegSrc, + p=p, C=C, W=W, Wp=Wp, Bp=Bp, idx=idx, mask=mask, + Bstack=Bstack, BstackT=BstackT, Apat=Apat, mabs=mabs, signm=signm, + Zf=Zf, degn=degn, degk=degk, ) @@ -195,15 +212,56 @@ def _tables_to_jnp(order: int, dtype: Any) -> dict[str, Array]: return {k: jnp.asarray(t[k], dtype=dtype) for k in _TABLE_KEYS} +def pack_centred(coeffs: Array, *, order: int) -> Array: + """``[N, C]`` packed coefficients -> ``[N, Bp*Wp]`` centred rows (zeros on padding). + + Parameters + ---------- + coeffs : Array + Packed coefficients, ``[N, (p+1)^2]``. + order : int + Expansion order. Static. + + Returns + ------- + Array + ``[N, Bp*Wp]``. + """ + t = m2l_real_csr_tables(int(order)) + idx = jnp.asarray(t["idx"]) + mask = jnp.asarray(t["mask"]) + rows = jnp.where(mask[None], coeffs[:, idx], jnp.zeros((), coeffs.dtype)) + return rows.reshape(coeffs.shape[0], t["Bp"] * t["Wp"]) + + +def unpack_centred(rows: Array, *, order: int) -> Array: + """Inverse of :func:`pack_centred`: ``[N, Bp*Wp]`` -> ``[N, C]``. + + Parameters + ---------- + rows : Array + Centred rows, ``[N, Bp*Wp]``. + order : int + Expansion order. Static. + + Returns + ------- + Array + ``[N, (p+1)^2]`` packed coefficients. + """ + t = m2l_real_csr_tables(int(order)) + mask = np.asarray(t["mask"]) + flat_slots = np.nonzero(mask.reshape(-1))[0] + packed_idx = np.asarray(t["idx"]).reshape(-1)[flat_slots] + out = jnp.zeros((rows.shape[0], t["C"]), dtype=rows.dtype) + return out.at[:, packed_idx].set(rows[:, flat_slots]) + + # --------------------------------------------------------------------------- math # Every helper below is written for BOTH the Pallas kernel (values loaded from # refs) and the pure-jnp twin: plain broadcast-multiply + sum, no dot, no gather. -def _matvec(mat: Array, vec: Array) -> Array: - return jnp.sum(mat * vec[None, :], axis=1) - - def _bapply(bstack: Array, rows: Array) -> Array: """``out[l, i] = sum_j bstack[l, i, j] rows[l, j]`` (block-diagonal by degree).""" return jnp.sum(bstack * rows[:, None, :], axis=-1) @@ -215,36 +273,32 @@ def _dz(rows: Array, cosv: Array, sinv: Array, apat: Array) -> Array: return rows * cosv[None, :] + jnp.sum(apat[None, :, :] * sv[:, None, :], axis=-1) -def _m2l_pair(mult: Array, delta3: tuple, t: dict[str, Array], *, bp: int, wp: int) -> Array: - """Full real M2L for one pair from the packed row ``mult`` and ``delta = c_t - c_s``. +def _m2l_pair_rows(rows: Array, delta3: tuple, t: dict[str, Array]) -> Array: + """Full real M2L for one pair in the centred layout. Parameters ---------- - mult : Array - Padded source multipole row ``(Cp,)``. + rows : Array + Source multipole in centred rows, ``(Bp, Wp)``. delta3 : tuple ``(x, y, z)`` scalars, target centre minus source centre. t : dict[str, Array] Tables from :func:`_tables_to_jnp` at the working dtype. - bp : int - Padded degree count ``Bp``. Static. - wp : int - Padded row width ``Wp``. Static. Returns ------- Array - Padded local contribution ``(Cp,)``. + Local contribution in centred rows, ``(Bp, Wp)``. """ x, y, z = delta3 - dtype = mult.dtype + dtype = rows.dtype rho2 = x * x + y * y rho = jnp.sqrt(rho2) az = jnp.arctan2(x, y) ax = jnp.arctan2(rho, z) r = jnp.sqrt(rho2 + z * z) r = jnp.maximum(r, jnp.asarray(1.0e-30, dtype=dtype)) - rinv = 1.0 / r + log_rinv = -jnp.log(r) mabs = t["mabs"] cos_az = jnp.cos(mabs * az) sin_az = jnp.sin(mabs * az) @@ -253,26 +307,23 @@ def _m2l_pair(mult: Array, delta3: tuple, t: dict[str, Array], *, bp: int, wp: i apat = t["Apat"] # world -> z (multipole): B Dz(-ax) B Dz(az), applied right to left - rows = _matvec(t["Ppack"], mult).reshape(bp, wp) - rows = _dz(rows, cos_az, sin_az, apat) - rows = _bapply(t["Bstack"], rows) - rows = _dz(rows, cos_ax, -sin_ax, apat) - rows = _bapply(t["Bstack"], rows) - mrf = _matvec(t["Uunpack"], rows.reshape(bp * wp)) - - # z-core, separable radius powers: r^-(n+k+1) = rinv^(n+1) * rinv^k - log_rinv = jnp.log(rinv) - rinv_out = jnp.exp(t["DegOut"] * log_rinv) - rinv_src = jnp.exp(t["DegSrc"] * log_rinv) - lz = rinv_out * _matvec(t["Zsf"], rinv_src * mrf) + v = _dz(rows, cos_az, sin_az, apat) + v = _bapply(t["Bstack"], v) + v = _dz(v, cos_ax, -sin_ax, apat) + v = _bapply(t["Bstack"], v) + + # z-core: same m, degree x degree; r^-(n+k+1) = rinv^(n+1) rinv^k + rinv_n = jnp.exp(t["degn"] * log_rinv) # (Bp,) + rinv_k = jnp.exp(t["degk"] * log_rinv) # (Bp,) + v = v * rinv_k[:, None] + v = jnp.sum(t["Zf"][:, :, None] * v[None, :, :], axis=1) # (Bp, Wp) + v = v * rinv_n[:, None] * t["signm"][None, :] # z -> world (local) = transpose of the multipole block: Dz(-az) B^T Dz(ax) B^T - rows = _matvec(t["Ppack"], lz).reshape(bp, wp) - rows = _bapply(t["BstackT"], rows) - rows = _dz(rows, cos_ax, sin_ax, apat) - rows = _bapply(t["BstackT"], rows) - rows = _dz(rows, cos_az, -sin_az, apat) - return _matvec(t["Uunpack"], rows.reshape(bp * wp)) + v = _bapply(t["BstackT"], v) + v = _dz(v, cos_ax, sin_ax, apat) + v = _bapply(t["BstackT"], v) + return _dz(v, cos_az, -sin_az, apat) # ----------------------------------------------------------------- pure-jnp twin @@ -296,17 +347,18 @@ def m2l_real_csr_pair_jax(multipoles: Array, deltas: Array, *, order: int) -> Ar ``[N, C]`` local contributions, one per pair (NOT reduced by target). """ tb = m2l_real_csr_tables(int(order)) - C, Cp, Bp, Wp = tb["C"], tb["Cp"], tb["Bp"], tb["Wp"] + Bp, Wp = tb["Bp"], tb["Wp"] mult = jnp.asarray(multipoles) dtype = mult.dtype t = _tables_to_jnp(int(order), dtype) - mult_p = jnp.pad(mult, ((0, 0), (0, Cp - C))) + rows = pack_centred(mult, order=int(order)).reshape(-1, Bp, Wp) d = jnp.asarray(deltas, dtype=dtype) - def one(m, dd): - return _m2l_pair(m, (dd[0], dd[1], dd[2]), t, bp=Bp, wp=Wp) + def one(rw, dd): + return _m2l_pair_rows(rw, (dd[0], dd[1], dd[2]), t) - return jax.vmap(one)(mult_p, d)[:, :C] + out_rows = jax.vmap(one)(rows, d).reshape(-1, Bp * Wp) + return unpack_centred(out_rows, order=int(order)) def csr_by_target( @@ -416,7 +468,7 @@ def _m2l_real_csr_kernel( Parameters ---------- mult_ref : KernelRef - Whole padded multipole table ``[n, Cp]`` (gathered by source id). + Whole centred multipole table ``[n, Bp*Wp]`` (gathered by source id). cent_ref : KernelRef Whole padded centre table ``[n, 4]``. src_ref : KernelRef @@ -427,7 +479,7 @@ def _m2l_real_csr_kernel( Segment length per target ``[n]``. *table_and_out_refs : KernelRef The ``_TABLE_KEYS`` constants (whole arrays) followed by the output ref - ``[1, Cp]``. + ``[1, Bp*Wp]``. bp : int ``Bp``. Static. wp : int @@ -436,7 +488,7 @@ def _m2l_real_csr_kernel( Returns ------- None - Writes the target's local row. + Writes the target's centred local row. """ table_refs = table_and_out_refs[: len(_TABLE_KEYS)] (out_ref,) = table_and_out_refs[len(_TABLE_KEYS) :] @@ -447,19 +499,18 @@ def _m2l_real_csr_kernel( ctx = cent_ref[tgt, 0] cty = cent_ref[tgt, 1] ctz = cent_ref[tgt, 2] - cp = int(out_ref.shape[1]) - acc0 = jnp.zeros((cp,), dtype=out_ref.dtype) + acc0 = jnp.zeros((bp, wp), dtype=out_ref.dtype) def body(k, acc): sid = src_ref[start + k] - m = mult_ref[sid, :] + rows = mult_ref[sid, :].reshape(bp, wp) dx = ctx - cent_ref[sid, 0] dy = cty - cent_ref[sid, 1] dz_ = ctz - cent_ref[sid, 2] - return acc + _m2l_pair(m, (dx, dy, dz_), t, bp=bp, wp=wp) + return acc + _m2l_pair_rows(rows, (dx, dy, dz_), t) acc = lax.fori_loop(0, cnt, body, acc0) - out_ref[0, :] = acc + out_ref[0, :] = acc.reshape(bp * wp) def m2l_real_csr_pallas( @@ -472,7 +523,7 @@ def m2l_real_csr_pallas( active_pair_count: Optional[Array] = None, interpret: bool = False, backend: str = "triton", - num_warps: int = 1, + num_warps: int = 4, ) -> Array: """Local coefficient increments from a flat far-pair list, one Pallas program per target. @@ -496,7 +547,10 @@ def m2l_real_csr_pallas( backend : str Pallas GPU lowering, ``"triton"`` by default. num_warps : int - Warps per program; the row width is ``Cp`` lanes, one warp suffices. + Warps per program. The working tile is ``(Bp, Wp)`` = 128 elements at + p <= 7 and the two rotation stacks are 2 x ``Bp*Wp*Wp`` constants held in + registers, so 4 warps (128 threads) keeps the per-thread register count + low; 1 warp spilled. Returns ------- @@ -504,7 +558,7 @@ def m2l_real_csr_pallas( ``[n, C]`` local increments, same dtype as ``multipoles``. """ tb = m2l_real_csr_tables(int(order)) - C, Cp, Bp, Wp = tb["C"], tb["Cp"], tb["Bp"], tb["Wp"] + C, Bp, Wp = tb["C"], tb["Bp"], tb["Wp"] mult = jnp.asarray(multipoles) dtype = mult.dtype n = int(mult.shape[0]) @@ -513,7 +567,7 @@ def m2l_real_csr_pallas( cent = jnp.asarray(centers, dtype=dtype) if cent.ndim != 2 or int(cent.shape[1]) != 3 or int(cent.shape[0]) != n: raise ValueError("centers must have shape (n, 3) aligned with multipoles") - mult_p = jnp.pad(mult, ((0, 0), (0, Cp - C))) + mult_c = pack_centred(mult, order=int(order)) # [n, Bp*Wp] cent_p = jnp.pad(cent, ((0, 0), (0, 1))) src_sorted, offsets, counts = csr_by_target( sources, targets, total_nodes=n, active_pair_count=active_pair_count @@ -531,25 +585,24 @@ def bs_full(arr: Array) -> pl.BlockSpec: kernel = functools.partial(_m2l_real_csr_kernel, bp=Bp, wp=Wp) backend_kwargs = pallas_backend_kwargs(backend, interpret) if "compiler_params" in backend_kwargs: - # one Cp-lane row per program: a single warp is the right launch shape backend_kwargs["compiler_params"] = type(backend_kwargs["compiler_params"])( num_warps=int(num_warps) ) - out = pl.pallas_call( + out_rows = pl.pallas_call( kernel, grid=(n,), in_specs=[ - bs_full(mult_p), + bs_full(mult_c), bs_full(cent_p), bs_full(src_sorted), bs_full(offsets), bs_full(counts), *[bs_full(a) for a in table_arrays], ], - out_specs=pl.BlockSpec((1, Cp), lambda i: (i, 0)), - out_shape=jax.ShapeDtypeStruct((n, Cp), dtype), + out_specs=pl.BlockSpec((1, Bp * Wp), lambda i: (i, 0)), + out_shape=jax.ShapeDtypeStruct((n, Bp * Wp), dtype), interpret=bool(interpret), **backend_kwargs, name=f"m2l_real_csr_p{int(order)}", - )(mult_p, cent_p, src_sorted, offsets, counts, *table_arrays) - return out[:, :C] + )(mult_c, cent_p, src_sorted, offsets, counts, *table_arrays) + return unpack_centred(out_rows, order=int(order)) From 76bcd5214f03e6ccaa51f95ab321b5203e7a4c9e Mon Sep 17 00:00:00 2001 From: TobiBu Date: Wed, 9 Sep 2026 18:06:56 +0200 Subject: [PATCH 07/12] test(m2l-csr): axis contracts + pack/unpack round trip; wiring test on the slow list (28 s on CPU) Co-Authored-By: Claude Fable 5.1 --- tests/slow_tests.txt | 1 + .../operators/test_m2l_real_csr_pallas.py | 28 +++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/tests/slow_tests.txt b/tests/slow_tests.txt index c47a8af0..4772e928 100644 --- a/tests/slow_tests.txt +++ b/tests/slow_tests.txt @@ -85,6 +85,7 @@ tests/unit/runtime/test_large_n_grad_reverse_path.py::test_reverse_pass_is_finit tests/unit/runtime/test_large_n_grad_reverse_path.py::test_the_far_field_is_load_bearing tests/unit/runtime/test_large_n_grad_reverse_path.py::test_the_lane_under_test_is_actually_the_large_n_lane tests/unit/runtime/test_large_n_grad_reverse_path.py::test_the_pairs_layout_is_rejected_rather_than_silently_dropped +tests/unit/runtime/test_m2l_csr_lane_wiring.py::test_csr_lane_is_taken_and_matches_the_chunked_lane tests/unit/test_large_n_config_thresholds.py::test_fp32_farfield_gradient_is_finite tests/unit/test_large_n_config_thresholds.py::test_grad_works_with_grouped_interactions_requested tests/unit/test_large_n_fast_path_policy.py::test_large_n_accel_eval_requires_fast_lane_state diff --git a/tests/unit/operators/test_m2l_real_csr_pallas.py b/tests/unit/operators/test_m2l_real_csr_pallas.py index ac749857..77d5a225 100644 --- a/tests/unit/operators/test_m2l_real_csr_pallas.py +++ b/tests/unit/operators/test_m2l_real_csr_pallas.py @@ -188,3 +188,31 @@ def test_csr_pallas_gpu_matches_rot_scale(order): ) assert np.all(np.isfinite(np.asarray(got))) assert _relerr(got, ref) < 3e-4 + + +# ---------------------------------------------------------------- axis contracts + + +def test_csr_pallas_rejects_a_coefficient_count_of_another_order(): + mult, centers, src, tgt = _case(3, np.float32, seed=30) + with pytest.raises(ValueError, match="coefficients"): + m2l_real_csr_pallas(jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), + jnp.asarray(tgt), order=4, interpret=True) + + +def test_csr_pallas_rejects_misaligned_centers(): + mult, centers, src, tgt = _case(3, np.float32, seed=31) + with pytest.raises(ValueError, match="centers"): + m2l_real_csr_pallas(jnp.asarray(mult), jnp.asarray(centers[:-1]), jnp.asarray(src), + jnp.asarray(tgt), order=3, interpret=True) + + +def test_pack_unpack_centred_round_trip(): + from jaccpot.pallas.m2l_real_csr import pack_centred, unpack_centred + + for order in (2, 4, 6): + c = sh_size(order) + x = jnp.asarray(np.random.default_rng(order).standard_normal((5, c)).astype(np.float32)) + rows = pack_centred(x, order=order) + assert rows.shape[1] & (rows.shape[1] - 1) == 0 # pow2 row width + np.testing.assert_array_equal(np.asarray(unpack_centred(rows, order=order)), np.asarray(x)) From f00948c9f55ac49fb5d740eb412f8ec219cbe402 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Wed, 9 Sep 2026 18:08:01 +0200 Subject: [PATCH 08/12] bench(m2l-csr): clear jit caches between the degree-batched variants (the knob is trace-time) Co-Authored-By: Claude Fable 5.1 --- bench/m2l_csr_microbench.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/bench/m2l_csr_microbench.py b/bench/m2l_csr_microbench.py index 1b256b51..27638d46 100644 --- a/bench/m2l_csr_microbench.py +++ b/bench/m2l_csr_microbench.py @@ -102,6 +102,10 @@ def body(acc, i): rows = {} for db in ("0", "1"): os.environ["JACCPOT_M2L_DEGREE_BATCHED"] = db + # the knob is read at trace time and the cascade's inner jits are + # cached across variants -- without this the second variant reuses + # the first's trace (measured: bit-identical output and time) + jax.clear_caches() fn = jax.jit(pure) out_p, t_p, med_p = timed(fn, mult, centers, src_p, tgt_p) rows[db] = (out_p, t_p, med_p) From 28ef6b173844e975d51cd119354efd244bc176b3 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Thu, 10 Sep 2026 00:29:11 +0200 Subject: [PATCH 09/12] docs: small-leaves record -- attributed baseline, the three kernels/fixes, per-step table, capacity rules Co-Authored-By: Claude Fable 5.1 --- docs/small_leaves_2026-09.md | 91 ++++++++++++++++++++++++++++++++++++ 1 file changed, 91 insertions(+) create mode 100644 docs/small_leaves_2026-09.md diff --git a/docs/small_leaves_2026-09.md b/docs/small_leaves_2026-09.md new file mode 100644 index 00000000..03052846 --- /dev/null +++ b/docs/small_leaves_2026-09.md @@ -0,0 +1,91 @@ +# Small leaves on the fused single-GPU lane (2026-09-07 .. 2026-09-10) + +Plan: `~/.claude/plans/ok-please-draft-a-effervescent-crown.md` (Odisseo box). Branch +`perf/small-leaves-p1`. Harness: `Odisseo-bench-multigpu/benchmark_multigpu/codes/{fit_smallleaf_caps, +smallleaf_baseline,smallleaf_dynamics}.py`, `FAST_LANE_ENV_BY_LEAF` in `compare_force.py`; results under +`artifacts/smallleaf/`. + +## Why + +At N=200k Plummer, p=4, theta 0.6 the fused lane at leaf 256 sums 58 % of N directly per target +(pkdgrav3: 0.24 %). The direct share falls as ~W^0.9 with the leaf size, so smaller leaves are the +only route to pkdgrav3's arithmetic -- and before this work leaf 64 was 3.1x SLOWER per step than +leaf 256 with 3x fewer pair evaluations. + +## What the attributed baseline said + +`strict_run_v2` per step, N=200k, theta 0.6, idle A100 (`baseline_base_*.json`): + +| leaf | step ms | eval-only ms | leaf-pair kernel | downward (M2L+L2L) | launches/step | +|---|---|---|---|---|---| +| 256 | 131.8 | 70.0 | 79 | 41 | 7.8k | +| 128 | 187.5 | 43.2 | 48 | 130 | 15.8k | +| 64 | 409.6 | 26.5 | 22 | 372 | 32k | +| 32 | 1168 | 19.6 | 13 | ~1100 | -- | + +The Pallas leaf-pair kernel has **no small-W cliff** (its time is exactly the direct-volume ratio), +the self-leaf scan was already negligible after #334, and the whole penalty was the far field: +one XLA scatter fusion launched once per 4096-pair M2L chunk at 686 us (`_chunk_segment_scatter_add`: +`segment_sum` with hundreds of duplicates per address plus ~3800 zero-adds onto node 0 per chunk, +both serialised by the atomic-add lowering) = 169 ms of the 410 ms leaf-64 step, and behind it a +21k-launch/step storm from the chunked rotation cascade. + +## What was built + +1. **Self-leaf block folded into the leaf-pair Pallas kernel** (`nearfield_fused_leaf.py`, + `include_self`; flag `JACCPOT_NEARFIELD_LEAFPAIR_FOLD_SELF`, default on, forward prepacked lane + only). Time-neutral (the scan was already batched), removes the per-leaf launch family. +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 + 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`): + + | pairs | p | CSR ns/pair | pure-JAX | degree-batched | + |---|---|---|---|---| + | 1M | 4 | 19.8 | 69.4 | 55.6 | + | 6M | 4 | 12.1 | 68.9 | 55.3 | + | 1M | 6 | 12.4 | 126.1 | 86.2 | + | 6M | 6 | 12.0 | 124.9 | 85.6 | + + Parity 1e-6 (fp32) against `m2l_rot_scale_real_batch`; interpret parity < 1e-10 (fp64) for orders + 2-6. In the step at leaf 64 the kernel costs 12.9 ms where the pure-JAX lane cost ~360. + +## Where it landed (per step, theta 0.6, N=200k, idle A100, foreign-process-free rows only) + +| leaf | baseline | fold + scatter | + CSR M2L | eval-only | aggL2 vs fp64 direct | +|---|---|---|---|---|---| +| 256 | 131.8 | 126.4 | **120.3** | 69 | 8.27e-4 | +| 128 | 187.5 | 145.2 | **122.3** | 42 | 1.15e-3 | +| 64 | 409.6 | 236.5 | **176.7** | 24 | 1.20e-3 | +| 32 | 1168 | 649.0 | **509.4** | 15 | 1.30e-3 | + +Theta 0.8 with CSR: leaf 64 113.5 (baseline 192.0), leaf 32 364.1 (baseline OOM'd on a shared card). + +**The far field is solved; the per-step leaf optimum did not move.** With M2L at 13 ms, the leaf-64 +step is ~100 ms of traced dual-tree walk and list compaction (pair-queue-sized scatter fusions at +~800 us x ~55 launches, ~200 memcpys, sorts) that scale with the leaf count and the traced +`max_pair_queue` (1M at leaf 64, 2M at leaf 32). Eval-only (force at fixed lists) is 3-5x faster at +leaf 32-64 than at 256, so a workload that refreshes lists every k steps, or a walk that costs +O(leaves) rather than O(queue), would collect the gain. Gate G3 (<= 35 ms per step at 200k) is +missed; G1, G2 and G2a are met. Next lever: the traced walk's queue buffers (yggdrax side). + +## Capacity rules for small leaves (fitted, `FAST_LANE_ENV_BY_LEAF`) + +* compact far-pair cap: pow2 >= 1.5x the far-pair count (101k / 363k / 992k / 2.34M at 256/128/64/32). +* neighbour-EDGE cap: the traced refresh pads to `num_leaves x traced_neighbour_cap`, and the #333 + carry-over sets that cap to pow2(1.5 x longest eager row + 1) with the longest row = + `num_leaves - 1` -> 2^21 already fails at leaf 128 (2^23), 2^25 at 64, 2^27 at 32. +* `max_neighbors_per_leaf` must be explicit above the 2048 clamp; `max_interactions_per_node` + needs 16384 at leaf 32 (and at leaf 64 for theta 0.4). + +## Traps recorded + +* The perfetto trace of a leaf-32 step hung for 44 h after the timing (use `--no-trace` there). +* Back-to-back configs on one card get rejected by the guard on the previous process's stale + utilisation reading; the queue scripts wait for 0 % and retry. +* `JACCPOT_M2L_DEGREE_BATCHED` is read at trace time; a microbench that flips it must + `jax.clear_caches()` or the second variant reuses the first trace (bit-identical output). +* The `l2l_only` detail diag mode is inconsistent with the cumulative modes -- do not attribute from it. From 5a97ea309ccd1b21afe391097e1bc66f62224dd5 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Thu, 10 Sep 2026 00:33:24 +0200 Subject: [PATCH 10/12] bench(gpu_gate): keep _MUST_RUN_SM80 to its one gated module; the new sm_80 tests run in the ordinary suite Co-Authored-By: Claude Fable 5.1 --- bench/gpu_gate.py | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/bench/gpu_gate.py b/bench/gpu_gate.py index 5f23aa9b..3ff623c9 100644 --- a/bench/gpu_gate.py +++ b/bench/gpu_gate.py @@ -91,16 +91,14 @@ "test_fused_pallas_complex_m2l_matches_the_pure_jax_lane_in_gradient[False]", "test_the_production_real_fused_m2l_kernel_carries_the_axis_derivative", "test_the_production_complex_fused_m2l_kernel_carries_the_axis_derivative", - # The near-field leaf-pair kernel had NO gate entry before the self-fold - # (plan "small leaves", Phase 1); its Triton lowering is what this checks. - "test_leafpair_include_self_gpu_matches_reference[None]", - "test_leafpair_include_self_gpu_matches_reference[3]", - # The target-tiled CSR M2L kernel (plan "small leaves", Phase 2): dynamic-trip - # loop, row gathers by id, atan2 and pow2 tiles all lower only on sm_80. - "test_csr_pallas_gpu_matches_rot_scale[2]", - "test_csr_pallas_gpu_matches_rot_scale[4]", - "test_csr_pallas_gpu_matches_rot_scale[6]", ) +# NOT registered here, by the registry's own contract: `tests/unit/test_gpu_gate_checks.py` +# derives the gated set from ONE module (test_transverse_degeneracy_jvp.py) and +# fails on any fragment outside it. The sm_80 tests of the near-field self-fold +# (test_pallas_nearfield_fused.py::test_leafpair_include_self_gpu_matches_reference) +# and of the CSR M2L kernel (test_m2l_real_csr_pallas.py::test_csr_pallas_gpu_matches_rot_scale) +# therefore run under the ordinary GPU suite; widening the registry to several +# modules is a separate change. # Measured on an A100 sm_80 / jax 0.10.2 and documented in ARCHITECTURE.md §9, # which carries the per-entry reasoning and the deterministic-ops column. A hit From 213ce3a7269d0ed9a2a932419e5bde4d9baf6b21 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 9 Sep 2026 22:46:44 +0000 Subject: [PATCH 11/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- bench/m2l_csr_microbench.py | 66 ++++++++--- jaccpot/pallas/m2l_real_csr.py | 30 ++++- .../operators/test_m2l_real_csr_pallas.py | 105 +++++++++++++----- .../unit/runtime/test_m2l_csr_lane_wiring.py | 16 ++- 4 files changed, 166 insertions(+), 51 deletions(-) diff --git a/bench/m2l_csr_microbench.py b/bench/m2l_csr_microbench.py index 27638d46..c681bced 100644 --- a/bench/m2l_csr_microbench.py +++ b/bench/m2l_csr_microbench.py @@ -79,7 +79,9 @@ def timed(fn, *a): src_j, tgt_j = jnp.asarray(src), jnp.asarray(tgt) counts = np.bincount(tgt, minlength=n) - csr = jax.jit(lambda m, c, s, t: m2l_real_csr_pallas(m, c, s, t, order=order)) + csr = jax.jit( + lambda m, c, s, t: m2l_real_csr_pallas(m, c, s, t, order=order) + ) out_k, t_k, med_k = timed(csr, mult, centers, src_j, tgt_j) chunk = int(args.chunk) @@ -94,9 +96,21 @@ def body(acc, i): valid = idx < P sc = s[idx] tc = t[idx] - contrib = m2l_rot_scale_real_batch(m[sc], c[tc] - c[sc], order=order) - return _chunk_segment_scatter_add(acc, contrib, tc, valid, chunk_size=chunk), None - acc, _ = lax.scan(body, jnp.zeros((n, C), jnp.float32), jnp.arange(n_chunks, dtype=jnp.int32)) + contrib = m2l_rot_scale_real_batch( + m[sc], c[tc] - c[sc], order=order + ) + return ( + _chunk_segment_scatter_add( + acc, contrib, tc, valid, chunk_size=chunk + ), + None, + ) + + acc, _ = lax.scan( + body, + jnp.zeros((n, C), jnp.float32), + jnp.arange(n_chunks, dtype=jnp.int32), + ) return acc rows = {} @@ -112,18 +126,40 @@ def body(acc, i): del fn ref = np.asarray(rows["0"][0], np.float64) assert np.all(np.isfinite(ref)), "pure-JAX reference has non-finite rows" - assert np.all(np.isfinite(np.asarray(out_k))), "CSR kernel produced non-finite rows" - rel = float(np.linalg.norm(np.asarray(out_k, np.float64) - ref) / np.linalg.norm(ref)) - rel_db = float(np.linalg.norm(np.asarray(rows["1"][0], np.float64) - ref) / np.linalg.norm(ref)) - row = dict(order=order, pairs=P, nodes=n, longest_row=int(counts.max()), - csr_ns_per_pair=1e9 * t_k / P, csr_ms=1e3 * t_k, csr_median_ms=1e3 * med_k, - pure_ns_per_pair=1e9 * rows["0"][1] / P, pure_ms=1e3 * rows["0"][1], - pure_db_ns_per_pair=1e9 * rows["1"][1] / P, pure_db_ms=1e3 * rows["1"][1], - rel_l2_csr_vs_pure=rel, rel_l2_db_vs_pure=rel_db, chunk=chunk) + assert np.all( + np.isfinite(np.asarray(out_k)) + ), "CSR kernel produced non-finite rows" + rel = float( + np.linalg.norm(np.asarray(out_k, np.float64) - ref) + / np.linalg.norm(ref) + ) + rel_db = float( + np.linalg.norm(np.asarray(rows["1"][0], np.float64) - ref) + / np.linalg.norm(ref) + ) + row = dict( + order=order, + pairs=P, + nodes=n, + longest_row=int(counts.max()), + csr_ns_per_pair=1e9 * t_k / P, + csr_ms=1e3 * t_k, + csr_median_ms=1e3 * med_k, + pure_ns_per_pair=1e9 * rows["0"][1] / P, + pure_ms=1e3 * rows["0"][1], + pure_db_ns_per_pair=1e9 * rows["1"][1] / P, + pure_db_ms=1e3 * rows["1"][1], + rel_l2_csr_vs_pure=rel, + rel_l2_db_vs_pure=rel_db, + chunk=chunk, + ) results.append(row) - print(f"p={order} pairs={P/1e6:.0f}M longest_row={counts.max()}: CSR {row['csr_ns_per_pair']:.1f} ns/pair " - f"({row['csr_ms']:.1f} ms) | pure-JAX {row['pure_ns_per_pair']:.1f} ns/pair | degree-batched " - f"{row['pure_db_ns_per_pair']:.1f} ns/pair | rel-L2 csr {rel:.2e}, db {rel_db:.2e}", flush=True) + print( + f"p={order} pairs={P/1e6:.0f}M longest_row={counts.max()}: CSR {row['csr_ns_per_pair']:.1f} ns/pair " + f"({row['csr_ms']:.1f} ms) | pure-JAX {row['pure_ns_per_pair']:.1f} ns/pair | degree-batched " + f"{row['pure_db_ns_per_pair']:.1f} ns/pair | rel-L2 csr {rel:.2e}, db {rel_db:.2e}", + flush=True, + ) if args.out: with open(args.out, "w") as fh: json.dump(results, fh, indent=2) diff --git a/jaccpot/pallas/m2l_real_csr.py b/jaccpot/pallas/m2l_real_csr.py index ea9247ca..4cb0308d 100644 --- a/jaccpot/pallas/m2l_real_csr.py +++ b/jaccpot/pallas/m2l_real_csr.py @@ -162,7 +162,8 @@ def m2l_real_csr_tables(order: int) -> dict: with jax.ensure_compile_time_eval(): for ell in range(p + 1): b = np.asarray( - compute_real_B_matrix_multipole(ell, dtype=jnp.float64), dtype=np.float64 + compute_real_B_matrix_multipole(ell, dtype=jnp.float64), + dtype=np.float64, ) Bstack[ell, p - ell : p + ell + 1, p - ell : p + ell + 1] = b BstackT = np.swapaxes(Bstack, -1, -2).copy() @@ -201,9 +202,21 @@ def m2l_real_csr_tables(order: int) -> dict: degn[: p + 1] = np.arange(p + 1) + 1 degk[: p + 1] = np.arange(p + 1) return dict( - p=p, C=C, W=W, Wp=Wp, Bp=Bp, idx=idx, mask=mask, - Bstack=Bstack, BstackT=BstackT, Apat=Apat, mabs=mabs, signm=signm, - Zf=Zf, degn=degn, degk=degk, + p=p, + C=C, + W=W, + Wp=Wp, + Bp=Bp, + idx=idx, + mask=mask, + Bstack=Bstack, + BstackT=BstackT, + Apat=Apat, + mabs=mabs, + signm=signm, + Zf=Zf, + degn=degn, + degk=degk, ) @@ -394,7 +407,9 @@ def csr_by_target( P = int(src.shape[0]) valid = (src >= 0) & (tgt >= 0) if active_pair_count is not None: - valid = valid & (jnp.arange(P, dtype=jnp.int32) < jnp.asarray(active_pair_count, jnp.int32)) + valid = valid & ( + jnp.arange(P, dtype=jnp.int32) < jnp.asarray(active_pair_count, jnp.int32) + ) key = jnp.where(valid, tgt, jnp.asarray(total_nodes, jnp.int32)) perm = jnp.argsort(key, stable=True) src_sorted = jnp.where(valid[perm], src[perm], 0) @@ -441,7 +456,10 @@ def m2l_real_csr_jax( tgt = jnp.asarray(targets, jnp.int32) valid = (src >= 0) & (tgt >= 0) if active_pair_count is not None: - valid = valid & (jnp.arange(src.shape[0], dtype=jnp.int32) < jnp.asarray(active_pair_count, jnp.int32)) + valid = valid & ( + jnp.arange(src.shape[0], dtype=jnp.int32) + < jnp.asarray(active_pair_count, jnp.int32) + ) s = jnp.where(valid, src, 0) tt = jnp.where(valid, tgt, 0) deltas = centers[tt] - centers[s] diff --git a/tests/unit/operators/test_m2l_real_csr_pallas.py b/tests/unit/operators/test_m2l_real_csr_pallas.py index 77d5a225..66acad5c 100644 --- a/tests/unit/operators/test_m2l_real_csr_pallas.py +++ b/tests/unit/operators/test_m2l_real_csr_pallas.py @@ -61,14 +61,18 @@ def _reference(mult, centers, src, tgt, order, active=None): valid &= np.arange(src.shape[0]) < active s, t = src[valid], tgt[valid] deltas = centers[t] - centers[s] - contrib = np.asarray(m2l_rot_scale_real_batch(jnp.asarray(mult[s]), jnp.asarray(deltas), order=order)) + contrib = np.asarray( + m2l_rot_scale_real_batch(jnp.asarray(mult[s]), jnp.asarray(deltas), order=order) + ) out = np.zeros_like(mult, dtype=np.float64) np.add.at(out, t, contrib.astype(np.float64)) return out def _relerr(a, ref): - return float(np.linalg.norm(np.asarray(a, np.float64) - ref) / (np.linalg.norm(ref) + 1e-30)) + return float( + np.linalg.norm(np.asarray(a, np.float64) - ref) / (np.linalg.norm(ref) + 1e-30) + ) def test_csr_by_target_partitions_the_live_pairs(): @@ -108,8 +112,12 @@ def test_csr_pallas_interpret_matches_rot_scale_f64(order): mult, centers, src, tgt = _case(order, np.float64, seed=order) ref = _reference(mult, centers, src, tgt, order) got = m2l_real_csr_pallas( - jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), jnp.asarray(tgt), - order=order, interpret=True, + jnp.asarray(mult), + jnp.asarray(centers), + jnp.asarray(src), + jnp.asarray(tgt), + order=order, + interpret=True, ) assert got.shape == mult.shape assert _relerr(got, ref) < 1e-10 @@ -122,11 +130,19 @@ def test_csr_pallas_interpret_matches_twin_and_rot_scale_f32(order): mult, centers, src, tgt = _case(order, np.float32, seed=10 + order) ref = _reference(mult, centers, src, tgt, order) got = m2l_real_csr_pallas( - jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), jnp.asarray(tgt), - order=order, interpret=True, + jnp.asarray(mult), + jnp.asarray(centers), + jnp.asarray(src), + jnp.asarray(tgt), + order=order, + interpret=True, ) twin = m2l_real_csr_jax( - jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), jnp.asarray(tgt), order=order + jnp.asarray(mult), + jnp.asarray(centers), + jnp.asarray(src), + jnp.asarray(tgt), + order=order, ) assert _relerr(got, ref) < 3e-4 assert _relerr(got, np.asarray(twin, np.float64)) < 1e-5 @@ -140,8 +156,13 @@ def test_csr_pallas_interpret_active_pair_count_truncates(): active = 23 ref = _reference(mult, centers, src, tgt, order, active=active) got = m2l_real_csr_pallas( - jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), jnp.asarray(tgt), - order=order, active_pair_count=jnp.asarray(active, jnp.int32), interpret=True, + jnp.asarray(mult), + jnp.asarray(centers), + jnp.asarray(src), + jnp.asarray(tgt), + order=order, + active_pair_count=jnp.asarray(active, jnp.int32), + interpret=True, ) assert _relerr(got, ref) < 1e-10 @@ -153,10 +174,16 @@ def test_csr_pallas_interpret_on_axis_deltas_are_exact(): order = 4 mult, centers, src, tgt = _case(order, np.float64, seed=8, on_axis=True) ref = _reference(mult, centers, src, tgt, order) - got = np.asarray(m2l_real_csr_pallas( - jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), jnp.asarray(tgt), - order=order, interpret=True, - )) + got = np.asarray( + m2l_real_csr_pallas( + jnp.asarray(mult), + jnp.asarray(centers), + jnp.asarray(src), + jnp.asarray(tgt), + order=order, + interpret=True, + ) + ) assert np.all(np.isfinite(got)) assert _relerr(got, ref) < 1e-10 @@ -165,10 +192,18 @@ def test_csr_pallas_under_jit_with_traced_active_count(): """The wrapper (sort + CSR + pallas_call) must trace: caps are static, counts traced.""" order = 3 mult, centers, src, tgt = _case(order, np.float32, seed=11) - fn = jax.jit(lambda m, c, s, t, a: m2l_real_csr_pallas(m, c, s, t, order=order, - active_pair_count=a, interpret=True)) - got = fn(jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), jnp.asarray(tgt), - jnp.asarray(40, jnp.int32)) + fn = jax.jit( + lambda m, c, s, t, a: m2l_real_csr_pallas( + m, c, s, t, order=order, active_pair_count=a, interpret=True + ) + ) + got = fn( + jnp.asarray(mult), + jnp.asarray(centers), + jnp.asarray(src), + jnp.asarray(tgt), + jnp.asarray(40, jnp.int32), + ) ref = _reference(mult, centers, src, tgt, order) assert _relerr(got, ref) < 3e-4 @@ -183,8 +218,12 @@ def test_csr_pallas_gpu_matches_rot_scale(order): mult, centers, src, tgt = _case(order, np.float32, n=40, pairs=400, seed=20 + order) ref = _reference(mult, centers, src, tgt, order) got = m2l_real_csr_pallas( - jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), jnp.asarray(tgt), - order=order, interpret=False, + jnp.asarray(mult), + jnp.asarray(centers), + jnp.asarray(src), + jnp.asarray(tgt), + order=order, + interpret=False, ) assert np.all(np.isfinite(np.asarray(got))) assert _relerr(got, ref) < 3e-4 @@ -196,15 +235,27 @@ def test_csr_pallas_gpu_matches_rot_scale(order): def test_csr_pallas_rejects_a_coefficient_count_of_another_order(): mult, centers, src, tgt = _case(3, np.float32, seed=30) with pytest.raises(ValueError, match="coefficients"): - m2l_real_csr_pallas(jnp.asarray(mult), jnp.asarray(centers), jnp.asarray(src), - jnp.asarray(tgt), order=4, interpret=True) + m2l_real_csr_pallas( + jnp.asarray(mult), + jnp.asarray(centers), + jnp.asarray(src), + jnp.asarray(tgt), + order=4, + interpret=True, + ) def test_csr_pallas_rejects_misaligned_centers(): mult, centers, src, tgt = _case(3, np.float32, seed=31) with pytest.raises(ValueError, match="centers"): - m2l_real_csr_pallas(jnp.asarray(mult), jnp.asarray(centers[:-1]), jnp.asarray(src), - jnp.asarray(tgt), order=3, interpret=True) + m2l_real_csr_pallas( + jnp.asarray(mult), + jnp.asarray(centers[:-1]), + jnp.asarray(src), + jnp.asarray(tgt), + order=3, + interpret=True, + ) def test_pack_unpack_centred_round_trip(): @@ -212,7 +263,11 @@ def test_pack_unpack_centred_round_trip(): for order in (2, 4, 6): c = sh_size(order) - x = jnp.asarray(np.random.default_rng(order).standard_normal((5, c)).astype(np.float32)) + x = jnp.asarray( + np.random.default_rng(order).standard_normal((5, c)).astype(np.float32) + ) rows = pack_centred(x, order=order) assert rows.shape[1] & (rows.shape[1] - 1) == 0 # pow2 row width - np.testing.assert_array_equal(np.asarray(unpack_centred(rows, order=order)), np.asarray(x)) + np.testing.assert_array_equal( + np.asarray(unpack_centred(rows, order=order)), np.asarray(x) + ) diff --git a/tests/unit/runtime/test_m2l_csr_lane_wiring.py b/tests/unit/runtime/test_m2l_csr_lane_wiring.py index 6f180eb0..0e9fdf28 100644 --- a/tests/unit/runtime/test_m2l_csr_lane_wiring.py +++ b/tests/unit/runtime/test_m2l_csr_lane_wiring.py @@ -59,13 +59,19 @@ def test_csr_lane_is_taken_and_matches_the_chunked_lane(monkeypatch): def solve(): # the default (non-strict) runtime: CPU-friendly, real basis, same downward solver = FastMultipoleMethod( - basis="real", theta=0.6, - G=1.0, softening=1e-3, working_dtype=jnp.float32, + basis="real", + theta=0.6, + G=1.0, + softening=1e-3, + working_dtype=jnp.float32, advanced=FMMAdvancedConfig( tree=TreeConfig(mode="static_radix", leaf_target=32), - farfield=FarFieldConfig(mode="auto"), nearfield=NearFieldConfig(mode="auto"), - mac_type="dehnen"), - fixed_order=3) + farfield=FarFieldConfig(mode="auto"), + nearfield=NearFieldConfig(mode="auto"), + mac_type="dehnen", + ), + fixed_order=3, + ) acc = solver.compute_accelerations( jnp.asarray(pos), jnp.asarray(mass), leaf_size=32, max_order=3, theta=0.6 ) From 9ade84bba1e2fdd1b62e44dcc9b92e4b35ce8774 Mon Sep 17 00:00:00 2001 From: TobiBu Date: Thu, 10 Sep 2026 10:03:17 +0200 Subject: [PATCH 12/12] style(m2l-csr): annotate the two closures, pydoclint sections, black CI on #340: test_type_annotation_guard flagged the inner `one`/`body` closures; pydoclint wanted Raises sections on m2l_real_csr_tables / m2l_real_csr_pallas and full docstrings on _bapply/_dz; black reformatted four files. Co-Authored-By: Claude Fable 5.1 --- jaccpot/pallas/m2l_real_csr.py | 49 +++++++++++++++++++++++++++++++--- 1 file changed, 45 insertions(+), 4 deletions(-) diff --git a/jaccpot/pallas/m2l_real_csr.py b/jaccpot/pallas/m2l_real_csr.py index 4cb0308d..8c66f44f 100644 --- a/jaccpot/pallas/m2l_real_csr.py +++ b/jaccpot/pallas/m2l_real_csr.py @@ -140,6 +140,11 @@ def m2l_real_csr_tables(order: int) -> dict: preserves ``m``, so it is a degree x degree operator per column. ``degn [Bp]``: ``n + 1`` (radius exponent of the output degree); ``degk [Bp]``: ``k`` (radius exponent of the source degree). + + Raises + ------ + ValueError + If ``order`` is negative. """ p = int(order) if p < 0: @@ -276,12 +281,42 @@ def unpack_centred(rows: Array, *, order: int) -> Array: def _bapply(bstack: Array, rows: Array) -> Array: - """``out[l, i] = sum_j bstack[l, i, j] rows[l, j]`` (block-diagonal by degree).""" + """Apply the per-degree constant blocks: ``out[l, i] = sum_j bstack[l, i, j] rows[l, j]``. + + Parameters + ---------- + bstack : Array + Block-diagonal-by-degree operator stack, ``(Bp, Wp, Wp)``. + rows : Array + Centred coefficient rows, ``(Bp, Wp)``. + + Returns + ------- + Array + ``(Bp, Wp)``. + """ return jnp.sum(bstack * rows[:, None, :], axis=-1) def _dz(rows: Array, cosv: Array, sinv: Array, apat: Array) -> Array: - """``Dz(t)`` on every degree row at once: ``cos(|m|t) v + A (sin(|m|t) v)``.""" + """``Dz(t)`` on every degree row at once: ``cos(|m|t) v + A (sin(|m|t) v)``. + + Parameters + ---------- + rows : Array + Centred coefficient rows, ``(Bp, Wp)``. + cosv : Array + ``cos(|m| t)`` per column, ``(Wp,)``. + sinv : Array + ``sin(|m| t)`` per column, ``(Wp,)``; pass its negative for ``Dz(-t)``. + apat : Array + The constant antisymmetric sine pattern, ``(Wp, Wp)``. + + Returns + ------- + Array + ``(Bp, Wp)``. + """ sv = rows * sinv[None, :] return rows * cosv[None, :] + jnp.sum(apat[None, :, :] * sv[:, None, :], axis=-1) @@ -367,7 +402,7 @@ def m2l_real_csr_pair_jax(multipoles: Array, deltas: Array, *, order: int) -> Ar rows = pack_centred(mult, order=int(order)).reshape(-1, Bp, Wp) d = jnp.asarray(deltas, dtype=dtype) - def one(rw, dd): + def one(rw: Array, dd: Array) -> Array: return _m2l_pair_rows(rw, (dd[0], dd[1], dd[2]), t) out_rows = jax.vmap(one)(rows, d).reshape(-1, Bp * Wp) @@ -519,7 +554,7 @@ def _m2l_real_csr_kernel( ctz = cent_ref[tgt, 2] acc0 = jnp.zeros((bp, wp), dtype=out_ref.dtype) - def body(k, acc): + def body(k: Array, acc: Array) -> Array: sid = src_ref[start + k] rows = mult_ref[sid, :].reshape(bp, wp) dx = ctx - cent_ref[sid, 0] @@ -574,6 +609,12 @@ def m2l_real_csr_pallas( ------- Array ``[n, C]`` local increments, same dtype as ``multipoles``. + + Raises + ------ + ValueError + If ``multipoles`` does not carry ``(p+1)^2`` coefficients or ``centers`` + is not ``(n, 3)`` aligned with it. """ tb = m2l_real_csr_tables(int(order)) C, Bp, Wp = tb["C"], tb["Bp"], tb["Wp"]