From 3eb68a371de3e1813294850f9717de69808a37c1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Fri, 27 Sep 2024 11:42:17 +0000 Subject: [PATCH] Capability as `usize`s --- launcher/src/gpu.rs | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/launcher/src/gpu.rs b/launcher/src/gpu.rs index 88be0775..755d246a 100644 --- a/launcher/src/gpu.rs +++ b/launcher/src/gpu.rs @@ -1,9 +1,9 @@ use std::sync::LazyLock; -pub static COMPUTE_CAPABILITY: LazyLock> = +pub static COMPUTE_CAPABILITY: LazyLock> = LazyLock::new(get_cuda_capability); -fn get_cuda_capability() -> Option<(isize, isize)> { +fn get_cuda_capability() -> Option<(usize, usize)> { use pyo3::prelude::*; let py_get_capability = |py: Python| -> PyResult<(isize, isize)> { @@ -13,7 +13,11 @@ fn get_cuda_capability() -> Option<(isize, isize)> { }; match pyo3::Python::with_gil(py_get_capability) { - Ok(capability) => Some(capability), + Ok((major, minor)) if major < 0 || minor < 0 => { + tracing::warn!("Ignoring negative GPU compute capabilities: {major}.{minor}"); + None + } + Ok((major, minor)) => Some((major as usize, minor as usize)), Err(err) => { tracing::warn!("Cannot determine GPU compute capability: {}", err); None