Skip to main content

henad_explore/testing/
device.rs

1//! A headless device for tests, at the limits and features that a test requests.
2
3#[cfg(not(target_arch = "wasm32"))]
4use henad_compute::gpu::GpuContext;
5use henad_compute::gpu::GpuNeeds;
6
7/// Limits and features to request for a test device.
8#[derive(Debug, Clone)]
9#[non_exhaustive]
10pub struct TestDeviceRequest {
11    /// Needs that the baseline is raised to meet, `None` for the baseline itself.
12    #[cfg_attr(
13        target_arch = "wasm32",
14        expect(dead_code, reason = "the native headless_test_device alone reads a request")
15    )]
16    needs: Option<GpuNeeds>,
17    features: wgpu::Features,
18}
19
20impl TestDeviceRequest {
21    /// Returns a request for `wgpu::Limits::default()`, the WebGPU baseline a browser offers.
22    pub fn baseline() -> Self {
23        Self {
24            needs: None,
25            features: wgpu::Features::empty(),
26        }
27    }
28
29    /// Returns a request for the baseline raised by [`henad_compute::gpu::limits::raise`] for `needs`, as a host's
30    /// device is raised.
31    ///
32    /// Note that a browser without those limits rejects a model that needs them.
33    pub fn raised(needs: GpuNeeds) -> Self {
34        Self {
35            needs: Some(needs),
36            features: wgpu::Features::empty(),
37        }
38    }
39
40    /// Requests `features` as well. `headless_test_device` returns `None` for this request on an adapter without them.
41    ///
42    /// A caller can name the type of `features` as [`henad_compute::gpu::wgpu::Features`], without a direct `wgpu`
43    /// dependency.
44    pub fn features(mut self, features: wgpu::Features) -> Self {
45        self.features |= features;
46        self
47    }
48}
49
50/// Returns a headless device for `request`, or `None` on a machine without a device.
51///
52/// The device carries its adapter's [`henad_compute::runtime_info::RuntimeInfo`]. Returns `None` as well when the
53/// adapter lacks an optional feature that the request lists, even under `HENAD_REQUIRE_GPU`. An adapter that offers
54/// less than the WebGPU baseline yields no device for a baseline or a raised request.
55///
56/// Note that each test takes its own device. Clones of a context share its fault sink, and a check of the kit
57/// attributes every fault it finds there to itself. A device error that another test leaves in a shared sink is then
58/// dropped, or reported as that check's failure.
59///
60/// # Panics
61///
62/// Panics when `HENAD_REQUIRE_GPU` is set to anything but empty or `0` and no device is available.
63#[cfg(not(target_arch = "wasm32"))]
64pub fn headless_test_device(request: &TestDeviceRequest) -> Option<GpuContext> {
65    use henad_compute::fault::FaultSink;
66    use henad_compute::runtime_info::RuntimeInfo;
67
68    let instance = wgpu::Instance::new(wgpu::InstanceDescriptor::new_without_display_handle());
69    let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions {
70        power_preference: wgpu::PowerPreference::HighPerformance,
71        ..Default::default()
72    }));
73    let adapter = match adapter {
74        Ok(adapter) => adapter,
75        Err(error) => return unavailable(&format!("no adapter is available: {error}")),
76    };
77    // A software adapter is not required to offer optional features.
78    if !adapter.features().contains(request.features) {
79        return None;
80    }
81    let raise = |baseline: &wgpu::Limits, needs| henad_compute::gpu::limits::raise(&adapter, baseline, needs);
82    let required_limits = match requested_limits(request, &adapter.limits(), raise) {
83        Ok(limits) => limits,
84        Err(limit) => {
85            return unavailable(&format!(
86                "adapter '{}' offers less than the WebGPU baseline in {limit}",
87                adapter.get_info().name
88            ));
89        }
90    };
91    let device = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor {
92        label: Some("henad-test"),
93        required_features: request.features,
94        required_limits,
95        ..Default::default()
96    }));
97    match device {
98        Ok((device, queue)) => {
99            let runtime = RuntimeInfo::collect(&adapter, &device);
100            Some(
101                GpuContext::new(device, queue, wgpu::TextureFormat::Rgba8Unorm, FaultSink::new())
102                    .with_runtime_info(runtime),
103            )
104        }
105        Err(error) => unavailable(&format!(
106            "adapter '{}' gave no device: {error}",
107            adapter.get_info().name
108        )),
109    }
110}
111
112/// Returns the limits to request for `request` from an adapter that offers `available`. `raise` raises the baseline
113/// to the needs of a raised request.
114///
115/// The adapter is checked against the baseline before the raise, as [`crate::device::acquire_headless`] checks it.
116/// The raise clamps the storage buffer count to the adapter's count. Otherwise a test would get a device below the
117/// baseline.
118///
119/// # Errors
120///
121/// Returns the name of the first baseline limit that `available` does not meet.
122#[cfg(not(target_arch = "wasm32"))]
123fn requested_limits(
124    request: &TestDeviceRequest,
125    available: &wgpu::Limits,
126    raise: impl FnOnce(&wgpu::Limits, GpuNeeds) -> wgpu::Limits,
127) -> Result<wgpu::Limits, &'static str> {
128    crate::device::device_limits(available, |baseline| match request.needs {
129        Some(needs) => raise(baseline, needs),
130        None => baseline.clone(),
131    })
132}
133
134/// Returns `None` for a test to skip on, after checking that `HENAD_REQUIRE_GPU` allows it.
135///
136/// # Panics
137///
138/// Panics with `reason` when `HENAD_REQUIRE_GPU` is set to anything but empty or `0`.
139#[cfg(not(target_arch = "wasm32"))]
140fn unavailable(reason: &str) -> Option<GpuContext> {
141    assert!(!gpu_required(), "HENAD_REQUIRE_GPU is set, but {reason}");
142    None
143}
144
145/// Returns whether `HENAD_REQUIRE_GPU` is set to anything but empty or `0`.
146///
147/// An empty value is treated as unset. A workflow matrix can then leave the variable blank on a runner without a GPU.
148pub(super) fn gpu_required() -> bool {
149    std::env::var_os("HENAD_REQUIRE_GPU").is_some_and(|value| !value.is_empty() && value != "0")
150}
151
152#[cfg(all(test, not(target_arch = "wasm32")))]
153mod tests {
154    use henad_compute::gpu::GpuNeeds;
155
156    use super::{TestDeviceRequest, requested_limits};
157
158    /// An adapter one storage buffer short of the baseline yields no device for a baseline or a raised request. The
159    /// raise clamps the count to the adapter's 7, and a check of the raised limits would pass it.
160    #[test]
161    fn an_adapter_short_of_the_baseline_gives_no_raised_device() {
162        let available = wgpu::Limits {
163            max_storage_buffers_per_shader_stage: 7,
164            ..wgpu::Limits::default()
165        };
166        let clamp = |base: &wgpu::Limits, needs: GpuNeeds| wgpu::Limits {
167            max_storage_buffers_per_shader_stage: base
168                .max_storage_buffers_per_shader_stage
169                .max(needs.storage_buffers())
170                .min(available.max_storage_buffers_per_shader_stage),
171            ..base.clone()
172        };
173        let raised = TestDeviceRequest::raised(GpuNeeds::with_storage_buffers(8));
174        for request in [TestDeviceRequest::baseline(), raised.clone()] {
175            assert_eq!(
176                requested_limits(&request, &available, clamp).err(),
177                Some("max_storage_buffers_per_shader_stage"),
178                "{request:?}"
179            );
180        }
181        assert!(requested_limits(&raised, &wgpu::Limits::default(), clamp).is_ok());
182    }
183}