zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-07-23 01:51:18 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Avoid io timeout for large arrays (#1442)

Browse Source
This commit is contained in:
Awni Hannun 2024-09-27 13:32:14 -07:00 committed by GitHub
parent 718aea3f1d
commit 11354d5bff
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 7 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
mlx/backend/metal/primitives.cpp
View File

@ -200,13 +200,19 @@ void Full::eval_gpu(const std::vector<array>& inputs, array& out) {
void Load::eval_gpu(const std::vector<array>& inputs, array& out) {
out.set_data(allocator::malloc_or_wait(out.nbytes()));
auto read_task = [out = out,
offset = offset_,
reader = reader_,
swap_endianness = swap_endianness_]() mutable {
load(out, offset, reader, swap_endianness);
};
// Limit the size that the command buffer will wait on to avoid timing out
// on the event (<4 seconds).
if (out.nbytes() > (1 << 28)) {
read_task();
return;
}
auto fut = io::thread_pool().enqueue(std::move(read_task)).share();
auto signal_task = [out = out, fut = std::move(fut)]() {
fut.wait();