#[cfg(any(target_os = "linux", test))]
mod meminfo;
use candle_core::Device;
pub(crate) fn budget_bytes(device: &Device) -> Option<u64> {
match device {
Device::Cpu => physical_memory(),
Device::Metal(metal) => metal_working_set(metal),
Device::Cuda(cuda) => cuda_free(cuda),
}
}
#[cfg(feature = "metal")]
fn metal_working_set(device: &candle_core::MetalDevice) -> Option<u64> {
let raw = device.metal_device();
let budget = raw.recommended_max_working_set_size() as u64;
Some(budget.saturating_sub(raw.current_allocated_size() as u64))
}
#[cfg(not(feature = "metal"))]
fn metal_working_set(_: &candle_core::MetalDevice) -> Option<u64> {
None
}
#[cfg(feature = "cuda")]
fn cuda_free(device: &candle_core::CudaDevice) -> Option<u64> {
let (free, _total) = device.cuda_stream().context().mem_get_info().ok()?;
Some(free as u64)
}
#[cfg(not(feature = "cuda"))]
fn cuda_free(_: &candle_core::CudaDevice) -> Option<u64> {
None
}
#[cfg(target_os = "macos")]
fn physical_memory() -> Option<u64> {
let mut bytes = 0u64;
let mut len = std::mem::size_of::<u64>();
let name = c"hw.memsize".as_ptr();
let status = unsafe {
libc::sysctlbyname(
name,
(&raw mut bytes).cast(),
&mut len,
std::ptr::null_mut(),
0,
)
};
(status == 0).then_some(bytes)
}
#[cfg(target_os = "linux")]
fn physical_memory() -> Option<u64> {
meminfo::mem_total(&std::fs::read_to_string("/proc/meminfo").ok()?)
}
#[cfg(not(any(target_os = "macos", target_os = "linux")))]
fn physical_memory() -> Option<u64> {
None
}
#[cfg(test)]
mod tests;