Manage Metal objects with smart pointers

mlx/backend/metal/allocator.cpp

@@ -40,10 +40,10 @@ MetalAllocator::MetalAllocator()
             if (!buf->heap()) {
               residency_set_.erase(buf);
             }
+            auto pool = metal::new_scoped_memory_pool();
             buf->release();
           }),
       residency_set_(device_) {
-  auto pool = metal::new_scoped_memory_pool();
   const auto& info = gpu::device_info(0);
   auto memsize = std::get<size_t>(info.at("memory_size"));
   auto max_rec_size =
@@ -59,21 +59,15 @@ MetalAllocator::MetalAllocator()
   if (is_vm) {
     return;
   }
-  auto heap_desc = MTL::HeapDescriptor::alloc()->init();
+  auto pool = metal::new_scoped_memory_pool();
+  auto heap_desc = MTL::HeapDescriptor::alloc()->init()->autorelease();
   heap_desc->setResourceOptions(resource_options);
   heap_desc->setSize(heap_size_);
-  heap_ = device_->newHeap(heap_desc);
-  heap_desc->release();
-  residency_set_.insert(heap_);
+  heap_ = NS::TransferPtr(device_->newHeap(heap_desc));
+  residency_set_.insert(heap_.get());
 }
 
-MetalAllocator::~MetalAllocator() {
-  auto pool = metal::new_scoped_memory_pool();
-  if (heap_) {
-    heap_->release();
-  }
-  buffer_cache_.clear();
-}
+MetalAllocator::~MetalAllocator() = default;
 
 size_t MetalAllocator::set_cache_limit(size_t limit) {
   std::unique_lock lk(mutex_);
@@ -128,8 +122,6 @@ Buffer MetalAllocator::malloc(size_t size) {
   if (!buf) {
     size_t mem_required = get_active_memory() + get_cache_memory() + size;
 
-    auto pool = metal::new_scoped_memory_pool();
-
     // If we have a lot of memory pressure try to reclaim memory from the cache
     if (mem_required >= gc_limit_ || num_resources_ >= resource_limit_) {
       num_resources_ -=
@@ -167,7 +159,6 @@ Buffer MetalAllocator::malloc(size_t size) {
 
   // Maintain the cache below the requested limit
   if (get_cache_memory() > max_pool_size_) {
-    auto pool = metal::new_scoped_memory_pool();
     num_resources_ -= buffer_cache_.release_cached_buffers(
         get_cache_memory() - max_pool_size_);
   }
@@ -177,7 +168,6 @@ Buffer MetalAllocator::malloc(size_t size) {
 
 void MetalAllocator::clear_cache() {
   std::unique_lock lk(mutex_);
-  auto pool = metal::new_scoped_memory_pool();
   num_resources_ -= buffer_cache_.clear();
 }
 

mlx/backend/metal/allocator.h

@@ -51,16 +51,18 @@ class MetalAllocator : public allocator::Allocator {
   // the heap, a heap can have at most heap.size() / 256 buffers.
   static constexpr int small_size_ = 256;
   static constexpr int heap_size_ = 1 << 20;
-  MTL::Heap* heap_;
+
   MetalAllocator();
   ~MetalAllocator();
+
   friend MetalAllocator& allocator();
 
+  NS::SharedPtr<MTL::Heap> heap_;
+  ResidencySet residency_set_;
+
   // Caching allocator
   BufferCache<MTL::Buffer> buffer_cache_;
 
-  ResidencySet residency_set_;
-
   // Allocation stats
   size_t block_limit_;
   size_t gc_limit_;