Reset peak memory (#1074)

* reset peak memory

* fix linux

* nits in docs

docs/src/python/metal.rst

@@ -10,6 +10,7 @@ Metal
   device_info
   get_active_memory
   get_peak_memory
+  reset_peak_memory
   get_cache_memory
   set_memory_limit
   set_cache_limit

mlx/backend/metal/allocator.cpp

@@ -258,6 +258,9 @@ size_t get_active_memory() {
 size_t get_peak_memory() {
   return allocator().get_peak_memory();
 }
+void reset_peak_memory() {
+  allocator().reset_peak_memory();
+}
 size_t get_cache_memory() {
   return allocator().get_cache_memory();
 }

mlx/backend/metal/allocator.h

@@ -62,6 +62,10 @@ class MetalAllocator : public allocator::Allocator {
   size_t get_peak_memory() {
     return peak_memory_;
   };
+  void reset_peak_memory() {
+    std::unique_lock lk(mutex_);
+    peak_memory_ = 0;
+  };
   size_t get_cache_memory() {
     return buffer_cache_.cache_size();
   };

mlx/backend/metal/metal.h

@@ -20,11 +20,15 @@ size_t get_active_memory();
 
 /* Get the peak amount of used memory in bytes.
  *
- * The maximum memory used is recorded from the beginning of the program
- * execution.
+ * The maximum memory used recorded from the beginning of the program
+ * execution or since the last call to reset_peak_memory.
  * */
 size_t get_peak_memory();
 
+/* Reset the peak memory to zero.
+ * */
+void reset_peak_memory();
+
 /* Get the cache size in bytes.
  *
  * The cache includes memory not currently used that has not been returned

mlx/backend/no_metal/metal.cpp

@@ -36,6 +36,7 @@ size_t get_active_memory() {
 size_t get_peak_memory() {
   return 0;
 }
+void reset_peak_memory() {}
 size_t get_cache_memory() {
   return 0;
 }

python/src/metal.cpp

@@ -35,8 +35,14 @@ void init_metal(nb::module_& m) {
       R"pbdoc(
       Get the peak amount of used memory in bytes.
 
-      The maximum memory used is recorded from the beginning of the program
-      execution.
+      The maximum memory used recorded from the beginning of the program
+      execution or since the last call to :func:`reset_peak_memory`.
+      )pbdoc");
+  metal.def(
+      "reset_peak_memory",
+      &metal::reset_peak_memory,
+      R"pbdoc(
+      Reset the peak memory to zero.
       )pbdoc");
   metal.def(
       "get_cache_memory",

python/tests/test_metal.py

@@ -45,6 +45,9 @@ class TestMetal(mlx_tests.MLXTestCase):
         mx.metal.clear_cache()
         self.assertEqual(mx.metal.get_cache_memory(), 0)
 
+        mx.metal.reset_peak_memory()
+        self.assertEqual(mx.metal.get_peak_memory(), 0)
+
 
 if __name__ == "__main__":
     unittest.main()