pito_codecs/vp9/enc/rc/
threads.rs

1use std::sync::atomic::{AtomicUsize, Ordering};
2
3pub const MIN_TILE_WIDTH_B64: usize = 4;
4pub const MAX_TILE_WIDTH_B64: usize = 64;
5pub const LAG: usize = 2;
6
7#[derive(Clone, Debug, PartialEq, Eq)]
8pub struct TileLayout {
9    pub log2_cols: u8,
10    pub log2_rows: u8,
11    pub mi_cols: usize,
12    pub mi_rows: usize,
13    pub cols: Vec<(usize, usize)>,
14    pub rows: Vec<(usize, usize)>,
15}
16
17pub fn min_log2_tile_cols(sb_cols: usize) -> u8 {
18    let mut l = 0;
19    while (MAX_TILE_WIDTH_B64 << l) < sb_cols {
20        l += 1;
21    }
22    l
23}
24
25pub fn max_log2_tile_cols(sb_cols: usize) -> u8 {
26    let mut l = 1;
27    while (sb_cols >> l) >= MIN_TILE_WIDTH_B64 {
28        l += 1;
29    }
30    l - 1
31}
32
33pub fn tile_offset(i: usize, mis: usize, log2: u8) -> usize {
34    let sbs = mis.div_ceil(8);
35    (((i * sbs) >> log2) << 3).min(mis)
36}
37
38impl TileLayout {
39    pub fn new(mi_cols: usize, mi_rows: usize, threads: usize) -> TileLayout {
40        let want = threads.max(1).next_power_of_two().trailing_zeros() as u8;
41        TileLayout::with_log2(mi_cols, mi_rows, want, 0)
42    }
43
44    pub fn with_log2(mi_cols: usize, mi_rows: usize, log2_cols: u8, log2_rows: u8) -> TileLayout {
45        let sb_cols = mi_cols.div_ceil(8);
46        let lo = min_log2_tile_cols(sb_cols);
47        let hi = max_log2_tile_cols(sb_cols).max(lo);
48        let log2_cols = log2_cols.clamp(lo, hi);
49        let log2_rows = log2_rows.min(2);
50        let span = |n: usize, mis: usize, log2: u8| {
51            (0..n)
52                .map(|i| (tile_offset(i, mis, log2), tile_offset(i + 1, mis, log2)))
53                .collect::<Vec<_>>()
54        };
55        TileLayout {
56            log2_cols,
57            log2_rows,
58            mi_cols,
59            mi_rows,
60            cols: span(1 << log2_cols, mi_cols, log2_cols),
61            rows: span(1 << log2_rows, mi_rows, log2_rows),
62        }
63    }
64
65    pub fn sb_cols(&self) -> usize {
66        self.mi_cols.div_ceil(8)
67    }
68
69    pub fn sb_rows(&self) -> usize {
70        self.mi_rows.div_ceil(8)
71    }
72
73    pub fn tile_row_of(&self, sb_row: usize) -> usize {
74        let mi = sb_row * 8;
75        self.rows.iter().position(|&(_, e)| mi < e).unwrap_or(0)
76    }
77}
78
79pub struct Wavefront {
80    layout: TileLayout,
81    progress: Vec<AtomicUsize>,
82    next: AtomicUsize,
83}
84
85pub struct Row<'a> {
86    pub tile_col: usize,
87    pub tile_row: usize,
88    pub sb_row: usize,
89    pub sb_start: usize,
90    pub sb_end: usize,
91    above: Option<&'a AtomicUsize>,
92    me: &'a AtomicUsize,
93}
94
95impl Row<'_> {
96    pub fn wait(&self, sb_col: usize) {
97        let Some(above) = self.above else {
98            return;
99        };
100        let need = (sb_col + LAG).min(self.sb_end) - self.sb_start;
101        let mut spins = 0u32;
102        while above.load(Ordering::Acquire) < need {
103            if spins < 256 {
104                std::hint::spin_loop();
105                spins += 1;
106            } else {
107                std::thread::yield_now();
108            }
109        }
110    }
111
112    pub fn publish(&self, sb_col: usize) {
113        self.me.store(sb_col + 1 - self.sb_start, Ordering::Release);
114    }
115
116    pub fn sbs(&self) -> Sbs<'_, '_> {
117        Sbs {
118            row: self,
119            at: self.sb_start,
120        }
121    }
122}
123
124pub struct Sbs<'r, 'a> {
125    row: &'r Row<'a>,
126    at: usize,
127}
128
129impl Iterator for Sbs<'_, '_> {
130    type Item = usize;
131
132    fn next(&mut self) -> Option<usize> {
133        if self.at > self.row.sb_start {
134            self.row.publish(self.at - 1);
135        }
136        if self.at >= self.row.sb_end {
137            return None;
138        }
139        self.row.wait(self.at);
140        self.at += 1;
141        Some(self.at - 1)
142    }
143}
144
145struct Finish<'a>(&'a AtomicUsize, usize);
146
147impl Drop for Finish<'_> {
148    fn drop(&mut self) {
149        self.0.store(self.1, Ordering::Release);
150    }
151}
152
153impl Wavefront {
154    pub fn new(layout: TileLayout) -> Wavefront {
155        let n = layout.cols.len() * layout.sb_rows();
156        Wavefront {
157            layout,
158            progress: (0..n).map(|_| AtomicUsize::new(0)).collect(),
159            next: AtomicUsize::new(0),
160        }
161    }
162
163    pub fn layout(&self) -> &TileLayout {
164        &self.layout
165    }
166
167    pub fn jobs(&self) -> usize {
168        self.progress.len()
169    }
170
171    pub fn run<S: Send>(&self, workers: &mut [S], f: impl Fn(&mut S, &Row) + Sync) {
172        for p in &self.progress {
173            p.store(0, Ordering::Relaxed);
174        }
175        self.next.store(0, Ordering::Relaxed);
176        let work = |w: &mut S| {
177            loop {
178                let j = self.next.fetch_add(1, Ordering::Relaxed);
179                if j >= self.progress.len() {
180                    break;
181                }
182                let row = self.row(j);
183                let _done = Finish(row.me, row.sb_end - row.sb_start);
184                f(w, &row);
185            }
186        };
187        match workers {
188            [] => {}
189            [w] => work(w),
190            [first, rest @ ..] => rayon::in_place_scope(|s| {
191                for w in rest {
192                    let work = &work;
193                    s.spawn(move |_| work(w));
194                }
195                work(first);
196            }),
197        }
198    }
199
200    fn row(&self, j: usize) -> Row<'_> {
201        let tcs = self.layout.cols.len();
202        let (sb_row, tile_col) = (j / tcs, j % tcs);
203        let (c0, c1) = self.layout.cols[tile_col];
204        Row {
205            tile_col,
206            tile_row: self.layout.tile_row_of(sb_row),
207            sb_row,
208            sb_start: c0 / 8,
209            sb_end: c1.div_ceil(8),
210            above: (sb_row > 0).then(|| &self.progress[j - tcs]),
211            me: &self.progress[j],
212        }
213    }
214}
pito_codecs/vp9/enc/rc/threads.rs

pito_codecs/vp9/enc/rc/
threads.rs
