diff --git a/qdp/qdp-core/src/gpu/cuda_ffi.rs b/qdp/qdp-core/src/gpu/cuda_ffi.rs index dc8930584c..dbbe844db2 100644 --- a/qdp/qdp-core/src/gpu/cuda_ffi.rs +++ b/qdp/qdp-core/src/gpu/cuda_ffi.rs @@ -69,6 +69,10 @@ unsafe extern "C" { pub(crate) fn cudaMemGetInfo(free: *mut usize, total: *mut usize) -> i32; + /// Number of CUDA-capable devices. Used to detect whether a GPU is + /// actually present at runtime (see `cuda_runtime_available`). + pub(crate) fn cudaGetDeviceCount(count: *mut i32) -> i32; + pub(crate) fn cudaMemcpyAsync( dst: *mut c_void, src: *const c_void, @@ -153,6 +157,10 @@ mod no_cuda_stubs { QDP_CUDA_UNAVAILABLE } + pub(crate) unsafe fn cudaGetDeviceCount(_count: *mut i32) -> i32 { + QDP_CUDA_UNAVAILABLE + } + pub(crate) unsafe fn cudaMemcpyAsync( _dst: *mut c_void, _src: *const c_void, @@ -216,3 +224,19 @@ mod no_cuda_stubs { #[cfg(qdp_no_cuda)] pub(crate) use no_cuda_stubs::*; + +/// Returns `true` if a usable CUDA device is present at runtime. +/// +/// This is deliberately distinct from "the extension is importable": a +/// no-toolkit (`qdp_no_cuda`) build links CUDA stubs that return a non-zero +/// sentinel, and a host with the toolkit but no GPU reports zero devices. Both +/// yield `false` here, so callers (and the Python test suite) can gate GPU work +/// on an accurate signal rather than mere importability. +pub fn cuda_runtime_available() -> bool { + let mut count: i32 = 0; + // SAFETY: `cudaGetDeviceCount` writes a single `i32` through `count` and + // accesses no other memory. In a `qdp_no_cuda` build this resolves to the + // stub above, which returns the unavailable sentinel without dereferencing. + let rc = unsafe { cudaGetDeviceCount(&mut count) }; + rc == CUDA_SUCCESS && count > 0 +} diff --git a/qdp/qdp-core/src/lib.rs b/qdp/qdp-core/src/lib.rs index 822fba96e1..95259793cf 100644 --- a/qdp/qdp-core/src/lib.rs +++ b/qdp/qdp-core/src/lib.rs @@ -36,6 +36,7 @@ pub mod types; mod profiling; pub use error::{MahoutError, Result, cuda_error_to_string}; +pub use gpu::cuda_ffi::cuda_runtime_available; pub use gpu::memory::Precision; pub use reader::{FloatElem, NullHandling, handle_float32_nulls, handle_float64_nulls}; pub use types::{Dtype, Encoding}; diff --git a/qdp/qdp-python/qumat_qdp/__init__.py b/qdp/qdp-python/qumat_qdp/__init__.py index 2b75661417..81ffd66347 100644 --- a/qdp/qdp-python/qumat_qdp/__init__.py +++ b/qdp/qdp-python/qumat_qdp/__init__.py @@ -44,6 +44,27 @@ BACKEND = get_backend() + +def is_cuda_available() -> bool: + """Return whether a usable CUDA device is available to the native engine. + + Unlike the importability of the ``_qdp`` extension -- which only means it + was built, possibly against CUDA stubs on a host without the toolkit -- this + reflects whether GPU work can actually run: ``False`` for a stub build or a + host with no CUDA device, ``True`` only when the native runtime reports at + least one device. + """ + if _qdp_mod is None: + return False + probe = getattr(_qdp_mod, "cuda_available", None) + if probe is None: + return False + try: + return bool(probe()) + except Exception: + return False + + from qumat_qdp.api import ( LatencyResult, QdpBenchmark, @@ -68,6 +89,7 @@ "ThroughputResult", "TritonAmdEngine", "force_backend", + "is_cuda_available", "is_triton_amd_available", "run_throughput_pipeline_py", ] diff --git a/qdp/qdp-python/src/lib.rs b/qdp/qdp-python/src/lib.rs index 04d772a906..0bfd77189e 100644 --- a/qdp/qdp-python/src/lib.rs +++ b/qdp/qdp-python/src/lib.rs @@ -68,6 +68,18 @@ fn run_throughput_pipeline_py( )) } +/// Returns ``True`` if a usable CUDA device is available to the native engine. +/// +/// This reflects whether GPU work can actually run -- it is ``False`` for a +/// stub build (the extension built without the CUDA toolkit) or a host with no +/// CUDA device, and ``True`` only when the native runtime reports a device. +/// Importability of this module only means it was built, which a stub build +/// makes possible without a GPU. +#[pyfunction] +fn cuda_available() -> bool { + qdp_core::cuda_runtime_available() +} + /// Quantum Data Plane (QDP) Python module /// /// GPU-accelerated quantum data encoding with DLPack integration. @@ -78,6 +90,7 @@ fn _qdp(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; + m.add_function(wrap_pyfunction!(cuda_available, m)?)?; #[cfg(target_os = "linux")] m.add_class::()?; #[cfg(target_os = "linux")] diff --git a/testing/conftest.py b/testing/conftest.py index 95009cc3cf..010b7230af 100644 --- a/testing/conftest.py +++ b/testing/conftest.py @@ -47,9 +47,18 @@ def _gpu_available() -> bool: toolkit -- it links stub CUDA Runtime symbols (see qdp-core ``build.rs`` / ``cuda_ffi.rs``). So importing ``_qdp`` no longer implies a working GPU, and ``@pytest.mark.gpu`` tests must additionally check for a real device. - This mirrors the ``torch.cuda.is_available()`` guard used by the inline - skips throughout the QDP test modules. + + Prefers the native engine's own signal (``qumat_qdp.is_cuda_available()``), + which is correct even for a stub build: a host with a GPU + PyTorch but no + CUDA toolkit builds ``_qdp`` against stubs, where ``torch.cuda.is_available()`` + would wrongly report True. Falls back to torch when the helper is absent. """ + try: + from qumat_qdp import is_cuda_available + + return bool(is_cuda_available()) + except Exception: + pass try: import torch diff --git a/testing/qdp_python/test_fallback.py b/testing/qdp_python/test_fallback.py index 9b34ff12b4..ac9e7739d4 100644 --- a/testing/qdp_python/test_fallback.py +++ b/testing/qdp_python/test_fallback.py @@ -85,6 +85,29 @@ def test_get_torch(self): t = get_torch() assert t is not None # torch is available in test env + def test_is_cuda_available_returns_bool(self): + """is_cuda_available() returns a plain bool without crashing. + + On a stub build (the extension built without the CUDA toolkit) or a host + with no device it must be False -- and querying it must NOT abort the + process, which is the regression this guards (a stub build previously + aborted when GPU code ran). + """ + from qumat_qdp import is_cuda_available + + assert isinstance(is_cuda_available(), bool) + + def test_is_cuda_available_matches_extension(self): + """The helper mirrors the native ``_qdp.cuda_available()`` signal.""" + from qumat_qdp import is_cuda_available + + try: + import _qdp + except ImportError: + assert is_cuda_available() is False + else: + assert is_cuda_available() == bool(_qdp.cuda_available()) + # --------------------------------------------------------------------------- # Loader with explicit PyTorch backend