Multi-vendor GPU support via PyTorch + the array API (v1.1)

CHANGELOG.md

@@ -4,6 +4,43 @@ All notable changes to cuPeriod are documented here. The format is based on
 [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), and this project adheres to
 [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased]
+
+### Added
+
+- **Multi-vendor GPU support via PyTorch and the Python array API.** All seven
+  period-search methods (GLS, BLS, PDM, CE, String-Length, MHAOV, TLS) gain a portable
+  `torch` backend that runs on AMD (ROCm), Intel (XPU), and Apple (MPS) GPUs as well as a
+  real CPU path — so the accelerated code is no longer NVIDIA-only, and works even with no
+  GPU at all. Select it with `backend="torch"` (or `"torch:cpu"`, `"torch:cuda"`,
+  `"torch:mps"`, `"torch:xpu"`); `backend="auto"` now reaches a torch GPU on non-NVIDIA
+  machines after the cufinufft/cupy fast paths.
+  - GLS adds a NUFFT-free direct trig-sum path (the portable formulation; cufinufft
+    remains the NVIDIA fast path).
+  - BLS, PDM, CE, String-Length, MHAOV, and TLS run their vectorized kernels through the
+    array-API namespace; the cupy `RawKernel`s (BLS/PDM/CE/TLS), numba (BLS), and finufft
+    (GLS) remain the fast paths where present.
+  - New `device` and `precision` settings: `precision="auto"` is float64 everywhere it is
+    supported and float32 only where the device forces it (Apple MPS cannot do float64);
+    an explicit `precision="float64"` on MPS raises rather than silently downgrading.
+- `array-api-compat` is now a dependency; install the portable accelerator with the
+  `[torch]` extra (`pip install 'cuperiod[torch]'`).
+
+### Known limitations
+
+- **Non-NVIDIA GPU numerics are written-to-spec and CPU-validated, not yet hardware-
+  verified.** The torch CUDA/ROCm/MPS/XPU paths share the array-API body that is parity-
+  tested on the CPU torch device; on-device parity self-skips (`requires_torch_gpu`) until
+  such hardware is available.
+- **No fp64 capability probe on Intel XPU.** `precision="auto"` resolves to float64 on an
+  XPU; a device without native float64 will error at compute time rather than falling back
+  to float32. Pass `precision="float32"` explicitly on such a device.
+- **The torch GPU path does not auto-shrink to small VRAM.** A large period×bin grid on a
+  small consumer GPU can raise an out-of-memory error; reduce `batch_periods`.
+- **Tie-broken best-fit *extras* may differ across devices.** Where an `argmax` lands on an
+  exact tie (e.g. BLS `transit_time`, TLS `t0`/`duration` at non-transit periods), the
+  chosen index is device-dependent; the periodogram power and best period are unaffected.
+
 ## [1.0.0] — 2026-06-29
 
 First public release.

benchmarks/README.md

