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]

Is it worthwhile to acknowledge a bit of subtlety related to the raw beam grids here? For production runs, the raw beam grids typically have $N_{\rm pix} \sim 65000$ for degree sampling, so for smaller chunk sizes this can be a dominant term, but for larger chunks this contribution might not be a very large part. (This is a genuine question—I'm not sure whether it's worthwhile to discuss this here.)

[steven-murray]

Yes, worth it — added a paragraph to "Memory and chunking": the raw beam-grid term doesn't scale with chunk size, so it can dominate for small chunks and becomes negligible once chunk-scaled terms take over (72014fc).

[r-pascua]

I'm not seeing any discussion of the profiling script on the docs except for this new addition. At the very least, it's probably worthwhile to add some reference about it that explains the parameters. Even better would be a reference page on the CLI utilities.

[steven-murray]

Went with the "even better" option — added docs/cli.rst, a sphinx-click-generated reference for every matvis profile/hera-profile parameter (self-maintaining, since it reads the actual --help text), plus a real annotated JSON output example. Linked from here, the toctree, and profiling/README.md (0d6e9b7).

[r-pascua]

This feels very AI-coded and can probably be improved. I'm assuming the derived object (?) that is referenced several times is an entry in the output json, but I'm not about to run the profiler myself just to figure out what it's talking about. A bare minimum improvement to this section would be an example showing the contents of the json file produced via matvis profile ....

In line 206 ("``derived.gpu_time_per_integration`` ..."), I would replace "total times the" with "total multiplied by the". Using "times" to mean multiplication in a section talking about time is confusing.

[steven-murray]

Fixed the "times"/multiplication ambiguity (7048c43), and added an annotated summary-stats-*.json example on the new docs/cli.rst page rather than duplicating it inline here (0d6e9b7). I also added some extra text to the top of the performance.rst to make it clear where the CLI docs are.

[r-pascua]

"speed of light" micro-benchmarks? I see this phrase used again in the profiling README, and it feels weird. Does the "speed of light" modifier really help here? It feels like fluff that can potentially be confusing to readers (unless this is actually meaningful to the benchmarking community).

[steven-murray]

Apparently Claude agreed with you and suggested "roofline micro-benchmarks" instead... which sounds just as jargony to me, but it says is the standard terminology. I changed it to roofline but also added a parenthesis to explain what that means for us non-gpu-experts.

[r-pascua]

Update the issue reference to a link to the issue if possible.

[steven-murray]

fixed (72014fc)

[r-pascua]

"y covers planes" means what? It looks like y is just parallelizing over beam components, while x is parallelizing over sources.

[steven-murray]

Yup, reworded to "x indexes sources, y indexes (beam, feed, axis) combinations" (1ed75e6).

[r-pascua]

Same comment as before about C code as raw strings living in a python module.

[steven-murray]

And same fix :-)

[r-pascua]

Might be helpful to human readers to add a comment above this block indicating that this is where out-of-bounds points are assigned the beam values at the input boundaries.

[steven-murray]

Added, close to your wording (1ed75e6).

[r-pascua]

Is this comment all that helpful? It's more-or-less repeating stuff said earlier on when the bilinear interpolation routine was declared. Probably more helpful to say something like "Use the custom beam interpolation kernel. If provided a power beam and a complex output buffer, cast interpolated beam to complex on copy.".

[steven-murray]

Used your suggested wording exactly (fc5b8da).

[r-pascua]

Another comment about a seemingly arbitrary choice for block size. Was there some optimization done that resulted in this choice?

[steven-murray]

Same answer: conventional warp-multiple default (128 here), not tuned. Commented accordingly (fc5b8da).

[r-pascua]

Marking another section of C code declared in a python module.

[steven-murray]

Same action taken.

[r-pascua]

Block size setting again

[steven-murray]

Same comment again :-)

[r-pascua]

Is this actually a helpful comment?

[steven-murray]

As written, not really — it assumed you already knew why. Expanded it to be self-contained: nbeampix is pre-summed across all beams by get_desired_chunks, so multiplying by nbeam here would double-count beam memory. (This was a real bug earlier in this PR, which is why the comment exists at all) (1ed75e6)

[r-pascua]

I think the phrase "is indicative only" is throwing me off, perhaps also because I'm not quite sure what the "line-profiler stage table" is. Is it referring to the stats that are dumped to the json file? Is the correct interpretation of this something along the lines of "Because the GPU loop is asynchronous, the reported GPU times do not correspond to actual wall times"?

[steven-murray]

Yeah, this was not well worded. Reworded to say that the stages table times Python lines, but GPU work is queued asynchronously, so a line's measured time is often how long the host waited for already-queued GPU work to finish, not the cost of that line itself (1ed75e6).

[r-pascua]

Having this max(...) catch here feels weird to me? I suppose it's possible that there can be a scenario where the median wall time is less than the median gpu time, but because of that it feels like reporting the median isn't the right thing to do? Why are we not just reporting mean and standard deviation? Since the profiling already includes a dry run to remove one-off setup costs, do we actually expect there to be outliers that would contaminate the means?

[steven-murray]

About the max(...): gpu_time_per_integration is median(per-chunk-total) × nchunks, a biased estimator of the true per-integration GPU total when chunks vary in cost (e.g. the horizon cut leaves different source counts active in different chunks), while wall time sums the actual per-integration total — they can cross without either being wrong, hence the max(). Added a comment explaining that.

