Reset peak memory (#1074)

* reset peak memory

* fix linux

* nits in docs

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;
 }