semlith 0.44.0

//! The GPU's half of [`crate::accel`]: whether this machine has one, the
//! pinned components a lane needs, and the session a worker runs.
//!
//! Only a hardware adapter is used. A machine whose only device is a software
//! renderer — Mesa's lavapipe or llvmpipe, SwiftShader, Microsoft's WARP or
//! Basic Render Driver, a hypervisor's framebuffer — counts as having no GPU:
//! lavapipe aborts the process on any MatMul (an LLVM fatal error), and even a
//! renderer that ran would be slower than the CPU it runs on. It is refused by
//! vendor before anything is downloaded, and again by the worker before any
//! session is made.

use anyhow::{Context, Result, bail};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};

/// PCI vendors that are never a hardware GPU: hypervisor framebuffers and
/// software renderers. Microsoft's is WARP, Basic Render and Hyper-V's video;
/// Mesa's is lavapipe and llvmpipe; Google's is SwiftShader.
const SOFTWARE_VENDORS: &[u32] = &[
    0x1414,  // Microsoft: WARP, Basic Render Driver, Hyper-V video
    0x10005, // Mesa: lavapipe, llvmpipe
    0x1ae0,  // Google: SwiftShader
    0x1234,  // QEMU's standard VGA
    0x15ad,  // VMware SVGA
    0x1af4,  // virtio-gpu
    0x80ee,  // VirtualBox
    0x1b36,  // Red Hat QXL
];

/// Adapter names that give a software renderer away when a vendor id does not.
const SOFTWARE_NAMES: &[&str] = &[
    "llvmpipe",
    "lavapipe",
    "swiftshader",
    "warp",
    "basic render",
    "basic display",
    "hyper-v",
    "virtio",
    "vmware",
    "virtualbox",
];

pub fn is_software(vendor: u32, name: &str) -> bool {
    let name = name.to_ascii_lowercase();
    SOFTWARE_VENDORS.contains(&vendor) || SOFTWARE_NAMES.iter().any(|n| name.contains(n))
}

/// One physical device as Vulkan lists it.
#[derive(Debug, Clone)]
pub struct VulkanDevice {
    pub vendor: u32,
    /// `VkPhysicalDeviceType`: 1 integrated, 2 discrete, 3 virtual, 4 CPU.
    pub kind: u32,
    pub name: String,
}

/// The Vulkan device the WebGPU lane can run on: a discrete GPU first, then
/// any other that is not the CPU or a software renderer by vendor or name.
/// Only software devices, or none, is the reason the lane is unavailable.
pub fn hardware_vulkan(devices: &[VulkanDevice]) -> std::result::Result<&VulkanDevice, String> {
    let hardware = |d: &&VulkanDevice| d.kind != 4 && !is_software(d.vendor, &d.name);
    devices
        .iter()
        .filter(hardware)
        .find(|d| d.kind == 2)
        .or_else(|| devices.iter().find(hardware))
        .ok_or_else(|| {
            if devices.is_empty() {
                "Vulkan lists no device: the GPU maker's Vulkan driver is missing".to_string()
            } else {
                let names: Vec<&str> = devices.iter().map(|d| d.name.as_str()).collect();
                format!(
                    "Vulkan has only a software device ({}), e.g. the GPU maker's Vulkan driver \
                     is missing or the container lacks the graphics capability",
                    names.join(", ")
                )
            }
        })
}

/// The hardware GPU this machine has, by name, or why it has none.
pub fn detect() -> std::result::Result<String, String> {
    platform::detect()
}

#[cfg(target_os = "macos")]
mod platform {
    /// Every Apple silicon Mac has a Metal GPU, and it is the one Metal device
    /// a GitHub macOS arm64 runner exposes too.
    pub fn detect() -> std::result::Result<String, String> {
        if cfg!(target_arch = "aarch64") {
            Ok(format!("{} GPU (Metal)", crate::system::cpu_name()))
        } else {
            Err("no Metal GPU semlith supports on an Intel Mac".to_string())
        }
    }
}

#[cfg(target_os = "linux")]
mod platform {
    use super::*;

