Ask HDF5 where an image is, then read it without the lock

Two things every worker thread of an offline run did inside the global HDF5 mutex, per image. It opened /entry/data/data and asked it for its dataspace, its datatype and its creation plist, then asked those for the rank, the dimensions, the chunking and the compression. All of that is a property of the file and identical for all of its images, so it is now resolved once when the file is first touched. And it read the pixels - megabytes of them, with the lock held, which is what turned a worker per hardware thread into a queue. HDF5 can say where a chunk lives instead - address and byte count, a lookup in the chunk index with no read attached - so that is all it is asked for now, and the bytes are fetched after the lock is dropped, with a positional read that any number of threads can make through one handle at once. Chunk addresses count from the end of the user block, so its size is added; zero for anything this project writes, not for every file. A file that is not one chunk per image, or a chunk that was never written and exists only as a fill value, still goes the old way - only HDF5 knows what those read as. On a 16 Mpx rotation dataset with the process file being written, the per-image loop at 48 workers goes 12.4 s -> 6.8 s, and stops getting slower as workers are added: 8 workers were faster than 48 before, and are not now. Where no process file is written the same loop only improves ~1%, because this machine has 1.5 TB of RAM and held the whole 7 GB test set in page cache - the read was never the expensive part here. It is where the cache is cold or the filesystem is remote. Battery 9m45s, space group 21/24, no failures, unchanged. The Windows path uses ReadFile with an OVERLAPPED offset for the same reason pread is used elsewhere: it takes the offset as an argument rather than moving a shared file position, so the viewer keeps building under MSVC and gets the same concurrency. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-16 01:08:53 -04:00
parent 09cb9e2be9
commit f5b3193253
5 changed files with 225 additions and 42 deletions
@@ -4,11 +4,64 @@
 #include "HDF5ImageSource.h"
 #include "../common/JFJochException.h"

+#ifdef _WIN32
+#include <windows.h>
+#else
+#include <fcntl.h>
+#include <unistd.h>
+#endif
+
+// Positional reads: pread() on POSIX, ReadFile() with an OVERLAPPED offset on Windows. Both take the
+// offset as an argument instead of moving a shared file position, which is what lets every worker
+// thread read through one handle at the same time.
+HDF5ImageSource::RawFile::RawFile(const std::string &path) {
+#ifdef _WIN32
+    HANDLE h = CreateFileA(path.c_str(), GENERIC_READ, FILE_SHARE_READ | FILE_SHARE_WRITE, nullptr,
+                           OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr);
+    handle_ = (h == INVALID_HANDLE_VALUE) ? -1 : reinterpret_cast<intptr_t>(h);
+#else
+    handle_ = ::open(path.c_str(), O_RDONLY);
+#endif
+}
+
+HDF5ImageSource::RawFile::~RawFile() {
+    if (handle_ == -1)
+        return;
+#ifdef _WIN32
+    CloseHandle(reinterpret_cast<HANDLE>(handle_));
+#else
+    ::close(static_cast<int>(handle_));
+#endif
+}
+
+void HDF5ImageSource::RawFile::ReadAt(void *dst, size_t size, uint64_t address) const {
+    auto *out = static_cast<uint8_t *>(dst);
+    size_t done = 0;
+    while (done < size) {
+#ifdef _WIN32
+        OVERLAPPED ov{};
+        ov.Offset = static_cast<DWORD>((address + done) & 0xFFFFFFFFULL);
+        ov.OffsetHigh = static_cast<DWORD>((address + done) >> 32);
+        DWORD got = 0;
+        const bool ok = ReadFile(reinterpret_cast<HANDLE>(handle_), out + done,
+                                 static_cast<DWORD>(size - done), &got, &ov);
+        const long long n = ok ? static_cast<long long>(got) : -1;
+#else
+        const long long n = ::pread(static_cast<int>(handle_), out + done, size - done, address + done);
+#endif
+        if (n <= 0)
+            throw JFJochException(JFJochExceptionCategory::HDF5, "Error reading image chunk from file");
+        done += static_cast<size_t>(n);
+    }
+}
+
 void HDF5ImageSource::Configure(HDF5ImageLocator::Layout layout) {
+    dataset_cache_.clear();
    locator_.Configure(std::move(layout));
 }

 void HDF5ImageSource::Clear() {
+    dataset_cache_.clear();
    locator_.Clear();
 }

@@ -23,45 +76,90 @@ std::vector<HDF5DataSourceMessage> HDF5ImageSource::GetSourceMapping(uint64_t fi
    return locator_.GetSourceMapping(first_image, image_count, total_images, stride);
 }

