Multi-vendor GPU support via PyTorch + the array API (v1.1) - #13
Closed
tjayasinghe wants to merge 9 commits into
Closed
Multi-vendor GPU support via PyTorch + the array API (v1.1)#13tjayasinghe wants to merge 9 commits into
tjayasinghe wants to merge 9 commits into
Conversation
Foundations for multi-vendor GPU support (AMD/Intel/Mac + a real CPU path). GLS and BLS gain a portable `torch` backend that runs on any torch device (CUDA/ROCm/MPS/XPU/CPU) through one array-API code path, leaving the NVIDIA fast paths (cufinufft, cupy RawKernel) untouched. - core/_arrayapi.py (new): array-API dispatch + the ops outside the standard — scatter_add, to_host, device/precision resolution, torch device resolution. - core/backend.py: torch detection (torch_available / torch_devices / best_torch_device); register `torch` in available_backends (no torch import). - methods/base.py: vendor-aware resolve_backend — `auto` prefers cufinufft/cupy on CUDA, then torch on other GPUs, else the proven CPU path; parse + validate `torch` / `torch:<device>`. - methods/gls.py: NUFFT-free direct trig-sum path (lombscargle_power_torch), sharing the Zechmeister-Kürster assembly via _assemble_power_parts (real components, so MPS without complex128 works). NUFFT path unchanged. - methods/_bls_core.py: _bls_search rewritten to array-API-standard ops; numpy and torch both flow through array_api_compat namespaces. cupy RawKernel + numba unchanged. - config.py: `torch` backend + orthogonal device/precision selectors (shared mixin on GLS/BLS) + GLSSettings.direct_freq_batch. - pyproject: array-api-compat core dep; [torch] extra; torch in [dev]; mypy/ruff config; version 1.1.0.dev0. Precision policy: precision="auto" is float64 except on MPS (float32); float64 on MPS raises rather than silently downgrading. Output is always numpy float64. Validated on CPU/torch-CPU: GLS torch vs finufft ~2e-11; BLS torch vs astropy ~4e-11 (torch vs numpy bit-identical); period recovery confirmed end-to-end. 151 passed, 8 skipped (CUDA-only). Also fixes two pre-existing numpy-2 stub typing nits (tls/bls geomspace) the toolchain surfaced. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Fold the torch backend into the permanent benchmark suite so the cross-vendor path is tracked alongside the CPU and CUDA paths. - benchmark.py: bench_single + both scaling sweeps now time backend="torch" (device recorded in torch_backend) for methods that support it (GLS, BLS). New safe_best_time() records NaN for a backend that's unavailable on the host (no CUDA GPU, or no torch) instead of aborting the sweep — so the suite also runs CPU-only, where torch-CPU vs finufft/numba is now a meaningful comparison. Trial-period count (BLS/TLS) read from the CPU run (backend-independent). - make_report.py: §4 performance table surfaces torch_s / torch_backend; the numeric formatters are NaN-safe (blank cell when a backend is absent). - README: note the torch column in the performance benchmark. Verified on CPU-only: GLS cpu(finufft)=19ms, gpu=— (graceful), torch:cpu=81ms. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…ackend Extends the portable array-API path from GLS/BLS (PR1) to the remaining five methods, so all seven now run on any torch device (CUDA/ROCm/MPS/XPU/CPU) with the NVIDIA fast paths untouched. - CE / PDM: vectorized histogram/box kernels re-spelled to array-API ops (scatter_add shim, xp.astype/remainder/clip, dtype-following-periods); numpy + torch flow through array_api_compat namespaces. cupy RawKernels unchanged. - String-Length: sort/gather re-spelled with fancy indexing + slicing (no take_along_axis/diff, which aren't in every namespace). - MHAOV: einsum + linalg.solve are available in the compat namespaces, so the torch path is a thin add; raw numpy/cupy keep einsum. - TLS: matched filter re-spelled (notably explicit float dtype on every zeros(), since torch defaults to float32; clip instead of minimum-with-scalar, which torch rejects). - config: device/precision mixin + "torch" backend on all five settings. - New tests/test_torch_bonus.py: torch:cpu vs numpy parity + period recovery for all five (harmonic-aware, since String-Length/MHAOV legitimately lock onto a harmonic identically on both backends). Validated on torch:cpu: CE 1.3e-15, String-Length 2.8e-14, PDM 0.0, TLS 0.0 (bit-identical); MHAOV 2.7e-10 relative (linalg.solve rounding, peak unaffected). 160 passed, 8 skipped (CUDA-only); ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
PR4 of the multi-vendor GPU effort — the install UX for the torch backend. - cli: new `cuperiod doctor` command — reports installed backends, the NVIDIA CUDA fast paths, the portable torch backend and each device it sees (CUDA/ROCm/MPS/XPU/CPU) with the precision each uses, and what `backend="auto"` resolves to per method. It sets KMP_DUPLICATE_LIB_OK for its own read-only device probe only (it does no numerics), so it can't OMP-abort on Windows; the library still never sets it for compute paths. - docs/installation: `[torch]` extra, per-accelerator PyTorch wheel guidance (the plain wheel is CPU-only), the Apple-MPS float32 note, and a Windows OpenMP-clash warning with the KMP_DUPLICATE_LIB_OK workaround. - docs/guide/backends: torch in the selector table; a "portable PyTorch backend" section covering devices, the device/precision settings, and the fast-on-NVIDIA / portable-everywhere split. - docs/guide/cli: document `doctor`. Verified: `cuperiod doctor` runs clean on a CPU-only Windows box without any external env (torch:cpu -> float64; auto -> finufft/numba/numpy). Docs build with `-W`. 161 passed, 8 skipped; ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The PR5 local adversarial review surfaced three float32-only bugs in the portable torch backend, all masked because every torch test ran float64-on-CPU: - BLS cast absolute BJD time to float32 *before* subtracting t_min, destroying sub-0.25-day timing on real light curves (silent wrong periodogram; this is the default path on Apple MPS). Subtract the origin in float64 before the device cast and restore the absolute transit_time on the host (mirrors gls._prep). - MHAOV multiband dropped `precision`, silently ignoring it (and, on MPS, downcasting to float32 instead of raising). Forward precision to aov_multiband_power. - MHAOV's fixed 1e-10 diagonal ridge underflowed the ~N Gram diagonal in float32, so linalg.solve raised on a singular matrix at degenerate frequencies. Scale the ridge by eps(dtype)*n_points. Add float32 regression tests (three proven to fail pre-fix) and the requires_torch_gpu marker. Document the remaining hardware-gated limitations (XPU fp64 probe, torch-GPU VRAM auto-sizing, GPU argmax tie-breaks of extras) in the changelog. NVIDIA fast paths verified unchanged; 165 passed, ruff+mypy clean. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The portable String-Length body sorted folded phases with whichever argsort the backend namespace provided: the numpy path used raw numpy.argsort (quicksort, unstable) while the torch path used the array-API argsort (stable). On exactly-tied phases the two ordered points differently, so the string length diverged — and because quicksort's tie order is platform-dependent, the numpy-vs-torch parity test passed locally but failed in CI across OS/Python. Route the numpy and cupy paths through the array-API namespace too, so every backend uses the standard stable sort and ties keep a backend- and platform-independent order. Add a regression test that forces heavy phase ties (deterministically failing pre-fix). Surfaced by the PR5 review (the argsort tie divergence) and confirmed by CI. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Multi-vendor GPU (PR1): portable torch backend for GLS + BLS via the array API
Multi-vendor GPU (PR3): port PDM, CE, String-Length, MHAOV, TLS to torch
Multi-vendor GPU (PR4): `cuperiod doctor` + PyTorch install/backend docs
Owner
Author
|
Superseded by #14, which targets the new |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Makes cuPeriod's accelerated period search run beyond NVIDIA — on AMD (ROCm), Intel (XPU), and Apple (MPS) GPUs, and on a real CPU path — via a portable PyTorch backend dispatched through the Python array API (
array-api-compat). All seven methods (GLS, BLS, PDM, CE, String-Length, MHAOV, TLS) gain atorchbackend. The existing NVIDIA fast paths (cufinufft for GLS; cupyRawKernels for BLS/PDM/CE/TLS; numba for BLS) are unchanged and stay the default underautowhere present.This is the accumulated
dev/multi-vendor-gpuintegration branch for the v1.1 release. It bundles:cuperiod doctor+ PyTorch install/backend docs #12) —cuperiod doctorCLI + per-platform install/backend docs.Architecture
Two-tier: a portable array-API core (numpy/cupy/torch through one namespace) plus the unchanged NVIDIA fast paths.
precision="auto"is float64 everywhere supported, float32 only where the device forces it (MPS); an explicit float64 on MPS raises rather than silently downgrading. Results always return as numpy float64.PR5 review (gate before merge)
A local multi-agent adversarial review (5 reviewers) over the whole diff, plus CI. NVIDIA fast paths verified byte-for-byte unchanged; backward-compat and packaging clean. Real bugs were found and fixed — all float32-only, all masked because every torch test had run float64-on-CPU:
tau = t - t_min, destroying sub-0.25-day timing (silent wrong periodogram; the default path on Apple MPS). Fixed with a float64 host pre-shift, restoring the absolutetransit_timeon the host.precision(silent ignore; on MPS a silent downcast instead of the mandated raise). Fixed by forwarding it.1e-10ridge underflowed the ~N Gram diagonal in float32, crashinglinalg.solveat degenerate frequencies. Fixed to aneps(dtype)·n_points-scaled ridge.np.argsort(unstable) vs the torch path's array-API stable sort, so tied phases diverged across platforms (green locally, red in CI). Fixed by routing every backend through the array-API stable sort.Added float32 + tie-forcing regression tests (several proven to fail pre-fix) and a
requires_torch_gpumarker.Known limitations (documented in CHANGELOG)
Non-NVIDIA GPU numerics are written-to-spec and CPU-validated but not yet hardware-verified (no such hardware available): the XPU fp64 capability probe, torch-GPU VRAM auto-sizing, and the GPU
argmaxtie-break of best-fit extras are deferred and self-skip viarequires_torch_gpu.Testing
166 tests pass locally on the torch-CPU path; CI is green across ubuntu + windows × py3.11/3.12, lint + types, and docs.
🤖 Generated with Claude Code