    /// A DRM card whose PCI vendor is a real GPU maker, and a Vulkan loader
    /// for the plugin to reach it through.
    pub fn detect() -> std::result::Result<String, String> {
        let mut found = Vec::new();
        for entry in std::fs::read_dir("/sys/class/drm")
            .into_iter()
            .flatten()
            .flatten()
        {
            let name = entry.file_name().to_string_lossy().into_owned();
            if !name.starts_with("card") || name.contains('-') {
                continue;
            }
            let Ok(vendor) = std::fs::read_to_string(entry.path().join("device/vendor")) else {
                continue;
            };
            let Ok(vendor) = u32::from_str_radix(vendor.trim().trim_start_matches("0x"), 16) else {
                continue;
            };
            if !is_software(vendor, "") {
                found.push(vendor_name(vendor));
            }
        }
        let Some(first) = found.first() else {
            return Err("no hardware GPU found".to_string());
        };
        // The DRM card says a GPU is plugged in, not that the plugin can
        // reach it: in a container without the graphics capability, or with
        // no driver from the GPU's maker, Vulkan offers only Mesa's llvmpipe,
        // which the plugin then runs on while reporting the card's PCI id
        // (#197). Ask Vulkan itself.
        let devices = vulkan_devices().map_err(|why| format!("{first} found, but {why}"))?;
        super::hardware_vulkan(&devices).map(|device| device.name.clone())
    }

    /// Every physical device Vulkan lists, through the loader directly: no
    /// crate, and the instance is destroyed before returning.
    fn vulkan_devices() -> std::result::Result<Vec<super::VulkanDevice>, String> {
        use std::ffi::{c_char, c_void};
        #[repr(C)]
        struct InstanceCreateInfo {
            s_type: u32,
            next: *const c_void,
            flags: u32,
            application: *const c_void,
            layer_count: u32,
            layers: *const *const c_char,
            extension_count: u32,
            extensions: *const *const c_char,
        }
        type Create =
            unsafe extern "C" fn(*const InstanceCreateInfo, *const c_void, *mut *mut c_void) -> i32;
        type Destroy = unsafe extern "C" fn(*mut c_void, *const c_void);
        type Enumerate = unsafe extern "C" fn(*mut c_void, *mut u32, *mut *mut c_void) -> i32;
        type Properties = unsafe extern "C" fn(*mut c_void, *mut u8);

        // SAFETY: the loader's own exported Vulkan 1.0 entry points, called
        // with the signatures the Vulkan headers give them. The properties
        // buffer is 4 KiB, aligned for any field, and VkPhysicalDeviceProperties
        // is 824 bytes; the name is read only up to its 256-byte bound.
        unsafe {
            let handle = libc::dlopen(
                c"libvulkan.so.1".as_ptr(),
                libc::RTLD_NOW | libc::RTLD_LOCAL,
            );
            if handle.is_null() {
                return Err("no Vulkan loader (libvulkan.so.1) is installed".to_string());
            }
            let symbol = |name: &std::ffi::CStr| libc::dlsym(handle, name.as_ptr());
            let (create, destroy, enumerate, properties) = (
                symbol(c"vkCreateInstance"),
                symbol(c"vkDestroyInstance"),
                symbol(c"vkEnumeratePhysicalDevices"),
                symbol(c"vkGetPhysicalDeviceProperties"),
            );
            if [create, destroy, enumerate, properties]
                .iter()
                .any(|f| f.is_null())
            {
                libc::dlclose(handle);
                return Err("the Vulkan loader lacks its core entry points".to_string());
            }
            let create = std::mem::transmute::<*mut c_void, Create>(create);
            let destroy = std::mem::transmute::<*mut c_void, Destroy>(destroy);
            let enumerate = std::mem::transmute::<*mut c_void, Enumerate>(enumerate);
            let properties = std::mem::transmute::<*mut c_void, Properties>(properties);
            let info = InstanceCreateInfo {
                s_type: 1, // VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO
                next: std::ptr::null(),
                flags: 0,
                application: std::ptr::null(),
                layer_count: 0,
                layers: std::ptr::null(),
                extension_count: 0,
                extensions: std::ptr::null(),
            };
            let mut instance = std::ptr::null_mut();
            let result = create(&info, std::ptr::null(), &mut instance);
            if result != 0 || instance.is_null() {
                libc::dlclose(handle);
                return Err(format!(
                    "Vulkan would not start (vkCreateInstance returned {result}): no Vulkan \
                     driver works here"
                ));
            }
            let mut count = 0u32;
            let mut devices = Vec::new();
            if enumerate(instance, &mut count, std::ptr::null_mut()) == 0 && count > 0 {
                let mut handles = vec![std::ptr::null_mut(); count as usize];
                if enumerate(instance, &mut count, handles.as_mut_ptr()) >= 0 {
                    handles.truncate(count as usize);
                    for device in handles {
                        let mut buffer = [0u64; 512];
                        let bytes = buffer.as_mut_ptr().cast::<u8>();
                        properties(device, bytes);
                        let field =
                            |at: usize| u32::from_ne_bytes(*bytes.add(at).cast::<[u8; 4]>());
                        let name = std::slice::from_raw_parts(bytes.add(20), 256);
                        let end = name.iter().position(|&b| b == 0).unwrap_or(name.len());
                        devices.push(super::VulkanDevice {
                            vendor: field(8),
                            kind: field(16),
                            name: String::from_utf8_lossy(&name[..end]).into_owned(),
                        });
                    }
                }
            }
            destroy(instance, std::ptr::null());
            libc::dlclose(handle);
            Ok(devices)
        }
    }

