Multi-vendor GPU (PR1): portable torch backend for GLS + BLS via the array API - #10
Merged
Merged
Conversation
tjayasinghe
commented
Jun 30, 2026
tjayasinghe
commented
Owner
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>
This was referenced Jun 30, 2026
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.