gled-core 2.32.1

//! Choosing the GPU: on a machine with several (an integrated and a discrete
//! one, or one through several drivers), gled measures each once and keeps
//! the fastest that works. The choice is cached per set of GPUs, so it is
//! measured again only when the hardware or the drivers change.

use std::{
    collections::BTreeMap,
    fs,
    sync::mpsc,
    thread,
    time::{Duration, Instant},
};

use serde::{Deserialize, Serialize};
use tracing::{info, warn};
use wgpu::{Adapter, DeviceType};

use crate::paths;

/// How long one GPU may take for the measurement before it counts as
/// broken. A driver that hangs must not keep gled from starting.
const BENCHMARK_TIMEOUT: Duration = Duration::from_secs(10);
/// Runs of the measured workload, after one run to warm up.
const BENCHMARK_RUNS: u32 = 10;
/// Workgroups of the workload per run, in two dimensions: one dimension is
/// limited to 65535 workgroups.
const BENCHMARK_GROUPS: (u32, u32) = (4096, 16);
/// Values the workload computes per run: 64 per workgroup.
const BENCHMARK_SIZE: u32 = BENCHMARK_GROUPS.0 * BENCHMARK_GROUPS.1 * 64;

/// An adapter as the cache names it.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
struct AdapterName {
    name: String,
    backend: String,
    driver: String,
}

impl AdapterName {
    fn of(adapter: &Adapter) -> Self {
        let info = adapter.get_info();
        Self {
            name: info.name,
            backend: format!("{:?}", info.backend),
            driver: format!("{} {}", info.driver, info.driver_info),
        }
    }
}

/// The choice made for each set of adapters, by the set's key.
type Cache = BTreeMap<String, AdapterName>;

/// The adapter gled renders with, among those `usable` accepts (for the
/// desktop: those that can draw into the window). `None` if none works.
pub fn choose(adapters: &[Adapter], usable: impl Fn(&Adapter) -> bool) -> Option<Adapter> {
    let mut candidates: Vec<&Adapter> = adapters.iter().filter(|adapter| usable(adapter)).collect();
    // Adapters with less than WebGPU's capabilities (OpenGL, old drivers)
    // may lack what the effects need; only if there is nothing else.
    if candidates
        .iter()
        .any(|adapter| adapter.get_downlevel_capabilities().is_webgpu_compliant())
    {
        candidates.retain(|adapter| adapter.get_downlevel_capabilities().is_webgpu_compliant());
    }
    // A software renderer only if there is nothing else.
    if candidates
        .iter()
        .any(|adapter| adapter.get_info().device_type != DeviceType::Cpu)
    {
        candidates.retain(|adapter| adapter.get_info().device_type != DeviceType::Cpu);
    }
    match candidates.as_slice() {
        [] => return None,
        [only] => return Some((*only).clone()),
        _ => {}
    }

    let key = config_key(&candidates);
    let mut cache = read_cache();
    if let Some(chosen) = cache.get(&key)
        && let Some(adapter) = candidates
            .iter()
            .find(|adapter| AdapterName::of(adapter) == *chosen)
    {
        info!(gpu = %chosen.name, backend = %chosen.backend, "using the GPU chosen before");
        return Some((*adapter).clone());
    }

    info!(
        "measuring {} GPUs to choose the fastest, once",
        candidates.len()
    );
    let fastest = candidates
        .iter()
        .filter_map(|adapter| {
            let name = AdapterName::of(adapter);
            match measure(adapter) {
                Some(time) => {
                    info!(gpu = %name.name, backend = %name.backend, ?time, "GPU measured");
                    Some((time, *adapter))
                }
                None => {
                    warn!(gpu = %name.name, backend = %name.backend, "GPU does not work, skipped");
                    None
                }
            }
        })
        .min_by_key(|(time, _)| *time)
        .map(|(_, adapter)| (*adapter).clone())?;

    let chosen = AdapterName::of(&fastest);
    info!(gpu = %chosen.name, backend = %chosen.backend, "chose the fastest GPU");
    cache.insert(key, chosen);
    write_cache(&cache);
    Some(fastest)
}

/// Names the set of adapters, so a new GPU or driver is measured again.
fn config_key(adapters: &[&Adapter]) -> String {
    let mut names: Vec<AdapterName> = adapters
        .iter()
        .map(|adapter| AdapterName::of(adapter))
        .collect();
    names.sort();
    names
        .iter()
        .map(|name| format!("{} / {} / {}", name.name, name.backend, name.driver))
        .collect::<Vec<_>>()
        .join(" | ")
}

