alpha2_gpu/
watchdog.rs

1//! How long one submission may keep the GPU busy before the OS may reset the device.
2//!
3//! Windows TDR resets a GPU after two seconds by default, macOS after a similar time, and
4//! browsers report either as device loss. A loss right after a submission that ran past
5//! [`hung_threshold_ms`] is attributed to that submission. [`submission_budget_ms`] is the target
6//! per submission.
7
8use std::sync::atomic::{AtomicU32, Ordering};
9
10/// The assumed watchdog in milliseconds: Windows' default TDR, the shortest of the platforms.
11pub const WATCHDOG_MS: u32 = 2_000;
12
13/// Zero when unset.
14static OVERRIDE_MS: AtomicU32 = AtomicU32::new(0);
15
16/// The watchdog in milliseconds: [`set_watchdog_ms`], else a positive `ALPHA2_GPU_WATCHDOG_MS`
17/// (native only), else [`WATCHDOG_MS`].
18#[must_use]
19pub fn watchdog_ms() -> u32 {
20    match OVERRIDE_MS.load(Ordering::Relaxed) {
21        0 => environment_watchdog_ms().unwrap_or(WATCHDOG_MS),
22        set => set,
23    }
24}
25
26#[cfg(not(target_arch = "wasm32"))]
27fn environment_watchdog_ms() -> Option<u32> {
28    std::env::var("ALPHA2_GPU_WATCHDOG_MS")
29        .ok()?
30        .trim()
31        .parse::<u32>()
32        .ok()
33        .filter(|value| *value > 0)
34}
35
36#[cfg(target_arch = "wasm32")]
37const fn environment_watchdog_ms() -> Option<u32> {
38    None
39}
40
41/// Overrides the watchdog for this process; zero restores the default. Lets browser tests,
42/// which have no environment variables, shorten it.
43pub fn set_watchdog_ms(value: u32) {
44    OVERRIDE_MS.store(value, Ordering::Relaxed);
45}
46
47/// Half the watchdog. A submission this long before a device loss is taken as its cause.
48#[must_use]
49pub fn hung_threshold_ms() -> u32 {
50    watchdog_ms() / 2
51}
52
53/// An eighth of the watchdog per submission, leaving room for other work on the GPU.
54#[must_use]
55pub fn submission_budget_ms() -> u32 {
56    watchdog_ms() / 8
57}
58
59#[cfg(test)]
60mod tests {
61    #[test]
62    fn the_threshold_and_budget_are_fractions_of_the_watchdog() {
63        super::set_watchdog_ms(4_000);
64        assert_eq!(super::watchdog_ms(), 4_000);
65        assert_eq!(super::hung_threshold_ms(), 2_000);
66        assert_eq!(super::submission_budget_ms(), 500);
67        super::set_watchdog_ms(0);
68        assert!(super::watchdog_ms() > 0);
69    }
70}
alpha2_gpu/watchdog.rs

alpha2_gpu/
watchdog.rs
