diff --git a/io/io.cc b/io/io.cc index 8614735d..87916614 100644 --- a/io/io.cc +++ b/io/io.cc @@ -13,16 +13,30 @@ // See the License for the specific language governing permissions and // limitations under the License. -// Safe to be first, does not include POSIX headers. -#include +// Feature-test macros must be defined before including any system header, +// including C++ headers such as , which may include . +// Replacement for the _BSD_SOURCE specified by preadv documentation (glibc). +#ifndef _DEFAULT_SOURCE +#define _DEFAULT_SOURCE +#endif +// Safe to include here: does not include any system headers. #include "hwy/detect_compiler_arch.h" + +#if HWY_OS_LINUX // Request POSIX 2008, including `pread()` and `posix_fadvise()`. This also -// implies `_POSIX_C_SOURCE`. +// implies `_POSIX_C_SOURCE`. Not done on BSD/Apple, where it would hide +// extensions such as `preadv` and `madvise`. #if !defined(_XOPEN_SOURCE) || _XOPEN_SOURCE < 700 #undef _XOPEN_SOURCE #define _XOPEN_SOURCE 700 // SUSv4 #endif +// FilePosix should only be compiled on non-Windows. It is easier to +// check this in source code because we support multiple build systems. Note +// that IOBatch at the end of this TU is still compiled on all platforms. +#include + +#endif // HWY_OS_LINUX // Make `off_t` 64-bit even on 32-bit systems. Works for Android >= r15c. #undef _FILE_OFFSET_BITS @@ -32,7 +46,9 @@ #include #include +#include #include +#include #include "io/io.h" #include "hwy/base.h" // HWY_ASSERT @@ -51,43 +67,39 @@ #define GEMMA_IO_FADVISE 0 #endif -// FilePosix should only be compiled on non-Windows. It is easier to -// check this in source code because we support multiple build systems. Note -// that IOBatch at the end of this TU is still compiled on all platforms. #if !HWY_OS_WIN #if GEMMA_IO_PREADV -// Replacement for the _BSD_SOURCE specified by preadv documentation. -#ifndef _DEFAULT_SOURCE -#define _DEFAULT_SOURCE -#endif -#include #include // preadv -#endif // GEMMA_IO_PREADV +#endif #include // open #include // IOV_MAX #include -#include // SEEK_END - unistd isn't enough for IDE. #include // Old OSX may require sys/types.h before sys/mman.h. #include // mmap -#include // O_RDONLY -#include // read, write, close +#include // fstat, O_RDONLY +#include // pread, pwrite, close // Placeholder for internal header, do not modify. #include "util/allocator.h" namespace gcpp { +HWY_INLINE_VAR constexpr size_t kMaxBytesPerSyscall = + 0x7FFFF000; // Linux limit + class FilePosix : public File { - int fd_ = 0; + int fd_ = -1; public: - explicit FilePosix(int fd) : fd_(fd) { HWY_ASSERT(fd > 0); } + explicit FilePosix(int fd) : fd_(fd) { HWY_ASSERT(fd >= 0); } ~FilePosix() override { - if (fd_ != 0) { - HWY_ASSERT(close(fd_) != -1); + if (fd_ >= 0) { + // Do not retry on EINTR: Linux releases the fd regardless, and retrying + // could close an fd that another thread was just given. + HWY_ASSERT(close(fd_) == 0 || errno == EINTR); } } @@ -96,20 +108,24 @@ class FilePosix : public File { uint64_t FileSize() const override { static_assert(sizeof(off_t) == 8, "64-bit off_t required"); - const off_t size = lseek(fd_, 0, SEEK_END); - if (size < 0) { + struct stat st; + if (fstat(fd_, &st) != 0 || st.st_size < 0) { return 0; } - return static_cast(size); + return static_cast(st.st_size); } bool Read(uint64_t offset, uint64_t size, void* to) const override { uint8_t* bytes = reinterpret_cast(to); uint64_t pos = 0; - for (;;) { + while (pos < size) { // pread seems to be faster than lseek + read when parallelized. - const auto bytes_read = pread(fd_, bytes + pos, size - pos, offset + pos); - if (bytes_read <= 0) { + const size_t want = static_cast( + HWY_MIN(size - pos, uint64_t{kMaxBytesPerSyscall})); + const ssize_t bytes_read = + pread(fd_, bytes + pos, want, static_cast(offset + pos)); + if (bytes_read < 0) { + if (errno == EINTR) continue; HWY_WARN( "Read failure at pos %zu within size %zu with offset %zu and " "errno %d\n", @@ -117,9 +133,9 @@ class FilePosix : public File { static_cast(offset), errno); break; } - pos += bytes_read; + if (bytes_read == 0) break; // EOF before reading `size` bytes. + pos += static_cast(bytes_read); HWY_ASSERT(pos <= size); - if (pos == size) break; } return pos == size; // success if managed to read desired size } @@ -127,10 +143,13 @@ class FilePosix : public File { bool Write(const void* from, uint64_t size, uint64_t offset) override { const uint8_t* bytes = reinterpret_cast(from); uint64_t pos = 0; - for (;;) { - const auto bytes_written = - pwrite(fd_, bytes + pos, size - pos, offset + pos); - if (bytes_written <= 0) { + while (pos < size) { + const size_t want = static_cast( + HWY_MIN(size - pos, uint64_t{kMaxBytesPerSyscall})); + const ssize_t bytes_written = + pwrite(fd_, bytes + pos, want, static_cast(offset + pos)); + if (bytes_written < 0) { + if (errno == EINTR) continue; HWY_WARN( "Write failure at pos %zu within size %zu with offset %zu and " "errno %d\n", @@ -138,15 +157,19 @@ class FilePosix : public File { static_cast(offset), errno); break; } - pos += bytes_written; + if (bytes_written == 0) break; + pos += static_cast(bytes_written); HWY_ASSERT(pos <= size); - if (pos == size) break; } return pos == size; // success if managed to write desired size } MapPtr Map() override { - const size_t mapping_size = FileSize(); + const uint64_t file_size = FileSize(); + if (file_size == 0 || file_size > static_cast(~size_t{0})) { + return MapPtr(); + } + const size_t mapping_size = static_cast(file_size); // No `MAP_POPULATE` because we do not want to wait for I/O, and // `MAP_NONBLOCK` is not guaranteed. `MAP_HUGETLB` fails. `MAP_SHARED` is // more efficient than `MAP_PRIVATE`; the main difference is that the former @@ -174,13 +197,14 @@ HWY_MAYBE_UNUSED extern std::unique_ptr OpenFileGoogle( const Path& filename, const char* mode); std::unique_ptr OpenFileOrNull(const Path& filename, const char* mode) { - std::unique_ptr file; // OpenFileGoogle omitted - if (file) return file; - const bool is_read = mode[0] != 'w'; + // Tries POSIX first because local paths (including FUSE mounts) are faster. const int flags = is_read ? O_RDONLY : O_CREAT | O_RDWR | O_TRUNC; - const int fd = open(filename.path.c_str(), flags, 0644); - if (fd < 0) return file; + int fd; + do { + fd = open(filename.path.c_str(), flags, 0644); + } while (fd < 0 && errno == EINTR); + if (fd < 0) return std::unique_ptr(); #if GEMMA_IO_FADVISE if (is_read) { @@ -230,11 +254,15 @@ IOBatch::IOBatch(uint64_t offset, size_t key_idx) spans_.reserve(kMaxSpans); } -// Returns true if the batch was full; if so, call again on the new batch. bool IOBatch::Add(void* mem, size_t bytes) { - if (spans_.size() >= kMaxSpans) return false; - if (total_bytes_ + bytes > 0x7FFFF000) return false; // Linux limit - spans_.push_back({.mem = mem, .bytes = bytes}); + // Always accept the first span, otherwise callers would loop forever. + if (!spans_.empty()) { + if (spans_.size() >= kMaxSpans) return false; +#if !HWY_OS_WIN + if (total_bytes_ + bytes > kMaxBytesPerSyscall) return false; +#endif + } + spans_.push_back(IOSpan{mem, bytes}); total_bytes_ += bytes; return true; } @@ -247,19 +275,46 @@ uint64_t IOBatch::Read(const File& file) const { HWY_ASSERT(!spans_.empty()); #if GEMMA_IO_PREADV - if (file.Handle() != -1) { - ssize_t bytes_read; - for (;;) { - bytes_read = - preadv(file.Handle(), reinterpret_cast(spans_.data()), - static_cast(spans_.size()), offset_); - if (bytes_read >= 0) break; - if (errno == EINTR) continue; // signal: retry - HWY_WARN("preadv(%d) for %4zu spans from offset %12zu failed, errno %d.", - file.Handle(), spans_.size(), offset_, errno); - return 0; + const int fd = file.Handle(); + if (fd != -1) { + static_assert(sizeof(IOSpan) == sizeof(iovec), "IOSpan must match iovec"); + const IOSpan* spans = spans_.data(); + size_t num_spans = spans_.size(); + std::vector remaining; // Only used after a short read. + uint64_t pos = 0; + while (pos < total_bytes_) { + const ssize_t bytes_read = preadv( + fd, reinterpret_cast(spans), + static_cast(num_spans), static_cast(offset_ + pos)); + if (bytes_read < 0) { + if (errno == EINTR) continue; // signal: retry + HWY_WARN( + "preadv(%d) for %4zu spans from offset %12zu failed, errno %d.", fd, + num_spans, static_cast(offset_ + pos), errno); + return 0; + } + if (bytes_read == 0) return 0; // EOF before reading `total_bytes_`. + pos += static_cast(bytes_read); + HWY_ASSERT(pos <= total_bytes_); + if (pos == total_bytes_) break; + + // Short read: rebuild the list of not yet (fully) read spans. + remaining.clear(); + uint64_t skip = pos; + for (const IOSpan& span : spans_) { + if (skip >= span.bytes) { + skip -= span.bytes; + continue; + } + const size_t skip_s = static_cast(skip); + remaining.push_back(IOSpan{static_cast(span.mem) + skip_s, + span.bytes - skip_s}); + skip = 0; + } + spans = remaining.data(); + num_spans = remaining.size(); } - return static_cast(bytes_read); + return total_bytes_; } #endif // GEMMA_IO_PREADV diff --git a/io/io.h b/io/io.h index d051715a..0eceb84c 100644 --- a/io/io.h +++ b/io/io.h @@ -95,15 +95,16 @@ class IOBatch { // Reserves memory in `spans_`. `key_idx` identifies the blob/tensor. explicit IOBatch(uint64_t offset, size_t key_idx); - // The next `bytes` will be read from file into `mem`. - // Returns true if the batch was full; if so, call again on the new batch. + // The next `bytes` will be read from file into `mem`. Returns false if the + // batch is full; if so, call again on a new batch. Always succeeds for the + // first span, regardless of its size, so this cannot loop forever. bool Add(void* mem, size_t bytes); uint64_t Offset() const { return offset_; } uint64_t TotalBytes() const { return total_bytes_; } size_t KeyIdx() const { return key_idx_; } - // Returns the total number of bytes read, or 0 if any I/O failed. + // Returns `TotalBytes()` if all spans were completely read, otherwise 0. // Thread-compatible. uint64_t Read(const File& file) const; diff --git a/io/io_win.cc b/io/io_win.cc index 34773d33..3c8e3586 100644 --- a/io/io_win.cc +++ b/io/io_win.cc @@ -21,6 +21,9 @@ #include #include +#include +#include + #include "io/io.h" #include "util/allocator.h" #include "hwy/base.h" // HWY_ASSERT @@ -33,12 +36,77 @@ #include namespace gcpp { +namespace { + +// Converts UTF-8 to UTF-16 (including the terminating null), or returns an +// empty vector on failure. +std::vector Utf8ToWide(const std::string& utf8) { + std::vector wide; + if (utf8.size() > static_cast(INT32_MAX)) return wide; + const int len_with_null = static_cast(utf8.size()) + 1; + const int num_wide = MultiByteToWideChar( + CP_UTF8, MB_ERR_INVALID_CHARS, utf8.c_str(), len_with_null, nullptr, 0); + if (num_wide <= 0) return wide; + wide.resize(static_cast(num_wide)); + if (MultiByteToWideChar(CP_UTF8, MB_ERR_INVALID_CHARS, utf8.c_str(), + len_with_null, wide.data(), num_wide) != num_wide) { + wide.clear(); + } + return wide; +} + +// Manual-reset event for waiting on one overlapped request. Each concurrent +// request requires its own event because waiting on the file handle itself +// cannot distinguish between multiple outstanding requests. +class OverlappedEvent { + public: + OverlappedEvent() + : event_(CreateEventW(/*lpEventAttributes=*/nullptr, + /*bManualReset=*/TRUE, /*bInitialState=*/FALSE, + /*lpName=*/nullptr)) {} + ~OverlappedEvent() { + if (event_ != nullptr) HWY_ASSERT(CloseHandle(event_)); + } + OverlappedEvent(const OverlappedEvent&) = delete; + OverlappedEvent& operator=(const OverlappedEvent&) = delete; + HANDLE Get() const { return event_; } + + private: + HANDLE event_; +}; + +// Issues a single ReadFile (`is_write` = false) or WriteFile at `offset` and +// waits for it to complete. Returns the number of bytes transferred, or 0 on +// failure or EOF. +DWORD TransferAt(HANDLE file, HANDLE event, bool is_write, uint64_t offset, + void* bytes, DWORD want) { + OVERLAPPED overlapped = {}; + overlapped.Offset = static_cast(offset & 0xFFFFFFFFu); + overlapped.OffsetHigh = static_cast(offset >> 32); + // ReadFile/WriteFile reset the event to non-signaled when they begin. + overlapped.hEvent = event; + // For overlapped handles, the byte count must be obtained from + // GetOverlappedResult, hence pass nullptr. + const BOOL ok = is_write ? WriteFile(file, bytes, want, nullptr, &overlapped) + : ReadFile(file, bytes, want, nullptr, &overlapped); + if (!ok && GetLastError() != ERROR_IO_PENDING) return 0; // incl. EOF + DWORD got = 0; + if (!GetOverlappedResult(file, &overlapped, &got, /*bWait=*/TRUE)) return 0; + return got; +} + +} // namespace + +// Opened with FILE_FLAG_OVERLAPPED so that concurrent Read/Write calls on the +// same handle can proceed in parallel. Otherwise, the I/O manager would +// serialize them via a per-file-object lock, even with explicit offsets. Each +// call still waits for its own completion, hence the interface is synchronous. class FileWin : public File { HANDLE hFile_ = INVALID_HANDLE_VALUE; public: - FileWin(HANDLE hFile) : hFile_(hFile) { + explicit FileWin(HANDLE hFile) : hFile_(hFile) { HWY_ASSERT(hFile != INVALID_HANDLE_VALUE); } ~FileWin() override { @@ -51,82 +119,80 @@ class FileWin : public File { bool IsAppendOnly() const override { return false; } uint64_t FileSize() const override { - DWORD hi; - const DWORD lo = GetFileSize(hFile_, &hi); - if (lo == INVALID_FILE_SIZE) return 0; - return (static_cast(hi) << 32) | lo; + LARGE_INTEGER size; + if (!GetFileSizeEx(hFile_, &size) || size.QuadPart < 0) return 0; + return static_cast(size.QuadPart); } bool Read(uint64_t offset, uint64_t size, void* to) const override { - uint8_t* bytes = reinterpret_cast(to); - OVERLAPPED overlapped = {0}; - // Loop is required because ReadFile[Ex] size argument is 32-bit. - while (size != 0) { - overlapped.Offset = offset & 0xFFFFFFFF; - overlapped.OffsetHigh = (offset >> 32) & 0xFFFFFFFF; - const DWORD want = - static_cast(HWY_MIN(size, uint64_t{0xFFFFFFFF})); - DWORD got; - if (!ReadFile(hFile_, bytes, want, &got, &overlapped)) { - if (GetLastError() != ERROR_HANDLE_EOF) { - return false; - } - } - offset += got; - bytes += got; - size -= got; - } - return true; // read everything => success + return Transfer(/*is_write=*/false, offset, size, to); } bool Write(const void* from, uint64_t size, uint64_t offset) override { - const uint8_t* bytes = reinterpret_cast(from); - OVERLAPPED overlapped = {0}; - // Loop is required because WriteFile[Ex] size argument is 32-bit. - while (size != 0) { - overlapped.Offset = offset & 0xFFFFFFFF; - overlapped.OffsetHigh = (offset >> 32) & 0xFFFFFFFF; - const DWORD want = - static_cast(HWY_MIN(size, uint64_t{0xFFFFFFFF})); - DWORD got; - if (!WriteFile(hFile_, bytes, want, &got, &overlapped)) { - if (GetLastError() != ERROR_HANDLE_EOF) { - return false; - } - } - offset += got; - bytes += got; - size -= got; - } - return true; // wrote everything => success + // WriteFile does not modify the buffer. + return Transfer(/*is_write=*/true, offset, size, const_cast(from)); } MapPtr Map() override { - if (hFile_ == INVALID_HANDLE_VALUE) return MapPtr(); + const uint64_t file_size = FileSize(); + if (file_size == 0 || file_size > static_cast(~size_t{0})) { + return MapPtr(); + } // Size=0 means the entire file. HANDLE hMapping = - CreateFileMappingA(hFile_, nullptr, PAGE_READONLY, 0, 0, nullptr); + CreateFileMappingW(hFile_, nullptr, PAGE_READONLY, 0, 0, nullptr); + if (hMapping == nullptr) return MapPtr(); // Offset zero and size=0 means the entire file. void* ptr = MapViewOfFile(hMapping, FILE_MAP_READ, 0, 0, 0); - if (!ptr) return MapPtr(); + if (!ptr) { + HWY_ASSERT(CloseHandle(hMapping)); + return MapPtr(); + } return MapPtr(static_cast(ptr), DeleterFunc([hMapping](void* ptr) { HWY_ASSERT(UnmapViewOfFile(ptr)); HWY_ASSERT(CloseHandle(hMapping)); })); } + + private: + bool Transfer(bool is_write, uint64_t offset, uint64_t size, + void* bytes) const { + if (size == 0) return true; + // One event per call (not per chunk) suffices because chunks are + // sequential. Creating it is cheap relative to disk I/O. + const OverlappedEvent event; + if (event.Get() == nullptr) return false; + uint8_t* pos = static_cast(bytes); + // Loop is required because the ReadFile/WriteFile size argument is 32-bit. + while (size != 0) { + const DWORD want = static_cast(HWY_MIN(size, uint64_t{1} << 30)); + const DWORD got = + TransferAt(hFile_, event.Get(), is_write, offset, pos, want); + if (got == 0) return false; + pos += got; + offset += got; + size -= got; + } + return true; + } }; // FileWin std::unique_ptr OpenFileOrNull(const Path& filename, const char* mode) { const bool is_read = mode[0] != 'w'; - const DWORD flags = - FILE_ATTRIBUTE_NORMAL | (is_read ? FILE_FLAG_SEQUENTIAL_SCAN : 0); - const DWORD access = is_read ? GENERIC_READ : GENERIC_WRITE; - const DWORD share = is_read ? FILE_SHARE_READ : 0; + const DWORD flags = FILE_ATTRIBUTE_NORMAL | FILE_FLAG_OVERLAPPED | + (is_read ? FILE_FLAG_SEQUENTIAL_SCAN : 0); + const DWORD access = is_read ? GENERIC_READ : (GENERIC_READ | GENERIC_WRITE); + const DWORD share = + is_read ? FILE_SHARE_READ + : (FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE); const DWORD create = is_read ? OPEN_EXISTING : CREATE_ALWAYS; - const HANDLE hFile = CreateFileA(filename.path.c_str(), access, share, - nullptr, create, flags, nullptr); + // Path is UTF-8; CreateFileA would interpret it in the ANSI code page. + const std::vector wide = Utf8ToWide(filename.path); + if (wide.empty()) return std::unique_ptr(); + const HANDLE hFile = + CreateFileW(wide.data(), access, share, nullptr, create, flags, nullptr); if (hFile == INVALID_HANDLE_VALUE) return std::unique_ptr(); return std::make_unique(hFile); }