Skip to main content

cgpu/
budget.rs

1use crate::error::GpuError;
2
3pub const DEFAULT_SAFE_VRAM_BYTES: usize = 128 * 1024 * 1024;
4pub const DEFAULT_FILL_RATIO: f32 = 0.92;
5pub const MIN_FILL_RATIO: f32 = 0.90;
6pub const MAX_FILL_RATIO: f32 = 0.95;
7
8#[derive(Clone, Copy, Debug, PartialEq)]
9pub struct GpuMemoryBudget {
10    safe_vram_bytes: usize,
11    fill_ratio: f32,
12}
13
14impl Default for GpuMemoryBudget {
15    fn default() -> Self {
16        Self {
17            safe_vram_bytes: DEFAULT_SAFE_VRAM_BYTES,
18            fill_ratio: DEFAULT_FILL_RATIO,
19        }
20    }
21}
22
23impl GpuMemoryBudget {
24    pub fn new(safe_vram_bytes: usize, fill_ratio: f32) -> Result<Self, GpuError> {
25        if safe_vram_bytes == 0 {
26            return Err(GpuError::InvalidBudget(
27                "safe_vram_bytes must be greater than zero",
28            ));
29        }
30
31        if !(MIN_FILL_RATIO..=MAX_FILL_RATIO).contains(&fill_ratio) {
32            return Err(GpuError::InvalidBudget(
33                "fill_ratio must stay between 0.90 and 0.95",
34            ));
35        }
36
37        Ok(Self {
38            safe_vram_bytes,
39            fill_ratio,
40        })
41    }
42
43    pub fn safe_vram_bytes(self) -> usize {
44        self.safe_vram_bytes
45    }
46
47    pub fn fill_ratio(self) -> f32 {
48        self.fill_ratio
49    }
50
51    pub fn target_batch_bytes(self) -> usize {
52        ((self.safe_vram_bytes as f64) * f64::from(self.fill_ratio)).floor() as usize
53    }
54}