perf: overhaul GPU hot path (7.7x at HERA scale) + revive parity tests

[steven-murray]

Description

This PR overhauls the GPU hot path for large simulations. At the production-relevant scale (350 antennas, 350 unique beams, polarized, single precision, gridded beams), per-chunk GPU time drops 505 ms → 65 ms (7.7×) and steady-state wall time per integration drops 15.2 s → 2.06 s, with GPU utilization rising from ~35% to ~95%. It also revives the CPU/GPU parity test suite (silently skipped for years), fixes single-precision GPU support with gridded beams, adds an fp32-vs-fp64 validation test, adds a reusable benchmarking/profiling kit, and adds user-facing performance documentation.

Benchmarks below were measured on an RTX A2000 laptop GPU (4 GB, Ampere) with per-chunk CUDA events at the "production-slice" config: nant=nbeam=350, polarized, gridded beams, 10⁶ sources in 30 chunks, fp32, ERFA coordinates. Reproduce with profiling/run-canonical.sh.

Headline numbers (per-chunk CUDA-event timing)

Stage	Before (ms)	After (ms)	Speedup
beam interp	224	9.9	22×
tau	23	2.1	11×
Z	13	5.7	2.3×
matprod	124	46.8	2.7×
chunk total	505	65.3	7.7×

After the change, ~70% of GPU time is the cuBLAS cherk kernel itself, i.e. the pipeline sits at the library roofline (verified against a bare-cuBLAS benchmark of the same shape). Larger relative gains are expected on data-centre GPUs (V100/A100), since the removed host-side overheads don't shrink with faster hardware.

How the baseline was diagnosed

nsys tracing at the production-slice shape (not toy scale — at small sizes ~90% of wall time is fixed setup, which is why earlier optimization attempts found nothing) showed:

