rightkit_gpu/
budget.rs

1//! Device memory shared by concurrent GPU jobs (upstream issue #1764). Adapted from Ember `crates/gpu`.
2//!
3//! Each job asks for room for what it is expected to hold on the device and waits while the
4//! renders in flight leave too little; a render always gets in when nothing else is in flight,
5//! however large it is (the device's own limits still apply). wgpu reports no memory size, so the
6//! budget starts from the largest buffer the device allows and adapts to the device it meets:
7//! it is halved when a render runs out of device memory and grows back a step with each render
8//! that succeeds (the additive-increase / multiplicative-decrease rule of TCP congestion control).
9//! A 2 GB card settles near what it holds instead of failing every render past its memory.
10//! Renders get in in the order they asked (a fair semaphore): a large preview is not passed over
11//! for ever by a stream of thumbnails that each fit.
12
13use std::collections::VecDeque;
14use std::sync::{Condvar, Mutex, PoisonError};
15use std::time::Duration;
16
17/// The budget never shrinks below this: one render always fits anyway, so this only bounds how
18/// far a run of failures can push it.
19const FLOOR: u64 = 256 << 20;
20
21/// Device memory a render is expected to hold: the source, two working copies and the output,
22/// as float RGBA, at the source's size.
23pub fn footprint(width: usize, height: usize) -> u64 {
24    (width as u64)
25        .saturating_mul(height as u64)
26        .saturating_mul(16 * 4)
27}
28
29/// Where the budget starts: 1.5 × the largest buffer the device allows (about the device's
30/// memory on small cards), at least 1 GiB.
31pub fn initial(max_buffer: u64) -> u64 {
32    (max_buffer / 2).saturating_mul(3).max(1 << 30)
33}
34
35struct State {
36    /// Bytes held by renders in flight.
37    used: u64,
38    /// Renders in flight.
39    running: u32,
40    /// A retry holds the device exclusively until its permit drops.
41    exclusive: bool,
42    /// Bytes renders may hold together now.
43    limit: u64,
44    /// What the budget grows back to.
45    ceiling: u64,
46    /// Tickets of the renders waiting, in the order they asked.
47    queue: VecDeque<u64>,
48    next_ticket: u64,
49}
50
51pub struct MemoryBudget {
52    state: Mutex<State>,
53    room: Condvar,
54}
55
56/// Room on the device for one render, released when dropped.
57pub struct Permit<'a> {
58    budget: &'a MemoryBudget,
59    bytes: u64,
60    alone: bool,
61}
62
63impl MemoryBudget {
64    pub fn new(ceiling: u64) -> MemoryBudget {
65        MemoryBudget {
66            state: Mutex::new(State {
67                used: 0,
68                running: 0,
69                exclusive: false,
70                limit: ceiling,
71                ceiling,
72                queue: VecDeque::new(),
73                next_ticket: 0,
74            }),
75            room: Condvar::new(),
76        }
77    }
78
79    /// Wait for room for `bytes`; `alone` holds the device exclusively (the retry of a
80    /// render that ran out of memory). `None` when `stop` turns true while waiting (the process is
81    /// ending): the caller renders on the CPU.
82    pub fn admit(&self, bytes: u64, alone: bool, stop: impl Fn() -> bool) -> Option<Permit<'_>> {
83        let mut st = self.state.lock().unwrap_or_else(PoisonError::into_inner);
84        let ticket = st.next_ticket;
85        st.next_ticket = st.next_ticket.wrapping_add(1);
86        st.queue.push_back(ticket);
87        while st.queue.front() != Some(&ticket)
88            || st.exclusive
89            || (st.running > 0 && (alone || st.used.saturating_add(bytes) > st.limit))
90        {
91            if stop() {
92                st.queue.retain(|t| *t != ticket);
93                // the render behind may be the next to fit
94                self.room.notify_all();
95                return None;
96            }
97            st = self
98                .room
99                .wait_timeout(st, Duration::from_millis(100))
100                .unwrap_or_else(PoisonError::into_inner)
101                .0;
102        }
103        st.queue.pop_front();
104        // the next in line may fit beside this one
105        self.room.notify_all();
106        st.used = st.used.saturating_add(bytes);
107        st.running += 1;
108        st.exclusive = alone;
109        Some(Permit {
110            budget: self,
111            bytes,
112            alone,
113        })
114    }
115
116    /// A render ran out of device memory: halve the budget (not below what one render needs).
117    pub fn out_of_memory(&self) {
118        let mut st = self.state.lock().unwrap_or_else(PoisonError::into_inner);
119        st.limit = (st.limit / 2).max(FLOOR);
120    }
121
122    /// A render succeeded: grow the budget back by a sixteenth of its ceiling.
123    pub fn succeeded(&self) {
124        let mut st = self.state.lock().unwrap_or_else(PoisonError::into_inner);
125        st.limit = st.limit.saturating_add(st.ceiling / 16).min(st.ceiling);
126    }
127}
128
129impl Drop for Permit<'_> {
130    fn drop(&mut self) {
131        let mut st = self
132            .budget
133            .state
134            .lock()
135            .unwrap_or_else(PoisonError::into_inner);
136        st.used = st.used.saturating_sub(self.bytes);
137        st.running = st.running.saturating_sub(1);
138        if self.alone {
139            st.exclusive = false;
140        }
141        self.budget.room.notify_all();
142    }
143}
rightkit_gpu/budget.rs

rightkit_gpu/
budget.rs