    fn vendor_name(vendor: u32) -> String {
        match vendor {
            0x10de => "NVIDIA GPU".to_string(),
            0x1002 => "AMD GPU".to_string(),
            0x8086 => "Intel GPU".to_string(),
            other => format!("GPU (PCI vendor {other:04x})"),
        }
    }
}

#[cfg(windows)]
mod platform {
    use super::*;

    /// The display adapters Windows lists, less the software ones. Read
    /// through PowerShell once per lane start, which is seconds apart from
    /// anything a user waits on.
    pub fn detect() -> std::result::Result<String, String> {
        let out = std::process::Command::new("powershell")
            .args([
                "-NoProfile",
                "-NonInteractive",
                "-Command",
                "Get-CimInstance Win32_VideoController | ForEach-Object { $_.Name + '|' + $_.PNPDeviceID }",
            ])
            .output()
            .map_err(|e| format!("could not list display adapters: {e}"))?;
        let text = String::from_utf8_lossy(&out.stdout);
        for line in text.lines() {
            let Some((name, pnp)) = line.split_once('|') else {
                continue;
            };
            let vendor = pnp
                .to_ascii_uppercase()
                .split("VEN_")
                .nth(1)
                .and_then(|rest| u32::from_str_radix(rest.get(..4)?, 16).ok())
                .unwrap_or(0);
            if vendor != 0 && !is_software(vendor, name) {
                return Ok(name.trim().to_string());
            }
        }
        Err("no hardware GPU found".to_string())
    }
}

#[cfg(not(any(target_os = "macos", target_os = "linux", windows)))]
mod platform {
    pub fn detect() -> std::result::Result<String, String> {
        Err("no GPU support on this platform".to_string())
    }
}

/// Why the CUDA lane cannot run in this build, when it cannot.
pub fn cuda_unavailable() -> String {
    cuda_unavailable_here().unwrap_or_else(|| "the CUDA lane is not part of this build".to_string())
}

/// Why CUDA cannot be turned on at all on this platform. Linux only in this
/// release: an NVIDIA card on Windows runs through WebGPU (D3D12).
pub fn cuda_unavailable_here() -> Option<String> {
    crate::cuda::unavailable_here()
}

/// Where the CUDA pack lives in the model cache.
pub fn cuda_dir(cache: &Path) -> PathBuf {
    crate::cuda::pack_dir(cache)
}

/// What turning CUDA on downloads, in bytes, said before it starts.
pub fn cuda_pack_bytes() -> u64 {
    crate::cuda::PACK_BYTES
}

// ------------------------------------------------------------ the components

/// The WebGPU plugin release, and the wheel it comes in for this platform:
/// URL, SHA-256 and size, from PyPI's own listing for 0.4.0.
pub const WEBGPU_VERSION: &str = "0.4.0";

struct Wheel {
    url: &'static str,
    sha256: &'static str,
    size: u64,
}