@@ -18,6 +18,12 @@ The rendered results are in **[REPORT.md](REPORT.md)** with figures in `figures/
 | Period recovery | all 7 | VSX literature period | find the real period? |
 | Performance | all 7 | astropy / PyAstronomy | how much faster, and how does it scale? |
 
+The performance benchmark also times the **portable `torch` backend** (`backend="torch"`;
+the resolved device is shown in the `torch_backend` column) for the ported methods — GLS
+and BLS — alongside the CPU and CUDA paths, so the cross-vendor path (AMD/Intel/Mac/CPU) is
+tracked. A backend absent on the host (no CUDA GPU, or no torch) is recorded blank rather
+than failing the sweep.
+
 ## Data
 
 * **`dataset/light_curves.parquet`** — 72 real ASAS-SN g-band light curves

benchmarks/benchmark.py

@@ -26,6 +26,9 @@
 import cuperiod as cup  # noqa: E402
 
 from _common import RESULTS, load_dataset  # noqa: E402
+from cuperiod.core.backend import torch_available  # noqa: E402
+from cuperiod.core.errors import BackendUnavailableError  # noqa: E402
+from cuperiod.methods.base import get_method  # noqa: E402
 
 # bounded period window for the box methods (BLS/TLS) — independent of the star,
 # so the trial-period count can never blow up on a short-period target.
@@ -60,6 +63,23 @@ def best_time(fn, repeat=3):
     return min(_timed(fn) for _ in range(repeat))
 
 
+def safe_best_time(fn, repeat=3):
+    """:func:`best_time`, but a backend unavailable here records NaN instead of raising.
+
+    Lets the sweep run on machines missing a backend — no CUDA GPU (the ``gpu`` column),
+    or no torch / no torch GPU device (the ``torch`` column) — leaving that cell blank.
+    """
+    try:
+        return best_time(fn, repeat)
+    except BackendUnavailableError:
+        return float("nan")
+
+
+def supports_torch(method):
+    """Whether ``method`` has the portable torch backend and torch is importable."""
+    return torch_available() and "torch" in get_method(method).all_backends
+
+
 def _timed(fn):
     t0 = time.perf_counter(); fn(); return time.perf_counter() - t0
 
@@ -91,13 +111,23 @@ def bench_single(t, y, e):
     for m, st in FREQ_SETTINGS.items():
         be = cup.periodogram((t, y, e), m, backend="cpu", grid=grid, settings=st()).backend
         tc = best_time(lambda: cup.periodogram((t, y, e), m, backend="cpu", grid=grid, settings=st()))
-        tg = best_time(lambda: cup.periodogram((t, y, e), m, backend="gpu", grid=grid, settings=st()))
+        tg = safe_best_time(lambda: cup.periodogram((t, y, e), m, backend="gpu", grid=grid, settings=st()))
+        has_torch = supports_torch(m)
+        tt = (safe_best_time(lambda: cup.periodogram((t, y, e), m, backend="torch", grid=grid, settings=st()))
+              if has_torch else np.nan)
+        tbe = (cup.periodogram((t, y, e), m, backend="torch", grid=grid, settings=st()).backend
+               if has_torch and np.isfinite(tt) else "—")
         tr = best_time(reftime[m], repeat=1) if m in reftime else np.nan
         rows.append(dict(method=m, n_grid=grid.size, cpu_backend=be, cpu_s=tc, gpu_s=tg,
-                         ref_s=tr, ref=refname.get(m, "—"), gpu_speedup=tc / tg,
-                         gpu_vs_ref=(tr / tg if np.isfinite(tr) else np.nan),
+                         torch_s=tt, torch_backend=tbe,
+                         ref_s=tr, ref=refname.get(m, "—"),
+                         gpu_speedup=(tc / tg if np.isfinite(tg) else np.nan),
+                         torch_speedup=(tc / tt if np.isfinite(tt) else np.nan),
+                         gpu_vs_ref=(tr / tg if np.isfinite(tr) and np.isfinite(tg) else np.nan),
                          cpu_vs_ref=(tr / tc if np.isfinite(tr) else np.nan)))
-        print(f"  {m:12s} cpu({be})={tc:.3f}s gpu={tg:.4f}s ({tc/tg:.0f}x)", flush=True)
+        gstr = f"gpu={tg:.4f}s ({tc/tg:.0f}x)" if np.isfinite(tg) else "gpu=—"
+        tstr = f" torch({tbe})={tt:.3f}s" if np.isfinite(tt) else ""
+        print(f"  {m:12s} cpu({be})={tc:.3f}s {gstr}{tstr}", flush=True)
     # box methods on a fixed, bounded period window. cpu_s is cuPeriod's *default*
     # CPU backend: "cpu" resolves to the multicore numba box search for BLS (or
     # astropy if numba is not installed), to numpy for TLS.
@@ -107,8 +137,13 @@ def bench_single(t, y, e):
              cup.TLSSettings(min_period_days=BOX_PMIN, max_period_days=BOX_PMAX))
         be = cup.periodogram((t, y, e), m, backend="cpu", settings=S).backend
         tc = best_time(lambda: cup.periodogram((t, y, e), m, backend="cpu", settings=S), repeat=2)
-        tg = best_time(lambda: cup.periodogram((t, y, e), m, backend="gpu", settings=S), repeat=2)
-        ng = cup.periodogram((t, y, e), m, backend="gpu", settings=S).power.size
+        tg = safe_best_time(lambda: cup.periodogram((t, y, e), m, backend="gpu", settings=S), repeat=2)
+        ng = cup.periodogram((t, y, e), m, backend="cpu", settings=S).power.size  # backend-independent
+        has_torch = supports_torch(m)
+        tt = (safe_best_time(lambda: cup.periodogram((t, y, e), m, backend="torch", settings=S), repeat=2)
+              if has_torch else np.nan)
+        tbe = (cup.periodogram((t, y, e), m, backend="torch", settings=S).backend
+               if has_torch and np.isfinite(tt) else "—")
         # BLS reference = astropy's compiled BoxLeastSquares; also time the pure-numpy
         # GPU-parity reference port to document it is not the product path.
         ref_s = (best_time(lambda: cup.periodogram((t, y, e), "BLS", backend="astropy", settings=S), repeat=1)
@@ -117,13 +152,18 @@ def bench_single(t, y, e):
         port_s = (best_time(lambda: cup.periodogram((t, y, e), "BLS", backend="numpy", settings=S), repeat=1)
                   if m == "BLS" else np.nan)
         rows.append(dict(method=m, n_grid=int(ng), cpu_backend=be, cpu_s=tc, gpu_s=tg,
-                         ref_s=ref_s, ref=ref, gpu_speedup=tc / tg,
-                         gpu_vs_ref=(ref_s / tg if np.isfinite(ref_s) else np.nan),
+                         torch_s=tt, torch_backend=tbe,
+                         ref_s=ref_s, ref=ref,
+                         gpu_speedup=(tc / tg if np.isfinite(tg) else np.nan),
+                         torch_speedup=(tc / tt if np.isfinite(tt) else np.nan),
+                         gpu_vs_ref=(ref_s / tg if np.isfinite(ref_s) and np.isfinite(tg) else np.nan),
                          cpu_vs_ref=(ref_s / tc if np.isfinite(ref_s) else np.nan),
                          cpu_port_s=port_s))
+        gstr = f"gpu={tg:.4f}s (gpu {tc/tg:.0f}x)" if np.isfinite(tg) else "gpu=—"
+        tstr = f" torch({tbe})={tt:.3f}s" if np.isfinite(tt) else ""
         extra = (f"  [vs astropy {ref_s/tc:.0f}x faster; numpy-port {port_s:.1f}s]"
                  if m == "BLS" else "")
-        print(f"  {m:12s} cpu({be})={tc:.3f}s gpu={tg:.4f}s (gpu {tc/tg:.0f}x){extra}  [{ng:,} periods]",
+        print(f"  {m:12s} cpu({be})={tc:.3f}s {gstr}{tstr}{extra}  [{ng:,} periods]",
               flush=True)
     df = pd.DataFrame(rows)
     df.to_parquet(RESULTS / "bench_single.parquet", index=False)
@@ -138,8 +178,12 @@ def bench_scaling_npoints(t, y, e):
         for m in ["GLS", "PDM", "MHAOV"]:
             st = FREQ_SETTINGS[m]
             tc = best_time(lambda: cup.periodogram((tt, yy, ee), m, backend="cpu", grid=grid, settings=st()), repeat=1)
-            tg = best_time(lambda: cup.periodogram((tt, yy, ee), m, backend="gpu", grid=grid, settings=st()), repeat=1)
-            rows.append(dict(axis="npoints", method=m, n=n, cpu_s=tc, gpu_s=tg, speedup=tc / tg))
+            tg = safe_best_time(lambda: cup.periodogram((tt, yy, ee), m, backend="gpu", grid=grid, settings=st()), repeat=1)
+            ttor = (safe_best_time(lambda: cup.periodogram((tt, yy, ee), m, backend="torch", grid=grid, settings=st()), repeat=1)
+                    if supports_torch(m) else np.nan)
+            rows.append(dict(axis="npoints", method=m, n=n, cpu_s=tc, gpu_s=tg, torch_s=ttor,
+                             speedup=(tc / tg if np.isfinite(tg) else np.nan),
+                             torch_speedup=(tc / ttor if np.isfinite(ttor) else np.nan)))
         print(f"  N={n:>6}: done", flush=True)
     df = pd.DataFrame(rows)
     df.to_parquet(RESULTS / "bench_npoints.parquet", index=False)
@@ -153,8 +197,12 @@ def bench_scaling_grid(t, y, e):
         for m in ["GLS", "PDM", "MHAOV"]:
             st = FREQ_SETTINGS[m]
             tc = best_time(lambda: cup.periodogram((t, y, e), m, backend="cpu", grid=grid, settings=st()), repeat=1)
-            tg = best_time(lambda: cup.periodogram((t, y, e), m, backend="gpu", grid=grid, settings=st()), repeat=1)
-            rows.append(dict(axis="grid", method=m, n=n, cpu_s=tc, gpu_s=tg, speedup=tc / tg))
+            tg = safe_best_time(lambda: cup.periodogram((t, y, e), m, backend="gpu", grid=grid, settings=st()), repeat=1)
+            ttor = (safe_best_time(lambda: cup.periodogram((t, y, e), m, backend="torch", grid=grid, settings=st()), repeat=1)
+                    if supports_torch(m) else np.nan)
+            rows.append(dict(axis="grid", method=m, n=n, cpu_s=tc, gpu_s=tg, torch_s=ttor,
+                             speedup=(tc / tg if np.isfinite(tg) else np.nan),
+                             torch_speedup=(tc / ttor if np.isfinite(ttor) else np.nan)))
         print(f"  grid={n:>7}: done", flush=True)
     df = pd.DataFrame(rows)
     df.to_parquet(RESULTS / "bench_grid.parquet", index=False)

benchmarks/make_report.py

@@ -432,21 +432,27 @@ def main():
                      f"({b.cpu_s*1e3:.0f} ms vs {b.ref_s:.1f} s on this light curve), matching it "
                      f"to floating-point{par}. The GPU then adds another {b.gpu_speedup:.0f}× "
                      f"({b.ref_s/b.gpu_s:.0f}× over astropy).\n")
-        cols = ["method", "cpu_backend", "cpu_s", "gpu_s", "ref", "ref_s", "cpu_vs_ref", "gpu_speedup"]
+        cols = ["method", "cpu_backend", "cpu_s", "gpu_s", "torch_s", "torch_backend",
+                "ref", "ref_s", "cpu_vs_ref", "gpu_speedup"]
         cols = [c for c in cols if c in s.columns]
+        nan_dash = lambda fmt: (lambda v: ("—" if not np.isfinite(v) else fmt(v)))
         L.append(md_table(s, cols, {
-            "cpu_s": lambda v: f"{v:.3f}", "gpu_s": lambda v: f"{v:.4f}",
-            "gpu_speedup": lambda v: f"{v:.0f}x",
-            "cpu_vs_ref": lambda v: ("—" if not np.isfinite(v) else f"{v:.0f}x"),
-            "ref_s": lambda v: ("—" if not np.isfinite(v) else f"{v:.2f}"),
+            "cpu_s": lambda v: f"{v:.3f}",
+            "gpu_s": nan_dash(lambda v: f"{v:.4f}"),
+            "torch_s": nan_dash(lambda v: f"{v:.3f}"),
+            "gpu_speedup": nan_dash(lambda v: f"{v:.0f}x"),
+            "cpu_vs_ref": nan_dash(lambda v: f"{v:.0f}x"),
+            "ref_s": nan_dash(lambda v: f"{v:.2f}"),
             "method": ml}))
         L.append("\n*cpu_backend* = what `backend=\"cpu\"` resolves to — the fast default a user "
                  "gets: finufft (GLS), the multicore numba box search (BLS), numpy (the rest). "
                  "*ref* = the established external tool; *cpu_vs_ref* = how much faster cuPeriod's "
