From 9c5141c7c0f47015f4d827551955fc3b86cef684 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 28 Sep 2026 02:00:50 +0200 Subject: [PATCH] CPU pixel pipeline: decode, preprocess and ring pass per bitshuffle block; flag rings in the local sweep The CPU-only image loop is DRAM-bound on 16M frames: decode, preprocess, the adaptive finder's plain ring pass and FlagRings each streamed the whole frame through memory. - JFJochDecompressHperfBlocks hands each decoded bitshuffle block to a callback; with no output buffer the block is unshuffled into a reused block-sized scratch (JFJochDecompressBlocks). - MXAnalysisWithoutFPGA::PreprocessCPU preprocesses each block into the int32 buffer (ImagePreprocessorCPU::AnalyzeBlock) and, when the fused CPU finder runs, puts it through the plain ring pass + fused azint (AdaptiveSpotFinderCPU::AccumulateRingsBlock) while it is in cache. Detect() then starts from those sums. The per-worker decompression buffer is no longer allocated for bitshuffled data. - FlagRings becomes FlagRow, called by DetectAt's first pass for row y+NBX just before that row enters the vertical sums; first_pass_needed is marked from each row's candidates at the same point. Exact: blocks arrive in pixel order, so the float azint sums see the same pixels in the same order; the per-pixel expressions are unchanged; everything else is integer. p.hkl, p.mtz, p_P1.mtz and p_unmerged.mtz byte-identical to rc173 on myob, cytc, lyso, sparse (CPU-only build), GPU myob identical (the GPU path does not take this route). CPU-only, 32 workers, under the gpulock on a shared (loaded) machine, base -> fused, two rounds (second in reversed order): myob 155.9 -> 97.7 s, 139.9 -> 88.1 s (loop 46.1 -> 26.8 s/pass; user 3690 -> 2221 s) cytc 220.8 -> 156.6 s, 216.6 -> 155.1 s (user 5160 -> 4128 s) lyso 134.5 -> 123.3 s, 69.5 -> 64.8 s peak RSS myob 15.2 -> 10.8 GB, cytc 12.6 -> 10.4 GB, lyso 7.0 -> 6.7 GB Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C --- compression/JFJochDecompress.h | 86 +++++++-- image_analysis/MXAnalysisWithoutFPGA.cpp | 48 +++-- image_analysis/MXAnalysisWithoutFPGA.h | 10 ++ .../ImagePreprocessorCPU.cpp | 45 ++--- .../ImagePreprocessorCPU.h | 7 +- .../spot_finding/AdaptiveSpotFinderCPU.cpp | 165 ++++++++++-------- .../spot_finding/AdaptiveSpotFinderCPU.h | 24 ++- .../spot_finding/ImageSpotFinderCPU.cpp | 57 ++++-- .../spot_finding/ImageSpotFinderCPU.h | 12 +- tests/ZSTDCompressorTest.cpp | 28 +++ 10 files changed, 332 insertions(+), 150 deletions(-) diff --git a/compression/JFJochDecompress.h b/compression/JFJochDecompress.h index 73c0518f4..c5a90df63 100644 --- a/compression/JFJochDecompress.h +++ b/compression/JFJochDecompress.h @@ -21,13 +21,19 @@ extern "C" { uint64_t bshuf_read_uint64_BE(const void* buf); }; -inline size_t JFJochDecompressHperfPtr(uint8_t *output, - CompressionAlgorithm algorithm, - const uint8_t *source, - size_t source_size, - size_t nelements, - size_t elem_size, - size_t block_size) { +// Decode the bitshuffle blocks one after another and hand each to block_done(first_element, data, +// nelements), in element order. With output set, each block is decoded in place into output; with +// output = nullptr, into a block-sized buffer that is reused for the next block, so a caller that +// consumes the block straight away reads it from cache and the whole image is never written out. +template +size_t JFJochDecompressHperfBlocks(uint8_t *output, + CompressionAlgorithm algorithm, + const uint8_t *source, + size_t source_size, + size_t nelements, + size_t elem_size, + size_t block_size, + F &&block_done) { if ((algorithm != CompressionAlgorithm::BSHUF_LZ4) && (algorithm != CompressionAlgorithm::BSHUF_ZSTD) && (algorithm != CompressionAlgorithm::BSHUF_ZSTD_RLE) && @@ -39,10 +45,11 @@ inline size_t JFJochDecompressHperfPtr(uint8_t *output, std::vector decompressed_block(block_size * elem_size); std::vector scratch(block_size * elem_size); + std::vector unshuffled(output ? 0 : block_size * elem_size); const uint8_t *src_ptr = source; const uint8_t *const source_end = source + source_size; - uint8_t *dst_ptr = output; + size_t first_element = 0; const size_t num_full_blocks = nelements / block_size; const size_t reminder_size = nelements - num_full_blocks * block_size; @@ -90,15 +97,17 @@ inline size_t JFJochDecompressHperfPtr(uint8_t *output, throw JFJochException(JFJochExceptionCategory::Compression, "Algorithm not supported"); } + uint8_t *dst_ptr = output ? output + first_element * elem_size : unshuffled.data(); if (JFJochBitUnshuffleBlock(reinterpret_cast(dst_ptr), decompressed_block.data(), scratch.data(), current_nelements, elem_size) < 0) throw JFJochException(JFJochExceptionCategory::Compression, "bitshuffle block decode error"); + block_done(first_element, dst_ptr, current_nelements); src_ptr += compressed_size; - dst_ptr += decompressed_size; + first_element += decompressed_size / elem_size; }; for (size_t i = 0; i < num_full_blocks; ++i) @@ -111,13 +120,26 @@ inline size_t JFJochDecompressHperfPtr(uint8_t *output, if (leftover_bytes > 0) { if (static_cast(source_end - src_ptr) < leftover_bytes) throw JFJochException(JFJochExceptionCategory::Compression, "Truncated trailing bytes"); + uint8_t *dst_ptr = output ? output + first_element * elem_size : unshuffled.data(); memcpy(dst_ptr, src_ptr, leftover_bytes); + block_done(first_element, dst_ptr, leftover_bytes / elem_size); src_ptr += leftover_bytes; } return static_cast(src_ptr - source); } +inline size_t JFJochDecompressHperfPtr(uint8_t *output, + CompressionAlgorithm algorithm, + const uint8_t *source, + size_t source_size, + size_t nelements, + size_t elem_size, + size_t block_size) { + return JFJochDecompressHperfBlocks(output, algorithm, source, source_size, nelements, elem_size, block_size, + [](size_t, const uint8_t *, size_t) {}); +} + // Plain LZ4, HDF5 filter 32004, as DECTRIS Eiger firmware 1.x wrote it. The framing is the same as // bitshuffle's - a 64-bit big-endian total size, a 32-bit big-endian block size IN BYTES, then each // block prefixed by its 32-bit big-endian compressed size - and the only difference is that no bit @@ -166,6 +188,40 @@ inline void JFJochDecompressLZ4Ptr(uint8_t *output, } } +// Check the 12-byte bitshuffle header and return the block size (in elements) it gives. +inline size_t JFJochBitshuffleHeaderBlockSize(const uint8_t *source, size_t source_size, + size_t nelements, size_t elem_size) { + // The header must be there before it can be read, and before source_size - 12 is handed to the + // decompressors. + if (source_size < 12) + throw JFJochException(JFJochExceptionCategory::Compression, "Buffer too short for the bitshuffle header"); + if (bshuf_read_uint64_BE(const_cast(source)) != nelements * elem_size) + throw JFJochException(JFJochExceptionCategory::Compression, "Mismatch in size"); + return bshuf_read_uint32_BE(source + 8) / elem_size; +} + +inline bool JFJochIsBitshuffle(CompressionAlgorithm algorithm) { + return (algorithm == CompressionAlgorithm::BSHUF_LZ4) + || (algorithm == CompressionAlgorithm::BSHUF_ZSTD) + || (algorithm == CompressionAlgorithm::BSHUF_ZSTD_RLE) + || (algorithm == CompressionAlgorithm::BSHUF_ZSTD_RLE_HUFF); +} + +// A bitshuffle-compressed image, handed block by block to block_done(first_element, data, nelements) +// in element order (see JFJochDecompressHperfBlocks), without the image ever being written out whole. +template +void JFJochDecompressBlocks(CompressionAlgorithm algorithm, + const uint8_t *source, + size_t source_size, + size_t nelements, + size_t elem_size, + F &&block_done) { + const size_t block_size = JFJochBitshuffleHeaderBlockSize(source, source_size, nelements, elem_size); + if (JFJochDecompressHperfBlocks(nullptr, algorithm, source + 12, source_size - 12, + nelements, elem_size, block_size, block_done) != source_size - 12) + throw JFJochException(JFJochExceptionCategory::Compression, "Decompression error"); +} + inline void JFJochDecompressPtr(uint8_t *output, CompressionAlgorithm algorithm, const uint8_t *source, @@ -179,16 +235,8 @@ inline void JFJochDecompressPtr(uint8_t *output, } size_t block_size = 0; - if (algorithm != CompressionAlgorithm::NO_COMPRESSION) { - // The 12-byte bitshuffle header must be there before it can be read, and before - // source_size - 12 is handed to the decompressors below. - if (source_size < 12) - throw JFJochException(JFJochExceptionCategory::Compression, "Buffer too short for the bitshuffle header"); - if (bshuf_read_uint64_BE(const_cast(source)) != nelements * elem_size) - throw JFJochException(JFJochExceptionCategory::Compression, "Mismatch in size"); - auto tmp = bshuf_read_uint32_BE(source + 8); - block_size = tmp / elem_size; - } + if (algorithm != CompressionAlgorithm::NO_COMPRESSION) + block_size = JFJochBitshuffleHeaderBlockSize(source, source_size, nelements, elem_size); switch (algorithm) { case CompressionAlgorithm::NO_COMPRESSION: diff --git a/image_analysis/MXAnalysisWithoutFPGA.cpp b/image_analysis/MXAnalysisWithoutFPGA.cpp index a861f50bf..ce3d2b173 100644 --- a/image_analysis/MXAnalysisWithoutFPGA.cpp +++ b/image_analysis/MXAnalysisWithoutFPGA.cpp @@ -50,7 +50,9 @@ MXAnalysisWithoutFPGA::MXAnalysisWithoutFPGA(const DiffractionExperiment &in_exp if (get_gpu_count() == 0) { #endif preprocessor_buffer = std::make_unique(experiment.GetPixelsNum()); - preprocessor = std::make_unique(in_experiment, in_mask); + auto cpu_preprocessor = std::make_unique(in_experiment, in_mask); + preprocessor_cpu = cpu_preprocessor.get(); + preprocessor = std::move(cpu_preprocessor); bragg_engine = std::make_unique(in_experiment); if (experiment.ROI().size() >= 1) roi = std::make_unique(experiment); @@ -128,7 +130,23 @@ void MXAnalysisWithoutFPGA::Analyze(DataMessage &output, } const auto compression_end_time = std::chrono::steady_clock::now(); - if (!decoded_on_device) { + // The fused GPU engine (rugnux offline, GPU, adaptive detection) produces the azimuthal profile as + // a byproduct of spot finding, so the separate azint pass is skipped in that case and the profile is + // lifted out of the finder below. + const bool fused = enable_fused_adaptive_gpu && spot_finding_settings.enable + && spot_finding_settings.adaptive_threshold && fused_adaptive != nullptr; + const bool fused_cpu = spot_finding_settings.enable && spot_finding_settings.adaptive_threshold + && fused_adaptive_cpu != nullptr; + + if (preprocessor_cpu) { + const auto preprocessing_start_time = std::chrono::steady_clock::now(); + ret = PreprocessCPU(output.image, fused_cpu); + const auto preprocessing_end_time = std::chrono::steady_clock::now(); + output.preprocessing_time_s = std::chrono::duration(preprocessing_end_time - preprocessing_start_time).count(); + // Decoded inside the preprocessing pass, whose time includes it. + if (output.image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION) + output.compression_time_s = 0.0f; + } else if (!decoded_on_device) { const uint8_t *image_ptr = Decompress(output.image); const auto decompressed_time = std::chrono::steady_clock::now(); if (output.image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION) @@ -149,14 +167,6 @@ void MXAnalysisWithoutFPGA::Analyze(DataMessage &output, output.preprocessing_time_s = total_s - decompress_s; } - // The fused GPU engine (rugnux offline, GPU, adaptive detection) produces the azimuthal profile as - // a byproduct of spot finding, so the separate azint pass is skipped in that case and the profile is - // lifted out of the finder below. - const bool fused = enable_fused_adaptive_gpu && spot_finding_settings.enable - && spot_finding_settings.adaptive_threshold && fused_adaptive != nullptr; - const bool fused_cpu = spot_finding_settings.enable && spot_finding_settings.adaptive_threshold - && fused_adaptive_cpu != nullptr; - if (!fused && !fused_cpu) { const auto azint_start_time = std::chrono::steady_clock::now(); AzInt().Run(*preprocessor_buffer, profile); @@ -397,6 +407,24 @@ const uint8_t *MXAnalysisWithoutFPGA::Decompress(const CompressedImage &image) { return image.GetUncompressedPtr(decompression_buffer); } +ImageStatistics MXAnalysisWithoutFPGA::PreprocessCPU(const CompressedImage &image, bool ring_pass) { + ImageStatistics ret{}; + if (ring_pass) + fused_adaptive_cpu->BeginRings(); + const auto process_block = [&](size_t first, const uint8_t *data, size_t n) { + preprocessor_cpu->AnalyzeBlock(*preprocessor_buffer, first, data, n, image.GetMode(), ret); + if (ring_pass) + fused_adaptive_cpu->AccumulateRingsBlock(*preprocessor_buffer, first, n); + }; + if (JFJochIsBitshuffle(image.GetCompressionAlgorithm())) + JFJochDecompressBlocks(image.GetCompressionAlgorithm(), image.GetCompressed(), image.GetCompressedSize(), + image.GetWidth() * image.GetHeight() * image.GetNumChannels(), image.GetByteDepth(), + process_block); + else + process_block(0, Decompress(image), npixels); + return ret; +} + void MXAnalysisWithoutFPGA::RunROIOnly(DataMessage &output) { output.roi.clear(); if (roi) diff --git a/image_analysis/MXAnalysisWithoutFPGA.h b/image_analysis/MXAnalysisWithoutFPGA.h index 20d4d5cb8..670a8b1e5 100644 --- a/image_analysis/MXAnalysisWithoutFPGA.h +++ b/image_analysis/MXAnalysisWithoutFPGA.h @@ -20,6 +20,7 @@ #include "IndexAndRefine.h" #include "image_preprocessing/ImagePreprocessor.h" #include "image_preprocessing/ImagePreprocessorBuffer.h" +#include "image_preprocessing/ImagePreprocessorCPU.h" class CudaStream; class AdaptiveSpotFinderGPU; @@ -32,6 +33,8 @@ class MXAnalysisWithoutFPGA { std::vector decompression_buffer; std::unique_ptr preprocessor; + // The preprocessor, where it is the CPU one. + ImagePreprocessorCPU *preprocessor_cpu = nullptr; size_t npixels; size_t xpixels; @@ -68,6 +71,13 @@ class MXAnalysisWithoutFPGA { // not compressed) and return where it landed. const uint8_t *Decompress(const CompressedImage &image); + // The CPU preprocessing. A bitshuffled image is decoded a block (~32 KB) at a time and each block + // is preprocessed while it is in cache, then - with ring_pass - put through the adaptive finder's + // plain ring pass as well, so the image is never written out whole before it is preprocessed and + // the preprocessed pixels are read back from cache, not from memory. The blocks come in pixel + // order, so every sum is taken in the same order as by separate passes. + ImageStatistics PreprocessCPU(const CompressedImage &image, bool ring_pass); + // Pixels outside the resolution limits, bit-packed. Built by the integration mapping, which is // shared by every worker's engine and hands out the same mask to all of them. std::shared_ptr> mask_resolution; diff --git a/image_analysis/image_preprocessing/ImagePreprocessorCPU.cpp b/image_analysis/image_preprocessing/ImagePreprocessorCPU.cpp index 8803a73db..cc4c9e263 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorCPU.cpp +++ b/image_analysis/image_preprocessing/ImagePreprocessorCPU.cpp @@ -11,57 +11,62 @@ ImagePreprocessorCPU::ImagePreprocessorCPU(const DiffractionExperiment &experime } ImageStatistics ImagePreprocessorCPU::Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *image_ptr, CompressedImageMode image_mode) { + if (processed_image.size() != npixels) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Processed image size mismatch"); + + ImageStatistics ret{}; + AnalyzeBlock(processed_image, 0, image_ptr, npixels, image_mode, ret); + return ret; +} + +void ImagePreprocessorCPU::AnalyzeBlock(ImagePreprocessorBuffer &processed_image, size_t first, const uint8_t *input, + size_t n, CompressedImageMode image_mode, ImageStatistics &stats) { switch (image_mode) { case CompressedImageMode::Int8: - return Analyze(processed_image, image_ptr, INT8_MIN, INT8_MAX); + return AnalyzeBlock(processed_image, first, input, n, INT8_MIN, INT8_MAX, stats); case CompressedImageMode::Int16: - return Analyze(processed_image, image_ptr, INT16_MIN, INT16_MAX); + return AnalyzeBlock(processed_image, first, input, n, INT16_MIN, INT16_MAX, stats); case CompressedImageMode::Int32: - return Analyze(processed_image, image_ptr, INT32_MIN, INT32_MAX); + return AnalyzeBlock(processed_image, first, input, n, INT32_MIN, INT32_MAX, stats); case CompressedImageMode::Uint8: - return Analyze(processed_image, image_ptr, UINT8_MAX, UINT8_MAX); + return AnalyzeBlock(processed_image, first, input, n, UINT8_MAX, UINT8_MAX, stats); case CompressedImageMode::Uint16: - return Analyze(processed_image, image_ptr, UINT16_MAX, UINT16_MAX); + return AnalyzeBlock(processed_image, first, input, n, UINT16_MAX, UINT16_MAX, stats); case CompressedImageMode::Uint32: - return Analyze(processed_image, image_ptr, UINT32_MAX, UINT32_MAX); + return AnalyzeBlock(processed_image, first, input, n, UINT32_MAX, UINT32_MAX, stats); default: throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "RGB/float mode not supported"); } } template -ImageStatistics ImagePreprocessorCPU::Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *input, T err_pixel_val, T sat_pixel_val) { - - if (processed_image.size() != npixels) - throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Processed image size mismatch"); - +void ImagePreprocessorCPU::AnalyzeBlock(ImagePreprocessorBuffer &processed_image, size_t first, const uint8_t *input, + size_t n, T err_pixel_val, T sat_pixel_val, ImageStatistics &ret) { auto image = reinterpret_cast(input); - - ImageStatistics ret{}; + int32_t *out = processed_image.data() + first; if (sat_pixel_val > saturation_limit) sat_pixel_val = static_cast(saturation_limit); - for (int i = 0; i < npixels; i++) { - if (mask_1bit[i] != 0) { - processed_image[i] = INT32_MIN; + for (size_t i = 0; i < n; i++) { + if (mask_1bit[first + i] != 0) { + out[i] = INT32_MIN; ++ret.masked_pixel_count; } else if (image[i] == err_pixel_val) { // Error/invalid marker = the pixel type's extreme value (0xFFFFFFFF for EIGER uint32). // Tested before saturation, since for unsigned types the marker also exceeds sat_pixel_val // (which is clipped above to the HDF5 saturation_value). - processed_image[i] = INT32_MIN; + out[i] = INT32_MIN; ++ret.error_pixel_count; } else if (image[i] >= sat_pixel_val) { - processed_image[i] = INT32_MAX; + out[i] = INT32_MAX; ++ret.saturated_pixel_count; } else { - processed_image[i] = static_cast(image[i]); + out[i] = static_cast(image[i]); if (image[i] > ret.max_value) ret.max_value = image[i]; if (image[i] < ret.min_value) ret.min_value = image[i]; } } - return ret; } diff --git a/image_analysis/image_preprocessing/ImagePreprocessorCPU.h b/image_analysis/image_preprocessing/ImagePreprocessorCPU.h index aa326d77b..f14706d64 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorCPU.h +++ b/image_analysis/image_preprocessing/ImagePreprocessorCPU.h @@ -9,8 +9,13 @@ class ImagePreprocessorCPU : public ImagePreprocessor { std::vector mask_1bit; - template ImageStatistics Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *input, T err_value, T sat_value); + template void AnalyzeBlock(ImagePreprocessorBuffer &processed_image, size_t first, const uint8_t *input, + size_t n, T err_value, T sat_value, ImageStatistics &stats); public: ImagePreprocessorCPU(const DiffractionExperiment &experiment, const PixelMask &mask); ImageStatistics Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *decompressed_image, CompressedImageMode image_mode) override; + // Analyze for the n pixels from pixel `first` on, whose values `input` holds, adding to stats. The + // blocks of an image can come in any order and together give Analyze()'s result. + void AnalyzeBlock(ImagePreprocessorBuffer &processed_image, size_t first, const uint8_t *input, size_t n, + CompressedImageMode image_mode, ImageStatistics &stats); }; diff --git a/image_analysis/spot_finding/AdaptiveSpotFinderCPU.cpp b/image_analysis/spot_finding/AdaptiveSpotFinderCPU.cpp index 93b7b9597..a23b12d99 100644 --- a/image_analysis/spot_finding/AdaptiveSpotFinderCPU.cpp +++ b/image_analysis/spot_finding/AdaptiveSpotFinderCPU.cpp @@ -2,7 +2,6 @@ // SPDX-License-Identifier: GPL-3.0-only #include -#include #include #include "AdaptiveSpotFinderCPU.h" @@ -34,71 +33,87 @@ void AdaptiveSpotFinderCPU::GetProfile(AzimuthalIntegrationProfile &profile) con // Per-ring background statistics with iterated peak exclusion following peakfinder8: // Barty et al. (2014) J. Appl. Cryst. 47, 1118-1131 -// Accumulate per-ring mean/variance from the raw (photon) image. clip_k <= 0 -> use every valid -// pixel (first pass); clip_k > 0 -> keep only pixels within clip_k sigma of the current ring mean, -// which removes the Bragg peaks from the background estimate. -void AdaptiveSpotFinderCPU::AccumulateRings(const ImagePreprocessorBuffer &image, float clip_k) { +// Per-ring mean/variance from the raw (photon) image: a plain pass over every valid pixel, then +// sigma-clip passes keeping only pixels within clip_k sigma of the current ring mean, which removes +// the Bragg peaks from the background estimate. +void AdaptiveSpotFinderCPU::ResetRings() { + std::fill(ring_sum.begin(), ring_sum.end(), 0); + std::fill(ring_sum2.begin(), ring_sum2.end(), 0); + std::fill(ring_cnt.begin(), ring_cnt.end(), 0); + std::fill(ring_hist.begin(), ring_hist.end(), 0); + ring_overflow.clear(); + if (fuse_azint) { + std::fill(azint_sum.begin(), azint_sum.end(), 0.0f); + std::fill(azint_sum2.begin(), azint_sum2.end(), 0.0f); + std::fill(azint_count.begin(), azint_count.end(), 0); + } +} + +void AdaptiveSpotFinderCPU::BeginRings() { + ResetRings(); + rings_from_blocks = true; +} + +// The plain pass over pixels [first, first + n). +void AdaptiveSpotFinderCPU::AccumulateRingsBlock(const ImagePreprocessorBuffer &image, size_t first, size_t n) { const auto &pixel_to_bin = mapping.GetPixelToBin(); const size_t nbins = ring_sum.size(); - const size_t npix = static_cast(width) * height; + const float *corrections = mapping.Corrections().data(); + + for (size_t pxl = first; pxl < first + n; ++pxl) { + const int32_t v = image[pxl]; + if (v == INT32_MIN || v == INT32_MAX) continue; // bad / saturated + const uint16_t b = pixel_to_bin[pxl]; + if (b >= nbins) continue; // masked / out of range (UINT16_MAX) + if (fuse_azint) { + const float val = static_cast(v) * corrections[pxl]; + const float val_sq = val * val; + azint_sum[b] += val; + azint_sum2[b] += val_sq; + ++azint_count[b]; + } + ring_sum[b] += v; + ring_sum2[b] += static_cast(static_cast(v) * v); + ring_cnt[b] += 1; + if (v >= 0 && v < HIST_VALUES) + ring_hist[b * HIST_VALUES + v] += 1; + else + ring_overflow.emplace_back(b, v); + } +} + +// A sigma-clip pass over the plain pass's values: each distinct value of a ring meets the same test the +// pixels holding it would, and its pixels are added as a count. +void AdaptiveSpotFinderCPU::ClipRings(float clip_k) { + const size_t nbins = ring_sum.size(); std::fill(ring_sum.begin(), ring_sum.end(), 0); std::fill(ring_sum2.begin(), ring_sum2.end(), 0); std::fill(ring_cnt.begin(), ring_cnt.end(), 0); - if (clip_k <= 0.0f) { - std::fill(ring_hist.begin(), ring_hist.end(), 0); - ring_overflow.clear(); - const float *corrections = mapping.Corrections().data(); - if (fuse_azint) { - std::fill(azint_sum.begin(), azint_sum.end(), 0.0f); - std::fill(azint_sum2.begin(), azint_sum2.end(), 0.0f); - std::fill(azint_count.begin(), azint_count.end(), 0); - } - for (size_t pxl = 0; pxl < npix; ++pxl) { - const int32_t v = image[pxl]; - if (v == INT32_MIN || v == INT32_MAX) continue; // bad / saturated - const uint16_t b = pixel_to_bin[pxl]; - if (b >= nbins) continue; // masked / out of range (UINT16_MAX) - if (fuse_azint) { - const float val = static_cast(v) * corrections[pxl]; - const float val_sq = val * val; - azint_sum[b] += val; - azint_sum2[b] += val_sq; - ++azint_count[b]; - } - ring_sum[b] += v; - ring_sum2[b] += static_cast(static_cast(v) * v); - ring_cnt[b] += 1; - if (v >= 0 && v < HIST_VALUES) - ring_hist[b * HIST_VALUES + v] += 1; - else - ring_overflow.emplace_back(b, v); - } - } else { - // A sigma-clip pass over the plain pass's values: each distinct value of a ring meets the same - // test the pixels holding it would, and its pixels are added as a count. - const auto keep = [&](uint16_t b, int32_t v) { - const float lo = ring_mean[b] - clip_k * ring_sigma[b]; - const float hi = ring_mean[b] + clip_k * ring_sigma[b]; - return !(v < lo || v > hi); // exclude peaks / outliers - }; - for (size_t b = 0; b < nbins; ++b) - for (int32_t v = 0; v < HIST_VALUES; ++v) { - const uint32_t n = ring_hist[b * HIST_VALUES + v]; - if (n == 0 || !keep(static_cast(b), v)) continue; - ring_sum[b] += static_cast(n) * v; - ring_sum2[b] += static_cast(n) * static_cast(static_cast(v) * v); - ring_cnt[b] += n; - } - for (const auto &[b, v] : ring_overflow) { - if (!keep(b, v)) continue; - ring_sum[b] += v; - ring_sum2[b] += static_cast(static_cast(v) * v); - ring_cnt[b] += 1; + const auto keep = [&](uint16_t b, int32_t v) { + const float lo = ring_mean[b] - clip_k * ring_sigma[b]; + const float hi = ring_mean[b] + clip_k * ring_sigma[b]; + return !(v < lo || v > hi); // exclude peaks / outliers + }; + for (size_t b = 0; b < nbins; ++b) + for (int32_t v = 0; v < HIST_VALUES; ++v) { + const uint32_t n = ring_hist[b * HIST_VALUES + v]; + if (n == 0 || !keep(static_cast(b), v)) continue; + ring_sum[b] += static_cast(n) * v; + ring_sum2[b] += static_cast(n) * static_cast(static_cast(v) * v); + ring_cnt[b] += n; } + for (const auto &[b, v] : ring_overflow) { + if (!keep(b, v)) continue; + ring_sum[b] += v; + ring_sum2[b] += static_cast(static_cast(v) * v); + ring_cnt[b] += 1; } +} +void AdaptiveSpotFinderCPU::UpdateRingStatistics() { + const size_t nbins = ring_sum.size(); for (size_t b = 0; b < nbins; ++b) { if (ring_cnt[b] > 0) { const double m = static_cast(ring_sum[b]) / ring_cnt[b]; @@ -114,9 +129,16 @@ void AdaptiveSpotFinderCPU::Detect(const ImagePreprocessorBuffer &image, const size_t nbins = ring_sum.size(); // --- Stage A: robust per-ring background (one plain pass + two sigma-clip passes) --- - AccumulateRings(image, 0.0f); - AccumulateRings(image, 3.0f); - AccumulateRings(image, 3.0f); + if (!rings_from_blocks) { + ResetRings(); + AccumulateRingsBlock(image, 0, static_cast(width) * height); + } + rings_from_blocks = false; + UpdateRingStatistics(); + ClipRings(3.0f); + UpdateRingStatistics(); + ClipRings(3.0f); + UpdateRingStatistics(); // --- Stage B: per-ring threshold from the single portable knob E (false pixels / frame) --- int64_t n_total = 0; @@ -153,11 +175,13 @@ void AdaptiveSpotFinderCPU::Detect(const ImagePreprocessorBuffer &image, : adaptive_threshold::RingThreshold(ring_mean[b], ring_sigma[b], p, z); // --- Stage C: the ring threshold, intersected with the classic local-box SNR test --- - FlagRings(image); + std::fill(ring_bits.begin(), ring_bits.end(), 0); if (settings.signal_to_noise_threshold <= 0.0f) { // No local test asked for: the ring threshold alone decides, as the fixed photon floor // alone would in the classic finder. + for (int32_t row = 0; row < height; row++) + FlagRow(image, row); output_buffer = ring_bits; return; } @@ -165,17 +189,17 @@ void AdaptiveSpotFinderCPU::Detect(const ImagePreprocessorBuffer &image, // The ring threshold IS the photon floor here, so the local pass must not apply another one. SpotFindingSettings local = settings; local.photon_count_threshold = 0; - // Only the ring pixels survive the intersection, so the local test is asked of those alone. - DetectAt(image, local, ring_bits); + // Only the ring pixels survive the intersection, so the local test is asked of those alone. They + // are flagged row by row as the local pass reaches each row, which saves reading the image for it. + DetectAt(image, local, ring_bits, [&](int32_t row) { FlagRow(image, row); }); } -void AdaptiveSpotFinderCPU::FlagRings(const ImagePreprocessorBuffer &image) { +void AdaptiveSpotFinderCPU::FlagRow(const ImagePreprocessorBuffer &image, int32_t row) { const auto &pixel_to_bin = mapping.GetPixelToBin(); const size_t nbins = ring_thr.size(); - const size_t npix = static_cast(width) * height; + const size_t first = static_cast(row) * width; - std::bitset<32> out = 0; - for (size_t pxl = 0; pxl < npix; ++pxl) { + for (size_t pxl = first; pxl < first + width; ++pxl) { const int32_t v = image[pxl]; const uint16_t b = pixel_to_bin[pxl]; bool strong = false; @@ -184,14 +208,7 @@ void AdaptiveSpotFinderCPU::FlagRings(const ImagePreprocessorBuffer &image) { else if (v != INT32_MIN && b < nbins && v >= ring_thr[b]) strong = true; - const int32_t bit = pxl % 32; if (strong) - out.set(bit); - if (bit == 31) { - ring_bits[pxl / 32] = out.to_ulong(); - out.reset(); - } + ring_bits[pxl / 32] |= 1U << (pxl % 32); } - if (npix % 32 != 0) - ring_bits[OutputSize() - 1] = out.to_ulong(); } diff --git a/image_analysis/spot_finding/AdaptiveSpotFinderCPU.h b/image_analysis/spot_finding/AdaptiveSpotFinderCPU.h index f741988b9..d0a82aa5f 100644 --- a/image_analysis/spot_finding/AdaptiveSpotFinderCPU.h +++ b/image_analysis/spot_finding/AdaptiveSpotFinderCPU.h @@ -71,7 +71,7 @@ class AdaptiveSpotFinderCPU : public ImageSpotFinderCPU { std::vector ring_hist; std::vector> ring_overflow; // (ring, value) - // The azimuthal-integration profile, taken in the plain AccumulateRings pass when asked for: the + // The azimuthal-integration profile, taken in the plain ring pass when asked for: the // same pixels in the same order and the same arithmetic as AzIntEngineCPU, so the same sums, and // one pass over the image less (the CPU twin of the fused GPU engine). bool fuse_azint = false; @@ -79,15 +79,31 @@ class AdaptiveSpotFinderCPU : public ImageSpotFinderCPU { std::vector azint_sum2; std::vector azint_count; - void AccumulateRings(const ImagePreprocessorBuffer &image, float clip_k); - // Fill ring_bits from the thresholds of the current frame. - void FlagRings(const ImagePreprocessorBuffer &image); + // Set by BeginRings(): the plain ring pass of the next Detect() is being accumulated block by block. + bool rings_from_blocks = false; + + // Zero the sums of the plain ring pass (and of the fused profile). + void ResetRings(); + // One sigma-clip pass over the plain pass's values. + void ClipRings(float clip_k); + // ring_mean / ring_sigma from the current sums. + void UpdateRingStatistics(); + // Set the ring_bits of the pixels of one row from the thresholds of the current frame. ring_bits + // is zeroed before the first row. + void FlagRow(const ImagePreprocessorBuffer &image, int32_t row); public: explicit AdaptiveSpotFinderCPU(const AzimuthalIntegrationMapping &mapping); void Detect(const ImagePreprocessorBuffer &image, const SpotFindingSettings &settings) override; [[nodiscard]] const std::vector &GetRingBackground() const override { return ring_bkg; } + // The plain ring pass taken while the image is being preprocessed, so the pixels are read while + // still in cache: BeginRings(), then AccumulateRingsBlock() over every block of the image in pixel + // order - the order keeps the profile's float sums the same - and the next Detect() of that image + // starts from these sums instead of passing over the image for them. + void BeginRings(); + void AccumulateRingsBlock(const ImagePreprocessorBuffer &image, size_t first, size_t n); + void FuseAzimuthalIntegration(bool enable) { fuse_azint = enable; } // The profile of the last Detect(), when fused. void GetProfile(AzimuthalIntegrationProfile &profile) const; diff --git a/image_analysis/spot_finding/ImageSpotFinderCPU.cpp b/image_analysis/spot_finding/ImageSpotFinderCPU.cpp index e2a41bc15..cc6946cc4 100644 --- a/image_analysis/spot_finding/ImageSpotFinderCPU.cpp +++ b/image_analysis/spot_finding/ImageSpotFinderCPU.cpp @@ -48,24 +48,12 @@ bool StrongInWindow(int64_t pxl_val, int64_t sum, int64_t sum2, int64_t valid, } // namespace void ImageSpotFinderCPU::DetectAt(const ImagePreprocessorBuffer &image, const SpotFindingSettings &settings, - const std::vector &candidates) { + const std::vector &candidates, + const std::function &fill_row) { candidate_windows.clear(); - // The first pass's bits are read only inside the window of a candidate (and at the candidate), so - // it tests only the pixels of the row/column blocks such a window reaches; the bits it leaves - // unset there are never read. - const int32_t nblocks = (width + 31) / 32; - first_pass_needed.assign(static_cast(height) * nblocks, 0); - for (size_t w = 0; w < candidates.size(); w++) - for (uint32_t bits = candidates[w]; bits; bits &= bits - 1) { - const int64_t pxl = static_cast(w) * 32 + std::countr_zero(bits); - if (pxl >= static_cast(width) * height) - break; - const int32_t line = static_cast(pxl / width), col = static_cast(pxl % width); - for (int32_t y = std::max(line - NBX, 0); y <= std::min(line + NBX, height - 1); y++) - for (int32_t b = std::max(col - NBX, 0) / 32; b <= std::min(col + NBX, width - 1) / 32; b++) - first_pass_needed[static_cast(y) * nblocks + b] = 1; - } - DetectPass(image, settings, nullptr, first_pass_buffer, candidates.data()); + // Filled in by DetectPass as the candidates of each row become known. + first_pass_needed.assign(static_cast(height) * ((width + 31) / 32), 0); + DetectPass(image, settings, nullptr, first_pass_buffer, candidates.data(), fill_row); std::fill(output_buffer.begin(), output_buffer.end(), 0); const float strong2 = settings.signal_to_noise_threshold * settings.signal_to_noise_threshold; @@ -107,10 +95,39 @@ void ImageSpotFinderCPU::DetectPass(const ImagePreprocessorBuffer &image, const SpotFindingSettings &settings, const uint32_t *prev_strong, std::vector &out_buffer, - const uint32_t *candidates) { + const uint32_t *candidates, + const std::function &fill_row) { for (int i = 0; i < OutputSize(); i++) out_buffer[i] = 0; + // The first pass's bits are read only inside the window of a candidate (and at the candidate), so + // with candidates it tests only the pixels of the row/column blocks such a window reaches; the bits + // it leaves unset there are never read. Those blocks are marked from each row's candidates as the + // row enters the vertical sums below, which is before any row within NBX of it is tested - so the + // candidates themselves can be filled in a row at a time (fill_row), while the row is in cache. + const int32_t nblocks = (width + 31) / 32; + const auto new_row = [&](int32_t r) { + if (!candidates) + return; + if (fill_row) + fill_row(r); + const int32_t first = r * width; + const int32_t last = first + width - 1; + for (int32_t w = first / 32; w <= last / 32; w++) { + uint32_t bits = candidates[w]; + if (w == first / 32) + bits &= UINT32_MAX << (first % 32); + if (w == last / 32 && last % 32 != 31) + bits &= (1U << (last % 32 + 1)) - 1; + for (; bits; bits &= bits - 1) { + const int32_t col = w * 32 + std::countr_zero(bits) - first; + for (int32_t y = std::max(r - NBX, 0); y <= std::min(r + NBX, height - 1); y++) + for (int32_t b = std::max(col - NBX, 0) / 32; b <= std::min(col + NBX, width - 1) / 32; b++) + first_pass_needed[static_cast(y) * nblocks + b] = 1; + } + } + }; + // A pixel found strong by the previous pass reads as INT32_MAX, which the accumulation below // already skips and the acceptance test below already takes as strong - the same substitution // the GPU kernel makes when it reads prev_out. @@ -147,6 +164,7 @@ void ImageSpotFinderCPU::DetectPass(const ImagePreprocessorBuffer &image, std::vector valid_vert(width, 0); for (int line = 0; line < NBX; line++) { + new_row(line); for (int col = 0; col < width; col++) { auto pxl = line * width + col; @@ -160,6 +178,8 @@ void ImageSpotFinderCPU::DetectPass(const ImagePreprocessorBuffer &image, } for (int line = 0; line < height; line++) { + if (line < height - NBX) + new_row(line + NBX); for (int col = 0; col < width; col++) { if (line < height - NBX) { auto pxl = (line + NBX) * width + col; @@ -188,7 +208,6 @@ void ImageSpotFinderCPU::DetectPass(const ImagePreprocessorBuffer &image, // candidate lies in one). Each run starts its window from the vertical sums over the // columns it covers there - integers, so the same sums the running window holds at // that column - and slides it as below; the bits outside the runs stay 0, as below. - const int32_t nblocks = (width + 31) / 32; const uint8_t *needed = first_pass_needed.data() + static_cast(line) * nblocks; for (int32_t b = 0; b < nblocks;) { if (!needed[b]) { b++; continue; } diff --git a/image_analysis/spot_finding/ImageSpotFinderCPU.h b/image_analysis/spot_finding/ImageSpotFinderCPU.h index 39e87a270..5b754b2da 100644 --- a/image_analysis/spot_finding/ImageSpotFinderCPU.h +++ b/image_analysis/spot_finding/ImageSpotFinderCPU.h @@ -5,6 +5,7 @@ #include +#include #include #include "ImageSpotFinder.h" @@ -26,10 +27,11 @@ class ImageSpotFinderCPU : public ImageSpotFinder { // fall into, and reported strong. // // With candidates set (SNR test only), the window sums of every candidate pixel are also recorded - // in candidate_windows, in pixel order. + // in candidate_windows, in pixel order, and fill_row (see DetectAt) is called. void DetectPass(const ImagePreprocessorBuffer &image, const SpotFindingSettings &settings, const uint32_t *prev_strong, std::vector &out, - const uint32_t *candidates = nullptr); + const uint32_t *candidates = nullptr, + const std::function &fill_row = {}); // Local-box sums (centre pixel included) of one candidate pixel, as the first pass saw them. struct CandidateWindow { @@ -47,8 +49,12 @@ protected: // the image: a candidate's second-pass window is its first-pass window with the pixels the first // pass found strong taken out, and those are few, so they are subtracted from the sums the first // pass recorded. Integer sums and the same test, so the bits are exactly Detect's. + // + // fill_row(row), when given, sets the candidates of that row: it is called once for every row, in + // order, before any candidate of the row is read, when the pass is about to read the row anyway. void DetectAt(const ImagePreprocessorBuffer &image, const SpotFindingSettings &settings, - const std::vector &candidates); + const std::vector &candidates, + const std::function &fill_row = {}); public: ImageSpotFinderCPU(int32_t width, int32_t height); diff --git a/tests/ZSTDCompressorTest.cpp b/tests/ZSTDCompressorTest.cpp index 5a497601e..fc6c76ee0 100644 --- a/tests/ZSTDCompressorTest.cpp +++ b/tests/ZSTDCompressorTest.cpp @@ -291,6 +291,34 @@ TEST_CASE("JFJochCompressor_JFJochDecompressor_LZ4","[ZSTD]") { REQUIRE(memcmp(image.data(), output.data(), x.GetPixelsNum() * sizeof(int32_t)) == 0); } +TEST_CASE("JFJochDecompressBlocks_matches_JFJochDecompress","[ZSTD]") { + // An element count that leaves a partial block and a few unshuffled trailing elements. + const size_t nelements = 3 * 8192 + 1000 + 5; + std::vector image(nelements); + std::mt19937 g(1); + std::poisson_distribution dist(3.0); + for (auto &i: image) + i = dist(g); + + for (auto algorithm: {CompressionAlgorithm::BSHUF_LZ4, CompressionAlgorithm::BSHUF_ZSTD}) { + JFJochBitShuffleCompressor compressor(algorithm); + std::vector tmp(nelements * sizeof(uint32_t) * 2 + 12); + tmp.resize(compressor.Compress(tmp.data(), tmp.size(), image)); + + std::vector output(nelements, UINT32_MAX); + size_t next = 0; + JFJochDecompressBlocks(algorithm, reinterpret_cast(tmp.data()), tmp.size(), + nelements, sizeof(uint32_t), + [&](size_t first, const uint8_t *data, size_t n) { + REQUIRE(first == next); + memcpy(output.data() + first, data, n * sizeof(uint32_t)); + next = first + n; + }); + REQUIRE(next == nelements); + REQUIRE(output == image); + } +} + TEST_CASE("JFJochDecompressor_None","[ZSTD]") { DiffractionExperiment x(DetJF4M()); x.Compression(CompressionAlgorithm::NO_COMPRESSION).BitDepthImage(32).PixelSigned(true);