pagesift_gpu/
pool.rs

1//! 显存池：`cuMemAlloc` 一次要几百微秒，每层都分配会吃掉连续执行省下的时间。
2//! 所有内核都排在同一条流上按序执行，所以缓冲一还回来就能给后面的内核用，不用等它们跑完。
3
4use std::collections::{HashMap, HashSet};
5
6use crate::Gpu;
7use crate::buffer::DeviceBuffer;
8
9/// 显存不够的错误都带这个前缀，调用方据此减小批次，而不是放弃 GPU。
10pub const OOM_TAG: &str = "out of GPU memory";
11
12/// 错误是否是显存不足（减小批次还有救）。
13pub fn is_oom(e: &str) -> bool {
14    e.starts_with(OOM_TAG)
15}
16
17/// 把元素数向上取整到容量档：每个二进制数量级分 8 档，浪费不超过 1/8，
18/// 宽度不同的批次也能复用同一批缓冲。
19pub(crate) fn capacity_class(len: usize) -> usize {
20    if len == 0 {
21        return 0;
22    }
23    // [推导] 每档差 2^(位数 − 4)，即 1/8 到 1/16 个数量级：8 档/倍频
24    let unit = 1usize << (usize::BITS - len.leading_zeros()).saturating_sub(4);
25    len.div_ceil(unit) * unit
26}
27
28/// 空闲缓冲按容量档存放；同时记录显存用量。
29pub(crate) struct Pool<'g> {
30    free: HashMap<usize, Vec<DeviceBuffer<'g>>>,
31    /// 这一批用过的容量档。
32    used: HashSet<usize>,
33    /// 分配过、还没还给驱动的字节数（在用的加空闲的）。
34    held: u64,
35    peak: u64,
36    in_use: u64,
37    in_use_peak: u64,
38}
39
40impl<'g> Pool<'g> {
41    pub fn new() -> Pool<'g> {
42        Pool {
43            free: HashMap::new(),
44            used: HashSet::new(),
45            held: 0,
46            peak: 0,
47            in_use: 0,
48            in_use_peak: 0,
49        }
50    }
51
52    /// 取一块至少 `cap` 个元素的缓冲。`cap` 必须已经是容量档。显存不够时先把空闲的还给驱动再试。
53    pub fn take(&mut self, gpu: &'g Gpu, cap: usize) -> Result<DeviceBuffer<'g>, String> {
54        self.used.insert(cap);
55        let b = match self.free.get_mut(&cap).and_then(Vec::pop) {
56            Some(b) => b,
57            None => {
58                let b = match gpu.alloc(cap) {
59                    Ok(b) => b,
60                    Err(_) => {
61                        self.trim();
62                        gpu.alloc(cap)
63                            .map_err(|e| format!("{OOM_TAG}: need {} bytes ({e})", bytes(cap)))?
64                    }
65                };
66                self.held += bytes(cap);
67                self.peak = self.peak.max(self.held);
68                b
69            }
70        };
71        self.in_use += bytes(cap);
72        self.in_use_peak = self.in_use_peak.max(self.in_use);
73        Ok(b)
74    }
75
76    /// 还回来。
77    pub fn give(&mut self, cap: usize, b: DeviceBuffer<'g>) {
78        self.in_use -= bytes(cap);
79        self.free.entry(cap).or_default().push(b);
80    }
81
82    /// 把空闲缓冲全还给驱动。
83    pub fn trim(&mut self) {
84        self.held -= self.idle();
85        self.free.clear();
86    }
87
88    /// 一批结束：这批没用到的容量档还给驱动。
89    pub fn end_batch(&mut self) {
90        let used = std::mem::take(&mut self.used);
91        let mut freed = 0;
92        self.free.retain(|cap, v| {
93            let keep = used.contains(cap);
94            if !keep {
95                freed += bytes(*cap) * v.len() as u64;
96            }
97            keep
98        });
99        self.held -= freed;
100    }
101
102    pub fn peak(&self) -> u64 {
103        self.peak
104    }
105
106    pub fn in_use_peak(&self) -> u64 {
107        self.in_use_peak
108    }
109
110    pub fn reset_peak(&mut self) {
111        self.peak = self.held;
112        self.in_use_peak = self.in_use;
113    }
114
115    /// 空闲缓冲占的字节数（需要时能还给驱动）。
116    pub fn idle(&self) -> u64 {
117        self.free
118            .iter()
119            .map(|(&cap, v)| bytes(cap) * v.len() as u64)
120            .sum()
121    }
122}
123
124pub(crate) fn bytes(len: usize) -> u64 {
125    len as u64 * size_of::<f32>() as u64
126}
127
128#[cfg(test)]
129mod tests {
130    use super::capacity_class;
131
132    #[test]
133    fn 容量档不小于请求且浪费有界() {
134        assert_eq!(capacity_class(0), 0);
135        for len in 1..100_000usize {
136            let c = capacity_class(len);
137            assert!(c >= len);
138            assert!((c - len) * 8 <= len.max(8), "{len} → {c}");
139        }
140        // 相近的宽度落在同一档，才能复用
141        assert_eq!(capacity_class(1000), capacity_class(1020));
142    }
143}
pagesift_gpu/pool.rs

pagesift_gpu/
pool.rs
