diff --git a/.github/workflows/cuda_nightly.yml b/.github/workflows/cuda_nightly.yml index d7c64b29..4c85bc94 100644 --- a/.github/workflows/cuda_nightly.yml +++ b/.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() diff --git a/alkahest-core/src/jit/mod.rs b/alkahest-core/src/jit/mod.rs index 5a95f016..6e345baa 100644 --- a/alkahest-core/src/jit/mod.rs +++ b/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; diff --git a/alkahest-core/src/jit/nvptx.rs b/alkahest-core/src/jit/nvptx.rs index 7a5399d6..13160dd4 100644 --- a/alkahest-core/src/jit/nvptx.rs +++ b/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 diff --git a/alkahest-core/src/lib.rs b/alkahest-core/src/lib.rs index e768c3b9..3eabedc2 100644 --- a/alkahest-core/src/lib.rs +++ b/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}; diff --git a/alkahest-core/tests/groebner_cuda.rs b/alkahest-core/tests/groebner_cuda.rs index 36344f85..d3cdc7bf 100644 --- a/alkahest-core/tests/groebner_cuda.rs +++ b/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()); + } +} diff --git a/alkahest-core/tests/nvptx_gpu.rs b/alkahest-core/tests/nvptx_gpu.rs index 9905bbe6..0af595f2 100644 --- a/alkahest-core/tests/nvptx_gpu.rs +++ b/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" + ); +} diff --git a/alkahest-py/src/lib.rs b/alkahest-py/src/lib.rs index bf98fab7..b21fdc72 100644 --- a/alkahest-py/src/lib.rs +++ b/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::()?; 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")] diff --git a/docs/mdbook/src/gpu.md b/docs/mdbook/src/gpu.md index 00e1b2f7..940a7a99 100644 --- a/docs/mdbook/src/gpu.md +++ b/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 diff --git a/python/alkahest/__init__.py b/python/alkahest/__init__.py index afb21645..92e0f243 100644 --- a/python/alkahest/__init__.py +++ b/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"] diff --git a/tests/test_agent_contract.py b/tests/test_agent_contract.py index 0f282986..a1301ddf 100644 --- a/tests/test_agent_contract.py +++ b/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"), diff --git a/tests/test_cuda.py b/tests/test_cuda.py index 3f8eb220..2970460c 100644 --- a/tests/test_cuda.py +++ b/tests/test_cuda.py @@ -38,9 +38,9 @@ import alkahest as ak import pytest -# The two names the native module defines *under* `--features cuda`. Both must -# be reachable from the public package, or neither. -CUDA_ENTRY_POINTS = ("CudaCompiledFn", "compile_cuda") +# The names the native module defines *under* `--features cuda`. All must be +# reachable from the public package, or none. +CUDA_ENTRY_POINTS = ("CudaCompiledFn", "compile_cuda", "cuda_device_count") # `CudaError` is not in that list on purpose: the native module registers it # unconditionally, like every other exception class, so it is bound on every @@ -424,3 +424,53 @@ def test_single_point_batch(pool): assert len(got) == 1 assert math.isclose(got[0], 10.0, rel_tol=1e-15) + + +@requires_cuda_build +def test_device_count_is_answerable_without_probing_by_exception(): + """The count must be a plain question with a plain answer. + + Before this existed the only way to learn the valid ordinal range was to + launch on a guess and catch ``E-CUDA-003`` — the workaround + ``docs/mdbook/src/gpu.md`` documented for exactly as long as there was no + verified implementation to replace it. + + Never raising is part of the contract: every "no GPU here" shape reports + ``0``. A probe that raised on a driverless machine would be unusable in the + one place callers most need it. + """ + n = ak.cuda_device_count() + assert isinstance(n, int) + assert n >= 0 + + +@requires_gpu +def test_device_count_agrees_with_the_ordinals_that_launch(pool): + """The count is only worth anything if launches agree with it. + + Checked in both directions: every ordinal below the count runs and returns + the same values, and the first ordinal at the count is refused. A count + that over-reports sends callers at a device that cannot run; one that + under-reports hides hardware they paid for. + """ + n = ak.cuda_device_count() + assert n >= 1, "a device answered the probe, so the count cannot be zero" + + x = pool.symbol("x") + fn = ak.compile_cuda(x * x + pool.integer(1), [x]) + pts = [0.0, 1.0, -2.5, 4.0] + expected = [1.0, 2.0, 7.25, 17.0] + + first = None + for dev in range(n): + got = fn.call_batch_on(dev, [pts]) + assert got == pytest.approx(expected, rel=1e-15) + if first is None: + first = got + else: + # Identical PTX on identical hardware: agreement should be exact, + # not merely close. + assert got == first, f"device {dev} disagreed bitwise with device 0" + + with pytest.raises(ak.CudaError): + fn.call_batch_on(n, [pts])