diff --git a/mlx-rs/src/lib.rs b/mlx-rs/src/lib.rs index 1cc36f97..53459f34 100644 --- a/mlx-rs/src/lib.rs +++ b/mlx-rs/src/lib.rs @@ -292,6 +292,7 @@ pub mod fast; pub mod fft; pub mod linalg; pub mod losses; +pub mod memory; pub mod module; pub mod nested; pub mod nn; diff --git a/mlx-rs/src/memory.rs b/mlx-rs/src/memory.rs new file mode 100644 index 00000000..3edc6938 --- /dev/null +++ b/mlx-rs/src/memory.rs @@ -0,0 +1,115 @@ +//! GPU memory management. +//! +//! These functions wrap the global memory-management APIs from `mlx-c`. They +//! affect the active device, which on Apple Silicon is the unified-memory +//! Metal allocator. +//! +//! # Example +//! +//! ```rust,no_run +//! use mlx_rs::memory; +//! +//! // Cap MLX's GPU allocator at 8 GiB. Returns the previous limit. +//! let prev = memory::set_memory_limit(8 * 1024 * 1024 * 1024).unwrap(); +//! println!("previous limit: {} bytes", prev); +//! +//! // Read the current limit back. +//! let cur = memory::get_memory_limit().unwrap(); +//! println!("current limit: {} bytes", cur); +//! ``` + +use std::panic::Location; + +use crate::error::{Exception, Result}; + +/// Get the current memory limit, in bytes. +/// +/// This is the cap that MLX's allocator will keep its working set under. +/// Allocations beyond the limit will block until scheduled tasks free memory. +/// The default limit is 1.5× the device's recommended max working-set size. +#[track_caller] +pub fn get_memory_limit() -> Result { + init_error_handler(); + let mut current: usize = 0; + let status = unsafe { mlx_sys::mlx_get_memory_limit(&mut current as *mut usize) }; + if status != 0 { + return Err(last_error_or_unknown( + "mlx_get_memory_limit failed without setting an error", + )); + } + Ok(current) +} + +/// Set the memory limit, in bytes. Returns the previous limit. +/// +/// Calls to allocate will block on scheduled tasks if the limit is exceeded. +/// The limit defaults to 1.5× the device's recommended max working-set size, +/// so you usually only need this to lower the cap (e.g. when sharing a +/// machine with other GPU workloads). +#[track_caller] +pub fn set_memory_limit(limit: usize) -> Result { + init_error_handler(); + let mut previous: usize = 0; + let status = unsafe { mlx_sys::mlx_set_memory_limit(&mut previous as *mut usize, limit) }; + if status != 0 { + return Err(last_error_or_unknown( + "mlx_set_memory_limit failed without setting an error", + )); + } + Ok(previous) +} + +#[inline] +fn init_error_handler() { + crate::error::INIT_ERR_HANDLER + .with(|init| init.call_once(crate::error::setup_mlx_error_handler)); +} + +#[track_caller] +fn last_error_or_unknown(fallback_msg: &'static str) -> Exception { + match crate::error::get_and_clear_last_mlx_error() { + Some(raw) => Exception::from(raw), + None => Exception { + what: fallback_msg.to_string(), + location: Location::caller(), + }, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // Memory limit is global state on the active device, so we exercise + // get/set + the "returns previous" contract in a single test to avoid + // races with parallel test runners that could interleave writes. + #[test] + fn test_get_and_set_memory_limit() { + let original = get_memory_limit().unwrap(); + + // 1 GiB and 2 GiB are well under the default cap of 1.5x the + // device's recommended working set, so they should fit on every + // machine that can run MLX at all. + let a = 1usize << 30; + let b = 2usize << 30; + + let prev_a = set_memory_limit(a).unwrap(); + assert_eq!(prev_a, original, "first set should return the prior limit"); + assert_eq!( + get_memory_limit().unwrap(), + a, + "get should reflect the most recent set" + ); + + let prev_b = set_memory_limit(b).unwrap(); + assert_eq!( + prev_b, a, + "consecutive set should return the immediately previous limit" + ); + assert_eq!(get_memory_limit().unwrap(), b); + + // Restore. + set_memory_limit(original).unwrap(); + assert_eq!(get_memory_limit().unwrap(), original); + } +}