Skip to content

Fix review findings across numerics, distributed, GPU and AD - #55

Merged
subhk merged 13 commits into
mainfrom
fix/review-high-severity
Oct 7, 2026
Merged

subhk merged 13 commits into
mainfrom
fix/review-high-severity

Conversation

@subhk

@subhk subhk commented Oct 5, 2026 •

Copy link
Copy Markdown
Member

Summary

Fixes the findings of a full review of the package, one commit per area. The most serious gave silently wrong results, corrupted memory or crashed long runs:

  • Legendre values went wrong above lmax ≈ 1900 (Float32 point evaluation above lmax ≈ 240): the recurrence seed underflowed to a denormal.
  • A rejected cfg.lmax/mmax/mres assignment still stored the new value, and the next transform wrote out of bounds under @inbounds.
  • Distributed calls leaked MPI communicators and died after about 1,000 calls. Vector/QST batches scrambled any layout but a θ-split one.
  • Rotation gradients were 33–54% off for complex targets, and the GPU rotation cache could return another rotation's blocks.

Commits

Commit Fixes
fix: keep the Legendre recurrence accurate beyond lmax ≈ 1900 exponent tracking in the CPU recurrence, as on the GPU
fix: validate config changes first and make create_config defaults exact setter validation; exact default nlat on non-Gauss grids; 2 GiB cap on automatic tables
fix: sweep Wigner matrices by degree and guard mmax < lmax rotations O(lmax⁴) → O(lmax³) rotations; pencil rotations no longer drop orders above mmax; Real angles
fix: reuse MPI topologies and honour pencil layouts in distributed calls communicator leak; batch and fixed-order layouts; :quad axisymmetric synthesis; Float32 and real inputs
fix: accurate Float32 GPU tables, bounded caches, pinned staging buffers Float32 accuracy (2.8e-4 → 2e-7); weakly owned, bounded caches; cache key; :quad evaluators; gpu_*_safe; point evaluation; pinning on Julia 1.11+
fix: reverse-mode AD on PencilArrays, exact rotation adjoints distributed Zygote/ForwardDiff through the public API; exact rotation adjoint; robust pullbacks; zgrad_* memory
fix!: use the Euclidean gradient convention for vorticity losses consistent gradients across the API (behaviour change)
fix: real FFTs into strided and padded buffers use_rfft=true into views and strided scratch
fix: bounds-check @sht_loop before its @inbounds loop out-of-range writes such as the one in examples/loop_demo.jl
fix: define every export and lift the Julia 1.12 compat cap DistAnalysisPlan/DistPlan defined; precompile warnings; julia = "1.10"
ci: test :quad again, run GPU kernels on the CPU backend, fix benchmark SHTNSKIT_PHI_SCALE hid :quad; new gpu-host-reference job; benchmark used SH_mul_mx wrongly
docs: document where conventions differ from SHTns toroidal sign, odd negative orders in LM_cplx, SH_mul_mx layout

Behaviour changes

Worth a release note, and probably a minor version bump:

  • grad_loss_vorticity_Tlm and loss_and_grad_vorticity_Tlm return m > 0 entries doubled: they now give the Euclidean gradient, like grad_enstrophy_Tlm and Zygote.
  • create_config on :regular, :regular_poles and :driscoll_healy grids without nlat (or with exactly nlat = lmax+2) gives 2lmax+1 latitudes (2(lmax+1) for Driscoll–Healy). Code that hard-codes the old size gets a DimensionMismatch.
  • create_regular_config no longer precomputes Legendre tables above 2 GiB; prepare_plm_tables! still does.
  • Some calls now raise ArgumentError instead of returning wrong data:
    • assigning an mmax the grid can't resolve;
    • l-split spectral batches;
    • order-mixing rotations on mmax < lmax layouts (pencil and GPU);
    • rank-varying cotangents on replicated distributed energies.
  • Float32 input to the 1-D and 2-D distributed analysis returns ComplexF32.
  • A few pullbacks return one tangent per positional argument, where they returned an extra keyword tangent before.
  • Pencils from create_spectral_pencil/create_spatial_pencil share a cached topology per communicator.

