From de9f299b801ff34b828b988be8215b776d3006ae Mon Sep 17 00:00:00 2001 From: "Yu, Guangye" <106960996+guangyey@users.noreply.github.com> Date: Thu, 16 Oct 2025 13:56:07 -0700 Subject: [PATCH] Add auto-detection for Intel GPU on Windows (#16280) This PR enables `--torch-backend=auto` to automatically detect Intel GPUs. It follows up on [#14386](https://github.com/astral-sh/uv/pull/14386). On Windows, detection is implemented by querying the `Win32_VideoController` class via the [WMI crate](https://github.com/ohadravid/wmi-rs/tree/v0.16.0). Currently, Intel GPUs (XPU) do not depend on specific driver or toolkit versions to determine which PyTorch wheel to use. --- Cargo.lock | 15 +++++++++ Cargo.toml | 1 + crates/uv-torch/Cargo.toml | 3 ++ crates/uv-torch/src/accelerator.rs | 50 ++++++++++++++++++++++++++++++ crates/uv-torch/src/backend.rs | 5 ++- docs/guides/integration/pytorch.md | 8 ----- 6 files changed, 71 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 53d4db9dc..e5adda24c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6714,6 +6714,7 @@ dependencies = [ "uv-pep440", "uv-platform-tags", "uv-static", + "wmi", ] [[package]] @@ -7470,6 +7471,20 @@ dependencies = [ "bitflags 2.9.4", ] +[[package]] +name = "wmi" +version = "0.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d9189bc72f0e4d814d812216ec06636ce3ea5597ff5f1ff9f9f0e5ec781c027" +dependencies = [ + "futures", + "log", + "serde", + "thiserror 2.0.17", + "windows 0.61.3", + "windows-core 0.61.2", +] + [[package]] name = "writeable" version = "0.6.1" diff --git a/Cargo.toml b/Cargo.toml index 3597757eb..c8a07e260 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -201,6 +201,7 @@ which = { version = "8.0.0", features = ["regex"] } windows = { version = "0.59.0", features = ["Win32_Globalization", "Win32_Security", "Win32_System_Console", "Win32_System_Kernel", "Win32_System_Diagnostics_Debug", "Win32_Storage_FileSystem", "Win32_System_Registry", "Win32_System_IO", "Win32_System_Ioctl"] } windows-registry = { version = "0.5.0" } wiremock = { version = "0.6.4" } +wmi = { version = "0.16.0", default-features = false } xz2 = { version = "0.1.7" } zeroize = { version = "1.8.1" } zip = { version = "2.2.3", default-features = false, features = ["deflate", "zstd", "bzip2", "lzma", "xz"] } diff --git a/crates/uv-torch/Cargo.toml b/crates/uv-torch/Cargo.toml index 407052d8c..ae04fed04 100644 --- a/crates/uv-torch/Cargo.toml +++ b/crates/uv-torch/Cargo.toml @@ -25,5 +25,8 @@ thiserror = { workspace = true } tracing = { workspace = true } url = { workspace = true } +[target.'cfg(all(target_os = "windows"))'.dependencies] +wmi = { workspace = true } + [lints] workspace = true diff --git a/crates/uv-torch/src/accelerator.rs b/crates/uv-torch/src/accelerator.rs index 696adc9a1..b96b2b67b 100644 --- a/crates/uv-torch/src/accelerator.rs +++ b/crates/uv-torch/src/accelerator.rs @@ -6,6 +6,11 @@ use tracing::debug; use uv_pep440::Version; use uv_static::EnvVars; +#[cfg(windows)] +use serde::Deserialize; +#[cfg(windows)] +use wmi::{COMLibrary, WMIConnection}; + #[derive(Debug, thiserror::Error)] pub enum AcceleratorError { #[error(transparent)] @@ -60,6 +65,7 @@ impl Accelerator { /// 5. `nvidia-smi --query-gpu=driver_version --format=csv,noheader`. /// 6. `rocm_agent_enumerator`, which lists the AMD GPU architectures. /// 7. `/sys/bus/pci/devices`, filtering for the Intel GPU via PCI. + /// 8. Windows Managmeent Instrumentation (WMI), filtering for the Intel GPU via PCI. pub fn detect() -> Result, AcceleratorError> { // Constants used for PCI device detection. const PCI_BASE_CLASS_MASK: u32 = 0x00ff_0000; @@ -187,6 +193,50 @@ impl Accelerator { Err(e) => return Err(e.into()), } + // Detect Intel GPU via WMI on Windows + #[cfg(windows)] + { + #[derive(Deserialize, Debug)] + #[serde(rename = "Win32_VideoController")] + #[serde(rename_all = "PascalCase")] + struct VideoController { + #[serde(rename = "PNPDeviceID")] + pnp_device_id: Option, + name: Option, + } + + match COMLibrary::new() { + Ok(com_library) => match WMIConnection::new(com_library) { + Ok(wmi_connection) => match wmi_connection.query::() { + Ok(gpu_controllers) => { + for gpu_controller in gpu_controllers { + if let Some(pnp_device_id) = &gpu_controller.pnp_device_id { + if pnp_device_id + .contains(&format!("VEN_{PCI_VENDOR_ID_INTEL:04X}")) + { + debug!( + "Detected Intel GPU from WMI: PNPDeviceID={}, Name={:?}", + pnp_device_id, gpu_controller.name + ); + return Ok(Some(Self::Xpu)); + } + } + } + } + Err(e) => { + debug!("Failed to query WMI for video controllers: {e}"); + } + }, + Err(e) => { + debug!("Failed to create WMI connection: {e}"); + } + }, + Err(e) => { + debug!("Failed to initialize COM library: {e}"); + } + } + } + debug!("Failed to detect GPU driver version"); Ok(None) diff --git a/crates/uv-torch/src/backend.rs b/crates/uv-torch/src/backend.rs index b9fe58a93..d67fed91d 100644 --- a/crates/uv-torch/src/backend.rs +++ b/crates/uv-torch/src/backend.rs @@ -466,11 +466,10 @@ impl TorchStrategy { ))), }, Self::Xpu { os, source } => match os { - Os::Manylinux { .. } => Either::Right(Either::Right(Either::Left( + Os::Manylinux { .. } | Os::Windows => Either::Right(Either::Right(Either::Left( std::iter::once(TorchBackend::Xpu.index_url(*source)), ))), - Os::Windows - | Os::Musllinux { .. } + Os::Musllinux { .. } | Os::Macos { .. } | Os::FreeBsd { .. } | Os::NetBsd { .. } diff --git a/docs/guides/integration/pytorch.md b/docs/guides/integration/pytorch.md index 0553c4449..5bbb6bb94 100644 --- a/docs/guides/integration/pytorch.md +++ b/docs/guides/integration/pytorch.md @@ -460,12 +460,4 @@ $ # With an environment variable. $ UV_TORCH_BACKEND=cu126 uv pip install torch torchvision ``` -On Windows, Intel GPU (XPU) is not automatically selected with `--torch-backend=auto`, but you can -manually specify it using `--torch-backend=xpu`: - -```shell -$ # Manual selection for Intel GPU. -$ uv pip install torch torchvision --torch-backend=xpu -``` - At present, `--torch-backend` is only available in the `uv pip` interface.