fn wheel() -> Option<Wheel> {
    if cfg!(target_os = "macos") {
        Some(Wheel {
            url: "https://files.pythonhosted.org/packages/c3/9f/1e09865bbe4c9202ff1334545535be32dfdbeffc46e7d4c36c5de67ffafe/onnxruntime_ep_webgpu-0.4.0-py3-none-macosx_14_0_universal2.whl",
            sha256: "8d0a91d44b43d931c3068b9134aed9770f27803e353c2b58c93e1700bb68bcca",
            size: 5_200_679,
        })
    } else if cfg!(all(target_os = "linux", target_arch = "x86_64")) {
        Some(Wheel {
            url: "https://files.pythonhosted.org/packages/c1/96/2a18a45079250afcd825687aef2895de266d5abc1cff9f52d2dede7598f2/onnxruntime_ep_webgpu-0.4.0-py3-none-manylinux_2_28_x86_64.whl",
            sha256: "5d8c961eda91a88961cad1fb7621f389f006b18497d6ea9ecf96f729f066cbfb",
            size: 6_892_120,
        })
    } else if cfg!(all(target_os = "linux", target_arch = "aarch64")) {
        Some(Wheel {
            url: "https://files.pythonhosted.org/packages/4e/ca/00c70322c19913c81a6bb2239aca83781b97a818b55c2ec3d25cec72dc4c/onnxruntime_ep_webgpu-0.4.0-py3-none-manylinux_2_28_aarch64.whl",
            sha256: "17f660db53b1a509c63e721ac6784a2daa1c4e91a968523ee9986505ee5b0a55",
            size: 6_096_222,
        })
    } else if cfg!(all(windows, target_arch = "x86_64")) {
        Some(Wheel {
            url: "https://files.pythonhosted.org/packages/d7/a4/c98a9e9433b3eeb576977b26c5c1cd0364f15ba9d198bb16101e7563ab06/onnxruntime_ep_webgpu-0.4.0-py3-none-win_amd64.whl",
            sha256: "7db646669d1a2390551da675115ea125e68cb3d4d2c3a2d6e463372bf8d6ea87",
            size: 13_149_847,
        })
    } else {
        None
    }
}

/// The fp16 variant of the pinned granite revision. The int8 graph does not
/// load on WebGPU, so the GPU lanes run this one.
pub const FP16_FILES: &[(&str, &str, u64)] = &[
    (
        "onnx/model_fp16.onnx",
        "ee200de55cb2f94e858aabca54be7697a9c0805a14c858ee26ad0922b05f57d7",
        200_792,
    ),
    (
        "onnx/model_fp16.onnx_data",
        "28d16e29cd623f25cc6fa0968700c5bc31036466091a5fa06d1353c1777f050e",
        97_402_880,
    ),
];

/// The plugin library's file name on this platform.
pub fn plugin_file() -> &'static str {
    if cfg!(target_os = "macos") {
        "libonnxruntime_providers_webgpu.dylib"
    } else if cfg!(windows) {
        "onnxruntime_providers_webgpu.dll"
    } else {
        "libonnxruntime_providers_webgpu.so"
    }
}

/// Everything the WebGPU lane needs, in one directory under the model cache:
/// the plugin library (and on Windows the two shader compilers beside it) and
/// the fp16 weights, each verified by digest. Fetched on first use, when a
/// hardware GPU was found and the lane is on; refused under `--airgap`
/// unless it was pre-seeded.
pub fn fetch_webgpu(cache: &Path, progress: &mut dyn FnMut(u8)) -> Result<PathBuf> {
    let dir = crate::accel::component_dir(cache, &format!("webgpu-{WEBGPU_VERSION}"));
    let stamp = dir.join(".semlith-verified");
    if stamp.exists() && dir.join(plugin_file()).exists() {
        return Ok(dir);
    }
    if crate::embed::airgap() {
        bail!(
            "{} is set and the WebGPU components are not in {} — pre-seed them on a \
             connected machine, or drop --airgap",
            crate::embed::AIRGAP_ENV,
            crate::plain(&dir.display().to_string())
        );
    }
    let Some(wheel) = wheel() else {
        bail!("unavailable — no WebGPU plugin is published for this platform");
    };
    crate::embed::check_cache_dir(cache)?;
    std::fs::create_dir_all(&dir).with_context(|| format!("creating {}", dir.display()))?;

    let total = wheel.size + FP16_FILES.iter().map(|(_, _, size)| size).sum::<u64>();
    let mut done = 0u64;
    let mut report = |bytes: u64| {
        done += bytes;
        progress(((done * 100) / total.max(1)).min(100) as u8);
    };

    let wheel_path = dir.join("plugin.whl");
    download(wheel.url, &wheel_path, wheel.sha256, &mut report)?;
    extract_plugin(&wheel_path, &dir)?;
    let _ = std::fs::remove_file(&wheel_path);

    fetch_fp16_into(&dir, &mut report)?;
    crate::home::write_private(&stamp, WEBGPU_VERSION.as_bytes())?;
    Ok(dir)
}

