open-standards/wai/wai-ops-gpu/src/device.rs

//! Adapters and devices: discovery, the pipelines, error scopes, and waiting.
//!
//! [`discover`] opens every adapter of the platform's primary GPU interfaces (Metal,
//! Vulkan, DX12; a browser's WebGPU on `wasm32-unknown-unknown`), compiles the kernels
//! on it and runs the arithmetic probe ([`crate::probe`]) and WGSL's integer semantics
//! ([`crate::semantics`]) against the host. An adapter that cannot host the kernels,
//! whose arithmetic disagrees with the host's on any word, or that gives another result
//! than WGSL's for an operation the kernels may use, is listed with the reason and never
//! runs a graph.
//!
//! Every module is created with `create_shader_module`, which keeps the translator's
//! run-time checks on: bounds checks, and the guarded integer division and remainder
//! (`wgpu::ShaderRuntimeChecks::int_div_checks`). The unchecked form is never used
//! (`tests/wgsl_rules.rs` holds the crate to that).

use std::borrow::Cow;
use std::collections::HashMap;
use std::num::NonZeroU64;
use std::sync::{Arc, Mutex, MutexGuard};

use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll, Waker};

use crate::kernels::{OP_ENTRIES, REC_WORDS, Sources};

/// How long a native wait on one submission may take.
pub const POLL_TIMEOUT: core::time::Duration = core::time::Duration::from_secs(120);

/// Why the device did not give a result. A graph or a decode that this executor does
/// not run is [`GpuError::Fallback`]; everything else is the device's.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum GpuError {
    /// The graph or the decode is outside what this executor runs (`crate::Fallback`):
    /// the caller runs it on the CPU and says so.
    Fallback(crate::Fallback),
    /// The device reported an error, was lost, or a wait failed.
    Device(String),
    /// A native wait did not finish within [`POLL_TIMEOUT`].
    Timeout,
    /// A native future was still pending after one poll.
    Pending,
}

impl core::fmt::Display for GpuError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            GpuError::Fallback(r) => write!(f, "runs on the CPU: {r}"),
            GpuError::Device(m) => write!(f, "device: {m}"),
            GpuError::Timeout => write!(f, "a wait on the device timed out"),
            GpuError::Pending => write!(f, "a native future was pending after one poll"),
        }
    }
}

impl std::error::Error for GpuError {}

/// What the platform reports about an adapter, as a run records it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Adapter {
    /// `metal`, `vulkan`, `dx12` or `webgpu`.
    pub backend: String,
    pub name: String,
    pub vendor: u32,
    pub device: u32,
    /// `discrete`, `integrated`, `virtual`, `cpu` or `other`.
    pub device_type: String,
    pub driver: String,
    pub driver_info: String,
    /// The granted storage-binding and buffer limits, in bytes.
    pub max_storage_binding: u64,
    pub max_buffer: u64,
}

impl Adapter {
    /// `gpu-<backend>-<vendor>` or, for a software adapter, `swgpu-<backend>-<vendor>`,
    /// the vendor as lowercase hexadecimal.
    pub fn key(&self) -> String {
        let class = if self.device_type == "cpu" { "swgpu" } else { "gpu" };
        format!("{class}-{}-{:x}", self.backend, self.vendor)
    }

    /// Whether the adapter is a GPU: not a software rasteriser.
    pub fn is_gpu(&self) -> bool {
        self.device_type != "cpu"
    }

    /// One line for a report: backend, adapter, type, driver.
    pub fn describe(&self) -> String {
        format!(
            "backend={} adapter={:?} type={} vendor={:#x} device={:#x} driver={:?} driver_info={:?}",
            self.backend, self.name, self.device_type, self.vendor, self.device, self.driver, self.driver_info
        )
    }
}

fn backend_label(b: wgpu::Backend) -> Option<&'static str> {
    match b {
        wgpu::Backend::Metal => Some("metal"),
        wgpu::Backend::Vulkan => Some("vulkan"),
        wgpu::Backend::Dx12 => Some("dx12"),
        wgpu::Backend::BrowserWebGpu => Some("webgpu"),
        _ => None,
    }
}

fn device_type_label(t: wgpu::DeviceType) -> &'static str {
    match t {
        wgpu::DeviceType::DiscreteGpu => "discrete",
        wgpu::DeviceType::IntegratedGpu => "integrated",
        wgpu::DeviceType::VirtualGpu => "virtual",
        wgpu::DeviceType::Cpu => "cpu",
        wgpu::DeviceType::Other => "other",
    }
}

#[derive(Default)]
pub(crate) struct Sink {
    messages: Vec<String>,
    lost: Option<String>,
}