-CompressedImage HDF5ImageSource::ReadImageAt(std::vector<uint8_t> &buffer,
-                                             const HDF5ImageLocator::Location &loc) const {
-    return LoadImageDataset(buffer, *loc.file, loc.local_index);
-}
+const HDF5ImageSource::OpenDataset &
+HDF5ImageSource::GetDataset(const HDF5ImageLocator::Location &loc) const {
+    if (auto it = dataset_cache_.find(loc.file.get()); it != dataset_cache_.end())
+        return it->second;

-CompressedImage HDF5ImageSource::LoadImageDataset(std::vector<uint8_t> &tmp, HDF5Object &file, hsize_t number) {
-    std::vector<hsize_t> start = {static_cast<hsize_t>(number), 0, 0};
+    OpenDataset entry;
+    entry.file = loc.file;
+    entry.dataset = std::make_unique<HDF5DataSet>(*loc.file, "/entry/data/data");

-    const std::string dataset_name = "/entry/data/data";
-
-    HDF5DataSet dataset(file, dataset_name);
-    HDF5DataSpace dataspace(dataset);
-    HDF5DataType datatype(dataset);
-    HDF5Dcpl dcpl(dataset);
+    HDF5DataSpace dataspace(*entry.dataset);
+    HDF5DataType datatype(*entry.dataset);
+    HDF5Dcpl dcpl(*entry.dataset);

    if (dataspace.GetNumOfDimensions() != 3)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "/entry/data/data dataset must be 3D");

-    auto dim = dataspace.GetDimensions();
-
-    CompressionAlgorithm algorithm = CompressionAlgorithm::NO_COMPRESSION;
-    auto chunk_size = dcpl.GetChunking();
-
-    if ((chunk_size.size() == 3) && (chunk_size[0] == 1) && (chunk_size[1] == dim[1]) && (chunk_size[2] == dim[2])) {
-        dataset.ReadDirectChunk(tmp, start);
-        algorithm = dcpl.GetCompression();
-    } else {
-        dataset.ReadVectorToU8(tmp, start, {1, dim[1], dim[2]});
-        algorithm = CompressionAlgorithm::NO_COMPRESSION;
-    }
-
    if (datatype.IsFloat())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Float datasets not supported at this time");

-    return {
-        tmp, dim[2], dim[1],
-        CalcImageMode(datatype.GetElemSize(), datatype.IsFloat(), datatype.IsSigned()),
-        algorithm
-    };
+    auto dim = dataspace.GetDimensions();
+    entry.height = dim[1];
+    entry.width = dim[2];
+    entry.mode = CalcImageMode(datatype.GetElemSize(), datatype.IsFloat(), datatype.IsSigned());
+
+    auto chunk_size = dcpl.GetChunking();
+    entry.direct_chunk = (chunk_size.size() == 3) && (chunk_size[0] == 1)
+                         && (chunk_size[1] == dim[1]) && (chunk_size[2] == dim[2]);
+    if (entry.direct_chunk)
+        entry.algorithm = dcpl.GetCompression();
+
+    if (entry.direct_chunk && !loc.path.empty()) {
+        entry.raw = std::make_unique<RawFile>(loc.path);
+        if (!entry.raw->IsOpen())
+            entry.raw.reset();
+        hid_t fcpl = H5Fget_create_plist(loc.file->GetID());
+        if (fcpl >= 0) {
+            hsize_t user_block = 0;
+            if (H5Pget_userblock(fcpl, &user_block) >= 0)
+                entry.user_block = user_block;
+            H5Pclose(fcpl);
+        }
+    }
+
+    return dataset_cache_.emplace(loc.file.get(), std::move(entry)).first->second;
+}
+
+std::optional<HDF5ImageSource::DirectChunk>
+HDF5ImageSource::PrepareDirectRead(const HDF5ImageLocator::Location &loc) const {
+    const auto &ds = GetDataset(loc);
+    if (!ds.raw)
+        return {};
+
+    const hsize_t coord[3] = {static_cast<hsize_t>(loc.local_index), 0, 0};
+    unsigned filter_mask = 0;
+    haddr_t address = HADDR_UNDEF;
+    hsize_t size = 0;
+    if (H5Dget_chunk_info_by_coord(ds.dataset->GetID(), coord, &filter_mask, &address, &size) < 0)
+        return {};
+    // A chunk nobody ever wrote has no address and no bytes; only HDF5 knows it reads as the fill
+    // value, so hand those back to it.
+    if (address == HADDR_UNDEF || size == 0)
+        return {};
+
+    return DirectChunk{ds.raw.get(), ds.user_block + address, static_cast<uint32_t>(size),
+                       ds.width, ds.height, ds.mode, ds.algorithm};
+}
+
+CompressedImage HDF5ImageSource::ReadDirect(std::vector<uint8_t> &buffer, const DirectChunk &chunk) {
+    buffer.resize(chunk.size);
+    chunk.file->ReadAt(buffer.data(), chunk.size, chunk.address);
+    return {buffer, chunk.width, chunk.height, chunk.mode, chunk.algorithm};
+}
+
+CompressedImage HDF5ImageSource::ReadImageAt(std::vector<uint8_t> &buffer,
+                                             const HDF5ImageLocator::Location &loc) const {
+    const auto &ds = GetDataset(loc);
+    const std::vector<hsize_t> start = {static_cast<hsize_t>(loc.local_index), 0, 0};
+
+    if (ds.direct_chunk)
+        ds.dataset->ReadDirectChunk(buffer, start);
+    else
+        ds.dataset->ReadVectorToU8(buffer, start, {1, ds.height, ds.width});
+
+    return {buffer, ds.width, ds.height, ds.mode, ds.algorithm};
 }