• Beam interpolation was 90% host overhead: nbeam·nfeed·nax = 1400 separate map_coordinates launches per chunk, plus per-beam coordinate-array allocation.
• The GEMM ran on the default stream (the direct cupy_backends cublas calls never set the handle's stream), serializing against the compute stream.
• A hidden complex128 GEMM per chunk: antpos * 1j promoted the phase-factor matmul to double precision even at precision=1, adding an fp64 GEMM + cast per chunk and a large temporary that caused OOMs.
• The multi-stream design was a no-op: one stream per chunk, but every stage shares one set of buffers across chunks, so real overlap would be a data race — a per-chunk Device().synchronize() in the Z stage was what kept it correct, and it stalled the pipeline.

Changes

1. matprod via cherk/zherk (gpu/_cublas.py): V = Z·Z^H is a Hermitian rank-k update — half the FLOPs of a general GEMM; a tiny kernel mirrors the computed triangle. General products (GPUVectorDot) use cublasCgemm3m (measured 2.1× over cgemm at matvis shapes). Both bound from libcublas via ctypes (cupy doesn't expose them), running on the current cupy stream. Falls back to cgemm/zgemm if the library can't be bound.
2. One fused bilinear beam-interpolation kernel (gpu/beams.py): a single launch over all (beam, feed, axis) planes and sources; map_coordinates remains the fallback for order > 1.
3. Fused Z kernel (gpu/getz.py): Z = A·√I·exptau in one elementwise pass (with the beam_idx gather), replacing 4 broadcast copies + a Python loop over all antennas + a full-device sync per chunk.
4. Single compute stream (gpu/gpu.py): correct with zero device synchronization in the loop; the host queues many chunks ahead. Stages carry NVTX ranges; wall/event stats are exposed via LAST_RUN_STATS for the profiler CLI.
5. tau precision fix (core/tau.py): keep the phase matmul at the requested precision.

Bugs fixed

• tests/test_cpu_vs_gpu.py still guarded on importorskip("pycuda") (removed in v1.3.0), so the main CPU/GPU parity suite had been silently skipped for years. Re-enabled with cupy and extended to fp32 — which immediately exposed the next three bugs:
• fp32 + gridded beams crashed on GPU (beam_data.set() dtype mismatch).
• gpu.py computed its own nsrc_alloc, disagreeing with CoordinateRotation (which ignores source_buffer for chunks ≤ 1000 sources) → shape-mismatch crashes in small chunked runs.
• Stale kx/ky spline options in the parity test (pyuvdata's map_coordinates interpolator takes order).

Validation

• Full test suite passes (193 passed), including the revived parity suite across polarized/unpolarized × analytic/gridded × fp32/fp64 × chunking × source-buffer configurations.
• New tests/test_precision.py: fp32 agrees with fp64 to 1e-5 of the peak visibility on both backends — the accuracy gate for running production in single precision.
• tests/test_cublas.py extended to rectangular shapes, both dtypes, and out=/beta= accumulation (exercises the herk + mirror path).
• Interpolation kernel verified against map_coordinates at fp32 tolerance.

Follow-up work

Filed as issues: #131 (validation on V100/ilifu), #132 (sum_chunks accumulation/pinned-memory transfer), #133 (horizon-cut compaction without a host sync), #134 (hoisting frequency-independent work out of the per-frequency loop), #135 (per-baseline source coarsening).

Checklist

I have

• Added a test covering your new feature adequately?
• Added a docstring (or note to a docstring) describing your feature?
• (Optional): Added a tutorial / section to a tutorial showing usage of your new feature? (new "Performance" docs page)
• Important: this feature does not break API compatibility.

🤖 Generated with Claude Code

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 99.39%. Comparing base (e931e83) to head (e781f13).
⚠️ Report is 17 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #130      +/-   ##
==========================================
+ Coverage   98.88%   99.39%   +0.50%     
==========================================
  Files          22       23       +1     
  Lines         989     1153     +164     
  Branches      103      145      +42     
==========================================
+ Hits          978     1146     +168     
+ Misses          6        4       -2     
+ Partials        5        3       -2     

Flag	Coverage Δ
unittests	88.03% <99.58%> (+9.56%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[r-pascua]

Since this was generated with Claude, do you think it's easy to also report the theoretical minimum GPU time/integration? I'm guessing the "effective R" indirectly reports this for the A2000, but it might be a useful number to have around if someone down the line wants a sense of how much could potentially be gained through further optimizations.

[steven-murray]

The t_gemm formula a few lines down already is that calculation once you plug in your own measured R. I added a sentence making the connection explicit instead of leaving it for the reader to infer (72014fc).

[r-pascua]

I'm not quite sure this answers what I was asking. By "theoretical minimum," I meant "If the GPU was used 100% efficiently (i.e., the realized FLOPs are the same as the GPU's maximum FLOPs according to the device specifications, and the kernels are compute limited instead of bandwidth limited), then how long would the 'canonical run' take?".

[r-pascua]

Can we add a simple comment here explaining why we always run cublas.setPointerMode? (And in other places where this try: finally structure appears)

[r-pascua]

I think this can be phrased a bit more clearly. I'd suggest something like

Whether the provided ``beam`` is in power units or E-field units. If not provided,
then it is inferred based on whether the provided ``beam`` is real- or complex-valued.
Failing to set ``power_beam=True`` and providing a power beam with cross-polarized
components will result in the interpolation routine treating the beam as if it were an
E-field beam instead of a power beam (i.e., no square root will be taken after interpolation).

[r-pascua]

assert

[r-pascua]

Rather than "See ...", can we just use the combine_docstring decorator to copy the parameters over?

[r-pascua]

I know that I brought this up somewhere else, but isn't this an API change? I believe Claude made some argument about this being changed in something that preceded this PR, but... it's changed here?

[r-pascua]

Is there an equivalent way to do this with combine_docstrings?

[r-pascua]

See my comment about the comment in core.coords.CoordinateRotation.select_chunk.

[r-pascua]

Why are we not using an f-string here?

[r-pascua]

Does this comment need to live here? It's a bit long, and it's a copy-paste of what's already in the rtd. Can we replace it with something like # Initialize a dictionary to track timing summary statistics.?

[r-pascua]

Similar comment to above. I'd recommend adding a comment above line 238 that is along the lines of # Estimate the per-integration GPU time based on the median per-chunk GPU time.

[r-pascua]

A comment somewhere explaining the expected shape of d0 would help in understanding whether this operation actually makes sense. Maybe in prepare_for_map_coords?

[r-pascua]

Might be better to assign the spline interpolation order to a variable since it's used in two places? Especially considering that we're planning on implementing a higher order interpolation kernel in the future.

[r-pascua]

Potential tests to add:

• Consistency check between GPU interpolation kernel results and linear order map_coordinates interpolation.
• Check that out-of-bounds interpolation points are assigned the boundary beam values.
• Explicit check that power beams are cast to complex beams of equivalent precision by interpolation routine.

May need to think more carefully about whether this provides complete meaningful coverage.

[r-pascua]

Why doesn't this just do

cb.zdotz(cp.asarray(a), out=out, beta=1.0)
np.testing.assert_allclose(out.get(), expected, rtol=1e-4)

[r-pascua]

Wouldn't it be cleaner to instead do

expected = cb.complex_matmul(cp.asarray(a), cp.asarray(a))
np.testing.assert_equal(c.get(), expected.get())

[r-pascua]

With the current version of things, this will be fine, but there is potential for this test to fail if, for some reason, _load_cublas_ext searches for 2 or fewer libraries. Probably not an actual issue, but could be a confusing test failure in the future if it ends up breaking.

One option is to update the _cublas module to define a _libcublas_sonames list, adjust the conditional on this line to if len(attempted) < max_attempts: (with max_attempts defined in the preceding line), then check that len(attempted) == min(max_attempts, len(cb._libcublas_sonames)) at the end of this test function.

[r-pascua]

Isn't this checking one of the lines that has the # pragma: no cover flag?

[r-pascua]

This seems like it would be a helpful function to expose to other test modules. Can it be easily accessed by other modules if it is placed in the conftest module?

[r-pascua]

Can't this be combined with the previous test through an appropriate parametrization on nbeam? (Although, if updated, it should be done in a more intuitive way via a parameter called all_beams_unique or something.)

[r-pascua]

This explanation doesn't feel very helpful. Also, it looks like this is just testing that nothing is shuffled when no beam indexing is provided? Surely that can be absorbed into the earlier test with a parameter.

[r-pascua]

I think one test case that is missing from this is running a check when 1 < nbeam < nant. While this is probably an unusual use case, I think it is still an important one to cover.

[r-pascua]

Flagging another hardcoded variable that is referenced twice.

[r-pascua]

Do we need this comment?

[r-pascua]

Some parts of this comment feel a little oddly worded. Can't the (~|V| ...) parenthetical just be restated to say "the autocorrelation amplitude"? Also doesn't this set a tolerance that is quite a bit larger than the expected accumulated roundoff error? I'm just guessing here, but it feels like something along the lines of scale = max_source_flux * n_src would be closer to the right thing.

[r-pascua]

Are these instructions to AI agents? Just curious

[r-pascua]

Maybe replace "plane" with "combination"

[r-pascua]

I think this utilization number is one of the things I was looking to see reported in the new "Performance" section of the docs. Right now it is currently hidden in the changelog, and I think it is important enough to be highlighted somewhere earlier in the page.

[r-pascua]

"large hidden temporary that" seems like it's missing a word

[r-pascua]

Also standard deviations

[r-pascua]

OK, I think I have gotten around to all of the responses and completed another look through all of the changes. I think this is very close to being done, just a few more tests I would like to see and some tidying up of comments/docs.