/// The fp16 weights alone, for the CUDA lane, which needs them and not the
/// WebGPU plugin. Same directory, same digests.
pub fn fetch_fp16(cache: &Path, progress: &mut dyn FnMut(u8)) -> Result<PathBuf> {
    let dir = crate::accel::component_dir(cache, &format!("webgpu-{WEBGPU_VERSION}"));
    let total: u64 = FP16_FILES.iter().map(|(_, _, size)| size).sum();
    let present = FP16_FILES.iter().all(|(file, _, _)| {
        Path::new(file)
            .file_name()
            .is_some_and(|name| dir.join(name).exists())
    });
    if present {
        return Ok(dir);
    }
    if crate::embed::airgap() {
        bail!(
            "{} is set and the fp16 model is not in {}",
            crate::embed::AIRGAP_ENV,
            crate::plain(&dir.display().to_string())
        );
    }
    crate::embed::check_cache_dir(cache)?;
    std::fs::create_dir_all(&dir).with_context(|| format!("creating {}", dir.display()))?;
    let mut done = 0u64;
    fetch_fp16_into(&dir, &mut |bytes| {
        done += bytes;
        progress(((done * 100) / total.max(1)).min(100) as u8);
    })?;
    Ok(dir)
}

fn fetch_fp16_into(dir: &Path, report: &mut dyn FnMut(u64)) -> Result<()> {
    let base = format!(
        "https://huggingface.co/{}/resolve/{}",
        crate::embed::GRANITE_REPO,
        crate::embed::GRANITE_REVISION
    );
    for (file, sha256, _) in FP16_FILES {
        let name = Path::new(file).file_name().expect("a file name");
        download(&format!("{base}/{file}"), &dir.join(name), sha256, report)?;
    }
    Ok(())
}

/// The plugin wheel's size on this platform, for the Privacy page.
pub fn plugin_bytes() -> u64 {
    wheel().map(|w| w.size).unwrap_or(0)
}

/// Stream `url` to `to`, hashing as it goes, and refuse it unless the digest
/// is the pinned one. A refused file is deleted, never kept.
pub(crate) fn download(
    url: &str,
    to: &Path,
    sha256: &str,
    progress: &mut dyn FnMut(u64),
) -> Result<()> {
    use sha2::{Digest, Sha256};
    if to.exists() && crate::embed::digest(&std::fs::read(to)?) == sha256 {
        progress(std::fs::metadata(to)?.len());
        return Ok(());
    }
    let agent: ureq::Agent = ureq::Agent::config_builder()
        .https_only(true)
        .timeout_connect(Some(std::time::Duration::from_secs(20)))
        .timeout_global(Some(std::time::Duration::from_secs(1800)))
        .build()
        .into();
    crate::add::note_outbound("pack", url);
    let mut response = agent
        .get(url)
        .call()
        .with_context(|| format!("fetching {url}"))?;
    let partial = to.with_extension("partial");
    let mut out = std::fs::File::create(&partial)
        .with_context(|| format!("writing {}", partial.display()))?;
    let mut reader = response.body_mut().as_reader();
    let mut hasher = Sha256::new();
    let mut buffer = vec![0u8; 1 << 16];
    loop {
        let n = reader
            .read(&mut buffer)
            .with_context(|| format!("reading {url}"))?;
        if n == 0 {
            break;
        }
        hasher.update(&buffer[..n]);
        out.write_all(&buffer[..n])?;
        progress(n as u64);
    }
    out.flush()?;
    drop(out);
    let actual: String = hasher
        .finalize()
        .iter()
        .map(|b| format!("{b:02x}"))
        .collect();
    if actual != sha256 {
        let _ = std::fs::remove_file(&partial);
        bail!(
            "{url} does not match the digest semlith pins for it (expected {sha256}, got {actual}); \
             it was deleted"
        );
    }
    std::fs::rename(&partial, to).with_context(|| format!("placing {}", to.display()))?;
    crate::home::tighten_file(to);
    Ok(())
}

/// The plugin's libraries, out of the wheel's `onnxruntime_ep_webgpu/`.
fn extract_plugin(wheel: &Path, dir: &Path) -> Result<()> {
    let file = std::fs::File::open(wheel)?;
    let mut archive = zip::ZipArchive::new(file).context("reading the plugin wheel")?;
    let mut placed = 0;
    for i in 0..archive.len() {
        let mut entry = archive.by_index(i)?;
        let name = entry.name()?.into_owned();
        let Some(base) = name.strip_prefix("onnxruntime_ep_webgpu/") else {
            continue;
        };
        let library = [".dylib", ".so", ".dll"]
            .iter()
            .any(|ext| base.ends_with(ext));
        if !library || base.contains('/') || base.contains("..") {
            continue;
        }
        let target = dir.join(base);
        let mut out = std::fs::File::create(&target)?;
        std::io::copy(&mut entry, &mut out)?;
        placed += 1;
    }
    if placed == 0 || !dir.join(plugin_file()).exists() {
        bail!("the plugin wheel held no {}", plugin_file());
    }
    Ok(())
}