subhk and others added 13 commits October 5, 2026 13:16
Plm_norm_row! seeded P̄_m^m ∝ sin^m θ with no exponent scaling. Once the
seed underflowed it stuck at the smallest denormal instead of shrinking,
and the recurrence scaled that wrong value up: at lmax=2047 the
synthesis→analysis round trip was off by up to 2.9, and synthesis_point
of a unit (2500, 1500) mode returned 2.6e9 instead of about 1.66.
Float32 point and latitude evaluation already failed from lmax ≈ 240.

Carry a shared power-of-two exponent for the recurrence pair, as the GPU
table builder does. The lmax=2047 round trip is now accurate to 2.1e-13,
and P̄ matches a BigFloat reference to about 3e-13 up to lmax=4000. The
magnitude is checked once per block of degrees, so transforms slow by 1%
at lmax=255 and 6% at lmax=1023.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Setting cfg.lmax, cfg.mmax or cfg.mres stored the new value before
validating it, so a rejected `cfg.mmax = 30` still left mmax == 30, and
nlon ≥ 2mmax+1 was never checked. Either way the next synthesis wrote
past its Fourier buffer inside @inbounds. The setters now validate the
whole candidate layout before storing anything.

create_config's default nlat = lmax+2 suits Gauss grids only: the
equiangular round trips were 7-22% off. The default now becomes 2lmax+1
(at least 2) on equiangular grids and 2(lmax+1) on Driscoll-Healy. It is
recognised by value, so an explicit nlat = lmax+2 behaves the same;
create_regular_config still builds an lmax+2 grid.

create_regular_config precomputed Legendre tables at any size, 32 GiB at
lmax=1023. Above 2 GiB it now computes rows on the fly after a one-time
note; prepare_plm_tables! still builds the tables.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Every rotation rebuilt each degree's Wigner-d matrix from d⁰, an
O(lmax⁴) cost: SH_Yrotate took 4.1 s at lmax=192. The CPU and pencil
rotations and the GPU host blocks now advance d^l from d^(l-1), which
gives bitwise-identical results in 89 ms at lmax=192.

