diff --git a/io/io.cc b/io/io.cc index 2f479b21..de568f92 100644 --- a/io/io.cc +++ b/io/io.cc @@ -14,6 +14,8 @@ // limitations under the License. // Safe to be first, does not include POSIX headers. +#include + #include "hwy/detect_compiler_arch.h" // Request POSIX 2008, including `pread()` and `posix_fadvise()`. This also // implies `_POSIX_C_SOURCE`. @@ -28,6 +30,7 @@ #include +#include #include #include @@ -261,7 +264,43 @@ uint64_t IOBatch::Read(const File& file) const { } #endif // GEMMA_IO_PREADV - // preadv disabled or no handle: use normal reads (higher kernel overhead). + // 1. Fast Path: Check if destination memory across all spans is contiguous. + bool is_contiguous = true; + const uint8_t* next_expected = + reinterpret_cast(spans_[0].mem); + for (const IOSpan& span : spans_) { + if (reinterpret_cast(span.mem) != next_expected) { + is_contiguous = false; + break; + } + next_expected += span.bytes; + } + + // Zero-copy direct read: 1 single bulk RPC directly into destination memory. + if (is_contiguous) { + if (file.Read(offset_, total_bytes_, spans_[0].mem)) { + return total_bytes_; + } + } else { + // 2. Slow Path (Strided Padding): Single bulk read + scatter memcpy. + std::unique_ptr temp_buf(new (std::nothrow) + uint8_t[total_bytes_]); + if (temp_buf && file.Read(offset_, total_bytes_, temp_buf.get())) { + const uint8_t* src = temp_buf.get(); + for (const IOSpan& span : spans_) { + std::memcpy(span.mem, src, span.bytes); + src += span.bytes; + } + return total_bytes_; + } else { + HWY_WARN( + "Failed to read %zu bytes from offset %zu in file with %zu spans", + total_bytes_, offset_, spans_.size()); + } + } + + // 3. Fallback if bulk read or temp allocation fails: per-span sequential + // reads. uint64_t total = 0; uint64_t offset = offset_; for (const IOSpan& span : spans_) {