// --------------------------------------------------------------- the session

/// What a worker embeds with.
pub enum Session {
    /// The int8 CPU session, run in a worker: the verification lane.
    Cpu(Box<crate::session::CpuSession>, String),
    /// fp16 on the GPU through ONNX Runtime directly, because fastembed has no
    /// way to hand a session a plugin device.
    Gpu(Box<GpuSession>),
    /// fp16 on an NVIDIA card through the CUDA pack's own ONNX Runtime.
    Cuda(Box<crate::cuda::Session>),
    /// The Neural Engine, or the GPU, through Core ML: the pack's native
    /// models rather than an ONNX graph.
    #[cfg(target_os = "macos")]
    CoreMl(Box<crate::coreml::Session>),
    /// fp16 on an RTX-class NVIDIA card through NVIDIA's TensorRT for RTX
    /// plugin.
    Trt(Box<crate::trt::Session>),
    /// fp16 on an Intel GPU or NPU through Intel's OpenVINO plugin.
    OpenVino(Box<crate::openvino::Session>),
    /// granite as GGUF f16 in a `llama-server` the worker runs.
    Llama(Box<crate::llama::Session>),
}

pub struct GpuSession {
    session: ort::session::Session,
    device: String,
}

/// Names the WebGPU adapter to prefer, by a substring of its name, over the
/// one semlith would choose. `semlith doctor --gpu` prints the choice.
pub const ADAPTER_ENV: &str = "SEMLITH_GPU_ADAPTER";

impl Session {
    /// `progress(loaded, of)` is told as a lane with several models loads
    /// them: the Neural Engine's first load compiles each for this Mac.
    pub fn open(
        lane: &str,
        dir: Option<&Path>,
        adapter: Option<&str>,
        progress: &mut dyn FnMut(usize, usize),
    ) -> Result<Self> {
        // Only the Core ML lanes load several models and say so.
        #[cfg(not(target_os = "macos"))]
        let _ = progress;
        let cache = crate::model_cache_dir()?;
        // Before anything touches ONNX Runtime: the CUDA worker loads the
        // pack's GPU core, and a CPU core loaded first would make that a
        // silent no-op.
        if lane == "cuda" {
            let pack = dir.context("the cuda lane needs its pack directory")?;
            let model = crate::accel::component_dir(&cache, &format!("webgpu-{WEBGPU_VERSION}"))
                .join("model_fp16.onnx");
            return Ok(Session::Cuda(Box::new(crate::cuda::Session::open(
                pack, &model,
            )?)));
        }
        #[cfg(target_os = "macos")]
        if matches!(lane, "ane" | "gpu-coreml") {
            let pack = dir.context("a Core ML lane needs its pack directory")?;
            let kind = if lane == "ane" {
                crate::coreml::Kind::NeuralEngine
            } else {
                crate::coreml::Kind::Gpu
            };
            return Ok(Session::CoreMl(Box::new(crate::coreml::Session::open(
                pack, kind, progress,
            )?)));
        }
        // llama.cpp is a server of its own; this process never loads ONNX
        // Runtime for it.
        if lane == "llama" {
            let pack = dir.context("the llama lane needs its pack directory")?;
            return Ok(Session::Llama(Box::new(crate::llama::Session::open(pack)?)));
        }
        crate::embed::link_runtime()?;
        match lane {
            "worker" => {
                let model = crate::session::CpuSession::open(
                    &cache,
                    crate::embed::Variant::Int8,
                    crate::embed::embed_threads(),
                    true,
                )?;
                Ok(Session::Cpu(Box::new(model), crate::system::cpu_name()))
            }
            "gpu" => {
                let dir = dir.context("the gpu lane needs its component directory")?;
                Ok(Session::Gpu(Box::new(GpuSession::open(dir, adapter)?)))
            }
            "trt" => {
                let pack = dir.context("the trt lane needs its pack directory")?;
                Ok(Session::Trt(Box::new(crate::trt::Session::open(pack)?)))
            }
            "openvino" => {
                let pack = dir.context("the openvino lane needs its pack directory")?;
                Ok(Session::OpenVino(Box::new(crate::openvino::Session::open(
                    pack,
                )?)))
            }
            other => bail!("unavailable — no lane called {other}"),
        }
    }

    pub fn device(&self) -> String {
        match self {
            Session::Cpu(_, name) => name.clone(),
            Session::Gpu(gpu) => gpu.device.clone(),
            Session::Cuda(cuda) => cuda.device(),
            #[cfg(target_os = "macos")]
            Session::CoreMl(coreml) => coreml.device(),
            Session::Trt(trt) => trt.device(),
            Session::OpenVino(openvino) => openvino.device(),
            Session::Llama(llama) => llama.device(),
        }
    }

