zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-31 07:58:14 +08:00
Code Issues Actions 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

25
python/src/load.cpp
View File

@@ -138,6 +138,21 @@ class PyFileReader : public io::Reader {
void read(char* data, size_t n) override {
nb::gil_scoped_acquire gil;
_read(data, n);
}
void read(char* data, size_t n, size_t offset) override {
nb::gil_scoped_acquire gil;
seek_func_(offset, (int)std::ios_base::beg);
_read(data, n);
}
std::string label() const override {
return "python file object";
}
private:
void _read(char* data, size_t n) {
auto memview = PyMemoryView_FromMemory(data, n, PyBUF_WRITE);
nb::object bytes_read = readinto_func_(nb::handle(memview));
@@ -146,16 +161,6 @@ class PyFileReader : public io::Reader {
}
}
void read(char* data, size_t n, size_t offset) override {
seek(offset);
read(data, n);
}
std::string label() const override {
return "python file object";
}
private:
nb::object pyistream_;
nb::object readinto_func_;
nb::object seek_func_;