From 585cc777d522851ca645faf5d4e83d70e97a0d31 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 2 Nov 2026 18:09:39 +0100 Subject: [PATCH] GPU analysis of converted JUNGFRAU frames with a per-run pre-pass With rugnux enabled, a serial (non-rotation) collection in the FPGA workflow is analysed as the DECTRIS frames are: the frame transformation threads compress each converted frame and hand it to StillsGPUAnalysis, whose threads run MXAnalysisWithoutFPGA (GPU spot finding, indexing, integration) and pass the result back to the receiver to be written. Before any frame is analysed, the first 100 are run through rugnux as `rugnux --no-merge` on them (RunStillsPrePass, in memory): its beam centre (the direct beam, from the scattered background and, where frames index, the stills geometry trial or the bundle refinement with a given cell) and its mask (plus the beam-stop shadow and the defective pixels) are what every frame of the collection is then analysed with. Frames arriving meanwhile are held. The start message waits for the pre-pass, so it states that centre and mask; the END message carries the centre as refined_beam_center_x/y. The start message dataset now carries the detector modules, unit cell and space group, which the pre-pass needs. tools/jfjoch_stills_gpu_replay replays a stored serial collection through the same analysis (pre-pass, indexed fraction, sustained rate). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01UqXYdJKPUFLoeAchjzv3x4 --- reader/JFJochStartMessageReader.cpp | 6 + receiver/CMakeLists.txt | 2 + receiver/JFJochReceiver.cpp | 37 +++--- receiver/JFJochReceiver.h | 4 + receiver/JFJochReceiverFPGA.cpp | 113 ++++++++++++++++++- receiver/JFJochReceiverFPGA.h | 11 +- receiver/JFJochReceiverService.cpp | 5 +- receiver/StillsGPUAnalysis.cpp | 137 +++++++++++++++++++++++ receiver/StillsGPUAnalysis.h | 85 ++++++++++++++ receiver/StillsPrePass.cpp | 64 +++++++++++ receiver/StillsPrePass.h | 32 ++++++ rugnux/Rugnux.h | 3 + rugnux/RugnuxLive.cpp | 39 +++++++ rugnux/RugnuxLive.h | 5 + tests/JFJochReceiverProcessingTest.cpp | 81 ++++++++++++++ tools/CMakeLists.txt | 4 + tools/jfjoch_stills_gpu_replay.cpp | 149 +++++++++++++++++++++++++ 17 files changed, 755 insertions(+), 22 deletions(-) create mode 100644 receiver/StillsGPUAnalysis.cpp create mode 100644 receiver/StillsGPUAnalysis.h create mode 100644 receiver/StillsPrePass.cpp create mode 100644 receiver/StillsPrePass.h create mode 100644 tools/jfjoch_stills_gpu_replay.cpp diff --git a/reader/JFJochStartMessageReader.cpp b/reader/JFJochStartMessageReader.cpp index f0c5fc14a..6ff183b7d 100644 --- a/reader/JFJochStartMessageReader.cpp +++ b/reader/JFJochStartMessageReader.cpp @@ -55,6 +55,12 @@ std::shared_ptr DatasetFromStartMessage(const StartMessage detector.MinCountTime(std::chrono::microseconds(0)); detector.ReadOutTime(std::chrono::nanoseconds(0)); x.Detector(detector); + // The modules one by one, where the start message states them (a JUNGFRAU's, as imaged). + if (!start.detector_modules.empty()) + x.Modules(DetectorModules(start.detector_modules)); + x.SetUnitCell(start.unit_cell); + if (start.space_group_number) + x.SpaceGroupNumber(static_cast(*start.space_group_number)); x.FrameTime(std::chrono::duration_cast(std::chrono::duration(start.frame_time)), std::chrono::duration_cast(std::chrono::duration(start.count_time))); diff --git a/receiver/CMakeLists.txt b/receiver/CMakeLists.txt index 4eaf0898a..9096e604c 100644 --- a/receiver/CMakeLists.txt +++ b/receiver/CMakeLists.txt @@ -12,6 +12,8 @@ ADD_LIBRARY(JFJochReceiver STATIC JFJochReceiverOutput.h JFJochReceiverLite.h JFJochReceiverLite.cpp JFJochReceiverRugnux.h JFJochReceiverRugnux.cpp + StillsPrePass.h StillsPrePass.cpp + StillsGPUAnalysis.h StillsGPUAnalysis.cpp RugnuxLiveStatus.h RugnuxLiveStatus.cpp RugnuxStructure.h RugnuxStructure.cpp JFJochReceiver.cpp diff --git a/receiver/JFJochReceiver.cpp b/receiver/JFJochReceiver.cpp index 8ecc77425..9ffa7badc 100644 --- a/receiver/JFJochReceiver.cpp +++ b/receiver/JFJochReceiver.cpp @@ -107,28 +107,33 @@ JFJochReceiverStatus JFJochReceiver::GetStatus() const { } -void JFJochReceiver::SendStartMessage() { +StartMessage JFJochReceiver::MakeStartMessage(const DiffractionExperiment &x, const PixelMask &mask, + const AzimuthalIntegrationMapping &mapping) { StartMessage message{}; - experiment.FillMessage(message); + x.FillMessage(message); message.arm_date = time_UTC(std::chrono::system_clock::now()); - message.az_int_q_bin_count = az_int_mapping->GetQBinCount(); - message.az_int_bin_to_q = az_int_mapping->GetBinToQ(); - message.az_int_bin_to_two_theta = az_int_mapping->GetBinToTwoTheta(); - message.az_int_phi_bin_count = az_int_mapping->GetAzimuthalBinCount(); - if (az_int_mapping->GetAzimuthalBinCount() > 1) { - message.az_int_bin_to_phi = az_int_mapping->GetBinToPhi(); - message.az_int_map = az_int_mapping->GetPixelToBin(); + message.az_int_q_bin_count = mapping.GetQBinCount(); + message.az_int_bin_to_q = mapping.GetBinToQ(); + message.az_int_bin_to_two_theta = mapping.GetBinToTwoTheta(); + message.az_int_phi_bin_count = mapping.GetAzimuthalBinCount(); + if (mapping.GetAzimuthalBinCount() > 1) { + message.az_int_bin_to_phi = mapping.GetBinToPhi(); + message.az_int_map = mapping.GetPixelToBin(); } message.writer_notification_zmq_addr = image_pusher.GetWriterNotificationSocketAddress(); - message.rois = experiment.ROI().ExportMetadata(); - if (!experiment.ROI().empty()) - message.roi_map = experiment.ExportROIMap(); - message.max_spot_count = experiment.GetMaxSpotCount(); + message.rois = x.ROI().ExportMetadata(); + if (!x.ROI().empty()) + message.roi_map = x.ExportROIMap(); + message.max_spot_count = x.GetMaxSpotCount(); + message.pixel_mask["default"] = mask.GetMask(x); + return message; +} +void JFJochReceiver::SendStartMessage() { + SendStartMessage(MakeStartMessage(experiment, pixel_mask, *az_int_mapping)); +} - std::vector nexus_mask; - message.pixel_mask["default"] = pixel_mask.GetMask(experiment); - +void JFJochReceiver::SendStartMessage(StartMessage message) { SaveStartMessageToImageBuffer(message); if (push_images_to_writer) diff --git a/receiver/JFJochReceiver.h b/receiver/JFJochReceiver.h index 93fb213ae..f0f76db20 100644 --- a/receiver/JFJochReceiver.h +++ b/receiver/JFJochReceiver.h @@ -93,7 +93,11 @@ protected: void UpdateMaxDelay(uint64_t delay); JFJochReceiverStatus GetStatus() const; + // The start message of this collection, stating that experiment, mask and azimuthal mapping. + StartMessage MakeStartMessage(const DiffractionExperiment &x, const PixelMask &mask, + const AzimuthalIntegrationMapping &mapping); void SendStartMessage(); + void SendStartMessage(StartMessage message); void SendEndMessage(); void SaveStartMessageToImageBuffer(const StartMessage &msg); // The last word on the end message before it is sent: a receiver whose analysis is not the diff --git a/receiver/JFJochReceiverFPGA.cpp b/receiver/JFJochReceiverFPGA.cpp index d22f64de4..4f8969587 100644 --- a/receiver/JFJochReceiverFPGA.cpp +++ b/receiver/JFJochReceiverFPGA.cpp @@ -6,6 +6,9 @@ #include #include "ImageMetadata.h" +#include "JFJochReceiverLite.h" +#include "../common/CUDAWrapper.h" +#include "../image_analysis/MXAnalysisWithoutFPGA.h" JFJochReceiverFPGA::JFJochReceiverFPGA(const DiffractionExperiment &in_experiment, const PixelMask &in_pixel_mask, @@ -19,7 +22,8 @@ JFJochReceiverFPGA::JFJochReceiverFPGA(const DiffractionExperiment &in_experimen ImageBuffer &in_send_buf_ctrl, ZMQPreviewSocket *in_zmq_preview_socket, ZMQMetadataSocket *in_zmq_metadata_socket, - IndexerThreadPool *indexing_thread_pool) + IndexerThreadPool *indexing_thread_pool, + bool in_gpu_analysis) : JFJochReceiver(in_experiment, in_send_buf_ctrl, in_image_sender, @@ -39,7 +43,9 @@ JFJochReceiverFPGA::JFJochReceiverFPGA(const DiffractionExperiment &in_experimen pedestal_nthreads((experiment.GetStorageCellNumber() > 2) ? 1 : 4), summation_nthreads(2), frame_transformation_ready((experiment.GetImageNum() > 0) ? frame_transformation_nthreads : 0), - data_acquisition_ready(ndatastreams) { + data_acquisition_ready(ndatastreams), + gpu_analysis(in_gpu_analysis && experiment.GetImageNum() > 0 + && experiment.GetDetectorMode() == DetectorMode::Standard) { // The FPGA integration core has a fixed bin-count limit. Reject configurations it cannot // handle before starting acquisition, unless azimuthal integration is forced onto the CPU. const auto &azint_settings = experiment.GetAzimuthalIntegrationSettings(); @@ -118,8 +124,21 @@ JFJochReceiverFPGA::JFJochReceiverFPGA(const DiffractionExperiment &in_experimen } logger.Debug("Pedestal threads ready ({} threads/module*SC)", pedestal_nthreads); } else if (experiment.GetImageNum() > 0) { - SendStartMessage(); - SendCalibration(); + if (gpu_analysis) { + // The start message and the calibration go out once the pre-pass is done. + gpu = std::make_unique( + experiment, pixel_mask, MakeStartMessage(experiment, pixel_mask, *az_int_mapping), + indexing_thread_pool, std::min(16, frame_transformation_nthreads), logger, + [this](const StillsGPUAnalysis::Setup &s) { + SendStartMessage(MakeStartMessage(s.experiment, s.mask, *s.mapping)); + SendCalibration(); + }, + [this](DataMessage &m, const AzimuthalIntegrationProfile &p) { GPUSink(m, p); }, + [this] { return GetSpotFindingSettings(); }); + } else { + SendStartMessage(); + SendCalibration(); + } for (int i = 0; i < experiment.GetImageNum(); i++) images_to_go.Put(i); @@ -420,6 +439,23 @@ void JFJochReceiverFPGA::FrameTransformationThread(uint32_t threadid) { experiment.GetImageMode(), CompressionAlgorithm::NO_COMPRESSION); + if (gpu_analysis) { + // Compressed here, as it is written, and analysed by GPUAnalysisThread. + auto held = std::make_shared(); + held->image.resize(experiment.GetImageBufferLocationSize()); + held->image.resize(transformation.CompressImage(held->image.data(), held->image.size())); + held->message = std::move(message); + held->message.image = CompressedImage(held->image.data(), held->image.size(), + experiment.GetXPixelsNum(), experiment.GetYPixelsNum(), + experiment.GetImageMode(), + experiment.GetCompressionAlgorithm()); + compressed_size += held->image.size(); + uncompressed_size += experiment.GetModulesNum() * RAW_MODULE_SIZE * experiment.GetByteDepthImage(); + ++images_collected; + gpu->Hold(held); + continue; + } + // No-op unless the CPU backend is forced; fills az_int_profile_image from the assembled image analyzer->RunAzimuthalIntegration(transformation.GetImage(), az_int_profile_image); @@ -598,6 +634,9 @@ void JFJochReceiverFPGA::FinalizeMeasurement() { logger.Debug("All processing threads done"); } + if (gpu) + gpu->Finish(); + current_status.SetProgress(1.0); current_status.SetStatus(GetStatus()); @@ -703,3 +742,69 @@ void JFJochReceiverFPGA::LoadCalibrationToFPGA(uint16_t data_stream) { // Initialize data processing acquisition_device[data_stream].InitializeDataProcessing(experiment, *az_int_mapping); } + +void JFJochReceiverFPGA::GPUSink(DataMessage &message, const AzimuthalIntegrationProfile &profile) { + try { + const auto status = image_buffer.GetStatus(); + message.receiver_buf_available = status.available_slots; + message.receiver_buf_in_preparation = status.preparation_slots; + message.receiver_buf_in_sending = status.sending_slots; + if (message.image.GetCompressedSize() > 0) + message.compression_ratio = static_cast(experiment.GetPixelsNum() * experiment.GetByteDepthImage()) + / static_cast(message.image.GetCompressedSize()); + + if (message.image_collection_efficiency == 1.0f) { + saturated_pixels.Add(message.saturated_pixel_count); + error_pixels.Add(message.error_pixel_count); + if (message.roi.contains("beam")) { + roi_beam_npixel.Add(message.roi["beam"].pixels); + roi_beam_sum.Add(message.roi["beam"].sum); + } + } + + scan_result.Add(message); + + if (!serialmx_filter.ApplyFilter(message)) + ++images_skipped; + else if (auto loc = image_buffer.GetImageSlot(); loc == nullptr) + writer_queue_full = true; + else { + auto writer_buffer = (uint8_t *) loc->GetImage(); + CBORStream2Serializer serializer(writer_buffer, experiment.GetImageBufferLocationSize()); + if (JFJochReceiverLite::SerializeImageOrDrop(serializer, message, *loc, + experiment.GetImageBufferLocationSize(), logger)) { + loc->SetImageNumber(message.number); + loc->SetImageSize(serializer.GetBufferSize()); + loc->SetIndexed(message.indexing_result.value_or(false)); + loc->ReadyToSend(); + + if (zmq_preview_socket != nullptr) + zmq_preview_socket->SendImage(writer_buffer, serializer.GetBufferSize()); + if (zmq_metadata_socket != nullptr) + zmq_metadata_socket->AddDataMessage(message); + + if (push_images_to_writer) { + if (image_pusher.SendImage(*loc)) + ++images_sent; + } else + loc->release(); + UpdateMaxImageSent(message.number); + } + } + + plots.Add(message, profile); + current_status.SetProgress(GetProgress()); + current_status.SetStatus(GetStatus()); + } catch (const JFJochException &e) { + logger.ErrorException(e); + Cancel(e); + } +} + +void JFJochReceiverFPGA::AddToEndMessage(EndMessage &message) { + if (!gpu) + return; + message.unit_cell = gpu->GetConsensusUnitCell(); + message.refined_beam_center_x = gpu->GetBeamX_pxl(); + message.refined_beam_center_y = gpu->GetBeamY_pxl(); +} diff --git a/receiver/JFJochReceiverFPGA.h b/receiver/JFJochReceiverFPGA.h index d4bcdf3d0..4de57d708 100644 --- a/receiver/JFJochReceiverFPGA.h +++ b/receiver/JFJochReceiverFPGA.h @@ -40,6 +40,7 @@ #include "../image_analysis/MXAnalysisAfterFPGA.h" #include "../common/MovingAverage.h" #include "JFJochReceiverOutput.h" +#include "StillsGPUAnalysis.h" class JFJochReceiverFPGA : public JFJochReceiver { const JFCalibration *calibration; @@ -86,6 +87,13 @@ class JFJochReceiverFPGA : public JFJochReceiver { void Cancel(const JFJochException &e) override; void FinalizeMeasurement(); + // Serial data with rugnux on: the converted frames are analysed on the GPU (StillsGPUAnalysis), and + // the start message waits for its pre-pass. + const bool gpu_analysis; + std::unique_ptr gpu; + void GPUSink(DataMessage &message, const AzimuthalIntegrationProfile &profile); + void AddToEndMessage(EndMessage &message) override; + void SendCalibration(); void SendPedestal(const std::string &prefix, const std::vector &v, int gain, int sc); void RetrievePedestal(std::vector &output) const; @@ -102,7 +110,8 @@ public: ImageBuffer &image_buffer, ZMQPreviewSocket *zmq_preview_socket, ZMQMetadataSocket *zmq_metadata_socket, - IndexerThreadPool *indexing_thread_pool); + IndexerThreadPool *indexing_thread_pool, + bool gpu_analysis = false); ~JFJochReceiverFPGA(); JFJochReceiverFPGA(const JFJochReceiverFPGA &other) = delete; JFJochReceiverFPGA& operator=(const JFJochReceiverFPGA &other) = delete; diff --git a/receiver/JFJochReceiverService.cpp b/receiver/JFJochReceiverService.cpp index fce06c914..95bef030b 100644 --- a/receiver/JFJochReceiverService.cpp +++ b/receiver/JFJochReceiverService.cpp @@ -114,7 +114,10 @@ void JFJochReceiverService::Start(const DiffractionExperiment &experiment, image_buffer, zmq_preview_socket.get(), zmq_metadata_socket.get(), - indexer_thread_pool.get()); + indexer_thread_pool.get(), + // Serial data with rugnux on: the frames are + // analysed on the GPU, after a pre-pass. + rugnux_enable && !RugnuxRotationData(experiment)); break; case DetectorType::DECTRIS: if (puller) diff --git a/receiver/StillsGPUAnalysis.cpp b/receiver/StillsGPUAnalysis.cpp new file mode 100644 index 000000000..fba263555 --- /dev/null +++ b/receiver/StillsGPUAnalysis.cpp @@ -0,0 +1,137 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "StillsGPUAnalysis.h" + +#include + +#include "../common/CUDAWrapper.h" +#include "../image_analysis/MXAnalysisWithoutFPGA.h" + +StillsGPUAnalysis::StillsGPUAnalysis(const DiffractionExperiment &in_experiment, const PixelMask &in_mask, + const StartMessage &in_start, IndexerThreadPool *in_indexing_thread_pool, + int64_t nthreads, Logger &in_logger, Ready in_ready, Sink in_sink, + Settings in_settings) + : experiment(in_experiment), mask(in_mask), start(in_start), indexing_thread_pool(in_indexing_thread_pool), + logger(in_logger), ready(std::move(in_ready)), sink(std::move(in_sink)), settings(std::move(in_settings)) { + prepass_future = std::async(std::launch::async, &StillsGPUAnalysis::PrePassThread, this); + for (int64_t i = 0; i < nthreads; i++) + analysis_futures.emplace_back(std::async(std::launch::async, &StillsGPUAnalysis::AnalysisThread, this)); + logger.Info("GPU analysis: {} threads, the first {} images held for the pre-pass", nthreads, PREPASS_IMAGES); +} + +StillsGPUAnalysis::~StillsGPUAnalysis() { + if (!finished) + Finish(); +} + +void StillsGPUAnalysis::Hold(const std::shared_ptr &image) { + { + std::unique_lock ul(prepass_mutex); + if (!prepass_go) { + if (image->message.number < PREPASS_IMAGES) + prepass_images[image->message.number] = image; + ++images_held; + // Every one of the first images is there - or, where some never come, twice as many are. + prepass_go = (prepass_images.size() == std::min(PREPASS_IMAGES, experiment.GetImageNum())) + || (images_held >= 2 * PREPASS_IMAGES); + } + } + prepass_cv.notify_all(); + queue.PutBlocking(image); +} + +void StillsGPUAnalysis::StartPrePass() { + { + std::unique_lock ul(prepass_mutex); + prepass_go = true; + } + prepass_cv.notify_all(); +} + +void StillsGPUAnalysis::Finish() { + finished = true; + StartPrePass(); + prepass_future.get(); + for (size_t i = 0; i < analysis_futures.size(); i++) + queue.PutBlocking(nullptr); + for (auto &f: analysis_futures) + f.get(); +} + +void StillsGPUAnalysis::PrePassThread() { + std::map> sample; + { + std::unique_lock ul(prepass_mutex); + prepass_cv.wait(ul, [this] { return prepass_go; }); + sample.swap(prepass_images); + } + + const auto masked = [](const PixelMask &m) { + return std::count_if(m.GetMask().begin(), m.GetMask().end(), [](uint32_t v) { return v != 0; }); + }; + setup.experiment = experiment; + setup.mask = mask; + if (!sample.empty()) { + try { + const auto t0 = std::chrono::steady_clock::now(); + const auto r = RunStillsPrePass(start, sample, static_cast(std::thread::hardware_concurrency())); + setup.experiment.BeamX_pxl(r.beam_x_pxl).BeamY_pxl(r.beam_y_pxl); + setup.mask = r.mask; + logger.Info("Pre-pass on {} images ({:.1f} s): beam centre ({:.2f}, {:.2f}) -> ({:.2f}, {:.2f}) px, " + "masked pixels {} -> {}", r.images, + std::chrono::duration(std::chrono::steady_clock::now() - t0).count(), + experiment.GetBeamX_pxl(), experiment.GetBeamY_pxl(), r.beam_x_pxl, r.beam_y_pxl, + masked(mask), masked(r.mask)); + } catch (const std::exception &e) { + logger.Warning("Pre-pass failed, the images are analysed with the stated beam centre and mask - {}", + e.what()); + } + } + sample.clear(); + setup.mapping = std::make_unique(setup.experiment, setup.mask); + indexer = std::make_unique(setup.experiment, indexing_thread_pool, false, true); + + ready(setup); + prepass_done.count_down(); +} + +void StillsGPUAnalysis::AnalysisThread() { + try { + pin_gpu(); + } catch (const JFJochException &e) { + logger.Warning("Error pinning GPU {}", e.what()); + } + + prepass_done.wait(); + + MXAnalysisWithoutFPGA analysis(setup.experiment, *setup.mapping, setup.mask, *indexer, + /*enable_fused_adaptive_gpu=*/true); + + while (auto held = queue.GetBlocking()) { + DataMessage &message = held->message; + AzimuthalIntegrationProfile profile(*setup.mapping); + const auto start_time = std::chrono::steady_clock::now(); + try { + analysis.Analyze(message, profile, settings()); + } catch (const JFJochException &e) { + // The frame is still written, without its analysis. + logger.ErrorException(e); + } + message.processing_time_s = std::chrono::duration(std::chrono::steady_clock::now() + - start_time).count(); + sink(message, profile); + } +} + +std::optional StillsGPUAnalysis::GetConsensusUnitCell() const { + return indexer->GetConsensusUnitCell(); +} + +float StillsGPUAnalysis::GetBeamX_pxl() const { + return setup.experiment.GetBeamX_pxl(); +} + +float StillsGPUAnalysis::GetBeamY_pxl() const { + return setup.experiment.GetBeamY_pxl(); +} diff --git a/receiver/StillsGPUAnalysis.h b/receiver/StillsGPUAnalysis.h new file mode 100644 index 000000000..1f043219e --- /dev/null +++ b/receiver/StillsGPUAnalysis.h @@ -0,0 +1,85 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "StillsPrePass.h" +#include "../common/DiffractionExperiment.h" +#include "../common/AzimuthalIntegrationMapping.h" +#include "../common/AzimuthalIntegrationProfile.h" +#include "../common/Logger.h" +#include "../common/ThreadSafeFIFO.h" +#include "../image_analysis/IndexAndRefine.h" + +// Serial frames analysed on the GPU as the DECTRIS ones are (MXAnalysisWithoutFPGA), after a pre-pass on +// the first PREPASS_IMAGES (RunStillsPrePass) has decided the beam centre and the mask they are analysed +// with. Every frame given before the pre-pass is done is held; `ready` is called once it is, with the +// geometry and mask the frames are then analysed at, before any of them is analysed; `sink` gets each +// frame once analysed, from one of the analysis threads. +class StillsGPUAnalysis { +public: + constexpr static int64_t PREPASS_IMAGES = 100; + + struct Setup { + DiffractionExperiment experiment; + PixelMask mask; + std::unique_ptr mapping; + }; + using Ready = std::function; + using Sink = std::function; + using Settings = std::function; + + // `start`: the collection's start message, as it would be sent without the pre-pass. + StillsGPUAnalysis(const DiffractionExperiment &experiment, const PixelMask &mask, const StartMessage &start, + IndexerThreadPool *indexing_thread_pool, int64_t nthreads, Logger &logger, + Ready ready, Sink sink, Settings settings); + ~StillsGPUAnalysis(); + + void Hold(const std::shared_ptr &image); + // No frame comes after this: the pre-pass runs on what came, if it has not, and every frame is + // analysed before it returns. + void Finish(); + + // After Finish. + [[nodiscard]] std::optional GetConsensusUnitCell() const; + [[nodiscard]] float GetBeamX_pxl() const; + [[nodiscard]] float GetBeamY_pxl() const; + +private: + const DiffractionExperiment &experiment; + const PixelMask &mask; + const StartMessage start; + IndexerThreadPool *indexing_thread_pool; + Logger &logger; + Ready ready; + Sink sink; + Settings settings; + + Setup setup; + std::unique_ptr indexer; + + std::mutex prepass_mutex; + std::condition_variable prepass_cv; + bool prepass_go = false; + int64_t images_held = 0; + std::map> prepass_images; + std::latch prepass_done{1}; + + ThreadSafeFIFO> queue; + std::future prepass_future; + std::vector> analysis_futures; + bool finished = false; + + void StartPrePass(); + void PrePassThread(); + void AnalysisThread(); +}; diff --git a/receiver/StillsPrePass.cpp b/receiver/StillsPrePass.cpp new file mode 100644 index 000000000..3719ac156 --- /dev/null +++ b/receiver/StillsPrePass.cpp @@ -0,0 +1,64 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "StillsPrePass.h" + +#include "../common/JFJochException.h" +#include "../reader/JFJochReader.h" +#include "../reader/JFJochStartMessageReader.h" +#include "../rugnux/RugnuxLive.h" + +namespace { + // The held images as a collection of their own, numbered from 0 as the collection numbers them. + class HeldImageReader : public JFJochReader { + const std::map> &images; + uint64_t number_of_images; + + bool LoadImage_i(std::shared_ptr &dataset, DataMessage &message, + std::vector &, int64_t image_number, bool) override { + const auto it = images.find(image_number); + if (!dataset || it == images.end()) + return false; + message = it->second->message; + return true; + } + public: + HeldImageReader(const StartMessage &start, const std::map> &in_images) + : images(in_images), number_of_images(in_images.empty() ? 0 : in_images.rbegin()->first + 1) { + SetStartMessage(DatasetFromStartMessage(start, default_experiment)); + } + + [[nodiscard]] uint64_t GetNumberOfImages() const override { return number_of_images; } + void Close() override {} + + bool ReadRawImage(int64_t image_number, JFJochReaderRawImage &ret) override { + const auto it = images.find(image_number); + if (it == images.end()) + return false; + ret.image = it->second->message.image; + return true; + } + + std::vector ReadSpots(int64_t) const override { return {}; } + }; +} + +StillsPrePassResult RunStillsPrePass(const StartMessage &start, const std::map> &images, + int nthreads) { + HeldImageReader reader(start, images); + const auto dataset = reader.GetDataset(); + const RugnuxLiveSetup setup = MakeRugnuxStillsPrePassSetup(*dataset, nthreads); + + Rugnux rugnux(reader, setup.experiment, *dataset->pixel_mask, setup.config); + const ProcessResult result = rugnux.Run(); + Rugnux::ReleasePinnedPool(); + + return StillsPrePassResult{ + // Where the beam lands: the receiver analyses with the detector square to it, as stated, and a tilt + // the run refined moves its PONI, not this. + .beam_x_pxl = static_cast(result.used_direct_beam_x_pxl), + .beam_y_pxl = static_cast(result.used_direct_beam_y_pxl), + .mask = rugnux.GetPixelMask(), + .images = images.size() + }; +} diff --git a/receiver/StillsPrePass.h b/receiver/StillsPrePass.h new file mode 100644 index 000000000..efcb44fb4 --- /dev/null +++ b/receiver/StillsPrePass.h @@ -0,0 +1,32 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include + +#include "../common/JFJochMessages.h" +#include "../common/PixelMask.h" + +// An image the receiver holds until the pre-pass is done: its message, with the image pointing into +// `image`, the compressed frame it owns. +struct HeldImage { + DataMessage message; + std::vector image; +}; + +struct StillsPrePassResult { + float beam_x_pxl; + float beam_y_pxl; + PixelMask mask; + size_t images = 0; +}; + +// The first frames of a serial collection through rugnux, as `rugnux --no-merge` on them: the beam centre +// it decides (from the scattered background, and from the frames that index where there are any) and the +// mask it ends with (the collection's, plus the beam-stop shadow and the defective pixels). `images` by +// number; the start message is the collection's. +StillsPrePassResult RunStillsPrePass(const StartMessage &start, const std::map> &images, + int nthreads); diff --git a/rugnux/Rugnux.h b/rugnux/Rugnux.h index 4a8bea15f..02ed28edb 100644 --- a/rugnux/Rugnux.h +++ b/rugnux/Rugnux.h @@ -1889,6 +1889,9 @@ public: static void ReleasePinnedPool(); // The pre-scan's powder measurement, for a report on a run that ended without a result. const PowderRings &GetPowderRings() const { return powder_; } + // The mask the run ended with: the input's, plus the beam-stop shadow and the defective pixels the + // pre-scan found. + const PixelMask &GetPixelMask() const { return pixel_mask_; } // Request cancellation; safe to call from any thread (the worker loop checks between images). // fast_live makes it a cut: a pass that is past its images still merges and makes its outputs, the diff --git a/rugnux/RugnuxLive.cpp b/rugnux/RugnuxLive.cpp index 90cf5f755..bade93af5 100644 --- a/rugnux/RugnuxLive.cpp +++ b/rugnux/RugnuxLive.cpp @@ -65,6 +65,45 @@ RugnuxLiveSetup MakeRugnuxLiveSetup(const JFJochReaderDataset &dataset, int nthr return s; } +RugnuxLiveSetup MakeRugnuxStillsPrePassSetup(const JFJochReaderDataset &dataset, int nthreads) { + RugnuxLiveSetup s; + + // The experiment as rugnux_cli.cpp sets it up for serial data with no options given. + DiffractionExperiment &x = s.experiment; + x = dataset.experiment; + ApplyRugnuxExperimentDefaults(x); + x.NumTriggers(1); + x.MaxSpotCount(RUGNUX_MAX_SPOT_COUNT); + + IndexingSettings indexing; + indexing.Algorithm(IndexingAlgorithmEnum::Auto); + indexing.RotationIndexing(false); + indexing.GeomRefinementAlgorithm(GeomRefinementAlgorithmEnum::BeamCenter); + indexing.IndexIceRings(false); + x.ImportIndexingSettings(indexing); + x.DetectIceRings(RugnuxDetectIceRings(dataset.file_detect_ice_rings, false)); + ApplyRugnuxIntegrationRadii(x, false); + + ProcessConfig &c = s.config; + c.mode = ProcessMode::FullAnalysis; + c.nthreads = nthreads; + c.nthreads_auto = true; + c.export_unmerged = false; + c.spot_finding = RugnuxDefaultSpotFinding(false); + // The frames come in photons: the receiver converted them. + ApplyRugnuxPixelGain(x, c.spot_finding, 1.0f); + c.rotation_indexing = false; + c.detect_beam_stop = RUGNUX_DETECT_BEAM_STOP_FRAMES; + c.beam_center_check = true; + c.adaptive_integration_radius = true; + c.measure_spot_budget = true; + // As `rugnux -C`: the geometry bundle adjustment needs the cell. + if (x.GetUnitCell().has_value()) + c.refine_geometry = RUGNUX_REFINE_GEOMETRY_FRAMES; + c.run_scaling = false; + return s; +} + void WarmUpRugnuxLiveGPU(const RugnuxLiveSetup &setup) { #ifdef JFJOCH_USE_CUDA // As Rugnux sets them for its first-pass pool (the measured centre's indexing and pass 2 use it as diff --git a/rugnux/RugnuxLive.h b/rugnux/RugnuxLive.h index cda7d1fb7..6372d38a4 100644 --- a/rugnux/RugnuxLive.h +++ b/rugnux/RugnuxLive.h @@ -20,6 +20,11 @@ struct RugnuxLiveSetup { RugnuxLiveSetup MakeRugnuxLiveSetup(const JFJochReaderDataset &dataset, int nthreads, std::shared_ptr model); +// The first frames of a serial collection, processed inside the broker before the rest are analysed: +// what a bare `rugnux --no-merge` (with -C where the collection states a cell) on those frames sets. +// The run's beam centre and mask are what the broker then analyses the whole collection with. +RugnuxLiveSetup MakeRugnuxStillsPrePassSetup(const JFJochReaderDataset &dataset, int nthreads); + // Builds on every card, at the same time, the GPU FFT indexers a live run of this setup is known to // need - its first pass's and its short-axis pass's (the settings Rugnux gives them) - and the cuFFT // plans of its beam-centre capture, and with them the card's CUDA context and cuFFT's one-time diff --git a/tests/JFJochReceiverProcessingTest.cpp b/tests/JFJochReceiverProcessingTest.cpp index ac1e72260..ed36b2fce 100644 --- a/tests/JFJochReceiverProcessingTest.cpp +++ b/tests/JFJochReceiverProcessingTest.cpp @@ -107,6 +107,87 @@ TEST_CASE("JFJochIntegrationTest_ZMQ_crystal_spot_and_index", "[JFJochReceiver]" REQUIRE_NOTHROW(writer_future.get()); } +TEST_CASE("JFJochIntegrationTest_ZMQ_crystal_spot_and_index_GPUStills", "[JFJochReceiver]") { + Logger logger(Catch::getResultCapture().getCurrentTestName()); + + RegisterHDF5Filter(); + + const uint16_t nthreads = 4; + + DiffractionExperiment experiment(DetJF4M()); + experiment.ImagesPerTrigger(5).NumTriggers(1).UseInternalPacketGenerator(true).ImagesPerFile(2) + .FilePrefix("crystal_test_gpu_stills").JungfrauConvPhotonCnt(false).SetFileWriterFormat(FileWriterFormat::NXmxVDS).OverwriteExistingFiles(true) + .DetectorDistance_mm(75).BeamY_pxl(1136).BeamX_pxl(1090).IncidentEnergy_keV(12.4) + .SetUnitCell(UnitCell{.a = 36.9, .b = 78.95, .c = 78.95, .alpha =90, .beta = 90, .gamma = 90}); + experiment.SampleTemperature_K(123.0).RingCurrent_mA(115); + + PixelMask pixel_mask(experiment); + + // Load example image + HDF5ReadOnlyFile data("../../tests/test_data/compression_benchmark.h5"); + HDF5DataSet dataset(data, "/entry/data/data"); + HDF5DataSpace file_space(dataset); + + REQUIRE(file_space.GetDimensions()[2] == experiment.GetXPixelsNum()); + REQUIRE(file_space.GetDimensions()[1] == experiment.GetYPixelsNum()); + std::vector image_conv (file_space.GetDimensions()[1] * file_space.GetDimensions()[2]); + + std::vector start = {4,0,0}; + std::vector file_size = {1, file_space.GetDimensions()[1], file_space.GetDimensions()[2]}; + dataset.ReadVector(image_conv, start, file_size); + + std::vector image_raw_geom(experiment.GetModulesNum() * RAW_MODULE_SIZE); + ConvertedToRawGeometry(experiment, image_raw_geom.data(), image_conv.data()); + logger.Info("Loaded image"); + + // Setup acquisition device + AcquisitionDeviceGroup aq_devices; + std::unique_ptr test = std::make_unique(0, 64); + for (int m = 0; m < experiment.GetModulesNum(); m++) + test->SetInternalGeneratorFrame((uint16_t *) image_raw_geom.data() + m * RAW_MODULE_SIZE, m); + + aq_devices.Add(std::move(test)); + + ZMQStream2Pusher pusher({"ipc://*"}); + pusher.WriterNotificationSocket("ipc://*"); + + ZMQImagePuller puller(pusher.GetAddress()[0]); + StreamWriter writer(logger, puller); + auto writer_future = std::async(std::launch::async, &StreamWriter::Run, &writer); + + JFJochReceiverService service(aq_devices, logger, pusher); + service.NumThreads(nthreads); + service.Indexing(experiment.GetIndexingSettings()); + // Serial data with rugnux on: the converted frames are analysed on the GPU, after the pre-pass. + service.RugnuxLive(true, nullptr, 4, 60); + + // No progress value at the start of measurement + REQUIRE(!service.GetProgress().has_value()); + + SpotFindingSettings settings = DiffractionExperiment::DefaultDataProcessingSettings(); + settings.signal_to_noise_threshold = 2.5; + settings.photon_count_threshold = 5; + settings.min_pix_per_spot = 1; + settings.max_pix_per_spot = 200; + settings.high_resolution_limit = 2.0; + settings.low_resolution_limit = 50.0; + service.SetSpotFindingSettings(settings); + + service.Start(experiment, pixel_mask, nullptr); + auto receiver_out = service.Stop(); + + CHECK(receiver_out.efficiency == 1.0); + REQUIRE(receiver_out.status.indexing_rate); + CHECK(receiver_out.status.indexing_rate.value() == 1.0); + CHECK(receiver_out.status.images_sent == experiment.GetImageNum()); + CHECK(!receiver_out.status.cancelled); + + // No progress value at the end of measurement + REQUIRE(!service.GetProgress().has_value()); + + REQUIRE_NOTHROW(writer_future.get()); +} + TEST_CASE("JFJochIntegrationTest_ZMQ_crystal_spot_and_index_refinement", "[JFJochReceiver]") { Logger logger(Catch::getResultCapture().getCurrentTestName()); diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 12c0e427f..18abcf779 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt @@ -46,6 +46,10 @@ IF(NOT JFJOCH_PORTABLE_ONLY) ADD_EXECUTABLE(jfjoch_stream2_replay jfjoch_stream2_replay.cpp) TARGET_LINK_LIBRARIES(jfjoch_stream2_replay JFJochReader JFJochZMQ CBORStream2FrameSerialize JFJochCommon) + # Not installed: replays a stored serial dataset through the GPU analysis of converted JUNGFRAU frames. + ADD_EXECUTABLE(jfjoch_stills_gpu_replay jfjoch_stills_gpu_replay.cpp) + TARGET_LINK_LIBRARIES(jfjoch_stills_gpu_replay JFJochReceiver JFJochWriter JFJochReader) + ADD_EXECUTABLE(jfjoch_pcie_status jfjoch_pcie_status.cpp) TARGET_LINK_LIBRARIES(jfjoch_pcie_status JFJochDevice JFJochCommon) INSTALL(TARGETS jfjoch_pcie_status RUNTIME COMPONENT jfjoch) diff --git a/tools/jfjoch_stills_gpu_replay.cpp b/tools/jfjoch_stills_gpu_replay.cpp new file mode 100644 index 000000000..7744e00f4 --- /dev/null +++ b/tools/jfjoch_stills_gpu_replay.cpp @@ -0,0 +1,149 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +// Replays the frames of a stored serial collection through the broker's GPU analysis of converted +// JUNGFRAU frames (StillsGPUAnalysis): the pre-pass on the first images, then every frame analysed with the +// beam centre and mask it decided. Reports the pre-pass, the indexed fraction and the sustained rate, and +// writes the mask the frames were analysed with (_mask.u32, uint32, row by row). + +#include +#include +#include +#include +#include +#include +#include + +#include "../common/Logger.h" +#include "../image_analysis/indexing/IndexerThreadPool.h" +#include "../reader/JFJochHDF5Reader.h" +#include "../receiver/StillsGPUAnalysis.h" +#include "../writer/HDF5Objects.h" + +void print_usage() { + std::cout << "Usage: jfjoch_stills_gpu_replay [options] " << std::endl; + std::cout << " -o output prefix (default stills_gpu_replay)" << std::endl; + std::cout << " -s first image (default 0)" << std::endl; + std::cout << " -n images to replay (default all)" << std::endl; + std::cout << " -r images given per second (default 0 = as fast as they are read)" << std::endl; + std::cout << " -N analysis threads (default 16)" << std::endl; + std::cout << " -x , -y beam centre, overriding the file" << std::endl; + std::cout << " -C unit cell of the collection" << std::endl; + std::cout << " -c leave the file's unit cell out of the collection" << std::endl; +} + +int main(int argc, char **argv) { + Logger logger("jfjoch_stills_gpu_replay"); + RegisterHDF5Filter(); + + std::string prefix = "stills_gpu_replay"; + int64_t first = 0; + std::optional images; + double rate = 0.0; + int nthreads = 16; + std::optional beam_x, beam_y; + std::optional cell; + bool no_cell = false; + + int opt; + while ((opt = getopt(argc, argv, "o:s:n:r:N:x:y:C:c")) != -1) { + switch (opt) { + case 'o': prefix = optarg; break; + case 's': first = atoll(optarg); break; + case 'n': images = atoll(optarg); break; + case 'r': rate = atof(optarg); break; + case 'N': nthreads = atoi(optarg); break; + case 'x': beam_x = atof(optarg); break; + case 'y': beam_y = atof(optarg); break; + case 'c': no_cell = true; break; + case 'C': { + UnitCell c{}; + if (sscanf(optarg, "%f,%f,%f,%f,%f,%f", &c.a, &c.b, &c.c, &c.alpha, &c.beta, &c.gamma) != 6) { + print_usage(); + return EXIT_FAILURE; + } + cell = c; + break; + } + default: print_usage(); return EXIT_FAILURE; + } + } + if (optind + 1 != argc) { + print_usage(); + return EXIT_FAILURE; + } + + JFJochHDF5Reader reader; + reader.ReadFile(argv[optind]); + const auto dataset = reader.GetDataset(); + const int64_t n = std::min(images.value_or(reader.GetNumberOfImages() - first), + reader.GetNumberOfImages() - first); + const auto stored = reader.GetStoredPixelFormat(); + + DiffractionExperiment x(dataset->experiment); + x.ImagesPerTrigger(n).NumTriggers(1); + if (beam_x) + x.BeamX_pxl(*beam_x); + if (beam_y) + x.BeamY_pxl(*beam_y); + if (no_cell) + x.SetUnitCell(std::nullopt); + if (cell) + x.SetUnitCell(cell); + const PixelMask &mask = *dataset->pixel_mask; + + StartMessage start; + x.FillMessage(start); + start.number_of_images = n; + start.bit_depth_image = stored.bit_depth; + start.pixel_signed = stored.is_signed; + start.pixel_mask["default"] = mask.GetMask(x); + + IndexerThreadPool pool(x.GetIndexingSettings()); + + std::atomic analysed = 0, indexed = 0; + std::chrono::steady_clock::time_point t_ready; + const auto t0 = std::chrono::steady_clock::now(); + { + StillsGPUAnalysis gpu(x, mask, start, &pool, nthreads, logger, + [&](const StillsGPUAnalysis::Setup &s) { + t_ready = std::chrono::steady_clock::now(); + const auto &m = s.mask.GetMask(); + std::ofstream f(prefix + "_mask.u32", std::ios::binary); + f.write(reinterpret_cast(m.data()), m.size() * sizeof(uint32_t)); + logger.Info("Ready after {:.1f} s: beam centre {:.2f} {:.2f}", + std::chrono::duration(t_ready - t0).count(), + s.experiment.GetBeamX_pxl(), s.experiment.GetBeamY_pxl()); + }, + [&](DataMessage &m, const AzimuthalIntegrationProfile &) { + ++analysed; + if (m.indexing_result.value_or(false)) + ++indexed; + }, + [] { return DiffractionExperiment::DefaultDataProcessingSettings(); }); + + JFJochReaderRawImage raw; + for (int64_t i = 0; i < n; i++) { + if (rate > 0) + std::this_thread::sleep_until(t0 + std::chrono::duration(static_cast(i) / rate)); + reader.ReadRawImage(first + i, raw); + auto held = std::make_shared(); + held->image.assign(raw.image.GetCompressed(), raw.image.GetCompressed() + raw.image.GetCompressedSize()); + held->message.number = i; + held->message.image = CompressedImage(held->image.data(), held->image.size(), raw.image.GetWidth(), + raw.image.GetHeight(), raw.image.GetMode(), + raw.image.GetCompressionAlgorithm()); + held->message.image_collection_efficiency = 1.0f; + gpu.Hold(held); + } + logger.Info("Gave {} images in {:.1f} s", n, + std::chrono::duration(std::chrono::steady_clock::now() - t0).count()); + gpu.Finish(); + } + const auto t1 = std::chrono::steady_clock::now(); + const double after_ready = std::chrono::duration(t1 - t_ready).count(); + logger.Info("Analysed {} images, {} indexed ({:.1f}%); {:.1f} s after the pre-pass = {:.1f} Hz; total {:.1f} s", + analysed.load(), indexed.load(), 100.0 * indexed / std::max(1, analysed), after_ready, + analysed / after_ready, std::chrono::duration(t1 - t0).count()); + return EXIT_SUCCESS; +}