From d9950cf7be6ef87e71849bd880ad62bdb6cc1a97 Mon Sep 17 00:00:00 2001 From: Subhajit Kar Date: Sun, 13 Sep 2026 06:55:21 +0300 Subject: [PATCH 1/3] fix: honour phi_scale end to end, guard truncating rotations, drop dead FFT cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Convention, AD and dead-code fixes found in a full-codebase audit. Every change is behaviour-preserving under the default `:dft` configuration; the breaking moves are confined to `:quad` and to hand-built configurations. phi_scale was only half-implemented. `synthesis` scaled its Fourier bins by `phi_inv_scale(cfg)` while `analysis` applied a fixed `cfg.cphi`, so under `:quad` the pair was not mutually inverse — `analysis(synthesis(alm))` came back as `alm/2π` exactly. Analysis now applies `cphi*nlon/phi_inv_scale(cfg)`, which is `cphi` under `:dft`, threaded through the batch, complex-packed, planned, distributed and adjoint paths. The direct evaluators (`synthesis_point`, `synthesis_axisym`, `SH_to_lat`, `SHqst_to_point`, `SH_to_grad_point`, `SHqst_to_lat` and the PencilArray locals) applied no φ factor at all and so disagreed with the grid they sample by 2π; they now carry the same factor. An unset `phi_scale` used to fall back to a grid-type guess, which made a regular grid from the exported keyword constructor disagree by 2π with the identical grid from `create_regular_config`; unset now means `:dft`, as every constructor emits. A Wigner-d rotation through a general β couples every |m'| <= l, so a layout with mmax < lmax silently dropped the components it could not hold — 14.8 % of the field energy at lmax=8/mmax=5, 24.0 % at mmax=3, with no error. Such a rotation now raises. The two non-mixing angles (β = 0, π) still work at any mmax, so pure Z-rotations expressed as ZYZ are unaffected. The turbo transforms walked a bare `0:mmax` instead of `0:mres:mmax`, making `analysis_turbo`/`synthesis_turbo` disagree with their core counterparts by O(1) on any `mres > 1` config, and nested `@threads :static` illegally inside an outer threaded region. Both now share the core's cached m ordering and threading predicate. Three rotation pullbacks (`SH_Yrotate`, `shtns_rotation_apply_cplx`, `shtns_rotation_apply_real`) read their primal input lazily, so an in-place call or a reused buffer silently corrupted the angle gradient; each now snapshots at primal time. The `analysis`/`synthesis`/sphtor rrules declared fewer keywords than their primals, which made ChainRules skip them entirely once a caller passed `use_rfft` or `fft_scratch`. `synthesis_qst`/`analysis_qst` had no rrule at all and crashed inside FFTW under Zygote. The `SHTNSKIT_CACHE_PENCILFFTS` knob and the four `*_fft_plan_cache*` controls forwarded to a cache in the parallel extension whose only reader had no call sites, storing NamedTuple placeholders the FFT wrappers ignored. The dead cache is deleted and the controls now address the one in `src/fftutils.jl` that the serial and distributed transforms actually use. `DistributedSpectralPlan2D` no longer attaches a finalizer, since `MPI_Comm_free` is collective and garbage collection is rank-local. Also: structural `SHTConfig` fields keep their derived state consistent or refuse assignment; the exported keyword constructor validates its invariants; `spatial_view` is exported as the missing bridge in the padding API; and the regular grids' `nlat >= 2*lmax + 1` exactness threshold is documented with measured numbers. Validation: 76,848 serial assertions; 1,714 parallel-grid assertions; MPI at 4 ranks (comprehensive, extended, audit-fixes, 2d-alignment) and 2 ranks (local correctness, comm cleanup, plan preflight, transpose operand preflight, transpose SHT, dealiased transpose, plan allocation). Julia 1.11.1, one thread, no CUDA or ROCm hardware present; Pkg.test/JET/Aqua and the GPU and parity runners were not re-run. Readiness evidence refreshed accordingly: the host-transfer allowlist regenerated (685 -> 676 entries) and the audited tree digest recomputed, with this run recorded separately from the historical readiness runs. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Dbj4TtuW11gke9JwSNs7Nz --- .github/workflows/ci.yml | 1 + CHANGELOG.md | 120 +++++++ docs/Distributed_SHTnsKit_Guide.md | 7 +- docs/phi_scaling.md | 66 +++- docs/src/grids.md | 20 ++ ext/ParallelLocal.jl | 12 +- ext/ParallelTransforms.jl | 60 ++-- ext/ParallelTransposeTransforms.jl | 4 +- ext/SHTnsKitAMDGPUExt.jl | 22 +- ext/SHTnsKitAdvancedADExt.jl | 95 +++++- ext/SHTnsKitGPUExt.jl | 22 +- ext/SHTnsKitLoopVecExt.jl | 73 ++++- ext/SHTnsKitParallelAMDGPUExt.jl | 4 +- ext/SHTnsKitParallelCUDAExt.jl | 4 +- ext/SHTnsKitParallelExt.jl | 307 +----------------- ext/SHTnsKitZygoteExt.jl | 8 +- src/SHTnsKit.jl | 98 +++--- src/batch_transforms.jl | 25 +- src/complex_packed.jl | 7 +- src/config.jl | 152 ++++++++- src/core_transforms.jl | 5 +- src/energy_diagnostics.jl | 4 + src/fftutils.jl | 118 +++++-- src/layout.jl | 11 +- src/legendre.jl | 67 +++- src/local.jl | 14 +- src/plan.jl | 4 +- src/qst_transforms.jl | 24 +- src/rotations.jl | 74 ++++- src/sphtor_transforms.jl | 36 +- src/transforms.jl | 57 ++-- src/vorticity_diagnostics.jl | 17 +- .../host_transfer_allowlist.toml | 152 +++------ test/fixtures/compatibility/task16_gate.toml | 8 +- .../compatibility/task16_local_commands.txt | 25 +- .../compatibility/task16_local_summary.log | 35 +- test/parallel/runtests.jl | 2 + test/parallel/test_mpi_2d_alignment.jl | 112 +++++++ test/serial/test_basic_transforms.jl | 80 ++++- test/serial/test_batch_qst.jl | 15 +- test/serial/test_batch_transforms.jl | 23 +- test/serial/test_configuration.jl | 113 +++++-- test/serial/test_energy_diagnostics.jl | 13 + test/serial/test_gradients.jl | 71 ++++ test/serial/test_indexing.jl | 18 + test/serial/test_local.jl | 35 ++ test/serial/test_plan.jl | 47 +++ test/serial/test_rotation_gradients.jl | 96 ++++++ test/serial/test_rotations.jl | 148 +++++++++ test/serial/test_turbo.jl | 49 +++ test/serial/test_vector_transforms.jl | 31 +- 51 files changed, 1899 insertions(+), 712 deletions(-) create mode 100644 test/parallel/test_mpi_2d_alignment.jl diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 68155489..b98e138f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -204,6 +204,7 @@ jobs: mpiexec -n 4 julia --project=. test/parallel/test_mpi_plan_preflight.jl mpiexec -n 2 julia --project=. test/parallel/test_mpi_transpose_operand_preflight.jl mpiexec -n 2 julia --project=. test/parallel/test_mpi_comm_cleanup.jl + mpiexec -n 4 julia --project=. test/parallel/test_mpi_2d_alignment.jl mpiexec -n 2 julia --project=. test/parallel/test_mpi_ad_tangent_spaces.jl mpiexec -n 8 julia --project=. test/parallel/test_mpi_audit_fixes.jl env: diff --git a/CHANGELOG.md b/CHANGELOG.md index 40aa0694..b2483c30 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,8 +37,91 @@ coefficient was `1/2π` too small — they inverted neither `synthesis_axisym` n the m=0 column of the full `analysis`. They now agree with both. Anything that compensated for the old scale downstream must drop that compensation. +**`analysis` and the direct evaluators now honour `phi_scale`; an unset +`phi_scale` resolves to `:dft`.** Two halves of one convention had drifted apart. + +*`analysis` ignored `phi_scale` entirely.* `synthesis` scales its Fourier bins by +`phi_inv_scale(cfg)`, but `analysis` applied a fixed `cfg.cphi`, so under `:quad` +the pair was not mutually inverse: `analysis(cfg, synthesis(cfg, alm))` came back +as `alm / 2π` exactly. Analysis now applies `cphi · nlon / phi_inv_scale(cfg)`, +which is `cphi` under `:dft` and restores the inverse property under `:quad`. The +same factor was threaded through the batch, complex-packed, planned, distributed +and adjoint analysis paths so they all agree. + +*The direct evaluators applied no φ factor at all.* `synthesis_point`, +`synthesis_point_cplx`, `synthesis_axisym`, `synthesis_axisym_l`, `SH_to_lat`, +`SH_to_lat_cplx`, `SHqst_to_point`, `SH_to_grad_point`, `SHqst_to_lat` and the +PencilArray local evaluations each disagreed with the grid they claim to sample by +exactly 2π under `:quad`. They now carry the same `phi_inv_scale(cfg)/nlon` factor +`synthesis` does. + +*An unset `phi_scale` (`:auto`) used to fall back to a grid-type guess* — +`grid_type == :gauss ? nlon : nlon/2π` — so a regular grid built through the +exported `SHTConfig(; …)` keyword constructor disagreed by 2π with the identical +grid from `create_regular_config`, which sets `:dft` explicitly. Every constructor +emits `:dft`, so an unset value now means `:dft` too. + +**Default `:dft` behaviour is unchanged in all three cases**; only `:quad` and +hand-built `:auto` configurations move, and they move to the values that make +`analysis` and `synthesis` inverses. + +**Order-mixing rotations now reject `mmax < lmax` instead of truncating.** +A Wigner-d rotation through a general `β` couples `Y_l^m` to every `Y_l^{m'}` with +`|m'| ≤ l`. When storage stopped at `mmax < lmax`, the `|m'| > mmax` components +were silently dropped — measured at `lmax = 8`, that discarded **14.8 %** of the +field's energy at `mmax = 5` and **24.0 %** at `mmax = 3`, with no error and no +warning. `SH_Yrotate`, `SH_Yrotate90`, `SH_Xrotate90` and the Euler-angle API now +raise an `ArgumentError` on such a configuration. Pure Z-rotations are unaffected: +`β ≡ 0` is diagonal and `β ≡ π` is anti-diagonal (`m' = -m`), so both still work +at any `mmax`. + +*Porting:* use `mmax == lmax` for anything but a Z-rotation. Results that appeared +to work before were missing the truncated energy. + +**Structural `SHTConfig` fields are no longer silently inconsistent.** +Assigning `cfg.lmax = 10` left `size(cfg.Nlm) == (7, 7)` while the transforms index +it as `(lmax+1, mmax+1)` under `@inbounds` — an out-of-bounds read of a live array. +`lmax`, `mmax` and `mres` now rebuild the derived spectral layout (`Nlm`, `nlm`, +`li`, `mi`, cached scale matrix and m-ordering) and drop the now-stale Legendre +tables; `nlat`, `nlon`, `grid_type`, `nlm`, `li`, `mi` and `nspat` raise an +`ArgumentError` pointing at the `create_*_config` constructors, because there is no +grid-type-independent way to regenerate the quadrature in place. + +**The exported `SHTConfig(; …)` keyword constructor validates its invariants.** +It previously checked nothing, so a hand-built configuration could violate +`nlon ≥ 2*mmax+1` and then silently synthesize an all-zero field for any mode it +could not resolve, or hand `use_rfft=true` a raw `BoundsError`. It now enforces the +same constraints the `create_*_config` helpers always have. The exported keyword +signature is otherwise unchanged. + ### Fixed +- **`analysis_turbo` / `synthesis_turbo` ignored `mres`.** Both walked a bare + `0:mmax` instead of `0:mres:mmax`, so `analysis_turbo` populated — and + `synthesis_turbo` consumed — coefficient columns an `mres > 1` transform has no + storage for. The disagreement with `analysis`/`synthesis` was O(1), not + roundoff. Both now share the core's cached `m` ordering. The turbo pair also + nested `@threads :static`, which is illegal inside an outer threaded region; + they now fall back to a serial loop there, using the same predicate the core + orchestrators use. +- **Rotation pullbacks read coefficients the primal had already overwritten.** + `SH_Yrotate`, `shtns_rotation_apply_cplx` and `shtns_rotation_apply_real` + captured their primal *input* and read it lazily, so an in-place call + (`Rlm === Qlm`) — or any caller reusing the buffer before the pullback ran — + silently corrupted the angle gradient. Each rule now snapshots what it needs at + primal time. Applies to both the ChainRules and Zygote adjoints. +- **`rrule`s for `analysis`/`synthesis`/`analysis_sphtor`/`synthesis_sphtor` + declared fewer keyword arguments than their primals.** Passing `use_rfft` or + `fft_scratch` — even at its default — made ChainRules skip the rule entirely and + fall through to source tracing. The keywords select a different FFT + implementation of the same linear operator, so the adjoints are unchanged. +- **`synthesis_qst` and `analysis_qst` had no `rrule` at all**, so + differentiating a QST pipeline fell through to Zygote's source tracing and + crashed inside FFTW. Each adjoint is the existing scalar and sphtor adjoints + side by side. +- **`shtns_rotation_apply_real` reported a bare size mismatch for `mres > 1` + configurations**, leaving the caller to reverse-engineer why. The message now + names `mres` and states the restriction, matching `dist_SH_Yrotate`. - **Silent precision loss in batch QST/sphtor transforms.** `analysis_qst_batch`, `_synthesis_qst_batch` and the sphtor batch pair derived their output element type from one input array instead of promoting across all of them, truncating @@ -95,6 +178,43 @@ compensated for the old scale downstream must drop that compensation. ### Internal +- **Removed the dead parallel FFT-plan cache.** The `SHTNSKIT_CACHE_PENCILFFTS` + environment variable and the `fft_plan_cache_enabled` / `set_fft_plan_cache!` / + `enable_fft_plan_cache!` / `disable_fft_plan_cache!` controls forwarded to a + cache in the parallel extension whose only reader, `_get_or_plan`, had no call + sites — the "plans" it stored were `NamedTuple` placeholders the FFT wrappers + ignored, so every knob was a no-op. The cache the transforms actually use is now + in `src/fftutils.jl`, shared by the serial and distributed paths, and the same + four controls address it without requiring the extension to be loaded. + `SHTNSKIT_FFT_PLAN_CACHE` is the current spelling; the old name still works. +- **`DistributedSpectralPlan2D` no longer attaches a finalizer.** `close` frees + the plan's `l_comm` / `m_comm` sub-communicators, and `MPI_Comm_free` is + collective; a finalizer runs at whatever point that rank's garbage collector + fires, which is rank-local and nondeterministic. Cleanup is explicit only — call + `close(plan)` collectively. Leaking two communicators until `MPI_Finalize` is + strictly better than a nondeterministic collective. +- **`spatial_view(cfg, A)` is exported**, the missing bridge in the padding API: + `allocate_padded_spatial` returns an array with `nlat_padded ≥ nlat` rows while + every transform requires exactly `nlat`, so the padded buffer could not be passed + to `analysis` at all. The view keeps the padded column stride, so it preserves + what the padding is for. +- **`set_batch_size!` is documented as advisory.** `howmany` / `spec_dist` mirror + the SHTns C batch descriptors and are stored for interoperability, but the Julia + batch entry points take the field count from `size(fields, 3)`; the old docstring + claimed otherwise. +- **Regular-grid quadrature exactness is documented.** Fejér and Clenshaw–Curtis + rules with `nlat` nodes are exact only through degree `nlat - 1`, and analysis + integrates degree `2*lmax`, so the equiangular grids need `nlat ≥ 2*lmax + 1` — + where Gauss–Legendre needs `nlat = lmax + 1`. Below that threshold nothing warns + and `analysis ∘ synthesis` is not an identity (7.2e-2 relative error at + `lmax = 8, nlat = 10`). Both `create_regular_config` and `docs/src/grids.md` now + say so with measured numbers. +- New regression coverage: planned-vs-`cfg` conformance across the tables / rfft / + Robert-form matrix; the `mmax < lmax` rotation guard; angle-gradient survival + under an in-place primal; evaluator `phi_scale` agreement; turbo `mres`; and a + 4-rank `test_mpi_2d_alignment.jl` for the 2D spectral-plan alignment + preconditions, wired into CI. + - `pack_lm!`/`pack_lm`/`unpack_lm!`/`unpack_lm` in `src/layout.jl` replace six open-coded copies of the packed↔dense `(l,m)` mapping. The `m % mres` guard had to be fixed three separate times across those copies. diff --git a/docs/Distributed_SHTnsKit_Guide.md b/docs/Distributed_SHTnsKit_Guide.md index 334a4701..0693d5e1 100644 --- a/docs/Distributed_SHTnsKit_Guide.md +++ b/docs/Distributed_SHTnsKit_Guide.md @@ -148,7 +148,7 @@ SHTnsKit.dist_synthesis!(spln, fθφ_out, PencilArray(Alm)) Enable plan caching across calls (optional) ```bash -export SHTNSKIT_CACHE_PENCILFFTS=1 +export SHTNSKIT_FFT_PLAN_CACHE=1 # legacy alias: SHTNSKIT_CACHE_PENCILFFTS ``` --- @@ -242,7 +242,10 @@ Both use the distributed transform paths internally and return gradients in the ``` - Robert form: for vector transforms, set `robert_form=true` in your config to stabilize polar behavior. - Normalization/phase: match `cfg.norm` and `cfg.cs_phase` to your data; conversions are handled internally on input/output. -- FFT plan caching: `ENV["SHTNSKIT_CACHE_PENCILFFTS"] = "1"` to reuse PencilFFTs plans. +- FFT plan caching: `ENV["SHTNSKIT_FFT_PLAN_CACHE"] = "1"` (default) to reuse the + per-rank φ-FFT plans. Equivalently `SHTnsKit.enable_fft_plan_cache!()` / + `disable_fft_plan_cache!()`. The cache is shared with the serial transforms; + cap it with `SHTnsKit.fft_plan_cache_max!(n)`. - Y-rotation strategy: truncated gather typically reduces bandwidth; switch to allgather for high-m–dominated spectra. --- diff --git a/docs/phi_scaling.md b/docs/phi_scaling.md index 35a9d19c..ce17ddd4 100644 --- a/docs/phi_scaling.md +++ b/docs/phi_scaling.md @@ -20,25 +20,32 @@ The `phi_scale` field in `SHTConfig` controls how the longitude (φ) dimension i - Regular grids with poles (simple trapezoidal rule) - **Rationale**: Adjusts for the φ integration measure `dφ` where ∫₀²ᵖ f dφ ≈ (2π/nlon) Σ f_j -### `:auto` - Automatic Selection -- **Behavior**: Chooses based on `grid_type` - - `:gauss` → `:dft` - - `:driscoll_healy` → `:dft` - - `:regular`, `:regular_poles` → `:quad` +### `:auto` - Unset +- **Behavior**: treated as `:dft`. +- Historically this keyed on `grid_type` and handed every non-Gauss grid `:quad`, + so a regular grid built through the exported `SHTConfig(; ...)` keyword + constructor (which defaulted to `:auto`) disagreed by 2π with the identical + grid from `create_regular_config`, which sets `:dft` explicitly. Both + constructors emit `:dft`, so an unset value now means `:dft` as well. ## Configuration ### In Code ```julia -# Explicit control -cfg = create_gauss_config(lmax, nlat; phi_scale=:dft) -cfg = create_regular_config(lmax, nlat; phi_scale=:quad) +# Every constructor emits :dft; there is no phi_scale keyword on them. +cfg = create_gauss_config(lmax, nlat) # :dft +cfg = create_regular_config(lmax, nlat) # :dft -# Automatic (recommended) -cfg = create_gauss_config(lmax, nlat) # Uses :dft -cfg = create_regular_config(lmax, nlat) # Uses :quad +# To opt into the quadrature convention, set it on the config: +cfg.phi_scale = :quad ``` +!!! note + An earlier version of this page showed `create_regular_config(lmax, nlat; + phi_scale=:quad)` and claimed regular grids default to `:quad`. Neither was + true: those constructors take no `phi_scale` keyword (that call raises) and + both set `:dft`. + ### Via Environment Variable ```bash # Override for all grids @@ -60,18 +67,24 @@ function phi_inv_scale(cfg::SHTConfig) return cfg.nlon end - # 2. Use config-specified mode + # 2. Use config-specified mode. `:dft` and anything unset (`:auto`) both + # mean the DFT convention — every constructor emits `:dft`, so an unset + # value is not a signal to guess from the grid type. if cfg.phi_scale === :quad return cfg.nlon / (2π) - elseif cfg.phi_scale === :dft - return cfg.nlon + else + return Float64(cfg.nlon) end - - # 3. Fall back to grid-type heuristic - return cfg.grid_type == :gauss ? cfg.nlon : cfg.nlon / (2π) end ``` +!!! note "The old grid-type fallback is gone" + An unset `phi_scale` used to fall back to `grid_type == :gauss ? nlon : nlon/2π`. + That made a regular grid built through the exported keyword `SHTConfig(...)` + constructor (whose `phi_scale` defaults to `:auto`) disagree by 2π with the + identical grid from `create_regular_config`, which sets `:dft` explicitly. + Unset now resolves to `:dft`, matching every constructor. + ## Why This Matters Incorrect φ-scaling leads to round-trip errors: @@ -90,3 +103,22 @@ The φ-scaling factor must match the quadrature weight convention to ensure: - Commit fc1d114: Regular grids changed to quadrature scaling (`nlon/(2π)`) - Commit 2441db0: Formalized with auto-detection - Current: Explicit `phi_scale` field for clarity and control + + +## Invariant + +Whichever mode is selected, `analysis` and `synthesis` are mutual inverses: + +```julia +analysis(cfg, synthesis(cfg, alm)) ≈ alm # exact under :dft and :quad +``` + +`synthesis` scales its Fourier bins by `phi_inv_scale(cfg)` and the inverse FFT +divides by `nlon`, a net spatial factor `σ`; `analysis` carries `cphi/σ` so the +two cancel. `:quad` therefore changes the scale of the *spatial* field (by 1/2π) +without changing what a round trip returns. Point and latitude evaluators +(`synthesis_point`, `SH_to_lat`, `SHqst_to_lat`, …) apply the same `σ`, so they +always agree with the grid `synthesis` produces. + +Before this was fixed, `analysis` ignored `phi_scale` entirely, so under `:quad` +a round trip returned `alm/2π` and every evaluator was 2π off from the grid. diff --git a/docs/src/grids.md b/docs/src/grids.md index e79474a3..34ad810c 100644 --- a/docs/src/grids.md +++ b/docs/src/grids.md @@ -37,6 +37,26 @@ Gauss–Legendre quadrature with `nlat` points integrates polynomials in quadrature is exact at the intended band limit when `nlat = 2*(lmax + 1)`; `nlat` must be even. +!!! warning "Equiangular grids need `nlat ≥ 2*lmax + 1`" + Fejér and Clenshaw–Curtis rules with `nlat` nodes are exact only through + degree ``n_\mathrm{lat} - 1``, and analysis integrates a product of two + degree-`lmax` Legendre functions — degree ``2\,l_\mathrm{max}``. So the + two regular grids reproduce `alm` from `analysis(synthesis(alm))` only from + `nlat = 2*lmax + 1` upward, where Gauss–Legendre needs just `nlat = lmax + 1`. + + Below the threshold nothing warns; the answer is simply wrong. Measured + relative round-trip error at `lmax = 8`: + + | `nlat` | `:gauss` | regular midpoint | regular with poles | + |---:|---:|---:|---:| + | 10 | 8e-16 | 7.2e-2 | 3.6e-1 | + | 14 | 8e-16 | 1.3e-2 | 7.1e-3 | + | 16 | 8e-16 | 4.2e-3 | 1.3e-3 | + | 17 (`2*lmax+1`) | 8e-16 | 7e-16 | 7e-16 | + + Size the grid deliberately, or use Gauss–Legendre (or Driscoll–Healy, exact + at `nlat = 2*(lmax+1)`) when round-trip accuracy matters. + !!! tip "Which grid should I choose?" Start with Gauss–Legendre unless you need to exchange data with a particular equiangular layout. Use regular midpoint for cell-centred data, regular with diff --git a/ext/ParallelLocal.jl b/ext/ParallelLocal.jl index 208cf6a5..d2d2a7ea 100644 --- a/ext/ParallelLocal.jl +++ b/ext/ParallelLocal.jl @@ -80,6 +80,8 @@ end `Q_p`/`S_p`/`T_p` use the coefficient convention configured by `cfg` (see [`dist_SH_to_lat`](@ref)). +When `cfg.robert_form` is enabled, the tangential components include the +`sin(θ)` factor used by full-grid vector synthesis; the radial component is unchanged. """ function SHTnsKit.dist_SHqst_to_point(cfg::SHTnsKit.SHTConfig, Q_p::PencilArray, S_p::PencilArray, T_p::PencilArray, cost::Real, phi::Real) return SHTnsKit.SHqst_to_point(cfg, Q_p, S_p, T_p, cost, phi) @@ -91,6 +93,8 @@ end `Q_p`/`S_p`/`T_p` use the coefficient convention configured by `cfg` (see [`dist_SH_to_lat`](@ref)). +When `cfg.robert_form` is enabled, the tangential components include the +`sin(θ)` factor used by full-grid vector synthesis; the radial component is unchanged. """ function SHTnsKit.dist_SHqst_to_lat(cfg::SHTnsKit.SHTConfig, Q_p::PencilArray, S_p::PencilArray, T_p::PencilArray, cost::Real; nphi::Int=cfg.nlon, ltr::Int=cfg.lmax, mtr::Int=cfg.mmax) @@ -313,6 +317,12 @@ function _pencil_local_qst(cfg, Q::PencilArray, S::PencilArray, Vt .*= sinth Vp .*= sinth end + # Same φ convention factor the serial evaluators apply (1 under :dft; + # 1/2π under :quad), so a PencilArray input agrees with the dense one. + sphi_scale = RT(SHTnsKit._evaluator_phi_scale(cfg)) + if sphi_scale != 1 + Vr .*= sphi_scale; Vt .*= sphi_scale; Vp .*= sphi_scale + end combined = vcat(Vr, Vt, Vp) _record_local_payload!(length(combined)) MPI.Allreduce!(combined, +, comm) @@ -1135,7 +1145,7 @@ function _analysis_mode_pencil(cfg::SHTnsKit.SHTConfig, im::Int, RT = typeof(real(zero(CT))) P = Vector{Float64}(undef, ltr + 1) rank = MPI.Comm_rank(comm) - phi_scale = axisymmetric ? cfg.cphi * cfg.nlon : cfg.cphi + phi_scale = axisymmetric ? SHTnsKit._analysis_phi_scale(cfg) * cfg.nlon : SHTnsKit._analysis_phi_scale(cfg) for root in 0:(MPI.Comm_size(comm) - 1) send = zeros(CT, counts[root + 1]) @inbounds for (i, θindex) in pairs(θglobals) diff --git a/ext/ParallelTransforms.jl b/ext/ParallelTransforms.jl index fe160b9b..a9b32dc7 100644 --- a/ext/ParallelTransforms.jl +++ b/ext/ParallelTransforms.jl @@ -1100,7 +1100,7 @@ function _analysis_owned_block!(block::AbstractMatrix{CT}, cfg::SHTnsKit.SHTConf fill!(block, zero(CT)) P = Vector{Float64}(undef, cfg.lmax + 1) RT = typeof(real(zero(CT))) - cphi = RT(cfg.cphi) + cphi = RT(SHTnsKit._analysis_phi_scale(cfg)) use_tbl = cfg.use_plm_tables && !isempty(cfg.NP_tables) @inbounds for (local_m, m_index) in pairs(m_indices) m = m_index - 1 @@ -1374,7 +1374,7 @@ Decompose latitude instead: `SHTnsKit.create_spatial_pencil(cfg; comm)` or `Penc if use_packed_storage # Original inline loop for packed storage (not the hot path). # Uses normalized rows (Plm_norm_row! / NP_tables): Nlm is already baked in. - xv = cfg.x; cphi = cfg.cphi # hoist field reads out of the loops below (cfg is mutable, so not auto-hoisted) + xv = cfg.x; cphi = SHTnsKit._analysis_phi_scale(cfg) # hoisted; inverts synthesis under any phi_scale # Stride by mres: create_packed_storage_info only assigns lm_to_packed for # m % mres == 0, leaving every other entry 0. Walking all m under @inbounds # therefore wrote Alm_local[0] — one element before the buffer. @@ -1441,7 +1441,7 @@ Decompose latitude instead: `SHTnsKit.create_spatial_pencil(cfg; comm)` or `Penc if !use_packed_storage # Apply φ scaling (cphi = 2π/nlon). Nlm is NOT applied here: the # normalized recurrence Plm_norm_row! already bakes Nlm into P̄. - cphi = cfg.cphi # hoist field read out of the normalization loop (cfg is mutable) + cphi = SHTnsKit._analysis_phi_scale(cfg) # hoisted; inverts synthesis under any phi_scale @inbounds for m in 0:cfg.mres:mmax @simd ivdep for l in m:lmax Alm_local[l+1, m+1] *= cphi @@ -1451,7 +1451,7 @@ Decompose latitude instead: `SHTnsKit.create_spatial_pencil(cfg; comm)` or `Penc else # θ is not distributed - no reduction needed, just apply φ scaling if !use_packed_storage - cphi = cfg.cphi # hoist field read out of the normalization loop (cfg is mutable) + cphi = SHTnsKit._analysis_phi_scale(cfg) # hoisted; inverts synthesis under any phi_scale @inbounds for m in 0:cfg.mres:mmax @simd ivdep for l in m:lmax Alm_local[l+1, m+1] *= cphi @@ -1515,7 +1515,7 @@ function _dist_analysis_plan_core!(plan::DistAnalysisPlan, end # φ scaling (Nlm already baked into the normalized Legendre rows) - cphi = cfg.cphi + cphi = SHTnsKit._analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) @inbounds for m in 0:cfg.mres:mmax @simd ivdep for l in m:lmax plan.Alm_work[l+1, m+1] *= cphi @@ -2054,7 +2054,7 @@ function _analysis_sphtor_owned_block!(Sblock::AbstractMatrix{CT}, for l in max(1, m):lcap d = RT(dtheta[l + 1]) term = complex(zero(RT), RT(m * over_sin[l + 1])) - factor = wi * RT(cfg.cphi) / RT(l * (l + 1)) + factor = wi * RT(SHTnsKit._analysis_phi_scale(cfg)) / RT(l * (l + 1)) Sblock[l + 1, local_m] += factor * (Ft * d + conj(term) * Fp) Tblock[l + 1, local_m] += @@ -2383,7 +2383,7 @@ function _analysis_sphtor_mode_pencil(cfg::SHTnsKit.SHTConfig, l >= max(1, physical_m) || continue derivative = RT(dtheta[l + 1]) term = complex(zero(RT), RT(physical_m * over_sin[l + 1])) - coefficient = weight * RT(cfg.cphi) / RT(l * (l + 1)) + coefficient = weight * RT(SHTnsKit._analysis_phi_scale(cfg)) / RT(l * (l + 1)) Ssend[k] += coefficient * (Ftheta * derivative + conj(term) * Fphi) Tsend[k] += coefficient * @@ -2571,7 +2571,7 @@ function _legacy_dist_analysis_sphtor(cfg::SHTnsKit.SHTConfig, dPdtheta = Vector{Float64}(undef, lmax + 1) P_over_sth = Vector{Float64}(undef, lmax + 1) Pbuf = Vector{Float64}(undef, lmax + 2) # scratch for normalized dθ recurrence - scaleφ = cfg.cphi + scaleφ = SHTnsKit._analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Main vector analysis loop — via the same function barriers the planned # path uses (concrete argument types; the previous inline loop boxed its @@ -2763,12 +2763,12 @@ function _dist_analysis_sphtor_plan_core!( _sphtor_analysis_loop_tbl!(cfg, plan.Slm_work, plan.Tlm_work, cfg.NP_tables, cfg.NdP_tables, plan.Ftθm, plan.Fpθm, plan.θ_globals, plan.sθ_cache, plan.weights_cache, - cfg.robert_form, cfg.cphi, lmax, mmax) + cfg.robert_form, SHTnsKit._analysis_phi_scale(cfg), lmax, mmax) else _sphtor_analysis_loop_otf!(plan.Slm_work, plan.Tlm_work, plan.P, plan.dPdtheta, plan.P_over_sth, plan.Pbuf, plan.Ftθm, plan.Fpθm, plan.x_cache, plan.sθ_cache, plan.inv_sθ_cache, - plan.weights_cache, cfg.robert_form, cfg.cphi, + plan.weights_cache, cfg.robert_form, SHTnsKit._analysis_phi_scale(cfg), cfg.mres, lmax, mmax) end @@ -3930,7 +3930,7 @@ function dist_analysis_distributed(cfg::SHTnsKit.SHTConfig, fθφ::PencilArray; # Compute local contributions to ALL coefficients (same as standard analysis) local_contrib = zeros(ComplexF64, lmax + 1, mmax + 1) - scaleφ = cfg.cphi + scaleφ = SHTnsKit._analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Pre-cache weights weights_cache = Vector{Float64}(undef, nθ_local) @@ -4177,6 +4177,19 @@ end return false end +""" + close(plan::DistributedSpectralPlan2D) + +Free the plan's `l_comm` / `m_comm` sub-communicators. + +!!! warning "Collective" + `MPI_Comm_free` is a collective operation: **every** rank holding this plan + must call `close` on it, and in the same order relative to other collectives. + It is therefore not attached to a finalizer — garbage collection is + rank-local and would call it at divergent points. Closing twice is safe + (the second call is a no-op); never closing only leaks the communicators + until `MPI_Finalize`. +""" function Base.close(plan::DistributedSpectralPlan2D) plan.closed && return nothing plan.closed = true @@ -4520,12 +4533,21 @@ function create_distributed_spectral_plan_2d(lmax::Int, mmax::Int, comm::MPI.Com m_group_nlm, with_scratch, scratch, scratch_context, false ) - finalizer(plan) do p - try - close(p) - catch - end - end + # NO finalizer here, deliberately. + # + # `close(plan)` frees the `l_comm` / `m_comm` sub-communicators, and + # `MPI_Comm_free` is a COLLECTIVE: every rank in the communicator must + # call it, in the same order relative to other collectives. A finalizer + # runs whenever that rank's garbage collector happens to fire, which is + # rank-local and nondeterministic — so attaching one meant ranks could + # enter `Comm_free` at completely different points in the program, or + # some not at all before the next collective. That is a hang or worse, + # and it contradicts the policy `ParallelPlans.jl` states for its own + # `Comm_split` bookkeeping. + # + # Leaking two communicators until `MPI_Finalize` is strictly better than + # a nondeterministic collective, so cleanup is explicit only: call + # `close(plan)` (collectively) when you are done with it. return plan catch _safe_comm_free(l_comm) @@ -4880,7 +4902,7 @@ function _dist_analysis_2d_safe(cfg::SHTnsKit.SHTConfig, fθφ::PencilArray; x_cache[ii] = cfg.x[iglob] end - scaleφ = cfg.cphi + scaleφ = SHTnsKit._analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Use NP_tables (already normalized P̄) if available; fall back to OTF normalized rows. use_tbl = use_tables && cfg.use_plm_tables && !isempty(cfg.NP_tables) P = Vector{Float64}(undef, lmax + 1) @@ -5388,7 +5410,7 @@ function _dist_analysis_2d_aligned(cfg::SHTnsKit.SHTConfig, fθφ::PencilArray; copyto!(Fθm, Fθm_temp) end - scaleφ = cfg.cphi + scaleφ = SHTnsKit._analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Use NP_tables (already normalized P̄) if available; fall back to OTF normalized rows. use_tbl = use_tables && cfg.use_plm_tables && !isempty(cfg.NP_tables) diff --git a/ext/ParallelTransposeTransforms.jl b/ext/ParallelTransposeTransforms.jl index 528670bc..17d11ce6 100644 --- a/ext/ParallelTransposeTransforms.jl +++ b/ext/ParallelTransposeTransforms.jl @@ -515,7 +515,7 @@ function _dist_transpose_analysis_unchecked!( fill!(A, zero(eltype(A))) w = plan.cfg.w - scaleφ = plan.cfg.cphi # 2π/nlon — converts unnormalized rFFT sum to integral + scaleφ = SHTnsKit._analysis_phi_scale(plan.cfg) # 2π/nlon under :dft; inverts synthesis under any phi_scale lmax = plan.lmax nlat = plan.nlat nlev = plan.nlev @@ -657,7 +657,7 @@ function _dist_transpose_analysis_sphtor_unchecked!( fill!(T, zero(eltype(T))) w = plan.cfg.w - scaleφ = plan.cfg.cphi # 2π/nlon — converts unnormalized rFFT sum to integral + scaleφ = SHTnsKit._analysis_phi_scale(plan.cfg) # 2π/nlon under :dft; inverts synthesis under any phi_scale lmax = plan.lmax nlat = plan.nlat nlev = plan.nlev diff --git a/ext/SHTnsKitAMDGPUExt.jl b/ext/SHTnsKitAMDGPUExt.jl index 7840894a..f16d9ce5 100644 --- a/ext/SHTnsKitAMDGPUExt.jl +++ b/ext/SHTnsKitAMDGPUExt.jl @@ -627,7 +627,7 @@ function _amdgpu_scalar_analysis_direct!(owner, cfg::SHTConfig, end backend = ROCBackend() scalar_analysis_kernel!(backend)( - output, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + output, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres, cfg.lmax; ndrange=(cfg.lmax + 1, cfg.mmax + 1), ) @@ -716,7 +716,7 @@ function _amdgpu_scalar_analysis(cfg::SHTConfig, field::AMDGPU.AnyROCArray; backend = ROCBackend() canonical = AMDGPU.zeros(CT, cfg.lmax + 1, cfg.mmax + 1) analyze! = scalar_analysis_kernel!(backend) - analyze!(canonical, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + analyze!(canonical, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres, lcap; ndrange=(lcap + 1, min(cfg.mmax, lcap) + 1)) @@ -849,7 +849,7 @@ function _amdgpu_vector_analysis_direct!(owner, cfg::SHTConfig, vector_analysis_kernel!(ROCBackend())( Sout, Tout, workspace.Ftheta, workspace.Fphi, tables.dtheta, tables.over_sin, tables.weights, tables.scales, - tables.x, RT(cfg.cphi), lcap, min(cfg.mmax, lcap), cfg.mres, + tables.x, RT(SHTnsKit._analysis_phi_scale(cfg)), lcap, min(cfg.mmax, lcap), cfg.mres, cfg.robert_form; ndrange=(lcap + 1, min(cfg.mmax, lcap) + 1), ) AMDGPU.synchronize() @@ -1044,7 +1044,7 @@ function _amdgpu_vector_mode_analysis(cfg::SHTConfig, stored_im::Integer, S = AMDGPU.zeros(CT, lcap - physical_m + 1); Tlm = similar(S) vector_mode_analysis_kernel!(ROCBackend())( S, Tlm, Vt, Vp, tables.dtheta, tables.over_sin, tables.weights, - tables.scales, tables.x, RT(cfg.cphi), physical_m, lcap, + tables.scales, tables.x, RT(SHTnsKit._analysis_phi_scale(cfg)), physical_m, lcap, cfg.robert_form; ndrange=length(S), ) AMDGPU.synchronize() @@ -1181,7 +1181,7 @@ function _amdgpu_vector_batch_analysis(cfg::SHTConfig, S = AMDGPU.zeros(CT, cfg.lmax + 1, cfg.mmax + 1, nfields); Tlm = similar(S) vector_batch_analysis_kernel!(ROCBackend())( S, Tlm, Ft, Fp, tables.dtheta, tables.over_sin, tables.weights, - tables.scales, tables.x, RT(cfg.cphi), cfg.lmax, cfg.mmax, + tables.scales, tables.x, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres, cfg.robert_form; ndrange=size(S), ) AMDGPU.synchronize() @@ -1537,7 +1537,7 @@ function analysis_axisym(::SHTnsKit.GPU, cfg::SHTConfig, length(field) == cfg.nlat || throw(DimensionMismatch( "field must have length nlat=$(cfg.nlat)", )) - return _amdgpu_mode_analysis(cfg, 0, field, cfg.lmax, cfg.cphi * cfg.nlon) + return _amdgpu_mode_analysis(cfg, 0, field, cfg.lmax, SHTnsKit._analysis_phi_scale(cfg) * cfg.nlon) end analysis_axisym(cfg::SHTConfig, field::AMDGPU.AnyROCArray{T,1}) where {T<:Real} = analysis_axisym(SHTnsKit.GPU(), cfg, field) @@ -1559,7 +1559,7 @@ function analysis_axisym_l(::SHTnsKit.GPU, cfg::SHTConfig, length(field) == cfg.nlat || throw(DimensionMismatch( "field must have length nlat=$(cfg.nlat)", )) - return _amdgpu_mode_analysis(cfg, 0, field, lcap, cfg.cphi * cfg.nlon) + return _amdgpu_mode_analysis(cfg, 0, field, lcap, SHTnsKit._analysis_phi_scale(cfg) * cfg.nlon) end analysis_axisym_l(cfg::SHTConfig, field::AMDGPU.AnyROCArray{T,1}, ltr::Integer) where {T<:Real} = @@ -1599,7 +1599,7 @@ function analysis_packed_ml(::SHTnsKit.GPU, cfg::SHTConfig, im::Int, length(mode) == cfg.nlat || throw(DimensionMismatch( "mode must have length nlat=$(cfg.nlat)", )) - return _amdgpu_mode_analysis(cfg, physical_m, mode, lcap, cfg.cphi) + return _amdgpu_mode_analysis(cfg, physical_m, mode, lcap, SHTnsKit._analysis_phi_scale(cfg)) end analysis_packed_ml(cfg::SHTConfig, im::Int, mode::AMDGPU.AnyROCArray{T,1}, ltr::Integer) where {T<:Complex} = @@ -1637,7 +1637,7 @@ function _amdgpu_analysis_packed_cplx(cfg::SHTConfig, mcap = min(cfg.mmax, lcap) kernel! = complex_packed_analysis_kernel!(ROCBackend()) kernel!(packed, fourier, tables.Plm, tables.weights, tables.scales, - RT(cfg.cphi), cfg.nlon, lcap, cfg.mmax, mcap; + RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.nlon, lcap, cfg.mmax, mcap; ndrange=(lcap + 1, 2mcap + 1)) AMDGPU.synchronize() return packed @@ -1723,7 +1723,7 @@ function _amdgpu_batch_analysis(cfg::SHTConfig, fields::AMDGPU.AnyROCArray; FFTW.fft!(fourier, 2) canonical = AMDGPU.zeros(CT, cfg.lmax + 1, cfg.mmax + 1, size(fields, 3)) kernel! = scalar_batch_analysis_kernel!(ROCBackend()) - kernel!(canonical, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + kernel!(canonical, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres; ndrange=size(canonical)) configured = canonical ./ reshape( tables.scales, cfg.lmax + 1, cfg.mmax + 1, 1, @@ -1765,7 +1765,7 @@ function _amdgpu_batch_analysis_direct!(cfg::SHTConfig, end backend = ROCBackend() scalar_batch_analysis_kernel!(backend)( - output, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + output, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres; ndrange=size(output), ) coefficient_batch_conversion_kernel!(backend)( diff --git a/ext/SHTnsKitAdvancedADExt.jl b/ext/SHTnsKitAdvancedADExt.jl index a9460050..674b5319 100644 --- a/ext/SHTnsKitAdvancedADExt.jl +++ b/ext/SHTnsKitAdvancedADExt.jl @@ -42,8 +42,13 @@ import SHTnsKit: wigner_d_matrix_deriv A isa ChainRulesCore.AbstractZero ? _coeff_zeros(cfg) : _to_complex(A) - function ChainRulesCore.rrule(::typeof(SHTnsKit.analysis), cfg::SHTnsKit.SHTConfig, f) - y = SHTnsKit.analysis(cfg, f) + # `fft_scratch` / `use_rfft` pick a different FFT implementation of the SAME + # linear operator, so the adjoint is unchanged — but a pullback must still + # ACCEPT them. Declaring fewer kwargs than the primal made ChainRules skip + # this rule entirely the moment a caller passed one, even at its default. + function ChainRulesCore.rrule(::typeof(SHTnsKit.analysis), cfg::SHTnsKit.SHTConfig, f; + fft_scratch=nothing, use_rfft::Bool=false) + y = SHTnsKit.analysis(cfg, f; fft_scratch, use_rfft) project_f = ProjectTo(f) function pullback(ȳ) ȳA = _to_complex(ȳ) @@ -61,8 +66,9 @@ import SHTnsKit: wigner_d_matrix_deriv # `_adjoint_synthesis` helper instead. (See `test_adjoint_consistency` # in the test suite for an FD verification.) function ChainRulesCore.rrule(::typeof(SHTnsKit.synthesis), cfg::SHTnsKit.SHTConfig, - alm; real_output::Bool=true) - y = SHTnsKit.synthesis(cfg, alm; real_output) + alm; real_output::Bool=true, + fft_scratch=nothing, use_rfft::Bool=false) + y = SHTnsKit.synthesis(cfg, alm; real_output, fft_scratch, use_rfft) project_alm = ProjectTo(alm) function pullback(ȳ) ȳ_mat = ChainRulesCore.unthunk(ȳ) # materialize Thunk/InplaceableThunk @@ -171,8 +177,9 @@ import SHTnsKit: wigner_d_matrix_deriv # Keep local alias for any direct callers of the ext symbol. const _adjoint_analysis_sphtor = SHTnsKit._adjoint_analysis_sphtor - function ChainRulesCore.rrule(::typeof(SHTnsKit.analysis_sphtor), cfg::SHTnsKit.SHTConfig, Vt, Vp) - Slm, Tlm = SHTnsKit.analysis_sphtor(cfg, Vt, Vp) + function ChainRulesCore.rrule(::typeof(SHTnsKit.analysis_sphtor), cfg::SHTnsKit.SHTConfig, Vt, Vp; + use_rfft::Bool=false) + Slm, Tlm = SHTnsKit.analysis_sphtor(cfg, Vt, Vp; use_rfft) project_Vt = ProjectTo(Vt) project_Vp = ProjectTo(Vp) function pullback(ṠTl) @@ -194,8 +201,8 @@ import SHTnsKit: wigner_d_matrix_deriv const _adjoint_synthesis_sphtor = SHTnsKit._adjoint_synthesis_sphtor function ChainRulesCore.rrule(::typeof(SHTnsKit.synthesis_sphtor), cfg::SHTnsKit.SHTConfig, - Slm, Tlm; real_output::Bool=true) - Vt, Vp = SHTnsKit.synthesis_sphtor(cfg, Slm, Tlm; real_output) + Slm, Tlm; real_output::Bool=true, use_rfft::Bool=false) + Vt, Vp = SHTnsKit.synthesis_sphtor(cfg, Slm, Tlm; real_output, use_rfft) project_Slm = ProjectTo(Slm) project_Tlm = ProjectTo(Tlm) function pullback(Ṽ) @@ -214,6 +221,45 @@ import SHTnsKit: wigner_d_matrix_deriv return (Vt, Vp), pullback end + # QST (3-component) transforms. `synthesis_qst` is the scalar synthesis of Q + # alongside the sphtor synthesis of (S,T), and `analysis_qst` is the mirror, + # so each adjoint is just the two existing adjoints side by side. Without + # these, differentiating a QST pipeline fell through to Zygote's source + # tracing and crashed inside FFTW. + function ChainRulesCore.rrule(::typeof(SHTnsKit.synthesis_qst), cfg::SHTnsKit.SHTConfig, + Qlm, Slm, Tlm; real_output::Bool=true, + use_rfft::Bool=false) + Vr, Vt, Vp = SHTnsKit.synthesis_qst(cfg, Qlm, Slm, Tlm; real_output, use_rfft) + project_Q = ProjectTo(Qlm); project_S = ProjectTo(Slm); project_T = ProjectTo(Tlm) + function pullback(V̄) + zsp() = zeros(Float64, cfg.nlat, cfg.nlon) + V̄r = ChainRulesCore.unthunk(V̄[1]); V̄r = V̄r isa ChainRulesCore.AbstractZero ? zsp() : V̄r + V̄t = ChainRulesCore.unthunk(V̄[2]); V̄t = V̄t isa ChainRulesCore.AbstractZero ? zsp() : V̄t + V̄p = ChainRulesCore.unthunk(V̄[3]); V̄p = V̄p isa ChainRulesCore.AbstractZero ? zsp() : V̄p + Q̄ = SHTnsKit._adjoint_synthesis(cfg, V̄r; real_output=real_output) + S̄, T̄ = SHTnsKit._adjoint_synthesis_sphtor(cfg, V̄t, V̄p; real_output=real_output) + return NoTangent(), NoTangent(), project_Q(Q̄), project_S(S̄), project_T(T̄), + (; real_output=NoTangent(), use_rfft=NoTangent()) + end + return (Vr, Vt, Vp), pullback + end + + function ChainRulesCore.rrule(::typeof(SHTnsKit.analysis_qst), cfg::SHTnsKit.SHTConfig, + Vr, Vt, Vp; use_rfft::Bool=false) + Qlm, Slm, Tlm = SHTnsKit.analysis_qst(cfg, Vr, Vt, Vp; use_rfft) + project_Vr = ProjectTo(Vr); project_Vt = ProjectTo(Vt); project_Vp = ProjectTo(Vp) + function pullback(Ā) + Q̄ = _materialize_coeff(Ā[1], cfg) + S̄ = _materialize_coeff(Ā[2], cfg) + T̄ = _materialize_coeff(Ā[3], cfg) + V̄r = _adjoint_analysis(cfg, Q̄) + V̄t, V̄p = _adjoint_analysis_sphtor(cfg, S̄, T̄) + return NoTangent(), NoTangent(), project_Vr(V̄r), project_Vt(V̄t), project_Vp(V̄p), + (; use_rfft=NoTangent()) + end + return (Qlm, Slm, Tlm), pullback + end + # Complex packed (LM_cplx layout — both signs of m stored explicitly). # # Both transforms are ℂ-linear in their argument (no `real()` anywhere), so @@ -265,7 +311,7 @@ import SHTnsKit: wigner_d_matrix_deriv CT = complex(float(eltype(ā_int))) F̄ = zeros(CT, nlat, nlon) P = Vector{Float64}(undef, lmax + 1) - scaleφ = cfg.cphi + scaleφ = SHTnsKit._analysis_phi_scale(cfg) xv = cfg.x; wv = cfg.w for am in 0:mmax colp = am + 1 @@ -370,16 +416,20 @@ function ChainRulesCore.rrule(::typeof(SHTnsKit.SH_Zrotate), cfg::SHTnsKit.SHTCo end function ChainRulesCore.rrule(::typeof(SHTnsKit.SH_Yrotate), cfg::SHTnsKit.SHTConfig, Qlm, alpha::Real, Rlm) + # Snapshot BEFORE the primal: an in-place rotation (Rlm === Qlm) overwrites + # Qlm, and callers may reuse either buffer before the pullback runs. The dα + # formula needs the *input* coefficients, so they must be preserved here. + Qlm_saved = copy(Qlm) y = SHTnsKit.SH_Yrotate(cfg, Qlm, alpha, Rlm) function pullback(ȳ) inverse = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax) SHTnsKit.shtns_rotation_set_angles_ZYZ(inverse, 0.0, -alpha, 0.0) - Q̄ = similar(Qlm) + Q̄ = similar(Qlm_saved) _configured_rotation_adjoint!(cfg, inverse, ȳ, Q̄) # angle gradient via d/dβ of Wigner-d at β=alpha dα = zero(float(alpha)) lmax, mmax = cfg.lmax, cfg.mmax - Qlm_canonical = SHTnsKit._internal_coefficients(Qlm, cfg) + Qlm_canonical = SHTnsKit._internal_coefficients(Qlm_saved, cfg) ȳ_canonical = SHTnsKit._analysis_cotangent_to_canonical(ȳ, cfg) for l in 0:lmax mm = min(l, mmax) @@ -439,10 +489,14 @@ end # Adjoint for complex rotation using conjugate-transpose of Wigner-D function ChainRulesCore.rrule(::typeof(SHTnsKit.shtns_rotation_apply_cplx), r::SHTnsKit.SHTRotation, Zlm, Rlm) + # Snapshot BEFORE the primal — the per-l blocks are read then written, so an + # in-place call (Rlm === Zlm) is legal and overwrites the input the angle + # gradients below depend on. See the same fix on SH_Zrotate / SH_Yrotate. + Zlm_saved = copy(Zlm) y = SHTnsKit.shtns_rotation_apply_cplx(r, Zlm, Rlm) function pullback(ȳ) lmax, mmax = r.lmax, r.mmax - Z̄ = similar(Zlm) + Z̄ = similar(Zlm_saved) fill!(Z̄, zero(eltype(Z̄))) α, β, γ = _rotation_rrule_angles(r) # This pullback reimplements the Wigner engine, which works in the Y_l^m @@ -457,7 +511,7 @@ function ChainRulesCore.rrule(::typeof(SHTnsKit.shtns_rotation_apply_cplx), r::S # `ȳ` is the closure's own argument and would be safe, but is renamed too # so the pair reads the same way. ε = SHTnsKit._lmcplx_ybasis_signs(lmax, mmax) - Zε = ε .* Zlm + Zε = ε .* Zlm_saved ȳε = ε .* ȳ gα = 0.0; gβ = 0.0; gγ = 0.0 for l in 0:lmax @@ -533,11 +587,16 @@ end # Adjoint for real packed rotation: extend to full, apply cplx adjoint, fold back function ChainRulesCore.rrule(::typeof(SHTnsKit.shtns_rotation_apply_real), r::SHTnsKit.SHTRotation, Qlm, Rlm) + # Snapshot BEFORE the primal: the input is fully expanded into a scratch + # LM_cplx buffer before anything is written back, so an in-place call + # (Rlm === Qlm) is legal and clobbers the coefficients the angle gradients + # below reconstruct `b` from. Same fix as SH_Zrotate / SH_Yrotate. + Qlm_saved = copy(Qlm) y = SHTnsKit.shtns_rotation_apply_real(r, Qlm, Rlm) function pullback(ȳ) lmax, mmax = r.lmax, r.mmax # Extend cotangent on packed to full complex - Zbar_full = zeros(eltype(Qlm), SHTnsKit.nlm_cplx_calc(lmax, mmax, 1)) + Zbar_full = zeros(eltype(Qlm_saved), SHTnsKit.nlm_cplx_calc(lmax, mmax, 1)) for l in 0:lmax mm = min(l, mmax) # m = 0 @@ -588,11 +647,11 @@ function ChainRulesCore.rrule(::typeof(SHTnsKit.shtns_rotation_apply_real), r::S b[mp + l + 1] = (SHTnsKit.LM_cplx_index(lmax, mmax, l, mp) >= 0) ? (begin # reconstruct from packed Qlm if mp == 0 - Qlm[SHTnsKit.LM_index(lmax, 1, l, 0) + 1] + Qlm_saved[SHTnsKit.LM_index(lmax, 1, l, 0) + 1] elseif mp > 0 - Qlm[SHTnsKit.LM_index(lmax, 1, l, mp) + 1] + Qlm_saved[SHTnsKit.LM_index(lmax, 1, l, mp) + 1] else - (-1)^(-mp) * conj(Qlm[SHTnsKit.LM_index(lmax, 1, l, -mp) + 1]) + (-1)^(-mp) * conj(Qlm_saved[SHTnsKit.LM_index(lmax, 1, l, -mp) + 1]) end end) : 0 b[mp + l + 1] *= cis(-mp * γ) @@ -614,7 +673,7 @@ function ChainRulesCore.rrule(::typeof(SHTnsKit.shtns_rotation_apply_real), r::S end end # Fold back to packed positive-m: q̄(m) = Z̄(m) + (-1)^m conj(Z̄(-m)) - Q̄ = zeros(eltype(Qlm), length(Qlm)) + Q̄ = zeros(eltype(Qlm_saved), length(Qlm_saved)) for l in 0:lmax mm = min(l, mmax) # m=0 diff --git a/ext/SHTnsKitGPUExt.jl b/ext/SHTnsKitGPUExt.jl index a2151294..4c0621bc 100644 --- a/ext/SHTnsKitGPUExt.jl +++ b/ext/SHTnsKitGPUExt.jl @@ -628,7 +628,7 @@ function _cuda_scalar_analysis_direct!(owner, cfg::SHTConfig, end backend = CUDABackend() scalar_analysis_kernel!(backend)( - output, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + output, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres, cfg.lmax; ndrange=(cfg.lmax + 1, cfg.mmax + 1), ) @@ -716,7 +716,7 @@ function _cuda_scalar_analysis(cfg::SHTConfig, field::CUDA.AnyCuArray; backend = CUDABackend() canonical = CUDA.zeros(CT, cfg.lmax + 1, cfg.mmax + 1) analyze! = scalar_analysis_kernel!(backend) - analyze!(canonical, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + analyze!(canonical, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres, lcap; ndrange=(lcap + 1, min(cfg.mmax, lcap) + 1)) @@ -846,7 +846,7 @@ function _cuda_vector_analysis_direct!(owner, cfg::SHTConfig, vector_analysis_kernel!(CUDABackend())( Sout, Tout, workspace.Ftheta, workspace.Fphi, tables.dtheta, tables.over_sin, tables.weights, tables.scales, - tables.x, RT(cfg.cphi), lcap, min(cfg.mmax, lcap), cfg.mres, + tables.x, RT(SHTnsKit._analysis_phi_scale(cfg)), lcap, min(cfg.mmax, lcap), cfg.mres, cfg.robert_form; ndrange=(lcap + 1, min(cfg.mmax, lcap) + 1), ) CUDA.synchronize() @@ -1052,7 +1052,7 @@ function _cuda_vector_mode_analysis(cfg::SHTConfig, stored_im::Integer, Tout = similar(Sout) vector_mode_analysis_kernel!(CUDABackend())( Sout, Tout, Vt, Vp, tables.dtheta, tables.over_sin, - tables.weights, tables.scales, tables.x, RT(cfg.cphi), physical_m, + tables.weights, tables.scales, tables.x, RT(SHTnsKit._analysis_phi_scale(cfg)), physical_m, lcap, cfg.robert_form; ndrange=length(Sout), ) CUDA.synchronize() @@ -1196,7 +1196,7 @@ function _cuda_vector_batch_analysis(cfg::SHTConfig, Tout = similar(Sout) vector_batch_analysis_kernel!(CUDABackend())( Sout, Tout, Ft, Fp, tables.dtheta, tables.over_sin, - tables.weights, tables.scales, tables.x, RT(cfg.cphi), cfg.lmax, + tables.weights, tables.scales, tables.x, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres, cfg.robert_form; ndrange=size(Sout), ) CUDA.synchronize() @@ -1532,7 +1532,7 @@ function analysis_axisym(::SHTnsKit.GPU, cfg::SHTConfig, length(field) == cfg.nlat || throw(DimensionMismatch( "field must have length nlat=$(cfg.nlat)", )) - return _cuda_mode_analysis(cfg, 0, field, cfg.lmax, cfg.cphi * cfg.nlon) + return _cuda_mode_analysis(cfg, 0, field, cfg.lmax, SHTnsKit._analysis_phi_scale(cfg) * cfg.nlon) end analysis_axisym(cfg::SHTConfig, field::CUDA.AnyCuArray{T,1}) where {T<:Real} = analysis_axisym(SHTnsKit.GPU(), cfg, field) @@ -1553,7 +1553,7 @@ function analysis_axisym_l(::SHTnsKit.GPU, cfg::SHTConfig, length(field) == cfg.nlat || throw(DimensionMismatch( "field must have length nlat=$(cfg.nlat)", )) - return _cuda_mode_analysis(cfg, 0, field, lcap, cfg.cphi * cfg.nlon) + return _cuda_mode_analysis(cfg, 0, field, lcap, SHTnsKit._analysis_phi_scale(cfg) * cfg.nlon) end analysis_axisym_l(cfg::SHTConfig, field::CUDA.AnyCuArray{T,1}, ltr::Integer) where {T<:Real} = analysis_axisym_l(SHTnsKit.GPU(), cfg, field, ltr) @@ -1591,7 +1591,7 @@ function analysis_packed_ml(::SHTnsKit.GPU, cfg::SHTConfig, im::Int, length(mode) == cfg.nlat || throw(DimensionMismatch( "mode must have length nlat=$(cfg.nlat)", )) - return _cuda_mode_analysis(cfg, physical_m, mode, lcap, cfg.cphi) + return _cuda_mode_analysis(cfg, physical_m, mode, lcap, SHTnsKit._analysis_phi_scale(cfg)) end analysis_packed_ml(cfg::SHTConfig, im::Int, mode::CUDA.AnyCuArray{T,1}, ltr::Integer) where {T<:Complex} = @@ -1628,7 +1628,7 @@ function _cuda_analysis_packed_cplx(cfg::SHTConfig, mcap = min(cfg.mmax, lcap) kernel! = complex_packed_analysis_kernel!(CUDABackend()) kernel!(packed, fourier, tables.Plm, tables.weights, tables.scales, - RT(cfg.cphi), cfg.nlon, lcap, cfg.mmax, mcap; + RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.nlon, lcap, cfg.mmax, mcap; ndrange=(lcap + 1, 2mcap + 1)) CUDA.synchronize() return packed @@ -1710,7 +1710,7 @@ function _cuda_batch_analysis(cfg::SHTConfig, fields::CUDA.AnyCuArray; gpu_fft!(fourier, 2) canonical = CUDA.zeros(CT, cfg.lmax + 1, cfg.mmax + 1, size(fields, 3)) kernel! = scalar_batch_analysis_kernel!(CUDABackend()) - kernel!(canonical, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + kernel!(canonical, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres; ndrange=size(canonical)) configured = canonical ./ reshape( tables.scales, cfg.lmax + 1, cfg.mmax + 1, 1, @@ -1752,7 +1752,7 @@ function _cuda_batch_analysis_direct!(cfg::SHTConfig, end backend = CUDABackend() scalar_batch_analysis_kernel!(backend)( - output, fourier, tables.Plm, tables.weights, RT(cfg.cphi), + output, fourier, tables.Plm, tables.weights, RT(SHTnsKit._analysis_phi_scale(cfg)), cfg.lmax, cfg.mmax, cfg.mres; ndrange=size(output), ) coefficient_batch_conversion_kernel!(backend)( diff --git a/ext/SHTnsKitLoopVecExt.jl b/ext/SHTnsKitLoopVecExt.jl index 4701577f..2fa2a865 100644 --- a/ext/SHTnsKitLoopVecExt.jl +++ b/ext/SHTnsKitLoopVecExt.jl @@ -7,6 +7,41 @@ using Base.Threads: @threads # Turbo-optimized variants live under the SHTnsKit namespace so users can call # SHTnsKit.analysis_turbo, synthesis_turbo, etc., when LoopVectorization is loaded. +""" + _turbo_threads_ok() -> Bool + +Whether this call may start its own threaded loop. `@threads :static` cannot be +nested or started concurrently, so a turbo transform invoked from inside an +outer threaded region (or from a worker task) must run serially on that worker. +Delegates to the same predicate the core orchestrators use, so the two agree. +""" +@inline _turbo_threads_ok() = SHTnsKit._use_internal_mloop_threads() + +"""Run a `for` loop with the given `@threads` schedule when safe, else serially.""" +macro _lv_threads(sched, loop) + loop isa Expr && loop.head === :for || + throw(ArgumentError("@_lv_threads requires a for loop")) + return esc(quote + if $(_turbo_threads_ok)() + Base.Threads.@threads $sched $loop + else + $loop + end + end) +end + +""" + _turbo_m_order(cfg) -> Vector{Int} + +The azimuthal orders this config actually represents: `0, mres, 2mres, …`. +Iterating a bare `0:mmax` here made `analysis_turbo` populate — and +`synthesis_turbo` consume — coefficient columns that an `mres > 1` transform has +no storage for, so the turbo pair silently disagreed with `analysis`/`synthesis` +(the disagreement is O(1), not roundoff). Shares the core's cached, load-balanced +ordering so the two stay in lockstep. +""" +@inline _turbo_m_order(cfg::SHTnsKit.SHTConfig) = SHTnsKit.cached_m_order(cfg) + """ SHTnsKit.analysis_turbo(cfg::SHTnsKit.SHTConfig, f::AbstractMatrix) @@ -18,27 +53,30 @@ function SHTnsKit.analysis_turbo(cfg::SHTnsKit.SHTConfig, f::AbstractMatrix) size(f, 1) == nlat || throw(DimensionMismatch("first dim must be nlat=$(nlat)")) size(f, 2) == nlon || throw(DimensionMismatch("second dim must be nlon=$(nlon)")) - fC = complex.(f) - Fφ = SHTnsKit.fft_phi(fC) + # One buffer + the shared cached FFTW plan, matching the core `analysis`. + # The old `fft_phi(complex.(f))` made a complex copy AND re-planned an + # out-of-place FFT on every call. + Fφ = SHTnsKit.fft_phi!(Matrix{complex(float(eltype(f)))}(undef, nlat, nlon), f) lmax, mmax = cfg.lmax, cfg.mmax CT = eltype(Fφ) alm = Matrix{CT}(undef, lmax + 1, mmax + 1) fill!(alm, 0.0 + 0.0im) - scaleφ = cfg.cphi + scaleφ = SHTnsKit._analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Bind cfg fields to locals so the @tturbo loops below operate on plain arrays. # LoopVectorization can't analyze property access (cfg.Nlm) inside @tturbo, and cfg is mutable. xv = cfg.x; wv = cfg.w # Adaptive threading: use nested parallelism for better load balancing n_threads = Threads.nthreads() - if mmax + 1 < n_threads ÷ 2 && nlat > 32 + m_order = _turbo_m_order(cfg) + if length(m_order) < n_threads ÷ 2 && nlat > 32 # Few m modes: parallelize over latitude points with thread-local accumulators n_tid = Threads.maxthreadid() CT = eltype(alm) # eltype-derived accumulators (this branch uses plain loops, not @tturbo) thread_alm = [Vector{CT}(undef, lmax + 1) for _ in 1:n_tid] thread_P_bufs = [Vector{Float64}(undef, lmax + 1) for _ in 1:n_tid] # per-thread Legendre scratch (hoisted out of the latitude loop) - for m in 0:mmax + for m in m_order col = m + 1 for t in 1:n_tid fill!(thread_alm[t], zero(CT)) @@ -46,7 +84,7 @@ function SHTnsKit.analysis_turbo(cfg::SHTnsKit.SHTConfig, f::AbstractMatrix) if cfg.use_plm_tables && length(cfg.NP_tables) == mmax + 1 # NP_tables[col][l+1, i] = P̄_l^m already; no extra Nlm multiply tbl = cfg.NP_tables[m + 1] - @threads :static for i in 1:nlat # :static pins iterations → threadid() stable (no buffer race under task migration) + @_lv_threads :static for i in 1:nlat # :static pins iterations → threadid() stable (no buffer race under task migration) tid = Threads.threadid() local_acc = thread_alm[tid] Fi = Fφ[i, col] @@ -56,7 +94,7 @@ function SHTnsKit.analysis_turbo(cfg::SHTnsKit.SHTConfig, f::AbstractMatrix) end end else - @threads :static for i in 1:nlat # :static pins iterations → threadid() stable (no buffer race under task migration) + @_lv_threads :static for i in 1:nlat # :static pins iterations → threadid() stable (no buffer race under task migration) tid = Threads.threadid() local_acc = thread_alm[tid] thread_P = thread_P_bufs[tid] @@ -79,7 +117,8 @@ function SHTnsKit.analysis_turbo(cfg::SHTnsKit.SHTConfig, f::AbstractMatrix) end else # Standard m-parallel approach with dynamic scheduling - @threads :dynamic for m in 0:mmax + @_lv_threads :dynamic for idx in eachindex(m_order) + m = m_order[idx] col = m + 1 if cfg.use_plm_tables && length(cfg.NP_tables) == mmax + 1 # NP_tables[col][l+1, i] = P̄_l^m already; no extra Nlm multiply @@ -124,9 +163,11 @@ function SHTnsKit.synthesis_turbo(cfg::SHTnsKit.SHTConfig, alm::AbstractMatrix; alm_int = SHTnsKit._internal_coefficients(alm, cfg) nlat, nlon = cfg.nlat, cfg.nlon - CT = eltype(alm_int) + # Fourier bins are complex even when every spectral coefficient is real + # (matches the core `_synthesis`). + CT = complex(float(eltype(alm_int))) Fφ = Matrix{CT}(undef, nlat, nlon) - fill!(Fφ, 0.0 + 0.0im) + fill!(Fφ, zero(CT)) inv_scaleφ = SHTnsKit.phi_inv_scale(cfg) # Bind cfg fields to locals so the @tturbo loops below operate on plain arrays. @@ -135,17 +176,18 @@ function SHTnsKit.synthesis_turbo(cfg::SHTnsKit.SHTConfig, alm::AbstractMatrix; # Adaptive threading: use nested parallelism for better load balancing n_threads = Threads.nthreads() - if mmax + 1 < n_threads ÷ 2 && nlat > 32 + m_order = _turbo_m_order(cfg) + if length(m_order) < n_threads ÷ 2 && nlat > 32 # Few m modes: parallelize over latitude points instead n_tid = Threads.maxthreadid() thread_P_bufs = [Vector{Float64}(undef, lmax + 1) for _ in 1:n_tid] # per-thread Legendre scratch (hoisted out of the latitude loop) G = Vector{CT}(undef, nlat) # shared latitude scratch (few-m branch: outer m-loop is serial) - for m in 0:mmax + for m in m_order col = m + 1 if cfg.use_plm_tables && length(cfg.NP_tables) == mmax + 1 # NP_tables[col][l+1, i] = P̄_l^m already; no extra Nlm multiply tbl = cfg.NP_tables[m + 1] - @threads for i in 1:nlat + @_lv_threads :dynamic for i in 1:nlat g_re = 0.0 g_im = 0.0 @tturbo warn_check_args=false for l in m:lmax @@ -157,7 +199,7 @@ function SHTnsKit.synthesis_turbo(cfg::SHTnsKit.SHTConfig, alm::AbstractMatrix; G[i] = complex(g_re, g_im) end else - @threads :static for i in 1:nlat # :static pins iterations → threadid() stable + @_lv_threads :static for i in 1:nlat # :static pins iterations → threadid() stable thread_P = thread_P_bufs[Threads.threadid()] SHTnsKit.Plm_norm_row!(thread_P, xv[i], lmax, m) g_re = 0.0 @@ -186,7 +228,8 @@ function SHTnsKit.synthesis_turbo(cfg::SHTnsKit.SHTConfig, alm::AbstractMatrix; # Per-m latitude scratch as columns of one pre-allocated matrix (each m owns # a distinct column → race-free, no per-iteration allocation). Gcols = Matrix{CT}(undef, nlat, mmax + 1) - @threads :dynamic for m in 0:mmax + @_lv_threads :dynamic for idx in eachindex(m_order) + m = m_order[idx] col = m + 1 thread_G = view(Gcols, :, col) if cfg.use_plm_tables && length(cfg.NP_tables) == mmax + 1 diff --git a/ext/SHTnsKitParallelAMDGPUExt.jl b/ext/SHTnsKitParallelAMDGPUExt.jl index 0b6d74cf..9a040698 100644 --- a/ext/SHTnsKitParallelAMDGPUExt.jl +++ b/ext/SHTnsKitParallelAMDGPUExt.jl @@ -66,7 +66,7 @@ function ParallelExt._dist_transpose_gpu_analysis!( AMDGPU.ROCBackend(), ) kernel!(parent(output), parent(plan.F_buf), tables.Plm, tables.weights, - tables.scales, RT(plan.cfg.cphi), _first_m(plan), plan.lmax, + tables.scales, RT(SHTnsKit._analysis_phi_scale(plan.cfg)), _first_m(plan), plan.lmax, plan.mmax, plan.cfg.mres, plan.lmax; ndrange=size(parent(output))) AMDGPU.synchronize() @@ -111,7 +111,7 @@ function ParallelExt._dist_transpose_gpu_vector_analysis!( ) kernel!(parent(Sout), parent(Tout), parent(plan.F_buf), parent(plan.F_buf2), tables.dtheta, tables.over_sin, - tables.weights, tables.scales, tables.x, RT(plan.cfg.cphi), + tables.weights, tables.scales, tables.x, RT(SHTnsKit._analysis_phi_scale(plan.cfg)), _first_m(plan), plan.lmax, plan.mmax, plan.cfg.mres, plan.cfg.robert_form; ndrange=size(parent(Sout))) diff --git a/ext/SHTnsKitParallelCUDAExt.jl b/ext/SHTnsKitParallelCUDAExt.jl index cadceadf..569c4af5 100644 --- a/ext/SHTnsKitParallelCUDAExt.jl +++ b/ext/SHTnsKitParallelCUDAExt.jl @@ -58,7 +58,7 @@ function ParallelExt._dist_transpose_gpu_analysis!( CUDA.CUDABackend(), ) kernel!(parent(output), parent(plan.F_buf), tables.Plm, tables.weights, - tables.scales, RT(plan.cfg.cphi), _first_m(plan), plan.lmax, + tables.scales, RT(SHTnsKit._analysis_phi_scale(plan.cfg)), _first_m(plan), plan.lmax, plan.mmax, plan.cfg.mres, plan.lmax; ndrange=size(parent(output))) CUDA.synchronize() @@ -103,7 +103,7 @@ function ParallelExt._dist_transpose_gpu_vector_analysis!( ) kernel!(parent(Sout), parent(Tout), parent(plan.F_buf), parent(plan.F_buf2), tables.dtheta, tables.over_sin, - tables.weights, tables.scales, tables.x, RT(plan.cfg.cphi), + tables.weights, tables.scales, tables.x, RT(SHTnsKit._analysis_phi_scale(plan.cfg)), _first_m(plan), plan.lmax, plan.mmax, plan.cfg.mres, plan.cfg.robert_form; ndrange=size(parent(Sout))) diff --git a/ext/SHTnsKitParallelExt.jl b/ext/SHTnsKitParallelExt.jl index 6e18befa..37ba147b 100644 --- a/ext/SHTnsKitParallelExt.jl +++ b/ext/SHTnsKitParallelExt.jl @@ -63,7 +63,9 @@ DEBUGGING TIPS ENVIRONMENT VARIABLES -------------------- -- SHTNSKIT_CACHE_PENCILFFTS: "1" (default) to cache FFT plans, "0" to disable +- SHTNSKIT_FFT_PLAN_CACHE: "1" (default) to cache φ-FFT plans, "0" to disable. + Read by SHTnsKit proper — the cache is shared with the serial transforms. + (legacy alias: SHTNSKIT_CACHE_PENCILFFTS) - SHTNSKIT_VERBOSE_STORAGE: "1" to print storage optimization info ================================================================================ =# @@ -75,8 +77,9 @@ Parallel extension module providing MPI-distributed spherical harmonic transform See module-level comments for architecture overview and debugging tips. # Module state -The extension keeps its FFT plan caches, locks, and cache controls in direct -module constants. +The extension keeps its locks and caches in direct module constants. The φ-FFT +plan cache itself lives in SHTnsKit proper (src/fftutils.jl) and is shared with +the serial transforms. """ using Base.Threads # Threads.@threads and locks/macros @@ -132,119 +135,20 @@ import SHTnsKit # Core spherical harmonic functionality end # ===== MODULE STATE ===== -const _CACHE_PENCILFFTS = Ref(get(ENV, "SHTNSKIT_CACHE_PENCILFFTS", "1") == "1") -const _pfft_cache = IdDict{Any,Any}() -const _PFFT_CACHE_MAX = Ref(parse(Int, get(ENV, "SHTNSKIT_PFFT_CACHE_MAX", "64"))) -const _cache_lock = Threads.ReentrantLock() -const _fftw_cache_lock = Threads.ReentrantLock() - -""" - pfft_cache_max!(n::Int) -> Int - -Set the maximum number of cached FFT plans (shared across all grids and -communicators). Set `n <= 0` to disable the cap entirely. Returns the previous -value. -""" -function pfft_cache_max!(n::Int) - prev = _PFFT_CACHE_MAX[] - _PFFT_CACHE_MAX[] = n - return prev -end +# The φ-FFT plan cache (and its `enable_fft_plan_cache!` / `disable_fft_plan_cache!` +# / `set_fft_plan_cache!` / `fft_plan_cache_enabled` controls) lives in +# src/fftutils.jl and is shared with the serial transforms. This extension used to +# carry a second, parallel-only cache keyed on the pencil decomposition — but the +# only function that ever consulted it, `_get_or_plan`, had no call sites, and the +# "plans" it stored were `NamedTuple` placeholders that the FFT wrappers ignored. +# Every knob pointed at it was therefore inert. It has been deleted rather than +# left in place; recover it from git history if a PencilFFTs-level cache is ever +# actually needed. # Compat helper: `ceildiv` was added in Julia 1.11 const _ceildiv = isdefined(Base, :ceildiv) ? Base.ceildiv : (a, b) -> cld(a, b) ceildiv(a::Integer, b::Integer) = _ceildiv(a, b) -function _fft_plan_cache_enabled_impl() - return _CACHE_PENCILFFTS[] -end - -function _fft_plan_cache_set_impl(flag::Bool; clear::Bool=true) - _CACHE_PENCILFFTS[] = flag - if !flag && clear - lock(_cache_lock) do - empty!(_pfft_cache) - end - end - return flag -end - -function _fft_plan_cache_enable_impl() - return _fft_plan_cache_set_impl(true) -end - -function _fft_plan_cache_disable_impl(; clear::Bool=true) - return _fft_plan_cache_set_impl(false; clear=clear) -end - -@inline function _decomp_hash(A) - if hasfield(typeof(A), :pencil) - pencil = getfield(A, :pencil) - if hasfield(typeof(pencil), :decomposition) - return hash(getfield(pencil, :decomposition)) - elseif hasfield(typeof(pencil), :plan) - return hash(getfield(pencil, :plan)) - end - end - return hash(size(A)) -end - -# Generate cache key based on array characteristics for FFT plan reuse -function _cache_key(kind::Symbol, A) - # Basic array characteristics - base_key = (kind, size(A,1), size(A,2), eltype(A)) - - # Add communicator size with robust error handling - comm_size = try - MPI.Comm_size(communicator(A)) - catch - 1 # Default to single process - end - - # Decomposition hash — no try/catch to avoid closure-box allocations on hot path. - decomp_hash = _decomp_hash(A) - - return (base_key..., comm_size, decomp_hash) -end - -function _get_or_plan(kind::Symbol, A) - # If caching disabled, create plan directly without storing - if !_CACHE_PENCILFFTS[] - return kind === :fft ? plan_fft(A; dims=2) : # Forward FFT along longitude (dim 2) - kind === :ifft ? plan_ifft(A; dims=2) : # Inverse FFT along longitude - kind === :rfft ? (try plan_rfft(A; dims=2) catch; nothing end) : # Real-to-complex FFT - kind === :irfft ? (try plan_irfft(A; dims=2) catch; nothing end) : # Complex-to-real IFFT - error("unknown plan kind") - end - - # Thread-safe caching with optimized lookup - key = _cache_key(kind, A) - - # Thread-safe plan lookup and creation - return lock(_cache_lock) do - # Double-check pattern: another thread might have created the plan - if haskey(_pfft_cache, key) - return _pfft_cache[key] - end - - # Create new plan and cache it for future use - plan = kind === :fft ? plan_fft(A; dims=2) : # Forward FFT along longitude - kind === :ifft ? plan_ifft(A; dims=2) : # Inverse FFT along longitude - kind === :rfft ? (try plan_rfft(A; dims=2) catch; nothing end) : # Real-to-complex FFT - kind === :irfft ? (try plan_irfft(A; dims=2) catch; nothing end) : # Complex-to-real IFFT - error("unknown plan kind") - - # Enforce the soft cap: flush before inserting so the fresh entry survives. - cap = _PFFT_CACHE_MAX[] - if cap > 0 && length(_pfft_cache) >= cap - empty!(_pfft_cache) - end - _pfft_cache[key] = plan - return plan - end -end - - # ===== PENCIL GRID SUGGESTION ===== @inline function _infer_comm_size(comm_or_nprocs::Any) comm_or_nprocs === nothing && return 1 @@ -404,44 +308,6 @@ end # Use FFTW for 1D FFTs along the longitude dimension (not PencilFFTs which is for multi-D) # PencilArrays provides the distributed array framework, FFTW provides the FFTs -# Cache for FFTW 1D plans (key includes inplace flag) -const _fftw_plan_cache = Dict{Tuple{Symbol, Int, DataType, Bool}, Any}() - -""" - get_fftw_plan(kind, n, T) -> plan - -Get or create a cached FFTW plan for 1D transforms. -""" -function get_fftw_plan(kind::Symbol, n::Int, ::Type{T}; inplace::Bool=false) where T - key = (kind, n, T, inplace) - lock(_fftw_cache_lock) do - if haskey(_fftw_plan_cache, key) - return _fftw_plan_cache[key] - end - - # Create sample array for planning - if kind == :fft - sample = zeros(Complex{real(T)}, n) - plan = inplace ? FFTW.plan_fft!(sample) : FFTW.plan_fft(sample) - elseif kind == :ifft - sample = zeros(Complex{real(T)}, n) - plan = inplace ? FFTW.plan_ifft!(sample) : FFTW.plan_ifft(sample) - elseif kind == :rfft - sample = zeros(real(T), n) - plan = FFTW.plan_rfft(sample) # rfft is always out-of-place - elseif kind == :irfft - # For irfft, input size is n÷2+1 - sample = zeros(Complex{real(T)}, n ÷ 2 + 1) - plan = FFTW.plan_irfft(sample, n) - else - error("Unknown FFT kind: $kind") - end - - _fftw_plan_cache[key] = plan - return plan - end -end - """ fft_along_dim2!(output, input) @@ -530,71 +396,6 @@ function ifft_along_dim2!(output::AbstractMatrix{Complex{T}}, input::AbstractMat return output end -# Local FFT wrappers used by the extension's plan cache. -function plan_fft(A::PencilArray; dims=:) - # Return a placeholder that indicates we'll use FFTW on local data - return (kind=:fft, local_size=size(parent(A))) -end - -function plan_ifft(A::PencilArray; dims=:) - return (kind=:ifft, local_size=size(parent(A))) -end - -function fft(A::PencilArray, p) - local_data = parent(A) - nlat, nlon = size(local_data) - output = similar(local_data, Complex{Float64}) - fft_along_dim2!(output, local_data) - return output -end - -function ifft(A::PencilArray, p) - local_data = parent(A) - nlat, nlon = size(local_data) - output = similar(local_data) - ifft_along_dim2!(output, local_data) - return output -end - -# RFFT/IRFFT variants -function plan_rfft(A::PencilArray; dims=:) - return (kind=:rfft, local_size=size(parent(A))) -end - -function plan_irfft(A::PencilArray; dims=:) - return (kind=:irfft, local_size=size(parent(A))) -end - -function rfft(A::PencilArray, p) - local_data = parent(A) - nlat, nlon = size(local_data) - nk = nlon ÷ 2 + 1 - output = Matrix{ComplexF64}(undef, nlat, nk) - @inbounds for i in 1:nlat - row = Vector{Float64}(collect(view(local_data, i, :))) - fft_result = FFTW.rfft(row) - for j in 1:nk - output[i, j] = fft_result[j] - end - end - return output -end - -function irfft(A::AbstractMatrix{<:Complex}, p) - nlat, nk = size(A) - # Assume original nlon was 2*(nk-1) for even-length arrays - nlon = 2 * (nk - 1) - output = Matrix{Float64}(undef, nlat, nlon) - @inbounds for i in 1:nlat - row = Vector{ComplexF64}(collect(view(A, i, :))) - ifft_result = FFTW.irfft(row, nlon) - for j in 1:nlon - output[i, j] = ifft_result[j] - end - end - return output -end - # ===== OPTIMIZED DISTRIBUTED FFT USING TRANSPOSE ===== # When φ is distributed, use a single all-to-all transpose instead of per-row Allgatherv. # This reduces the number of MPI calls from O(nlat) to O(1). @@ -710,38 +511,6 @@ function distributed_rfft_phi!(Fθm_out::AbstractMatrix{Complex{T}}, return Fθm_out end -""" - distributed_irfft_phi!(local_out, Fθm, θ_range, φ_range, nlon, comm) - -Complex-to-real inverse FFT for distributed synthesis. `Fθm` is `(nlat_local, -nlon÷2+1)` and must be identical on every rank in a given θ-slab (caller -responsibility — typical pattern replicates the Fourier buffer). After local -`irfft` to full `(nlat_local, nlon)` real, the function slices this rank's -local φ window into `local_out`. -""" -function distributed_irfft_phi!(local_out::AbstractMatrix{<:Real}, - Fθm::AbstractMatrix{<:Complex}, - θ_range::AbstractRange, φ_range::AbstractRange, - nlon::Int, comm) - nlat_local = length(θ_range) - nlon_local = length(φ_range) - size(Fθm, 2) == nlon ÷ 2 + 1 || throw(DimensionMismatch("Fθm must have nlon÷2+1 columns")) - size(Fθm, 1) == nlat_local || throw(DimensionMismatch("Fθm must have nlat_local rows")) - size(local_out) == (nlat_local, nlon_local) || throw(DimensionMismatch("local_out must be (nlat_local, nlon_local)")) - - spatial_full = Matrix{eltype(local_out)}(undef, nlat_local, nlon) - spatial_full .= FFTW.irfft(Fθm, nlon, 2) - - φ_start = first(φ_range) - @inbounds for j in 1:nlon_local - for i in 1:nlat_local - local_out[i, j] = spatial_full[i, φ_start + j - 1] - end - end - - return local_out -end - """ distributed_ifft_phi!(local_out, Fθm, θ_range, φ_range, nlon, comm) @@ -801,52 +570,6 @@ function efficient_spectral_reduce!(local_data::AbstractVector, comm) return local_data end -""" - bandwidth_aware_broadcast!(data, root, comm) - -Bandwidth-aware broadcasting that adapts to network topology and data size. -Uses pipeline broadcasting for large data and tree broadcasting for small data. -""" -function bandwidth_aware_broadcast!(data::AbstractArray, root::Int, comm) - nprocs = MPI.Comm_size(comm) - data_size_mb = (sizeof(data)) / (1024 * 1024) - - if nprocs > 32 && data_size_mb > 10.0 - # Use pipeline broadcast for large data on large clusters - pipeline_broadcast!(data, root, comm) - else - # Use standard tree broadcast for smaller cases - MPI.Bcast!(data, root, comm) - end - - return data -end - -""" - pipeline_broadcast!(data, root, comm) - -Pipeline broadcast that overlaps communication with local copying for better bandwidth utilization. -""" -function pipeline_broadcast!(data::AbstractArray, root::Int, comm) - rank = MPI.Comm_rank(comm) - nprocs = MPI.Comm_size(comm) - - # Determine pipeline parameters - pipeline_stages = min(nprocs, 8) # Limit pipeline depth - chunk_size = max(1, length(data) ÷ pipeline_stages) - - for stage in 1:pipeline_stages - start_idx = (stage - 1) * chunk_size + 1 - end_idx = stage == pipeline_stages ? length(data) : stage * chunk_size - chunk_view = view(data, start_idx:end_idx) - - # Pipeline broadcast of this chunk - MPI.Bcast!(chunk_view, root, comm) - end - - return data -end - # Note: Avoid forwarding Base.zeros(Pencil) to PencilArrays.zeros to prevent # potential recursion when PencilArrays.zeros may call Base.zeros internally. diff --git a/ext/SHTnsKitZygoteExt.jl b/ext/SHTnsKitZygoteExt.jl index 70b94a47..7bf35d16 100644 --- a/ext/SHTnsKitZygoteExt.jl +++ b/ext/SHTnsKitZygoteExt.jl @@ -286,16 +286,20 @@ Zygote.@adjoint function SHTnsKit.SH_Zrotate(cfg::SHTnsKit.SHTConfig, Qlm::Abstr end Zygote.@adjoint function SHTnsKit.SH_Yrotate(cfg::SHTnsKit.SHTConfig, Qlm::AbstractVector{<:Complex}, alpha::Real, Rlm::AbstractVector{<:Complex}) + # Snapshot BEFORE the primal: an in-place rotation (Rlm === Qlm) overwrites + # Qlm, and callers may reuse either buffer before the pullback runs. The dα + # formula needs the *input* coefficients, so they must be preserved here. + Qlm_saved = copy(Qlm) y = SHTnsKit.SH_Yrotate(cfg, Qlm, alpha, Rlm) function back(ȳ) inverse = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax) SHTnsKit.shtns_rotation_set_angles_ZYZ(inverse, 0.0, -alpha, 0.0) - Q̄ = similar(Qlm) + Q̄ = similar(Qlm_saved) _zyg_configured_rotation_adjoint!(cfg, inverse, ȳ, Q̄) # angle gradient via derivative of Wigner-d at beta=alpha dα = zero(float(alpha)) lmax, mmax = cfg.lmax, cfg.mmax - Qlm_canonical = SHTnsKit._internal_coefficients(Qlm, cfg) + Qlm_canonical = SHTnsKit._internal_coefficients(Qlm_saved, cfg) ȳ_canonical = SHTnsKit._analysis_cotangent_to_canonical(ȳ, cfg) for l in 0:lmax mm = min(l, mmax) diff --git a/src/SHTnsKit.jl b/src/SHTnsKit.jl index 6c64747a..77b40794 100644 --- a/src/SHTnsKit.jl +++ b/src/SHTnsKit.jl @@ -101,7 +101,9 @@ ENVIRONMENT VARIABLES -------------------- - SHTNSKIT_PHI_SCALE: "dft" or "quad" for φ scaling convention - SHTNSKIT_VERBOSE_STORAGE: "1" to print storage optimization info -- SHTNSKIT_CACHE_PENCILFFTS: "0" to disable FFT plan caching (parallel ext) +- SHTNSKIT_FFT_PLAN_CACHE: "0" to disable φ-FFT plan caching (serial and distributed) + (legacy alias: SHTNSKIT_CACHE_PENCILFFTS) +- SHTNSKIT_FFT_PLAN_CACHE_MAX: cap on distinct cached plans (default 64; ≤0 = no cap) DEBUGGING TIPS -------------- @@ -146,13 +148,45 @@ function phi_inv_scale(cfg::SHTConfig) end if cfg.phi_scale === :quad return cfg.nlon / (2π) - elseif cfg.phi_scale === :dft - return Float64(cfg.nlon) else - return cfg.grid_type == :gauss ? Float64(cfg.nlon) : cfg.nlon / (2π) + # `:dft` and anything unset (`:auto`). The old fallback keyed on + # `grid_type` and handed every non-Gauss grid `nlon/2π` — so a regular + # grid built through the exported keyword constructor (which defaulted to + # `:auto`) disagreed by 2π with the identical grid from + # `create_regular_config`, which sets `:dft` explicitly. Both constructors + # emit `:dft`, so `:dft` is the right default for an unset value. + return Float64(cfg.nlon) end end +""" + _analysis_phi_scale(cfg) -> Float64 + +The φ quadrature factor `analysis` must apply so that it inverts `synthesis`. + +`synthesis` scales its Fourier bins by `phi_inv_scale(cfg)` and the inverse FFT +divides by `nlon`, a net spatial factor of `σ = phi_inv_scale(cfg)/nlon`. For the +pair to be mutually inverse, analysis must carry `cphi/σ`. + +Under the default `:dft` mode `σ = 1` and this is just `cphi = 2π/nlon`, exactly +what analysis always used — so nothing changes for any configuration the +`create_*_config` constructors produce. Under `:quad` (`σ = 1/2π`) the old fixed +`cphi` made `analysis(synthesis(alm))` come back as `alm/2π`: the two halves of +the transform pair simply disagreed about the convention, with synthesis honouring +`phi_scale` and analysis ignoring it. +""" +@inline _analysis_phi_scale(cfg::SHTConfig) = cfg.cphi * cfg.nlon / phi_inv_scale(cfg) + +""" + _evaluator_phi_scale(cfg) -> Float64 + +The factor a point/latitude evaluator must apply to reproduce the value +`synthesis` writes on the grid: `phi_inv_scale(cfg)/nlon`, i.e. 1 under `:dft` +and `1/2π` under `:quad`. Omitting it made every direct evaluator disagree with +the very grid it claims to sample by a factor of 2π. +""" +@inline _evaluator_phi_scale(cfg::SHTConfig) = phi_inv_scale(cfg) / cfg.nlon + include("buffer_utils.jl") # Common buffer allocation patterns include("kernels.jl") # Legendre accumulation kernels include("plan.jl") # Transform planning and optimization @@ -185,6 +219,7 @@ export set_allow_padding!, disable_padding!, is_padding_enabled # Memory p export get_nlat_padded, get_spat_dist, compute_optimal_padding # Padding queries export allocate_padded_spatial, allocate_padded_spatial_batch # Padded array allocation export copy_to_padded!, copy_from_padded!, estimate_padding_overhead # Padding utilities +export spatial_view # Padded buffer → transform-shaped view # ===== BASIC TRANSFORMS ===== # The `*_cplx` helpers are intentionally separate from `real_output=false` @@ -340,52 +375,15 @@ function DistQstPlan(args...; kwargs...) return getproperty(ext, :DistQstPlan)(args...; kwargs...) end -function fft_plan_cache_enabled() - ext = _parallel_ext_module() - return ext === nothing ? false : getproperty(ext, :_fft_plan_cache_enabled_impl)() -end - -function set_fft_plan_cache!(flag::Bool; clear::Bool=true) - ext = _parallel_ext_module() - ext === nothing && error("Parallel extension not loaded") - return getproperty(ext, :_fft_plan_cache_set_impl)(flag; clear=clear) -end - -function enable_fft_plan_cache!() - ext = _parallel_ext_module() - ext === nothing && error("Parallel extension not loaded") - return getproperty(ext, :_fft_plan_cache_enable_impl)() -end - -function disable_fft_plan_cache!(; clear::Bool=true) - ext = _parallel_ext_module() - ext === nothing && error("Parallel extension not loaded") - return getproperty(ext, :_fft_plan_cache_disable_impl)(; clear=clear) -end - -Base.@doc """ - fft_plan_cache_enabled() -> Bool - -Return whether distributed FFT plan caching is currently enabled. -""" fft_plan_cache_enabled - -Base.@doc """ - set_fft_plan_cache!(flag::Bool; clear::Bool=true) - -Enable or disable caching of distributed FFT plans. When disabling and `clear=true`, cached plans are freed. -""" set_fft_plan_cache! - -Base.@doc """ - enable_fft_plan_cache!() - -Convenience wrapper to enable distributed FFT plan caching. -""" enable_fft_plan_cache! - -Base.@doc """ - disable_fft_plan_cache!(; clear::Bool=true) - -Disable distributed FFT plan caching. Pass `clear=false` to retain existing cache entries. -""" disable_fft_plan_cache! +# NOTE: the φ-FFT plan cache and its `fft_plan_cache_enabled` / +# `set_fft_plan_cache!` / `enable_fft_plan_cache!` / `disable_fft_plan_cache!` +# controls live in src/fftutils.jl. They used to forward to the parallel +# extension, where the cache they addressed had no readers at all — `_get_or_plan` +# was never called from anywhere, so every one of these knobs (and the +# `SHTNSKIT_CACHE_PENCILFFTS` environment variable the distributed guide +# advertises) was a no-op. The cache that every transform actually uses, serial +# and distributed alike, is the one in fftutils.jl, so the controls now address +# that and no longer require the extension to be loaded. # ===== PENCIL GRID SUGGESTION ===== function _suggest_pencil_grid_fallback(comm_or_nprocs::Any, nlat::Integer, nlon::Integer; diff --git a/src/batch_transforms.jl b/src/batch_transforms.jl index 7d65174b..72798e7b 100644 --- a/src/batch_transforms.jl +++ b/src/batch_transforms.jl @@ -228,9 +228,19 @@ end """ set_batch_size!(cfg::SHTConfig, howmany::Int; spec_dist::Int=0) -Configure batch processing for multiple fields. After calling this function, -batch transform functions (`analysis_batch`, `synthesis_batch`, etc.) will -process `howmany` fields simultaneously. +Record the SHTns-style batched-layout descriptors on the configuration. + +!!! note "Advisory metadata" + `howmany` and `spec_dist` mirror the SHTns C API's batch descriptors and are + stored for interoperability, but the Julia batch entry points do **not** + consult them: `analysis_batch`, `synthesis_batch` and their sphtor/QST + siblings take the field count from `size(fields, 3)` of the array you pass. + Setting a batch size neither constrains nor accelerates those calls, and + passing a different number of fields is allowed and works. + +Configure the batch descriptors carried by `cfg` for callers that mirror the +SHTns C layout. See the note above for what the Julia batch transforms actually +use. # Arguments - `cfg`: SHTConfig to modify @@ -245,11 +255,10 @@ process `howmany` fields simultaneously. # Example ```julia cfg = create_gauss_config(32, 34) -set_batch_size!(cfg, 4) # Process 4 fields at once +set_batch_size!(cfg, 4) # recorded on cfg; does not bind the call below -# Now use batch transforms fields = rand(cfg.nlat, cfg.nlon, 4) -alms = analysis_batch(cfg, fields) +alms = analysis_batch(cfg, fields) # count comes from size(fields, 3) ``` """ function set_batch_size!(cfg::SHTConfig, howmany::Int; spec_dist::Int=0) @@ -322,7 +331,7 @@ function analysis_batch(cfg::SHTConfig, fields::AbstractArray{<:Real,3}; use_rff _batch_fft_phi!(Fφ_batch, fields) end - scaleφ = cfg.cphi + scaleφ = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Hoist cfg field reads to locals: cfg is mutable, so reads inside the m/l loops below aren't auto-hoisted. w = cfg.w @@ -423,7 +432,7 @@ function analysis_batch!(cfg::SHTConfig, alm_out::AbstractArray{<:Complex,3}, # mutating an output view that overlaps the spatial input is safe. fill!(alm_out, zero(eltype(alm_out))) - scaleφ = cfg.cphi + scaleφ = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Hoist cfg field reads to locals: cfg is mutable, so reads inside the m/l loops below aren't auto-hoisted. w = cfg.w diff --git a/src/complex_packed.jl b/src/complex_packed.jl index 017e9980..16d57246 100644 --- a/src/complex_packed.jl +++ b/src/complex_packed.jl @@ -179,7 +179,7 @@ function _analysis_packed_cplx(cfg::SHTConfig, z::AbstractMatrix{<:Complex}, # fft_phi re-plans each call). eltype preserved for the AD/DFT fallback. Fφ = fft_phi!(Matrix{complex(float(eltype(z)))}(undef, size(z)...), z) P = Vector{Float64}(undef, lmax + 1) - scaleφ = cfg.cphi + scaleφ = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) xv = cfg.x; wv = cfg.w # hoist field reads out of the m/l loops (cfg is mutable, so not auto-hoisted) # Read both signs of m from the FFT output and store them in LM_cplx order. @@ -285,7 +285,10 @@ function synthesis_point_cplx(cfg::SHTConfig, alm::AbstractVector{<:Complex}, co acc += gp * cis(convert(PT, am) * ph) am > 0 && (acc += gn * cis(-convert(PT, am) * ph)) end - return acc + # Same φ convention factor `synthesis_packed_cplx` carries (1 under :dft; + # 1/2π under :quad), so a point evaluation matches the grid it samples. + sφ = _evaluator_phi_scale(cfg) + return sφ == 1 ? acc : acc * sφ end diff --git a/src/config.jl b/src/config.jl index 67ff205c..55def3e6 100644 --- a/src/config.jl +++ b/src/config.jl @@ -248,6 +248,33 @@ function SHTConfig(; nlat_padded::Integer = 0, spat_dist::Integer = 0, ) + # Validate the invariants the transforms rely on. This constructor is + # exported and was previously unchecked, so a hand-built configuration could + # violate e.g. `nlon >= 2*mmax+1` and then silently synthesize an all-zero + # field for any mode it could not resolve. + lmax >= 0 || throw(ArgumentError("lmax must be ≥ 0, got $lmax")) + mmax >= 0 || throw(ArgumentError("mmax must be ≥ 0, got $mmax")) + mmax <= lmax || throw(ArgumentError("mmax must be ≤ lmax, got mmax=$mmax, lmax=$lmax")) + mres >= 1 || throw(ArgumentError("mres must be ≥ 1, got $mres")) + nlat >= 1 || throw(ArgumentError("nlat must be ≥ 1, got $nlat")) + nlon >= 2*mmax + 1 || throw(ArgumentError( + "nlon must be ≥ 2*mmax+1 = $(2*mmax+1) to resolve every azimuthal order, got nlon=$nlon")) + length(θ) == nlat || throw(DimensionMismatch("θ must have nlat=$nlat entries, got $(length(θ))")) + length(x) == nlat || throw(DimensionMismatch("x must have nlat=$nlat entries, got $(length(x))")) + length(w) == nlat || throw(DimensionMismatch("w must have nlat=$nlat entries, got $(length(w))")) + length(st) == nlat || throw(DimensionMismatch("st must have nlat=$nlat entries, got $(length(st))")) + length(φ) == nlon || throw(DimensionMismatch("φ must have nlon=$nlon entries, got $(length(φ))")) + size(Nlm) == (lmax + 1, mmax + 1) || throw(DimensionMismatch( + "Nlm must be ($(lmax+1), $(mmax+1)), got $(size(Nlm))")) + let expected_nlm = nlm_calc(lmax, mmax, mres) + nlm == expected_nlm || throw(ArgumentError( + "nlm must be nlm_calc(lmax, mmax, mres) = $expected_nlm, got $nlm")) + end + length(li) == nlm || throw(DimensionMismatch("li must have nlm=$nlm entries, got $(length(li))")) + length(mi) == nlm || throw(DimensionMismatch("mi must have nlm=$nlm entries, got $(length(mi))")) + nspat == nlat * nlon || throw(ArgumentError( + "nspat must be nlat*nlon = $(nlat*nlon), got $nspat")) + # Concretize vectors/matrices so SHTGrid/SHTNorm/SHTTables fields always # carry the exact declared types. Lets callers pass ranges, views, etc. grid = SHTGrid(collect(Float64, θ), collect(Float64, φ), @@ -392,11 +419,63 @@ function Base.setproperty!(cfg::SHTConfig, name::Symbol, val) return setfield!(getfield(cfg, :_scratch), :otf_Pb, val) elseif name === :_m_order return setfield!(getfield(cfg, :_scratch), :m_order, val) + # ----- structural fields: keep derived state consistent ----- + elseif name === :lmax || name === :mmax || name === :mres + setfield!(cfg, name, Int(val)) + _rebuild_spectral_layout!(cfg) + return val + elseif name === :nlat || name === :nlon || name === :grid_type || + name === :nlm || name === :li || name === :mi || name === :nspat + throw(ArgumentError( + "`cfg.$(name)` cannot be reassigned: the quadrature nodes, weights and " * + "packed index tables are all derived from it, and there is no grid-type-" * + "independent way to regenerate them in place. Build a new configuration " * + "with `create_gauss_config` / `create_regular_config` / `create_config` " * + "instead. (`lmax`, `mmax` and `mres` may be assigned; the spectral layout " * + "is rebuilt for you.)")) else return setfield!(cfg, name, val) end end +""" + _rebuild_spectral_layout!(cfg::SHTConfig) + +Regenerate everything derived from `lmax`/`mmax`/`mres` after one of them is +reassigned: the normalization table, the packed mode count and its `li`/`mi` +lookups, and the cached norm-scale matrix and m-ordering. + +Without this, assigning `cfg.lmax = 10` left `size(cfg.Nlm) == (7, 7)` while the +transforms indexed it as `(lmax+1, mmax+1)` under `@inbounds` — an out-of-bounds +read of a live array. Precomputed Legendre tables are sized `(lmax+1, nlat)` and +cannot survive the change either, so they are dropped; transforms fall back to +the on-the-fly path until `prepare_plm_tables!` is called again. +""" +function _rebuild_spectral_layout!(cfg::SHTConfig) + lmax = getfield(cfg, :lmax); mmax = getfield(cfg, :mmax); mres = getfield(cfg, :mres) + lmax >= 0 || throw(ArgumentError("lmax must be ≥ 0, got $lmax")) + mmax >= 0 || throw(ArgumentError("mmax must be ≥ 0, got $mmax")) + mmax <= lmax || throw(ArgumentError("mmax must be ≤ lmax, got mmax=$mmax, lmax=$lmax")) + mres >= 1 || throw(ArgumentError("mres must be ≥ 1, got $mres")) + + setfield!(cfg, :nlm, nlm_calc(lmax, mmax, mres)) + li, mi = build_li_mi(lmax, mmax, mres) + setfield!(cfg, :li, li) + setfield!(cfg, :mi, mi) + + nrm = getfield(cfg, :_norm) + nrm.Nlm = Nlm_table(lmax, mmax) + nrm.scale_matrix[] = Matrix{Float64}(undef, 0, 0) # rebuilt lazily at the new size + + tbl = getfield(cfg, :_tables) + tbl.enabled = false + tbl.plm = Matrix{Float64}[]; tbl.dplm = Matrix{Float64}[] + tbl.NP = Matrix{Float64}[]; tbl.NdP = Matrix{Float64}[] + + empty!(getfield(cfg, :_scratch).m_order) # rebuilt lazily by cached_m_order + return cfg +end + function Base.propertynames(::SHTConfig, private::Bool=false) return (fieldnames(SHTConfig)..., :θ, :φ, :x, :w, :st, :cphi, @@ -940,15 +1019,60 @@ The first nlat rows contain the actual data; remaining rows are padding. ```julia cfg = create_gauss_config(64, 66) set_allow_padding!(cfg) -field = allocate_padded_spatial(cfg) -# field has size (nlat_padded, nlon), use field[1:cfg.nlat, :] for data +field = allocate_padded_spatial(cfg) # (nlat_padded, nlon) +copy_to_padded!(field, data, cfg) +alm = analysis(cfg, spatial_view(cfg, field)) # transforms need exactly nlat rows ``` + +The transforms require an array with exactly `nlat` rows, so a padded buffer must +go through [`spatial_view`](@ref) (or an equivalent `view(field, 1:cfg.nlat, :)`). +The view keeps the padded column stride, preserving what the padding is for. """ function allocate_padded_spatial(cfg::SHTConfig, T::Type=Float64) nlat_p = get_nlat_padded(cfg) return zeros(T, nlat_p, cfg.nlon) end +""" + spatial_view(cfg::SHTConfig, A::AbstractArray) -> SubArray + +The `(nlat, nlon)` (or `(nlat, nlon, nfields)`) window of a padded spatial buffer, +in the form the transforms accept. + +`allocate_padded_spatial` deliberately returns an array whose first dimension is +`nlat_padded ≥ nlat`, while every transform requires exactly `nlat` rows — so the +padded buffer could not be handed to `analysis` at all, which left the whole +padding API unusable. This is the missing bridge: + +```julia +cfg = create_gauss_config(64, 66) +set_allow_padding!(cfg) +pad = allocate_padded_spatial(cfg) +copy_to_padded!(pad, field, cfg) +alm = analysis(cfg, spatial_view(cfg, pad)) # padded stride preserved +``` + +The view does not throw the padding away: a `SubArray` over the leading rows keeps +the padded column stride (`nlat_padded`, not `nlat`), which is exactly the +cache-conflict avoidance the padding exists for. Results are bit-identical to the +unpadded layout. +""" +function spatial_view(cfg::SHTConfig, A::AbstractMatrix) + size(A, 1) >= cfg.nlat || throw(DimensionMismatch( + "padded spatial array needs ≥ nlat=$(cfg.nlat) rows, got $(size(A, 1))")) + size(A, 2) == cfg.nlon || throw(DimensionMismatch( + "spatial array second dim must be nlon=$(cfg.nlon), got $(size(A, 2))")) + return view(A, 1:cfg.nlat, 1:cfg.nlon) +end + +function spatial_view(cfg::SHTConfig, A::AbstractArray{<:Any,3}) + size(A, 1) >= cfg.nlat || throw(DimensionMismatch( + "padded spatial batch needs ≥ nlat=$(cfg.nlat) rows, got $(size(A, 1))")) + size(A, 2) == cfg.nlon || throw(DimensionMismatch( + "spatial batch second dim must be nlon=$(cfg.nlon), got $(size(A, 2))")) + return view(A, 1:cfg.nlat, 1:cfg.nlon, :) +end + """ allocate_padded_spatial_batch(cfg::SHTConfig, nfields::Int, T::Type=Float64) -> Array{T} @@ -1018,6 +1142,27 @@ Fejér's first rule on `θ = (i+0.5)π/nlat` nodes by default; set By default associated Legendre tables are precomputed, which mirrors SHTns' regular-grid behaviour and improves performance. +!!! warning "Needs `nlat ≥ 2*lmax + 1` to be exact" + Fejér and Clenshaw–Curtis rules with `nlat` nodes integrate polynomials in + `x = cosθ` through degree `nlat - 1` exactly, and analysis integrates a + product of two degree-`lmax` Legendre functions — degree `2*lmax`. So these + grids are exact only from `nlat = 2*lmax + 1` upward, whereas Gauss–Legendre + needs just `nlat = lmax + 1`. + + Below that threshold `analysis(cfg, synthesis(cfg, alm))` does not return + `alm`. Measured relative error at `lmax = 8`: + + | `nlat` | `:regular` | `:regular_poles` | + |---:|---:|---:| + | 10 | 7.2e-2 | 3.6e-1 | + | 14 | 1.3e-2 | 7.1e-3 | + | 16 | 4.2e-3 | 1.3e-3 | + | 17 (`2*lmax+1`) | 7e-16 | 7e-16 | + + Nothing warns at run time, so size the grid deliberately: pass + `nlat ≥ 2*lmax + 1` here, or use [`create_gauss_config`](@ref) (exact at + `nlat = lmax + 1`) or `use_dh_weights=true` (exact at `nlat = 2*(lmax+1)`). + # Driscoll-Healy Quadrature Set `use_dh_weights=true` to use Driscoll-Healy quadrature for exact spherical @@ -1215,6 +1360,7 @@ function prepare_plm_tables!(cfg::SHTConfig) # Working arrays for computing one row at a time P = Vector{Float64}(undef, lmax + 1) # P̄_l^m(x) normalized values dPdtheta = Vector{Float64}(undef, lmax + 1) # dP̄_l^m/dθ normalized θ-derivatives + Pbuf = Vector{Float64}(undef, lmax + 2) # extended P̄ row reused across all nlat*(mmax+1) calls # Compute tables for each azimuthal order m using bounded normalized recurrence # plm_tables[m+1][l+1, i] = P̄_l^m(x_i) (orthonormal; no overflow at high lmax) @@ -1227,7 +1373,7 @@ function prepare_plm_tables!(cfg::SHTConfig) # Compute normalized Legendre polynomials at each latitude point for i in 1:nlat s_i = sqrt(max(0.0, 1.0 - cfg.x[i]^2)) - Plm_norm_and_dPdtheta_row!(P, dPdtheta, cfg.x[i], lmax, m) + Plm_norm_and_dPdtheta_row!(P, dPdtheta, cfg.x[i], lmax, m, Pbuf) # Store normalized values — no Nlm multiply needed (P̄ already = Nlm * rawP) @inbounds @views tbl[:, i] .= P # P̄_l^m(x_i) for l=0:lmax diff --git a/src/core_transforms.jl b/src/core_transforms.jl index 63876e6d..c77a2d23 100644 --- a/src/core_transforms.jl +++ b/src/core_transforms.jl @@ -591,7 +591,8 @@ function _adjoint_analysis(cfg::SHTConfig, Alm̄::AbstractMatrix; Fφ = Matrix{CT}(undef, nlat_local, nlon) fill!(Fφ, zero(eltype(Fφ))) lmax, mmax = cfg.lmax, cfg.mmax - φadj = 2π # nlon (ifft adjoint) × cphi (2π/nlon) = 2π + # nlon (ifft adjoint) × the analysis φ factor. = 2π under :dft. + φadj = cfg.nlon * _analysis_phi_scale(cfg) use_tbl = has_fused_scalar_tables(cfg) P = use_tbl ? nothing : Vector{Float64}(undef, lmax + 1) for m in 0:cfg.mres:mmax @@ -625,7 +626,7 @@ end """Scalar analysis orchestrator. Parallelizes Legendre integration over m-modes.""" function _analysis_scalar_mloop!(alm::AbstractMatrix, cfg::SHTConfig, Fph::AbstractMatrix) lmax, mmax = cfg.lmax, cfg.mmax - scale_phi = cfg.cphi + scale_phi = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) m_order = cached_m_order(cfg) if has_fused_scalar_tables(cfg) _analysis_scalar_mloop_tbl!(alm, cfg, Fph, m_order, scale_phi) diff --git a/src/energy_diagnostics.jl b/src/energy_diagnostics.jl index a526fd3f..8575b853 100644 --- a/src/energy_diagnostics.jl +++ b/src/energy_diagnostics.jl @@ -90,6 +90,9 @@ This represents the L² norm of the field, which is conserved under orthonormal spherical harmonic transforms (Parseval's identity). """ function energy_scalar(cfg::SHTConfig, alm::AbstractMatrix; real_field::Bool=true) + # The accumulation loop is `@inbounds`; without this a too-small matrix + # silently returned a wrong number instead of raising, unlike every sibling. + validate_spectral_dimensions(alm, cfg, "alm") lmax, mmax = cfg.lmax, cfg.mmax scale_matrix = _diagnostic_scale_matrix(cfg) # Type-stable accumulator (stays inferrable for Float32 / ForwardDiff.Dual inputs). @@ -109,6 +112,7 @@ For a vector field V = ∇×(T Y_l^m êᵣ) + ∇ₕ(S Y_l^m), the kinetic energ KE = (1/2) ∫ |V|² dΩ = (1/2) Σ [l(l+1)|S_lm|² + l(l+1)|T_lm|²] """ function energy_vector(cfg::SHTConfig, Slm::AbstractMatrix, Tlm::AbstractMatrix; real_field::Bool=true) + validate_spectral_pair_dimensions(Slm, Tlm, cfg, ("Slm", "Tlm")) lmax, mmax = cfg.lmax, cfg.mmax scale_matrix = _diagnostic_scale_matrix(cfg) E = zero(promote_type(Float64, real(float(eltype(Slm))), real(float(eltype(Tlm))))) diff --git a/src/fftutils.jl b/src/fftutils.jl index b553989a..c4ebc661 100644 --- a/src/fftutils.jl +++ b/src/fftutils.jl @@ -87,10 +87,77 @@ const _FFT_BACKEND_DFT = 2 # planning allocation cost without changing the public `SHTPlan` contract. # Strides are part of the key because views with the same shape may need # different FFTW plans. +# +# This is THE φ-FFT plan cache for the whole package: the serial transforms, the +# batch helpers and the distributed extension's per-rank FFTs all route through +# `fft_phi!`/`rfft_phi!`/`ifft_phi!`/`irfft_phi!` and therefore through here. The +# public `enable_fft_plan_cache!` / `disable_fft_plan_cache!` / `set_fft_plan_cache!` +# / `fft_plan_cache_enabled` knobs control THIS cache. const _LOCAL_FFT_PLAN_CACHE = Dict{Tuple{Symbol,DataType,NTuple{2,Int},NTuple{2,Int},Int}, Any}() const _LOCAL_FFT_PLAN_CACHE_LOCK = ReentrantLock() +# `SHTNSKIT_CACHE_PENCILFFTS` is the legacy spelling kept for scripts that already +# set it; `SHTNSKIT_FFT_PLAN_CACHE` is the current name. +const _FFT_PLAN_CACHE_ENABLED = + Ref(get(ENV, "SHTNSKIT_FFT_PLAN_CACHE", + get(ENV, "SHTNSKIT_CACHE_PENCILFFTS", "1")) != "0") + +# Soft cap on distinct cached plans. Transform sizes are few and stable in +# practice, but a long-lived process that sweeps many grid shapes would otherwise +# grow this dictionary without bound. Flush wholesale when the cap is reached. +const _FFT_PLAN_CACHE_MAX = + Ref(parse(Int, get(ENV, "SHTNSKIT_FFT_PLAN_CACHE_MAX", "64"))) + +""" + fft_plan_cache_enabled() -> Bool + +Whether φ-FFT plans are cached and reused across calls. +""" +fft_plan_cache_enabled() = _FFT_PLAN_CACHE_ENABLED[] + +""" + set_fft_plan_cache!(flag::Bool; clear::Bool=true) -> Bool + +Enable or disable φ-FFT plan caching. When disabling with `clear=true` (the +default) the existing entries are dropped as well. Returns `flag`. +""" +function set_fft_plan_cache!(flag::Bool; clear::Bool=true) + _FFT_PLAN_CACHE_ENABLED[] = flag + if !flag && clear + lock(_LOCAL_FFT_PLAN_CACHE_LOCK) do + empty!(_LOCAL_FFT_PLAN_CACHE) + end + end + return flag +end + +""" + enable_fft_plan_cache!() -> Bool + +Convenience wrapper for `set_fft_plan_cache!(true)`. +""" +enable_fft_plan_cache!() = set_fft_plan_cache!(true) + +""" + disable_fft_plan_cache!(; clear::Bool=true) -> Bool + +Convenience wrapper for `set_fft_plan_cache!(false; clear)`. +""" +disable_fft_plan_cache!(; clear::Bool=true) = set_fft_plan_cache!(false; clear=clear) + +""" + SHTnsKit.fft_plan_cache_max!(n::Int) -> Int + +Set the maximum number of cached φ-FFT plans; `n <= 0` removes the cap. +Returns the previous value. +""" +function fft_plan_cache_max!(n::Int) + prev = _FFT_PLAN_CACHE_MAX[] + _FFT_PLAN_CACHE_MAX[] = n + return prev +end + function fft_phi_backend() v = _FFT_BACKEND[] v == _FFT_BACKEND_FFTW && return :fftw @@ -102,32 +169,43 @@ end return (kind, eltype(A), (size(A, 1), size(A, 2)), (stride(A, 1), stride(A, 2)), nlon) end +"""Build one FFTW plan for `A`; see `_cached_local_fft_plan` for the flag rationale.""" +function _build_local_fft_plan(kind::Symbol, A::AbstractMatrix, nlon::Int) + # UNALIGNED is required, not an optimization choice. The cache key + # covers eltype/size/strides but NOT the base pointer's alignment, so + # a plan built for one array gets reused for another that FFTW may + # consider differently aligned. Without UNALIGNED that reuse throws + # `ArgumentError`, which the callers catch and answer by falling back + # to the pure-Julia O(nlat·nlon²) DFT — an order-of-magnitude + # slowdown with no error surfaced. Forfeiting the aligned SIMD + # codelets is much cheaper than forfeiting the FFT. The batch helpers + # in batch_transforms.jl pass the same flag for the same reason. + flags = FFTW.ESTIMATE | FFTW.UNALIGNED + return if kind === :fft + plan_fft!(A, 2; flags) + elseif kind === :ifft + plan_ifft!(A, 2; flags) + elseif kind === :rfft + plan_rfft(A, 2; flags) + elseif kind === :irfft + plan_irfft(A, nlon, 2; flags) + else + throw(ArgumentError("unknown FFT plan kind: $kind")) + end +end + function _cached_local_fft_plan(kind::Symbol, A::AbstractMatrix, nlon::Int=0) + _FFT_PLAN_CACHE_ENABLED[] || return _build_local_fft_plan(kind, A, nlon) key = _local_fft_plan_key(kind, A, nlon) lock(_LOCAL_FFT_PLAN_CACHE_LOCK) try plan = get(_LOCAL_FFT_PLAN_CACHE, key, nothing) if plan === nothing - # UNALIGNED is required, not an optimization choice. The cache key - # covers eltype/size/strides but NOT the base pointer's alignment, so - # a plan built for one array gets reused for another that FFTW may - # consider differently aligned. Without UNALIGNED that reuse throws - # `ArgumentError`, which the callers catch and answer by falling back - # to the pure-Julia O(nlat·nlon²) DFT — an order-of-magnitude - # slowdown with no error surfaced. Forfeiting the aligned SIMD - # codelets is much cheaper than forfeiting the FFT. The batch helpers - # in batch_transforms.jl pass the same flag for the same reason. - flags = FFTW.ESTIMATE | FFTW.UNALIGNED - plan = if kind === :fft - plan_fft!(A, 2; flags) - elseif kind === :ifft - plan_ifft!(A, 2; flags) - elseif kind === :rfft - plan_rfft(A, 2; flags) - elseif kind === :irfft - plan_irfft(A, nlon, 2; flags) - else - throw(ArgumentError("unknown FFT plan kind: $kind")) + plan = _build_local_fft_plan(kind, A, nlon) + # Enforce the soft cap: flush before inserting so the fresh entry survives. + cap = _FFT_PLAN_CACHE_MAX[] + if cap > 0 && length(_LOCAL_FFT_PLAN_CACHE) >= cap + empty!(_LOCAL_FFT_PLAN_CACHE) end _LOCAL_FFT_PLAN_CACHE[key] = plan end diff --git a/src/layout.jl b/src/layout.jl index 4cec7b4a..5918eddb 100644 --- a/src/layout.jl +++ b/src/layout.jl @@ -227,17 +227,22 @@ the given packed index. Useful for algorithms that need to determine the azimuthal symmetry of a coefficient from its packed storage location. """ -function im_from_lm(lm::Int, lmax::Int, mres::Int) +function im_from_lm(lm::Int, lmax::Int, mres::Int; mmax::Int=lmax) # Validate packed index lm ≥ 0 || throw(ArgumentError("lm must be ≥ 0")) + (0 ≤ mmax ≤ lmax) || throw(ArgumentError("require 0 ≤ mmax ≤ lmax")) # Search through m-blocks to find the one containing this packed index im = 0 # Current reduced m-index being tested base = 0 # Base offset for current m-block - im_max = lmax ÷ mres # maximum valid reduced m-index + # Bound by the last block the LAYOUT actually has, which is set by mmax, not + # by lmax. Using `lmax ÷ mres` let a past-the-end index resolve to an order + # the configuration does not store (`im_from_lm(30, 8, 1)` returned 4 for an + # mmax=3 layout whose valid indices stop at 29) instead of raising. + im_max = mmax ÷ mres # maximum valid reduced m-index while true - im > im_max && throw(ArgumentError("lm=$lm is out of range for lmax=$lmax, mres=$mres")) + im > im_max && throw(ArgumentError("lm=$lm is out of range for lmax=$lmax, mmax=$mmax, mres=$mres")) # Size of current m-block (number of l-modes for this m) block = lmax - im*mres + 1 diff --git a/src/legendre.jl b/src/legendre.jl index c311a24f..e9a435c9 100644 --- a/src/legendre.jl +++ b/src/legendre.jl @@ -245,17 +245,46 @@ At poles (sinθ → 0): """ function Plm_norm_and_dPdtheta_row!(P::AbstractVector{T}, dPdtheta::AbstractVector{T}, x::T, lmax::Int, m::Int) where {T<:Real} - # Fill P with orthonormal P̄_l^m - Plm_norm_row!(P, x, lmax, m) + # Allocating fallback: provide a local Pbuf and delegate. + Pbuf = zeros(T, lmax + 2) + return Plm_norm_and_dPdtheta_row!(P, dPdtheta, x, lmax, m, Pbuf) +end + +""" + Plm_norm_and_dPdtheta_row!(P, dPdtheta, x, lmax, m, Pbuf) + +Buffer-taking form: `Pbuf` is caller-supplied scratch of length ≥ `lmax+2`, +avoiding a heap allocation per call in hot loops (`prepare_plm_tables!` calls +this `nlat*(mmax+1)` times). +""" +function Plm_norm_and_dPdtheta_row!(P::AbstractVector{T}, dPdtheta::AbstractVector{T}, + x::T, lmax::Int, m::Int, + Pbuf::AbstractVector{T}) where {T<:Real} + length(Pbuf) >= lmax + 2 || throw(ArgumentError("Pbuf must have length ≥ lmax+2")) @inbounds begin fill!(dPdtheta, zero(T)) - lmax < m && return P, dPdtheta sinth = sqrt(max(zero(T), one(T) - x*x)) + interior = lmax >= m && sinth >= POLE_TOLERANCE_FACTOR * eps(T) + + if interior + # The dθ recurrence needs P̄ up to lmax+1, and the first lmax+1 of + # those ARE the row we must return. Computing `P` with its own + # `Plm_norm_row!` call and then `Pbuf` with a second one ran the + # (inherently serial) recurrence twice for one row — about 40 % of + # this function's cost. Run it once into the longer buffer instead. + Plm_norm_row!(Pbuf, x, lmax + 1, m) + fill!(P, zero(T)) + copyto!(P, 1, Pbuf, 1, lmax + 1) + else + Plm_norm_row!(P, x, lmax, m) + end + + lmax < m && return P, dPdtheta # Handle poles (sinθ ≈ 0) - if sinth < POLE_TOLERANCE_FACTOR * eps(T) + if !interior if m == 1 # Analytic limit: dP̄_l^1/dθ|_{θ=0} = N_{l,1} * (−l(l+1)/2) # N_{l,1} = sqrt[(2l+1)/(4π) / (l*(l+1))], so @@ -278,11 +307,8 @@ function Plm_norm_and_dPdtheta_row!(P::AbstractVector{T}, dPdtheta::AbstractVect return P, dPdtheta end - # Standard case: use the normalized dθ recurrence. - # We need P̄_{l+1}^m for the l-term, so compute P̄ to lmax+1. - # Use a local buffer for the extended row. - Pbuf = zeros(T, lmax + 2) # indices 1..lmax+2 → degrees 0..lmax+1 - Plm_norm_row!(Pbuf, x, lmax + 1, m) + # Standard case: `Pbuf` already holds P̄ up to degree lmax+1 from the + # single recurrence pass above. # d(l,m) = sqrt((l^2 - m^2) / (4l^2 - 1)) [0 for l=0] inv_sinth = one(T) / sinth @@ -334,16 +360,29 @@ function Plm_norm_dPdtheta_over_sinth_row!(P::AbstractVector{T}, dPdtheta::Abstr x::T, lmax::Int, m::Int, Pbuf::AbstractVector{T}) where {T<:Real} length(Pbuf) >= lmax + 2 || throw(ArgumentError("Pbuf must have length ≥ lmax+2")) - Plm_norm_row!(P, x, lmax, m) @inbounds begin fill!(dPdtheta, zero(T)) fill!(P_over_sinth, zero(T)) - lmax < m && return P, dPdtheta, P_over_sinth sinth = sqrt(max(zero(T), one(T) - x*x)) + interior = lmax >= m && sinth >= POLE_TOLERANCE_FACTOR * eps(T) + + if interior + # Single recurrence pass: the dθ formula needs P̄ to degree lmax+1 and + # the first lmax+1 entries of that row ARE `P`. Filling `P` and `Pbuf` + # with two separate `Plm_norm_row!` calls ran the serial recurrence + # twice per row (~40 % of this function's cost). + Plm_norm_row!(Pbuf, x, lmax + 1, m) + fill!(P, zero(T)) + copyto!(P, 1, Pbuf, 1, lmax + 1) + else + Plm_norm_row!(P, x, lmax, m) + end - if sinth < POLE_TOLERANCE_FACTOR * eps(T) + lmax < m && return P, dPdtheta, P_over_sinth + + if !interior if m == 1 for l in 1:lmax mag = T(0.5) * T(_INV_SQRT_4PI) * sqrt(T(2l + 1) * T(l) * T(l + 1)) @@ -362,8 +401,8 @@ function Plm_norm_dPdtheta_over_sinth_row!(P::AbstractVector{T}, dPdtheta::Abstr return P, dPdtheta, P_over_sinth end - # Standard case: compute P̄ to lmax+1 for dθ recurrence - Plm_norm_row!(Pbuf, x, lmax + 1, m) + # Standard case: `Pbuf` already holds P̄ to degree lmax+1 from the single + # recurrence pass above. inv_sinth = one(T) / sinth diff --git a/src/local.jl b/src/local.jl index dcfc38c3..3a2f0707 100644 --- a/src/local.jl +++ b/src/local.jl @@ -68,6 +68,8 @@ function SH_to_lat(cfg::SHTConfig, Qlm::AbstractVector{<:Complex}, cost::Real; n vals[j+1] += 2 * real(gm * cis(PT(2π * m * j / nphi))) end end + sφ = _evaluator_phi_scale(cfg) # 1 under :dft; 1/2π under :quad + sφ == 1 || (vals .*= sφ) return vals end @@ -126,6 +128,8 @@ function SH_to_lat_cplx(cfg::SHTConfig, alm_packed::AbstractVector{<:Complex}, c vals[j+1] += gm * phase + gn * conj(phase) end end + sφ = _evaluator_phi_scale(cfg) # 1 under :dft; 1/2π under :quad + sφ == 1 || (vals .*= sφ) return vals end @@ -203,7 +207,8 @@ function SHqst_to_point(cfg::SHTConfig, Qlm::AbstractVector{<:Complex}, Slm::Abs vt *= sinth vp *= sinth end - return real(vr), real(vt), real(vp) + sφ = _evaluator_phi_scale(cfg) # 1 under :dft; 1/2π under :quad + return real(vr) * sφ, real(vt) * sφ, real(vp) * sφ end """ @@ -273,7 +278,8 @@ function SH_to_grad_point(cfg::SHTConfig, DrSlm::AbstractVector{<:Complex}, Slm: vt *= sinth vp *= sinth end - return real(vr), real(vt), real(vp) + sφ = _evaluator_phi_scale(cfg) # 1 under :dft; 1/2π under :quad + return real(vr) * sφ, real(vt) * sφ, real(vp) * sφ end """ @@ -367,6 +373,10 @@ function SHqst_to_lat(cfg::SHTConfig, Qlm::AbstractVector{<:Complex}, Slm::Abstr Vt .*= sinth Vp .*= sinth end + sφ = _evaluator_phi_scale(cfg) # 1 under :dft; 1/2π under :quad + if sφ != 1 + Vr .*= sφ; Vt .*= sφ; Vp .*= sφ + end return Vr, Vt, Vp end diff --git a/src/plan.jl b/src/plan.jl index bab074ae..97aa886e 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -212,7 +212,7 @@ function analysis_sphtor!(plan::SHTPlan, Slm_out::AbstractMatrix, Tlm_out::Abstr size(Tlm_out,1)==cfg.lmax+1 && size(Tlm_out,2)==cfg.mmax+1 || throw(DimensionMismatch("Tlm_out dims")) lmax, mmax = cfg.lmax, cfg.mmax - scaleφ = cfg.cphi + scaleφ = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) fill!(Slm_out, zero(eltype(Slm_out))); fill!(Tlm_out, zero(eltype(Tlm_out))) # Two passes over (Vt, Vp): each packs the component into a real/complex @@ -421,7 +421,7 @@ function analysis!(plan::SHTPlan, alm_out::AbstractMatrix, f::AbstractMatrix) )) lmax, mmax = cfg.lmax, cfg.mmax - scaleφ = cfg.cphi + scaleφ = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) if plan.use_rfft eltype(f) <: Real || throw(ArgumentError("use_rfft plan requires real-valued f")) diff --git a/src/qst_transforms.jl b/src/qst_transforms.jl index 8610ef84..3f3432d6 100644 --- a/src/qst_transforms.jl +++ b/src/qst_transforms.jl @@ -331,12 +331,12 @@ function _synthesis_qst_l(cfg::SHTConfig, Qlm::AbstractMatrix, Slm::AbstractMatr end """ - analysis_qst_ml(cfg, im, Vr_m, Vt_m, Vp_m, ltr) -> (Ql, Sl, Tl) + analysis_qst_ml(cfg, mval, Vr_m, Vt_m, Vp_m, ltr) -> (Ql, Sl, Tl) -Mode-limited transform for specific azimuthal mode im. +Mode-limited transform for specific azimuthal mode mval. """ -function analysis_qst_ml(cfg::SHTConfig, im::Integer, Vr_m::AbstractVector{<:Complex}, Vt_m::AbstractVector{<:Complex}, Vp_m::AbstractVector{<:Complex}, ltr::Integer) - stored_im, _, lcap = _validate_stored_order(cfg, im, ltr) +function analysis_qst_ml(cfg::SHTConfig, mval::Integer, Vr_m::AbstractVector{<:Complex}, Vt_m::AbstractVector{<:Complex}, Vp_m::AbstractVector{<:Complex}, ltr::Integer) + stored_im, _, lcap = _validate_stored_order(cfg, mval, ltr) # Transform each component for this specific mode Ql = analysis_packed_ml(cfg, stored_im, Vr_m, lcap) Sl, Tl = analysis_sphtor_ml(cfg, stored_im, Vt_m, Vp_m, lcap) @@ -345,23 +345,23 @@ function analysis_qst_ml(cfg::SHTConfig, im::Integer, Vr_m::AbstractVector{<:Com return Ql, Sl, Tl end -function analysis_qst_ml(::CPU, cfg::SHTConfig, im::Integer, +function analysis_qst_ml(::CPU, cfg::SHTConfig, mval::Integer, Vr::AbstractVector{<:Complex}, Vt::AbstractVector{<:Complex}, Vp::AbstractVector{<:Complex}, ltr::Integer) for value in (Vr, Vt, Vp) _require_cpu_storage(:analysis_qst_ml, value) end - return analysis_qst_ml(cfg, im, Vr, Vt, Vp, ltr) + return analysis_qst_ml(cfg, mval, Vr, Vt, Vp, ltr) end """ - synthesis_qst_ml(cfg, im, Ql, Sl, Tl, ltr) -> (Vr_m, Vt_m, Vp_m) + synthesis_qst_ml(cfg, mval, Ql, Sl, Tl, ltr) -> (Vr_m, Vt_m, Vp_m) -Mode-limited synthesis for specific azimuthal mode im. +Mode-limited synthesis for specific azimuthal mode mval. """ -function synthesis_qst_ml(cfg::SHTConfig, im::Integer, Ql::AbstractVector{<:Complex}, Sl::AbstractVector{<:Complex}, Tl::AbstractVector{<:Complex}, ltr::Integer) - stored_im, _, lcap = _validate_stored_order(cfg, im, ltr) +function synthesis_qst_ml(cfg::SHTConfig, mval::Integer, Ql::AbstractVector{<:Complex}, Sl::AbstractVector{<:Complex}, Tl::AbstractVector{<:Complex}, ltr::Integer) + stored_im, _, lcap = _validate_stored_order(cfg, mval, ltr) # Each fixed-mode sub-transform converts its component to canonical once. Vr_m = synthesis_packed_ml(cfg, stored_im, Ql, lcap) Vt_m, Vp_m = synthesis_sphtor_ml(cfg, stored_im, Sl, Tl, lcap) @@ -369,12 +369,12 @@ function synthesis_qst_ml(cfg::SHTConfig, im::Integer, Ql::AbstractVector{<:Comp return Vr_m, Vt_m, Vp_m end -function synthesis_qst_ml(::CPU, cfg::SHTConfig, im::Integer, +function synthesis_qst_ml(::CPU, cfg::SHTConfig, mval::Integer, Q::AbstractVector{<:Complex}, S::AbstractVector{<:Complex}, Tlm::AbstractVector{<:Complex}, ltr::Integer) for value in (Q, S, Tlm) _require_cpu_storage(:synthesis_qst_ml, value) end - return synthesis_qst_ml(cfg, im, Q, S, Tlm, ltr) + return synthesis_qst_ml(cfg, mval, Q, S, Tlm, ltr) end diff --git a/src/rotations.jl b/src/rotations.jl index 05cb307f..1af0e325 100644 --- a/src/rotations.jl +++ b/src/rotations.jl @@ -97,6 +97,22 @@ Currently supports fast rotation around the Z-axis by angle `alpha` in radians. Rotate a real-field SH expansion around the Z-axis by angle `alpha`. Input and output are packed `Qlm` vectors (LM order, m ≥ 0). In-place supported if `Rlm === Qlm`. + +# Sign convention + +`R_lm = Q_lm · exp(-i m α)`. This is the **active** rotation of the field by `+α` +about `ẑ`: the rotated field is `g(θ, φ) = f(θ, φ - α)`. Equivalently, a feature +at longitude `φ₀` moves to `φ₀ + α`. + +Three things pin this sign and would all break if it were flipped to `+imα` +(which is the *passive* convention, `f(θ, φ + α)`): + + * the spatial rotation above, verified directly in + `test/serial/test_rotations.jl` against an FFT-grid φ shift; + * the general Wigner engine — `shtns_rotation_apply_real` with + `ZYZ(α, 0, 0)` (or `ZYZ(0, 0, α)`) must equal this function, and it builds + `diag(e^{-imα}) · d(β) · diag(e^{-imγ})`; + * the distributed twins `dist_SH_Zrotate` and every rotation `rrule`. """ function SH_Zrotate(::CPU, cfg::SHTConfig, Qlm::AbstractVector{<:Complex}, alpha::Real, Rlm::AbstractVector{<:Complex}) @@ -582,12 +598,61 @@ function _rotation_host_blocks(r::SHTRotation, ::Type{T}) where {T<:AbstractFloa return (; offsets, values, input_scales, output_scales, alpha=α, gamma=γ) end +""" + _require_full_m_range(r::SHTRotation, β::Real) + +Reject an order-mixing rotation on a layout that cannot hold every order it +produces. + +A Wigner-d rotation through a general `β` couples `Y_l^m` to every `Y_l^{m′}` +with `|m′| ≤ l`. If the storage stops at `mmax < lmax`, the `|m′| > mmax` +components have nowhere to go and were silently dropped — measured at +`lmax = 8`, that quietly discarded **14.8 %** of the field’s energy at +`mmax = 5` and **24.0 %** at `mmax = 3`, with no error and no warning. + +The two degenerate angles are exempt because their `d^l` is not order-mixing: +`β ≡ 0` is diagonal, and `β ≡ π` is anti-diagonal (`m′ = -m`), so `|m′| = |m|` +and a truncated layout still holds the result. That keeps pure Z-rotations +expressed as `ZYZ(α, 0, γ)` working at any `mmax`. + +Mirrors the `mres > 1` restriction stated by `dist_SH_Yrotate` and the packed +distributed rotations, for the same reason. +""" +function _require_full_m_range(r::SHTRotation, β::Real) + r.mmax >= r.lmax && return nothing + abs(sin(float(β))) <= 1e-12 && return nothing # β ≡ 0 (mod π): no m mixing + throw(ArgumentError( + "rotation with β=$(β) mixes azimuthal orders, but this configuration " * + "stores only m ≤ mmax=$(r.mmax) with lmax=$(r.lmax); the |m| > mmax " * + "components such a rotation generates cannot be represented and would " * + "be silently discarded. Use a configuration with mmax == lmax for " * + "Y/X rotations and general Euler angles. Pure Z-rotations (β ≡ 0 mod π) " * + "are unaffected and still work at any mmax.")) +end + +""" + _rotation_packed_length_check(v, expected, name, r) + +Length guard for the packed rotation inputs, with an `mres`-aware message. +""" +function _rotation_packed_length_check(v::AbstractVector, expected::Int, + name::AbstractString, r::SHTRotation) + length(v) == expected && return nothing + throw(DimensionMismatch( + "$name has length $(length(v)), expected $expected for the packed (mres=1) " * + "layout at lmax=$(r.lmax), mmax=$(r.mmax). A Y/X rotation mixes azimuthal " * + "orders, so it cannot be represented in an mres-strided layout at all — an " * + "mres>1 config produces a shorter packed vector and lands here. Use mres=1 " * + "for rotations other than SH_Zrotate.")) +end + function _rotation_apply_cplx_canonical!(r::SHTRotation, Zlm::AbstractVector{<:Complex}, Rlm::AbstractVector{<:Complex}) r.lmax ≥ 0 || return Rlm RT = typeof(real(zero(eltype(Rlm)))) α, β, γ = _rotation_zyz_angles(r, RT) + _require_full_m_range(r, β) # Pre-allocate working arrays at maximum size to avoid per-l allocations nmax = 2 * r.lmax + 1 @@ -692,8 +757,13 @@ function shtns_rotation_apply_real(::CPU, r::SHTRotation, _require_cpu_storage(:shtns_rotation_apply_real, Qlm) _require_cpu_storage(:shtns_rotation_apply_real, Rlm) expected = nlm_calc(r.lmax, r.mmax, 1) - length(Qlm) == expected || throw(DimensionMismatch("LM packed size mismatch")) - length(Rlm) == expected || throw(DimensionMismatch("LM packed size mismatch")) + # A length mismatch here is almost always an `mres > 1` config reaching a + # rotation that mixes orders, which no mres-strided layout can represent. + # Say that, rather than reporting a bare size mismatch the caller has to + # reverse-engineer. (`dist_SH_Yrotate` and the packed distributed rotations + # state the same restriction up front.) + _rotation_packed_length_check(Qlm, expected, "Qlm", r) + _rotation_packed_length_check(Rlm, expected, "Rlm", r) eltype(Qlm) === eltype(Rlm) || throw(ArgumentError( "rotation input and output element types must match", )) diff --git a/src/sphtor_transforms.jl b/src/sphtor_transforms.jl index 8458ad47..970b07f8 100644 --- a/src/sphtor_transforms.jl +++ b/src/sphtor_transforms.jl @@ -404,7 +404,7 @@ function _adjoint_analysis_sphtor(cfg::SHTConfig, Slm̄::AbstractMatrix, Tlm̄:: dPdtheta = Vector{Float64}(undef, lmax + 1) P_over_sinth = Vector{Float64}(undef, lmax + 1) Pbuf = Vector{Float64}(undef, lmax + 2) # scratch for extended P̄ row (avoids per-call alloc) - φadj = 2π + φadj = cfg.nlon * _analysis_phi_scale(cfg) # = 2π under :dft for m in 0:cfg.mres:mmax col = m + 1 @@ -533,7 +533,7 @@ function _analysis_sphtor_mloop!(Slm::AbstractMatrix, Tlm::AbstractMatrix, ltr::Int=cfg.lmax) lmax, mmax = cfg.lmax, cfg.mmax ltr_eff = min(ltr, lmax) - scale_phi = cfg.cphi + scale_phi = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) m_order = cached_m_order(cfg) if has_fused_vector_tables(cfg) @@ -997,39 +997,39 @@ synthesis_tor_l_cplx(::CPU, cfg::SHTConfig, Tlm::AbstractMatrix, synthesis_tor_l(CPU(), cfg, Tlm, ltr; real_output=false) """ - synthesis_sph_ml(cfg::SHTConfig, im::Int, Sl::AbstractVector{<:Complex}, ltr::Int) + synthesis_sph_ml(cfg::SHTConfig, mval::Int, Sl::AbstractVector{<:Complex}, ltr::Int) Mode-limited spheroidal-only synthesis wrapper. """ -function synthesis_sph_ml(cfg::SHTConfig, im::Integer, Sl::AbstractVector{<:Complex}, ltr::Integer) +function synthesis_sph_ml(cfg::SHTConfig, mval::Integer, Sl::AbstractVector{<:Complex}, ltr::Integer) # Mode-limited wrappers use zero-vector views for the missing component; - # this avoids an O(ltr-im) allocation on repeated per-mode calls. + # this avoids an O(ltr-mval) allocation on repeated per-mode calls. Tl_zero = _zero_spectral_vector(eltype(Sl), length(Sl)) - return synthesis_sphtor_ml(cfg, im, Sl, Tl_zero, ltr) + return synthesis_sphtor_ml(cfg, mval, Sl, Tl_zero, ltr) end -function synthesis_sph_ml(::CPU, cfg::SHTConfig, im::Integer, +function synthesis_sph_ml(::CPU, cfg::SHTConfig, mval::Integer, Sl::AbstractVector{<:Complex}, ltr::Integer) _require_cpu_storage(:synthesis_sph_ml, Sl) - return synthesis_sph_ml(cfg, im, Sl, ltr) + return synthesis_sph_ml(cfg, mval, Sl, ltr) end """ - synthesis_tor_ml(cfg::SHTConfig, im::Int, Tl::AbstractVector{<:Complex}, ltr::Int) + synthesis_tor_ml(cfg::SHTConfig, mval::Int, Tl::AbstractVector{<:Complex}, ltr::Int) Mode-limited toroidal-only synthesis wrapper. """ -function synthesis_tor_ml(cfg::SHTConfig, im::Integer, Tl::AbstractVector{<:Complex}, ltr::Integer) +function synthesis_tor_ml(cfg::SHTConfig, mval::Integer, Tl::AbstractVector{<:Complex}, ltr::Integer) # Mode-limited wrappers use zero-vector views for the missing component; - # this avoids an O(ltr-im) allocation on repeated per-mode calls. + # this avoids an O(ltr-mval) allocation on repeated per-mode calls. Sl_zero = _zero_spectral_vector(eltype(Tl), length(Tl)) - return synthesis_sphtor_ml(cfg, im, Sl_zero, Tl, ltr) + return synthesis_sphtor_ml(cfg, mval, Sl_zero, Tl, ltr) end -function synthesis_tor_ml(::CPU, cfg::SHTConfig, im::Integer, +function synthesis_tor_ml(::CPU, cfg::SHTConfig, mval::Integer, Tl::AbstractVector{<:Complex}, ltr::Integer) _require_cpu_storage(:synthesis_tor_ml, Tl) - return synthesis_tor_ml(cfg, im, Tl, ltr) + return synthesis_tor_ml(cfg, mval, Tl, ltr) end """ @@ -1081,7 +1081,7 @@ function analysis_sphtor_ml(cfg::SHTConfig, stored_im::Integer, Vt_m::AbstractVe dPdtheta = Vector{Float64}(undef, ltr + 1) P_over_sinth = Vector{Float64}(undef, ltr + 1) Pbuf = Vector{Float64}(undef, ltr + 2) # scratch for extended P̄ row (avoids per-call alloc) - scaleφ = cfg.cphi + scaleφ = _analysis_phi_scale(cfg) # inverts synthesis under any phi_scale (= cphi under :dft) # Integrate using Legendre polynomials and derivatives (pole-safe) for i in 1:nlat @@ -1229,12 +1229,12 @@ synthesis_grad_l(::CPU, cfg::SHTConfig, Slm::AbstractMatrix, synthesis_sph_l(CPU(), cfg, Slm, ltr; kwargs...) """ - synthesis_grad_ml(cfg::SHTConfig, im::Int, Sl::AbstractVector{<:Complex}, ltr::Int) + synthesis_grad_ml(cfg::SHTConfig, mval::Int, Sl::AbstractVector{<:Complex}, ltr::Int) Mode-limited gradient synthesis alias. """ -function synthesis_grad_ml(cfg::SHTConfig, im::Integer, Sl::AbstractVector{<:Complex}, ltr::Integer) - return synthesis_sph_ml(cfg, im, Sl, ltr) +function synthesis_grad_ml(cfg::SHTConfig, mval::Integer, Sl::AbstractVector{<:Complex}, ltr::Integer) + return synthesis_sph_ml(cfg, mval, Sl, ltr) end synthesis_grad_ml(::CPU, cfg::SHTConfig, im::Integer, diff --git a/src/transforms.jl b/src/transforms.jl index b495d146..de599702 100644 --- a/src/transforms.jl +++ b/src/transforms.jl @@ -108,7 +108,7 @@ function analysis_axisym(cfg::SHTConfig, Vr::AbstractVector{<:Real}) # applied explicitly here. Without it this is NOT the inverse of # `synthesis_axisym` (which matches `synthesis` exactly) and disagrees with # the m=0 column of `analysis` by 1/2π. - scaleφ = cfg.cphi * cfg.nlon + scaleφ = _analysis_phi_scale(cfg) * cfg.nlon # inverts synthesis under any phi_scale (= cphi under :dft) @inbounds for l in 0:lmax Ql[l+1] *= scaleφ end @@ -239,6 +239,11 @@ function synthesis_axisym(cfg::SHTConfig, Qlm::AbstractVector{<:Complex}) Vr[i] = val end + # Same φ convention factor the full `synthesis` carries (1 under :dft; + # 1/2π under :quad), so this stays the m=0 column of `synthesis` and the + # exact inverse of `analysis_axisym` in either mode. + sφ = _evaluator_phi_scale(cfg) + sφ == 1 || (Vr .*= RT(sφ)) return Vr end @@ -270,7 +275,7 @@ function analysis_axisym_l(cfg::SHTConfig, Vr::AbstractVector{<:Real}, ltr::Inte end # Same φ quadrature factor as `analysis_axisym` — see the comment there. - scaleφ = cfg.cphi * cfg.nlon + scaleφ = _analysis_phi_scale(cfg) * cfg.nlon # inverts synthesis under any phi_scale (= cphi under :dft) @inbounds for l in eachindex(Ql) Ql[l] *= scaleφ end @@ -307,25 +312,28 @@ function synthesis_axisym_l(cfg::SHTConfig, Qlm::AbstractVector{<:Complex}, ltr: Vr[i] = val end + # Same φ convention factor as `synthesis_axisym` — see the comment there. + sφ = _evaluator_phi_scale(cfg) + sφ == 1 || (Vr .*= RT(sφ)) return Vr end """ - analysis_packed_ml(cfg, im, Vr_m, ltr) -> Vector{<:Complex} + analysis_packed_ml(cfg, mval, Vr_m, ltr) -> Vector{<:Complex} Transform spatial field for one stored azimuthal order to spherical harmonic -coefficients. `im` is the zero-based stored-order index, so the physical order -is `m = im * cfg.mres`. `Vr_m` contains complex spatial values for that mode. +coefficients. `mval` is the zero-based stored-order index, so the physical order +is `m = mval * cfg.mres`. `Vr_m` contains complex spatial values for that mode. Returns coefficients Q_l for degrees l = m..ltr. """ -function analysis_packed_ml(cfg::SHTConfig, im::Int, Vr_m::AbstractVector{<:Complex}, ltr::Integer) +function analysis_packed_ml(cfg::SHTConfig, mval::Int, Vr_m::AbstractVector{<:Complex}, ltr::Integer) nlat = cfg.nlat length(Vr_m) == nlat || throw(DimensionMismatch("Vr_m length must be nlat=$(nlat)")) - im >= 0 || throw(ArgumentError("im must be >= 0")) - im <= cfg.mmax ÷ cfg.mres || throw(ArgumentError("im must be <= mmax/mres=$(cfg.mmax ÷ cfg.mres)")) - m = im * cfg.mres + mval >= 0 || throw(ArgumentError("mval must be >= 0")) + mval <= cfg.mmax ÷ cfg.mres || throw(ArgumentError("mval must be <= mmax/mres=$(cfg.mmax ÷ cfg.mres)")) + m = mval * cfg.mres ltr = _validate_degree_limit(cfg, ltr) - ltr >= m || throw(ArgumentError("ltr must be >= im*mres=$(m)")) + ltr >= m || throw(ArgumentError("ltr must be >= mval*mres=$(m)")) num_l = ltr - m + 1 CT = complex(float(real(eltype(Vr_m)))) # AD/Float32-safe output eltype @@ -333,7 +341,7 @@ function analysis_packed_ml(cfg::SHTConfig, im::Int, Vr_m::AbstractVector{<:Comp fill!(Ql, zero(CT)) P = Vector{Float64}(undef, ltr + 1) - scaleφ = cfg.cphi # Match full transform normalization + scaleφ = _analysis_phi_scale(cfg) # Match full transform normalization xv = cfg.x; wv = cfg.w # hoist field reads out of the i/l loops (cfg is mutable, so not auto-hoisted) for i in 1:nlat @@ -352,20 +360,20 @@ function analysis_packed_ml(cfg::SHTConfig, im::Int, Vr_m::AbstractVector{<:Comp end """ - synthesis_packed_ml(cfg, im, Ql, ltr) -> Vector{<:Complex} + synthesis_packed_ml(cfg, mval, Ql, ltr) -> Vector{<:Complex} Transform spherical harmonic coefficients for specific mode m to spatial field. -`im` is the zero-based stored-order index (`m = im * cfg.mres`); `Ql` +`mval` is the zero-based stored-order index (`m = mval * cfg.mres`); `Ql` contains coefficients for degrees l = m..ltr. Returns complex spatial values for that azimuthal mode. """ -function synthesis_packed_ml(cfg::SHTConfig, im::Int, Ql::AbstractVector{<:Complex}, ltr::Integer) +function synthesis_packed_ml(cfg::SHTConfig, mval::Int, Ql::AbstractVector{<:Complex}, ltr::Integer) nlat = cfg.nlat - im >= 0 || throw(ArgumentError("im must be >= 0")) - im <= cfg.mmax ÷ cfg.mres || throw(ArgumentError("im must be <= mmax/mres=$(cfg.mmax ÷ cfg.mres)")) - m = im * cfg.mres + mval >= 0 || throw(ArgumentError("mval must be >= 0")) + mval <= cfg.mmax ÷ cfg.mres || throw(ArgumentError("mval must be <= mmax/mres=$(cfg.mmax ÷ cfg.mres)")) + m = mval * cfg.mres ltr = _validate_degree_limit(cfg, ltr) - ltr >= m || throw(ArgumentError("ltr must be >= im*mres=$(m)")) + ltr >= m || throw(ArgumentError("ltr must be >= mval*mres=$(m)")) expected_len = ltr - m + 1 length(Ql) == expected_len || throw(DimensionMismatch("Ql length must be $(expected_len)")) @@ -432,15 +440,15 @@ function synthesis_axisym_l(::CPU, cfg::SHTConfig, _require_cpu_storage(:synthesis_axisym_l, coefficients) return synthesis_axisym_l(cfg, coefficients, ltr) end -function analysis_packed_ml(::CPU, cfg::SHTConfig, im::Int, +function analysis_packed_ml(::CPU, cfg::SHTConfig, mval::Int, mode::AbstractVector{<:Complex}, ltr::Integer) _require_cpu_storage(:analysis_packed_ml, mode) - return analysis_packed_ml(cfg, im, mode, ltr) + return analysis_packed_ml(cfg, mval, mode, ltr) end -function synthesis_packed_ml(::CPU, cfg::SHTConfig, im::Int, +function synthesis_packed_ml(::CPU, cfg::SHTConfig, mval::Int, coefficients::AbstractVector{<:Complex}, ltr::Integer) _require_cpu_storage(:synthesis_packed_ml, coefficients) - return synthesis_packed_ml(cfg, im, coefficients, ltr) + return synthesis_packed_ml(cfg, mval, coefficients, ltr) end """ @@ -538,7 +546,10 @@ function synthesis_point(cfg::SHTConfig, Qlm::AbstractMatrix{<:Complex}, cost::R result += 2 * real(gm * phase) end - return result + # Same φ convention factor the grid `synthesis` carries (1 under :dft; + # 1/2π under :quad), so a point evaluation matches the grid it samples. + sφ = _evaluator_phi_scale(cfg) + return sφ == 1 ? result : result * sφ end function synthesis_point(::CPU, cfg::SHTConfig, diff --git a/src/vorticity_diagnostics.jl b/src/vorticity_diagnostics.jl index 2c3a9c35..4b408d5a 100644 --- a/src/vorticity_diagnostics.jl +++ b/src/vorticity_diagnostics.jl @@ -210,7 +210,14 @@ function grad_loss_vorticity_Tlm(cfg::SHTConfig, Tlm::AbstractMatrix, ζ_target: gζlm = analysis(cfg, residual) # Analysis includes the loss's quadrature weights, but the adjoint also # carries synthesis's inverse-FFT scale (1 for :dft, 1/2π for :quad). - synthesis_scale = phi_inv_scale(cfg) / cfg.nlon + # `analysis` is used here as a stand-in for the synthesis adjoint: with + # σ = phi_inv_scale(cfg)/nlon the spatial factor of synthesis, and + # g = ∂L/∂ζ_grid = cphi·w·r, the true adjoint is + # ∂L/∂ζlm = σ · Λᴴ(g) = σ · analysis_dft(r), + # where `analysis_dft` is analysis with the fixed cphi factor. Analysis now + # carries cphi/σ instead, i.e. `analysis(r) = analysis_dft(r)/σ`, so the + # compensation picks up σ a second time. Exactly 1 in the default :dft mode. + synthesis_scale = (phi_inv_scale(cfg) / cfg.nlon)^2 # Apply chain rule: ∂L/∂T_lm = ∂L/∂ζ_lm * ∂ζ_lm/∂T_lm lmax, mmax = cfg.lmax, cfg.mmax @@ -218,7 +225,9 @@ function grad_loss_vorticity_Tlm(cfg::SHTConfig, Tlm::AbstractMatrix, ζ_target: gT = similar(Tlm) fill!(gT, 0.0) - for m in 0:mmax, l in max(1,m):lmax + # Stride by mres like every other diagnostic: orders that are not multiples + # of mres have no packed storage, so a gradient entry there is meaningless. + for m in 0:cfg.mres:mmax, l in max(1,m):lmax L2 = l * (l + 1) # Note: negative sign from ζ = -l(l+1)T gT[l+1, m+1] = -L2 * synthesis_scale * _convention_metric(scale_matrix, l, m) * gζlm[l+1, m+1] end @@ -239,13 +248,13 @@ function loss_and_grad_vorticity_Tlm(cfg::SHTConfig, Tlm::AbstractMatrix, ζ_tar # Backward pass for gradient gζlm = analysis(cfg, residual) - synthesis_scale = phi_inv_scale(cfg) / cfg.nlon + synthesis_scale = (phi_inv_scale(cfg) / cfg.nlon)^2 # see grad_loss_vorticity_Tlm lmax, mmax = cfg.lmax, cfg.mmax scale_matrix = _diagnostic_scale_matrix(cfg) gT = similar(Tlm) fill!(gT, 0.0) - for m in 0:mmax, l in max(1,m):lmax + for m in 0:cfg.mres:mmax, l in max(1,m):lmax L2 = l * (l + 1) gT[l+1, m+1] = -L2 * synthesis_scale * _convention_metric(scale_matrix, l, m) * gζlm[l+1, m+1] end diff --git a/test/fixtures/compatibility/host_transfer_allowlist.toml b/test/fixtures/compatibility/host_transfer_allowlist.toml index b0dab7e9..d000b082 100644 --- a/test/fixtures/compatibility/host_transfer_allowlist.toml +++ b/test/fixtures/compatibility/host_transfer_allowlist.toml @@ -1,5 +1,5 @@ [audit] -entry_count = 685 +entry_count = 676 scanner = "deterministic non-overlapping spelling scanner; longer transfer spellings win" scope = "src/device_utils.jl and every ext/*.jl file" @@ -4229,6 +4229,30 @@ snippet_sha256 = "11154a1e24a1c10bab3991372fe0861ea6f27729071474ecfaf5c7b46aaa64 token = "copy" [[entry]] classification = "cpu_only" +key = "ext/SHTnsKitAdvancedADExt.jl|copy|4f6fc4b2fb87282c2fdcfdb386afaac0c96c5d90f06c74191664013078bf75d7|1" +path = "ext/SHTnsKitAdvancedADExt.jl" +reason = "CPU-only automatic-differentiation storage; vendor PencilArray inputs must be rejected before forward work or materialization" +same_snippet_ordinal = 1 +snippet_sha256 = "4f6fc4b2fb87282c2fdcfdb386afaac0c96c5d90f06c74191664013078bf75d7" +token = "copy" +[[entry]] +classification = "cpu_only" +key = "ext/SHTnsKitAdvancedADExt.jl|copy|fcc1f975f3afbc60122289577da76a1629d323077ea5611f3f89d25ba94c433d|1" +path = "ext/SHTnsKitAdvancedADExt.jl" +reason = "CPU-only automatic-differentiation storage; vendor PencilArray inputs must be rejected before forward work or materialization" +same_snippet_ordinal = 1 +snippet_sha256 = "fcc1f975f3afbc60122289577da76a1629d323077ea5611f3f89d25ba94c433d" +token = "copy" +[[entry]] +classification = "cpu_only" +key = "ext/SHTnsKitAdvancedADExt.jl|copy|fcc1f975f3afbc60122289577da76a1629d323077ea5611f3f89d25ba94c433d|2" +path = "ext/SHTnsKitAdvancedADExt.jl" +reason = "CPU-only automatic-differentiation storage; vendor PencilArray inputs must be rejected before forward work or materialization" +same_snippet_ordinal = 2 +snippet_sha256 = "fcc1f975f3afbc60122289577da76a1629d323077ea5611f3f89d25ba94c433d" +token = "copy" +[[entry]] +classification = "cpu_only" key = "ext/SHTnsKitAdvancedADExt.jl|typed_array|6a3f945a436cfadc7d0d88754bc5e5305471ff778353833a998abc86020f20a3|1" path = "ext/SHTnsKitAdvancedADExt.jl" reason = "CPU-only automatic-differentiation storage; vendor PencilArray inputs must be rejected before forward work or materialization" @@ -4517,6 +4541,14 @@ snippet_sha256 = "0678e972f4053722b1aa5d0b265619b74a559c977aeda46711ae43d0c07bac token = "matrix" [[entry]] classification = "cpu_only" +key = "ext/SHTnsKitLoopVecExt.jl|matrix|2c5b53249ab10c4b6f67ff033f0d9e1e91daed69ccbadac17c7c37c0a6bc697a|1" +path = "ext/SHTnsKitLoopVecExt.jl" +reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" +same_snippet_ordinal = 1 +snippet_sha256 = "2c5b53249ab10c4b6f67ff033f0d9e1e91daed69ccbadac17c7c37c0a6bc697a" +token = "matrix" +[[entry]] +classification = "cpu_only" key = "ext/SHTnsKitLoopVecExt.jl|matrix|307d2a260f79d8bc361a96501373d64779ff7e0904c01cbef4e95780cb11c72a|1" path = "ext/SHTnsKitLoopVecExt.jl" reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" @@ -5188,22 +5220,6 @@ same_snippet_ordinal = 1 snippet_sha256 = "a24f0dcc14d951e081d408beb6d1acd0f4b8f86eba91bd1f035e69f9b399a2f1" token = "vector" [[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|collect|8433b6662c3ee812f9d5d61be28a22a4c30b89340af8dab6e7d989764d414287|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 1 -snippet_sha256 = "8433b6662c3ee812f9d5d61be28a22a4c30b89340af8dab6e7d989764d414287" -token = "collect" -[[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|collect|ee6ffc0f24762d622b08b06e934ab308c06696b4f9772672f4be640735ef83a5|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 1 -snippet_sha256 = "ee6ffc0f24762d622b08b06e934ab308c06696b4f9772672f4be640735ef83a5" -token = "collect" -[[entry]] classification = "cpu_only" key = "ext/SHTnsKitParallelExt.jl|matrix|023d49d77839bf58bb374b6d6f6731bde879032d5c394bb3724750c7449bc2f0|1" path = "ext/SHTnsKitParallelExt.jl" @@ -5221,22 +5237,6 @@ snippet_sha256 = "1ac610d54c86c71678c376f4d15e96d723b3620896bbbec9914099de0a5527 token = "matrix" [[entry]] classification = "cpu_only" -key = "ext/SHTnsKitParallelExt.jl|matrix|1ac610d54c86c71678c376f4d15e96d723b3620896bbbec9914099de0a552729|2" -path = "ext/SHTnsKitParallelExt.jl" -reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" -same_snippet_ordinal = 2 -snippet_sha256 = "1ac610d54c86c71678c376f4d15e96d723b3620896bbbec9914099de0a552729" -token = "matrix" -[[entry]] -classification = "cpu_only" -key = "ext/SHTnsKitParallelExt.jl|matrix|5ded71272574f235927bbe03113b5064dc30aef85f7c63746d6cf16394d8c7eb|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" -same_snippet_ordinal = 1 -snippet_sha256 = "5ded71272574f235927bbe03113b5064dc30aef85f7c63746d6cf16394d8c7eb" -token = "matrix" -[[entry]] -classification = "cpu_only" key = "ext/SHTnsKitParallelExt.jl|matrix|ee8d49b69190631fea6f53ce92fca583041b30b039bf9433f864b4b17fe0e0cb|1" path = "ext/SHTnsKitParallelExt.jl" reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" @@ -5244,22 +5244,6 @@ same_snippet_ordinal = 1 snippet_sha256 = "ee8d49b69190631fea6f53ce92fca583041b30b039bf9433f864b4b17fe0e0cb" token = "matrix" [[entry]] -classification = "cpu_only" -key = "ext/SHTnsKitParallelExt.jl|matrix|f93b14319dca66cd64aefa6f03aef33dd45906d1efec458970c7d8b7686b8c45|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" -same_snippet_ordinal = 1 -snippet_sha256 = "f93b14319dca66cd64aefa6f03aef33dd45906d1efec458970c7d8b7686b8c45" -token = "matrix" -[[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|parent|58349167721d1a9bba7d82a45f001f1713623da1b1d945cd7fd7b5d2220272b0|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 1 -snippet_sha256 = "58349167721d1a9bba7d82a45f001f1713623da1b1d945cd7fd7b5d2220272b0" -token = "parent" -[[entry]] classification = "metadata_or_storage_preserving" key = "ext/SHTnsKitParallelExt.jl|parent|6c8ce9fa30f5cb7ab5f6604ca8adc56134475f81939910b6f8673263e5727b5d|1" path = "ext/SHTnsKitParallelExt.jl" @@ -5276,54 +5260,6 @@ same_snippet_ordinal = 1 snippet_sha256 = "82071ef968af13d92eff3e6d2972b5e8cfb8a6c508d32f3442eed291de6c3608" token = "parent" [[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|parent|afdf86c69aa5206e89b5435a67f6b02a00fd99b8f9037bc58414169e05fef202|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 1 -snippet_sha256 = "afdf86c69aa5206e89b5435a67f6b02a00fd99b8f9037bc58414169e05fef202" -token = "parent" -[[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|parent|afdf86c69aa5206e89b5435a67f6b02a00fd99b8f9037bc58414169e05fef202|2" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 2 -snippet_sha256 = "afdf86c69aa5206e89b5435a67f6b02a00fd99b8f9037bc58414169e05fef202" -token = "parent" -[[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|parent|afdf86c69aa5206e89b5435a67f6b02a00fd99b8f9037bc58414169e05fef202|3" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 3 -snippet_sha256 = "afdf86c69aa5206e89b5435a67f6b02a00fd99b8f9037bc58414169e05fef202" -token = "parent" -[[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|parent|b026f8b420c9cf2cebde304c52634b3b660d8980c84d292bcb38f0ef35a921fb|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 1 -snippet_sha256 = "b026f8b420c9cf2cebde304c52634b3b660d8980c84d292bcb38f0ef35a921fb" -token = "parent" -[[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|parent|da2163864051cb7dbb4972619469a9e138391299cb05766ebabaacd422148cc8|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 1 -snippet_sha256 = "da2163864051cb7dbb4972619469a9e138391299cb05766ebabaacd422148cc8" -token = "parent" -[[entry]] -classification = "metadata_or_storage_preserving" -key = "ext/SHTnsKitParallelExt.jl|parent|e7146bfbd323b0a2294e29c6c4cd684d22e89b8cabb06f79d782705e4548fc65|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "metadata collection or storage-preserving parent/copy operation; verify that the destination retains the source storage family" -same_snippet_ordinal = 1 -snippet_sha256 = "e7146bfbd323b0a2294e29c6c4cd684d22e89b8cabb06f79d782705e4548fc65" -token = "parent" -[[entry]] classification = "cpu_only" key = "ext/SHTnsKitParallelExt.jl|vector|4fa452b93b779a5d536c482a3ffcf627ff0134c6afb925a707fc312f872d483a|1" path = "ext/SHTnsKitParallelExt.jl" @@ -5349,14 +5285,6 @@ snippet_sha256 = "7f628a4fb8031eb06efa9ec7c386e111cfe896aabecc8c954a393f02c5b975 token = "vector" [[entry]] classification = "cpu_only" -key = "ext/SHTnsKitParallelExt.jl|vector|8433b6662c3ee812f9d5d61be28a22a4c30b89340af8dab6e7d989764d414287|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" -same_snippet_ordinal = 1 -snippet_sha256 = "8433b6662c3ee812f9d5d61be28a22a4c30b89340af8dab6e7d989764d414287" -token = "vector" -[[entry]] -classification = "cpu_only" key = "ext/SHTnsKitParallelExt.jl|vector|a93a2bfa0355e3c7cd001f93b48c1ff7a08709afb7fb408782aedb91c19250dc|1" path = "ext/SHTnsKitParallelExt.jl" reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" @@ -5381,19 +5309,19 @@ snippet_sha256 = "a93a2bfa0355e3c7cd001f93b48c1ff7a08709afb7fb408782aedb91c19250 token = "vector" [[entry]] classification = "cpu_only" -key = "ext/SHTnsKitParallelExt.jl|vector|ee6ffc0f24762d622b08b06e934ab308c06696b4f9772672f4be640735ef83a5|1" -path = "ext/SHTnsKitParallelExt.jl" -reason = "explicit host workspace in a CPU-only algorithm; never route a vendor-backed mathematical result through this allocation" +key = "ext/SHTnsKitZygoteExt.jl|copy|11154a1e24a1c10bab3991372fe0861ea6f27729071474ecfaf5c7b46aaa6421|1" +path = "ext/SHTnsKitZygoteExt.jl" +reason = "CPU-only automatic-differentiation storage; vendor PencilArray inputs must be rejected before forward work or materialization" same_snippet_ordinal = 1 -snippet_sha256 = "ee6ffc0f24762d622b08b06e934ab308c06696b4f9772672f4be640735ef83a5" -token = "vector" +snippet_sha256 = "11154a1e24a1c10bab3991372fe0861ea6f27729071474ecfaf5c7b46aaa6421" +token = "copy" [[entry]] classification = "cpu_only" -key = "ext/SHTnsKitZygoteExt.jl|copy|11154a1e24a1c10bab3991372fe0861ea6f27729071474ecfaf5c7b46aaa6421|1" +key = "ext/SHTnsKitZygoteExt.jl|copy|fcc1f975f3afbc60122289577da76a1629d323077ea5611f3f89d25ba94c433d|1" path = "ext/SHTnsKitZygoteExt.jl" reason = "CPU-only automatic-differentiation storage; vendor PencilArray inputs must be rejected before forward work or materialization" same_snippet_ordinal = 1 -snippet_sha256 = "11154a1e24a1c10bab3991372fe0861ea6f27729071474ecfaf5c7b46aaa6421" +snippet_sha256 = "fcc1f975f3afbc60122289577da76a1629d323077ea5611f3f89d25ba94c433d" token = "copy" [[entry]] classification = "cpu_only" diff --git a/test/fixtures/compatibility/task16_gate.toml b/test/fixtures/compatibility/task16_gate.toml index 0d3398ae..c494a1f8 100644 --- a/test/fixtures/compatibility/task16_gate.toml +++ b/test/fixtures/compatibility/task16_gate.toml @@ -6,13 +6,13 @@ platform = "aarch64-apple-darwin" [local_gate] evidence_kind = "local_readiness" certifying = false -audited_tree_digest = "1bae4f3a0acc933d30369fe3e11d2333379b40dbc55e4dbadf9c705bb860cbd3" +audited_tree_digest = "2cb7eee1644d1df5613a0b39e188ee634a23aceafef5f49b418cde59717e3643" audited_scope = ["Project.toml", ".github/workflows", "docs/src/shtns37-parity.md", "ext", "src", "test"] audited_exclusions = ["test/fixtures/compatibility/task16_gate.toml", "test/fixtures/compatibility/task16_local_commands.txt", "test/fixtures/compatibility/task16_local_summary.log"] commands_file = "test/fixtures/compatibility/task16_local_commands.txt" -commands_sha256 = "524f0421f4f3aa013f8e4805b5ae2a0d5ffe5c858ff8140f062a77e6c8447712" +commands_sha256 = "9f61a016bc43244a56c24ef18d2a1c0a9aefd2be09d24745efd55d82278b7ea8" summary_log = "test/fixtures/compatibility/task16_local_summary.log" -summary_log_sha256 = "c41fe35538fca5cd7b927bbff92c39412816c5f46339fba4a3ab5da518a59818" +summary_log_sha256 = "63082905f4212b90b67c09834f35473c19e8dac945e42e9d48efdd494b50b251" # Historical readiness counts; current review runs are recorded in the summary. pkg_exit_code = 0 serial_pass = 76347 @@ -26,7 +26,7 @@ mpi_cpu_fixture_pass_per_rank = 66 [host_transfer_audit] allowlist = "test/fixtures/compatibility/host_transfer_allowlist.toml" scanner = "test/support/host_transfer_inventory.jl" -entry_count = 685 +entry_count = 676 allowscalar_match_count = 0 legacy_cuda_host_result_bridges = 2 diff --git a/test/fixtures/compatibility/task16_local_commands.txt b/test/fixtures/compatibility/task16_local_commands.txt index b74eb75a..9a4c22f5 100644 --- a/test/fixtures/compatibility/task16_local_commands.txt +++ b/test/fixtures/compatibility/task16_local_commands.txt @@ -5,7 +5,7 @@ certifying=false manual_cuda_tool_command=python3 -B -m unittest discover -s test/support -p test_manual_cuda.py -v manual_cuda_contract_command=/Users/subha/.julia/juliaup/julia-1.12.4+0.aarch64.apple.darwin14/bin/julia --startup-file=no --project=. test/serial/test_shtns37_contract.jl manual_cuda_numerical_source_unchanged_from=3b8b2fe118900f2431208088932631db4a162908 -audited_tree_digest=1bae4f3a0acc933d30369fe3e11d2333379b40dbc55e4dbadf9c705bb860cbd3 +audited_tree_digest=2cb7eee1644d1df5613a0b39e188ee634a23aceafef5f49b418cde59717e3643 audited_scope=Project.toml,.github/workflows,docs/src/shtns37-parity.md,ext,src,test audited_exclusions=test/fixtures/compatibility/task16_gate.toml,test/fixtures/compatibility/task16_local_commands.txt,test/fixtures/compatibility/task16_local_summary.log pkg_command=/usr/bin/env JULIA_DEPOT_PATH=/private/tmp/julia_depot_shtnskit:/Users/subha/.julia SHTNSKIT_RUN_JET_TESTS=1 SHTNSKIT_RUN_AQUA_TESTS=1 /Users/subha/.julia/juliaup/julia-1.12.4+0.aarch64.apple.darwin14/bin/julia --startup-file=no --project=. -e 'using Pkg; Pkg.test()' @@ -67,3 +67,26 @@ qst_communicator_full_parity_file=test/parity/runtests_mpi.jl with no ARGS qst_communicator_mpi_timeout_seconds=900 per test file, including JIT compilation qst_communicator_audit_command=using SHTnsKit; include("test/serial/test_final_parity_gate.jl") qst_communicator_full_parity_ranks=2,4 + +# Convention/AD/dead-cache audit integration, 2026-09-13. Historical and prior +# review runs above are NOT re-runs; the keys below record only what was +# executed for this change. +audit2_date=2026-09-13 +audit2_base_commit=f06d13d1 +audit2_julia=/Users/subha/.julia/juliaup/julia-1.11.1+0.aarch64.apple.darwin14/bin/julia +audit2_julia_version=1.11.1 +audit2_platform=arm64-apple-darwin22.4.0 +audit2_threads=1 +audit2_serial_env=/private/tmp/shtkit_testenv (Test, Random, ForwardDiff, Zygote, ChainRulesCore, LoopVectorization; SHTnsKit dev'd from the worktree) +audit2_mpi_env=/private/tmp/shtkit_mpienv (MPI, PencilArrays, PencilFFTs; SHTnsKit dev'd from the worktree) +audit2_serial_command=julia --startup-file=no --project=/private/tmp/shtkit_testenv test/serial/runtests.jl +audit2_parallel_grid_command=julia --startup-file=no --project=/private/tmp/shtkit_mpienv test/parallel/runtests.jl +audit2_mpi_launcher=MPI.mpiexec() from the MPI env with the same Julia executable and project +audit2_mpi_files_4rank=test/parallel/test_mpi_comprehensive.jl,test/parallel/test_mpi_extended.jl,test/parallel/test_mpi_audit_fixes.jl,test/parallel/test_mpi_2d_alignment.jl +audit2_mpi_files_2rank=test/parallel/test_mpi_parallel_local_correctness.jl,test/parallel/test_mpi_comm_cleanup.jl,test/parallel/test_mpi_plan_preflight.jl,test/parallel/test_mpi_transpose_operand_preflight.jl,test/parallel/test_transpose_sht.jl,test/parallel/test_disttranspose_dealiased.jl,test/parallel/test_dist_plan_alloc.jl +audit2_host_transfer_regen_command=julia --startup-file=no --project=/private/tmp/shtkit_testenv test/support/generate_host_transfer_allowlist.jl +audit2_pkg_test_rerun=false +audit2_jet_rerun=false +audit2_aqua_rerun=false +audit2_gpu_rerun=false +audit2_parity_runtests_rerun=false diff --git a/test/fixtures/compatibility/task16_local_summary.log b/test/fixtures/compatibility/task16_local_summary.log index 0524c64e..44c8d5c8 100644 --- a/test/fixtures/compatibility/task16_local_summary.log +++ b/test/fixtures/compatibility/task16_local_summary.log @@ -9,7 +9,7 @@ manual_cuda_no_hardware_guard_rejected=true manual_cuda_julia_syntax_pass=3 manual_cuda_github_status_published=false manual_cuda_numerical_source_unchanged_from=3b8b2fe118900f2431208088932631db4a162908 -audited_tree_digest=1bae4f3a0acc933d30369fe3e11d2333379b40dbc55e4dbadf9c705bb860cbd3 +audited_tree_digest=2cb7eee1644d1df5613a0b39e188ee634a23aceafef5f49b418cde59717e3643 platform=aarch64-apple-darwin julia=1.12.4 pkg_exit_code=0 @@ -107,3 +107,36 @@ qst_communicator_full_parity_ranks=2,4 qst_communicator_full_parity_exit_code=0 qst_communicator_qst_native_pass_per_rank=55 qst_communicator_shtns37_fixture_pass_per_rank=66 + +# Convention/AD/dead-cache audit integration results, 2026-09-13. Results above +# are historical and were not re-run. +audit2_date=2026-09-13 +audit2_base_commit=f06d13d1 +audit2_julia=1.11.1 +audit2_platform=arm64-apple-darwin22.4.0 +audit2_threads=1 +audit2_serial_pass=76848 +audit2_serial_fail=0 +audit2_serial_exit_code=0 +audit2_parallel_grid_pass=1714 +audit2_parallel_grid_exit_code=0 +audit2_mpi_ranks_4=4 +audit2_mpi_2d_alignment_pass_per_rank=15 +audit2_mpi_audit_fixes_pass_per_rank=41 +audit2_mpi_comprehensive_result=passed +audit2_mpi_extended_result=passed +audit2_mpi_ranks_2=2 +audit2_mpi_parallel_local_pass_per_rank=172 +audit2_mpi_comm_cleanup_pass_per_rank=4 +audit2_mpi_plan_preflight_pass_per_rank=8 +audit2_mpi_transpose_operand_preflight_pass_per_rank=25 +audit2_mpi_exit_code=0 +audit2_host_transfer_entry_count=676 +audit2_host_transfer_regenerated=true +audit2_pkg_test_rerun=false +audit2_jet_rerun=false +audit2_aqua_rerun=false +audit2_gpu_rerun=false +audit2_physical_gpu_execution=false +audit2_parity_runtests_rerun=false +audit2_note=Run on Julia 1.11.1, not the 1.12.4 of the historical readiness runs; no CUDA or ROCm hardware is present on this machine. diff --git a/test/parallel/runtests.jl b/test/parallel/runtests.jl index 7a80e252..7cb96935 100644 --- a/test/parallel/runtests.jl +++ b/test/parallel/runtests.jl @@ -13,6 +13,7 @@ # - test_disttranspose_dealiased.jl : DistTransposePlan on dealiased nlon>2*mmax+1 (mpiexec) # - test_dist_plan_alloc.jl : DistAnalysisPlan correctness + per-call allocation budget (mpiexec) # - test_mpi_audit_fixes.jl : Regressions for the 2026-08 distributed audit fixes (mpiexec) +# - test_mpi_2d_alignment.jl : 2D spectral-plan alignment preconditions (mpiexec, 4 ranks) # # To run MPI tests: # mpiexec -n 4 julia --project test/parallel/test_mpi_comprehensive.jl @@ -21,6 +22,7 @@ # mpiexec -n 2 julia --project test/parallel/test_disttranspose_dealiased.jl # mpiexec -n 2 julia --project test/parallel/test_dist_plan_alloc.jl # mpiexec -n 4 julia --project test/parallel/test_mpi_audit_fixes.jl +# mpiexec -n 4 julia --project test/parallel/test_mpi_2d_alignment.jl # mpiexec -n 1 julia --project test/parallel/test_parallel_ad_storage.jl using Test diff --git a/test/parallel/test_mpi_2d_alignment.jl b/test/parallel/test_mpi_2d_alignment.jl new file mode 100644 index 00000000..92180eb7 --- /dev/null +++ b/test/parallel/test_mpi_2d_alignment.jl @@ -0,0 +1,112 @@ +# SHTnsKit.jl - 2D spectral-plan alignment contracts (run with mpiexec -n 4) +# +# The optimized 2D routines reduce within `l_comm` (ranks sharing an `m_rank`) +# instead of over the whole communicator. That is only valid when the spatial θ +# split lines up with the spectral l split, so that every rank in such a group +# owns a DISTINCT θ slab. Where the precondition holds these paths agree with +# the safe ones to ~1e-14; where it does not they used to return silent garbage: +# +# dist_synthesis_distributed_2d_optimized max|err| = 84.4 (field of O(10)) +# dist_analysis_distributed_2d(assume_aligned=true) rel err = 1.15 +# +# Both now validate collectively and raise. This file pins both directions. +# +# (This is also why `_dist_analysis_2d_aligned` does NOT call +# `_keep_one_phi_partner!` the way its full-comm siblings do: within `l_comm` +# the alignment precondition already guarantees there are no duplicate slabs.) + +using Test, MPI, PencilArrays, PencilFFTs, SHTnsKit, LinearAlgebra, Random + +MPI.Initialized() || MPI.Init() +const COMM = MPI.COMM_WORLD +const RANK = MPI.Comm_rank(COMM) +const NP = MPI.Comm_size(COMM) +const EXT = Base.get_extension(SHTnsKit, :SHTnsKitParallelExt) + +NP == 4 || error("test_mpi_2d_alignment.jl expects 4 ranks, got $NP") + +const LMAX = 8 +const NLAT = LMAX + 2 +const NLON = 2 * LMAX + 2 +const CFG = create_gauss_config(LMAX, NLAT; nlon=NLON) + +Random.seed!(1234) +const FGLOB = randn(NLAT, NLON) +const AREF = analysis(CFG, FGLOB) + +const ALM = let a = zeros(ComplexF64, LMAX + 1, CFG.mmax + 1) + for m in 0:CFG.mmax, l in m:LMAX + a[l+1, m+1] = m == 0 ? Float64(l + 1) : ComplexF64(l + 1, m) + end + a +end +const FREF = synthesis(CFG, ALM) + +"""Build a spatial PencilArray on a `pθ × pφ` process grid, filled from `FGLOB`.""" +function spatial(pθ, pφ) + dims = pφ == 1 ? (1,) : (pθ == 1 ? (2,) : (1, 2)) + pen = Pencil((NLAT, NLON), dims, COMM) + f = PencilArray{Float64}(undef, pen) + lr = PencilArrays.range_local(pen) + loc = parent(f) + for (jj, jg) in enumerate(lr[2]), (ii, ig) in enumerate(lr[1]) + loc[ii, jj] = FGLOB[ig, jg] + end + f +end + +"""Largest spatial error against the serial reference, reduced over all ranks.""" +function spatial_err(out, prototype) + loc = out isa PencilArray ? parent(out) : out + lr = PencilArrays.range_local(PencilArrays.pencil(prototype)) + e = 0.0 + for (jj, jg) in enumerate(lr[2]), (ii, ig) in enumerate(lr[1]) + e = max(e, abs(loc[ii, jj] - FREF[ig, jg])) + end + MPI.Allreduce(e, max, COMM) +end + +@testset "2D plan alignment contracts (4 ranks)" begin + for (pθ, pφ) in ((4, 1), (1, 4), (2, 2)) + f = spatial(pθ, pφ) + plan = EXT.create_distributed_spectral_plan_2d(LMAX, CFG.mmax, COMM; p_l=pθ, p_m=pφ) + try + aligned, _ = EXT.validate_2d_distribution_alignment(plan, f) + # The verdict must be identical on every rank, or the guards below + # would throw on some ranks and enter a collective on others. + @test MPI.Allreduce(aligned, &, COMM) == MPI.Allreduce(aligned, |, COMM) + + @testset "pθ=$pθ × pφ=$pφ (aligned=$aligned)" begin + # The safe paths are correct on every decomposition. + Asafe = EXT.dist_analysis_distributed_2d(CFG, f; plan=plan, assume_aligned=false) + @test maximum(abs, EXT.gather_to_full_dense_2d(Asafe) .- AREF) < + 1e-12 * maximum(abs, AREF) + + dsa = EXT.create_distributed_spectral_array_2d(plan) + EXT.scatter_from_dense_2d!(dsa, ALM) + @test spatial_err(EXT.dist_synthesis_distributed_2d(CFG, dsa; prototype_θφ=f), f) < 1e-11 + + if aligned + # Optimized paths agree with the safe ones where the + # precondition holds. + Aal = EXT.dist_analysis_distributed_2d(CFG, f; plan=plan, assume_aligned=true) + @test maximum(abs, EXT.gather_to_full_dense_2d(Aal) .- AREF) < + 1e-12 * maximum(abs, AREF) + @test spatial_err( + EXT.dist_synthesis_distributed_2d_optimized(CFG, dsa; prototype_θφ=f), f) < 1e-11 + else + # ...and raise, rather than returning garbage, where it does not. + @test_throws ArgumentError EXT.dist_analysis_distributed_2d( + CFG, f; plan=plan, assume_aligned=true) + @test_throws ArgumentError EXT.dist_synthesis_distributed_2d_optimized( + CFG, dsa; prototype_θφ=f) + end + end + finally + close(plan) + end + end +end + +MPI.Barrier(COMM) +RANK == 0 && println("\nAll 2D alignment contract tests PASSED") diff --git a/test/serial/test_basic_transforms.jl b/test/serial/test_basic_transforms.jl index 4f18e287..a65db1f9 100644 --- a/test/serial/test_basic_transforms.jl +++ b/test/serial/test_basic_transforms.jl @@ -8,6 +8,17 @@ using SHTnsKit @isdefined(VERBOSE) || (const VERBOSE = get(ENV, "SHTNSKIT_TEST_VERBOSE", "0") == "1") +# Allocation budgets below are calibrated for a SINGLE-THREADED run, where they +# are exact. The shared m-loop orchestrators start `@threads` tasks whenever +# threads are available, and each threaded region costs a few hundred bytes to +# spawn — a constant independent of problem size, and not the kind of regression +# these budgets exist to catch (a per-row temporary or a dense zero spectrum +# costs tens of KB and grows with lmax). Allow a per-thread slack rather than +# letting the whole suite go red on any multi-core machine; note this makes the +# budgets coarse at `nthreads > 1`, so the tight check is the 1-thread run. +@isdefined(_thread_alloc_slack) || + (_thread_alloc_slack() = Threads.nthreads() > 1 ? 4_096 * Threads.nthreads() : 0) + @testset "Basic Scalar Transforms" begin @testset "Real coefficient matrices" begin for T in (Float32, Float64, Int), norm in (:orthonormal, :schmidt) @@ -295,7 +306,7 @@ using SHTnsKit # Julia/FFTW patch versions can impose a small constant allocation # floor around plan execution. Keep the budget below scratch-sized # allocations while allowing the observed Julia 1.10 Linux floor. - rfft_alloc_budget = 2_048 + rfft_alloc_budget = 2_048 + _thread_alloc_slack() @test @allocated(analysis!(cfg, alm_out, f; fft_scratch=rfft_scratch, use_rfft=true)) <= rfft_alloc_budget @test @allocated(synthesis!(cfg, f_out, alm_c; real_output=true, fft_scratch=rfft_scratch, use_rfft=true)) <= rfft_alloc_budget @@ -374,4 +385,71 @@ using SHTnsKit @test isapprox(Ql_rec_l, Ql_rec_full[1:ltr+1]; rtol=1e-10, atol=1e-12) end + + @testset "analysis inverts synthesis under every phi_scale" begin + # `synthesis` honoured `phi_scale` while `analysis` always applied a fixed + # `cphi`, so the two halves of the pair disagreed about the convention: + # under :quad `analysis(synthesis(alm))` came back as `alm/2π` exactly. + rng = MersenneTwister(7731) + for mode in (:dft, :quad) + cfg = create_gauss_config(6, 8) + cfg.phi_scale = mode + alm = zeros(ComplexF64, cfg.lmax + 1, cfg.mmax + 1) + for m in 0:cfg.mmax, l in m:cfg.lmax + alm[l+1, m+1] = m == 0 ? randn(rng) : complex(randn(rng), randn(rng)) + end + S = copy(alm); T = 0.5 .* alm; S[1,1] = 0; T[1,1] = 0 + + @test analysis(cfg, synthesis(cfg, alm)) ≈ alm rtol=1e-10 + Vt, Vp = synthesis_sphtor(cfg, S, T) + S2, T2 = analysis_sphtor(cfg, Vt, Vp) + @test S2 ≈ S rtol=1e-10 + @test T2 ≈ T rtol=1e-10 + + f = synthesis(cfg, alm) + @test analysis_batch(cfg, reshape(f, size(f)..., 1))[:, :, 1] ≈ alm rtol=1e-10 + + # the planned path must agree with the cfg form in both modes + plan = SHTPlan(cfg) + out = similar(alm) + analysis!(plan, out, f) + @test out ≈ alm rtol=1e-10 + + # a hand-built config must not disagree with the constructor's + # for the same grid: `:auto` used to mean `:quad` for non-Gauss grids + @test SHTnsKit.phi_inv_scale(create_regular_config(6, 10; nlon=14)) == + Float64(14) + end + end + + @testset "padded spatial buffers reach the transforms" begin + # `allocate_padded_spatial` returns nlat_padded rows and every transform + # demands exactly nlat, so the padding API had no usable path into a + # transform at all. `spatial_view` is that path, and it preserves the + # padded column stride the padding exists for. + cfg = create_gauss_config(16, 18) + set_allow_padding!(cfg) + @test get_nlat_padded(cfg) > cfg.nlat + + rng = MersenneTwister(4477) + f = randn(rng, cfg.nlat, cfg.nlon) + pad = allocate_padded_spatial(cfg) + copy_to_padded!(pad, f, cfg) + v = spatial_view(cfg, pad) + + @test size(v) == (cfg.nlat, cfg.nlon) + @test stride(v, 2) == get_nlat_padded(cfg) # padding retained + @test analysis(cfg, v) == analysis(cfg, f) # bit-identical + + batch = allocate_padded_spatial_batch(cfg, 3) + fb = randn(rng, cfg.nlat, cfg.nlon, 3) + for k in 1:3 + copy_to_padded!(view(batch, :, :, k), view(fb, :, :, k), cfg) + end + @test analysis_batch(cfg, spatial_view(cfg, batch)) == analysis_batch(cfg, fb) + + @test_throws DimensionMismatch spatial_view(cfg, zeros(cfg.nlat - 1, cfg.nlon)) + @test_throws DimensionMismatch spatial_view(cfg, zeros(cfg.nlat, cfg.nlon + 1)) + end + end diff --git a/test/serial/test_batch_qst.jl b/test/serial/test_batch_qst.jl index 9b2269fd..c55a624b 100644 --- a/test/serial/test_batch_qst.jl +++ b/test/serial/test_batch_qst.jl @@ -8,6 +8,17 @@ using SHTnsKit @isdefined(VERBOSE) || (const VERBOSE = get(ENV, "SHTNSKIT_TEST_VERBOSE", "0") == "1") +# Allocation budgets below are calibrated for a SINGLE-THREADED run, where they +# are exact. The shared m-loop orchestrators start `@threads` tasks whenever +# threads are available, and each threaded region costs a few hundred bytes to +# spawn — a constant independent of problem size, and not the kind of regression +# these budgets exist to catch (a per-row temporary or a dense zero spectrum +# costs tens of KB and grows with lmax). Allow a per-thread slack rather than +# letting the whole suite go red on any multi-core machine; note this makes the +# budgets coarse at `nthreads > 1`, so the tight check is the 1-thread run. +@isdefined(_thread_alloc_slack) || + (_thread_alloc_slack() = Threads.nthreads() > 1 ? 4_096 * Threads.nthreads() : 0) + function _real_alm(rng, lmax, mmax) a = randn(rng, ComplexF64, lmax + 1, mmax + 1) a[:, 1] .= real.(a[:, 1]) @@ -51,7 +62,7 @@ end synthesis_qst_batch(cfg, Qb, Sb, Tb) GC.gc() - @test @allocated(synthesis_qst_batch(cfg, Qb, Sb, Tb)) <= 25_000 + @test @allocated(synthesis_qst_batch(cfg, Qb, Sb, Tb)) <= 25_000 + _thread_alloc_slack() end @testset "analysis_qst_batch matches per-field analysis_qst" begin @@ -79,7 +90,7 @@ end analysis_qst_batch(cfg, Vr_b, Vt_b, Vp_b) GC.gc() - @test @allocated(analysis_qst_batch(cfg, Vr_b, Vt_b, Vp_b)) <= 28_000 + @test @allocated(analysis_qst_batch(cfg, Vr_b, Vt_b, Vp_b)) <= 28_000 + _thread_alloc_slack() end @testset "Batch QST roundtrip" begin diff --git a/test/serial/test_batch_transforms.jl b/test/serial/test_batch_transforms.jl index 654e7f9a..a322808a 100644 --- a/test/serial/test_batch_transforms.jl +++ b/test/serial/test_batch_transforms.jl @@ -7,6 +7,17 @@ using SHTnsKit @isdefined(VERBOSE) || (const VERBOSE = get(ENV, "SHTNSKIT_TEST_VERBOSE", "0") == "1") +# Allocation budgets below are calibrated for a SINGLE-THREADED run, where they +# are exact. The shared m-loop orchestrators start `@threads` tasks whenever +# threads are available, and each threaded region costs a few hundred bytes to +# spawn — a constant independent of problem size, and not the kind of regression +# these budgets exist to catch (a per-row temporary or a dense zero spectrum +# costs tens of KB and grows with lmax). Allow a per-thread slack rather than +# letting the whole suite go red on any multi-core machine; note this makes the +# budgets coarse at `nthreads > 1`, so the tight check is the 1-thread run. +@isdefined(_thread_alloc_slack) || + (_thread_alloc_slack() = Threads.nthreads() > 1 ? 4_096 * Threads.nthreads() : 0) + @testset "Batch Transforms" begin @testset "Scalar batch analysis" begin lmax = 6 @@ -105,10 +116,10 @@ using SHTnsKit synthesis_batch!(cfg, fields_out, alm_batch; fft_batch=rfft_batch, use_rfft=true) GC.gc() - @test @allocated(analysis_batch!(cfg, alm_batch, fields; fft_batch=fft_batch)) <= 8_000 - @test @allocated(synthesis_batch!(cfg, fields_out, alm_batch; fft_batch=fft_batch)) <= 8_000 - @test @allocated(analysis_batch!(cfg, alm_batch, fields; fft_batch=rfft_batch, use_rfft=true)) <= 8_000 - @test @allocated(synthesis_batch!(cfg, fields_out, alm_batch; fft_batch=rfft_batch, use_rfft=true)) <= 8_000 + @test @allocated(analysis_batch!(cfg, alm_batch, fields; fft_batch=fft_batch)) <= 8_000 + _thread_alloc_slack() + @test @allocated(synthesis_batch!(cfg, fields_out, alm_batch; fft_batch=fft_batch)) <= 8_000 + _thread_alloc_slack() + @test @allocated(analysis_batch!(cfg, alm_batch, fields; fft_batch=rfft_batch, use_rfft=true)) <= 8_000 + _thread_alloc_slack() + @test @allocated(synthesis_batch!(cfg, fields_out, alm_batch; fft_batch=rfft_batch, use_rfft=true)) <= 8_000 + _thread_alloc_slack() @inferred synthesis_batch(cfg, alm_batch) fields_kw = synthesis_batch(cfg, alm_batch; real_output=false) @@ -223,8 +234,8 @@ using SHTnsKit synthesis_sphtor_batch(cfg, Slm_batch, Tlm_batch) analysis_sphtor_batch(cfg, Vt_batch, Vp_batch) GC.gc() - @test @allocated(synthesis_sphtor_batch(cfg, Slm_batch, Tlm_batch)) <= 16_000 - @test @allocated(analysis_sphtor_batch(cfg, Vt_batch, Vp_batch)) <= 18_000 + @test @allocated(synthesis_sphtor_batch(cfg, Slm_batch, Tlm_batch)) <= 16_000 + _thread_alloc_slack() + @test @allocated(analysis_sphtor_batch(cfg, Vt_batch, Vp_batch)) <= 18_000 + _thread_alloc_slack() end @testset "QST batch transforms" begin diff --git a/test/serial/test_configuration.jl b/test/serial/test_configuration.jl index 7fe5e5cd..0c374f00 100644 --- a/test/serial/test_configuration.jl +++ b/test/serial/test_configuration.jl @@ -2,6 +2,7 @@ # Tests for grid configuration, indexing, and normalization using Test +using Random using SHTnsKit @isdefined(VERBOSE) || (const VERBOSE = get(ENV, "SHTNSKIT_TEST_VERBOSE", "0") == "1") @@ -329,33 +330,97 @@ using SHTnsKit @test size(spatial_scratch) == (nlat, nlon) end - @testset "FFT plan cache control" begin - # FFT plan cache requires the parallel extension - # Skip if not available - try - # Save initial state - initial_state = fft_plan_cache_enabled() - @test typeof(initial_state) == Bool - - # Test enable/disable - enable_fft_plan_cache!() - @test fft_plan_cache_enabled() == true - - disable_fft_plan_cache!() - @test fft_plan_cache_enabled() == false + @testset "structural field assignment keeps derived state consistent" begin + # `cfg.lmax = 10` used to leave `size(cfg.Nlm) == (7,7)` while the + # transforms index it as (lmax+1, mmax+1) under `@inbounds` — a live + # out-of-bounds read. lmax/mmax/mres now rebuild the spectral layout; + # the fields that cannot be made consistent are rejected outright. + cfg = create_gauss_config(6, 8) + prepare_plm_tables!(cfg) + @test SHTnsKit.has_fused_scalar_tables(cfg) + + cfg.lmax = 10 + @test size(cfg.Nlm) == (cfg.lmax + 1, cfg.mmax + 1) + @test cfg.nlm == SHTnsKit.nlm_calc(cfg.lmax, cfg.mmax, cfg.mres) + @test length(cfg.li) == cfg.nlm && length(cfg.mi) == cfg.nlm + @test !SHTnsKit.has_fused_scalar_tables(cfg) # stale tables dropped + + cfg2 = create_gauss_config(6, 8; mres=1) + cfg2.mres = 2 + @test cfg2.nlm == SHTnsKit.nlm_calc(6, cfg2.mmax, 2) + @test all(m -> m % 2 == 0, cfg2.mi) + + for bad in (:nlat, :nlon, :grid_type, :nlm, :nspat) + @test_throws ArgumentError setproperty!(create_gauss_config(4, 6), bad, + bad === :grid_type ? :regular : 99) + end + # An inconsistent spectral triple is rejected rather than stored. + @test_throws ArgumentError (c = create_gauss_config(6, 8); c.mmax = 99) + end - # Test set function - set_fft_plan_cache!(true) - @test fft_plan_cache_enabled() == true + @testset "kwarg SHTConfig constructor validates its invariants" begin + # This constructor is exported and used to check nothing, so a hand-built + # config could violate nlon >= 2*mmax+1 and then silently synthesize an + # all-zero field for any mode it could not resolve. + lmax = mmax = 6 + mk(; nlon=2*mmax + 2, nlat=8, kw...) = begin + θ = collect(range(0.1, 3.0; length=nlat)) + base = (; lmax, mmax, mres=1, nlat, nlon, + θ, φ=collect(range(0, 2π; length=nlon+1))[1:nlon], + x=cos.(θ), w=fill(2/nlat, nlat), st=sin.(θ), + Nlm=SHTnsKit.Nlm_table(lmax, mmax), cphi=2π/nlon, + nlm=SHTnsKit.nlm_calc(lmax, mmax, 1), + li=SHTnsKit.build_li_mi(lmax, mmax, 1)[1], + mi=SHTnsKit.build_li_mi(lmax, mmax, 1)[2], + nspat=nlat*nlon, norm=:orthonormal, cs_phase=true, + real_norm=false, robert_form=false) + SHTnsKit.SHTConfig(; base..., kw...) + end + @test mk() isa SHTnsKit.SHTConfig # the valid case still builds + @test_throws ArgumentError mk(nlon=8) # nlon < 2*mmax+1 + @test_throws ArgumentError mk(mres=0) + @test_throws ArgumentError mk(nspat=1) # nspat != nlat*nlon + @test_throws DimensionMismatch mk(w=fill(0.25, 4)) # w shorter than nlat + @test_throws DimensionMismatch mk(Nlm=zeros(2, 2)) # Nlm not (lmax+1, mmax+1) + end - set_fft_plan_cache!(false) - @test fft_plan_cache_enabled() == false + @testset "FFT plan cache control" begin + # These knobs used to forward to a cache in the parallel extension that + # nothing ever read (`_get_or_plan` had no call sites), so the whole + # documented feature was a no-op and this testset was wrapped in a + # try/catch that skipped it whenever MPI was absent. They now control the + # φ-FFT plan cache every transform actually goes through, serial or + # distributed, so assert the behaviour rather than just the flag. + initial_state = fft_plan_cache_enabled() + @test typeof(initial_state) == Bool - # Restore initial state - set_fft_plan_cache!(initial_state) - catch e - @info "Skipping FFT plan cache tests (requires parallel extension)" exception=e - end + enable_fft_plan_cache!() + @test fft_plan_cache_enabled() == true + disable_fft_plan_cache!() + @test fft_plan_cache_enabled() == false + set_fft_plan_cache!(true) + @test fft_plan_cache_enabled() == true + set_fft_plan_cache!(false) + @test fft_plan_cache_enabled() == false + + enable_fft_plan_cache!() + cfg_cache = create_gauss_config(4, 6) + field = randn(MersenneTwister(77), cfg_cache.nlat, cfg_cache.nlon) + ref = analysis(cfg_cache, field) + @test !isempty(SHTnsKit._LOCAL_FFT_PLAN_CACHE) + + disable_fft_plan_cache!() # clears by default + @test isempty(SHTnsKit._LOCAL_FFT_PLAN_CACHE) + @test analysis(cfg_cache, field) ≈ ref # same answer, unplanned + @test isempty(SHTnsKit._LOCAL_FFT_PLAN_CACHE) # and still bypassed + + enable_fft_plan_cache!() + @test analysis(cfg_cache, field) ≈ ref + @test !isempty(SHTnsKit._LOCAL_FFT_PLAN_CACHE) + + @test SHTnsKit.fft_plan_cache_max!(SHTnsKit.fft_plan_cache_max!(8)) == 8 + + set_fft_plan_cache!(initial_state) end @testset "Pencil grid suggestion" begin diff --git a/test/serial/test_energy_diagnostics.jl b/test/serial/test_energy_diagnostics.jl index 2c60a4a0..af2e5570 100644 --- a/test/serial/test_energy_diagnostics.jl +++ b/test/serial/test_energy_diagnostics.jl @@ -283,4 +283,17 @@ using SHTnsKit @test energy_vector(cfg, Slm, Tlm) >= 0 end end + + @testset "spectral energy rejects a mis-sized spectrum" begin + # The accumulation loops are `@inbounds`; a too-small matrix used to + # return a silently wrong number instead of raising. + cfg = create_gauss_config(6, 8) + small = zeros(ComplexF64, 2, 2) + ok = zeros(ComplexF64, cfg.lmax + 1, cfg.mmax + 1) + @test_throws DimensionMismatch energy_scalar(cfg, small) + @test_throws DimensionMismatch energy_vector(cfg, small, ok) + @test_throws DimensionMismatch energy_vector(cfg, ok, small) + @test energy_scalar(cfg, ok) == 0.0 + end + end diff --git a/test/serial/test_gradients.jl b/test/serial/test_gradients.jl index 354833b7..a43ee3de 100644 --- a/test/serial/test_gradients.jl +++ b/test/serial/test_gradients.jl @@ -402,4 +402,75 @@ end @test isapprox(dE_ad, dE_fd; rtol=5e-4, atol=1e-8) end + +end + + +if _HAS_CHAINRULES +@testset "QST rrules and implementation-detail kwargs" begin + # The QST family had no rrule at all, so differentiating a QST pipeline fell + # through to source tracing and crashed inside FFTW. And every rrule that + # declared fewer kwargs than its primal was skipped the moment a caller + # passed one, even at its default value. + cfg = create_gauss_config(5, 7; nlon=12) + rng = MersenneTwister(9412) + rnd() = begin + a = zeros(ComplexF64, cfg.lmax + 1, cfg.mmax + 1) + for m in 0:cfg.mmax, l in m:cfg.lmax + a[l+1, m+1] = m == 0 ? randn(rng) : complex(randn(rng), randn(rng)) + end + a + end + ε = 1e-6 + + @testset "synthesis_qst pullback vs finite differences" begin + Q, S, T = rnd(), rnd(), rnd() + Cr, Ct, Cp = randn(rng, cfg.nlat, cfg.nlon), randn(rng, cfg.nlat, cfg.nlon), randn(rng, cfg.nlat, cfg.nlon) + loss(q, s, t) = (V = synthesis_qst(cfg, q, s, t); + sum(Cr .* V[1]) + sum(Ct .* V[2]) + sum(Cp .* V[3])) + _, back = rrule(synthesis_qst, cfg, Q, S, T) + _, _, Q̄, S̄, T̄, _ = back((Cr, Ct, Cp)) + for (slot, ḡ, perturb) in ((:Q, Q̄, (h, q, s, t) -> (q .+ h, s, t)), + (:S, S̄, (h, q, s, t) -> (q, s .+ h, t)), + (:T, T̄, (h, q, s, t) -> (q, s, t .+ h))) + h = rnd() + fd = (loss(perturb(ε .* h, Q, S, T)...) - loss(perturb(-ε .* h, Q, S, T)...)) / (2ε) + @test real(sum(conj(ḡ) .* h)) ≈ fd rtol=1e-6 atol=1e-8 + end + end + + @testset "analysis_qst pullback vs finite differences" begin + Vr, Vt, Vp = randn(rng, cfg.nlat, cfg.nlon), randn(rng, cfg.nlat, cfg.nlon), randn(rng, cfg.nlat, cfg.nlon) + G1, G2, G3 = rnd(), rnd(), rnd() + loss(vr, vt, vp) = (A = analysis_qst(cfg, vr, vt, vp); + real(sum(conj(G1) .* A[1]) + sum(conj(G2) .* A[2]) + sum(conj(G3) .* A[3]))) + _, back = rrule(analysis_qst, cfg, Vr, Vt, Vp) + _, _, V̄r, V̄t, V̄p = back((G1, G2, G3)) + for (ḡ, perturb) in ((V̄r, (h, a, b, c) -> (a .+ h, b, c)), + (V̄t, (h, a, b, c) -> (a, b .+ h, c)), + (V̄p, (h, a, b, c) -> (a, b, c .+ h))) + h = randn(rng, cfg.nlat, cfg.nlon) + fd = (loss(perturb(ε .* h, Vr, Vt, Vp)...) - loss(perturb(-ε .* h, Vr, Vt, Vp)...)) / (2ε) + @test sum(real(ḡ) .* h) ≈ fd rtol=1e-6 atol=1e-8 + end + end + + @testset "rrules accept the primal's implementation-detail kwargs" begin + alm = rnd() + f = randn(rng, cfg.nlat, cfg.nlon) + Vt, Vp = randn(rng, cfg.nlat, cfg.nlon), randn(rng, cfg.nlat, cfg.nlon) + @test rrule(synthesis, cfg, alm; use_rfft=false) !== nothing + @test rrule(synthesis, cfg, alm; real_output=true, fft_scratch=nothing) !== nothing + @test rrule(analysis, cfg, f; use_rfft=false) !== nothing + @test rrule(analysis_sphtor, cfg, Vt, Vp; use_rfft=false) !== nothing + @test rrule(synthesis_sphtor, cfg, alm, alm; use_rfft=false) !== nothing + # and the rfft path gives the same adjoint as the complex one + ȳ = randn(rng, cfg.nlat, cfg.nlon) + _, b1 = rrule(synthesis, cfg, alm) + _, b2 = rrule(synthesis, cfg, alm; use_rfft=true) + @test b1(ȳ)[3] ≈ b2(ȳ)[3] rtol=1e-10 + end +end +else + @info "Skipping QST rrule tests (ChainRulesCore not available)" end diff --git a/test/serial/test_indexing.jl b/test/serial/test_indexing.jl index 5350283c..9fc918a0 100644 --- a/test/serial/test_indexing.jl +++ b/test/serial/test_indexing.jl @@ -210,4 +210,22 @@ using SHTnsKit cfg = create_gauss_config(lmax, lmax + 2) @test cfg.nlm == nlm_calc(lmax, lmax, mres) end + + @testset "im_from_lm is bounded by mmax, not lmax" begin + # The bound was `lmax ÷ mres`, so for an mmax < lmax layout a + # past-the-end index resolved to an order the configuration does not + # store instead of raising. + lmax, mmax = 8, 3 + n = nlm_calc(lmax, mmax, 1) + # every in-range index still maps to its own order + li, mi = build_li_mi(lmax, mmax, 1) + for k in 1:n + @test SHTnsKit.im_from_lm(k - 1, lmax, 1; mmax=mmax) == mi[k] + end + # one past the end now raises rather than inventing m = 4 + @test_throws ArgumentError SHTnsKit.im_from_lm(n, lmax, 1; mmax=mmax) + # default mmax=lmax keeps the old, correct behaviour for full layouts + @test SHTnsKit.im_from_lm(nlm_calc(lmax, lmax, 1) - 1, lmax, 1) == lmax + end + end diff --git a/test/serial/test_local.jl b/test/serial/test_local.jl index 9a8840b0..5c04477b 100644 --- a/test/serial/test_local.jl +++ b/test/serial/test_local.jl @@ -48,6 +48,41 @@ using Random end end + @testset "point/latitude evaluators honour phi_scale" begin + # `synthesis` scales its Fourier bins by `phi_inv_scale(cfg)` and the + # inverse FFT divides by nlon, a net factor of 1 under :dft but 1/2π + # under :quad. The direct evaluators applied no factor at all, so under + # :quad every one of them disagreed with the grid it claims to sample by + # exactly 2π. + for mode in (:dft, :quad) + cfg = create_gauss_config(6, 8) + cfg.phi_scale = mode + rng = MersenneTwister(5150) + A = zeros(ComplexF64, cfg.lmax + 1, cfg.mmax + 1) + for m in 0:cfg.mmax, l in m:cfg.lmax + A[l+1, m+1] = m == 0 ? randn(rng) : complex(randn(rng), randn(rng)) + end + packed = SHTnsKit.pack_lm(cfg, A) + zero_packed = zeros(ComplexF64, cfg.nlm) + + f = synthesis(cfg, A) + Vr, Vt, Vp = synthesis_qst(cfg, A, zero(A), zero(A)) + i, j = 2, 3 + + @test synthesis_point(cfg, A, cfg.x[i], cfg.φ[j]) ≈ f[i, j] rtol=1e-10 + @test SH_to_lat(cfg, packed, cfg.x[i]) ≈ f[i, :] rtol=1e-10 + @test SHqst_to_lat(cfg, packed, zero_packed, zero_packed, cfg.x[i])[1] ≈ + Vr[i, :] rtol=1e-10 + @test SHTnsKit.SHqst_to_point(cfg, packed, zero_packed, zero_packed, + cfg.x[i], cfg.φ[j])[1] ≈ Vr[i, j] rtol=1e-10 + + # axisymmetric pair + a0 = zeros(ComplexF64, cfg.lmax + 1); a0[3] = 0.7 + am = zeros(ComplexF64, cfg.lmax + 1, cfg.mmax + 1); am[3, 1] = 0.7 + @test synthesis_axisym(cfg, a0) ≈ synthesis(cfg, am)[:, 1] rtol=1e-10 + end + end + @testset "SH_to_lat matches synthesis at grid latitudes" begin lmax = 8 nlat = lmax + 2 diff --git a/test/serial/test_plan.jl b/test/serial/test_plan.jl index f3269036..06bca325 100644 --- a/test/serial/test_plan.jl +++ b/test/serial/test_plan.jl @@ -262,4 +262,51 @@ end @test all(isfinite, f_cplx) @test eltype(f_cplx) <: Complex end + + @testset "planned transforms match the cfg form in every buffer mode" begin + # The planned transforms exist to be a drop-in, faster replacement for + # the `cfg` forms. Pin that they agree numerically across the whole + # matrix of buffer modes: Legendre tables on/off, rfft on/off, and + # Robert form on/off — a divergence in any cell is a silent wrong + # answer for anyone who reaches for a plan. + rng = MersenneTwister(4821) + for tables in (false, true), use_rfft in (false, true), robert in (false, true) + cfgp = create_gauss_config(10, 12) + cfgp.robert_form = robert + tables ? SHTnsKit.prepare_plm_tables!(cfgp) : SHTnsKit.disable_plm_tables!(cfgp) + + alm = zeros(ComplexF64, cfgp.lmax + 1, cfgp.mmax + 1) + for m in 0:cfgp.mmax, l in m:cfgp.lmax + alm[l + 1, m + 1] = m == 0 ? randn(rng) : complex(randn(rng), randn(rng)) + end + S = copy(alm); T = 0.5 .* alm + S[1, 1] = 0; T[1, 1] = 0 + + f = synthesis(cfgp, alm) + Vt, Vp = synthesis_sphtor(cfgp, S, T) + alm_ref = analysis(cfgp, f) + S_ref, T_ref = analysis_sphtor(cfgp, Vt, Vp) + + plan = SHTPlan(cfgp; use_rfft=use_rfft) + tol = 1e-12 + + alm_out = similar(alm_ref) + analysis!(plan, alm_out, f) + @test maximum(abs, alm_out .- alm_ref) <= tol + + f_out = similar(f) + synthesis!(plan, f_out, alm) + @test maximum(abs, f_out .- f) <= tol + + S_out = similar(S_ref); T_out = similar(T_ref) + SHTnsKit.analysis_sphtor!(plan, S_out, T_out, Vt, Vp) + @test maximum(abs, S_out .- S_ref) <= tol + @test maximum(abs, T_out .- T_ref) <= tol + + Vt_out = similar(Vt); Vp_out = similar(Vp) + SHTnsKit.synthesis_sphtor!(plan, Vt_out, Vp_out, S, T) + @test maximum(abs, Vt_out .- Vt) <= tol + @test maximum(abs, Vp_out .- Vp) <= tol + end + end end diff --git a/test/serial/test_rotation_gradients.jl b/test/serial/test_rotation_gradients.jl index 1ac3d149..273c01c0 100644 --- a/test/serial/test_rotation_gradients.jl +++ b/test/serial/test_rotation_gradients.jl @@ -83,6 +83,102 @@ end end end +@testset "Angle gradients survive an in-place primal" begin + # `SH_Zrotate` is covered above. Its siblings capture the primal INPUT and + # read it lazily in the pullback, so the same hazard applies to them: an + # in-place call (Rlm === Qlm) overwrites those coefficients, and so does any + # caller that reuses the buffer before the pullback runs. Each rrule must + # snapshot what it needs at primal time. + rng = MersenneTwister(20931) + ε = 1e-6 + + @testset "SH_Yrotate dα" begin + cfg = create_gauss_config(4, 6; nlon=9) + Q = randn(rng, ComplexF64, cfg.nlm) + C = randn(rng, ComplexF64, cfg.nlm) + α = 0.7 + loss(q, a) = real(sum(conj(C) .* SH_Yrotate(cfg, copy(q), a, similar(q)))) + fd = (loss(Q, α + ε) - loss(Q, α - ε)) / (2ε) + + q = copy(Q) + _, back_sep = ChainRulesCore.rrule(SH_Yrotate, cfg, q, α, similar(q)) + @test back_sep(C)[4] ≈ fd rtol=1e-6 atol=1e-8 + + inplace = copy(Q) # Rlm === Qlm + _, back_ip = ChainRulesCore.rrule(SH_Yrotate, cfg, inplace, α, inplace) + @test back_ip(C)[4] ≈ fd rtol=1e-6 atol=1e-8 + + reused = copy(Q) # caller clobbers the input afterwards + _, back_reuse = ChainRulesCore.rrule(SH_Yrotate, cfg, reused, α, similar(reused)) + fill!(reused, 0) + @test back_reuse(C)[4] ≈ fd rtol=1e-6 atol=1e-8 + + if _HAS_ZYGOTE_ROT + zq = copy(Q) + _, zback = Zygote.pullback(SH_Yrotate, cfg, zq, α, zq) + @test zback(C)[3] ≈ fd rtol=1e-6 atol=1e-8 + end + end + + @testset "shtns_rotation_apply_cplx dβ" begin + lmax = mmax = 4 + mkrot(β) = (r = SHTnsKit.SHTRotation(lmax, mmax); + SHTnsKit.shtns_rotation_set_angles_ZYZ(r, 0.3, β, 0.2); r) + n = SHTnsKit.nlm_cplx_calc(lmax, mmax, 1) + Z = randn(rng, ComplexF64, n) + C = randn(rng, ComplexF64, n) + β = 0.7 + loss(r, z) = (R = similar(z); + SHTnsKit.shtns_rotation_apply_cplx(r, copy(z), R); + real(sum(conj(C) .* R))) + fd = (loss(mkrot(β + ε), Z) - loss(mkrot(β - ε), Z)) / (2ε) + + z = copy(Z) + _, back_sep = ChainRulesCore.rrule(SHTnsKit.shtns_rotation_apply_cplx, mkrot(β), z, similar(z)) + @test back_sep(C)[2].β ≈ fd rtol=1e-6 atol=1e-8 + + zi = copy(Z) # Rlm === Zlm + _, back_ip = ChainRulesCore.rrule(SHTnsKit.shtns_rotation_apply_cplx, mkrot(β), zi, zi) + @test back_ip(C)[2].β ≈ fd rtol=1e-6 atol=1e-8 + end + + @testset "shtns_rotation_apply_real dβ" begin + cfg = create_gauss_config(4, 7; nlon=11) + mkrot(β) = (r = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax); + SHTnsKit.shtns_rotation_set_angles_ZYZ(r, 0.3, β, 0.2); r) + Q = randn(rng, ComplexF64, cfg.nlm) + C = randn(rng, ComplexF64, cfg.nlm) + β = 0.7 + loss(r, q) = (R = similar(q); + SHTnsKit.shtns_rotation_apply_real(r, copy(q), R); + real(sum(conj(C) .* R))) + fd = (loss(mkrot(β + ε), Q) - loss(mkrot(β - ε), Q)) / (2ε) + + q = copy(Q) + _, back_sep = ChainRulesCore.rrule(SHTnsKit.shtns_rotation_apply_real, mkrot(β), q, similar(q)) + @test back_sep(C)[2].β ≈ fd rtol=1e-6 atol=1e-8 + + qi = copy(Q) # Rlm === Qlm + _, back_ip = ChainRulesCore.rrule(SHTnsKit.shtns_rotation_apply_real, mkrot(β), qi, qi) + @test back_ip(C)[2].β ≈ fd rtol=1e-6 atol=1e-8 + end +end + +@testset "Order-mixing rotations reject mres > 1 with a usable message" begin + cfg = create_gauss_config(4, 6; mres=2) + Q = randn(MersenneTwister(5), ComplexF64, cfg.nlm) + err = try + SH_Yrotate(cfg, Q, 0.3, similar(Q)); nothing + catch e + e + end + # Either guard is fine — SH_Yrotate's own mres check or the packed-length + # check inside the Wigner engine — as long as the message names `mres`, so + # the caller is not left reverse-engineering a bare size mismatch. + @test err isa Union{ArgumentError,DimensionMismatch} + @test occursin("mres", sprint(showerror, err)) +end + @testset "Complex-packed analysis rrule respects configured convention" begin lmax = 4 cfg = create_gauss_config(lmax, 7; nlon=11, norm=:schmidt, diff --git a/test/serial/test_rotations.jl b/test/serial/test_rotations.jl index 4387d0d6..4cdc9b76 100644 --- a/test/serial/test_rotations.jl +++ b/test/serial/test_rotations.jl @@ -33,6 +33,56 @@ using SHTnsKit end end + @testset "Z-rotation sign convention is pinned to physical ground truth" begin + # The `exp(-imα)` sign is not a free choice: flipping it to `exp(+imα)` + # (the passive convention) silently desynchronises SH_Zrotate from the + # actual spatial rotation, from the general Wigner engine, and from the + # distributed twins. Pin all three so it cannot drift again. + lmax = 6 + nlon = 2 * lmax + 2 # even, so an exact grid shift exists + cfg = create_gauss_config(lmax, lmax + 2; nlon=nlon) + rng = MersenneTwister(4242) + + A = zeros(ComplexF64, lmax + 1, cfg.mmax + 1) + for m in 0:cfg.mmax, l in m:lmax + A[l + 1, m + 1] = m == 0 ? randn(rng) : complex(randn(rng), randn(rng)) + end + Qlm = SHTnsKit.pack_lm(cfg, A) + + shift = 3 # α is an exact multiple of the φ spacing + α = 2π * shift / nlon + Rlm = similar(Qlm) + SH_Zrotate(cfg, Qlm, α, Rlm) + + # (1) Ground truth: rotating the FIELD by +α about ẑ is g(θ,φ) = f(θ,φ-α), + # which on this grid is a pure column shift. + f = synthesis(cfg, A) + g = similar(f) + for j in 1:nlon + g[:, j] .= @view f[:, mod(j - 1 - shift, nlon) + 1] + end + spatial = SHTnsKit.pack_lm(cfg, analysis(cfg, g)) + @test isapprox(Rlm, spatial; rtol=1e-10, atol=1e-12) + + # ...and the opposite sign is the OTHER rotation, not a rounding detail. + flipped = [Qlm[k] * cis(+cfg.mi[k] * α) for k in 1:cfg.nlm] + @test !isapprox(flipped, spatial; rtol=1e-3, atol=1e-6) + + # (2) The general ZYZ engine must agree for both the α and γ slots. + for angles in ((α, 0.0, 0.0), (0.0, 0.0, α)) + r = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax) + SHTnsKit.shtns_rotation_set_angles_ZYZ(r, angles...) + Rzyz = similar(Qlm) + SHTnsKit.shtns_rotation_apply_real(r, Qlm, Rzyz) + @test isapprox(Rzyz, Rlm; rtol=1e-10, atol=1e-12) + end + + # (3) The dense distributed twin must agree coefficient for coefficient. + Rdense = zeros(ComplexF64, lmax + 1, cfg.mmax + 1) + SHTnsKit.dist_SH_Zrotate(cfg, A, α, Rdense) + @test isapprox(SHTnsKit.pack_lm(cfg, Rdense), Rlm; rtol=1e-12, atol=1e-14) + end + @testset "Z-axis rotation in-place" begin lmax = 6 cfg = create_gauss_config(lmax, lmax + 2; nlon=2*lmax + 1) @@ -128,6 +178,104 @@ using SHTnsKit @test isapprox(Rlm4, Qlm; rtol=1e-8, atol=1e-10) end + @testset "X-rotation 90 is Rx(+pi/2), not its inverse" begin + # `SH_Xrotate90` must be the +90° turn its name promises. The existing + # "four 90° rotations = identity" check passes for either sign, so pin + # the direction explicitly against the ZXZ entry point. + lmax = 5 + cfg = create_gauss_config(lmax, lmax + 2; nlon=2*lmax + 2) + rng = MersenneTwister(8801) + Qlm = randn(rng, ComplexF64, cfg.nlm) + Qlm[1:lmax+1] .= real.(Qlm[1:lmax+1]) + + X = similar(Qlm) + SH_Xrotate90(cfg, Qlm, X) + + # Ground truth: a pure Rx(+π/2) via the ZXZ entry point, whose own + # converter documents Rx(β) = Rz(-π/2)·Ry(β)·Rz(π/2). + fwd = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax) + SHTnsKit.shtns_rotation_set_angles_ZXZ(fwd, 0.0, π/2, 0.0) + Xref = similar(Qlm) + SHTnsKit.shtns_rotation_apply_real(fwd, Qlm, Xref) + @test isapprox(X, Xref; rtol=1e-10, atol=1e-12) + + # ...and it is NOT the inverse rotation. + inv_rot = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax) + SHTnsKit.shtns_rotation_set_angles_ZXZ(inv_rot, 0.0, -π/2, 0.0) + Xinv = similar(Qlm) + SHTnsKit.shtns_rotation_apply_real(inv_rot, Qlm, Xinv) + @test !isapprox(X, Xinv; rtol=1e-3, atol=1e-6) + + # Three forward turns equal one backward turn. + a = similar(Qlm); b = similar(Qlm); c = similar(Qlm) + SH_Xrotate90(cfg, Qlm, a); SH_Xrotate90(cfg, a, b); SH_Xrotate90(cfg, b, c) + @test isapprox(c, Xinv; rtol=1e-9, atol=1e-11) + end + + @testset "angle-axis half-turns pick the right axis" begin + # At β≈π the observable Euler combination is α-γ read off the NEGATED + # first column; reusing the β≈0 formula there returns a rotation about + # the wrong axis (a 180° turn about x̂ came back as Ry(π)). Only exact + # half-turns are affected, so pin them. + cfg = create_gauss_config(5, 7; nlon=12) + rng = MersenneTwister(6041) + Q = randn(rng, ComplexF64, cfg.nlm) + Q[1:cfg.lmax+1] .= real.(Q[1:cfg.lmax+1]) + apply(r) = (R = similar(Q); SHTnsKit.shtns_rotation_apply_real(r, Q, R); R) + byaxis(θ, ax...) = (r = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax); + SHTnsKit.shtns_rotation_set_angle_axis(r, θ, ax...); apply(r)) + byzxz(β) = (r = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax); + SHTnsKit.shtns_rotation_set_angles_ZXZ(r, 0.0, β, 0.0); apply(r)) + byzyz(β) = (r = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax); + SHTnsKit.shtns_rotation_set_angles_ZYZ(r, 0.0, β, 0.0); apply(r)) + + @test isapprox(byaxis(π, 1.0, 0.0, 0.0), byzxz(π); rtol=1e-10, atol=1e-12) # x̂ half-turn + @test isapprox(byaxis(π, 0.0, 1.0, 0.0), byzyz(π); rtol=1e-10, atol=1e-12) # ŷ half-turn + @test !isapprox(byaxis(π, 1.0, 0.0, 0.0), byzyz(π); rtol=1e-3, atol=1e-6) # and they differ + @test isapprox(byaxis(π/2, 1.0, 0.0, 0.0), byzxz(π/2); rtol=1e-10, atol=1e-12) + @test isapprox(byaxis(0.0, 1.0, 0.0, 0.0), Q; rtol=1e-10, atol=1e-12) + + # A half-turn about any axis must square to the identity. + for ax in ((1.0,0.0,0.0), (0.0,1.0,0.0), (0.0,0.0,1.0), (1.0,1.0,0.0), (1.0,1.0,1.0)) + r = SHTnsKit.SHTRotation(cfg.lmax, cfg.mmax) + SHTnsKit.shtns_rotation_set_angle_axis(r, π, ax...) + once = similar(Q); SHTnsKit.shtns_rotation_apply_real(r, Q, once) + twice = similar(Q); SHTnsKit.shtns_rotation_apply_real(r, once, twice) + @test isapprox(twice, Q; rtol=1e-9, atol=1e-11) + end + end + + @testset "order-mixing rotations reject mmax < lmax" begin + # A Wigner-d rotation through a general β couples every |m'| ≤ l. With + # mmax < lmax the |m'| > mmax components were silently dropped — 14.8 % + # of the field energy at lmax=8/mmax=5, 24.0 % at mmax=3. + full = create_gauss_config(8, 10; mmax=8, nlon=18) + cut = create_gauss_config(8, 10; mmax=5, nlon=18) + rng = MersenneTwister(6042) + qf = randn(rng, ComplexF64, full.nlm); qf[1:full.lmax+1] .= real.(qf[1:full.lmax+1]) + qc = randn(rng, ComplexF64, cut.nlm); qc[1:cut.lmax+1] .= real.(qc[1:cut.lmax+1]) + + # mmax == lmax is unaffected and conserves energy. + Rf = similar(qf); SH_Yrotate(full, qf, 0.7, Rf) + w = Float64[full.mi[k] == 0 ? 1.0 : 2.0 for k in 1:full.nlm] + @test sum(w .* abs2.(Rf)) ≈ sum(w .* abs2.(qf)) rtol=1e-10 + + # mmax < lmax with an order-mixing β must raise, not truncate. + @test_throws ArgumentError SH_Yrotate(cut, qc, 0.7, similar(qc)) + @test_throws ArgumentError SH_Xrotate90(cut, qc, similar(qc)) + + # ...but the two non-mixing angles still work at reduced mmax: + # β ≡ 0 is diagonal (a pure Z-rotation) and β ≡ π is anti-diagonal. + for β in (0.0, π) + r = SHTnsKit.SHTRotation(cut.lmax, cut.mmax) + SHTnsKit.shtns_rotation_set_angles_ZYZ(r, 0.4, β, 0.0) + R = similar(qc) + @test (SHTnsKit.shtns_rotation_apply_real(r, qc, R); true) + wc = Float64[cut.mi[k] == 0 ? 1.0 : 2.0 for k in 1:cut.nlm] + @test sum(wc .* abs2.(R)) ≈ sum(wc .* abs2.(qc)) rtol=1e-10 + end + end + @testset "Wigner-d matrix orthogonality" begin # Test orthogonality: d^T d = I for real orthogonal case for l in 0:4 diff --git a/test/serial/test_turbo.jl b/test/serial/test_turbo.jl index 4888db71..cd0afbfb 100644 --- a/test/serial/test_turbo.jl +++ b/test/serial/test_turbo.jl @@ -113,5 +113,54 @@ else @test haskey(b, :synthesis_turbo) @test b.analysis_turbo > 0 end + + @testset "turbo honours cfg.mres" begin + # Regression: both turbo loops iterated a bare `0:mmax`, so for an + # mres>1 config `analysis_turbo` populated coefficient columns the + # transform has no storage for and `synthesis_turbo` summed them + # back in. The disagreement with `analysis`/`synthesis` was O(1), + # not roundoff. + for mres in (2, 3) + cfgm = create_gauss_config(6, 8; mres=mres, nlon=15) + rngm = MersenneTwister(9100 + mres) + fm = randn(rngm, cfgm.nlat, cfgm.nlon) + + a_ref = analysis(cfgm, fm) + a_turbo = SHTnsKit.analysis_turbo(cfgm, fm) + @test a_turbo ≈ a_ref atol=1e-12 + + # Orders absent from the layout must stay exactly zero. + for m in 0:cfgm.mmax + m % mres == 0 && continue + @test all(iszero, @view a_turbo[:, m + 1]) + end + + # Feed a spectrum that DOES carry junk on the skipped columns: + # synthesis must ignore it, exactly as the core does. + dirty = copy(a_ref) + for m in 0:cfgm.mmax + m % mres == 0 && continue + for l in m:cfgm.lmax + dirty[l + 1, m + 1] = 3.0 - 2.0im + end + end + @test SHTnsKit.synthesis_turbo(cfgm, dirty) ≈ synthesis(cfgm, dirty) atol=1e-12 + @test SHTnsKit.synthesis_turbo(cfgm, dirty) ≈ synthesis(cfgm, a_ref) atol=1e-12 + end + end + + @testset "turbo runs inside an outer threaded region" begin + # `@threads :static` cannot be nested; the turbo loops must fall back + # to the caller's thread the way the core orchestrators do. + cfgt = create_gauss_config(5, 7; nlon=11) + ft = randn(MersenneTwister(31), cfgt.nlat, cfgt.nlon) + ref = analysis(cfgt, ft) + out = Vector{Any}(undef, 2) + Threads.@threads for k in 1:2 + out[k] = SHTnsKit.analysis_turbo(cfgt, ft) + end + @test out[1] ≈ ref atol=1e-12 + @test out[2] ≈ ref atol=1e-12 + end end end diff --git a/test/serial/test_vector_transforms.jl b/test/serial/test_vector_transforms.jl index 989e0107..04686c45 100644 --- a/test/serial/test_vector_transforms.jl +++ b/test/serial/test_vector_transforms.jl @@ -8,6 +8,17 @@ using SHTnsKit @isdefined(VERBOSE) || (const VERBOSE = get(ENV, "SHTNSKIT_TEST_VERBOSE", "0") == "1") +# Allocation budgets below are calibrated for a SINGLE-THREADED run, where they +# are exact. The shared m-loop orchestrators start `@threads` tasks whenever +# threads are available, and each threaded region costs a few hundred bytes to +# spawn — a constant independent of problem size, and not the kind of regression +# these budgets exist to catch (a per-row temporary or a dense zero spectrum +# costs tens of KB and grows with lmax). Allow a per-thread slack rather than +# letting the whole suite go red on any multi-core machine; note this makes the +# budgets coarse at `nthreads > 1`, so the tight check is the 1-thread run. +@isdefined(_thread_alloc_slack) || + (_thread_alloc_slack() = Threads.nthreads() > 1 ? 4_096 * Threads.nthreads() : 0) + @testset "Vector Transforms (Spheroidal-Toroidal)" begin @testset "Sphtor roundtrip" begin lmax = 8 @@ -87,8 +98,8 @@ using SHTnsKit synthesis_tor(cfg, Tlm) GC.gc() - @test @allocated(synthesis_sph(cfg, Slm)) <= 15_000 - @test @allocated(synthesis_tor(cfg, Tlm)) <= 15_000 + @test @allocated(synthesis_sph(cfg, Slm)) <= 15_000 + _thread_alloc_slack() + @test @allocated(synthesis_tor(cfg, Tlm)) <= 15_000 + _thread_alloc_slack() end @testset "Gradient transform" begin @@ -453,4 +464,20 @@ using SHTnsKit @test isapprox(Gp_l, Gp_ref; rtol=1e-10, atol=1e-12) end + + @testset "mode-limited sphtor validates its order and truncation" begin + # Out-of-range arguments used to walk off the end of the norm-scale table + # and return ±Inf coefficients; the scalar twin already raised. + cfg = create_gauss_config(3, 5; nlon=8) + v = randn(MersenneTwister(3311), ComplexF64, cfg.nlat) + sl = zeros(ComplexF64, cfg.lmax + 1) + @test_throws ArgumentError SHTnsKit.analysis_sphtor_ml(cfg, 9, v, v, 12) + @test_throws ArgumentError SHTnsKit.analysis_sphtor_ml(cfg, -1, v, v, 3) + @test_throws ArgumentError SHTnsKit.analysis_sphtor_ml(cfg, 1, v, v, 99) + @test_throws ArgumentError SHTnsKit.synthesis_sphtor_ml(cfg, 9, sl, sl, 12) + # a valid call still returns finite coefficients + S, T = SHTnsKit.analysis_sphtor_ml(cfg, 1, v, v, 3) + @test all(isfinite, S) && all(isfinite, T) + end + end From b31a8e10ac91d0374b508dd1ffcc6ec427a550aa Mon Sep 17 00:00:00 2001 From: Subhajit Kar Date: Sun, 13 Sep 2026 07:11:05 +0300 Subject: [PATCH 2/3] test: give the two remaining thread-sensitive allocation budgets the same slack MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `test_qst_transforms.jl`'s `synthesis_qst_l` ceiling and `test_plan.jl`'s `cfg`-form `synthesis!` ceiling are calibrated for a single-threaded run. The shared m-loop orchestrators start `@threads` tasks whenever threads are available, and each threaded region costs a few hundred bytes to spawn — a constant independent of problem size, not the field-sized regression these budgets exist to catch. Both went red at `-t 4` (13392 > 13000 and 2464 > 128); the `test_plan.jl` one is pre-existing on main. Uses the same `_thread_alloc_slack()` helper the four other budget files already carry. The planned `synthesis!` / `synthesis_sphtor!` ceilings are left strict: those own their scratch and spawn no tasks, so they stay at 0 regardless of thread count. Serial suite: 76,848 pass at one thread and at four. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Dbj4TtuW11gke9JwSNs7Nz --- test/fixtures/compatibility/task16_gate.toml | 6 +++--- .../compatibility/task16_local_commands.txt | 2 +- .../compatibility/task16_local_summary.log | 4 +++- test/serial/test_plan.jl | 16 +++++++++++++++- test/serial/test_qst_transforms.jl | 13 ++++++++++++- 5 files changed, 34 insertions(+), 7 deletions(-) diff --git a/test/fixtures/compatibility/task16_gate.toml b/test/fixtures/compatibility/task16_gate.toml index c494a1f8..77353b2c 100644 --- a/test/fixtures/compatibility/task16_gate.toml +++ b/test/fixtures/compatibility/task16_gate.toml @@ -6,13 +6,13 @@ platform = "aarch64-apple-darwin" [local_gate] evidence_kind = "local_readiness" certifying = false -audited_tree_digest = "2cb7eee1644d1df5613a0b39e188ee634a23aceafef5f49b418cde59717e3643" +audited_tree_digest = "763b5c3413870230bd8b906ced62324a36885f9d026abbdf0ef6788dbecd2793" audited_scope = ["Project.toml", ".github/workflows", "docs/src/shtns37-parity.md", "ext", "src", "test"] audited_exclusions = ["test/fixtures/compatibility/task16_gate.toml", "test/fixtures/compatibility/task16_local_commands.txt", "test/fixtures/compatibility/task16_local_summary.log"] commands_file = "test/fixtures/compatibility/task16_local_commands.txt" -commands_sha256 = "9f61a016bc43244a56c24ef18d2a1c0a9aefd2be09d24745efd55d82278b7ea8" +commands_sha256 = "a973719289891efa284d3627e21f8881edc7763c40523205e759b9bded6bbdfc" summary_log = "test/fixtures/compatibility/task16_local_summary.log" -summary_log_sha256 = "63082905f4212b90b67c09834f35473c19e8dac945e42e9d48efdd494b50b251" +summary_log_sha256 = "1b4df56fbe299b2d4f51a55bafdac0048efade1a8aa4cc56420fe01f097854c8" # Historical readiness counts; current review runs are recorded in the summary. pkg_exit_code = 0 serial_pass = 76347 diff --git a/test/fixtures/compatibility/task16_local_commands.txt b/test/fixtures/compatibility/task16_local_commands.txt index 9a4c22f5..dd0876b2 100644 --- a/test/fixtures/compatibility/task16_local_commands.txt +++ b/test/fixtures/compatibility/task16_local_commands.txt @@ -5,7 +5,7 @@ certifying=false manual_cuda_tool_command=python3 -B -m unittest discover -s test/support -p test_manual_cuda.py -v manual_cuda_contract_command=/Users/subha/.julia/juliaup/julia-1.12.4+0.aarch64.apple.darwin14/bin/julia --startup-file=no --project=. test/serial/test_shtns37_contract.jl manual_cuda_numerical_source_unchanged_from=3b8b2fe118900f2431208088932631db4a162908 -audited_tree_digest=2cb7eee1644d1df5613a0b39e188ee634a23aceafef5f49b418cde59717e3643 +audited_tree_digest=763b5c3413870230bd8b906ced62324a36885f9d026abbdf0ef6788dbecd2793 audited_scope=Project.toml,.github/workflows,docs/src/shtns37-parity.md,ext,src,test audited_exclusions=test/fixtures/compatibility/task16_gate.toml,test/fixtures/compatibility/task16_local_commands.txt,test/fixtures/compatibility/task16_local_summary.log pkg_command=/usr/bin/env JULIA_DEPOT_PATH=/private/tmp/julia_depot_shtnskit:/Users/subha/.julia SHTNSKIT_RUN_JET_TESTS=1 SHTNSKIT_RUN_AQUA_TESTS=1 /Users/subha/.julia/juliaup/julia-1.12.4+0.aarch64.apple.darwin14/bin/julia --startup-file=no --project=. -e 'using Pkg; Pkg.test()' diff --git a/test/fixtures/compatibility/task16_local_summary.log b/test/fixtures/compatibility/task16_local_summary.log index 44c8d5c8..e59f6c21 100644 --- a/test/fixtures/compatibility/task16_local_summary.log +++ b/test/fixtures/compatibility/task16_local_summary.log @@ -9,7 +9,7 @@ manual_cuda_no_hardware_guard_rejected=true manual_cuda_julia_syntax_pass=3 manual_cuda_github_status_published=false manual_cuda_numerical_source_unchanged_from=3b8b2fe118900f2431208088932631db4a162908 -audited_tree_digest=2cb7eee1644d1df5613a0b39e188ee634a23aceafef5f49b418cde59717e3643 +audited_tree_digest=763b5c3413870230bd8b906ced62324a36885f9d026abbdf0ef6788dbecd2793 platform=aarch64-apple-darwin julia=1.12.4 pkg_exit_code=0 @@ -115,6 +115,8 @@ audit2_base_commit=f06d13d1 audit2_julia=1.11.1 audit2_platform=arm64-apple-darwin22.4.0 audit2_threads=1 +audit2_serial_pass_threads_4=76848 +audit2_serial_threads_4_exit_code=0 audit2_serial_pass=76848 audit2_serial_fail=0 audit2_serial_exit_code=0 diff --git a/test/serial/test_plan.jl b/test/serial/test_plan.jl index 06bca325..62c237e4 100644 --- a/test/serial/test_plan.jl +++ b/test/serial/test_plan.jl @@ -8,6 +8,15 @@ using SHTnsKit @isdefined(VERBOSE) || (const VERBOSE = get(ENV, "SHTNSKIT_TEST_VERBOSE", "0") == "1") +# Allocation budgets below are calibrated for a SINGLE-THREADED run, where they +# are exact. The shared m-loop orchestrators start `@threads` tasks whenever +# threads are available, and each threaded region costs a few hundred bytes to +# spawn — a constant independent of problem size, and not the kind of regression +# these budgets exist to catch. Allow a per-thread slack rather than letting the +# suite go red on any multi-core machine; the tight check is the 1-thread run. +@isdefined(_thread_alloc_slack) || + (_thread_alloc_slack() = Threads.nthreads() > 1 ? 4_096 * Threads.nthreads() : 0) + function _rand_real_alm(rng, lmax, mmax) alm = randn(rng, ComplexF64, lmax + 1, mmax + 1) alm[:, 1] .= real.(alm[:, 1]) @@ -161,10 +170,15 @@ end # runners. Keep the strict production-path ceiling while still ruling # out any field-sized allocation in coverage-enabled CI. allocation_limit = Base.JLOptions().code_coverage == 0 ? 128 : 2048 + # The planned forms own their scratch and never spawn tasks, so they stay + # at the strict ceiling regardless of thread count. The `cfg` form routes + # through the threaded m-loop orchestrator, so it pays the task-spawn + # constant; that is not a field-sized allocation and not what this guards. @test @allocated(synthesis!(plan, f, alm)) <= allocation_limit @test @allocated(synthesis!(plan_r, f, alm)) <= allocation_limit @test @allocated(synthesis_sphtor!(plan, Vt, Vp, Slm, Tlm)) <= allocation_limit - @test @allocated(synthesis!(cfg, f, alm; fft_scratch)) <= allocation_limit + @test @allocated(synthesis!(cfg, f, alm; fft_scratch)) <= + allocation_limit + _thread_alloc_slack() end @testset "Planned scalar matches the non-planned path exactly" begin diff --git a/test/serial/test_qst_transforms.jl b/test/serial/test_qst_transforms.jl index 0e0e9a4e..e2ae8f1f 100644 --- a/test/serial/test_qst_transforms.jl +++ b/test/serial/test_qst_transforms.jl @@ -7,6 +7,17 @@ using SHTnsKit @isdefined(VERBOSE) || (const VERBOSE = get(ENV, "SHTNSKIT_TEST_VERBOSE", "0") == "1") +# Allocation budgets below are calibrated for a SINGLE-THREADED run, where they +# are exact. The shared m-loop orchestrators start `@threads` tasks whenever +# threads are available, and each threaded region costs a few hundred bytes to +# spawn — a constant independent of problem size, and not the kind of regression +# these budgets exist to catch (a per-row temporary or a dense zero spectrum +# costs tens of KB and grows with lmax). Allow a per-thread slack rather than +# letting the whole suite go red on any multi-core machine; note this makes the +# budgets coarse at `nthreads > 1`, so the tight check is the 1-thread run. +@isdefined(_thread_alloc_slack) || + (_thread_alloc_slack() = Threads.nthreads() > 1 ? 4_096 * Threads.nthreads() : 0) + @testset "QST (3D Vector) Transforms" begin @testset "QST roundtrip" begin lmax = 6 @@ -90,7 +101,7 @@ using SHTnsKit # Array header/alignment overhead differs across Julia patch versions # and platforms (notably Windows CI). Keep this below an extra # component-sized scratch allocation while allowing that fixed floor. - qst_l_alloc_budget = 13_000 + qst_l_alloc_budget = 13_000 + _thread_alloc_slack() @test @allocated(synthesis_qst_l(cfg, Qlm, Slm, Tlm, ltr; real_output=true)) <= qst_l_alloc_budget end From 5ba07e5ee2b1b52c962cc8fa33928c4dec137081 Mon Sep 17 00:00:00 2001 From: Subhajit Kar Date: Sun, 13 Sep 2026 07:18:48 +0300 Subject: [PATCH 3/3] fix: keep the evaluator phi_scale factor from widening Float32 results MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_evaluator_phi_scale` returns a `Float64`, so multiplying a `Float32` point or latitude evaluation by it promoted every result to `Float64` — breaking the element type these functions promise their caller, and the `Dual` types AD needs to see through. `test/parity/runtests_cpu.jl` caught it: 64 failures on `@test all(value -> value isa T, qst_host)` in the local-evaluation parity suite, all element-type assertions, no numeric drift. Adds `_evaluator_phi_scale(cfg, ::Type{T})`, which narrows the scale to the evaluator's own real type, and uses it at every site: `SHqst_to_point`, `SH_to_grad_point`, `SH_to_lat`, `SH_to_lat_cplx`, `SHqst_to_lat`, `synthesis_point` and `synthesis_point_cplx`. The axisym pair and the PencilArray locals already converted explicitly and are unchanged. Pins the contract with a new testset that asserts the returned element type for every evaluator across `Float32`/`Float64` and both `phi_scale` modes, so the serial suite catches this without the parity runner. test/parity/runtests_cpu.jl now exits 0 (was 2148 passed / 64 failed; stock origin/main is 2212/2212, so this was a regression introduced by the evaluator phi_scale change in d9950cf7, not a pre-existing failure). Serial suite 76,884 pass at one thread and at four. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Dbj4TtuW11gke9JwSNs7Nz --- src/SHTnsKit.jl | 13 +++++++++ src/complex_packed.jl | 2 +- src/local.jl | 16 +++++++---- src/transforms.jl | 2 +- test/fixtures/compatibility/task16_gate.toml | 6 ++-- .../compatibility/task16_local_commands.txt | 2 +- .../compatibility/task16_local_summary.log | 7 +++-- test/serial/test_local.jl | 28 +++++++++++++++++++ 8 files changed, 63 insertions(+), 13 deletions(-) diff --git a/src/SHTnsKit.jl b/src/SHTnsKit.jl index 77b40794..5ca9ce1e 100644 --- a/src/SHTnsKit.jl +++ b/src/SHTnsKit.jl @@ -187,6 +187,19 @@ the very grid it claims to sample by a factor of 2π. """ @inline _evaluator_phi_scale(cfg::SHTConfig) = phi_inv_scale(cfg) / cfg.nlon +""" + _evaluator_phi_scale(cfg, ::Type{T}) -> T + +`_evaluator_phi_scale` narrowed to the evaluator's own real type. + +The evaluators promise their caller the element type their coefficients carry — +`Float32` coefficients give `Float32` values, and a `Dual` stays a `Dual`. The +untyped scale is a `Float64`, so multiplying by it silently widens every +`Float32` result to `Float64`. Convert once, at the boundary. +""" +@inline _evaluator_phi_scale(cfg::SHTConfig, ::Type{T}) where {T} = + convert(real(T), _evaluator_phi_scale(cfg)) + include("buffer_utils.jl") # Common buffer allocation patterns include("kernels.jl") # Legendre accumulation kernels include("plan.jl") # Transform planning and optimization diff --git a/src/complex_packed.jl b/src/complex_packed.jl index 16d57246..6fe1be9f 100644 --- a/src/complex_packed.jl +++ b/src/complex_packed.jl @@ -287,7 +287,7 @@ function synthesis_point_cplx(cfg::SHTConfig, alm::AbstractVector{<:Complex}, co end # Same φ convention factor `synthesis_packed_cplx` carries (1 under :dft; # 1/2π under :quad), so a point evaluation matches the grid it samples. - sφ = _evaluator_phi_scale(cfg) + sφ = _evaluator_phi_scale(cfg, typeof(acc)) return sφ == 1 ? acc : acc * sφ end diff --git a/src/local.jl b/src/local.jl index 3a2f0707..5e9ae307 100644 --- a/src/local.jl +++ b/src/local.jl @@ -68,7 +68,7 @@ function SH_to_lat(cfg::SHTConfig, Qlm::AbstractVector{<:Complex}, cost::Real; n vals[j+1] += 2 * real(gm * cis(PT(2π * m * j / nphi))) end end - sφ = _evaluator_phi_scale(cfg) # 1 under :dft; 1/2π under :quad + sφ = _evaluator_phi_scale(cfg, eltype(vals)) # 1 under :dft; 1/2π under :quad sφ == 1 || (vals .*= sφ) return vals end @@ -128,7 +128,7 @@ function SH_to_lat_cplx(cfg::SHTConfig, alm_packed::AbstractVector{<:Complex}, c vals[j+1] += gm * phase + gn * conj(phase) end end - sφ = _evaluator_phi_scale(cfg) # 1 under :dft; 1/2π under :quad + sφ = _evaluator_phi_scale(cfg, eltype(vals)) # 1 under :dft; 1/2π under :quad sφ == 1 || (vals .*= sφ) return vals end @@ -207,7 +207,10 @@ function SHqst_to_point(cfg::SHTConfig, Qlm::AbstractVector{<:Complex}, Slm::Abs vt *= sinth vp *= sinth end - sφ = _evaluator_phi_scale(cfg) # 1 under :dft; 1/2π under :quad + # Narrow the scale to the accumulator's own real type first: the untyped + # scale is Float64, and multiplying by it would widen a Float32 evaluation + # (and break the element type these functions promise their caller). + sφ = _evaluator_phi_scale(cfg, typeof(real(vr))) return real(vr) * sφ, real(vt) * sφ, real(vp) * sφ end @@ -278,7 +281,10 @@ function SH_to_grad_point(cfg::SHTConfig, DrSlm::AbstractVector{<:Complex}, Slm: vt *= sinth vp *= sinth end - sφ = _evaluator_phi_scale(cfg) # 1 under :dft; 1/2π under :quad + # Narrow the scale to the accumulator's own real type first: the untyped + # scale is Float64, and multiplying by it would widen a Float32 evaluation + # (and break the element type these functions promise their caller). + sφ = _evaluator_phi_scale(cfg, typeof(real(vr))) return real(vr) * sφ, real(vt) * sφ, real(vp) * sφ end @@ -373,7 +379,7 @@ function SHqst_to_lat(cfg::SHTConfig, Qlm::AbstractVector{<:Complex}, Slm::Abstr Vt .*= sinth Vp .*= sinth end - sφ = _evaluator_phi_scale(cfg) # 1 under :dft; 1/2π under :quad + sφ = _evaluator_phi_scale(cfg, eltype(Vr)) # 1 under :dft; 1/2π under :quad if sφ != 1 Vr .*= sφ; Vt .*= sφ; Vp .*= sφ end diff --git a/src/transforms.jl b/src/transforms.jl index de599702..08d71d71 100644 --- a/src/transforms.jl +++ b/src/transforms.jl @@ -548,7 +548,7 @@ function synthesis_point(cfg::SHTConfig, Qlm::AbstractMatrix{<:Complex}, cost::R # Same φ convention factor the grid `synthesis` carries (1 under :dft; # 1/2π under :quad), so a point evaluation matches the grid it samples. - sφ = _evaluator_phi_scale(cfg) + sφ = _evaluator_phi_scale(cfg, typeof(result)) return sφ == 1 ? result : result * sφ end diff --git a/test/fixtures/compatibility/task16_gate.toml b/test/fixtures/compatibility/task16_gate.toml index 77353b2c..b36bfb0e 100644 --- a/test/fixtures/compatibility/task16_gate.toml +++ b/test/fixtures/compatibility/task16_gate.toml @@ -6,13 +6,13 @@ platform = "aarch64-apple-darwin" [local_gate] evidence_kind = "local_readiness" certifying = false -audited_tree_digest = "763b5c3413870230bd8b906ced62324a36885f9d026abbdf0ef6788dbecd2793" +audited_tree_digest = "34289cc58102e077f0a420f00eafab9f7c3629255ebfc5ab9c749e972fc10f44" audited_scope = ["Project.toml", ".github/workflows", "docs/src/shtns37-parity.md", "ext", "src", "test"] audited_exclusions = ["test/fixtures/compatibility/task16_gate.toml", "test/fixtures/compatibility/task16_local_commands.txt", "test/fixtures/compatibility/task16_local_summary.log"] commands_file = "test/fixtures/compatibility/task16_local_commands.txt" -commands_sha256 = "a973719289891efa284d3627e21f8881edc7763c40523205e759b9bded6bbdfc" +commands_sha256 = "a9c74fca5943848f977389ff4519e17125037f1b0cae8b3e20f59e8942573cf5" summary_log = "test/fixtures/compatibility/task16_local_summary.log" -summary_log_sha256 = "1b4df56fbe299b2d4f51a55bafdac0048efade1a8aa4cc56420fe01f097854c8" +summary_log_sha256 = "eaa7e8adbbed059f6604bfa1dfc4e4c45c86eef17703869b1aade77b3aa3794d" # Historical readiness counts; current review runs are recorded in the summary. pkg_exit_code = 0 serial_pass = 76347 diff --git a/test/fixtures/compatibility/task16_local_commands.txt b/test/fixtures/compatibility/task16_local_commands.txt index dd0876b2..861917fd 100644 --- a/test/fixtures/compatibility/task16_local_commands.txt +++ b/test/fixtures/compatibility/task16_local_commands.txt @@ -5,7 +5,7 @@ certifying=false manual_cuda_tool_command=python3 -B -m unittest discover -s test/support -p test_manual_cuda.py -v manual_cuda_contract_command=/Users/subha/.julia/juliaup/julia-1.12.4+0.aarch64.apple.darwin14/bin/julia --startup-file=no --project=. test/serial/test_shtns37_contract.jl manual_cuda_numerical_source_unchanged_from=3b8b2fe118900f2431208088932631db4a162908 -audited_tree_digest=763b5c3413870230bd8b906ced62324a36885f9d026abbdf0ef6788dbecd2793 +audited_tree_digest=34289cc58102e077f0a420f00eafab9f7c3629255ebfc5ab9c749e972fc10f44 audited_scope=Project.toml,.github/workflows,docs/src/shtns37-parity.md,ext,src,test audited_exclusions=test/fixtures/compatibility/task16_gate.toml,test/fixtures/compatibility/task16_local_commands.txt,test/fixtures/compatibility/task16_local_summary.log pkg_command=/usr/bin/env JULIA_DEPOT_PATH=/private/tmp/julia_depot_shtnskit:/Users/subha/.julia SHTNSKIT_RUN_JET_TESTS=1 SHTNSKIT_RUN_AQUA_TESTS=1 /Users/subha/.julia/juliaup/julia-1.12.4+0.aarch64.apple.darwin14/bin/julia --startup-file=no --project=. -e 'using Pkg; Pkg.test()' diff --git a/test/fixtures/compatibility/task16_local_summary.log b/test/fixtures/compatibility/task16_local_summary.log index e59f6c21..16e75ea0 100644 --- a/test/fixtures/compatibility/task16_local_summary.log +++ b/test/fixtures/compatibility/task16_local_summary.log @@ -9,7 +9,7 @@ manual_cuda_no_hardware_guard_rejected=true manual_cuda_julia_syntax_pass=3 manual_cuda_github_status_published=false manual_cuda_numerical_source_unchanged_from=3b8b2fe118900f2431208088932631db4a162908 -audited_tree_digest=763b5c3413870230bd8b906ced62324a36885f9d026abbdf0ef6788dbecd2793 +audited_tree_digest=34289cc58102e077f0a420f00eafab9f7c3629255ebfc5ab9c749e972fc10f44 platform=aarch64-apple-darwin julia=1.12.4 pkg_exit_code=0 @@ -140,5 +140,8 @@ audit2_jet_rerun=false audit2_aqua_rerun=false audit2_gpu_rerun=false audit2_physical_gpu_execution=false -audit2_parity_runtests_rerun=false +audit2_parity_cpu_command=julia --startup-file=no --project= test/parity/runtests_cpu.jl +audit2_parity_cpu_local_evaluation_pass=2212 +audit2_parity_cpu_shtns37_manifest_pass=1082 +audit2_parity_mpi_rerun=false audit2_note=Run on Julia 1.11.1, not the 1.12.4 of the historical readiness runs; no CUDA or ROCm hardware is present on this machine. diff --git a/test/serial/test_local.jl b/test/serial/test_local.jl index 5c04477b..58da8940 100644 --- a/test/serial/test_local.jl +++ b/test/serial/test_local.jl @@ -83,6 +83,34 @@ using Random end end + @testset "evaluators preserve the coefficient element type" begin + # The φ convention factor is a Float64. Multiplying a Float32 evaluation + # by it silently widens every result to Float64, breaking the element + # type these functions promise their caller (and the Dual types AD needs + # to see through). Narrow the scale at the boundary instead. + for mode in (:dft, :quad), T in (Float32, Float64) + cfg = create_gauss_config(4, 6; nlon=9) + cfg.phi_scale = mode + CT = Complex{T} + A = zeros(CT, cfg.lmax + 1, cfg.mmax + 1) + A[2, 1] = CT(0.6); A[3, 2] = CT(0.4, 0.2) + packed = SHTnsKit.pack_lm(cfg, A) + zpack = zeros(CT, cfg.nlm) + x = T(cfg.x[2]); φ = T(cfg.φ[3]) + + @test synthesis_point(cfg, A, x, φ) isa T + @test eltype(SH_to_lat(cfg, packed, x)) === T + @test eltype(SH_to_lat_cplx(cfg, zeros(CT, SHTnsKit.nlm_cplx_calc(cfg.lmax, cfg.mmax, 1)), x)) === CT + @test all(v -> v isa T, SHTnsKit.SHqst_to_point(cfg, packed, zpack, zpack, x, φ)) + @test all(v -> v isa T, SHTnsKit.SH_to_grad_point(cfg, packed, zpack, x, φ)) + @test all(V -> eltype(V) === T, SHqst_to_lat(cfg, packed, zpack, zpack, x)) + @test eltype(synthesis_axisym(cfg, A[:, 1])) === T + @test eltype(SHTnsKit.synthesis_axisym_l(cfg, A[:, 1], cfg.lmax)) === T + @test SHTnsKit.synthesis_point_cplx( + cfg, zeros(CT, SHTnsKit.nlm_cplx_calc(cfg.lmax, cfg.mmax, 1)), x, φ) isa CT + end + end + @testset "SH_to_lat matches synthesis at grid latitudes" begin lmax = 8 nlat = lmax + 2