fn read_cache() -> Cache {
    let Some(path) = paths::gpu_choice_file() else {
        return Cache::new();
    };
    fs::read(&path)
        .ok()
        .and_then(|bytes| serde_json::from_slice(&bytes).ok())
        .unwrap_or_default()
}

fn write_cache(cache: &Cache) {
    let Some(path) = paths::gpu_choice_file() else {
        return;
    };
    if let Some(dir) = path.parent()
        && let Err(err) = fs::create_dir_all(dir)
    {
        warn!("Could not create {}: {err}", dir.display());
        return;
    }
    match serde_json::to_vec_pretty(cache) {
        Ok(json) => {
            if let Err(err) = fs::write(&path, json) {
                warn!(
                    "Could not remember the GPU choice in {}: {err}",
                    path.display()
                );
            }
        }
        Err(err) => warn!("Could not encode the GPU choice: {err}"),
    }
}

/// Time for the workload, `None` if the GPU fails or takes too long. Runs on
/// a thread of its own, which is left behind if the driver hangs.
fn measure(adapter: &Adapter) -> Option<Duration> {
    let adapter = adapter.clone();
    let (result_sender, result) = mpsc::channel();
    thread::Builder::new()
        .name("gled:gpu:measure".to_string())
        .spawn(move || {
            let _ = result_sender.send(run_benchmark(&adapter));
        })
        .ok()?;
    result.recv_timeout(BENCHMARK_TIMEOUT).ok().flatten()
}

/// A compute workload like the effects': trigonometry over a large buffer,
/// then reading a part back, as the output does.
fn run_benchmark(adapter: &Adapter) -> Option<Duration> {
    let (device, queue) =
        pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor::default())).ok()?;
    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some("gpu choice benchmark"),
        source: wgpu::ShaderSource::Wgsl(BENCHMARK_SHADER.into()),
    });
    let size = u64::from(BENCHMARK_SIZE) * 4;
    let data = device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("gpu choice data"),
        size,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        mapped_at_creation: false,
    });
    let readback = device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("gpu choice readback"),
        size: 4096,
        usage: wgpu::BufferUsages::MAP_READ | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
        label: Some("gpu choice pipeline"),
        layout: None,
        module: &module,
        entry_point: Some("main"),
        compilation_options: Default::default(),
        cache: None,
    });
    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("gpu choice bind group"),
        layout: &pipeline.get_bind_group_layout(0),
        entries: &[wgpu::BindGroupEntry {
            binding: 0,
            resource: data.as_entire_binding(),
        }],
    });

    let run = || -> Option<()> {
        let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
        {
            let mut pass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
            pass.set_pipeline(&pipeline);
            pass.set_bind_group(0, &bind_group, &[]);
            pass.dispatch_workgroups(BENCHMARK_GROUPS.0, BENCHMARK_GROUPS.1, 1);
        }
        encoder.copy_buffer_to_buffer(&data, 0, &readback, 0, 4096);
        queue.submit([encoder.finish()]);
        let (mapped_sender, mapped) = mpsc::channel();
        readback.map_async(wgpu::MapMode::Read, .., move |result| {
            let _ = mapped_sender.send(result.is_ok());
        });
        device.poll(wgpu::PollType::wait_indefinitely()).ok()?;
        let ok = mapped.recv().ok()?;
        readback.unmap();
        ok.then_some(())
    };

    run()?;
    let start = Instant::now();
    for _ in 0..BENCHMARK_RUNS {
        run()?;
    }
    Some(start.elapsed())
}

const BENCHMARK_SHADER: &str = r#"
@group(0) @binding(0) var<storage, read_write> data: array<f32>;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) id: vec3<u32>) {
    let i = id.y * 4096u * 64u + id.x;
    if (i >= arrayLength(&data)) {
        return;
    }
    var value = f32(i) * 0.001;
    for (var step = 0u; step < 256u; step++) {
        value = sin(value) * 1.7 + cos(value * 0.5);
    }
    data[i] = value;
}
"#;

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

    #[test]
    #[ignore = "needs a GPU"]
    fn the_workload_runs_on_the_local_gpus() {
        let instance =
            wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle_from_env());
        let adapters = pollster::block_on(instance.enumerate_adapters(wgpu::Backends::all()));
        assert!(!adapters.is_empty());
        for adapter in &adapters {
            let time = measure(adapter);
            println!("{:?}: {time:?}", AdapterName::of(adapter));
        }
        assert!(adapters.iter().any(|adapter| measure(adapter).is_some()));
    }
}