mirror of
https://github.com/ml-explore/mlx.git
synced 2025-10-20 17:38:09 +08:00
move memory APIs into top level mlx.core (#1982)
This commit is contained in:
125
python/src/memory.cpp
Normal file
125
python/src/memory.cpp
Normal file
@@ -0,0 +1,125 @@
|
||||
// Copyright © 2025 Apple Inc.
|
||||
|
||||
#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");
|
||||
}
|
Reference in New Issue
Block a user