From 1d16dcd2b9272c3537996643b7aba39d3ace73b8 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Tue, 25 Aug 2026 01:08:38 +0200 Subject: [PATCH] Read a chunk without zeroing it first, and hold no frame the fused decoder never writes Three costs before and around the image loop. Every image allocated a fresh buffer for its compressed chunk and resized it, which value-initialises, and the read then overwrote every byte. At a few megabytes a chunk the allocation is large enough to be mapped rather than reused, so the zeroing was page-fault bound and cost more than the read it preceded - twenty gigabytes of it over a long sweep. The buffer now uses an allocator that does not construct, and the two HDF5 read paths are templated on the allocator so every existing caller compiles unchanged. The rebind is deliberate: without it the vector base rebinds to the default allocator and the zeroing quietly returns. The bitshuffle decoder allocated a whole uncompressed frame in its constructor - seventy megabytes a worker, five hundred and fifty across the loop - for the route that decodes the shuffled image separately. That route is taken only when a bitshuffle block is too large for the fused kernel, which neither writer this pipeline reads produces, so on a real frame the buffer is allocated, never touched, and freed. It is now allocated where it is used. The comment two lines below already warned against sizing a buffer from the uncompressed size; the line above it had not been given the same treatment. The first call into cuFFT pays the library's one-time initialisation, and it landed in the middle of the first pass with nothing to overlap it. It is now forced on a background thread at startup, alongside the file open and the mapping build, in the manner the shadow finder already uses. Finally, the detector mask was copied into the start message whether or not a file would carry it, which a merging run does not. It is filled where a writer is constructed - both places one is constructed, the second being the fallback that writes a process file when nothing indexed. Faster on eleven of thirty-eight crystals and slower on none; the whole rotation test set falls from four minutes thirty to four minutes seventeen, with each binary repeating itself to within half a per cent. Space groups thirty-five of thirty-eight and no failures throughout, and every column of the comparison table is identical. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01EGpGdgmJ8MyY9pCGWjktyi --- common/CompressedImage.h | 21 ++++++++++++++++++ .../image_preprocessing/BSLZ4DecoderGPU.cu | 15 +++++++------ image_analysis/indexing/IndexerThreadPool.cpp | 10 +++++++++ image_analysis/indexing/IndexerThreadPool.h | 5 +++++ reader/HDF5ImageSource.cpp | 17 ++------------ reader/HDF5ImageSource.h | 19 +++++++++++++--- reader/JFJochReaderImage.h | 2 +- rugnux/Rugnux.cpp | 14 +++++++++--- rugnux/rugnux_cli.cpp | 6 +++++ writer/HDF5Objects.cpp | 17 -------------- writer/HDF5Objects.h | 22 +++++++++++++++++-- 11 files changed, 100 insertions(+), 48 deletions(-) diff --git a/common/CompressedImage.h b/common/CompressedImage.h index 79e2265e..188f47d3 100644 --- a/common/CompressedImage.h +++ b/common/CompressedImage.h @@ -5,12 +5,33 @@ #include #include +#include +#include #include +#include #include #include "../compression/CompressionAlgorithmEnum.h" #include "ColorScale.h" +// std::vector value-initialises whatever it resizes into. A buffer that the very next read fills in +// full has no use for that, and on a compressed image chunk it is megabytes of memset per image. This +// allocator default-initialises instead, which for bytes is no initialisation at all. The rebind is +// not optional: without it the container rebinds to std::allocator and the zeroing comes back. +template +struct NoInitAllocator : std::allocator { + NoInitAllocator() = default; + template NoInitAllocator(const NoInitAllocator &) {} + template struct rebind { using other = NoInitAllocator; }; + template void construct(U *p) { ::new (static_cast(p)) U; } + template void construct(U *p, Args &&...args) { + ::new (static_cast(p)) U(std::forward(args)...); + } +}; + +// Bytes of an image as they are stored: sized by the reader and then overwritten by the read. +using RawByteBuffer = std::vector>; + enum class CompressedImageMode {Int8, Int16, Int32, Uint8, Uint16, Uint32, RGB, Float16, Float32, Float64}; CompressedImageMode CalcImageMode(size_t byte_depth, bool is_float, bool is_signed); diff --git a/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu b/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu index 1bbe3ba3..df84b0fc 100644 --- a/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu +++ b/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu @@ -118,7 +118,6 @@ bool BSLZ4DecoderGPU::Supports(const CompressedImage &image) { BSLZ4DecoderGPU::BSLZ4DecoderGPU(size_t in_max_uncompressed_bytes, std::shared_ptr in_stream) : stream(std::move(in_stream)), max_uncompressed_bytes(in_max_uncompressed_bytes) { - gpu_shuffled = CudaDevicePtr(max_uncompressed_bytes); gpu_status = CudaDevicePtr(1); host_status = CudaHostPtr(1); // The compressed buffer and the descriptors are grown to fit the first image instead of being @@ -126,13 +125,15 @@ BSLZ4DecoderGPU::BSLZ4DecoderGPU(size_t in_max_uncompressed_bytes, std::shared_p // tens; sizing this from the UNCOMPRESSED size cost ~73 MB per worker to hold ~4 MB. } -// The stored pixel depth is fixed within a dataset, so in practice this runs once - but it is not -// promised anywhere, and sizing for the widest type instead would hold twice the memory a 16-bit -// detector needs. Grown with slack because cudaMalloc and cudaFree synchronise the whole device. +// gpu_shuffled holds a whole uncompressed frame - 72 MB at 18 Mpx, per worker - and only the +// DecodeShuffled() route ever writes it. That route is taken when a bitshuffle block is too large for +// the fused kernel, which neither writer this pipeline reads produces, so on a real frame the buffer +// is never touched. So allocate it the first time it is actually asked for, at the size the caller +// declared, rather than in the constructor. void BSLZ4DecoderGPU::EnsureUncompressedCapacity(size_t bytes) { - if (bytes <= max_uncompressed_bytes) + if (gpu_shuffled.get() && bytes <= max_uncompressed_bytes) return; - const size_t want = std::max(bytes, max_uncompressed_bytes + max_uncompressed_bytes / 2); + const size_t want = std::max(bytes, max_uncompressed_bytes); cuda_err(cudaStreamSynchronize(*stream)); gpu_shuffled = CudaDevicePtr(want); max_uncompressed_bytes = want; @@ -176,7 +177,6 @@ BSLZ4ShuffledImage BSLZ4DecoderGPU::PrepareChunk(const CompressedImage &image) { if (clen < 12) throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 chunk shorter than its header"); - EnsureUncompressedCapacity(total_bytes); if (be64(src) != total_bytes) throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 header size does not match the image"); @@ -259,6 +259,7 @@ BSLZ4ShuffledImage BSLZ4DecoderGPU::PrepareChunk(const CompressedImage &image) { } BSLZ4ShuffledImage BSLZ4DecoderGPU::DecodeShuffled(const CompressedImage &image) { + EnsureUncompressedCapacity(image.GetUncompressedSize()); BSLZ4ShuffledImage ret = PrepareChunk(image); if (ret.nblocks > 0) { lz4_decode_blocks<<<(ret.nblocks * 32 + 255) / 256, 256, 0, *stream>>>( diff --git a/image_analysis/indexing/IndexerThreadPool.cpp b/image_analysis/indexing/IndexerThreadPool.cpp index ee5d11e3..30e6952e 100644 --- a/image_analysis/indexing/IndexerThreadPool.cpp +++ b/image_analysis/indexing/IndexerThreadPool.cpp @@ -14,6 +14,16 @@ #include "FFTIndexerCPU.h" #endif +void WarmUpCuFFT() { +#ifdef JFJOCH_USE_CUDA + if (get_gpu_count() == 0) + return; + cufftHandle plan = 0; + if (cufftPlan1d(&plan, 1024, CUFFT_C2C, 1) == CUFFT_SUCCESS) + cufftDestroy(plan); +#endif +} + // The indexer for one RESOLVED algorithm, or nullptr if this build/host cannot serve it. static std::unique_ptr MakeIndexer(IndexingAlgorithmEnum algorithm, const IndexingSettings &settings) { #ifdef JFJOCH_USE_CUDA diff --git a/image_analysis/indexing/IndexerThreadPool.h b/image_analysis/indexing/IndexerThreadPool.h index 8d980683..2fb319e9 100644 --- a/image_analysis/indexing/IndexerThreadPool.h +++ b/image_analysis/indexing/IndexerThreadPool.h @@ -35,6 +35,11 @@ // that can never be dispatched to. enum class IndexerConstruction { Preconstruct, OnFirstUse }; +// Force cuFFT's one-time library initialisation - 0.3 s measured, paid by whichever call happens to +// be first - so it can be got out of the way on a background thread at startup rather than landing in +// the middle of the first pass with nothing overlapping it. No-op without a GPU. +void WarmUpCuFFT(); + class IndexerThread { struct TaskInput { const DiffractionExperiment &experiment; diff --git a/reader/HDF5ImageSource.cpp b/reader/HDF5ImageSource.cpp index c207ccfa..afca58fb 100644 --- a/reader/HDF5ImageSource.cpp +++ b/reader/HDF5ImageSource.cpp @@ -152,21 +152,8 @@ HDF5ImageSource::PrepareDirectRead(const HDF5ImageLocator::Location &loc) const ds.width, ds.height, ds.mode, ds.algorithm}; } -CompressedImage HDF5ImageSource::ReadDirect(std::vector &buffer, const DirectChunk &chunk) { +CompressedImage HDF5ImageSource::ReadDirect(RawByteBuffer &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 &buffer, - const HDF5ImageLocator::Location &loc) const { - const auto &ds = GetDataset(loc); - const std::vector start = {static_cast(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}; + return {buffer.data(), buffer.size(), chunk.width, chunk.height, chunk.mode, chunk.algorithm}; } diff --git a/reader/HDF5ImageSource.h b/reader/HDF5ImageSource.h index 97cdaf30..a84334e2 100644 --- a/reader/HDF5ImageSource.h +++ b/reader/HDF5ImageSource.h @@ -68,8 +68,21 @@ public: // file that holds a legacy/VDS image's per-image metadata. HDF5ImageLocator::Location Resolve(int64_t global) const; - // Read the pixels at a resolved location into a CompressedImage backed by `buffer`. - CompressedImage ReadImageAt(std::vector &buffer, const HDF5ImageLocator::Location &loc) const; + // Read the pixels at a resolved location into a CompressedImage backed by `buffer`. Templated on + // the allocator so a caller can hand over a buffer that does not zero what it is about to + // overwrite (RawByteBuffer). + template + CompressedImage ReadImageAt(std::vector &buffer, const HDF5ImageLocator::Location &loc) const { + const auto &ds = GetDataset(loc); + const std::vector start = {static_cast(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.data(), buffer.size(), ds.width, ds.height, ds.mode, ds.algorithm}; + } // Ask HDF5 where image `loc` is in the file rather than asking it for the image. This is a // lookup in the chunk index and nothing else - no read - so the mutex is held for a fraction of @@ -83,7 +96,7 @@ public: // Read what PrepareDirectRead() found. Touches no HDF5 and no shared state, so it needs no // mutex; this is the whole point of the two-step split. - static CompressedImage ReadDirect(std::vector &buffer, const DirectChunk &chunk); + static CompressedImage ReadDirect(RawByteBuffer &buffer, const DirectChunk &chunk); std::vector GetSourceMapping(uint64_t first_image, std::optional image_count, diff --git a/reader/JFJochReaderImage.h b/reader/JFJochReaderImage.h index 113270b8..b300bc9a 100644 --- a/reader/JFJochReaderImage.h +++ b/reader/JFJochReaderImage.h @@ -25,7 +25,7 @@ constexpr static int32_t MIN_REAL_PXL_VALUE = INT32_MIN + 3; constexpr static int32_t SATURATED_PXL_VALUE = INT32_MAX; struct JFJochReaderRawImage { - std::vector image_buffer; + RawByteBuffer image_buffer; CompressedImage image; }; diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index df714bb4..d068c9ca 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -1158,11 +1158,8 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b start_message.az_int_phi_bin_count = mapping.GetAzimuthalBinCount(); if (mapping.GetAzimuthalBinCount() > 1) start_message.az_int_bin_to_phi = mapping.GetBinToPhi(); - start_message.pixel_mask["default"] = pixel_mask_.GetMask(experiment_); if (full) { start_message.rois = experiment_.ROI().ExportMetadata(); - if (!experiment_.ROI().empty()) - start_message.roi_map = experiment_.ExportROIMap(); start_message.max_spot_count = experiment_.GetMaxSpotCount(); } start_message.master_suffix = "process"; @@ -1193,9 +1190,19 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b : (int64_t(1) << stored.bit_depth) - 1; // UINTx_MAX start_message.underload_value = stored.is_signed ? start_message.error_value.value() + 1 : 0; + // The full-detector mask is 4 B/px and the ROI map 2 B/px - tens of megabytes each - and nothing + // but the process file reads them. Copy them where a writer is actually built rather than here: + // the default merging path writes no process file at all. + auto fill_writer_maps = [&] { + start_message.pixel_mask["default"] = pixel_mask_.GetMask(experiment_); + if (full && !experiment_.ROI().empty()) + start_message.roi_map = experiment_.ExportROIMap(); + }; + std::unique_ptr writer; std::unique_ptr writer_queue; if (write_files && config_.write_process_h5) { + fill_writer_maps(); writer = std::make_unique(start_message, /*check_overwrite_at_start=*/true, /*trusted_path=*/true); // Deep enough that a worker never waits for the writer in the normal case, shallow enough that @@ -2976,6 +2983,7 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b if (!writer && full && write_files && !cancelled_ && result.indexing_rate.value_or(0.0f) <= 0.0f) { logger.Warning("No image indexed, so there are no merged reflections to write - writing " "{}_process.h5 with the per-image analysis instead", config_.output_prefix); + fill_writer_maps(); writer = std::make_unique(start_message, /*check_overwrite_at_start=*/true, /*trusted_path=*/true); } diff --git a/rugnux/rugnux_cli.cpp b/rugnux/rugnux_cli.cpp index 63c445a4..926a149b 100644 --- a/rugnux/rugnux_cli.cpp +++ b/rugnux/rugnux_cli.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -24,6 +25,7 @@ #include "../common/print_license.h" #include "../common/JFJochMath.h" #include "../image_analysis/geom_refinement/Calibrants.h" +#include "../image_analysis/indexing/IndexerThreadPool.h" #include "../image_analysis/LoadFCalcFromMtz.h" #include "../image_analysis/UpdateReflectionResolution.h" #include "../image_analysis/WriteReflections.h" @@ -1205,6 +1207,10 @@ static int RunRugnux(int argc, char **argv) { logger.Info("Using space group {} (number {})", space_group->hm, space_group_number.value()); } + // Off the critical path, so cuFFT's one-time initialisation overlaps the file open, the mask read + // and the beam-stop pre-scan instead of landing inside the first pass. + auto cufft_warmup = std::async(std::launch::async, WarmUpCuFFT); + // 1. Read Input File JFJochHDF5Reader reader; try { diff --git a/writer/HDF5Objects.cpp b/writer/HDF5Objects.cpp index 10847e37..a234f217 100644 --- a/writer/HDF5Objects.cpp +++ b/writer/HDF5Objects.cpp @@ -950,23 +950,6 @@ HDF5DataSet::~HDF5DataSet() { } } -void HDF5DataSet::ReadDirectChunk(std::vector &val, const std::vector &offset) { - if (offset.size() != ndim) - throw JFJochException(JFJochExceptionCategory::HDF5, "Inconsistent dimension settings"); - - hsize_t chunk_bytes; - if (H5Dget_chunk_storage_size(id, offset.data(), &chunk_bytes) < 0) - throw JFJochException(JFJochExceptionCategory::HDF5, "Error checking chunk size"); - - val.resize(chunk_bytes); - uint32_t filters; - size_t buf_size = val.size(); - if (H5Dread_chunk(id, H5P_DEFAULT, offset.data(), &filters, val.data(), &buf_size)) - throw JFJochException(JFJochExceptionCategory::HDF5, - "Error reading dataset (with direct chunks) to HDF5 file"); -} - - void RegisterHDF5Filter() { if (H5Zregister(bshuf_H5Filter) < 0) throw JFJochException(JFJochExceptionCategory::HDF5, "Cannot register Bitshuffle filter"); diff --git a/writer/HDF5Objects.h b/writer/HDF5Objects.h index 85250e9c..41fa261f 100644 --- a/writer/HDF5Objects.h +++ b/writer/HDF5Objects.h @@ -256,7 +256,24 @@ public: } HDF5DataSet& WriteDirectChunk(const void *val, hsize_t data_size, const std::vector& offset); - void ReadDirectChunk(std::vector &val, const std::vector& offset); + + // Templated on the allocator so a caller can hand over a buffer that does not zero what it is + // about to overwrite (RawByteBuffer). + template void ReadDirectChunk(std::vector &val, const std::vector& offset) { + if (offset.size() != ndim) + throw JFJochException(JFJochExceptionCategory::HDF5, "Inconsistent dimension settings"); + + hsize_t chunk_bytes; + if (H5Dget_chunk_storage_size(id, offset.data(), &chunk_bytes) < 0) + throw JFJochException(JFJochExceptionCategory::HDF5, "Error checking chunk size"); + + val.resize(chunk_bytes); + uint32_t filters; + size_t buf_size = val.size(); + if (H5Dread_chunk(id, H5P_DEFAULT, offset.data(), &filters, val.data(), &buf_size)) + throw JFJochException(JFJochExceptionCategory::HDF5, + "Error reading dataset (with direct chunks) to HDF5 file"); + } HDF5DataSet& Flush(); void SetExtent(const std::vector& dims); @@ -272,7 +289,8 @@ public: return output; } - void ReadVectorToU8(std::vector &v, const std::vector& slab_start, const std::vector& slab_size) const { + template + void ReadVectorToU8(std::vector &v, const std::vector& slab_start, const std::vector& slab_size) const { HDF5DataType data_type(*this); HDF5DataSpace mem_space(slab_size); HDF5DataSpace file_space(*this);