feat(cuda): cuda_device_count, and two device probes that still aborted

.github/workflows/cuda_nightly.yml

@@ -95,14 +95,31 @@ jobs:
               --test nvptx_gpu --test groebner_cuda
         continue-on-error: false
 
-      - name: Run compute-sanitizer racecheck
+      # Promoted to a gate: it is clean on real hardware (0 hazards, twice, on
+      # 2x RTX 3090), and it was never the false-positive risk the old comment
+      # assumed — but read the coverage note before trusting a green tick here.
+      #
+      # racecheck detects hazards on __shared__ memory, and no kernel in this
+      # workspace uses any: the hand-written `ELIMINATE_ROW_PTX` declares no
+      # `.shared`, `compile_cuda`'s generated PTX declares none, and launches
+      # pass `shared_mem_bytes: 0`. So today this step *cannot* fail, and its
+      # green tick is a statement about the absence of shared memory, not about
+      # the absence of races. It is a gate for the first kernel that introduces
+      # `__shared__` — which is precisely when it starts being able to fail, and
+      # when nobody would think to add it.
+      #
+      # Left enabled rather than deleted because the cost is ~4 s and the
+      # alternative is remembering. Named for what it checks so a reader does
+      # not over-read the result; `continue-on-error: false` so that the day it
+      # does have something to say, it is heard.
+      - name: Run compute-sanitizer racecheck (__shared__ hazards; none today)
         env:
           ALKAHEST_GPU_TESTS: "1"
         run: |
           compute-sanitizer --target-processes all --tool racecheck \
             cargo test --features cuda,groebner-cuda \
               --test nvptx_gpu --test groebner_cuda
-        continue-on-error: true  # racecheck may have false positives
+        continue-on-error: false
 
       - name: Upload sanitizer logs
         if: always()

alkahest-core/src/jit/mod.rs

@@ -67,7 +67,7 @@ fn registry() -> &'static PrimitiveRegistry {
 #[cfg(feature = "cuda")]
 pub mod nvptx;
 #[cfg(feature = "cuda")]
-pub use nvptx::{compile_cuda, CudaCompiledFn, CudaError};
+pub use nvptx::{compile_cuda, cuda_device_count, CudaCompiledFn, CudaError};
 
 #[cfg(feature = "cranelift")]
 mod cranelift_backend;

alkahest-core/src/jit/nvptx.rs

@@ -87,6 +87,34 @@ impl crate::errors::AlkahestError for CudaError {
     }
 }
 
+/// Number of CUDA devices visible to this process, or `0` when none are.
+///
+/// Answers "which ordinals may I pass to [`CudaCompiledFn::call_batch_on`]?"
+/// without the caller having to probe by launching and catching
+/// `E-CUDA-003`, which is what `docs/mdbook/src/gpu.md` had to recommend
+/// while no verified implementation existed.
+///
+/// Returns `0` rather than an error for every "no GPU here" shape — no
+/// driver, no device, driver too old. The distinction a caller acts on is
+/// "can I use a GPU", and each of those answers it identically; an ordinal
+/// is valid iff it is `< cuda_device_count()`.
+///
+/// `catch_unwind` is load-bearing for the same reason it is in
+/// `groebner_cuda.rs::gpu_available`: `cudarc` *panics* rather than
+/// returning `Err` when `libcuda.so` cannot be dlopen'd at all, which is the
+/// state of any machine with no driver installed. A capability probe that
+/// aborts the process on the exact configuration it exists to report on
+/// would be worse than useless.
+#[cfg(feature = "cuda")]
+pub fn cuda_device_count() -> usize {
+    std::panic::catch_unwind(|| {
+        cudarc::driver::CudaContext::device_count()
+            .map(|n| n.max(0) as usize)
+            .unwrap_or(0)
+    })
+    .unwrap_or(0)
+}
+
 /// A compiled CUDA kernel for evaluating an Alkahest expression on GPU.
 ///
 /// The generated PTX is self-contained (libdevice has been linked in during

alkahest-core/src/lib.rs

@@ -171,7 +171,7 @@ pub use stablehlo::emit_stablehlo;
 
 // V5-3 — NVPTX JIT backend
 #[cfg(feature = "cuda")]
-pub use jit::{compile_cuda, CudaCompiledFn, CudaError};
+pub use jit::{compile_cuda, cuda_device_count, CudaCompiledFn, CudaError};
 
 // Phase 22 — Ball arithmetic
 pub use ball::{AcbBall, ArbBall, IntervalEval, DEFAULT_PREC};
