XFELReader: a deflated frame is inflated and converted a piece at a time, while it is in the cache

The reader inflated a whole deflated chunk (4 MB for a 1 Mpx int32 frame, 2 MB for its uint16 per-frame
mask) into a buffer and then converted it in a second pass; with 16 reading threads those buffers never
stay in the cache, so the conversion ran at memory speed. A deflated, unshuffled chunk is now inflated in
128 kB pieces (InflateInPieces, the zlib stream API - the stream's checksum is checked as before), and
each piece is converted into the image, or marked in the mask, before the next is inflated. Shuffled,
bitshuffle/LZ4, HDF5-read and single-panel (placed) frames are converted as a whole, as before.

Bit-identical: the FNV pixel hashes of 2000 frames on the five reader configurations of the previous
commit are unchanged, and a new test reads 300 x 300 frames (several pieces, a final partial one) with a
per-frame mask and a gain pixel by pixel. Reader alone, 16 threads, 3000 frames: 2.43 -> 1.73 s (the
previous commit's state -> this). rugnux on the 14,445-frame CXI run, against rc178 (A/B, two runs each):
geometry refinement sample 11.6 -> 5.1 s, image pass 27.7 -> 12.9 s, whole run 67.7 -> 46.2 s; the
geometry refinement, module placement and tilt log lines are identical.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu

reader/JFJochXFELReader.cpp

@@ -558,7 +558,7 @@ void JFJochXFELReader::Place() {
                 v = mask_file.ReadVector<uint32_t>(m.data, {0, 0}, {frame_height_, frame_width_});
             } else
                 v = FileAt(0).ReadVector<uint32_t>(m.data, {0, 0}, {frame_height_, frame_width_});
-            MarkBad(m, v.data(), masked_by_file);
+            MarkBad(m, v.data(), 0, v.size(), masked_by_file);
         }
     }
 
@@ -761,7 +761,7 @@ void JFJochXFELReader::Place() {
 }
 
 template <class T>
