fix malloc or wait deadlock (#1976)

mlx/allocator.cpp

@@ -9,7 +9,7 @@
 namespace mlx::core::allocator {
 
 Buffer malloc(size_t size) {
-  auto buffer = allocator().malloc(size, /* allow_swap */ true);
+  auto buffer = allocator().malloc(size);
   if (size && !buffer.ptr()) {
     std::ostringstream msg;
     msg << "[malloc] Unable to allocate " << size << " bytes.";
@@ -22,7 +22,7 @@ void free(Buffer buffer) {
   allocator().free(buffer);
 }
 
-Buffer CommonAllocator::malloc(size_t size, bool) {
+Buffer CommonAllocator::malloc(size_t size) {
   void* ptr = std::malloc(size + sizeof(size_t));
   if (ptr != nullptr) {
     *static_cast<size_t*>(ptr) = size;
@@ -41,26 +41,4 @@ size_t CommonAllocator::size(Buffer buffer) const {
   return *static_cast<size_t*>(buffer.ptr());
 }
 
-Buffer malloc_or_wait(size_t size) {
-  auto buffer = allocator().malloc(size);
-
-  while (size && !buffer.ptr() && scheduler::n_active_tasks() > 0) {
-    scheduler::wait_for_one();
-    buffer = allocator().malloc(size);
-  }
-
-  // Try swapping if needed
-  if (size && !buffer.ptr()) {
-    buffer = allocator().malloc(size, /* allow_swap = */ true);
-  }
-
-  if (size && !buffer.ptr()) {
-    std::ostringstream msg;
-    msg << "[malloc_or_wait] Unable to allocate " << size << " bytes.";
-    throw std::runtime_error(msg.str());
-  }
-
-  return buffer;
-}
-
 } // namespace mlx::core::allocator