-                 "CPU is than that tool; *gpu_speedup* = GPU over cuPeriod's CPU. cuPeriod's CPU "
-                 "path already beats every reference tool it has (GLS, PDM, BLS) — so the GPU's "
-                 "marginal gain is small where the CPU is already fast (BLS, GLS) and large where "
-                 "it is not (PDM, MHAOV, TLS).\n")
+                 "CPU is than that tool; *gpu_speedup* = GPU over cuPeriod's CPU. *torch_s* = the "
+                 "portable PyTorch backend (device shown in *torch_backend*: cpu/cuda/mps/xpu) — "
+                 "the cross-vendor path that also runs on AMD/Intel/Mac; blank for methods not yet "
+                 "ported to it. cuPeriod's CPU path already beats every reference tool it has "
+                 "(GLS, PDM, BLS) — so the GPU's marginal gain is small where the CPU is already "
+                 "fast (BLS, GLS) and large where it is not (PDM, MHAOV, TLS).\n")
         if len(bls) and "cpu_port_s" in bls and np.isfinite(bls.cpu_port_s.iloc[0]):
             b = bls.iloc[0]
             L.append(f"\n> The pure-`numpy` BLS backend shares one array-module-generic source "

docs/guide/backends.md

@@ -10,13 +10,14 @@ floating-point round-off regardless of which one runs.
 pg = cup.periodogram(lc, "GLS", backend="auto")    # default
 ```
 
-The four selectors:
+The selectors:
 
 | `backend=` | Meaning |
 | --- | --- |
-| `"auto"` *(default)* | Use the GPU when a CUDA device and the `[gpu]` extra are present; otherwise the best CPU backend. The same code runs on any machine. |
+| `"auto"` *(default)* | The NVIDIA CUDA fast path when a CUDA device and the `[gpu]` extra are present; then the portable **torch** backend on another GPU (AMD/Intel/Mac) when the `[torch]` extra sees one; otherwise the best CPU backend. The same code runs on any machine. |
 | `"cpu"` | Force the best CPU backend for this method. |
-| `"gpu"` | Force the GPU backend. Raises {exc}`~cuperiod.BackendUnavailableError` if no GPU is usable. |
+| `"gpu"` | Force a GPU backend — the CUDA fast path, or torch on a non-NVIDIA GPU. Raises {exc}`~cuperiod.BackendUnavailableError` if no GPU is usable. |
+| `"torch"` / `"torch:<device>"` | Force the portable PyTorch backend; `<device>` is `cpu`, `cuda`, `mps`, or `xpu` (bare `"torch"` picks the best present). Needs the `[torch]` extra. |
 | a concrete name | Force a specific implementation, e.g. `"finufft"`, `"astropy"`, `"numpy"`, `"cupy"`, `"cufinufft"`, `"numba"`. |
 
 ## What each method can run
@@ -64,6 +65,41 @@ array-module-generic source with the CUDA kernel so the two validate to floating
 but it is slow (it trades memory traffic for the parallelism that makes the GPU fast).
 For CPU BLS use `numba` (the default with `[fast]`) or `astropy`, **not** `numpy`.
 
+In addition, **every method has a portable `torch` backend** (the `[torch]` extra) that
+runs the same array-API code on any torch device — see below.
+
+## The portable PyTorch backend
+
+`backend="torch"` runs a method through the [array API](https://data-apis.org/array-api/)
+on whichever torch device you have — CUDA, AMD **ROCm**, Apple **MPS**, Intel **XPU**, or
+**CPU** — so the accelerated code is no longer NVIDIA-only and works even with no GPU at
+all. The NVIDIA fast paths (cufinufft for GLS, the cupy kernels for BLS/PDM/CE/TLS) are
+untouched and remain what `"auto"`/`"gpu"` pick on CUDA; torch is the cross-vendor path
+for everything else.
+
+```python
+pg = cup.periodogram(lc, "GLS", backend="torch")        # best torch device
+pg = cup.periodogram(lc, "BLS", backend="torch:xpu")    # pin the Intel GPU
+pg = cup.periodogram(lc, "PDM", backend="torch:cpu", settings=cup.PDMSettings(device="cpu"))
+```
+
+Two orthogonal settings tune it (both environment-overridable, e.g. `CUPERIOD_GLS_DEVICE`):
+
+- **`device`** — `"auto"` (default; the best present), `"cpu"`, `"cuda"`, `"mps"`, `"xpu"`.
+  A `"torch:<device>"` backend string overrides it.
+- **`precision`** — `"auto"` (default) is float64 everywhere it is supported and float32
+  only where the device forces it (Apple MPS; some Intel GPUs). `"float64"` and `"float32"`
+  force it; `precision="float64"` on MPS raises rather than silently downgrading. Results
+  are always returned as float64 numpy arrays regardless of the device precision.
+
+:::{note}
+The portable path is correct everywhere (it matches the CPU reference to round-off — float
+methods bit-for-bit, MHAOV to ~1e-10 from its linear solve) but it is **not** the speed
+champion on CPU, where finufft (GLS) and the numba box search (BLS) are faster. Its value
+is reaching GPUs the CUDA fast paths can't — AMD, Intel, and Apple. On those devices its
+massively-parallel evaluation is the win. Run `cuperiod doctor` to see what you have.
+:::
+
 Inspect the live picture for your install:
 
 ```python

docs/guide/cli.md

@@ -8,7 +8,8 @@ Python API, so the CLI and library share one code path and give identical result
 cuperiod run         one light curve, one or more methods → prints the best periods
 cuperiod batch       many light curves with CPU or GPU workers → Parquet/CSV
 cuperiod methods     list registered methods and their backends
-cuperiod gpu-info    show the GPU and suggested worker counts
+cuperiod gpu-info    show the CUDA GPU and suggested worker counts
+cuperiod doctor      diagnose backends, torch devices, and the precision each uses
 cuperiod grid-info   show a method's trial grid for a light curve (no compute)
 ```
 
@@ -102,6 +103,18 @@ cuperiod gpu-info
 Shows the CUDA device (name, free/total memory, MPS status) and the suggested batch worker
 count per GPU-capable method. If no GPU is present, it says so and exits cleanly.
 
+## `doctor` — full environment diagnosis
+
+```bash
+cuperiod doctor
+```
+
+A one-stop "will the accelerated paths run here, and on what?" report: which backends are
+installed, the NVIDIA CUDA fast paths, the portable **torch** backend and each device it
+sees (CUDA/ROCm/MPS/XPU/CPU) with the precision it would use, and what `backend="auto"`
+resolves to for every method. Reach for it first when a GPU isn't being picked up or you're
+unsure which build of PyTorch you have.
+
 ## `grid-info` — inspect a grid
 
 ```bash