-void JFJochXFELReader::MarkBad(const xfel::Mask &m, const T *v, std::vector<uint8_t> &bad) const {
+void JFJochXFELReader::MarkBad(const xfel::Mask &m, const T *v, size_t first, size_t n, std::vector<uint8_t> &bad) const {
     int64_t min_fs = 0, max_fs = static_cast<int64_t>(frame_width_) - 1;
     int64_t min_ss = 0, max_ss = static_cast<int64_t>(frame_height_) - 1;
     for (const auto &p: geom_.panels)
@@ -773,12 +773,16 @@ void JFJochXFELReader::MarkBad(const xfel::Mask &m, const T *v, std::vector<uint
     // be m's or bad's own, every pixel would read them again, and the loop would not be vectorised.
     const uint32_t good = m.good, bad_bits = m.bad;
     const bool nonzero_good = m.nonzero_good;
-    const size_t width = max_fs - min_fs + 1;
-    for (int64_t ss = min_ss; ss <= max_ss; ss++) {
-        const size_t first = static_cast<size_t>(ss) * frame_width_ + min_fs;
-        const T *in = v + first;
-        uint8_t *out = bad.data() + first;
-        for (size_t i = 0; i < width; i++)
+    const size_t last = first + n;   // the pixels [first, last) of the frame are those v holds
+    for (int64_t ss = std::max<int64_t>(min_ss, first / frame_width_);
+         ss <= max_ss && static_cast<size_t>(ss) * frame_width_ < last; ss++) {
+        const size_t begin = std::max(static_cast<size_t>(ss) * frame_width_ + min_fs, first);
+        const size_t end = std::min(static_cast<size_t>(ss) * frame_width_ + max_fs + 1, last);
+        if (end <= begin)
+            continue;
+        const T *in = v + (begin - first);
+        uint8_t *out = bad.data() + begin;
+        for (size_t i = 0; i < end - begin; i++)
             out[i] |= xfel::Mask::Bad(in[i], good, bad_bits, nonzero_good);
     }
 }
@@ -822,9 +826,12 @@ void JFJochXFELReader::Close() {
     dataset_.reset();
 }
 
-template <class Buffer>
+template <class Buffer, class Piece>
 bool JFJochXFELReader::ReadStored(const Frame &frame, const Layout &layout, Buffer &stored,
-                                  bool keep_compressed) const {
+                                  bool keep_compressed, Piece piece) const {
+    constexpr bool in_pieces = !std::is_same_v<Piece, std::nullptr_t>;
+    const size_t n = frame_width_ * frame_height_;
+    const size_t size = layout.stored_size;
     const std::vector<hsize_t> start = frame.index >= 0
                                        ? std::vector<hsize_t>{static_cast<hsize_t>(frame.index), 0, 0}
                                        : std::vector<hsize_t>{0, 0};
@@ -833,6 +840,7 @@ bool JFJochXFELReader::ReadStored(const Frame &frame, const Layout &layout, Buff
     // Under the lock only where the chunk is (a lookup in its index); the read and the decoding not.
     std::shared_ptr<const HDF5ImageSource::RawFile> raw_file;
     uint64_t address = 0, chunk_size = 0;
+    bool read_by_hdf5 = false;
     {
         std::unique_lock ul(hdf5_mutex);
         HDF5DataSet dataset(FileAt(frame.file), frame.dataset);
@@ -858,10 +866,16 @@ bool JFJochXFELReader::ReadStored(const Frame &frame, const Layout &layout, Buff
                 stored.resize(v.size() * sizeof(float));
                 std::memcpy(stored.data(), v.data(), stored.size());
             }
-            return false;
+            read_by_hdf5 = true;
+        } else {
+            raw_file = RawFileAt(frame.file);
+            address += user_block_[frame.file];   // HDF5's chunk addresses count from the end of the user block
         }
-        raw_file = RawFileAt(frame.file);
-        address += user_block_[frame.file];   // HDF5's chunk addresses count from the end of the user block
+    }
+    if (read_by_hdf5) {
+        if constexpr (in_pieces)
+            piece(0, stored.data(), n);
+        return false;
     }
     if (keep_compressed && !layout.deflated) {
         stored.resize(chunk_size);
@@ -872,8 +886,11 @@ bool JFJochXFELReader::ReadStored(const Frame &frame, const Layout &layout, Buff
     compressed.resize(chunk_size);
     raw_file->ReadAt(compressed.data(), compressed.size(), address);
 
-    const size_t n = frame_width_ * frame_height_;
-    const size_t size = layout.stored_size;
+    if (layout.deflated && !layout.shuffled && in_pieces) {
+        if constexpr (in_pieces)
+            InflateInPieces(compressed, n * size, size, frame, piece);
+        return false;
+    }
     if (layout.deflated) {
         if (!layout.shuffled)
             stored.resize(n * size);
@@ -895,26 +912,65 @@ bool JFJochXFELReader::ReadStored(const Frame &frame, const Layout &layout, Buff
         stored.resize(n * size);
         JFJochDecompressPtr(stored.data(), layout.algorithm, compressed.data(), compressed.size(), n, size);
     }
+    if constexpr (in_pieces)
+        piece(0, stored.data(), n);
     return false;
 }
 
+template <class Piece>
+void JFJochXFELReader::InflateInPieces(const std::vector<uint8_t> &compressed, size_t bytes, size_t size,
+                                       const Frame &frame, Piece piece) const {
+    // A piece fits the core's own cache with the window inflate keeps, so it is converted before it leaves.
+    constexpr size_t PIECE_BYTES = 128 * 1024;
+    thread_local std::vector<uint8_t> out(PIECE_BYTES);
+    z_stream z{};
+    if (inflateInit(&z) != Z_OK)
+        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Cannot start to inflate");
+    z.next_in = const_cast<uint8_t *>(compressed.data());
+    z.avail_in = compressed.size();
+    size_t done = 0;
+    int ret = Z_OK;
+    while (ret == Z_OK) {
+        z.next_out = out.data();
+        z.avail_out = PIECE_BYTES;
+        ret = inflate(&z, Z_NO_FLUSH);
+        const size_t got = PIECE_BYTES - z.avail_out;
+        // A piece is whole pixels: inflate fills it before it returns, except at the end of the frame.
+        if ((ret != Z_OK && ret != Z_STREAM_END) || done + got > bytes || got % size != 0)
+            break;
+        if (got > 0)
+            piece(done / size, out.data(), got / size);
+        done += got;
+    }
+    inflateEnd(&z);
+    if (ret != Z_STREAM_END || done != bytes)
+        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
+                              "Cannot inflate " + frame.dataset + " of " + files_[frame.file]);
+}
+
 template <class Out, class In, class Buffer, class Convert>
-void JFJochXFELReader::PlaceInto(const In *in, const std::vector<uint8_t> &masked, Out err, Buffer &buffer,
+void JFJochXFELReader::PlaceInto(const Frame &frame, const std::vector<uint8_t> &masked, Out err, Buffer &buffer,
                                  Convert convert) const {
     buffer.resize(width_ * height_ * sizeof(Out));
     auto out = reinterpret_cast<Out *>(buffer.data());
-    if (blocks_.empty()) {   // the stored frame, as stored
-        const size_t n = width_ * height_;
-        if (masked.empty())
-            for (size_t i = 0; i < n; i++)
-                out[i] = static_cast<Out>(convert(in[i]));
-        else {
-            const uint8_t *m = masked.data();
-            for (size_t i = 0; i < n; i++)
-                out[i] = static_cast<Out>(Pick(m[i] != 0, convert(in[i]), err));
-        }
+    thread_local std::vector<uint8_t> stored;
+    if (blocks_.empty()) {   // the stored frame, as stored: converted a piece at a time, as it is read
+        ReadStored(frame, layout_, stored, false, [&](size_t first, const uint8_t *bytes, size_t n) {
+            const auto *in = reinterpret_cast<const In *>(bytes);
+            Out *o = out + first;
+            if (masked.empty())
+                for (size_t i = 0; i < n; i++)
+                    o[i] = static_cast<Out>(convert(in[i]));
+            else {
+                const uint8_t *m = masked.data() + first;
+                for (size_t i = 0; i < n; i++)
+                    o[i] = static_cast<Out>(Pick(m[i] != 0, convert(in[i]), err));
+            }
+        });
         return;
     }
+    ReadStored(frame, layout_, stored, false);
+    const auto *in = reinterpret_cast<const In *>(stored.data());
     std::fill(out, out + width_ * height_, Out{0});   // between the panels
     for (const auto &b: blocks_)
         for (int64_t ss = 0; ss < b.height; ss++) {
@@ -944,9 +1000,6 @@ CompressedImage JFJochXFELReader::DecodeInto(int64_t image_number, Buffer &buffe
         return CompressedImage(buffer.data(), buffer.size(), width_, height_, stored_mode,
                                compressed ? layout_.algorithm : CompressionAlgorithm::NO_COMPRESSION);
     }
-    thread_local std::vector<uint8_t> stored;
-    ReadStored(frame, layout_, stored, false);
-
     // A pixel a per-frame mask marks in this frame is handed out as no measurement.
     std::vector<uint8_t> masked;
     if (!frame_masks_.empty()) {
@@ -957,11 +1010,12 @@ CompressedImage JFJochXFELReader::DecodeInto(int64_t image_number, Buffer &buffe
             const Frame mask{frame.file, ReplaceEvent(m.data, frame.group), frame.index, frame.group};
             if (mask_layout.plain_type) {
                 thread_local std::vector<uint8_t> mask_stored;
-                ReadStored(mask, mask_layout, mask_stored, false);
-                if (mask_layout.stored_size == 2)
-                    MarkBad(m, reinterpret_cast<const uint16_t *>(mask_stored.data()), masked);
-                else
-                    MarkBad(m, reinterpret_cast<const uint32_t *>(mask_stored.data()), masked);
+                ReadStored(mask, mask_layout, mask_stored, false, [&](size_t first, const uint8_t *bytes, size_t n) {
+                    if (mask_layout.stored_size == 2)
+                        MarkBad(m, reinterpret_cast<const uint16_t *>(bytes), first, n, masked);
+                    else
+                        MarkBad(m, reinterpret_cast<const uint32_t *>(bytes), first, n, masked);
+                });
             } else {
                 std::vector<uint32_t> v;
                 {
@@ -971,7 +1025,7 @@ CompressedImage JFJochXFELReader::DecodeInto(int64_t image_number, Buffer &buffe
                                                                       {1, frame_height_, frame_width_})
                             : FileAt(frame.file).ReadVector<uint32_t>(mask.dataset, {0, 0}, {frame_height_, frame_width_});
                 }
-                MarkBad(m, v.data(), masked);
+                MarkBad(m, v.data(), 0, v.size(), masked);
             }
         }
     }
@@ -1002,25 +1056,24 @@ CompressedImage JFJochXFELReader::DecodeInto(int64_t image_number, Buffer &buffe
                 return Pick(v == std::numeric_limits<decltype(v)>::min(), to_int32(v), INT32_MIN);
             return to_int32(v);
         };
-        const uint8_t *p = stored.data();
         if (layout_.stored_size == 2 && layout_.stored_signed)
-            PlaceInto<int32_t>(reinterpret_cast<const int16_t *>(p), masked, INT32_MIN, buffer, convert);
+            PlaceInto<int32_t, int16_t>(frame, masked, INT32_MIN, buffer, convert);
         else if (layout_.stored_size == 2)
-            PlaceInto<int32_t>(reinterpret_cast<const uint16_t *>(p), masked, INT32_MIN, buffer, convert);
+            PlaceInto<int32_t, uint16_t>(frame, masked, INT32_MIN, buffer, convert);
         else if (layout_.stored_signed)
-            PlaceInto<int32_t>(reinterpret_cast<const int32_t *>(p), masked, INT32_MIN, buffer, convert);
+            PlaceInto<int32_t, int32_t>(frame, masked, INT32_MIN, buffer, convert);
         else
-            PlaceInto<int32_t>(reinterpret_cast<const uint32_t *>(p), masked, INT32_MIN, buffer, convert);
+            PlaceInto<int32_t, uint32_t>(frame, masked, INT32_MIN, buffer, convert);
     };
 
     // Floats become photons, rounded. Integers stay integers: divided only by a gain other than one, and a
     // 16-bit signed frame - as Cheetah writes - is handed out in its own type.
     if (layout_.stored_float || !layout_.plain_type)
-        PlaceInto<int32_t>(reinterpret_cast<const float *>(stored.data()), masked, INT32_MIN, buffer, photons);
+        PlaceInto<int32_t, float>(frame, masked, INT32_MIN, buffer, photons);
     else if (gain != 1.0f)
         integers(photons);
     else if (layout_.stored_size == 2 && layout_.stored_signed) {
-        PlaceInto<int16_t>(reinterpret_cast<const int16_t *>(stored.data()), masked, INT16_MIN, buffer, as_is);
+        PlaceInto<int16_t, int16_t>(frame, masked, INT16_MIN, buffer, as_is);
         return image(CompressedImageMode::Int16);
     } else
         integers([](auto v) { return static_cast<int32_t>(v); });

reader/JFJochXFELReader.h

@@ -3,6 +3,7 @@
 
 #pragma once
 
+#include <cstddef>
 #include <memory>
 #include <optional>
 #include <string>
@@ -76,9 +77,10 @@ class JFJochXFELReader : public JFJochReader {
     double clen_m_ = 0, distance_mm_ = 0, photon_energy_eV_ = 0;
     // Masks stored with one plane per frame (a Cheetah CXI's per-shot mask), applied as each frame is read.
     std::vector<xfel::Mask> frame_masks_;
-    // Marks in bad every stored pixel the mask m, read as v, leaves out (on its panel, where it has one).
+    // Marks in bad every stored pixel the mask m, read as v, leaves out (on its panel, where it has one); v holds
+    // the n pixels of the frame from first on.
     template <class T>
-    void MarkBad(const xfel::Mask &m, const T *v, std::vector<uint8_t> &bad) const;
+    void MarkBad(const xfel::Mask &m, const T *v, size_t first, size_t n, std::vector<uint8_t> &bad) const;
 
     // Where every frame is stored as one chunk of a plain little-endian number type - deflated (and
     // optionally byte-shuffled), as Cheetah writes it, or bitshuffle/LZ4 (or zstd), as Jungfraujoch does -
@@ -123,12 +125,20 @@ class JFJochXFELReader : public JFJochReader {
 
     // The frame's pixels, stored as layout says, as stored (plain_type), or as float, into stored; with
     // keep_compressed, a chunk decoded by CompressedImage (not deflated) is left as stored in the file, and
-    // true returned.
-    template <class Buffer>
-    bool ReadStored(const Frame &frame, const Layout &layout, Buffer &stored, bool keep_compressed) const;
-    // The stored pixels as the image handed out (as stored, or placed), as Out, each through convert; a pixel masked marks is err.
+    // true returned. Given piece, the pixels are handed to piece(first pixel, their bytes, pixels) instead: a
+    // deflated chunk a piece at a time as it is inflated (InflateInPieces), any other at once.
+    template <class Buffer, class Piece = std::nullptr_t>
+    bool ReadStored(const Frame &frame, const Layout &layout, Buffer &stored, bool keep_compressed,
+                    Piece piece = nullptr) const;
+    // The deflated stream compressed, bytes long once inflated, to piece a piece of whole pixels of size bytes
+    // at a time.
+    template <class Piece>
+    void InflateInPieces(const std::vector<uint8_t> &compressed, size_t bytes, size_t size, const Frame &frame,
+                         Piece piece) const;
+    // The frame's stored pixels, as In, as the image handed out (as stored, or placed), as Out, each through
+    // convert; a pixel masked marks is err.
     template <class Out, class In, class Buffer, class Convert>
-    void PlaceInto(const In *in, const std::vector<uint8_t> &masked, Out err, Buffer &buffer, Convert convert) const;
+    void PlaceInto(const Frame &frame, const std::vector<uint8_t> &masked, Out err, Buffer &buffer, Convert convert) const;
     template <class Buffer>
     CompressedImage DecodeInto(int64_t image_number, Buffer &buffer) const;
 

tests/XFELReaderTest.cpp

@@ -4,6 +4,7 @@
 #include <catch2/catch_all.hpp>
 #include <algorithm>
 #include <climits>
+#include <cmath>
 #include <cstdint>
 #include <cstring>
 #include <filesystem>
@@ -383,6 +384,78 @@ bad_p1/panel = p1
     std::filesystem::remove_all(dir);
 }
 
+TEST_CASE("JFJochXFELReader_DeflatedInPieces", "[portable]") {
+    // Frames large enough to be inflated in several pieces (a 300 x 300 int32 frame and its uint16 per-frame
+    // mask, deflated without the shuffle), read with a gain: every pixel is the stored value in photons,
+    // rounded half away from zero, or no measurement where the per-frame mask marks it.
+    const auto dir = std::filesystem::path("xfel_reader_in_pieces");
+    std::filesystem::create_directories(dir);
+    const std::string name = (dir / "frames.cxi").string();
+    constexpr size_t side = 300, n = side * side;
+    std::vector<int32_t> frames(2 * n);
+    for (size_t i = 0; i < frames.size(); i++)
+        frames[i] = static_cast<int32_t>(i % 1001) - 500;
+    std::vector<uint16_t> mask(2 * n, 0);
+    for (size_t i = 0; i < mask.size(); i += 997)
+        mask[i] = 0x4;
+    {
+        const hid_t f = H5Fcreate(name.c_str(), H5F_ACC_TRUNC, H5P_DEFAULT, H5P_DEFAULT);
+        for (const char *g: {"entry_1", "entry_1/data_1", "entry_1/instrument_1", "entry_1/instrument_1/source_1"})
+            H5Gclose(H5Gcreate2(f, g, H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT));
+        WriteFrames(f, "entry_1/data_1/data", H5T_NATIVE_INT32, frames.data(), {2, side, side}, false);
+        WriteFrames(f, "entry_1/data_1/mask", H5T_NATIVE_UINT16, mask.data(), {2, side, side}, false);
+        WriteVector(f, "entry_1/instrument_1/source_1/energy", std::vector<double>(2, 9000.0 * 1.602176634e-19));
+        H5Fclose(f);
+    }
+    const std::string geom = (dir / "frames.geom").string();
+    std::ofstream(geom) << R"(
+clen = 0.1
+res = 10000
+dim0 = %
+dim1 = ss
+dim2 = fs
+data = /entry_1/data_1/data
+mask = /entry_1/data_1/mask
+mask_good = 0x0
+mask_bad = 0xffff
+p0/min_fs = 0
+p0/max_fs = 149
+p0/min_ss = 0
+p0/max_ss = 299
+p0/fs = x
+p0/ss = y
+p0/corner_x = -150
+p0/corner_y = -150
+p1/min_fs = 150
+p1/max_fs = 299
+p1/min_ss = 0
+p1/max_ss = 299
+p1/fs = x
+p1/ss = y
+p1/corner_x = 0
+p1/corner_y = -150
+)";
+
+    JFJochXFELReader reader;
+    REQUIRE_NOTHROW(reader.ReadFiles({name}, geom, 2.0f));
+    REQUIRE(reader.GetNumberOfImages() == 2);
+    for (size_t image = 0; image < 2; image++) {
+        auto raw = reader.GetRawImage(static_cast<int64_t>(image));
+        REQUIRE(raw);
+        REQUIRE(raw->image.GetWidth() * raw->image.GetHeight() == n);
+        const auto *pixels = reinterpret_cast<const int32_t *>(raw->image_buffer.data());
+        size_t wrong = 0;
+        for (size_t i = 0; i < n; i++) {
+            const size_t k = image * n + i;
+            const int32_t expected = mask[k] ? INT32_MIN : static_cast<int32_t>(std::lround(frames[k] / 2.0f));
+            wrong += pixels[i] != expected;
+        }
+        CHECK(wrong == 0);
+    }
+    reader.Close();
+    std::filesystem::remove_all(dir);
+}
+
 namespace {
 
 // An LCLS-I XTC stream of the shape JFJochXTCReader reads: a Configure datagram, then one event (an