Repository navigation
Fix review findings across numerics, distributed, GPU and AD - #55
Merged
Merged
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
cfg.lmax/mmax/mresassignment still stored the new value, and the next transform wrote out of bounds under@inbounds.Commits
nlaton non-Gauss grids; 2 GiB cap on automatic tablesRealangles:quadaxisymmetric synthesis; Float32 and real inputs:quadevaluators;gpu_*_safe; point evaluation; pinning on Julia 1.11+zgrad_*memoryuse_rfft=trueinto views and strided scratchexamples/loop_demo.jlDistAnalysisPlan/DistPlandefined; precompile warnings;julia = "1.10"SHTNSKIT_PHI_SCALEhid:quad; newgpu-host-referencejob; benchmark usedSH_mul_mxwronglyLM_cplx,SH_mul_mxlayoutBehaviour changes
Worth a release note, and probably a minor version bump:
grad_loss_vorticity_Tlmandloss_and_grad_vorticity_Tlmreturn m > 0 entries doubled: they now give the Euclidean gradient, likegrad_enstrophy_Tlmand Zygote.create_configon:regular,:regular_polesand:driscoll_healygrids withoutnlat(or with exactlynlat = lmax+2) gives2lmax+1latitudes (2(lmax+1)for Driscoll–Healy). Code that hard-codes the old size gets aDimensionMismatch.create_regular_configno longer precomputes Legendre tables above 2 GiB;prepare_plm_tables!still does.ArgumentErrorinstead of returning wrong data:mmaxthe grid can't resolve;mmax < lmaxlayouts (pencil and GPU);create_spectral_pencil/create_spatial_pencilshare a cached topology per communicator.