Expose function to clear memory cache (#1032)

* expose function to clear memory cache

* fix linux build

* fix metal tests

mlx/backend/metal/allocator.cpp

@@ -209,6 +209,11 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
   return Buffer{static_cast<void*>(buf)};
 }
 
+void MetalAllocator::clear_cache() {
+  std::unique_lock lk(mutex_);
+  buffer_cache_.clear();
+}
+
 void MetalAllocator::free(Buffer buffer) {
   auto buf = static_cast<MTL::Buffer*>(buffer.ptr());
   std::unique_lock lk(mutex_);
@@ -242,6 +247,9 @@ size_t get_peak_memory() {
 size_t get_cache_memory() {
   return allocator().get_cache_memory();
 }
+void clear_cache() {
+  return allocator().clear_cache();
+}
 
 } // namespace metal
 

mlx/backend/metal/allocator.h

@@ -26,6 +26,7 @@ class BufferCache {
   size_t cache_size() {
     return pool_size_;
   }
+  void clear();
 
  private:
   struct BufferHolder {
@@ -37,7 +38,6 @@ class BufferCache {
     MTL::Buffer* buf;
   };
 
-  void clear();
   void add_at_head(BufferHolder* to_add);
   void remove_from_list(BufferHolder* to_remove);
 
@@ -67,6 +67,7 @@ class MetalAllocator : public allocator::Allocator {
   };
   size_t set_cache_limit(size_t limit);
   size_t set_memory_limit(size_t limit, bool relaxed);
+  void clear_cache();
 
  private:
   MTL::Device* device_;

mlx/backend/metal/metal.h

@@ -54,6 +54,9 @@ size_t set_memory_limit(size_t limit, bool relaxed = true);
  * */
 size_t set_cache_limit(size_t limit);
 
+/* Clear the memory cache. */
+void clear_cache();
+
 /** Capture a GPU trace, saving it to an absolute file `path` */
 void start_capture(std::string path = "");
 void stop_capture();

mlx/backend/no_metal/metal.cpp

@@ -4,7 +4,6 @@
 
 #include "mlx/backend/metal/metal.h"
 #include "mlx/backend/metal/metal_impl.h"
-
 namespace mlx::core::metal {
 
 bool is_available() {
@@ -48,5 +47,6 @@ size_t set_cache_limit(size_t) {
 }
 void start_capture(std::string path) {}
 void stop_capture() {}
+void clear_cache() {}
 
 } // namespace mlx::core::metal