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

13
mlx/io/load.h
View File

@@ -14,6 +14,8 @@ namespace mlx::core {
namespace io {
ThreadPool& thread_pool();
class Reader {
public:
virtual bool is_open() const = 0;
@@ -43,10 +45,8 @@ class Writer {
class ParallelFileReader : public Reader {
public:
explicit ParallelFileReader(std::string file_path, int num_threads)
: fd_(open(file_path.c_str(), O_RDONLY)),
label_(std::move(file_path)),
thread_pool_(ThreadPool(num_threads)) {}
explicit ParallelFileReader(std::string file_path)
: fd_(open(file_path.c_str(), O_RDONLY)), label_(std::move(file_path)) {}
~ParallelFileReader() override {
close(fd_);
@@ -79,11 +79,10 @@ class ParallelFileReader : public Reader {
}
private:
// 4MB
static constexpr size_t batch_size_ = (1 << 22);
static constexpr size_t batch_size_ = 1 << 25;
static ThreadPool thread_pool_;
int fd_;
std::string label_;
ThreadPool thread_pool_;
};
class FileWriter : public Writer {