pub(crate) fn lock<T>(m: &Mutex<T>) -> MutexGuard<'_, T> {
    m.lock().unwrap_or_else(|p| p.into_inner())
}

/// The compiled kernels on one device.
pub(crate) struct Pipelines {
    pub layout: wgpu::BindGroupLayout,
    /// `ops32`'s entry points, in [`OP_ENTRIES`] order.
    pub ops32: Vec<wgpu::ComputePipeline>,
    /// `ops64`'s.
    pub ops64: Vec<wgpu::ComputePipeline>,
    pub mv: wgpu::ComputePipeline,
    pub probe: wgpu::ComputePipeline,
    pub semantics: wgpu::ComputePipeline,
}

/// One adapter opened as a device, its kernels compiled and its arithmetic probed.
pub struct Gpu {
    pub(crate) device: wgpu::Device,
    pub(crate) queue: wgpu::Queue,
    pub(crate) pipes: Pipelines,
    pub(crate) adapter: Adapter,
    pub(crate) stride: u64,
    pub(crate) dummy: wgpu::Buffer,
    pub(crate) sink: Arc<Mutex<Sink>>,
    pub(crate) cache: Mutex<Vec<crate::prepare::Cached>>,
    pub(crate) pool: Mutex<HashMap<u64, Vec<wgpu::Buffer>>>,
    pub(crate) label: String,
}

impl core::fmt::Debug for Gpu {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("Gpu").field("adapter", &self.adapter).field("kernels", &self.label).finish()
    }
}

/// What discovery found: the adapters that passed, and why each other one did not.
#[derive(Debug, Default)]
pub struct Discovery {
    pub found: Vec<Gpu>,
    pub refused: Vec<(Adapter, String)>,
}

fn instance() -> wgpu::Instance {
    let mut d = wgpu::InstanceDescriptor::new_without_display_handle();
    d.backends = wgpu::Backends::PRIMARY;
    d.flags = wgpu::InstanceFlags::empty();
    d.backend_options.dx12.shader_compiler = wgpu::Dx12Compiler::Fxc;
    wgpu::Instance::new(d)
}

/// Every adapter the platform offers, each opened and probed in this process.
pub async fn discover() -> Discovery {
    discover_with(&Sources::shipped(), "shipped").await
}

/// [`discover`], with the kernels compiled from `src` and labelled `label` (a test's
/// mutant). The probe runs `src.probe`, and the semantics suite `src.semantics`.
pub async fn discover_with(src: &Sources, label: &str) -> Discovery {
    let mut d = Discovery::default();
    if wgpu::Instance::enabled_backend_features().is_empty() {
        return d;
    }
    let inst = instance();
    #[cfg(not(target_arch = "wasm32"))]
    let adapters = inst.enumerate_adapters(wgpu::Backends::PRIMARY).await;
    #[cfg(target_arch = "wasm32")]
    let adapters: Vec<wgpu::Adapter> = inst.request_adapter(&wgpu::RequestAdapterOptions::default()).await.into_iter().collect();
    for a in adapters {
        match open(a, src, label).await {
            Ok(g) => d.found.push(g),
            Err(e) => d.refused.push(*e),
        }
    }
    d
}

fn identity(info: &wgpu::AdapterInfo, limits: &wgpu::Limits) -> Adapter {
    Adapter {
        backend: backend_label(info.backend).map_or_else(|| format!("{:?}", info.backend).to_lowercase(), str::to_string),
        name: info.name.clone(),
        vendor: info.vendor,
        device: info.device,
        device_type: device_type_label(info.device_type).into(),
        driver: info.driver.clone(),
        driver_info: info.driver_info.clone(),
        max_storage_binding: limits.max_storage_buffer_binding_size,
        max_buffer: limits.max_buffer_size,
    }
}

