Fix array is_available race cases (#1468)

mlx/array.cpp

@@ -95,13 +95,29 @@ void array::detach() {
   array_desc_->primitive = nullptr;
 }
 
-void array::eval() {
-  // Ensure the array is ready to be read
-  if (status() == Status::scheduled) {
+bool array::is_available() const {
+  if (status() == Status::available) {
+    return true;
+  } else if (status() == Status::evaluated && event().is_signaled()) {
+    set_status(Status::available);
+    return true;
+  }
+  return false;
+}
+
+void array::wait() {
+  if (!is_available()) {
     event().wait();
     set_status(Status::available);
-  } else if (status() == Status::unscheduled) {
+  }
+}
+
+void array::eval() {
+  // Ensure the array is ready to be read
+  if (status() == Status::unscheduled) {
     mlx::core::eval({*this});
+  } else {
+    wait();
   }
 }
 

mlx/array.h

@@ -344,11 +344,33 @@ class array {
     return static_cast<T*>(array_desc_->data_ptr);
   }
 
-  enum Status { unscheduled, scheduled, available };
+  enum Status {
+    // The ouptut of a computation which has not been scheduled.
+    // For example, the status of `x` in `auto x = a + b`.
+    unscheduled,
 
-  bool is_available() const {
-    return status() == Status::available;
-  }
+    // The ouptut of a computation which has been scheduled but `eval_*` has
+    // not yet been called on the array's primitive. A possible
+    // status of `x` in `auto x = a + b; eval(x);`
+    scheduled,
+
+    // The array's `eval_*` function has been run, but the computation is not
+    // necessarily complete. The array will have memory allocated and if it is
+    // not a tracer then it will be detached from the graph.
+    evaluated,
+
+    // If the array is the output of a computation then the computation
+    // is complete. Constant arrays are always available (e.g. `array({1, 2,
+    // 3})`)
+    available
+  };
+
+  // Check if the array is safe to read.
+  bool is_available() const;
+
+  // Wait on the array to be available. After this `is_available` returns
+  // `true`.
+  void wait();
 
   Status status() const {
     return array_desc_->status;

mlx/backend/metal/event.cpp

@@ -27,4 +27,9 @@ void Event::signal() {
   static_cast<MTL::SharedEvent*>(raw_event().get())->setSignaledValue(value());
 }
 
+bool Event::is_signaled() const {
+  return static_cast<MTL::SharedEvent*>(raw_event().get())->signaledValue() >=
+      value();
+}
+
 } // namespace mlx::core

mlx/backend/metal/metal.cpp

@@ -74,7 +74,7 @@ std::function<void()> make_task(array arr, bool signal) {
       arr.detach();
     }
     for (auto& out : outputs) {
-      out.set_status(array::Status::available);
+      out.set_status(array::Status::evaluated);
     }
 
     if (signal || d.get_command_buffer_ops(s.index) >= MAX_OPS_PER_BUFFER) {

mlx/backend/no_metal/event.cpp

@@ -36,4 +36,11 @@ void Event::signal() {
   ec->cv.notify_all();
 }
 
+bool Event::is_signaled() const {
+  auto ec = static_cast<EventCounter*>(raw_event().get());
+  {
+    std::lock_guard<std::mutex> lk(ec->mtx);
+    return (ec->value > value());
+  }
+}
 } // namespace mlx::core

mlx/event.h

@@ -20,6 +20,9 @@ class Event {
   // Signal the event at its current value
   void signal();
 
+  // Check if the event has been signaled at its current value
+  bool is_signaled() const;
+
   // Check if the event is valid
   bool valid() const {
     return event_ != nullptr;

mlx/transforms.cpp

@@ -79,7 +79,7 @@ array eval_impl(std::vector<array> outputs, bool async) {
           continue;
         }
 
-        if (!in.is_available()) {
+        if (in.status() == array::Status::unscheduled) {
           if (async && in.is_tracer()) {
             throw std::invalid_argument(
                 "[async_eval] Not allowed inside a graph transformation.");
@@ -115,7 +115,8 @@ array eval_impl(std::vector<array> outputs, bool async) {
       }
 
       // All inputs are done being processed, process this array
-      if (a.is_available() && !a.is_tracer() && a.has_primitive()) {
+      if ((a.status() != array::Status::unscheduled) && !a.is_tracer() &&
+          a.has_primitive()) {
         // If the array is evaluated and is no longer a tracer, detach it
         a.detach();
       } else if (a.status() == array::Status::unscheduled) {
@@ -208,14 +209,12 @@ void eval(std::vector<array> outputs) {
         return x.status() == array::Status::unscheduled;
       })) {
     for (auto& x : outputs) {
-      if (!x.is_available()) {
-        x.event().wait();
-      }
+      x.wait();
     }
     return;
   }
 
-  eval_impl(std::move(outputs), false).event().wait();
+  eval_impl(std::move(outputs), false).wait();
 }
 
 std::pair<std::vector<array>, std::vector<array>> vjp(