mlx_rs/
memory.rs

1//! Process-global MLX allocator controls and observations.
2//!
3//! The state exposed here affects every thread and every evaluation in the process. These
4//! functions do not synchronize a sequence of reads and writes for callers.
5
6use crate::{error::Result, utils::guard::Guarded};
7
8/// Clear cached allocator buffers.
9///
10/// This affects all threads and evaluations. It is unrelated to
11/// [`transforms::compile::clear_cache`](crate::transforms::compile::clear_cache), which clears
12/// compiled functions.
13pub fn clear_cache() -> Result<()> {
14    <()>::try_from_op(|_| unsafe { mlx_sys::mlx_clear_cache() })
15}
16
17/// Return the number of bytes in active use by MLX.
18pub fn active_memory() -> Result<usize> {
19    usize::try_from_op(|res| unsafe { mlx_sys::mlx_get_active_memory(res) })
20}
21
22/// Return the number of bytes held in MLX's allocator cache.
23pub fn cache_memory() -> Result<usize> {
24    usize::try_from_op(|res| unsafe { mlx_sys::mlx_get_cache_memory(res) })
25}
26
27/// Return the peak number of bytes used by MLX.
28pub fn peak_memory() -> Result<usize> {
29    usize::try_from_op(|res| unsafe { mlx_sys::mlx_get_peak_memory(res) })
30}
31
32/// Reset the process-global peak-memory counter for all threads and evaluations.
33pub fn reset_peak_memory() -> Result<()> {
34    <()>::try_from_op(|_| unsafe { mlx_sys::mlx_reset_peak_memory() })
35}
36
37/// Return the process-global memory limit in bytes.
38pub fn memory_limit() -> Result<usize> {
39    usize::try_from_op(|res| unsafe { mlx_sys::mlx_get_memory_limit(res) })
40}
41
42/// Set the process-global memory limit in bytes and return its previous value.
43///
44/// The new limit affects all threads and evaluations.
45pub fn set_memory_limit(limit: usize) -> Result<usize> {
46    usize::try_from_op(|res| unsafe { mlx_sys::mlx_set_memory_limit(res, limit) })
47}
48
49/// Set the process-global allocator-cache limit in bytes and return its previous value.
50///
51/// The new limit affects all threads and evaluations.
52pub fn set_cache_limit(limit: usize) -> Result<usize> {
53    usize::try_from_op(|res| unsafe { mlx_sys::mlx_set_cache_limit(res, limit) })
54}
55
56/// Set the process-global wired-memory limit in bytes and return its previous value.
57///
58/// The new limit affects all threads and evaluations. Wired-memory limits are useful only with
59/// Metal on macOS 15 or later; other configurations may return an upstream error.
60pub fn set_wired_limit(limit: usize) -> Result<usize> {
61    usize::try_from_op(|res| unsafe { mlx_sys::mlx_set_wired_limit(res, limit) })
62}
mlx_rs/memory.rs

mlx_rs/
memory.rs
