Fix some leaks and races (#1629)

* fix leak and fix potential race

* more leak fixes

* fix one more

mlx/backend/metal/allocator.cpp

@@ -30,7 +30,7 @@ BufferCache::BufferCache(MTL::Device* device)
     : device_(device), head_(nullptr), tail_(nullptr), pool_size_(0) {}
 
 BufferCache::~BufferCache() {
-  auto thread_pool = metal::new_scoped_memory_pool();
+  auto pool = metal::new_scoped_memory_pool();
   clear();
 }
 
@@ -155,11 +155,13 @@ MetalAllocator::MetalAllocator()
 }
 
 size_t MetalAllocator::set_cache_limit(size_t limit) {
+  std::unique_lock lk(mutex_);
   std::swap(limit, max_pool_size_);
   return limit;
 };
 
 size_t MetalAllocator::set_memory_limit(size_t limit, bool relaxed) {
+  std::unique_lock lk(mutex_);
   std::swap(limit, block_limit_);
   relaxed_ = relaxed;
   gc_limit_ = std::min(
@@ -169,6 +171,7 @@ size_t MetalAllocator::set_memory_limit(size_t limit, bool relaxed) {
 };
 
 size_t MetalAllocator::set_wired_limit(size_t limit) {
+  std::unique_lock lk(mutex_);
   std::swap(limit, wired_limit_);
   residency_set_.resize(wired_limit_);
   return limit;
@@ -205,7 +208,7 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
       return Buffer{nullptr};
     }
 
-    auto thread_pool = metal::new_scoped_memory_pool();
+    auto pool = metal::new_scoped_memory_pool();
 
     // If we have a lot of memory pressure or are over the maximum cache size,
     // try to reclaim memory from the cache
@@ -226,7 +229,7 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
 
   // Maintain the cache below the requested limit
   if (get_cache_memory() >= max_pool_size_) {
-    auto thread_pool = metal::new_scoped_memory_pool();
+    auto pool = metal::new_scoped_memory_pool();
     buffer_cache_.release_cached_buffers(get_cache_memory() - max_pool_size_);
   }
 
@@ -237,6 +240,7 @@ Buffer MetalAllocator::malloc(size_t size, bool allow_swap /* = false */) {
 
 void MetalAllocator::clear_cache() {
   std::unique_lock lk(mutex_);
+  auto pool = metal::new_scoped_memory_pool();
   buffer_cache_.clear();
 }
 
@@ -252,7 +256,7 @@ void MetalAllocator::free(Buffer buffer) {
     buffer_cache_.recycle_to_cache(buf);
   } else {
     lk.unlock();
-    auto thread_pool = metal::new_scoped_memory_pool();
+    auto pool = metal::new_scoped_memory_pool();
     buf->release();
   }
 }

mlx/backend/metal/device.cpp

@@ -645,6 +645,9 @@ void new_stream(Stream stream) {
 
 std::unordered_map<std::string, std::variant<std::string, size_t>>
 device_info() {
+  auto init_device_info = []()
+      -> std::unordered_map<std::string, std::variant<std::string, size_t>> {
+    auto pool = new_scoped_memory_pool();
     auto raw_device = device(default_device()).mtl_device();
     auto arch = std::string(raw_device->architecture()->name()->utf8String());
 
@@ -660,6 +663,9 @@ device_info() {
         {"max_recommended_working_set_size",
          raw_device->recommendedMaxWorkingSetSize()},
         {"memory_size", memsize}};
+  };
+  static auto device_info_ = init_device_info();
+  return device_info_;
 }
 
 } // namespace mlx::core::metal

mlx/backend/metal/metal.cpp

@@ -94,6 +94,7 @@ std::function<void()> make_synchronize_task(
     Stream s,
     std::shared_ptr<std::promise<void>> p) {
   return [s, p = std::move(p)]() {
+    auto pool = new_scoped_memory_pool();
     auto& d = metal::device(s.device);
     auto cb = d.get_command_buffer(s.index);
     cb->retain();

mlx/backend/metal/resident.cpp

@@ -63,6 +63,7 @@ void ResidencySet::resize(size_t size) {
   size_t current_size = wired_set_->allocatedSize();
 
   if (current_size < size) {
+    auto pool = new_scoped_memory_pool();
     // Add unwired allocations to the set
     for (auto it = unwired_set_.begin(); it != unwired_set_.end();) {
       auto buf_size = (*it)->allocatedSize();
@@ -77,6 +78,7 @@ void ResidencySet::resize(size_t size) {
     wired_set_->commit();
     wired_set_->requestResidency();
   } else if (current_size > size) {
+    auto pool = new_scoped_memory_pool();
     // Remove wired allocations until under capacity
     auto allocations = wired_set_->allAllocations();
     auto num_allocations = wired_set_->allocationCount();
@@ -92,6 +94,7 @@ void ResidencySet::resize(size_t size) {
 
 ResidencySet::~ResidencySet() {
   if (wired_set_) {
+    auto pool = new_scoped_memory_pool();
     wired_set_->release();
   }
 }