@@ -269,7 +269,7 @@ pub mod stable {
     pub use crate::integrate::{integrate, integrate_definite, IntegrationError};
     pub use crate::jit::{compile, CompileCache, CompiledFn, JitError};
     #[cfg(feature = "cuda")]
-    pub use crate::jit::{compile_cuda, CudaCompiledFn, CudaError};
+    pub use crate::jit::{compile_cuda, cuda_device_count, CudaCompiledFn, CudaError};
     #[allow(deprecated)]
     pub use crate::kernel::pool_persist::PoolPersistError;
     pub use crate::kernel::pool_persist::{load_from, open_persistent, save_to, IoError};

alkahest-core/tests/groebner_cuda.rs

@@ -337,3 +337,59 @@ fn gpu_macaulay_reduce_kernel() {
         "CPU and GPU row reduction must agree mod p"
     );
 }
+
+/// The error-fallback branch: a device was requested, the driver refused it,
+/// and the run finished on the CPU with the failure recorded.
+///
+/// Distinct from [`assert_no_gpu`], which covers `device_id: None` — that path
+/// never touches the driver, so it leaves `first_gpu_error` empty and exercises
+/// neither the `Err(e)` arm of the `reduce_gpu` dispatch nor the latch that
+/// keeps the *first* error. Those lines had never executed anywhere: a
+/// `reductions_on_gpu` counter stuck at zero would make `ran_on_gpu()`
+/// permanently false and `assert_ran_on_gpu` a gate that could never pass,
+/// while a `first_gpu_error` that never latched would make a fallback
+/// indistinguishable from a clean CPU run — which is the defect the report
+/// exists to expose.
+#[test]
+fn requested_device_that_does_not_exist_falls_back_and_records_why() {
+    if !gpu_available() {
+        return;
+    }
+    // Far past any plausible ordinal on a real host; `CudaContext::new` fails
+    // rather than succeeding on some other card.
+    const ABSENT_DEVICE: usize = 4096;
+
+    let f = poly(&[(&[1, 0], 1), (&[0, 1], 1), (&[0, 0], -1)]);
+    let g = poly(&[(&[1, 0], 1), (&[0, 1], -1)]);
+    let order = MonomialOrder::Lex;
+
+    let (basis, backend) =
+        compute_groebner_basis_gpu(vec![f.clone(), g.clone()], order, Some(ABSENT_DEVICE))
+            .expect("a GPU failure must fall back, not abort the computation");
+
+    assert_eq!(backend.requested_device, Some(ABSENT_DEVICE));
+    assert_eq!(
+        backend.reductions_on_gpu, 0,
+        "nothing can have run on a device that does not exist: {backend:?}"
+    );
+    assert!(backend.reductions_on_cpu > 0, "{backend:?}");
+    assert!(backend.fell_back_to_cpu(), "{backend:?}");
+    assert!(
+        !backend.ran_on_gpu(),
+        "a run that fell back is not a GPU run: {backend:?}"
+    );
+    assert!(
+        backend.first_gpu_error.is_some(),
+        "the driver refusal must be recorded for a caller that never reads \
+         stderr: {backend:?}"
+    );
+
+    // The fallback is only acceptable because the answer is still right.
+    let basis_cpu = compute_groebner_basis(vec![f, g], order);
+    for p in &basis_cpu {
+        assert!(cpu_reduce(p, &basis, order).is_zero());
+    }
+    for p in &basis {
+        assert!(cpu_reduce(p, &basis_cpu, order).is_zero());
+    }
+}

alkahest-core/tests/nvptx_gpu.rs

