forked from ml-explore/mlx
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathmemory.cpp
More file actions
125 lines (102 loc) · 3.38 KB
/
memory.cpp
File metadata and controls
125 lines (102 loc) · 3.38 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
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");
}