async fn open(adapter: wgpu::Adapter, src: &Sources, label: &str) -> Result<Gpu, Box<(Adapter, String)>> {
    let info = adapter.get_info();
    let offered = adapter.limits();
    let mut who = identity(&info, &offered);
    if backend_label(info.backend).is_none() {
        return Err(Box::new((who, format!("backend {:?} is not one the kernels run on", info.backend))));
    }
    if !adapter.get_downlevel_capabilities().flags.contains(wgpu::DownlevelFlags::COMPUTE_SHADERS) {
        return Err(Box::new((who, "no compute shaders".into())));
    }
    // The WebGPU default limits (eight storage buffers per stage; the kernels bind five),
    // with the two buffer limits raised to the adapter's.
    let mut short: Option<String> = None;
    wgpu::Limits::default().check_limits_with_fail_fn(&offered, true, |name, needed, has| {
        short.get_or_insert_with(|| format!("limit {name}: the kernels need {needed}, the adapter offers {has}"));
    });
    if let Some(why) = short {
        return Err(Box::new((who, why)));
    }
    let required_limits = wgpu::Limits {
        max_storage_buffer_binding_size: offered.max_storage_buffer_binding_size,
        max_buffer_size: offered.max_buffer_size,
        ..wgpu::Limits::default()
    };
    let descriptor = wgpu::DeviceDescriptor { label: Some("wai-ops-gpu"), required_features: wgpu::Features::empty(), required_limits, ..Default::default() };
    let (device, queue) = match adapter.request_device(&descriptor).await {
        Ok(dq) => dq,
        Err(e) => return Err(Box::new((who, format!("request_device: {e}")))),
    };
    let sink = Arc::new(Mutex::new(Sink::default()));
    let errors = Arc::clone(&sink);
    device.on_uncaptured_error(Arc::new(move |e: wgpu::Error| lock(&errors).messages.push(e.to_string())));
    let lost = Arc::clone(&sink);
    device.set_device_lost_callback(move |reason, message| {
        let mut s = lock(&lost);
        if s.lost.is_none() {
            s.lost = Some(format!("device lost ({reason:?}): {message}"));
        }
    });
    let pipes = match pipelines(&device, src).await {
        Ok(p) => p,
        Err(e) => return Err(Box::new((who, format!("the kernels did not compile: {e}")))),
    };
    let granted = device.limits();
    who.max_storage_binding = granted.max_storage_buffer_binding_size;
    who.max_buffer = granted.max_buffer_size;
    let align = u64::from(granted.min_uniform_buffer_offset_alignment).max(4);
    let stride = ((REC_WORDS as u64 * 4).div_ceil(align)) * align;
    let dummy = device.create_buffer(&wgpu::BufferDescriptor { label: Some("wai-ops-gpu dummy"), size: 16, usage: wgpu::BufferUsages::STORAGE, mapped_at_creation: false });
    let gpu = Gpu {
        device,
        queue,
        pipes,
        adapter: who,
        stride,
        dummy,
        sink,
        cache: Mutex::new(Vec::new()),
        pool: Mutex::new(HashMap::new()),
        label: label.to_string(),
    };
    if let Err(why) = crate::probe::check(&gpu).await {
        return Err(Box::new((gpu.adapter.clone(), why)));
    }
    match crate::semantics::check(&gpu).await {
        Ok(()) => Ok(gpu),
        Err(why) => Err(Box::new((gpu.adapter.clone(), why))),
    }
}

/// The three error scopes a step of work runs inside, popped innermost first.
pub(crate) struct Scopes {
    internal: wgpu::ErrorScopeGuard,
    validation: wgpu::ErrorScopeGuard,
    memory: wgpu::ErrorScopeGuard,
}

impl Scopes {
    pub(crate) fn push(device: &wgpu::Device) -> Scopes {
        let memory = device.push_error_scope(wgpu::ErrorFilter::OutOfMemory);
        let validation = device.push_error_scope(wgpu::ErrorFilter::Validation);
        let internal = device.push_error_scope(wgpu::ErrorFilter::Internal);
        Scopes { internal, validation, memory }
    }

    pub(crate) async fn close(self) -> Result<(), GpuError> {
        let Scopes { internal, validation, memory } = self;
        let internal = internal.pop();
        let validation = validation.pop();
        let memory = memory.pop();
        let (internal, validation, memory) = (internal.await, validation.await, memory.await);
        match memory.or(validation).or(internal) {
            Some(e) => Err(GpuError::Device(e.to_string())),
            None => Ok(()),
        }
    }
}

