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}