use crate::error::GpuError;
pub const DEFAULT_SAFE_VRAM_BYTES: usize = 128 * 1024 * 1024;
pub const DEFAULT_FILL_RATIO: f32 = 0.92;
pub const MIN_FILL_RATIO: f32 = 0.90;
pub const MAX_FILL_RATIO: f32 = 0.95;
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct GpuMemoryBudget {
safe_vram_bytes: usize,
fill_ratio: f32,
}
impl Default for GpuMemoryBudget {
fn default() -> Self {
Self {
safe_vram_bytes: DEFAULT_SAFE_VRAM_BYTES,
fill_ratio: DEFAULT_FILL_RATIO,
}
}
}
impl GpuMemoryBudget {
pub fn new(safe_vram_bytes: usize, fill_ratio: f32) -> Result<Self, GpuError> {
if safe_vram_bytes == 0 {
return Err(GpuError::InvalidBudget(
"safe_vram_bytes must be greater than zero",
));
}
if !(MIN_FILL_RATIO..=MAX_FILL_RATIO).contains(&fill_ratio) {
return Err(GpuError::InvalidBudget(
"fill_ratio must stay between 0.90 and 0.95",
));
}
Ok(Self {
safe_vram_bytes,
fill_ratio,
})
}
pub fn safe_vram_bytes(self) -> usize {
self.safe_vram_bytes
}
pub fn fill_ratio(self) -> f32 {
self.fill_ratio
}
pub fn target_batch_bytes(self) -> usize {
((self.safe_vram_bytes as f64) * f64::from(self.fill_ratio)).floor() as usize
}
}