perf: JAX PointSolver step-0 containment without the (N,3,2) triangle gather

[Jammy2211]

Overview

Phase 4b of the point-source CPU speed-up epic (point-source-cpu-speed). Phase 4a (autolens_profiling#314, PR #321) measured step-0 containment at 1.21 ms of a 1.82 ms single-source solved likelihood on a RAL Xeon 8490H, of which ~0.90 ms is materialising the (23283, 3, 2) triangle array in ArrayTriangles.triangles. This task computes step-0 containment on the static lattice without that gather, keeping the result bit-identical (kept indices, image sets, fiducial log L 7.743201200876812, grad, vmap). No geometry, default or completeness change.

Plan

• Prototype three candidates behind a private switch and measure before choosing: (A) structured strided slices of the traced static vertex table, (B) per-component 1-D gathers, (C) drop only the no-op NaN pad/where.
• Plumb the chosen route: CoordinateArrayTriangles.with_vertices marks the returned ArrayTriangles (static pytree aux data); containing_indices branches on the marker; refinement steps 1-7 keep the existing path; ArrayTriangles.triangles unchanged.
• Tests: bit-identity fuzz vs the gather path (jit + vmap, several geometries), refinement-path-unchanged, and an HLO guard that the compiled step-0 containment has no f64[23283,3,2] gather (red on main, green on branch). PyAutoLens suite incl. the static-lattice tie test unchanged.
• Measure with the phase-3 interleaved A/B protocol in solver_config_sweep.py (new branch feature/point-source-cpu-p4b in autolens_profiling), RAL CPU 8490H + A100 no-regression row.
• Ship library-first: PyAutoArray PR, PyAutoLens only if changed, then the autolens_profiling data PR. Stop rule: if no candidate beats control by >= 1.3x on RAL with bit-identity, ship the data PR only as a documented no-go.

Detailed implementation plan

What the code does today (traced)

• ArrayTriangles.triangles (array/PyAutoArray/autoarray/structures/triangles/array.py:121-136) builds the triangle array. It pads the indices, gathers self.vertices[safe_indices] → (N,3,2), then applies a where against NaN. At step 0 no index is −1, so that where does nothing, but it is still traced.
• containing_indices (array.py:147-168) is its only consumer: shape.mask(self.triangles) → jnp.where(inside, size=15, fill=-1).
• Point.mask (shape.py:164-185 → _barycentric_contains :122-139) uses only six (N,) component vectors: a0 a1 b0 b1 c0 c1.
• The step-0 index map is closed-form, not scrambled.
  • static_vertex_table (coordinate_array.py:31-103) orders the vertices by lexicographic integer key.
  • Every key row has the same width W (59 for ±9.9/0.2).
  • vid = (ky-ky_min)*W + (kx-kx_min(ky))//4, where kx_min alternates with the parity of ky.
  • Triangles are row-major, with up and down interleaved by (cy+cx) parity.
• Refinement steps 1–7 use the same containing_indices, but on derived lattices with no vertex_table, where the indices are arange(3N). They must stay untouched.

Approach

When the triangle object came from a step-0 vertex_table, compute the barycentric test on the six component vectors taken by slicing and reshaping the traced vertex table, not by the general (N,3,2) gather. The values in are identical, so the booleans out are identical. Everything else keeps the current path.

1. Worktrees / claims: start_library for PyAutoArray (primary). Add PyAutoLens only if the solver needs a change; the expectation is none, since with_vertices already runs through PyAutoArray. autolens_profiling gets a data PR afterwards on a new branch feature/point-source-cpu-p4b. Register a new issue on PyAutoArray.

2. Prototype three candidates behind a private switch, measured before choosing:

   • (A) Structured slices. At static_vertex_table build time (NumPy, cached), precompute a small layout: W, the parity row offsets, and per-corner slice or stride descriptors. Then (V,2) → reshape (rows, W, 2), take strided slices per corner for the up/down classes, and evaluate _barycentric_contains on the lattice-shaped arrays. Reorder only the boolean mask back to triangle order (row-major, parity-interleaved) before jnp.where.
   • (B) Per-component 1-D gathers. vertices[:,0][idx[:,k]] for k in 0..2: six (N,) gathers, no (N,3,2) array, no NaN where. This is the simplest option.
   • (C) Drop the no-op NaN where/pad when a table is present, keeping the gather. This is the control-minus-overhead floor.

   Pick the fastest candidate that passes every gate. Prefer (B) if it is within ~10 % of (A), because it is simpler.

3. Plumb it:

   • CoordinateArrayTriangles.with_vertices passes the layout, or a structured=True marker, into the ArrayTriangles it returns.
   • ArrayTriangles.containing_indices branches on that marker. The attribute must be static in the pytree (aux data), not a traced child, so jit and vmap see a Python constant.
   • for_indexes, neighborhood and up_sample do not propagate the marker, matching how vertex_table is dropped today.
   • Keep ArrayTriangles.triangles itself unchanged for any other caller.

4. Tests (PyAutoArray, alongside test_coordinate_jax.py):

   • Bit-identity fuzz. Structured vs gather containment on about 10⁵ points: random, exactly on vertices, edge midpoints and centroids, plus the NaN-padded case. Compare kept index arrays with array_equal, under both jit and vmap.
   • HLO guard. The compiled step-0 containment has no gather producing f64[23283,3,2]. Copy the phase-2 no-sort guard pattern. The guard must fail on main.
   • Refinement untouched. A derived lattice takes the old path.
   • PyAutoLens. Run the whole suite, including test_static_lattice_jax.py and its tie test test__source_on_a_step_0_vertex_returns_the_two_true_images, unchanged.

5. Measure (autolens_profiling, scripts/point_source_image/likelihood_breakdown/solver_config_sweep.py):

   • Add a library route beside control.
   • Use the phase-3 interleaved A/B protocol: fresh closures + jax.clear_caches(), 20×20, bootstrap 90 % CI.
   • Record the step0_split before/after, the 200 prior + 200 stress completeness draws against the control (they must be identical), vmap 1/4/16, compile time (≤ +20 %) and XLA memory.
   • Run on RAL CPU pinned to an 8490H node, plus an A100 no-regression row, using the phase-4a submits.
   • The library route runs from branch clones on RAL; the mirror is not touched while Euclid arrays run. Assert source_revisions in the job.

6. Ship library-first.

   • PyAutoArray PR, then PyAutoLens (only if changed), then the autolens_profiling data PR with the campaign-note "Phase 4b" section.
   • The human runs /prm on each.
   • Release-gated follow-up: after the next release, the library route should read library_matches.

Execution

This Opus session plans and judges; the prototype, tests, RAL runs and ship run in model: "opus" subagents with progress files + Monitors. Nothing is armed beyond the turn. Stop rule: if no candidate beats the control by ≥ 1.3× on RAL while keeping bit-identity, ship the data PR only, as a documented no-go, with no library change.

Verification

• Kept indices match bit for bit on the fuzz set and on the 400 completeness draws.
• The fiducial 7.743201200876812 holds bit-exactly.
• grad and vmap equal control.
• The HLO guard is red on main and green on the branch.
• The PyAutoArray and PyAutoLens suites pass.
• autolens_workspace_test point-source jax_likelihood ×4 + jax_grad are identical to main.
• The RAL CPU speed-up has a 90 % CI; the A100 shows ≈1×, with no regression.

Affected Repositories

• PyAutoArray (primary)
• autolens_profiling (measurement / data PR)
• PyAutoLens (only if strictly needed; expected none)

Branch Survey

Repository	Current Branch	Dirty?
./PyAutoArray	main	clean
./PyAutoLens	main	clean
./autolens_profiling	main	1 untracked/modified path (not this task)

Suggested branch: feature/pointsolver-step0-gather (PyAutoArray); feature/point-source-cpu-p4b (autolens_profiling, based on feature/point-source-cpu-p4 until #321 merges)

Key Files

• autoarray/structures/triangles/array.py — ArrayTriangles.triangles, containing_indices
• autoarray/structures/triangles/coordinate_array.py — static_vertex_table, with_vertices
• autoarray/structures/triangles/shape.py — Point.mask, _barycentric_contains
• test_autoarray/structures/triangles/test_coordinate_jax.py — tests

Original Prompt

Prompt: https://github.com/PyAutoLabs/PyAutoMind/blob/main/active/pointsolver_step0_gather_containment.md

Click to expand starting prompt

Point-source CPU speed-up phase 4b — cut the step-0 triangle gather / containment in the JAX PointSolver

Type: feature
Target: autoarray
Repos:

• PyAutoArray
• PyAutoLens
• autolens_profiling
  Themes:
• point-source
• profiling
• jax-compile
  Difficulty: medium
  Autonomy: supervised
  Priority: high
  Status: formalised
  Consequence: judge
  Review-minutes: 20
  Unattended: ready
  Epic: point-source-cpu-speed
  Filed: 2026-09-26
  Parent: active/pointsolver_cpu_speed_phase_4.md (issue autolens_profiling#314)

Goal

Remove, or sharply cut, the ≈ 0.9 ms vertices[indices] materialisation at step 0 of the JAX
PointSolver. On the released code it is ≈ 49 % of the single-source likelihood. This is a pure
code lever. The solver geometry (±9.9″ / 0.2″ / 1e-3, MAX_CONTAINING_SIZE 15) does not
change, and the result must stay bit-identical:

• image sets and counts;
• the fiducial simple solved log L 7.743201200876812;
• jax.grad;
• vmap.

Human decision (2026-09-26): this is phase 4b, first of the phase-4a follow-ups.

Evidence (phase 4a, RAL CPU, Xeon 8490H euclid-ral-compute-10-4, 8 CPUs, fp64, JAX 0.10.2)

Ledger: lens/autolens_profiling/results/notes/point_source_cpu_campaign.md, section "Phase 4a".

• Re-baseline, job 356365 (results/breakdown/point_source_image/image_plane_hpc_ral_cpu_fp64_p4.json,
  5 runs): fused solved median 2.095 ms. Step 0 is 66 % of the call. Refinement steps 1–7 are 15 %,
  β* 8 %, magnification 7 % and χ² 4 %. The phase-3 FLOP estimate (refinement ≈ 60 %) was wrong in
  wall time. That cell's own step-0 "ray trace vs containment" split counts the gather as trace, so
  do not quote it.
• Sweep, job 356367 (results/breakdown/point_source_image/solver_config_sweep_hpc_ral_cpu_fp64.json,
  key step0_split.control): the control is 1.824 ms and step 0 is 1.48 ms (81 %). It divides into:
  • ray trace of the 11 859 static-lattice vertices: 0.27 ms;
  • containment: 1.21 ms (66 % of the likelihood). This is containing_indices: the gather +
    Point.mask + jnp.where.
  • Of the containment, the gather that materialises the (23 283, 3, 2) triangle array alone is
    ≈ 0.90 ms (triangle_materialisation_ms).
• Under ±2.5″/0.4 (243 rows), containment is 0.06 ms, which shows how much of the cost is the
  23 283-triangle size of the lattice.
• Why this lever over shrinking the grid: it recovers most of the extent/scale speed-up (the
  best admissible config, ±2.5/0.4, was 2.37× and 5.55× at vmap-16) with no completeness risk and no
  default change. The extent became a workspace choice, per draft/feature/autolens/pointsolver_extent_sanity_check.md
  and draft/feature/autolens_workspace/pointsolver_grid_extent_per_package.md.
• Revisions measured: PyAutoArray 3de624b5, PyAutoLens 86054bbc, autolens_profiling 6c45fec.
  The point-source path is byte-identical to 2026.9.26.1.

Candidate mechanisms (choose by measurement)

1. Exploit the regular step-0 lattice. Step 0 is the static phase-3 lattice
   (static_vertex_table, PyAutoArray autoarray/structures/triangles/coordinate_array.py; wired in
   PyAutoLens AbstractSolver._initial_triangles). Each triangle's vertices are fixed offsets into
   the lattice. Compute containment by row/column arithmetic or strided slicing of the traced
   vertex table instead of the fancy-index gather. For example, evaluate the barycentric / sign test
   on the up- and down-triangles of each lattice row as two dense slices.
2. A fused containment kernel. Do the sign test on the index arrays directly, so XLA never
   materialises the (N, 3, 2) array. Check the optimised HLO for the gather's disappearance.
3. Anything else that removes the materialisation while keeping the kept set bit-identical.
   Beware: phase 3's tie study showed that a changed float path at step-0 vertices moves the kept set.
   Pin the tie test test_static_lattice_jax.py::test__source_on_a_step_0_vertex_returns_the_two_true_images.

Protocol (as phase 3)

• A red control on main, then an interleaved A/B: fresh closures + jax.clear_caches(),
  20 rounds × 20 calls, rotated round-robin, bootstrap 90 % CI.
• The harness is lens/autolens_profiling/scripts/point_source_image/likelihood_breakdown/solver_config_sweep.py.
  Its step0_split prefixes (source centre / jnp.sum(plane.vertices) / materialised triangles /
  containing_indices) are the before/after instrument, and its 200 prior + 200 stress completeness
  draws are the regression set. Add a library route beside control.
• Gates: bit-identical log L on the stream and on the fiducial, positions, image counts, grad, and
  vmap 1 / 4 / 16. Also compile no worse than +20 %, and no memory regression.
• RAL CPU pinned to an 8490H node (check sinfo -p ral -N -o "%N %T %C"; idle* nodes are
  unreachable), plus an A100 no-regression row. The A100 is launch-bound (phase 3), so expect ~1×.
• Library-first ship: PyAutoArray → PyAutoLens → autolens_profiling data PR.

Traps

• PyAutoArray is currently claimed by task interferometer-transform-real-scatter. Check the
  claim at start_dev (run worktree_check_conflict, not a grep). A parallel claim is fine only if
  the file sets are disjoint.
• JAX caches jaxprs on function identity. A monkeypatch A/B needs a distinct function object and
  jax.clear_caches() before each compile.
• jax.grad through an AnalysisPoint needs autofit.jax.register_model(model); without it the
  gradient is silently all-zero.
• Fix, or at least do not quote, image_plane.py's step-0 prefix split. It counts the gather as ray
  trace.

[Jammy2211]

Checkpoint 2026-09-26 (session ended by the human; resumes tomorrow)

Done:

• Registered in Mind (pointsolver-step0-gather). Worktrees ~/Code/PyAutoLabs-wt/pointsolver-step0-gather/{PyAutoArray (feature/pointsolver-step0-gather off main 14d63360), autolens_profiling (feature/point-source-cpu-p4b off feature/point-source-cpu-p4 d2bd1ef)}.
• Local WIP commit 8755072f, not pushed. It adds a private switch array._STEP0_CONTAINMENT with four routes:
  • gather: the old path.
  • nopad: the (N,3,2) gather without the pad or NaN where.
  • components: six 1-D gathers.
  • structured: contiguous 2-D slices of the vertex table viewed as (n_pairs, W0+W1, 2) row pairs. The four (row, col) parity classes are interleaved back to row-major on the boolean mask, and the table is padded by one row when the number of key rows is odd.
• Layout details:
  • static_lattice_layout is lru-cached and made only of ints (Step0Layout, carried as pytree aux data). It is checked element-wise against the index map and falls back to grid=None.
  • It handles alternating key-row widths, which occur when the number of lattice columns is even.
  • It is plumbed through for_limits_and_scale(static_vertices=True) → with_vertices → containing_indices, and only for Point. Derived lattices drop it.
• Scratch check under jit+vmap: all 4 routes give identical kept indices to gather on ±9.9/0.2 and on two asymmetric geometries with alternating row widths, over 400 random, vertex and centroid points each. The formal tests are not written yet.

Next step to resume (plan steps 4 + laptop measurement):

1. Tests in test_autoarray/structures/triangles/test_coordinate_jax.py:
   • a bit-identity fuzz of about 1e5 points per route (edges and midpoints included), jit + vmap, over 3 or more geometries;
   • a refinement-path-unchanged test;
   • an HLO guard for no f64[23283,3,2] gather. It must pass the vertices as a traced arg, since a closure constant-folds. Show it red on main.
2. Full PyAutoArray and PyAutoLens suites against the branch.
3. Add --step0-route to solver_config_sweep.py in the p4b worktree and take the laptop A/B table.
4. Pick the default route. It is currently "structured", provisionally.

[Jammy2211]

Library PR Created: #580 (2026-09-27)

• Default route: structured, human-approved. PyAutoLens needs no PR: its full suite passes unchanged against the branch.
• Bit-identity: every route matches the old path on the laptop, RAL EPYC 7702 CPU and RAL A100: log L, positions, image counts, grad and vmap. Completeness is 200/200 prior and 200/200 stress.
• Speed: structured runs 1.62× gather on EPYC (3.184 → 1.971 ms) and 1.87× at vmap-16. The A100 shows no change (1.00×). Compile is +5 % and temporary memory −47 %.
• HLO guard: fails on main and passes on the new routes. The PR also fixes a WIP bug where Point subclasses were caught by an isinstance check.
• Heart: YELLOW (manifest drift, stale release rehearsal), acknowledged by the human.

Held (profiling half):

• The quotable Xeon 8490H job 357321 is pending on a saturated ral partition.
• autolens_profiling feature/point-source-cpu-p4b holds local, unpushed commits: the harness --step0-route option, the instrument fixes, the laptop/EPYC/A100 results and the ledger Phase 4b section.
• Resume: pull the 8490H JSON → update the ledger → make the A100 fiducial gate device-aware → rebase onto main after fix(inversion): make AbstractMeshGeometry picklable (xp module → _use_jax bool + property) #321 → ship_workspace the data PR.

[Jammy2211]

Smoke Test Results — 2026-09-27 (targeted: point solver)

PyAutoArray feature/pointsolver-step0-gather @ c13b2d7; other libraries on canonical main. Import proof: autoarray.__file__ = PyAutoLabs-wt/pointsolver-step0-gather/PyAutoArray/autoarray/__init__.py.

Workspace	Passed	Failed	Total
autolens_workspace	19	0	19
autolens_workspace_test	9	0	9
HowToLens	3	0	3

Scope

• autolens_workspace: point_source simulators, modeling, fit, plot and features (fluxes, time_delays, deblending, extra_galaxies, scaling_relation); multi_dataset imaging_and_point_source; guides/point_source_pairing; the smoke-listed cluster/lenstool data and modeling; group/simulator. The smoke profile sets PYAUTO_DISABLE_JAX=1 here.
• autolens_workspace_test (JAX): point_source/jax_likelihood/{point, image_plane, source_plane, fluxes_time_delays} (point.py and image_plane.py assert the exact literal log likelihood), visualization x3, cluster/likelihood_sanity, jax_grad/gradient.
• HowToLens: chapter_1 tutorial_4_point_sources, simulator/no_lens_light, simulator/group.
• Excluded: no_run.yaml entries, point_source/start_here (disabled, PyAutoFit#1179), and the unsmokeable cluster scripts.
• Run with the Hands run_python.py smoke profile in the cached Heart envs, on scratch worktrees of each workspace main. pyauto-heart smoke --prepare-only fails locally: smoke_install.sh uses the flat ./PyAutoNerves layout.

[Jammy2211]

PyAutoArray PR #580 has been merged as 4383ea8 via /prm. CI passed on every leg, the Heart freeze was clear, and the targeted smoke passed 31/31. The issue stays open for the profiling half: the queued 8490H job 357321, then the autolens_profiling data PR from feature/point-source-cpu-p4b after #321.

[Jammy2211]

The quotable RAL Xeon 8490H row is in (job 357321, euclid-ral-compute-10-4, all_gates_pass true). It's committed locally as autolens_profiling f1b4930.

route	median ms	× gather	vmap-16 ms/L	containment ms	temp KB
gather	2.260	1.00	1.931	1.223	1692
components	1.550	1.46	1.108	0.433	1506
structured (default)	1.570	1.44	0.851	0.494	902

• The node was loaded, with loadavg around 200 on 236 cores (phase 4a ran at about 1), so the absolute ms aren't comparable to phase 4a. The interleaved ratios stand.
• Containment falls by 60 %. structured and components are tied on the scalar call, and structured is best at vmap-16 (2.27× gather).
• Compile is +8.6 %, inside the +20 % gate.

Next: the data PR ships after autolens_profiling #321 merges (p4b is stacked on it): rebase, then /ship_workspace, then /prm.

[Jammy2211]

Shipped — 2026-09-27

PRs

• research: point-source CPU phase 4b — step-0 containment route A/B (PyAutoArray#579) autolens_profiling#330 (merged as 81af10f7)

Upstream PR

• perf(triangles): step-0 Point containment without the (N,3,2) gather (#579) #580 (merged as 4383ea8; labelled pending-release, not yet in a published release)

Summary

Four JAX PointSolver step-0 containment routes now sit behind array._STEP0_CONTAINMENT; the default is structured, which drops the (N,3,2) triangle gather. All four routes are bit-identical on every route tested. On the quotable RAL Xeon 8490H row (job 357321), structured makes the gather 1.44x faster on a single call and 2.27x at vmap-16. Containment time drops 60 %, compile time rises 8.6 % and temp memory falls 47 %. That node was loaded (load ~200). Supplementary rows: EPYC 1.62x, A100 1.00x (no change), laptop 1.38x. The PR also fixes the WIP bug (type(shape) is Point).

Session Notes

• Tests: a bit-identity fuzz, refinement tests and an HLO guard, each red on main before the fix. Suites: PyAutoArray 1732 passed, PyAutoLens 758 passed + 1 xfail. Targeted smoke 31/31.
• autolens_profiling: instrument fixes and a backend-aware fiducial. The WALL-BASIS rows were fixed after the first lint run went red.
• Heart was YELLOW twice (manifest drift, then a stale/incomplete rehearsal). The human acknowledged it both times.
• Follow-ups are filed in PyAutoMind: nopad is dead weight in _STEP0_CONTAINMENT and can go in a later library cleanup; the RAL scratch copies under /mnt/ral/jnightin/autolens_profiling_wt/ are left in place; a quiet-node 8490H re-run is still needed if absolute ms against phase 4a are ever wanted.