|
|
|
|
|
|
|
|
#include "mlx/memory.h" |
|
|
#include <nanobind/nanobind.h> |
|
|
|
|
|
namespace mx = mlx::core; |
|
|
namespace nb = nanobind; |
|
|
using namespace nb::literals; |
|
|
|
|
|
void init_memory(nb::module_& m) { |
|
|
m.def( |
|
|
"get_active_memory", |
|
|
&mx::get_active_memory, |
|
|
R"pbdoc( |
|
|
Get the actively used memory in bytes. |
|
|
|
|
|
Note, this will not always match memory use reported by the system because |
|
|
it does not include cached memory buffers. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"get_peak_memory", |
|
|
&mx::get_peak_memory, |
|
|
R"pbdoc( |
|
|
Get the peak amount of used memory in bytes. |
|
|
|
|
|
The maximum memory used recorded from the beginning of the program |
|
|
execution or since the last call to :func:`reset_peak_memory`. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"reset_peak_memory", |
|
|
&mx::reset_peak_memory, |
|
|
R"pbdoc( |
|
|
Reset the peak memory to zero. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"get_cache_memory", |
|
|
&mx::get_cache_memory, |
|
|
R"pbdoc( |
|
|
Get the cache size in bytes. |
|
|
|
|
|
The cache includes memory not currently used that has not been returned |
|
|
to the system allocator. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"set_memory_limit", |
|
|
&mx::set_memory_limit, |
|
|
"limit"_a, |
|
|
R"pbdoc( |
|
|
Set the memory limit. |
|
|
|
|
|
The memory limit is a guideline for the maximum amount of memory to use |
|
|
during graph evaluation. If the memory limit is exceeded and there is no |
|
|
more RAM (including swap when available) allocations will result in an |
|
|
exception. |
|
|
|
|
|
When metal is available the memory limit defaults to 1.5 times the |
|
|
maximum recommended working set size reported by the device. |
|
|
|
|
|
Args: |
|
|
limit (int): Memory limit in bytes. |
|
|
|
|
|
Returns: |
|
|
int: The previous memory limit in bytes. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"set_cache_limit", |
|
|
&mx::set_cache_limit, |
|
|
"limit"_a, |
|
|
R"pbdoc( |
|
|
Set the free cache limit. |
|
|
|
|
|
If using more than the given limit, free memory will be reclaimed |
|
|
from the cache on the next allocation. To disable the cache, set |
|
|
the limit to ``0``. |
|
|
|
|
|
The cache limit defaults to the memory limit. See |
|
|
:func:`set_memory_limit` for more details. |
|
|
|
|
|
Args: |
|
|
limit (int): The cache limit in bytes. |
|
|
|
|
|
Returns: |
|
|
int: The previous cache limit in bytes. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"set_wired_limit", |
|
|
&mx::set_wired_limit, |
|
|
"limit"_a, |
|
|
R"pbdoc( |
|
|
Set the wired size limit. |
|
|
|
|
|
.. note:: |
|
|
* This function is only useful on macOS 15.0 or higher. |
|
|
* The wired limit should remain strictly less than the total |
|
|
memory size. |
|
|
|
|
|
The wired limit is the total size in bytes of memory that will be kept |
|
|
resident. The default value is ``0``. |
|
|
|
|
|
Setting a wired limit larger than system wired limit is an error. You can |
|
|
increase the system wired limit with: |
|
|
|
|
|
.. code-block:: |
|
|
|
|
|
sudo sysctl iogpu.wired_limit_mb=<size_in_megabytes> |
|
|
|
|
|
Use :func:`device_info` to query the system wired limit |
|
|
(``"max_recommended_working_set_size"``) and the total memory size |
|
|
(``"memory_size"``). |
|
|
|
|
|
Args: |
|
|
limit (int): The wired limit in bytes. |
|
|
|
|
|
Returns: |
|
|
int: The previous wired limit in bytes. |
|
|
)pbdoc"); |
|
|
m.def( |
|
|
"clear_cache", |
|
|
&mx::clear_cache, |
|
|
R"pbdoc( |
|
|
Clear the memory cache. |
|
|
|
|
|
After calling this, :func:`get_cache_memory` should return ``0``. |
|
|
)pbdoc"); |
|
|
} |
|
|
|
