diff --git a/autoarray/inversion/mesh/interpolator/sibson.py b/autoarray/inversion/mesh/interpolator/sibson.py index 65a494734..53b75cc8f 100644 --- a/autoarray/inversion/mesh/interpolator/sibson.py +++ b/autoarray/inversion/mesh/interpolator/sibson.py @@ -11,6 +11,8 @@ model sample is rejected instead of silently using a truncated stencil. """ +import os + import numpy as np from autonerves import cached_property @@ -33,7 +35,126 @@ # autolens_workspace_test/scripts/misc/jax_assertions/delaunay_nn_caps.py. SIBSON_MAX_CAVITY_TRIANGLES = 32 SIBSON_MAX_NEIGHBORS = 32 -SIBSON_QUERY_CHUNK = 256 + + +def _positive_int_env(name, raw): + """Parse a positive-integer environment override; ``None`` when unset.""" + if raw is None: + return None + message = f"{name} must be a positive integer, got {raw!r}" + try: + value = int(raw) + except ValueError as error: + raise ValueError(message) from error + if value <= 0: + raise ValueError(message) + return value + + +def _bool_env(name, raw): + """Parse a ``"0"``/``"1"`` environment override; ``None`` when unset.""" + if raw is None: + return None + if raw not in ("0", "1"): + raise ValueError(f"{name} must be '0' or '1', got {raw!r}") + return raw == "1" + + +# ``jax.lax.map`` block size for the JAX Sibson query loop, bound onto the mesh +# as ``DelaunayNN.query_chunk``. It is a MEMORY GUARD and nothing else: it +# bounds the ``(C, 3, 2)`` per-query intermediates of the cavity -- the +# inserted Watson circumcircle centres and their contributions, of order +# 15-25 kB per query per lane at cap 32 -- so the live footprint stays bounded +# when the likelihood is vmapped over many live points. +# +# It is emphatically not a speed knob to be kept small. The walk and the two +# cavity loops inside a chunk are latency-bound, so each extra chunk serialises +# another complete set of kernel launches: halving the chunk roughly doubles +# the launch count of the pass (issue #532). Raise it as far as device memory +# allows for the mesh and vmap batch in use. +# +# The default is the A100 sweep of issue #532 (2026-09-08, RAL +# ``euclid-ral-gpu-2``, NVIDIA A100 80GB PCIe, fp64, jobs 342321/342322 and the +# array 342323_[0-3]; autolens_profiling +# ``results/notes/delaunay_nn_launch_latency.md``). Cell: the HST imaging +# DelaunayNN likelihood breakdown, Hilbert-1500 mesh, MGE-60 lens light, +# ConstantSplit regularization, 17,980 over-sampled data queries + 6,000 split +# points, ``--split-setup --vmap-batch 16``. Measured params->H prefix and the +# peak ``nvidia-smi`` memory sampled through the whole run at vmap 16: +# +# chunk | params->H unbatched | params->H per call @vmap 16 | peak VRAM +# 256 | 88.29 ms | 23.07 ms | 41,495 MiB +# 512 | 52.21 ms | 18.24 ms | 41,503 MiB +# 1024 | 35.32 ms | 19.57 ms | 41,503 MiB +# 2048 | 27.22 ms | 19.13 ms | 41,503 MiB +# 4096 | 24.66 ms | 16.45 ms | 41,503 MiB +# +# (control, ``main`` at ``d7c96762``, chunk 256: 143.90 ms / 24.32 ms / +# 41,495 MiB. The eager ``EXPECTED_LOG_EVIDENCE_HST = 29144.581944`` pin held +# on every row, so the chunk is bit-neutral as designed.) +# +# 4096 is fastest on both readings and the VRAM clause of the sweep's decision +# rule turned out uninformative: the ~41.5 GiB plateau is identical at every +# chunk *and* on the control, because it is the vmap-16 dense inversion block, +# not the cavity intermediates. The guard arithmetic says why there is room -- +# ``(C, 3, 2)`` fp64 intermediates at ~25 kB per query per lane over 4096 +# queries x 16 lanes is ~1.6 GB, ~2 % of an 80 GB card. +# +# Set ``PYAUTO_SIBSON_QUERY_CHUNK`` (a positive integer) to override the +# default at import time. That is both how the sweep above was run without +# editing source and the escape hatch for a smaller GPU or a much larger cell: +# lower it until the ``(C, 3, 2)`` intermediates fit, at a proportional cost in +# sequential ``lax.map`` trips. +SIBSON_QUERY_CHUNK = 4096 + +_QUERY_CHUNK_OVERRIDE = _positive_int_env( + "PYAUTO_SIBSON_QUERY_CHUNK", os.environ.get("PYAUTO_SIBSON_QUERY_CHUNK") +) +if _QUERY_CHUNK_OVERRIDE is not None: + SIBSON_QUERY_CHUNK = _QUERY_CHUNK_OVERRIDE + +# Candidate-edge loop strategy inside the cavity walk +# (:func:`_cavity_triangle_indexes_jax`). The three edges of a cavity triangle +# are either unrolled at trace time or run through a 3-trip ``fori_loop``. The +# two are BIT-IDENTICAL -- the same ``add_candidate`` calls for edges 0, 1, 2 in +# that order, so the same cavity insertion order and therefore the same stencil +# column order and the same floating-point summation order downstream (verified +# against a frozen reference of every ``jax_delaunay_nn`` output). They differ +# only in how the program is emitted, and which is faster is a property of the +# backend, so it is decided at trace time: +# +# * accelerator -- the cavity walk is launch-latency bound, and unrolling the +# inner loop removes ~28% of the kernel launches per chunk (1,244 -> 892 on +# the A100 HST / Hilbert-1500 cell, issue #532), so unroll. +# * CPU -- there is no launch cost to remove, and the unrolled body is three +# times the code inside a 32-trip loop. Measured with an interleaved +# in-process paired A/B of ``jax_delaunay_nn`` (N=1500, Q=17,980 data + +# 6,000 split, fp64, warm median of 25-30 alternating rounds): rolled +# 566.0 / 567.8 ms vs unrolled 614.4 / 614.7 ms, i.e. the rolled loop is +# ~8-9% faster. The same harness reads 0.998 between two copies of +# identical code, so that gap is real, and the standing no-CPU-slowdown +# constraint keeps the loop rolled here. +# +# ``PYAUTO_SIBSON_UNROLL_CANDIDATES`` ("1" or "0") forces one strategy for +# benchmarking; unset means decide from ``jax.default_backend()``. +SIBSON_UNROLL_CANDIDATES = _bool_env( + "PYAUTO_SIBSON_UNROLL_CANDIDATES", + os.environ.get("PYAUTO_SIBSON_UNROLL_CANDIDATES"), +) + + +def _sibson_unroll_candidates(): + """Whether to unroll the cavity walk's three candidate edges. + + ``SIBSON_UNROLL_CANDIDATES`` (the environment override) wins when set; + otherwise every backend but CPU unrolls. See the comment above. + """ + if SIBSON_UNROLL_CANDIDATES is not None: + return SIBSON_UNROLL_CANDIDATES + + import jax + + return jax.default_backend() != "cpu" def _cross(u, v): @@ -165,6 +286,8 @@ def _cavity_triangle_indexes_jax( import jax import jax.numpy as jnp + unroll = _sibson_unroll_candidates() + safe_seed = jnp.maximum(seed_simplex, 0) cavity = -jnp.ones((max_cavity_triangles,), dtype=jnp.int32) cavity = cavity.at[0].set(safe_seed) @@ -198,7 +321,16 @@ def add_candidate(edge, inner_carry): overflow = overflow | (accepted & ~has_space) return cavity, count, overflow - return jax.lax.fori_loop(0, 3, add_candidate, (cavity, count, overflow)) + # Unrolled or rolled by ``_sibson_unroll_candidates()``; the two are + # bit-identical by construction (same calls, edges 0, 1, 2, same + # order), and the choice is a backend performance question only -- + # see the comment on ``SIBSON_UNROLL_CANDIDATES``. + carry = (cavity, count, overflow) + if unroll: + for edge in range(3): + carry = add_candidate(edge, carry) + return carry + return jax.lax.fori_loop(0, 3, add_candidate, carry) return jax.lax.fori_loop( 0, @@ -434,10 +566,37 @@ def sibson_mappings_weights_from_tables( max_cavity_triangles=SIBSON_MAX_CAVITY_TRIANGLES, max_neighbors=SIBSON_MAX_NEIGHBORS, query_chunk=SIBSON_QUERY_CHUNK, + circumcircles=None, xp=np, ): """Calculate fixed-shape Sibson mappings and weights from Delaunay tables. + Loop structure (and why the chunk exists) + ----------------------------------------- + On the JAX path this is three nested loops: a ``jax.lax.map`` over + ``query_chunk``-sized blocks of queries, a ``fori_loop`` of + ``max_cavity_triangles`` trips walking the insertion cavity + (:func:`_cavity_triangle_indexes_jax`), and the three candidate edges of + each cavity triangle -- the last of which is unrolled. Only the outer + ``lax.map`` is sequential *in the queries*: every chunk re-runs the whole + cavity walk, so the kernel-launch count of one call is roughly + ``ceil(Q / query_chunk)`` times the launches of a single chunk. With the + cavity loop dominating that per-chunk count, a small chunk multiplies a + latency-bound program rather than saving time, which is why + ``query_chunk`` is documented as a memory guard on the ``(C, 3, 2)`` + per-cavity intermediates and nothing else (issue #532). The NumPy path + ignores ``query_chunk`` entirely and loops over queries in Python. + + Parameters + ---------- + circumcircles + Optional precomputed ``(centres, radii_squared, valid)`` triple for + ``simplices_padded``, as returned by + :func:`delaunay_circumcircles_from`. They depend only on the frozen + simplex table, so a caller interpolating several query sets against + one mesh (``jax_delaunay_nn``) computes them once and passes them in. + ``None`` computes them here, which is what every other caller does. + Returns ------- mappings, sizes, weights @@ -447,9 +606,9 @@ def sibson_mappings_weights_from_tables( Prototype diagnostics. Overflow or a Watson edge degeneracy produces NaN weights rather than silently returning an approximate stencil. """ - circumcentres, circumradii_squared, circumcircle_valid = ( - delaunay_circumcircles_from(points, simplices_padded, xp=xp) - ) + if circumcircles is None: + circumcircles = delaunay_circumcircles_from(points, simplices_padded, xp=xp) + circumcentres, circumradii_squared, circumcircle_valid = circumcircles def single(query, seed_simplex, outside_fallback_index): return _sibson_single_from_tables( @@ -639,38 +798,23 @@ def jax_delaunay_nn( Qhull returns only fixed-shape integer connectivity through a stopped ``pure_callback``. Point location, circumcircles, Sibson weights, dual areas, and split-cross coordinates all remain in the JAX graph. + + The data grid and the ``4N`` split-cross points are located *and* + interpolated in ONE concatenated pass, the same pattern + :func:`jax_delaunay` uses for its walk. Both the visibility walk and the + Sibson cavity loops are latency-bound, so two passes pay two full sets of + kernel launches while one concatenated pass pays roughly one (issue #532). + The circumcircles of the frozen simplex table do not depend on the + queries, so they are computed once here and handed to + :func:`sibson_mappings_weights_from_tables` instead of being recomputed + per pass. Split points are still seeded at their own nearest vertex and + keep their own outside-hull fallback: concatenation changes only how many + programs run, never a per-query result. """ import jax.numpy as jnp simplices_padded, simplex_neighbors, vertex_simplex = _jax_delaunay_tables(points) - def mappings_weights_for(query): - delaunay_mappings, simplex_indexes = pix_indexes_delaunay_walk_from( - query_points=query, - points=points, - simplices_padded=simplices_padded, - simplex_neighbors=simplex_neighbors, - vertex_simplex=vertex_simplex, - xp=jnp, - return_simplex_indexes=True, - ) - return sibson_mappings_weights_from_tables( - query_points=query, - points=points, - simplices_padded=simplices_padded, - simplex_neighbors=simplex_neighbors, - simplex_indexes=simplex_indexes, - outside_fallback_indexes=delaunay_mappings[:, 0], - max_cavity_triangles=max_cavity_triangles, - max_neighbors=max_neighbors, - query_chunk=query_chunk, - xp=jnp, - ) - - mappings, sizes, weights, cavity_sizes, overflow, degenerate = mappings_weights_for( - query_points - ) - valid = simplices_padded[:, 0] >= 0 simplices = simplices_padded.clip(min=0) p0 = points[simplices[:, 0]] @@ -689,6 +833,44 @@ def mappings_weights_for(query): area_weights=areas_factor * jnp.sqrt(areas), xp=jnp, ) + + n_query = query_points.shape[0] + all_query_points = jnp.concatenate([query_points, split_points]) + + delaunay_mappings, simplex_indexes = pix_indexes_delaunay_walk_from( + query_points=all_query_points, + points=points, + simplices_padded=simplices_padded, + simplex_neighbors=simplex_neighbors, + vertex_simplex=vertex_simplex, + xp=jnp, + return_simplex_indexes=True, + ) + + circumcircles = delaunay_circumcircles_from(points, simplices_padded, xp=jnp) + + outputs = sibson_mappings_weights_from_tables( + query_points=all_query_points, + points=points, + simplices_padded=simplices_padded, + simplex_neighbors=simplex_neighbors, + simplex_indexes=simplex_indexes, + outside_fallback_indexes=delaunay_mappings[:, 0], + max_cavity_triangles=max_cavity_triangles, + max_neighbors=max_neighbors, + query_chunk=query_chunk, + circumcircles=circumcircles, + xp=jnp, + ) + + ( + mappings, + sizes, + weights, + cavity_sizes, + overflow, + degenerate, + ) = (output[:n_query] for output in outputs) ( splitted_mappings, splitted_sizes, @@ -696,7 +878,7 @@ def mappings_weights_for(query): split_cavity_sizes, split_overflow, split_degenerate, - ) = mappings_weights_for(split_points) + ) = (output[n_query:] for output in outputs) return ( points, diff --git a/test_autoarray/inversion/pixelization/interpolator/test_delaunay_nn.py b/test_autoarray/inversion/pixelization/interpolator/test_delaunay_nn.py index daf673cad..2f0fbb20b 100644 --- a/test_autoarray/inversion/pixelization/interpolator/test_delaunay_nn.py +++ b/test_autoarray/inversion/pixelization/interpolator/test_delaunay_nn.py @@ -25,7 +25,7 @@ def test__mesh_is_public_and_selects_natural_neighbor_interpolator(): assert mesh.areas_factor == 0.4 assert mesh.max_cavity_triangles == 32 assert mesh.max_neighbors == 32 - assert mesh.query_chunk == 256 + assert mesh.query_chunk == 4096 assert mesh.interpolator_cls is InterpolatorDelaunayNN assert aa.InterpolatorDelaunayNN is InterpolatorDelaunayNN diff --git a/test_autoarray/inversion/pixelization/interpolator/test_sibson.py b/test_autoarray/inversion/pixelization/interpolator/test_sibson.py index 418a41e2e..73e78745a 100644 --- a/test_autoarray/inversion/pixelization/interpolator/test_sibson.py +++ b/test_autoarray/inversion/pixelization/interpolator/test_sibson.py @@ -1,11 +1,19 @@ +import hashlib + import numpy as np +import pytest +from autoarray.inversion.mesh.interpolator import sibson from autoarray.inversion.mesh.interpolator.delaunay import ( pix_indexes_delaunay_walk_from, scipy_delaunay_tri_only, ) from autoarray.inversion.mesh.interpolator.sibson import ( + _bool_env, + _positive_int_env, + _sibson_unroll_candidates, delaunay_circumcircles_from, + scipy_delaunay_nn, sibson_mappings_weights_from_tables, ) @@ -175,3 +183,154 @@ def test__cavity_cap__reports_overflow_instead_of_silent_approximation(): assert cavity_sizes[0] == 1 assert overflow[0] assert np.isnan(weights[0]).all() + + +def test__precomputed_circumcircles__match_the_default_call_bit_for_bit(): + """``sibson_mappings_weights_from_tables`` accepts the circumcircles of the + frozen simplex table so a caller interpolating several query sets against + one mesh computes them once (``jax_delaunay_nn``). Passing them must be a + pure hoist: every output identical, not merely close.""" + rng = np.random.default_rng(11) + points = rng.uniform(-1.0, 1.0, size=(80, 2)) + query = rng.uniform(-0.8, 0.8, size=(120, 2)) + + simplices, neighbors, vertex_simplex = scipy_delaunay_tri_only(points) + delaunay_mappings, simplex_indexes = pix_indexes_delaunay_walk_from( + query_points=query, + points=points, + simplices_padded=simplices, + simplex_neighbors=neighbors, + vertex_simplex=vertex_simplex, + xp=np, + return_simplex_indexes=True, + ) + kwargs = dict( + query_points=query, + points=points, + simplices_padded=simplices, + simplex_neighbors=neighbors, + simplex_indexes=simplex_indexes, + outside_fallback_indexes=delaunay_mappings[:, 0], + xp=np, + ) + + default = sibson_mappings_weights_from_tables(**kwargs) + hoisted = sibson_mappings_weights_from_tables( + circumcircles=delaunay_circumcircles_from(points, simplices, xp=np), + **kwargs, + ) + + assert len(default) == len(hoisted) == 6 + for expected, actual in zip(default, hoisted): + if np.issubdtype(expected.dtype, np.floating): + assert np.array_equal(expected, actual, equal_nan=True) + else: + assert np.array_equal(expected, actual) + + +def test__scipy_delaunay_nn__fixed_seed_regression(): + """Guard the NumPy Sibson path against silent drift. + + The integer connectivity (mappings, sizes, cavity sizes) is exact, so it is + hashed rather than inlined; the floating weights are pinned as a partition + of unity plus one stored data row and one stored split row. Values were + computed on PyAutoArray ``main`` before the issue #532 JAX changes, which + do not touch this path. + """ + rng = np.random.default_rng(24) + points = rng.uniform(-1.0, 1.0, size=(60, 2)) + query = rng.uniform(-0.7, 0.7, size=(40, 2)) + + ( + _, + _, + mappings, + sizes, + weights, + split_points, + splitted_mappings, + splitted_sizes, + splitted_weights, + cavity_sizes, + overflow, + degenerate, + split_cavity_sizes, + split_overflow, + split_degenerate, + ) = scipy_delaunay_nn(points, query, areas_factor=0.5) + + def integer_digest(*arrays): + hasher = hashlib.sha256() + for array in arrays: + hasher.update(np.ascontiguousarray(array, dtype=np.int64).tobytes()) + return hasher.hexdigest()[:16] + + assert split_points.shape == (4 * points.shape[0], 2) + assert integer_digest(mappings, sizes, cavity_sizes) == "b827766a6308f8a6" + assert ( + integer_digest(splitted_mappings, splitted_sizes, split_cavity_sizes) + == "09733405e7271ecd" + ) + + assert not overflow.any() + assert not degenerate.any() + assert not split_overflow.any() + assert not split_degenerate.any() + + np.testing.assert_allclose(weights.sum(axis=1), 1.0, atol=1.0e-12) + np.testing.assert_allclose(splitted_weights.sum(axis=1), 1.0, atol=1.0e-12) + + assert mappings[7][mappings[7] >= 0].tolist() == [9, 18, 20, 24, 34, 45] + np.testing.assert_allclose( + weights[7][:6], + [ + 0.1945365763239002, + 0.37915688518183144, + 0.1014855420340886, + 0.004891861361516876, + 0.18448730058462717, + 0.13544183451403582, + ], + atol=1.0e-14, + ) + assert splitted_mappings[13][splitted_mappings[13] >= 0].tolist() == [3, 6, 13, 35] + np.testing.assert_allclose( + splitted_weights[13][:4], + [ + 0.5667287604913306, + 0.0769230792268042, + 0.25906145975374095, + 0.0972867005281243, + ], + atol=1.0e-14, + ) + + +def test__env_override_parsers__accept_valid_and_reject_invalid(): + """``PYAUTO_SIBSON_QUERY_CHUNK`` and ``PYAUTO_SIBSON_UNROLL_CANDIDATES`` + are read once at import, so the parsing itself is what the unit tests can + reach; a bad value must fail loudly rather than fall back to the default.""" + assert _positive_int_env("CHUNK", None) is None + assert _positive_int_env("CHUNK", "64") == 64 + assert _positive_int_env("CHUNK", "1024") == 1024 + for bad in ("0", "-1", "notanint", "2.5", ""): + with pytest.raises(ValueError): + _positive_int_env("CHUNK", bad) + + assert _bool_env("UNROLL", None) is None + assert _bool_env("UNROLL", "1") is True + assert _bool_env("UNROLL", "0") is False + for bad in ("", "true", "yes", "2", "-1"): + with pytest.raises(ValueError): + _bool_env("UNROLL", bad) + + +def test__unroll_gate__module_override_wins_and_needs_no_backend(monkeypatch): + """With the override set, the gate answers without importing JAX -- the + unit suite must not pull in a backend. Both settings are bit-identical + code paths; only their emitted program differs (issue #532).""" + monkeypatch.setattr(sibson, "SIBSON_UNROLL_CANDIDATES", True) + assert _sibson_unroll_candidates() is True + + monkeypatch.setattr(sibson, "SIBSON_UNROLL_CANDIDATES", False) + assert _sibson_unroll_candidates() is False