async fn pipelines(device: &wgpu::Device, src: &Sources) -> Result<Pipelines, GpuError> {
    let storage = |binding: u32, read_only: bool| wgpu::BindGroupLayoutEntry {
        binding,
        visibility: wgpu::ShaderStages::COMPUTE,
        ty: wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Storage { read_only }, has_dynamic_offset: false, min_binding_size: None },
        count: None,
    };
    let scopes = Scopes::push(device);
    let layout = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        label: Some("wai-ops-gpu"),
        entries: &[
            wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::COMPUTE,
                ty: wgpu::BindingType::Buffer { ty: wgpu::BufferBindingType::Uniform, has_dynamic_offset: false, min_binding_size: NonZeroU64::new(REC_WORDS as u64 * 4) },
                count: None,
            },
            storage(1, true),
            storage(2, true),
            storage(3, true),
            storage(4, false),
            storage(5, false),
        ],
    });
    let pl = device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor { label: Some("wai-ops-gpu"), bind_group_layouts: &[Some(&layout)], immediate_size: 0 });
    let module = |label: &str, text: &str| device.create_shader_module(wgpu::ShaderModuleDescriptor { label: Some(label), source: wgpu::ShaderSource::Wgsl(Cow::Owned(text.to_string())) });
    let m32 = module("wai-ops-gpu ops32", &src.ops32);
    let m64 = module("wai-ops-gpu ops64", &src.ops64);
    let mmv = module("wai-ops-gpu move", &src.mv);
    let mpr = module("wai-ops-gpu probe", &src.probe);
    let msem = module("wai-ops-gpu semantics", &src.semantics);
    let pipe = |m: &wgpu::ShaderModule, entry: &str| {
        device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
            label: Some(entry),
            layout: Some(&pl),
            module: m,
            entry_point: Some(entry),
            compilation_options: Default::default(),
            cache: None,
        })
    };
    let ops32 = OP_ENTRIES.iter().map(|e| pipe(&m32, e)).collect();
    let ops64 = OP_ENTRIES.iter().map(|e| pipe(&m64, e)).collect();
    let mv = pipe(&mmv, "mv");
    let probe = pipe(&mpr, "probe");
    let semantics = pipe(&msem, "semantics");
    scopes.close().await?;
    Ok(Pipelines { layout, ops32, ops64, mv, probe, semantics })
}

impl Gpu {
    /// The adapter, as the platform reports it.
    pub fn adapter(&self) -> &Adapter {
        &self.adapter
    }

    /// The label of the kernels this device runs: `shipped`, or a test's mutant.
    pub fn kernels(&self) -> &str {
        &self.label
    }

    /// An error if the device has been lost, or reported anything outside a scope since
    /// the last check.
    pub(crate) fn settled(&self) -> Result<(), GpuError> {
        let mut s = lock(&self.sink);
        if let Some(why) = &s.lost {
            return Err(GpuError::Device(why.clone()));
        }
        if s.messages.is_empty() {
            return Ok(());
        }
        Err(GpuError::Device(std::mem::take(&mut s.messages).join("; ")))
    }

    pub(crate) fn buffer(&self, label: &str, size: u64, usage: wgpu::BufferUsages) -> wgpu::Buffer {
        self.device.create_buffer(&wgpu::BufferDescriptor { label: Some(label), size, usage, mapped_at_creation: false })
    }

    /// A storage buffer of at least `size` bytes from the pool, or a new one.
    pub(crate) fn take(&self, size: u64) -> wgpu::Buffer {
        let size = size.div_ceil(256).max(1) * 256;
        if let Some(b) = lock(&self.pool).get_mut(&size).and_then(Vec::pop) {
            return b;
        }
        self.buffer("wai-ops-gpu tensor", size, wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST)
    }

    /// Give buffers back to the pool, keeping at most 64 of each size.
    pub(crate) fn give(&self, bufs: Vec<wgpu::Buffer>) {
        let mut p = lock(&self.pool);
        for b in bufs {
            let v = p.entry(b.size()).or_default();
            if v.len() < 64 {
                v.push(b);
            }
        }
    }

    /// Read `len` bytes of `staging` (MAP_READ) once submission `index` has run.
    pub(crate) async fn read(&self, staging: &wgpu::Buffer, len: u64, _index: wgpu::SubmissionIndex) -> Result<Vec<u8>, GpuError> {
        let slot = Arc::new(Mutex::new(MapSlot::default()));
        let filled = Arc::clone(&slot);
        staging.map_async(wgpu::MapMode::Read, 0..len, move |r| {
            let waker = {
                let mut s = lock(&filled);
                s.result = Some(r);
                s.waker.take()
            };
            if let Some(w) = waker {
                w.wake();
            }
        });
        Mapped {
            slot,
            #[cfg(not(target_arch = "wasm32"))]
            device: self.device.clone(),
            #[cfg(not(target_arch = "wasm32"))]
            sink: Arc::clone(&self.sink),
            #[cfg(not(target_arch = "wasm32"))]
            index: _index,
        }
        .await?;
        let out = {
            let view = staging.get_mapped_range(0..len).map_err(|e| GpuError::Device(e.to_string()))?;
            view.to_vec()
        };
        staging.unmap();
        Ok(out)
    }
}

#[derive(Default)]
struct MapSlot {
    result: Option<Result<(), wgpu::BufferAsyncError>>,
    waker: Option<Waker>,
}

