From 1a28b69ee28672e4e6a3c7d2f962547c9fb514f6 Mon Sep 17 00:00:00 2001
From: Awni Hannun <awni@apple.com>
Date: Sun, 6 Apr 2025 17:38:25 -0700
Subject: [PATCH] only add to residency set once (#2049)

---
 mlx/backend/metal/allocator.cpp | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/mlx/backend/metal/allocator.cpp b/mlx/backend/metal/allocator.cpp
index d21c66ade..0a69dd261 100644
--- a/mlx/backend/metal/allocator.cpp
+++ b/mlx/backend/metal/allocator.cpp
@@ -272,9 +272,13 @@ Buffer MetalAllocator::malloc(size_t size) {
     if (!buf) {
       buf = device_->newBuffer(size, resource_options);
     }
+    if (!buf) {
+      return Buffer{nullptr};
+    }
     lk.lock();
-    if (buf) {
-      num_resources_++;
+    num_resources_++;
+    if (!buf->heap()) {
+      residency_set_.insert(buf);
     }
   }
 
@@ -288,10 +292,6 @@ Buffer MetalAllocator::malloc(size_t size) {
         get_cache_memory() - max_pool_size_);
   }
 
-  if (!buf->heap()) {
-    residency_set_.insert(buf);
-  }
-
   return Buffer{static_cast<void*>(buf)};
 }