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

Even Even Faster IO (#1374)

Browse Source 
* even more faster io

* make reader pool static

* make python reader thread safe

* one more optimization
This commit is contained in:
Awni Hannun
2024-08-29 16:05:40 -07:00
committed by GitHub
parent 28be4de7c2
commit dba2bd1105
6 changed files with 73 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files

26
mlx/io/load.cpp
View File

@@ -298,11 +298,18 @@ array load(std::shared_ptr<io::Reader> in_stream, StreamOrDevice s) {
/** Load array from file in .npy format */
array load(std::string file, StreamOrDevice s) {
return load(std::make_shared<io::ParallelFileReader>(std::move(file), 4), s);
return load(std::make_shared<io::ParallelFileReader>(std::move(file)), s);
}
namespace io {
ThreadPool& thread_pool() {
static ThreadPool pool_{4};
return pool_;
}
ThreadPool ParallelFileReader::thread_pool_{4};
void ParallelFileReader::read(char* data, size_t n) {
while (n != 0) {
auto m = ::read(fd_, data, std::min(n, static_cast<size_t>(INT32_MAX)));
@@ -330,11 +337,18 @@ void ParallelFileReader::read(char* data, size_t n, size_t offset) {
};
std::vector<std::future<bool>> futs;
while (n != 0) {
size_t m = std::min(batch_size_, n);
futs.emplace_back(thread_pool_.enqueue(readfn, offset, m, data));
data += m;
n -= m;
offset += m;
if (n < batch_size_) {
if (!readfn(offset, n, data)) {
throw std::runtime_error("[read] Unable to read from file.");
}
break;
} else {
size_t m = batch_size_;
futs.emplace_back(thread_pool_.enqueue(readfn, offset, m, data));
data += m;
n -= m;
offset += m;
}
}
for (auto& f : futs) {
if (!f.get()) {