Avoid invoking allocator::malloc when creating CUDA event (#2232)

mlx/backend/cuda/allocator.cpp

@@ -18,7 +18,10 @@ CudaAllocator::CudaAllocator()
     : buffer_cache_(
           getpagesize(),
           [](CudaBuffer* buf) { return buf->size; },
-          [this](CudaBuffer* buf) { cuda_free(buf); }) {
+          [this](CudaBuffer* buf) {
+            cuda_free(buf->data);
+            delete buf;
+          }) {
   // TODO: Set memory limit for multi-device.
   size_t free, total;
   CHECK_CUDA_ERROR(cudaMemGetInfo(&free, &total));
@@ -70,7 +73,8 @@ void CudaAllocator::free(Buffer buffer) {
     buffer_cache_.recycle_to_cache(buf);
   } else {
     lock.unlock();
-    cuda_free(buf);
+    cuda_free(buf->data);
+    delete buf;
   }
 }
 
@@ -87,6 +91,25 @@ void CudaAllocator::register_this_thread() {
   allowed_threads_.insert(std::this_thread::get_id());
 }
 
+void CudaAllocator::cuda_free(void* buf) {
+  // If cuda_free() is called from a unregistered thread, reschedule the call to
+  // worker.
+  {
+    std::lock_guard lock(worker_mutex_);
+    if (allowed_threads_.count(std::this_thread::get_id()) == 0) {
+      if (!worker_) {
+        worker_.reset(new Worker);
+      }
+      worker_->add_task([this, buf]() { this->cuda_free(buf); });
+      worker_->end_batch();
+      worker_->commit();
+      return;
+    }
+  }
+
+  cudaFree(buf);
+}
+
 size_t CudaAllocator::get_active_memory() const {
   return active_memory_;
 }
@@ -125,26 +148,6 @@ void CudaAllocator::clear_cache() {
   buffer_cache_.clear();
 }
 
-void CudaAllocator::cuda_free(CudaBuffer* buf) {
-  // If cuda_free() is called from a unregistered thread, reschedule the call to
-  // worker.
-  {
-    std::lock_guard lock(worker_mutex_);
-    if (allowed_threads_.count(std::this_thread::get_id()) == 0) {
-      if (!worker_) {
-        worker_.reset(new Worker);
-      }
-      worker_->add_task([this, buf]() { this->cuda_free(buf); });
-      worker_->end_batch();
-      worker_->commit();
-      return;
-    }
-  }
-
-  cudaFree(buf->data);
-  delete buf;
-}
-
 CudaAllocator& allocator() {
   // By creating the |allocator_| on heap, the destructor of CudaAllocator
   // will not be called on exit and buffers in the cache will be leaked. This