jungfrauandClaude Opus 5
2cbb3fc8b4
Every worker thread built a full set of analysis engines. Two of them are never asked for on the offline path: the fixed-threshold spot finder, because detection is adaptive by default, and the azimuthal integrator, because the fused adaptive finder produces the profile as a by-product. They are still needed elsewhere - the broker defaults to non-adaptive detection, and --no-adaptive-spots asks for the finder - so they are built on first use rather than removed. A lazily built finder takes the current resolution mask on construction; without that it would find spots outside the limits it was never told about. The bitshuffle decoder sized its output buffer for the widest pixel type there is rather than the one the images actually have, holding a second full frame per worker on 16-bit data. It is sized from the image now and grows if a later frame needs more. The shared-table checksum runs over eight interleaved lanes. FNV's multiply is a loop-carried dependency, so one chain retires a byte every few cycles whatever memory bandwidth is spare, and every worker hashes tens of megabytes of geometry tables as it builds its engines - about 5% of all CPU samples on a 16M-pixel detector. Measured on a 16M-pixel rotation dataset: cudaMalloc 11314 -> 9474 calls and, with cudaFree, 117 s -> 78 s of aggregate thread time; both synchronise the whole device, so that time is spent blocking every other worker. Whole battery 15m32s -> 12m30s. Data quality against main, over 24 crystals and eight statistics each: the same space group on all 24, and every difference smaller than what two runs of an IDENTICAL binary produce (measured: 13 of 24 crystals reproduce exactly run to run, worst R_meas swing 5.5 points, against 4.6 points for main vs this branch). The float atomics in the reductions have always made this so. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
379 lines
20 KiB
C++
379 lines
20 KiB
C++
// SPDX-FileCopyrightText: 2024 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "MXAnalysisWithoutFPGA.h"
|
|
|
|
#include <algorithm>
|
|
#include <spdlog/spdlog.h>
|
|
|
|
#include "spot_finding/StrongPixelSet.h"
|
|
#include "../compression/JFJochDecompress.h"
|
|
|
|
#include "spot_finding/SpotUtils.h"
|
|
#include "bragg_prediction/BraggPredictionFactory.h"
|
|
#include "image_preprocessing/ImagePreprocessorCPU.h"
|
|
|
|
#include "azint/AzIntEngineCPU.h"
|
|
#include "roi/ROIIntegrationCPU.h"
|
|
#include "spot_finding/ImageSpotFinderCPU.h"
|
|
#include "spot_finding/AdaptiveSpotFinderCPU.h"
|
|
#include "bragg_integration/BraggIntegrationEngineCPU.h"
|
|
#ifdef JFJOCH_USE_CUDA
|
|
#include "azint/AzIntEngineGPU.h"
|
|
#include "roi/ROIIntegrationGPU.h"
|
|
#include "spot_finding/ImageSpotFinderGPU.h"
|
|
#include "spot_finding/AdaptiveSpotFinderGPU.h"
|
|
#include "image_preprocessing/ImagePreprocessorGPU.h"
|
|
#include "image_preprocessing/ImagePreprocessorBufferGPU.h"
|
|
#include "bragg_integration/BraggIntegrationEngineGPU.h"
|
|
#include "../common/CUDAWrapper.h"
|
|
#endif
|
|
|
|
|
|
MXAnalysisWithoutFPGA::MXAnalysisWithoutFPGA(const DiffractionExperiment &in_experiment,
|
|
const AzimuthalIntegrationMapping &in_integration,
|
|
const PixelMask &in_mask,
|
|
IndexAndRefine &in_indexer,
|
|
bool in_enable_fused_adaptive_gpu)
|
|
: experiment(in_experiment),
|
|
integration(in_integration),
|
|
enable_fused_adaptive_gpu(in_enable_fused_adaptive_gpu),
|
|
npixels(experiment.GetPixelsNum()),
|
|
xpixels(experiment.GetXPixelsNum()),
|
|
indexer(in_indexer),
|
|
prediction(CreateBraggPrediction(experiment.IsRotationIndexing())),
|
|
mask(in_mask),
|
|
mask_resolution(experiment.GetPixelsNum(), false),
|
|
mask_high_res(-1),
|
|
mask_low_res(-1) {
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (get_gpu_count() == 0) {
|
|
#endif
|
|
preprocessor_buffer = std::make_unique<ImagePreprocessorBuffer>(experiment.GetPixelsNum());
|
|
preprocessor = std::make_unique<ImagePreprocessorCPU>(in_experiment, in_mask);
|
|
bragg_engine = std::make_unique<BraggIntegrationEngineCPU>(in_experiment);
|
|
if (experiment.ROI().size() >= 1)
|
|
roi = std::make_unique<ROIIntegrationCPU>(experiment);
|
|
#ifdef JFJOCH_USE_CUDA
|
|
} else {
|
|
stream = std::make_shared<CudaStream>();
|
|
// The host copy of the preprocessed image is only read when a CPU engine wants it, which is
|
|
// the same condition that drives copy_image_to_host below. Skipping it also skips page-locking
|
|
// 4 bytes per pixel per worker.
|
|
preprocessor_buffer = std::make_unique<ImagePreprocessorBufferGPU>(
|
|
experiment.GetPixelsNum(), /*host_mirror=*/!enable_fused_adaptive_gpu);
|
|
// The preprocessed image only has to come back to the host if a CPU engine reads it. Every
|
|
// engine built below runs on the GPU, except the CPU adaptive finder that is kept when the fused
|
|
// GPU engine is off - so that is the one case that needs the copy. Every caller currently passes
|
|
// enable_fused_adaptive_gpu = true, so on the GPU path the copy is off in practice.
|
|
preprocessor = std::make_unique<ImagePreprocessorGPU>(in_experiment, in_mask, stream,
|
|
/*copy_image_to_host=*/!enable_fused_adaptive_gpu);
|
|
bragg_engine = std::make_unique<BraggIntegrationEngineGPU>(in_experiment, stream);
|
|
if (experiment.ROI().size() >= 1)
|
|
roi = std::make_unique<ROIIntegrationGPU>(experiment, stream);
|
|
if (enable_fused_adaptive_gpu) {
|
|
// One GPU engine that computes the azimuthal profile and the adaptive spot mask in a single
|
|
// image pass. fused_adaptive aliases it so Analyze() can lift the profile out of it.
|
|
auto fused = std::make_unique<AdaptiveSpotFinderGPU>(integration, stream);
|
|
fused_adaptive = fused.get();
|
|
adaptiveSpotFinder = std::move(fused);
|
|
}
|
|
}
|
|
#endif
|
|
if (!adaptiveSpotFinder)
|
|
adaptiveSpotFinder = std::make_unique<AdaptiveSpotFinderCPU>(integration);
|
|
}
|
|
|
|
void MXAnalysisWithoutFPGA::Analyze(DataMessage &output,
|
|
AzimuthalIntegrationProfile &profile,
|
|
const SpotFindingSettings &spot_finding_settings) {
|
|
if ((output.image.GetWidth() != xpixels)
|
|
|| (output.image.GetWidth() * output.image.GetHeight() != npixels))
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
|
|
"Mismatch in pixel size");
|
|
|
|
// Decompress on the device where the preprocessor can, so only the compressed chunk crosses PCIe
|
|
// and the host does no decompression at all. AnalyzeCompressed says whether it took the image;
|
|
// when it declines (a CPU preprocessor, or an algorithm with no device decoder) fall through to
|
|
// the host route unchanged. The two produce the same preprocessed image.
|
|
const auto compression_start_time = std::chrono::steady_clock::now();
|
|
ImageStatistics ret{};
|
|
bool decoded_on_device = false;
|
|
try {
|
|
decoded_on_device = preprocessor->AnalyzeCompressed(*preprocessor_buffer, output.image, ret);
|
|
} catch (const JFJochException &e) {
|
|
// The device route must never be the reason a frame fails: whatever it could not handle, the
|
|
// host decoder gets its turn. If the data really is bad the host throws too and the caller
|
|
// sees the same error it saw before any of this existed - but a GPU-side problem costs speed
|
|
// rather than the acquisition.
|
|
spdlog::warn("Device decoding failed ({}), falling back to host decompression", e.what());
|
|
decoded_on_device = false;
|
|
}
|
|
const auto compression_end_time = std::chrono::steady_clock::now();
|
|
|
|
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)
|
|
output.compression_time_s = std::chrono::duration<float>(decompressed_time - compression_start_time).count();
|
|
|
|
const auto preprocessing_start_time = std::chrono::steady_clock::now();
|
|
ret = preprocessor->Analyze(*preprocessor_buffer, image_ptr, output.image.GetMode());
|
|
const auto preprocessing_end_time = std::chrono::steady_clock::now();
|
|
output.preprocessing_time_s = std::chrono::duration<float>(preprocessing_end_time - preprocessing_start_time).count();
|
|
} else {
|
|
// Decode and preprocess are one device operation here, but the decompression is still a real,
|
|
// separately measurable cost - the decoder brackets it with CUDA events - so it is still
|
|
// reported as one. Leaving compression_time_s unset instead would blank the broker's
|
|
// "compression" plot trace and fill /entry/profiling/compressionTime with NaN.
|
|
const float total_s = std::chrono::duration<float>(compression_end_time - compression_start_time).count();
|
|
const float decompress_s = std::min(preprocessor->GetLastDecompressionTime_s(), total_s);
|
|
output.compression_time_s = decompress_s;
|
|
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;
|
|
|
|
if (!fused) {
|
|
const auto azint_start_time = std::chrono::steady_clock::now();
|
|
AzInt().Run(*preprocessor_buffer, profile);
|
|
const auto azint_end_time = std::chrono::steady_clock::now();
|
|
output.azint_time_s = std::chrono::duration<float>(azint_end_time - azint_start_time).count();
|
|
}
|
|
|
|
if (roi)
|
|
roi->Run(*preprocessor_buffer, output.roi);
|
|
|
|
if (spot_finding_settings.enable) {
|
|
// Update resolution mask
|
|
if (mask_high_res != spot_finding_settings.high_resolution_limit
|
|
|| mask_low_res != spot_finding_settings.low_resolution_limit)
|
|
UpdateMaskResolution(spot_finding_settings);
|
|
|
|
ImageSpotFinder &finder = spot_finding_settings.adaptive_threshold
|
|
? static_cast<ImageSpotFinder &>(*adaptiveSpotFinder)
|
|
: FixedThresholdFinder();
|
|
const auto integrate_fn = [this](const std::vector<Reflection> &predicted, size_t npredicted,
|
|
int64_t image_number) {
|
|
return bragg_engine->Run(*preprocessor_buffer, predicted, npredicted, image_number);
|
|
};
|
|
|
|
// The radial background correction has to be decided BEFORE this image is integrated, so the
|
|
// ice score is taken here rather than with the other per-image quantities at the end of the
|
|
// function. It needs the peak-excluded per-ring background, which the adaptive finder has as
|
|
// soon as it has detected - so this must be called after detection and before integration.
|
|
// Where no such background exists (no adaptive finder), auto leaves the correction off: the
|
|
// plain profile carries the Bragg peaks and cannot support an absolute threshold.
|
|
const auto decide_radial_background = [this, &spot_finding_settings, &output]() {
|
|
if (!bragg_engine->IsBackgroundRadialAuto())
|
|
return;
|
|
// Same condition as the score at the end of this function: the adaptive finder holds its
|
|
// ring background from whenever it last ran, so requiring that it ran for THIS image is
|
|
// what keeps a stale curve out.
|
|
if (!spot_finding_settings.adaptive_threshold)
|
|
return;
|
|
const std::vector<float> &ring_bkg = adaptiveSpotFinder->GetRingBackground();
|
|
if (ring_bkg.empty())
|
|
return;
|
|
output.ice_ring_score = AzimuthalIntegrationProfile::IceRingScore(
|
|
ring_bkg, integration.GetQBinCount(), integration.Settings(),
|
|
spot_finding_settings.ice_ring_width_Q_recipA);
|
|
bragg_engine->BackgroundRadial(*output.ice_ring_score
|
|
>= experiment.GetScalingSettings().GetIceMinScore());
|
|
};
|
|
|
|
// A missing min-pix (std::nullopt) means "choose it per image". This applies only to the stills
|
|
// indexing path (each frame is indexed independently); rotation indexing builds one lattice from
|
|
// all frames, so it keeps the fixed min-pix and the single-pass finder.
|
|
const bool adaptive_min_pix = !spot_finding_settings.min_pix_per_spot.has_value()
|
|
&& spot_finding_settings.indexing
|
|
&& !experiment.IsRotationIndexing();
|
|
if (adaptive_min_pix) {
|
|
// Choose the per-image min-pix adaptively instead of a fixed one. min-pix filters connected
|
|
// components AFTER detection, so BOTH the detection (the expensive per-pixel pass) and the
|
|
// connected-component search run ONCE and only the filter is repeated; the azimuthal
|
|
// profile is the one Detect() computed.
|
|
// Index at 3/2/1 (index-only, no integration/accumulation), keep whichever maximises
|
|
// n_indexed^2 / n_total (indexed count weighted by indexed fraction) together with its spot
|
|
// list, and integrate that one.
|
|
const auto detect_start_time = std::chrono::steady_clock::now();
|
|
finder.Detect(*preprocessor_buffer, spot_finding_settings);
|
|
const auto &components = finder.ExtractComponents(*preprocessor_buffer, spot_finding_settings);
|
|
float spot_finding_time_s =
|
|
std::chrono::duration<float>(std::chrono::steady_clock::now() - detect_start_time).count();
|
|
float indexing_time_s = 0.0f;
|
|
|
|
SpotFindingSettings s = spot_finding_settings;
|
|
std::vector<DiffractionSpot> best_spots;
|
|
int best_mp = 0;
|
|
double best_score = -1.0;
|
|
for (int mp : {3, 2, 1}) {
|
|
s.min_pix_per_spot = mp;
|
|
const auto extract_start_time = std::chrono::steady_clock::now();
|
|
std::vector<DiffractionSpot> spots = ImageSpotFinder::Filter(components, s);
|
|
spot_finding_time_s +=
|
|
std::chrono::duration<float>(std::chrono::steady_clock::now() - extract_start_time).count();
|
|
SpotAnalyze(experiment, s, spots, output);
|
|
const bool indexed = indexer.IndexFrameOnly(output, s);
|
|
indexing_time_s += output.indexing_time_s.value_or(0.0f);
|
|
if (indexed) {
|
|
const double n_idx = static_cast<double>(output.spot_count_indexed.value_or(0));
|
|
const double n_tot = static_cast<double>(std::max<int64_t>(1, output.spot_count.value_or(1)));
|
|
const double score = n_idx * n_idx / n_tot;
|
|
if (score > best_score) {
|
|
best_score = score;
|
|
best_mp = mp;
|
|
best_spots = std::move(spots);
|
|
}
|
|
}
|
|
}
|
|
if (best_mp != 0) {
|
|
// Index and integrate the winning spot list; no spot finding left to do.
|
|
s.min_pix_per_spot = best_mp;
|
|
SpotAnalyze(experiment, s, best_spots, output);
|
|
decide_radial_background();
|
|
indexer.ProcessImage(output, s, *prediction, integrate_fn);
|
|
indexing_time_s += output.indexing_time_s.value_or(0.0f);
|
|
}
|
|
// Each indexer call reports only its own time, so the escalation's total is summed here.
|
|
output.spot_finding_time_s = spot_finding_time_s;
|
|
output.indexing_time_s = indexing_time_s;
|
|
} else {
|
|
const auto spot_finding_start_time = std::chrono::steady_clock::now();
|
|
const std::vector<DiffractionSpot> spots = finder.Run(*preprocessor_buffer, spot_finding_settings);
|
|
SpotAnalyze(experiment, spot_finding_settings, spots, output);
|
|
output.spot_finding_time_s = std::chrono::duration<float>(std::chrono::steady_clock::now() - spot_finding_start_time).count();
|
|
decide_radial_background();
|
|
if (spot_finding_settings.indexing)
|
|
indexer.ProcessImage(output, spot_finding_settings, *prediction, integrate_fn);
|
|
}
|
|
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (fused) {
|
|
// Lift the azimuthal profile the fused engine computed in the same detection pass; its azint
|
|
// cost is folded into spot_finding_time_s above.
|
|
profile.Clear(integration);
|
|
profile += fused_adaptive->GetProfile();
|
|
output.azint_time_s = 0.0f;
|
|
}
|
|
#endif
|
|
}
|
|
|
|
output.max_viable_pixel_value = ret.max_value;
|
|
output.min_viable_pixel_value = ret.min_value;
|
|
output.error_pixel_count = ret.error_pixel_count;
|
|
output.saturated_pixel_count = ret.saturated_pixel_count;
|
|
output.az_int_profile = profile.GetResult();
|
|
output.az_int_profile_count = profile.GetPixelCount();
|
|
output.az_int_profile_std = profile.GetStd();
|
|
|
|
output.bkg_estimate = profile.GetBkgEstimate(integration.Settings());
|
|
|
|
// The ice score wants a radial profile with the Bragg peaks taken OUT of it. The azimuthal profile
|
|
// is a plain per-ring mean, so a strong low-resolution reflection landing in a ring's bin is
|
|
// indistinguishable from ice sitting there - measured, that alone lifts clean crystals to a score
|
|
// of 1.5-4.2, right into the range real ice occupies. The adaptive spot finder already computes
|
|
// exactly what is wanted: a sigma-clipped per-ring background, in the same bins, from which the
|
|
// peaks have been removed (an ice ring is azimuthally smooth, so it survives the clip). It is in
|
|
// raw counts rather than corrected ones, which the score does not care about - it is a ratio to the
|
|
// background interpolated under the ring, and the corrections are smooth in radius.
|
|
const std::vector<float> &ring_bkg = adaptiveSpotFinder->GetRingBackground();
|
|
const bool have_ring_bkg = spot_finding_settings.enable && spot_finding_settings.adaptive_threshold
|
|
&& !ring_bkg.empty();
|
|
output.ice_ring_score = AzimuthalIntegrationProfile::IceRingScore(
|
|
have_ring_bkg ? ring_bkg : profile.GetResult1D(), integration.GetQBinCount(),
|
|
integration.Settings(), spot_finding_settings.ice_ring_width_Q_recipA);
|
|
}
|
|
|
|
ImageSpotFinder &MXAnalysisWithoutFPGA::FixedThresholdFinder() {
|
|
if (!spotFinder) {
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (stream)
|
|
spotFinder = std::make_unique<ImageSpotFinderGPU>(experiment.GetXPixelsNum(),
|
|
experiment.GetYPixelsNum(), stream);
|
|
else
|
|
#endif
|
|
spotFinder = std::make_unique<ImageSpotFinderCPU>(experiment.GetXPixelsNum(),
|
|
experiment.GetYPixelsNum());
|
|
// It missed every mask update that happened before it existed, so it takes the current one
|
|
// now. Without this it would find spots outside the resolution limits.
|
|
if (mask_high_res.has_value() || mask_low_res.has_value())
|
|
spotFinder->SetResolutionMask(mask_resolution);
|
|
}
|
|
return *spotFinder;
|
|
}
|
|
|
|
AzIntEngine &MXAnalysisWithoutFPGA::AzInt() {
|
|
if (!azint) {
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (stream)
|
|
azint = std::make_unique<AzIntEngineGPU>(integration, stream);
|
|
else
|
|
#endif
|
|
azint = std::make_unique<AzIntEngineCPU>(integration);
|
|
}
|
|
return *azint;
|
|
}
|
|
|
|
void MXAnalysisWithoutFPGA::RebuildROI() {
|
|
if (experiment.ROI().empty()) {
|
|
roi.reset();
|
|
return;
|
|
}
|
|
#ifdef JFJOCH_USE_CUDA
|
|
if (stream) {
|
|
roi = std::make_unique<ROIIntegrationGPU>(experiment, stream);
|
|
return;
|
|
}
|
|
#endif
|
|
roi = std::make_unique<ROIIntegrationCPU>(experiment);
|
|
}
|
|
|
|
void MXAnalysisWithoutFPGA::AnalyzeROIOnly(DataMessage &output) {
|
|
if ((output.image.GetWidth() != xpixels)
|
|
|| (output.image.GetWidth() * output.image.GetHeight() != npixels))
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
|
|
"Mismatch in pixel size");
|
|
|
|
const uint8_t *image_ptr = Decompress(output.image);
|
|
preprocessor->Analyze(*preprocessor_buffer, image_ptr, output.image.GetMode());
|
|
RunROIOnly(output);
|
|
}
|
|
|
|
const uint8_t *MXAnalysisWithoutFPGA::Decompress(const CompressedImage &image) {
|
|
// An uncompressed image is read straight out of the message and never touches decompression_buffer,
|
|
// so it stays in pageable memory - the buffer is only worth page-locking when it is actually used.
|
|
if (image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION)
|
|
preprocessor->PinInputBuffer(decompression_buffer, image.GetUncompressedSize());
|
|
return image.GetUncompressedPtr(decompression_buffer);
|
|
}
|
|
|
|
void MXAnalysisWithoutFPGA::RunROIOnly(DataMessage &output) {
|
|
output.roi.clear();
|
|
if (roi)
|
|
roi->Run(*preprocessor_buffer, output.roi);
|
|
}
|
|
|
|
void MXAnalysisWithoutFPGA::UpdateMaskResolution(const SpotFindingSettings &settings) {
|
|
mask_low_res = settings.low_resolution_limit;
|
|
mask_high_res = settings.high_resolution_limit;
|
|
// An unset limit masks nothing at that end. At the high-resolution end 0 does that on its own - no
|
|
// pixel has d < 0, and the detector's own edge is where the pixels stop anyway; at the low-resolution
|
|
// end every pixel lies above any finite stand-in, so it takes an infinite one.
|
|
const float high_res = mask_high_res.value_or(0.0f);
|
|
const float low_res = mask_low_res.value_or(INFINITY);
|
|
auto const &resolution_map = integration.Resolution();
|
|
for (int i = 0; i < mask_resolution.size(); i++)
|
|
mask_resolution[i] = (resolution_map[i] > low_res) || (resolution_map[i] < high_res);
|
|
|
|
// The finders keep their own copy (the GPU ones a bit-packed device copy), so the mask is handed
|
|
// over here - when the limits change - rather than with every image.
|
|
if (spotFinder)
|
|
spotFinder->SetResolutionMask(mask_resolution);
|
|
adaptiveSpotFinder->SetResolutionMask(mask_resolution);
|
|
}
|