Wired (#1510)

* expose residency sets as wire/unwire

* returns wired size

* fix

* runtime support check

* fix os check

* fix test

* fix no metal build

* docs

* nit

* nits in docs

* nits

python/src/metal.cpp

@@ -6,6 +6,7 @@
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/unordered_map.h>
 #include <nanobind/stl/variant.h>
+#include <nanobind/stl/vector.h>
 
 namespace nb = nanobind;
 using namespace nb::literals;
@@ -98,6 +99,38 @@ void init_metal(nb::module_& m) {
       Returns:
         int: The previous cache limit in bytes.
       )pbdoc");
+  metal.def(
+      "set_wired_limit",
+      &metal::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");
   metal.def(
       "clear_cache",
       &metal::clear_cache,

python/tests/test_metal.py

@@ -47,6 +47,14 @@ class TestMetal(mlx_tests.MLXTestCase):
         mx.metal.reset_peak_memory()
         self.assertEqual(mx.metal.get_peak_memory(), 0)
 
+        old_limit = mx.metal.set_wired_limit(1000)
+        old_limit = mx.metal.set_wired_limit(0)
+        self.assertEqual(old_limit, 1000)
+
+        max_size = mx.metal.device_info()["max_recommended_working_set_size"]
+        with self.assertRaises(ValueError):
+            mx.metal.set_wired_limit(max_size + 10)
+
 
 if __name__ == "__main__":
     unittest.main()