zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Fix array is_available race cases (#1468)

Browse Source
This commit is contained in:
Awni Hannun
2024-10-07 19:13:50 -07:00
committed by GitHub
parent 9b12093739
commit 3274c6a087
9 changed files with 77 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
mlx/array.cpp
View File

@@ -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();
}
}