    pub fn variant(&self) -> &'static str {
        match self {
            Session::Cpu(..) => "int8-cpu",
            Session::Gpu(_) => "fp16-webgpu",
            Session::Cuda(_) => "fp16-cuda",
            #[cfg(target_os = "macos")]
            Session::CoreMl(coreml) => coreml.variant(),
            Session::Trt(trt) => trt.variant(),
            Session::OpenVino(openvino) => openvino.variant(),
            Session::Llama(llama) => llama.variant(),
        }
    }

    /// One batch of token ids. The worker never tokenises: the run did,
    /// once, in its prepare stage.
    pub fn embed(&mut self, batch: &[&[u32]]) -> Result<Vec<Vec<f32>>> {
        match self {
            Session::Cpu(model, _) => model.embed(batch),
            Session::Gpu(gpu) => crate::session::run(&mut gpu.session, batch)
                .map_err(|e| anyhow::anyhow!("on the GPU: {e:#}")),
            Session::Cuda(cuda) => cuda.embed(batch),
            #[cfg(target_os = "macos")]
            Session::CoreMl(coreml) => coreml.embed(batch),
            Session::Trt(trt) => trt.embed(batch),
            Session::OpenVino(openvino) => openvino.embed(batch),
            Session::Llama(llama) => llama.embed(batch),
        }
    }
}

impl GpuSession {
    fn open(dir: &Path, named: Option<&str>) -> Result<Self> {
        use ort::AsPointer;
        use ort::environment::Environment;
        let env = Environment::current().context("starting ONNX Runtime")?;
        let library = env
            .register_ep_library("webgpu", dir.join(plugin_file()))
            .context("registering the WebGPU plugin")?;
        // Kept registered for the life of the worker, which is the life of the
        // session; unregistering at exit is the process ending.
        std::mem::forget(library);

        let mut devices = Vec::new();
        let mut candidates = Vec::new();
        let mut refused = Vec::new();
        for device in env.devices() {
            if device.ep().ok() == Some("CPUExecutionProvider") {
                continue;
            }
            let hardware = device.hardware_device();
            if hardware.ty() != ort::memory::DeviceType::GPU {
                continue;
            }
            // SAFETY: the pointer is the device ORT handed out for this
            // environment, alive as long as `env`, and the call only reads it.
            let vendor_id = unsafe { (ort::api().HardwareDevice_VendorId)(hardware.ptr()) };
            let name = hardware.vendor().unwrap_or("unknown").to_string();
            if is_software(vendor_id, &name) {
                refused.push(format!("{name} (vendor {vendor_id:04x})"));
                continue;
            }
            candidates.push(Adapter {
                name,
                vendor: vendor_id,
            });
            devices.push(device);
        }
        let Some((at, why)) = choose(&candidates, named) else {
            if refused.is_empty() {
                bail!("unavailable — the WebGPU plugin found no hardware GPU");
            }
            bail!(
                "unavailable — only a software renderer was found ({}), which is never used",
                refused.join(", ")
            );
        };
        // On Linux the plugin's own device names no model (an NVIDIA card
        // comes back as ""), and the refusal of a software-only Vulkan has to
        // hold in the worker too: `doctor --gpu` and a run both start here.
        // ponytail: names the Vulkan device chosen by the same rule, which on a
        // machine with two hardware GPUs may not be the one the plugin picked.
        #[cfg(target_os = "linux")]
        let label = {
            let vulkan =
                platform::detect().map_err(|why| anyhow::anyhow!("unavailable — {why}"))?;
            format!("{vulkan} (WebGPU; {why})")
        };
        #[cfg(not(target_os = "linux"))]
        let label = format!("{} GPU (WebGPU; {why})", candidates[at].name);
        let devices = vec![devices.swap_remove(at)];
        let session = ort::session::Session::builder()
            .map_err(|e| anyhow::anyhow!("{e}"))?
            .with_devices(devices, None)
            .map_err(|e| anyhow::anyhow!("{e}"))?
            .commit_from_file(dir.join("model_fp16.onnx"))
            .map_err(|e| anyhow::anyhow!("loading the fp16 model on the GPU: {e}"))?;
        Ok(Self {
            session,
            device: label,
        })
    }
}

/// One hardware adapter the WebGPU plugin offered.
#[derive(Debug, Clone)]
pub struct Adapter {
    pub name: String,
    pub vendor: u32,
}