@@ -26,7 +26,16 @@ const N_BW: usize = 16 << 20;
 /// remains usable on a developer machine with no GPU. Main CI never builds
 /// these features at all, so this cannot block ordinary PRs.
 fn device_available() -> bool {
-    let device_ok = cudarc::driver::CudaContext::new(0).is_ok();
+    // `catch_unwind` for the same reason as `groebner_cuda.rs::gpu_available`:
+    // `cudarc` *panics* rather than returning `Err` when `libcuda.so` cannot be
+    // dlopen'd at all, which is the state of any machine with no driver. That
+    // fix landed on the Gröbner suite only, so this one still aborted on
+    // exactly the driverless machines the doc comment above promises to
+    // support. Missing library and missing device must both mean "not
+    // available"; only a device *asserted* to exist and not usable is a
+    // failure.
+    let device_ok =
+        std::panic::catch_unwind(|| cudarc::driver::CudaContext::new(0).is_ok()).unwrap_or(false);
     let requested = std::env::var("ALKAHEST_GPU_TESTS").ok().as_deref() == Some("1");
     assert!(
         !requested || device_ok,
@@ -143,7 +152,10 @@ fn nvptx_bandwidth_sin_cos_16m() {
 
 #[test]
 fn nvptx_multi_device_both_3090s() {
-    let n_dev = cudarc::driver::CudaContext::device_count().unwrap_or(0) as usize;
+    // Via the guarded helper rather than `CudaContext::device_count()` direct:
+    // that call panics on a driverless machine for the same dlopen reason as
+    // `CudaContext::new`, so this test aborted where it meant to skip.
+    let n_dev = alkahest_cas::jit::nvptx::cuda_device_count();
     if n_dev < 2 {
         eprintln!("skipped: only {n_dev} CUDA device(s) present");
         return;
@@ -235,3 +247,43 @@ fn nvptx_polynomial_beats_cpu_jit() {
         gpu_time.as_secs_f64() * 1e3
     );
 }
+
+/// `cuda_device_count` must agree with the ordinals that actually launch.
+///
+/// The point of the function is to answer "which ordinals may I pass to
+/// `call_batch_on`?" without probing by launching and catching `E-CUDA-003`.
+/// That contract is only worth anything if the count and the launches agree,
+/// so this checks both directions: every ordinal below the count runs, and the
+/// first ordinal at the count does not.
+#[test]
+fn cuda_device_count_matches_the_ordinals_that_launch() {
+    if !device_available() {
+        return;
+    }
+    let n = alkahest_cas::jit::nvptx::cuda_device_count();
+    assert!(n > 0, "a device initialised, so the count cannot be zero");
+
+    let pool = ExprPool::new();
+    let x = pool.symbol("x", Domain::Real);
+    let expr = pool.add(vec![pool.mul(vec![x, x]), pool.integer(1)]);
+    let compiled = compile_cuda(expr, &[x], &pool).expect("compile_cuda");
+
+    let xs = [0.0f64, 1.0, -2.5, 4.0];
+    let want = [1.0f64, 2.0, 7.25, 17.0];
+
+    for dev in 0..n {
+        let mut got = vec![0.0f64; xs.len()];
+        compiled
+            .call_batch_on(dev, &[&xs[..]], &mut got)
+            .unwrap_or_else(|e| panic!("ordinal {dev} < count {n} must launch, got {e}"));
+        for (g, w) in got.iter().zip(want.iter()) {
+            assert!((g - w).abs() < 1e-12, "device {dev}: {g} vs {w}");
+        }
+    }
+
+    let mut got = vec![0.0f64; xs.len()];
+    assert!(
+        compiled.call_batch_on(n, &[&xs[..]], &mut got).is_err(),
+        "ordinal {n} == count must not launch"
+    );
+}

alkahest-py/src/lib.rs

@@ -9656,6 +9656,25 @@ fn py_compile_cuda(
     Ok(PyCudaCompiledFn { inner: compiled })
 }
 
+/// `alkahest.cuda_device_count() -> int`
+///
+/// Number of CUDA devices visible to this process; `0` when none are.
+///
+/// The valid arguments to :meth:`CudaCompiledFn.call_batch_on` are exactly
+/// ``range(cuda_device_count())``. Without this, the only way to discover the
+/// range was to launch on an ordinal and catch ``E-CUDA-003`` — the workaround
+/// `docs/mdbook/src/gpu.md` had to document while no verified implementation
+/// existed.
+///
+/// Never raises: every "no GPU here" shape (no driver, no device, driver too
+/// old) reports `0`, which is the single answer a caller acts on.
+#[cfg(feature = "cuda")]
+#[pyfunction]
+#[pyo3(name = "cuda_device_count")]
+fn py_cuda_device_count() -> usize {
+    alkahest_core::cuda_device_count()
+}
+
 // ---------------------------------------------------------------------------
 // V5-11 — Gröbner basis
 // ---------------------------------------------------------------------------
@@ -11466,6 +11485,7 @@ fn alkahest(m: &Bound<'_, PyModule>) -> PyResult<()> {
     {
         m.add_class::<PyCudaCompiledFn>()?;
         m.add_function(wrap_pyfunction!(py_compile_cuda, m)?)?;
+        m.add_function(wrap_pyfunction!(py_cuda_device_count, m)?)?;
     }
     // V5-11 — Gröbner basis / V1-16 — GroebnerBasis.compute
     #[cfg(feature = "groebner")]

docs/mdbook/src/gpu.md

@@ -55,7 +55,8 @@ Read these bits precisely — each says **what was linked**, and nothing more:
 - `cuda == True` guarantees `ak.compile_cuda` and `ak.CudaCompiledFn` exist and that
   PTX can be emitted on the host. It does **not** promise a GPU: the driver is loaded
   lazily, so a machine with no device compiles happily and fails at `call_batch` with
-  `E-CUDA-003`. The only way to find out is to launch something.
+  `E-CUDA-003`. Ask `ak.cuda_device_count()` — it reports `0` when there is no usable
+  device, and never raises.
 - `llvm_jit == True` on a `cuda` build even when *alkahest-py*'s own `jit` feature was
   never named, because `alkahest-core`'s `cuda` feature turns on `jit`. Cranelift and
   LLVM are not mutually exclusive; a CUDA build can link both.
@@ -99,29 +100,31 @@ DCE → PTX for `sm_86` (Ampere) → loaded through the CUDA driver by `cudarc`.
 
 ### Discovering the valid device ordinals
 
-There is **no `ak.cuda_device_count()`**. The only way to find the valid range for
-`call_batch_on` today is to try an ordinal and catch the refusal:
-
 ```python
-def device_count(fn, limit=16):
-    """Largest N such that ordinals 0..N-1 accept a launch."""
-    n = 0
-    while n < limit:
-        try:
-            fn.call_batch_on(n, [[0.0]] * fn.n_inputs)
-        except ak.CudaError:   # E-CUDA-003 — no such device
-            return n
-        n += 1
-    return n
+n = ak.cuda_device_count()          # 0 when there is no GPU here
+for dev in range(n):
+    fn.call_batch_on(dev, [[0.0]] * fn.n_inputs)
 ```
 
-That is a workaround, not an API, and it is recorded here rather than fixed because a
-`cuda_device_count` binding could not be verified by anything: `cuda` implies LLVM 15
-with NVPTX, so it cannot even be *compiled* on an ordinary dev box, no CI job builds
-the Python extension with the feature (see below), and running it needs a device. It
-belongs in the same change as the missing `maturin develop --features cuda` nightly
-step — shipping it before that would add exactly the kind of unverified surface that
-produced the capability overclaims this page now documents.
+The valid arguments to `call_batch_on` are exactly `range(ak.cuda_device_count())`.
+Both directions are pinned by tests that run on hardware
+(`tests/test_cuda.py::test_device_count_agrees_with_the_ordinals_that_launch` and
+`nvptx_gpu::cuda_device_count_matches_the_ordinals_that_launch`): every ordinal below
+the count launches, and the ordinal *at* the count is refused.
+
+`cuda_device_count()` **never raises**. Every "no GPU here" shape — no driver, no
+device, driver too old — reports `0`, because that is the single answer a caller acts
+on. This matters more than it looks: `cudarc` *panics* rather than returning `Err`
+when `libcuda.so` cannot be `dlopen`'d at all, so a naive binding would abort the
+process on precisely the machines a capability probe exists to report on. The same
+bug bit `groebner_cuda.rs::gpu_available` and `nvptx_gpu.rs::device_available`.
+
+Earlier releases documented a workaround here — launch on an ordinal and catch
+`E-CUDA-003` — because a `cuda_device_count` binding could not be verified by
+anything on an ordinary dev box: `cuda` implies LLVM 15 with NVPTX, so it could not
+even be *compiled*, and no CI job built the Python extension with the feature. It
+shipped once both could be done on real hardware, which is the standard the
+capability overclaims this page documents were failing.
 
 ### Limits worth knowing before you reach for it
 

python/alkahest/__init__.py

@@ -2212,8 +2212,8 @@ def __getattr__(name: str):
 # native raise. See `tests/test_cuda.py`.
 # ---------------------------------------------------------------------------
 try:  # pragma: no cover - exercised only on CUDA builds
-    from .alkahest import CudaCompiledFn, compile_cuda
+    from .alkahest import CudaCompiledFn, compile_cuda, cuda_device_count
 except ImportError:  # the overwhelmingly common case: no CUDA feature
     pass
 else:
-    __all__ += ["CudaCompiledFn", "compile_cuda"]
+    __all__ += ["CudaCompiledFn", "compile_cuda", "cuda_device_count"]

tests/test_agent_contract.py

@@ -240,7 +240,11 @@ def _probe_cuda():
 #: entry point missing (the `cuda` bug in `d139a46`) *and* a bit reading
 #: `False` on a build that really does have it.
 _FEATURE_EXCLUSIVE_NAMES = {
-    "cuda": (("alkahest", "compile_cuda"), ("alkahest", "CudaCompiledFn")),
+    "cuda": (
+        ("alkahest", "compile_cuda"),
+        ("alkahest", "CudaCompiledFn"),
+        ("alkahest", "cuda_device_count"),
+    ),
     "parallel": (
         ("alkahest.CompiledFn", "call_batch_raw_par"),
         ("alkahest.CompiledFn", "call_batch_buffer_par"),