Fix correctness, AD, type-stability, and MPI-efficiency bugs from full audit#26
Merged
Conversation
…l audit Two audit passes (serial+AD, then MPI/distributed) surfaced a set of confirmed defects; each fix below is verified (serial suite 67253/0/0; distributed AD FD-verified single-rank and on 2/3 ranks via mpiexec). Correctness - ext: distributed AD rrules called bare `globalindices`/`communicator` (defined only in the sibling ParallelExt module, not exported by PencilArrays) -> UndefVarError on first use; add local wrappers. - ParallelExt `adaptive_spectral_communication!` chose its MPI collective branch from rank-local sparsity -> divergent collectives -> deadlock; decide it collectively via Allreduce. - ParallelPlans `_validate_cfg_replicated` threw on only the mismatched ranks -> matching ranks hung at later collectives; Allreduce the mismatch so all ranks raise together. - config `prepare_plm_tables!` divided by sinθ at pole nodes -> NaN tables on pole-inclusive grids; detect poles, warn, keep on-the-fly path. - ParallelTransforms packed-storage + non-default norm called the matrix-only `convert_alm_norm!` on a packed Vector -> MethodError; scale the packed vector in place instead. Also guard `first([])` on a rank owning zero longitude columns. - ParallelLocal `dist_*_packed_cplx` filled the -m slot verbatim; store the Hermitian value and document that truly-complex fields are unsupported. - api_compat `shtns_init` ignored the documented SHT_NO_CS_PHASE / SHT_REAL_NORM flags; parse and thread them through. - ForwardDiffExt `fdgrad_*` distributed-array guard was unreachable for 2-D inputs (AbstractMatrix is more specific); move the guard onto the matrix methods. - normalization cache (`_ensure_norm_scale_matrix!`) was keyed on size only; invalidate it in setproperty! on norm/cs_phase change. - SH_mul_mx aliasing and @sht_loop field-access limitations documented; SH_to_lat_cplx now guards mres==1. Automatic differentiation - All packed rotation adjoints (SH_Zrotate/Yrotate/Yrotate90/Xrotate90, in both AdvancedADExt and ZygoteExt) returned wrong Qlm-gradients: the standard-inner-product adjoint is Q̄ = W·R⁻¹·(W⁻¹ȳ) with W=diag(wm), wm=2 for m>0; the code used the field-inner-product form (and Zrotate conjugated the cotangent, Xrotate90 used non-inverse angles). Add a new FD test (test_rotation_gradients.jl, 24 checks). - `synthesis_batch` rrule used `analysis` as the synthesis adjoint (off by Gauss weights); use `_adjoint_synthesis` per slice. Add batch-AD FD test. - Distributed sphtor pullback skipped per-element unthunking on the Tuple cotangent branch -> Matrix(::Thunk) failure; always unthunk. Type stability (breaks Float32 / ForwardDiff.Dual) - Derive accumulator/buffer eltype instead of hardcoding ComplexF64/Float64 in `_adjoint_synthesis`, `_adjoint_analysis`, batch synthesis, LoopVec few-m analysis, batch AD rrules, and the energy diagnostics. Efficiency - ParallelExt `efficient_spectral_reduce!` routed to a "hierarchical" path that wraps the same MPI.Allreduce! in two per-call Comm_split collectives for zero gain (tree_reduce! == Allreduce!); call Allreduce directly. - Skip the per-call reduce-comm Comm_split for the common theta decomposition (reduce over the full comm). - Batch the 3-way qst reduces and the sphtor pullback S̄/T̄ reduce into one collective; hoist the Wigner log-factorial buffer in the truncated Y-rotation gather. - Collapse the +/-m Legendre recompute to once per |m| in the complex packed synthesis/analysis/point evaluators. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… division The pole-NaN fix disabled precomputed tables entirely on pole-inclusive regular/DH grids, which regressed `create_regular_config(..., precompute_plm=true, include_poles=true)` (test asserts use_plm_tables). Value tables (plm/NP) have no sinθ division and are correct at poles, so table use stays enabled. Only the derivative tables' `-(dP̄/dθ)/sinθ` needs guarding: at an exact pole sinθ=0 gives 0/0 -> NaN. Set the pole entry to 0, which is kernel-consistent (the kernel's ×sinθ makes dθY = -0·NdP = 0 at the pole regardless) and avoids NaN propagation. Verified: DH pole grid keeps use_plm_tables=true, derivative tables are NaN-free, scalar roundtrip err 2.9e-15; full test/runtests.jl passes (serial 67253, parallel-grid 1673). Co-Authored-By: Claude Opus 4.8 (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
1 participant
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
Two audit passes over the codebase — first serial + AD, then the MPI/distributed layer — surfaced a set of confirmed defects. This PR fixes all of them. Every fix is verified: the serial suite passes 67253/0/0, distributed AD is FD-verified single-rank and on 2 and 3 ranks via real
mpiexec, and the rotation/batch AD fixes ship with new finite-difference tests.Correctness
SHTnsKitParallelADExtcalled bareglobalindices/communicator, which live only in the siblingSHTnsKitParallelExtmodule and are not exported by PencilArrays →UndefVarError. Added local wrappers. (FD-verified 1.1e-10.)adaptive_spectral_communication!selected its collective branch from rank-local sparsity → ranks diverged → hang. The decision is now made collectively viaAllreduce._validate_cfg_replicatedthrew on only the mismatched ranks, leaving rank 0 stuck at later collectives. It nowAllreduces the mismatch so all ranks raise together.prepare_plm_tables!divided bysinθ=0 on pole-inclusive grids. It now detects pole nodes, warns, and keeps the pole-safe on-the-fly path.MethodErrorcalling the matrix-onlyconvert_alm_norm!on a packed vector) — now scales the packed vector in place. Also guardsfirst([])on a rank owning zero longitude columns.dist_*_packed_cplxnow stores the Hermitian −m value (was verbatim +m) and documents that truly-complex fields are unsupported.shtns_initnow honors the documentedSHT_NO_CS_PHASE/SHT_REAL_NORMflags (were silently ignored).fdgrad_*distributed-array guard was unreachable for 2-D inputs (AbstractMatrixis more specific); moved onto the matrix methods.norm/cs_phasechange (was keyed on size only).SH_mul_mxno-alias and@sht_loopfield-access limitations;SH_to_lat_cplxnow guardsmres==1.Automatic differentiation
SH_Zrotate/SH_Yrotate/SH_Yrotate90/SH_Xrotate90(in bothAdvancedADExtandZygoteExt) returned incorrect Qlm-gradients. The correct standard-inner-product adjoint isQ̄ = W·R⁻¹·(W⁻¹ȳ)withW=diag(wm),wm=2form>0; the code used the field-inner-product form (Zrotate also conjugated the cotangent, Xrotate90 used non-inverse ZYZ angles). Newtest/serial/test_rotation_gradients.jl(24 FD checks).synthesis_batchrrule usedanalysisas the synthesis adjoint (off by Gauss weights); now uses_adjoint_synthesisper slice. New batch-AD FD test.Matrix(::Thunk)failure; now always unthunks.Type stability (Float32 /
ForwardDiff.Dual)ComplexF64/Float64in_adjoint_synthesis,_adjoint_analysis, batch synthesis, LoopVec few-m analysis, batch AD rrules, and the energy diagnostics.Efficiency
efficient_spectral_reduce!routed to a "hierarchical" reducer that wraps the sameMPI.Allreduce!in two per-callComm_splitcollectives for zero algorithmic gain (tree_reduce!is literallyAllreduce!). Now callsAllreduce!directly.Comm_splitfor the common θ-decomposition.S̄/T̄reduce into single collectives; hoists the Wigner log-factorial buffer in the truncated Y-rotation gather.|m|in the complex-packed synthesis/analysis/point evaluators.Verification
Notes
Project.tomlchange.Pencil((nlat,nlon),comm)splits φ, which routes through the anti-scalingAllgathervpath; scalable runs need explicit θ-decompositionPencil((nlat,nlon),(1,),comm)or the transpose plan. The code already warns.🤖 Generated with Claude Code