/// Whether a vendor's GPU is a discrete card. NVIDIA and AMD make discrete
/// cards (AMD's APUs are the exception this cannot see); Intel's, Apple's and
/// Qualcomm's are integrated.
fn discrete(vendor: u32) -> bool {
    matches!(vendor, 0x10de | 0x1002)
}

/// Which adapter the lane runs on, and why. A name the setting gives, by a
/// case-insensitive substring, wins; then a discrete GPU over an integrated
/// one, because the first adapter a machine lists is usually the one driving
/// the screen; then the first listed.
///
/// ponytail: discrete is read from the vendor. ORT's hardware-device metadata
/// carries the adapter's own answer on Windows; read it when an AMD APU turns
/// up on the wrong side of this.
pub fn choose(adapters: &[Adapter], named: Option<&str>) -> Option<(usize, String)> {
    if adapters.is_empty() {
        return None;
    }
    if let Some(named) = named.map(str::to_ascii_lowercase)
        && let Some(at) = adapters
            .iter()
            .position(|a| a.name.to_ascii_lowercase().contains(&named))
    {
        return Some((at, format!("named by the gpu_adapter setting ({named})")));
    }
    if adapters.len() == 1 {
        return Some((0, "the only hardware GPU".to_string()));
    }
    match adapters.iter().position(|a| discrete(a.vendor)) {
        Some(at) => Some((
            at,
            "a discrete GPU, preferred over an integrated one".to_string(),
        )),
        None => Some((0, "the first of several integrated GPUs".to_string())),
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn a_discrete_adapter_is_chosen_over_an_integrated_one_unless_one_is_named() {
        let adapters = vec![
            Adapter {
                name: "Intel(R) UHD Graphics 770".into(),
                vendor: 0x8086,
            },
            Adapter {
                name: "NVIDIA GeForce RTX 4070".into(),
                vendor: 0x10de,
            },
        ];
        let (at, why) = choose(&adapters, None).unwrap();
        assert_eq!(at, 1, "{why}");
        assert!(why.contains("discrete"));
        let (at, why) = choose(&adapters, Some("uhd")).unwrap();
        assert_eq!(at, 0);
        assert!(why.contains("named"));
        // A name nothing matches falls back to the rule.
        assert_eq!(choose(&adapters, Some("radeon")).unwrap().0, 1);
        assert!(choose(&[], None).is_none());
    }

    #[test]
    fn software_renderers_are_refused_by_vendor_or_name() {
        assert!(is_software(0x10005, "llvmpipe (LLVM 17.0.6, 256 bits)"));
        assert!(is_software(0x1414, "Microsoft Basic Render Driver"));
        assert!(is_software(0, "SwiftShader Device (Subzero)"));
        assert!(is_software(0x1414, "Microsoft Hyper-V Video"));
        assert!(!is_software(0x10de, "NVIDIA GeForce RTX 3060"));
        assert!(!is_software(0x106b, "Apple"));
        assert!(!is_software(0x1002, "AMD Radeon RX 7800 XT"));
    }

    fn device(kind: u32, vendor: u32, name: &str) -> VulkanDevice {
        VulkanDevice {
            vendor,
            kind,
            name: name.into(),
        }
    }

    #[test]
    fn a_vulkan_with_only_a_software_device_is_refused_by_name() {
        let llvmpipe = device(4, 0x10005, "llvmpipe (LLVM 15.0.7, 256 bits)");
        let why = hardware_vulkan(std::slice::from_ref(&llvmpipe)).unwrap_err();
        assert!(
            why.starts_with("Vulkan has only a software device (llvmpipe"),
            "{why}"
        );
        assert!(why.contains("graphics capability"));
        assert!(
            hardware_vulkan(&[])
                .unwrap_err()
                .contains("lists no device")
        );
        // A software renderer that calls itself something else is still refused.
        let disguised = device(1, 0x1af4, "Virtio-GPU Venus");
        assert!(hardware_vulkan(&[disguised]).is_err());

        let card = device(2, 0x10de, "NVIDIA RTX PRO 6000 Blackwell Server Edition");
        let igpu = device(1, 0x8086, "Intel(R) UHD Graphics 770");
        let both = [llvmpipe.clone(), igpu.clone(), card.clone()];
        assert_eq!(
            hardware_vulkan(&both).unwrap().name,
            card.name,
            "discrete first"
        );
        assert_eq!(
            hardware_vulkan(&[llvmpipe, igpu.clone()]).unwrap().name,
            igpu.name
        );
    }
}