mlx_c_rust/runtime/
memory.rs

1//! Safe wrappers over MLX's process-global memory and Metal controls.
2//!
3//! The embedding application decides when to call these; this module only
4//! makes the calls safe and typed. Every function here mirrors exactly one
5//! official MLX-C control.
6
7use crate::{MlxCError, MlxStream, error::check_status, raw};
8
9/// Applies the graph-evaluation active-memory guidance and returns the
10/// previous limit in bytes.
11pub fn set_memory_limit(active_memory_limit_bytes: usize) -> Result<usize, MlxCError> {
12    let mut previous_limit_bytes = 0;
13    // SAFETY: The output points to initialized writable storage and the
14    // non-terminating handler is installed before policy runs.
15    let status =
16        unsafe { raw::mlx_set_memory_limit(&mut previous_limit_bytes, active_memory_limit_bytes) };
17    check_status(status, "set the MLX active memory limit")?;
18    Ok(previous_limit_bytes)
19}
20
21/// Applies the allocator-cache retention limit and returns the previous limit
22/// in bytes.
23pub fn set_cache_limit(allocator_cache_limit_bytes: usize) -> Result<usize, MlxCError> {
24    let mut previous_limit_bytes = 0;
25    // SAFETY: The output points to initialized writable storage and the cache
26    // limit is bounded by the active-memory limit the caller validated.
27    let status =
28        unsafe { raw::mlx_set_cache_limit(&mut previous_limit_bytes, allocator_cache_limit_bytes) };
29    check_status(status, "set the MLX allocator cache memory limit")?;
30    Ok(previous_limit_bytes)
31}
32
33/// Reads the allocator limit currently enforced by MLX.
34pub fn memory_limit_bytes() -> Result<usize, MlxCError> {
35    read_metric("read the MLX memory limit", raw::mlx_get_memory_limit)
36}
37
38/// Samples process-local MLX active-memory bytes.
39pub fn active_memory_bytes() -> Result<usize, MlxCError> {
40    read_metric("read MLX active memory", raw::mlx_get_active_memory)
41}
42
43/// Samples process-local MLX allocator-cache bytes.
44pub fn cache_memory_bytes() -> Result<usize, MlxCError> {
45    read_metric("read MLX allocator-cache memory", raw::mlx_get_cache_memory)
46}
47
48/// Samples process-local MLX peak active-memory bytes.
49pub fn peak_memory_bytes() -> Result<usize, MlxCError> {
50    read_metric("read MLX peak memory", raw::mlx_get_peak_memory)
51}
52
53/// Resets MLX's process-local peak active-memory counter.
54pub fn reset_peak_memory() -> Result<(), MlxCError> {
55    // SAFETY: The process-global error handler is installed during runtime
56    // initialization and `mlx_reset_peak_memory` has no pointer arguments.
57    let status = unsafe { raw::mlx_reset_peak_memory() };
58    check_status(status, "reset MLX peak memory")
59}
60
61/// Releases reclaimable MLX allocator-cache allocations.
62pub fn clear_allocator_cache() -> Result<(), MlxCError> {
63    // SAFETY: The process-global error handler is installed during runtime
64    // initialization and `mlx_clear_cache` has no pointer arguments.
65    let status = unsafe { raw::mlx_clear_cache() };
66    check_status(status, "clear the MLX allocator cache")
67}
68
69/// Waits for submitted work on the given stream to complete.
70pub fn synchronize(stream: &MlxStream) -> Result<(), MlxCError> {
71    // SAFETY: The caller's stream owner keeps the stream live for the call.
72    let status = unsafe { raw::mlx_synchronize(stream.raw()) };
73    check_status(status, "synchronize the MLX GPU stream")
74}
75
76fn read_metric(
77    operation: &'static str,
78    metric_reader: unsafe extern "C" fn(*mut usize) -> i32,
79) -> Result<usize, MlxCError> {
80    let mut metric_bytes = 0;
81    // SAFETY: Every accepted reader is an MLX C memory getter with the same
82    // ABI and receives a valid pointer to writable `usize` storage.
83    let status = unsafe { metric_reader(&mut metric_bytes) };
84    check_status(status, operation)?;
85    Ok(metric_bytes)
86}
mlx_c_rust/runtime/memory.rs

mlx_c_rust/runtime/
memory.rs
