forked from ml-explore/mlx
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathmetal.cpp
More file actions
108 lines (99 loc) · 3.02 KB
/
metal.cpp
File metadata and controls
108 lines (99 loc) · 3.02 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
// Copyright © 2023-2024 Apple Inc.
#include <iostream>
#include <nanobind/nanobind.h>
#include <nanobind/stl/optional.h>
#include <nanobind/stl/string.h>
#include <nanobind/stl/unordered_map.h>
#include <nanobind/stl/variant.h>
#include <nanobind/stl/vector.h>
#include "mlx/backend/metal/metal.h"
#include "mlx/memory.h"
namespace mx = mlx::core;
namespace nb = nanobind;
using namespace nb::literals;
bool DEPRECATE(const std::string& old_fn, const std::string new_fn) {
std::cerr << old_fn << " is deprecated and will be removed in a future "
<< "version. Use " << new_fn << " instead." << std::endl;
return true;
}
#define DEPRECATE(oldfn, newfn) static bool dep = DEPRECATE(oldfn, newfn)
void init_metal(nb::module_& m) {
nb::module_ metal = m.def_submodule("metal", "mlx.metal");
metal.def(
"is_available",
&mx::metal::is_available,
R"pbdoc(
Check if the Metal back-end is available.
)pbdoc");
metal.def("get_active_memory", []() {
DEPRECATE("mx.metal.get_active_memory", "mx.get_active_memory");
return mx::get_active_memory();
});
metal.def("get_peak_memory", []() {
DEPRECATE("mx.metal.get_peak_memory", "mx.get_peak_memory");
return mx::get_peak_memory();
});
metal.def("reset_peak_memory", []() {
DEPRECATE("mx.metal.reset_peak_memory", "mx.reset_peak_memory");
mx::reset_peak_memory();
});
metal.def("get_cache_memory", []() {
DEPRECATE("mx.metal.get_cache_memory", "mx.get_cache_memory");
return mx::get_cache_memory();
});
metal.def(
"set_memory_limit",
[](size_t limit) {
DEPRECATE("mx.metal.set_memory_limt", "mx.set_memory_limit");
return mx::set_memory_limit(limit);
},
"limit"_a);
metal.def(
"set_cache_limit",
[](size_t limit) {
DEPRECATE("mx.metal.set_cache_limt", "mx.set_cache_limit");
return mx::set_cache_limit(limit);
},
"limit"_a);
metal.def(
"set_wired_limit",
[](size_t limit) {
DEPRECATE("mx.metal.set_wired_limt", "mx.set_wired_limit");
return mx::set_wired_limit(limit);
},
"limit"_a);
metal.def("clear_cache", []() {
DEPRECATE("mx.metal.clear_cache", "mx.clear_cache");
mx::clear_cache();
});
metal.def(
"start_capture",
&mx::metal::start_capture,
"path"_a,
R"pbdoc(
Start a Metal capture.
Args:
path (str): The path to save the capture which should have
the extension ``.gputrace``.
)pbdoc");
metal.def(
"stop_capture",
&mx::metal::stop_capture,
R"pbdoc(
Stop a Metal capture.
)pbdoc");
metal.def(
"device_info",
&mx::metal::device_info,
R"pbdoc(
Get information about the GPU device and system settings.
Currently returns:
* ``architecture``
* ``max_buffer_size``
* ``max_recommended_working_set_size``
* ``memory_size``
* ``resource_limit``
Returns:
dict: A dictionary with string keys and string or integer values.
)pbdoc");
}