/// Resolves when `map_async`'s callback has run. Natively it waits on its submission
/// inside the device poll, so it is ready on the first poll; in a browser the event
/// loop runs the callback, which wakes it.
struct Mapped {
    slot: Arc<Mutex<MapSlot>>,
    #[cfg(not(target_arch = "wasm32"))]
    device: wgpu::Device,
    #[cfg(not(target_arch = "wasm32"))]
    sink: Arc<Mutex<Sink>>,
    #[cfg(not(target_arch = "wasm32"))]
    index: wgpu::SubmissionIndex,
}

impl Future for Mapped {
    type Output = Result<(), GpuError>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        #[cfg(not(target_arch = "wasm32"))]
        if let Err(e) = wait(&self.device, &self.sink, self.index.clone()) {
            return Poll::Ready(Err(e));
        }
        let mut s = lock(&self.slot);
        match s.result.take() {
            Some(r) => Poll::Ready(r.map_err(|e| GpuError::Device(e.to_string()))),
            None => {
                s.waker = Some(cx.waker().clone());
                Poll::Pending
            }
        }
    }
}

#[cfg(not(target_arch = "wasm32"))]
fn wait(device: &wgpu::Device, sink: &Mutex<Sink>, index: wgpu::SubmissionIndex) -> Result<(), GpuError> {
    let poll = || device.poll(wgpu::PollType::Wait { submission_index: Some(index), timeout: Some(POLL_TIMEOUT) });
    match std::panic::catch_unwind(std::panic::AssertUnwindSafe(poll)) {
        Ok(Ok(_)) => Ok(()),
        Ok(Err(wgpu::PollError::Timeout)) => Err(GpuError::Timeout),
        Ok(Err(e)) => Err(GpuError::Device(e.to_string())),
        Err(payload) => {
            let message = payload.downcast_ref::<String>().cloned().or_else(|| payload.downcast_ref::<&str>().map(|m| m.to_string())).unwrap_or_else(|| "no message".into());
            let why = format!("the device poll failed: {message}");
            let mut s = lock(sink);
            if s.lost.is_none() {
                s.lost = Some(why.clone());
            }
            Err(GpuError::Device(why))
        }
    }
}

/// The calls for a native caller without an async runtime. Natively every future of
/// the GPU interface is already resolved and a buffer map completes inside the device
/// poll that waits for it, so one poll resolves each call.
#[cfg(not(target_arch = "wasm32"))]
pub mod blocking {
    use super::GpuError;
    use core::future::Future;
    use core::task::{Context, Poll, Waker};

    /// Polls `f` once with a waker that does nothing; a future still pending is
    /// [`GpuError::Pending`], never a spin.
    pub fn resolve<F: Future>(f: F) -> Result<F::Output, GpuError> {
        let mut f = core::pin::pin!(f);
        match f.as_mut().poll(&mut Context::from_waker(Waker::noop())) {
            Poll::Ready(v) => Ok(v),
            Poll::Pending => Err(GpuError::Pending),
        }
    }

    /// [`super::discover`].
    pub fn discover() -> super::Discovery {
        resolve(super::discover()).unwrap_or_default()
    }

    /// [`super::discover_with`].
    pub fn discover_with(src: &crate::kernels::Sources, label: &str) -> super::Discovery {
        resolve(super::discover_with(src, label)).unwrap_or_default()
    }

    /// [`crate::Gpu::run`].
    pub fn run(
        g: &crate::Gpu, l: &wai_ops::Loaded, inputs: &[wai_ops::Tensor], scalars: &[i64], seq: &mut crate::GpuSequence,
    ) -> Result<Result<Vec<wai_ops::Tensor>, wai_ops::Refusal>, GpuError> {
        resolve(g.run(l, inputs, scalars, seq))?
    }

    /// [`crate::Gpu::run_at`].
    pub fn run_at(
        g: &crate::Gpu, l: &wai_ops::Loaded, width: crate::Width, inputs: &[wai_ops::Tensor], scalars: &[i64], seq: &mut crate::GpuSequence,
    ) -> Result<Result<Vec<wai_ops::Tensor>, wai_ops::Refusal>, GpuError> {
        resolve(g.run_at(l, width, inputs, scalars, seq))?
    }

    /// [`crate::probe::run`].
    pub fn probe(g: &crate::Gpu, recs: &[crate::probe::ProbeIn]) -> Result<Vec<[u32; crate::probe::OUT_WORDS]>, GpuError> {
        resolve(crate::probe::run(g, recs))?
    }

    /// [`crate::semantics::run`].
    pub fn semantics(g: &crate::Gpu, cases: &[crate::semantics::Case]) -> Result<Vec<u32>, GpuError> {
        resolve(crate::semantics::run(g, cases))?
    }
}