diff --git a/bench/gpu_gate.py b/bench/gpu_gate.py index df3e070f..3ff623c9 100644 --- a/bench/gpu_gate.py +++ b/bench/gpu_gate.py @@ -93,10 +93,12 @@ "test_the_production_complex_fused_m2l_kernel_carries_the_axis_derivative", ) # 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 test of the near-field self-fold +# 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) -# runs under the ordinary GPU suite; widening the registry is a separate change. +# 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 diff --git a/bench/m2l_csr_microbench.py b/bench/m2l_csr_microbench.py new file mode 100644 index 00000000..c681bced --- /dev/null +++ b/bench/m2l_csr_microbench.py @@ -0,0 +1,170 @@ +"""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) + 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) + + 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 + # 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) + 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, + ) + 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/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. diff --git a/jaccpot/pallas/m2l_real_csr.py b/jaccpot/pallas/m2l_real_csr.py new file mode 100644 index 00000000..8c66f44f --- /dev/null +++ b/jaccpot/pallas/m2l_real_csr.py @@ -0,0 +1,667 @@ +"""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 (``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``. + +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 +``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", + "pack_centred", + "unpack_centred", +] + +_TABLE_KEYS = ( + "Bstack", + "BstackT", + "Apat", + "mabs", + "signm", + "Zf", + "degn", + "degk", +) + + +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. + + 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 + Expansion order ``p``. + + Returns + ------- + dict + NumPy float64 arrays (cast to the working dtype by the caller) plus the + 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]``. + + ``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. ``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). + + Raises + ------ + ValueError + If ``order`` is negative. + """ + p = int(order) + if p < 0: + raise ValueError("order must be >= 0") + C = sh_size(p) + W = 2 * p + 1 + Wp = _next_pow2(W) + Bp = _next_pow2(p + 1) + + 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): + 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 + # 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) + 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) + 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): + 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, + 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, + ) + + +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} + + +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 _bapply(bstack: Array, rows: Array) -> Array: + """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)``. + + 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) + + +def _m2l_pair_rows(rows: Array, delta3: tuple, t: dict[str, Array]) -> Array: + """Full real M2L for one pair in the centred layout. + + Parameters + ---------- + 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. + + Returns + ------- + Array + Local contribution in centred rows, ``(Bp, Wp)``. + """ + x, y, z = delta3 + 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)) + log_rinv = -jnp.log(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 + 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 + 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 + + +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)) + Bp, Wp = tb["Bp"], tb["Wp"] + mult = jnp.asarray(multipoles) + dtype = mult.dtype + t = _tables_to_jnp(int(order), dtype) + rows = pack_centred(mult, order=int(order)).reshape(-1, Bp, Wp) + d = jnp.asarray(deltas, dtype=dtype) + + 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) + return unpack_centred(out_rows, order=int(order)) + + +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 centred multipole table ``[n, Bp*Wp]`` (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, Bp*Wp]``. + bp : int + ``Bp``. Static. + wp : int + ``Wp``. Static. + + Returns + ------- + None + 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) :] + 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] + acc0 = jnp.zeros((bp, wp), dtype=out_ref.dtype) + + 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] + dy = cty - cent_ref[sid, 1] + dz_ = ctz - cent_ref[sid, 2] + return acc + _m2l_pair_rows(rows, (dx, dy, dz_), t) + + acc = lax.fori_loop(0, cnt, body, acc0) + out_ref[0, :] = acc.reshape(bp * wp) + + +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 = 4, +) -> 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 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 + ------- + 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"] + 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_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 + ) + 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: + backend_kwargs["compiler_params"] = type(backend_kwargs["compiler_params"])( + num_warps=int(num_warps) + ) + out_rows = pl.pallas_call( + kernel, + grid=(n,), + in_specs=[ + 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, 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_c, cent_p, src_sorted, offsets, counts, *table_arrays) + return unpack_centred(out_rows, order=int(order)) 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/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 new file mode 100644 index 00000000..66acad5c --- /dev/null +++ b/tests/unit/operators/test_m2l_real_csr_pallas.py @@ -0,0 +1,273 @@ +"""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 + + +# ---------------------------------------------------------------- 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) + ) 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..0e9fdf28 --- /dev/null +++ b/tests/unit/runtime/test_m2l_csr_lane_wiring.py @@ -0,0 +1,97 @@ +"""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