diff --git a/tests/applications/test_svgd_vs_exact.py b/tests/applications/test_svgd_vs_exact.py index 1db2f05..e00f5f4 100644 --- a/tests/applications/test_svgd_vs_exact.py +++ b/tests/applications/test_svgd_vs_exact.py @@ -12,6 +12,7 @@ import jax import jax.numpy as jnp +import numpy as np import pytest from yggdrax import DualTreeTraversalConfig @@ -241,3 +242,62 @@ def test_svgd_phi_cutoff_bandwidths_is_the_c_times_h_convention(): topo = build_svgd_topology(p, kernel_cutoff=6.0 * h, **kw) via_topo = svgd_phi_from_topology(p, sc, h, topo) assert float(jnp.max(jnp.abs(via_helper - via_topo))) < 1e-12 + + +# --- the two near-field accumulations -------------------------------------- +# +# The near field can be summed by scattering each unordered pair's two +# directions, or by a segmented reduction over the directed list. Which is +# faster is a property of the dtype (float32's `.at[].add()` is 43x float64's +# under index contention), so both exist. They must agree. + + +@pytest.mark.parametrize("backend", ["radix", "leaf_kdtree"]) +def test_accumulations_agree(backend): + p = jax.random.normal(jax.random.PRNGKey(5), (1200, 3)) * 1.2 + sc = p * 0.5 + h = 0.6 + topo = build_svgd_topology( + p, theta=0.4, leaf_size=16, backend=backend, traversal_config=_CFG + ) + by_scatter = svgd_phi_from_topology(p, sc, h, topo, accumulate="scatter") + by_segment = svgd_phi_from_topology(p, sc, h, topo, accumulate="segment") + assert float(jnp.max(jnp.abs(by_scatter - by_segment))) < 1e-12 + + # "auto" is one of the two, never a third thing. + by_auto = svgd_phi_from_topology(p, sc, h, topo, accumulate="auto") + assert float(jnp.max(jnp.abs(by_auto - by_scatter))) < 1e-12 + + +def test_segment_accumulation_is_exact_at_theta_zero(): + """The segmented path is a different summation order, not a different sum.""" + p = jax.random.normal(jax.random.PRNGKey(6), (900, 3)) * 1.2 + sc = p * 0.5 + h = float(median_heuristic(p)) + topo = build_svgd_topology( + p, theta=0.0, leaf_size=16, backend="radix", traversal_config=_CFG + ) + ref = exact_phi(p, sc, h) + out = svgd_phi_from_topology(p, sc, h, topo, accumulate="segment") + assert float(jnp.linalg.norm(out - ref) / jnp.linalg.norm(ref)) < 1e-10 + + +def test_unknown_accumulation_is_rejected(): + p = jax.random.normal(jax.random.PRNGKey(7), (200, 3)) + topo = build_svgd_topology( + p, theta=0.4, leaf_size=16, backend="radix", traversal_config=_CFG + ) + with pytest.raises(ValueError, match="accumulate must be"): + svgd_phi_from_topology(p, p * 0.5, 0.6, topo, accumulate="nonsense") + + +def test_directed_pair_list_is_sorted_and_twice_the_halved_one(): + """What makes the segmented reduction segmented.""" + p = jax.random.normal(jax.random.PRNGKey(8), (1500, 3)) * 1.2 + topo = build_svgd_topology( + p, theta=0.4, leaf_size=16, backend="radix", traversal_config=_CFG + ) + directed = np.asarray(topo.near_dir_target) + assert directed.shape[0] == 2 * int(topo.near_target_row.shape[0]) + assert directed.shape[0] == int(topo.num_near_leaf_pairs) + assert np.all(np.diff(directed) >= 0), "target rows must be non-decreasing" diff --git a/yggdrax/applications/svgd/exact.py b/yggdrax/applications/svgd/exact.py index bc1049c..cff3e2f 100644 --- a/yggdrax/applications/svgd/exact.py +++ b/yggdrax/applications/svgd/exact.py @@ -13,31 +13,72 @@ import jax.numpy as jnp from jaxtyping import Array, Float -from yggdrax.applications.svgd.kernel import stein_pair_terms - def exact_phi( particles: Float[Array, "n d"], scores: Float[Array, "n d"], h: float | Float[Array, ""], + *, + block_size: int | None = None, ) -> Float[Array, "n d"]: """Exact Stein update direction phi(x_i) for every particle, O(N^2). + The pair sum is contracted, not materialised. Writing it out, + + .. math:: + + \\phi_i = \\frac{1}{N}\\Big[(K S)_i + + \\big(x_i (K \\mathbf{1})_i - (K X)_i\\big) / h^2\\Big], + + with :math:`K_{ij} = \\exp(-\\lVert x_i - x_j\\rVert^2 / 2h^2)`, so the whole + update is one kernel matrix and **two matmuls**. The obvious form builds an + ``(n, n, d)`` tensor of per-pair terms, which is *d* times the memory, runs + at elementwise rather than GEMM throughput, and cannot be evaluated at all + beyond N ~ 2e4 -- which mattered, because this is the baseline the tree + update is judged against, and a weak baseline flatters the tree. + Args: particles: Particle positions, shape ``(n, d)``. scores: Target score at each particle, shape ``(n, d)``. h: Kernel bandwidth. + block_size: Targets per block. ``None`` does every target at once, + which costs an ``(n, n)`` kernel matrix; pass a block size to cap + that at ``(block_size, n)`` and reach large N. Returns: Update directions, shape ``(n, d)``. """ - n = particles.shape[0] - # target axis 0 (i), source axis 1 (j): terms[i, j] contributes to phi[i]. - x_t = particles[:, None, :] - x_s = particles[None, :, :] - s_s = scores[None, :, :] - terms = stein_pair_terms(x_t, x_s, s_s, h) # (n, n, d) - return jnp.sum(terms, axis=1) / n + n, d = particles.shape + if block_size is None: + block_size = n + block = max(1, min(int(block_size), n)) + + def _block(x_t: Array) -> Array: + """Contribution of every source to one block of targets.""" + # (B, n) kernel, then two contractions -- no (B, n, d) tensor is ever + # built. sum_j k_ij s_j is a matmul, and + # sum_j k_ij (x_i - x_j) = x_i * sum_j k_ij - sum_j k_ij x_j is another. + d2 = ( + jnp.sum(x_t * x_t, axis=-1)[:, None] + - 2.0 * (x_t @ particles.T) + + jnp.sum(particles * particles, axis=-1)[None, :] + ) + k = jnp.exp(-jnp.maximum(d2, 0.0) / (2.0 * h**2)) # (B, n) + attract = k @ scores + repulse = x_t * jnp.sum(k, axis=1)[:, None] - k @ particles + return attract + repulse / (h**2) + + if block >= n: + return _block(particles) / n + + pad = (-n) % block + padded = ( + particles + if pad == 0 + else jnp.concatenate([particles, jnp.zeros((pad, d), particles.dtype)]) + ) + out = jax.lax.map(_block, padded.reshape(-1, block, d)) + return out.reshape(-1, d)[:n] / n def svgd_step( diff --git a/yggdrax/applications/svgd/sampler.py b/yggdrax/applications/svgd/sampler.py index 5f5abee..a0b92a1 100644 --- a/yggdrax/applications/svgd/sampler.py +++ b/yggdrax/applications/svgd/sampler.py @@ -174,6 +174,8 @@ class SvgdTopology(NamedTuple): leaf_mask: Array # (L, max_leaf) 1.0 valid / 0.0 pad near_target_row: Array # (Q,) one row of each UNORDERED near leaf pair near_source_row: Array # (Q,) the other row; row_a < row_b throughout + near_dir_target: Array # (2Q,) DIRECTED pairs, ascending in target row + near_dir_source: Array # (2Q,) the source row of each directed pair far_tgt_slot: Array # (M,) sorted-slot of each far target particle far_src_start: Array # (M,) inclusive start slot of the far source node far_src_end: Array # (M,) inclusive end slot of the far source node @@ -356,6 +358,12 @@ def assemble_svgd_topology(walk: SvgdTraversal) -> SvgdTopology: "Stein partition is only complete when every near pair appears in " "both directions." ) + # Keep the directed list too: it is what the segment-sum accumulation needs, + # and it is already exactly what that wants -- np.repeat emits it ascending + # in target row, so there is nothing to sort. (Concatenating the halved list + # with its mirror would double it; it is *already* both directions.) + dir_target = near_target_row + dir_source = near_source_row near_target_row = near_target_row[upper] near_source_row = near_source_row[upper] @@ -415,6 +423,8 @@ def assemble_svgd_topology(walk: SvgdTraversal) -> SvgdTopology: leaf_mask=jnp.asarray(leaf_mask), near_target_row=jnp.asarray(near_target_row), near_source_row=jnp.asarray(near_source_row), + near_dir_target=jnp.asarray(dir_target), + near_dir_source=jnp.asarray(dir_source), far_tgt_slot=far_tgt_slot, far_src_start=far_src_start, far_src_end=far_src_end, @@ -540,12 +550,129 @@ def _near_chunk_bothways( return phi.at[slots_b].add(to_b * mask_b[..., None]) +def _near_chunk_to_target( + pos: Array, + sco: Array, + leaf_slots: Array, + leaf_mask: Array, + rows_t: Array, + rows_s: Array, + live: Array, + h: float | Float[Array, ""], +) -> Array: + """Return one chunk of directed near pairs' contribution to their targets. + + One direction only, so the kernel is evaluated twice per unordered pair -- + the opposite trade to :func:`_near_chunk_bothways`, and the right one when + the accumulation is a segmented reduction rather than a scatter. + + Args: + pos: Positions in sorted order, shape ``(n, d)``. + sco: Scores in sorted order, shape ``(n, d)``. + leaf_slots: Padded per-leaf slot blocks, shape ``(L, ml)``. + leaf_mask: Validity of ``leaf_slots``, shape ``(L, ml)``. + rows_t: Target leaf row of each pair, shape ``(chunk,)``. + rows_s: Source leaf row of each pair, shape ``(chunk,)``. + live: 1.0 for real pairs, 0.0 for the chunk's padding, ``(chunk,)``. + h: Kernel bandwidth. + + Returns: + Per-pair target contributions, shape ``(chunk, ml, d)``. + """ + slots_t, slots_s = leaf_slots[rows_t], leaf_slots[rows_s] + mask_s = leaf_mask[rows_s] * live[:, None] + x_t, x_s = pos[slots_t], pos[slots_s] + diff = x_t[:, :, None, :] - x_s[:, None, :, :] + k = jnp.exp(-jnp.sum(diff * diff, axis=-1) / (2.0 * h**2))[..., None] + terms = k * sco[slots_s][:, None, :, :] + k * diff / (h**2) + return jnp.sum(terms * mask_s[:, None, :, None], axis=2) + + +def _accumulate_near_by_segment( + pos: Array, + sco: Array, + topo: SvgdTopology, + h: float | Float[Array, ""], + chunk_pairs: int, +) -> Array: + """Accumulate the near field with a segmented reduction, not a scatter. + + Leaves tile ``[0, n)`` disjointly, so summing each target leaf's directed + pairs into a ``(L, ml, d)`` array and then *placing* it is a permutation -- + no index is written twice and no atomic is needed anywhere. + + This is the float32 path. ``.at[].add()`` costs 258 ms of the 279 ms float32 + near field at N = 1e5 against 6 ms of 29 ms in float64, because the scatter + indices repeat ~62x on average and float32 lowers to contended + ``atomicAdd``. Measured, near field only, N = 1e5 on an A100: + + ========================================== ========= ========= + strategy float64 float32 + ========================================== ========= ========= + gather + arithmetic only (the floor) 22.60 ms 20.47 ms + halved pairs, two scatters 28.60 ms 278.90 ms + directed pairs, segment_sum, chunked 56.20 ms 32.50 ms + ========================================== ========= ========= + + Hence the split: float64 keeps the scatter, float32 comes here. + + Args: + pos: Positions in sorted order, shape ``(n, d)``. + sco: Scores in sorted order, shape ``(n, d)``. + topo: The partition. + h: Kernel bandwidth. + chunk_pairs: Directed pairs per chunk. + + Returns: + The near-field contribution in sorted-slot order, shape ``(n, d)``. + """ + rows_t, rows_s = topo.near_dir_target, topo.near_dir_source + num_pairs = int(rows_t.shape[0]) + num_leaves, max_leaf = topo.leaf_slots.shape + chunk = min(int(chunk_pairs), num_pairs) + num_chunks = -(-num_pairs // chunk) + pad = num_chunks * chunk - num_pairs + live = jnp.ones((num_pairs,), dtype=pos.dtype) + if pad: + # Pad with the last leaf paired to itself and zero weight: the segment + # ids stay non-decreasing, which is what makes the reduction segmented. + tail = jnp.full((pad,), num_leaves - 1, dtype=rows_t.dtype) + rows_t = jnp.concatenate([rows_t, tail]) + rows_s = jnp.concatenate([rows_s, tail]) + live = jnp.concatenate([live, jnp.zeros((pad,), dtype=pos.dtype)]) + + @jax.checkpoint + def _step(carry, xs): + contrib = _near_chunk_to_target( + pos, sco, topo.leaf_slots, topo.leaf_mask, xs[0], xs[1], xs[2], h + ) + return ( + carry + + jax.ops.segment_sum( + contrib, xs[0], num_segments=num_leaves, indices_are_sorted=True + ), + None, + ) + + acc, _ = jax.lax.scan( + _step, + jnp.zeros((num_leaves, max_leaf, pos.shape[1]), dtype=pos.dtype), + ( + rows_t.reshape(num_chunks, chunk), + rows_s.reshape(num_chunks, chunk), + live.reshape(num_chunks, chunk), + ), + ) + return jnp.zeros_like(pos).at[topo.leaf_slots].add(acc * topo.leaf_mask[..., None]) + + def svgd_phi_from_topology( particles: Float[Array, "n d"], scores: Float[Array, "n d"], h: float | Float[Array, ""], topo: SvgdTopology, chunk_pairs: int | None = None, + accumulate: str = "scatter", ) -> Float[Array, "n d"]: """Tree-accelerated Stein update given a fixed partition (differentiable). @@ -556,11 +683,37 @@ def svgd_phi_from_topology( topo: Partition from :func:`build_svgd_topology`. chunk_pairs: Near leaf pairs per rematerialised chunk. ``None`` (default) picks the largest chunk whose ``(chunk, ml, ml, d)`` - tensor stays under 64 MB. + tensor stays under 256 MiB. + accumulate: How the near field is summed. ``"scatter"`` (default) sums + each unordered pair once and scatters both directions; + ``"segment"`` sums directed pairs with a segmented reduction, which + needs no atomics; ``"auto"`` picks ``"segment"`` for float32 and + ``"scatter"`` otherwise. **The default is deliberately the one that + is never worse under differentiation** -- see the note below. Returns: Update directions, shape ``(n, d)``. + + Note: + ``"segment"`` is much faster *forward* in float32 -- 11.64 -> 2.93 ms at + N = 1e4 and 586 -> 321 ms at 1e5 on an A100 -- because it replaces a + contended ``atomicAdd`` with a segmented reduction. It is **slower under + reverse mode** (fwd+grad 2392 -> 3172 ms at N = 1e5), because the + transpose of a gather is a scatter: reverse mode reintroduces exactly + the operation the forward pass removed, and doubles the arithmetic + besides, since the directed list evaluates each kernel twice. + + So a caller that only ever runs forward should ask for ``"auto"`` -- + which is what :func:`svgd_phi` and :func:`run_tree_svgd` do -- and a + caller that differentiates should leave the default alone. + + Raises: + ValueError: If ``accumulate`` is not one of the three names. """ + if accumulate not in ("scatter", "segment", "auto"): + raise ValueError( + f"accumulate must be 'scatter', 'segment' or 'auto'; got " f"{accumulate!r}" + ) n, d = particles.shape pos = particles[topo.order] # sorted order sco = scores[topo.order] @@ -578,14 +731,30 @@ def svgd_phi_from_topology( within = jnp.sum(terms * mask[:, None, :, None], axis=2) # (L, ml, d) phi = phi.at[topo.leaf_slots].add(within * mask[..., None]) - # cross-leaf near pairs: one entry per unordered pair, both directions - # scattered from one kernel evaluation, in rematerialised chunks. + # cross-leaf near pairs. Two accumulations, and which one is faster is a + # property of the dtype, not of the problem -- see _accumulate_near_by_segment + # for the table. float64's scatter is cheap and its segmented reduction is + # not; float32's scatter is 43x more expensive than float64's under index + # contention, and the reduction wins 8.6x. num_pairs = int(topo.near_target_row.shape[0]) if num_pairs > 0: max_leaf = int(topo.leaf_slots.shape[1]) if chunk_pairs is None: per_pair = max(1, max_leaf * max_leaf * d * particles.dtype.itemsize) chunk_pairs = max(1, _NEAR_CHUNK_BYTES // per_pair) + use_segment = accumulate == "segment" or ( + accumulate == "auto" and particles.dtype.itemsize <= 4 + ) + if use_segment: + return _finish( + phi + _accumulate_near_by_segment(pos, sco, topo, h, chunk_pairs), + pos, + sco, + topo, + h, + n, + d, + ) chunk = min(int(chunk_pairs), num_pairs) num_chunks = -(-num_pairs // chunk) pad = num_chunks * chunk - num_pairs @@ -625,6 +794,34 @@ def _step(carry, xs): ), ) + return _finish(phi, pos, sco, topo, h, n, d) + + +def _finish( + phi: Array, + pos: Array, + sco: Array, + topo: SvgdTopology, + h: float | Float[Array, ""], + n: int, + d: int, +) -> Array: + """Add the far field, average, and return to the caller's particle order. + + Shared by both near-field accumulations so they cannot drift apart. + + Args: + phi: Near-field accumulator in sorted-slot order, shape ``(n, d)``. + pos: Positions in sorted order, shape ``(n, d)``. + sco: Scores in sorted order, shape ``(n, d)``. + topo: The partition. + h: Kernel bandwidth. + n: Particle count. + d: Dimension. + + Returns: + Update directions in the caller's particle order, shape ``(n, d)``. + """ # --- far field: monopole (M2P) --- if topo.far_tgt_slot.shape[0] > 0: zero_x = jnp.zeros((1, d), pos.dtype) @@ -651,7 +848,7 @@ def _step(carry, xs): # pure array computation; jitting it collapses the eager per-op dispatch into one # kernel (~1.5x faster per step even when the partition shapes vary a little). _jit_svgd_phi_from_topology = jax.jit( - svgd_phi_from_topology, static_argnames=("chunk_pairs",) + svgd_phi_from_topology, static_argnames=("chunk_pairs", "accumulate") ) @@ -709,7 +906,10 @@ def svgd_phi( traversal_config=traversal_config, kernel_cutoff=_cutoff_radius(cutoff_bandwidths, h), ) - return _jit_svgd_phi_from_topology(particles, scores, h, topo) + # svgd_phi is the forward-only entry point, so it takes the accumulation + # that is fastest forward; svgd_phi_from_topology keeps the differentiable + # default for callers that wrap it in grad. + return _jit_svgd_phi_from_topology(particles, scores, h, topo, accumulate="auto") def tree_svgd_step(