On mean vs. median: yes, we do see outliers even after warmup — the V100 dev-run's tau stage showed median=0.44ms vs mean=1.61ms from a single-sample outlier in only 8 warmed-up samples (see the Performance docs page). Added std alongside median/mean in the JSON so we're not limited to picking one (1ed75e6).

[r-pascua]

Changing "the" to "one" is weird because right now it is still the only other alternative? Is the name change from VectorLoop to VectorDot that helpful? Isn't this also an API change?

[steven-murray]

The change from VectorLoop to VectorDot happened long ago -- this was just an incorrect vestige that Claude found and fixed. You're right about "the"/"one" — reverted that back to "the other option" since there is only one (1ed75e6).

[r-pascua]

I didn't catch this earlier, but this will silently do the wrong thing when interpolating a cross-polarized power beam (I believe also for the co-pols in the case where the power beam is constructed from an e-field beam with calc_cross_pols=True). In the mindset of "explicit is better than implicit", shouldn't the e-field or power behavior be controlled with an argument that is passed to the function instead of an inference of the beam type based on the beam dtype?

[steven-murray]

Traced this through carefully. It can't currently happen via the real pipeline: core/beams.py's _wrangle_beams requires beam_type == "efield" whenever polarized=True (power beams never reach the interpolator in that case), and prepare_beam_unpolarized calls as_power_beam(include_cross_pols=False, ...) for the unpolarized path — so cross-pol data is excluded before this function ever sees it. That said, your underlying point stands: gpu_beam_interpolation is a public function called directly by tests and roofline.py, so this is a fragile API regardless of whether our calling functions happen to be safe. Added an explicit power_beam: bool | None parameter (default None keeps the current inference for existing callers); GPUBeamInterpolator now passes it explicitly as not self.polarized. The reworked test suite (see the test_beam_interp_gpu.py thread) exercises the new parameter directly (fc5b8da).

[r-pascua]

These tests should instead use pytest fixtures for the input beam used for the tests. Maybe something like an E-field beam and its corresponding power beam sampled sparsely but over the full sphere. It may be a good idea to have a real-valued E-field beam and a complex-valued E-field beam to fully cover the bases.

[steven-murray]

Reworked the whole file around this. New fixtures build on conftest.py's existing uvbeam/uvbeam_unpol (the full test beam already used elsewhere in the suite) rather than inventing synthetic arrays. See the other three threads on this file for what replaced the old tests (d8b2d05).

[r-pascua]

A better test here is probably to check that the interpolated beam with ndimage.map_coordinates gives the same thing as interpolating the UVBeam object directly with the az_za_map_coordinates setting for the spatial interpolation (with the appropriate corresponding set of spline_opts)?

[steven-murray]

Done — test_order_gt_1_matches_uvbeam_interp now cross-checks gpu_beam_interpolation(..., order=2) against UVBeam.interp(interpolation_function="az_za_map_coordinates", spline_opts={"order": 2}), evaluated at non-node points (offset from the native grid) rather than at grid nodes — the old test's node-only check was a tautology any interpolator passes regardless of order, since it just reproduces its own input (d8b2d05).

[r-pascua]

Also not a fan of this test. I think for something that's testing "Does this take the square root when it's supposed to?", a better idea would be along the following lines:

• Provide the test function with a real-valued e-field beam, its associated power beam, and a parametrized efield_or_power argument (or something similar) that says whether to do the interpolation on the e-field beam or the power beam
• For both cases, do a no-op interpolation (i.e., evaluate the beam at the coordinates it already samples) and check that the returned beam is the same as the e-field beam

[steven-murray]

Implemented essentially as you described — test_noop_interpolation_matches_input, parametrized over efield_or_power, evaluates each beam at exactly its own grid nodes and checks the output reproduces the input (with sqrt applied for the power case) (d8b2d05).

[r-pascua]

This should probably be extended to test all pairs of input beam dtype and output beam dtype. Same idea as with earlier tests where the output should be compared to the provided e-field beam (but now cast to the same dtype as the output).

[steven-murray]

The same no-op test above is also parametrized over out_dtype (complex64/complex128/None), covering the input × output dtype matrix for both the power and e-field cases (d8b2d05).

[r-pascua]

Kind of weird deleting the old test_cublas.py file and making a new test_cublas_gpu.py file when the cublas tests are specifically for the GPU (i.e., the _gpu extension doesn't tell us anything new).

[steven-murray]

You were right that the filename itself carries no information — the actual reason was CI plumbing, not a naming preference: the self-hosted GPU job selected tests via pytest -k "gpu", a substring match against the filename, so a file without "gpu" in its name silently got zero GPU-runner coverage (which is exactly how this file lost coverage earlier in this PR). Fixed the root cause instead of just explaining it: registered a real pytest.mark.gpu marker, marked every GPU-only test explicitly, switched CI to pytest -m gpu, and renamed this back to test_cublas.py now that the filename doesn't need to encode "gpu" for CI selection. Verified the marker-based selection reproduces the old filename-based selection exactly (diffed both collection outputs, zero difference) (d8b2d05).

[steven-murray]

@r-pascua thanks VERY much for the detailed review. I think all the questions have been answered. To be honest, I set Claude on your review, and advised it to refactor the inline CUDA code into external files, as well as to fix the way that GPU tests were collected, and to focus on humanizing the documentation.

It did OK at humanizing but not perfectly. I performed a last gloss of the performance.rst myself and removed some stuff that seemed obnoxious.

My reply comments to you were generated by Claude but I read each one individually and amended them if they were too annoying, and some I replaced wholesale with my own comments.

Thanks again -- your comments have definitely made this PR much better already.