Pencil Y, X and Euler rotations on mmax < lmax layouts silently dropped
the orders above mmax, 9-27% of the energy. They now raise the serial
engine's error on every rank. The serial check uses the stored Float64
β, so an exact half-turn with Float32 coefficients is no longer
rejected. Distributed rotations also accept any Real angle (π, 1,
1//2), as the serial ones do.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
- Every distributed call built a new Pencil, whose Cartesian
  communicator and subcommunicators only GC finalizers release, so
  analysis on a PencilArray died at call 1,022 on 2 ranks with "too many
  communicators" (analysis_qst_batch at call 29). Per-call pencils now
  share one cached topology per communicator.
- Vector and QST batch transforms copied each field linearly into a
  θ-split pencil whatever the batch layout: 16% wrong on PencilArrays'
  default 3-D layout, and with uneven splits a BoundsError on some ranks
  while the rest waited in a collective. Fields keep the batch's own
  topology, and l-split spectral batches raise ArgumentError on every
  rank.
- Fixed-order vector transforms require a dimension-1, unpermuted split;
  other layouts read past local data on one rank while the others hung.
- Distributed synthesis_axisym(_l) applies the evaluator φ factor, so
  :quad results are no longer 2π too large.
- Dense dist_synthesis accepts real coefficients, the 1-D and 2-D
  distributed analyses keep Float32 precision, and
  dist_scalar_laplacian! checks its destination before transforming.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
- Tables are built from Float64 nodes and rounded once, and the vector
  kernels take sin θ from the exact node. Float32 synthesis error at
  lmax=127 drops from 2.8e-4 to 2e-7.
- Table caches hold their configuration weakly and drop its tables once
  it is collected, rotation blocks are capped at 1 GiB per device, and
  cache signatures no longer build a Tuple of all of Nlm per call.
- The rotation-block cache is keyed on the effective angles: a keyword
  and a setter rotation with the same stored fields shared an entry (up
  to 1.95 relative error). mmax < lmax rotations raise as on the CPU.
- The point and latitude evaluators and synthesis_axisym apply the
  evaluator φ factor; under :quad they were 2π too large.
- gpu_*_safe no longer count tables that are already resident, and point
  evaluation builds only the convention scales, not the full
  nlat×(lmax+1)×(mmax+1) table (8.6 GB at lmax=1023).
- Pinned host buffers pin their backing Memory. On Julia 1.11+ reshape
  returns a new wrapper, so the finalizer on the allocated Vector
  unpinned staging memory that the host mirrors still used.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
- Add rules for the public analysis/synthesis (scalar, sphtor, qst) and
  the energies on PencilArrays, which the untyped serial rules used to
  capture. zgrad_* work on PencilArrays, and fdgrad_* gather the field,
  since ForwardDiff cannot pass duals through MPI. Gradients match the
  serial ones on 2 and 3 ranks. Batch, packed and SH_mul_mx on
  PencilArrays raise ArgumentError instead of returning a wrong adjoint.
- The distributed rules accept the primal keywords (comm, ltr, Aminus).
  Replicated cotangents compare NaN as equal, and replicated energy
  cotangents must agree across ranks.
- Packed real rotations are only ℝ-linear, so the inverse-rotation
  adjoint was 33-54% off for cotangents with complex m=0 entries. One
  exact adjoint now serves every rotation rule, and the angle pullbacks
  sweep Wigner matrices by degree (16.6 s → 0.29 s at lmax=192).
- SH_mul_mx pullbacks snapshot their operands; reused buffers corrupted
  the mx gradient.
- Pullbacks accept ZeroTangent and thunks and return one tangent per
  positional argument. Energy rules stop zgrad_* from tracing scalar
  loops (about 33 GB → 21 MiB at lmax=255).
- Document the distributed AD semantics, including that every rank's
  loss must use each distributed output: Zygote skips the pullback of an
  unused output, which strands its MPI collectives.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
grad_loss_vorticity_Tlm and loss_and_grad_vorticity_Tlm returned m > 0
entries at half the gradient that grad_enstrophy_Tlm, grad_energy_* and
Zygote return, so summing them in a regularised inverse problem gave an
inconsistent gradient. They now agree with Zygote to about 1e-13 under
both :dft and :quad.

Behaviour change: callers see the m > 0 entries doubled.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
FFTW's out-of-place real plans assume a contiguous output, so
use_rfft=true into a padded spatial view or a strided scratch buffer
failed with a wrong-strides error, in scalar and batch transforms.
Non-contiguous destinations now receive a contiguous result.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@sht_loop ran its CPU body under @inbounds with no check, and
examples/loop_demo.jl wrote an 8×8 array over a 10×10 interior. The
macro now checks, before the loop, the array accesses that every
iteration makes: indices affine in the loop index at the corners of the
range, others at every point. Guarded accesses, non-array operands and
indices that call other functions are left alone, so valid loops run as
before and no index with side effects is evaluated early.

The docstring's GPU accumulation example raced; it is replaced and the
rule documented.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
- DistAnalysisPlan and DistPlan were exported but undefined in core
  (UndefVarError). Core forwarders now reach the parallel extension or
  report it missing. A stray `using SHTnsKit` in ParallelLocal.jl caused
  three precompile warnings, and the extension now extends
  SHTnsKit.local_size instead of shadowing it. Aqua's undefined_exports
  check is enabled again.
- julia = "1.10" replaces the 1.10-1.12 cap, which would have held the
  next registered release back on Julia 1.13. CI tests 1.13.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
- CI exported SHTNSKIT_PHI_SCALE=dft, which overrides cfg.phi_scale, so
  every :quad test re-ran :dft. The variable is gone, and the tests that
  switch phi_scale clear it themselves.
- A gpu-host-reference job runs the GPU wrappers and shared kernels on
  KernelAbstractions.CPU(); the GPU parity workflow needs hardware
  runners that are not registered.
- The benchmark passed a matrix to SH_mul_mx and hid the failure behind
  a catch. It now uses packed vectors and fails loudly.
- Run the distributed AD API suite in CI.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
norms.md lists the three differences: the sign of toroidal
coefficients, the (-1)^m of odd negative orders in LM_cplx storage, and
the SH_mul_mx operator layout (how Y_l^m couples to l±1 rather than what
each (l, m) receives), with a recipe that converts an SHTns operator
vector. The operator docstrings note the layout.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
test_mpi_ad_public_api.jl differentiates with ForwardDiff and Zygote, but
the focused MPI step runs in the package environment with only MPI,
PencilArrays, PencilFFTs and ChainRulesCore added, so it stopped with
"Package ForwardDiff not found in current path".

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
@subhk
subhk merged commit 9c197d2 into main Oct 7, 2026
17 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant