leonarski_fandClaude Opus 5.5
301762a8da
With a cell given, the live stills pre-pass ran again on all the frames held each time their number doubled, until half the image buffer was full. On a chip with a low hit rate the sample grew to tens of thousands of frames; the last run on them took long enough for the queue to fill, and every frame arriving until it ended was refused and not written. The sample now grows to PREPASS_MAX_IMAGES (4 x PREPASS_IMAGES) or half of what the queue holds, whichever is smaller, and frames that arrived during the previous run do not make the next one larger than that: the last run is short and the other half of the queue absorbs the frames arriving meanwhile. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HVeDQsU5a8AJUdJPgWevAK
420 lines
18 KiB
C++
420 lines
18 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "StillsGPUAnalysis.h"
|
|
|
|
#include <thread>
|
|
|
|
#include "../common/CUDAWrapper.h"
|
|
#include "../common/ThreadAffinity.h"
|
|
#include "../image_analysis/MXAnalysisWithoutFPGA.h"
|
|
#include "../rugnux/RugnuxDefaults.h"
|
|
|
|
StillsGPUAnalysis::StillsGPUAnalysis(const DiffractionExperiment &in_experiment, const PixelMask &in_mask,
|
|
const StartMessage &in_start, IndexerThreadPool *in_indexing_thread_pool,
|
|
int64_t nthreads, int64_t in_max_held, Logger &in_logger, Ready in_ready,
|
|
Sink in_sink, Settings in_settings, Failed in_failed,
|
|
RugnuxLiveStatus *in_status, bool in_merge)
|
|
: 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)),
|
|
setup_failed(std::move(in_failed)), status(in_status), merge_frames(in_merge),
|
|
max_held(std::max<int64_t>(1, in_max_held)), queue(max_held) {
|
|
live.images_needed = std::min(PREPASS_IMAGES, experiment.GetImageNum());
|
|
live.file_beam_x_pxl = experiment.GetBeamX_pxl();
|
|
live.file_beam_y_pxl = experiment.GetBeamY_pxl();
|
|
live.file_distance_mm = experiment.GetDetectorDistance_mm();
|
|
if (status) {
|
|
status->Start(experiment.GetImageNum(),
|
|
static_cast<double>(experiment.GetImageNum())
|
|
* std::chrono::duration<double>(experiment.GetImageTime()).count(),
|
|
std::chrono::steady_clock::now());
|
|
status->Step("Pre-pass");
|
|
Publish();
|
|
}
|
|
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, at least the first {} images held for the pre-pass", nthreads, PREPASS_IMAGES);
|
|
}
|
|
|
|
StillsGPUAnalysis::~StillsGPUAnalysis() {
|
|
try {
|
|
if (!finished)
|
|
Finish();
|
|
} catch (const std::exception &e) {
|
|
logger.Error("GPU analysis: {}", e.what());
|
|
}
|
|
if (status && outcome)
|
|
status->Finish(outcome->first, outcome->second);
|
|
}
|
|
|
|
void StillsGPUAnalysis::AnalysisError(const std::string &text) {
|
|
logger.Error("GPU analysis: {}", text);
|
|
std::unique_lock ul(error_mutex);
|
|
if (analysis_error.empty())
|
|
analysis_error = text;
|
|
}
|
|
|
|
void StillsGPUAnalysis::Publish() {
|
|
std::unique_lock ul(live_mutex);
|
|
if (live.prepass == RugnuxLivePrePass::Done || live.prepass == RugnuxLivePrePass::Failed) {
|
|
const double t = std::chrono::duration<double>(std::chrono::steady_clock::now() - ready_time).count();
|
|
if (t > 0 && live.images_analysed > 0)
|
|
live.analysis_rate_hz = static_cast<double>(live.images_analysed) / t;
|
|
}
|
|
status->Stills(live);
|
|
}
|
|
|
|
bool StillsGPUAnalysis::Hold(const std::shared_ptr<HeldImage> &image) {
|
|
if (status)
|
|
status->ImageReceived(std::chrono::steady_clock::now());
|
|
if (!queue.Put(image)) {
|
|
// As many frames wait as the buffer is to hold: the pre-pass runs on what it has, for the last time.
|
|
StartPrePass();
|
|
return false;
|
|
}
|
|
{
|
|
std::unique_lock ul(prepass_mutex);
|
|
if (!prepass_over) {
|
|
prepass_images[image->message.number] = image;
|
|
std::unique_lock ul_live(live_mutex);
|
|
live.images_held = static_cast<int64_t>(prepass_images.size());
|
|
}
|
|
}
|
|
prepass_cv.notify_all();
|
|
if (status)
|
|
Publish();
|
|
return true;
|
|
}
|
|
|
|
void StillsGPUAnalysis::StartPrePass() {
|
|
{
|
|
std::unique_lock ul(prepass_mutex);
|
|
prepass_last = 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();
|
|
if (status) {
|
|
last_image = status->LastImage();
|
|
Publish();
|
|
}
|
|
if (merge_future.valid()) {
|
|
if (status)
|
|
status->Step("Merge");
|
|
merge_queue.Put(-1);
|
|
merge_future.get();
|
|
if (merged && status) {
|
|
PublishMerge(*merged, "the collection");
|
|
status->Headline(merge_experiment.GetSpaceGroupOrP1().xhm(), merged->merge.statistics.overall.d_min,
|
|
std::nullopt);
|
|
}
|
|
}
|
|
if (status) {
|
|
if (setup_error)
|
|
status->Finish(RugnuxLiveState::Failed, "GPU analysis setup failed");
|
|
else {
|
|
// Reported by Written, once everything is written - the end message and the writer go on meanwhile.
|
|
status->Step("Writing");
|
|
std::unique_lock ul(error_mutex);
|
|
if (analysis_error.empty())
|
|
outcome.emplace(RugnuxLiveState::Done, "");
|
|
else
|
|
outcome.emplace(RugnuxLiveState::Failed, analysis_error);
|
|
}
|
|
}
|
|
if (setup_error)
|
|
std::rethrow_exception(setup_error);
|
|
}
|
|
|
|
void StillsGPUAnalysis::Written() {
|
|
if (!status || !outcome)
|
|
return;
|
|
if (last_image)
|
|
status->Latency(std::chrono::duration<double>(std::chrono::steady_clock::now() - last_image.value()).count());
|
|
status->Finish(outcome->first, outcome->second);
|
|
outcome.reset();
|
|
}
|
|
|
|
void StillsGPUAnalysis::PrePassThread() {
|
|
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;
|
|
const int64_t all_images = experiment.GetImageNum();
|
|
std::optional<StillsPrePassResult> result;
|
|
std::chrono::steady_clock::time_point prepass_start = std::chrono::steady_clock::now();
|
|
// At a low hit rate the first frames may not be enough to fit the geometry on, or may not hold a crystal: with
|
|
// a cell to fit, the pre-pass runs again on all that is held once that has doubled, until the fit has its
|
|
// frames, or the collection ended, or the sample is as large as it may grow (the last try): bounded so that
|
|
// the last run is short and the other half of the queue takes the frames arriving meanwhile.
|
|
const int64_t max_sample = std::min(PREPASS_MAX_IMAGES, max_held / 2);
|
|
for (int64_t needed = std::min(PREPASS_IMAGES, all_images); ;) {
|
|
std::map<int64_t, std::shared_ptr<HeldImage>> sample;
|
|
bool last;
|
|
{
|
|
std::unique_lock ul(prepass_mutex);
|
|
prepass_cv.wait(ul, [&] {
|
|
return prepass_last || static_cast<int64_t>(prepass_images.size()) >= needed;
|
|
});
|
|
// Frames that arrived during the previous run do not make this one larger than the bound.
|
|
for (const auto &[n, image]: prepass_images) {
|
|
if (static_cast<int64_t>(sample.size()) >= std::max(needed, max_sample))
|
|
break;
|
|
sample.emplace(n, image);
|
|
}
|
|
last = prepass_last || static_cast<int64_t>(sample.size()) >= std::min(all_images, max_sample);
|
|
}
|
|
if (sample.empty())
|
|
break;
|
|
{
|
|
std::unique_lock ul(live_mutex);
|
|
live.prepass = RugnuxLivePrePass::Running;
|
|
}
|
|
if (status)
|
|
Publish();
|
|
try {
|
|
const auto t0 = std::chrono::steady_clock::now();
|
|
auto r = RunStillsPrePass(start, sample, static_cast<int>(std::thread::hardware_concurrency()));
|
|
logger.Info("Pre-pass on {} images ({:.1f} s): beam centre ({:.2f}, {:.2f}) -> ({:.2f}, {:.2f}) px, "
|
|
"distance {:.2f} -> {:.2f} mm on {} frames, masked pixels {} -> {}", r.images,
|
|
std::chrono::duration<double>(std::chrono::steady_clock::now() - t0).count(),
|
|
experiment.GetBeamX_pxl(), experiment.GetBeamY_pxl(), r.beam_x_pxl, r.beam_y_pxl,
|
|
experiment.GetDetectorDistance_mm(), r.distance_mm, r.geometry_frames,
|
|
masked(mask), masked(r.mask));
|
|
result = std::move(r);
|
|
} catch (const std::exception &e) {
|
|
logger.Warning("Pre-pass on {} images failed - {}", sample.size(), e.what());
|
|
}
|
|
if (last || !experiment.GetUnitCell()
|
|
|| (result && result->geometry_frames >= RUGNUX_REFINE_GEOMETRY_FRAMES))
|
|
break;
|
|
needed = std::min({2 * static_cast<int64_t>(sample.size()), all_images, max_sample});
|
|
std::unique_lock ul(live_mutex);
|
|
live.prepass = RugnuxLivePrePass::Holding;
|
|
live.images_needed = needed;
|
|
}
|
|
{
|
|
// The frames are analysed from here on: their slots go back to the buffer once they are sent.
|
|
std::unique_lock ul(prepass_mutex);
|
|
prepass_over = true;
|
|
prepass_images.clear();
|
|
}
|
|
const bool failed = !result;
|
|
int geometry_frames = 0;
|
|
if (result) {
|
|
const auto &r = *result;
|
|
setup.experiment.BeamX_pxl(r.beam_x_pxl).BeamY_pxl(r.beam_y_pxl).DetectorDistance_mm(r.distance_mm);
|
|
setup.mask = r.mask;
|
|
geometry_frames = r.geometry_frames;
|
|
const auto &m = r.mask.GetMask();
|
|
std::unique_lock ul(live_mutex);
|
|
live.beam_x_pxl = r.beam_x_pxl;
|
|
live.beam_y_pxl = r.beam_y_pxl;
|
|
live.distance_mm = r.distance_mm;
|
|
live.shadow_pixels = std::count_if(m.begin(), m.end(), [](uint32_t v) {
|
|
return (v & (1u << PixelMask::BeamStopPixelBit)) != 0;
|
|
});
|
|
live.bad_pixels = std::count_if(m.begin(), m.end(), [](uint32_t v) {
|
|
return (v & (1u << PixelMask::HotPixelBit)) != 0;
|
|
});
|
|
} else
|
|
logger.Warning("Pre-pass failed, the images are analysed with the stated beam centre and mask");
|
|
try {
|
|
setup.mapping = std::make_unique<AzimuthalIntegrationMapping>(setup.experiment, setup.mask);
|
|
// Retained for the merge, which takes each image's reflections as soon as it is analysed.
|
|
indexer = std::make_unique<IndexAndRefine>(setup.experiment, indexing_thread_pool, merge_frames, true);
|
|
if (geometry_frames == 0) {
|
|
const std::string text = fmt::format(
|
|
"The pre-pass could not refine the geometry (too few frames indexed in what the buffer holds, or no unit cell given): the "
|
|
"whole run is analysed at the distance in the header, {:.2f} mm, beam centre ({:.2f}, {:.2f}) px",
|
|
setup.experiment.GetDetectorDistance_mm(), setup.experiment.GetBeamX_pxl(),
|
|
setup.experiment.GetBeamY_pxl());
|
|
logger.Warning("{}", text);
|
|
if (status)
|
|
status->Note(text);
|
|
}
|
|
if (merge_frames) {
|
|
merge_experiment = setup.experiment;
|
|
merge_experiment.ImportScalingSettings(RugnuxDefaultScalingSettings(false));
|
|
merge = std::make_unique<StillsLiveMerge>(merge_experiment, experiment.GetImageNum(),
|
|
setup.experiment.GetUnitCell(), std::nullopt,
|
|
setup.experiment.IsDetectIceRings(),
|
|
SpotFindingSettings().ice_ring_width_Q_recipA, logger);
|
|
merge_future = std::async(std::launch::async, &StillsGPUAnalysis::MergeThread, this);
|
|
}
|
|
ready(setup);
|
|
} catch (const std::exception &e) {
|
|
// Nothing can be written without the start message: the frames are dropped (AnalysisThread), and
|
|
// the collection is stopped.
|
|
const auto *j = dynamic_cast<const JFJochException *>(&e);
|
|
const JFJochException error = j ? *j : JFJochException(JFJochExceptionCategory::WrongDAQState,
|
|
std::string("GPU analysis setup failed: ") + e.what());
|
|
logger.ErrorException(error);
|
|
setup_error = j ? std::current_exception() : std::make_exception_ptr(error);
|
|
if (setup_failed)
|
|
setup_failed(error);
|
|
prepass_done.count_down();
|
|
return;
|
|
}
|
|
{
|
|
std::unique_lock ul(live_mutex);
|
|
live.prepass = failed ? RugnuxLivePrePass::Failed : RugnuxLivePrePass::Done;
|
|
live.prepass_s = std::chrono::duration<double>(std::chrono::steady_clock::now() - prepass_start).count();
|
|
ready_time = std::chrono::steady_clock::now();
|
|
}
|
|
if (status) {
|
|
status->Step("GPU analysis");
|
|
Publish();
|
|
}
|
|
prepass_done.count_down();
|
|
}
|
|
|
|
void StillsGPUAnalysis::AnalysisThread() {
|
|
try {
|
|
pin_gpu();
|
|
} catch (const std::exception &e) {
|
|
logger.Warning("Error pinning GPU {}", e.what());
|
|
}
|
|
|
|
prepass_done.wait();
|
|
|
|
if (setup_error) {
|
|
while (queue.GetBlocking()) {}
|
|
return;
|
|
}
|
|
|
|
// A thread whose analysis cannot be built still passes its frames on, unanalysed.
|
|
std::unique_ptr<MXAnalysisWithoutFPGA> analysis;
|
|
try {
|
|
analysis = std::make_unique<MXAnalysisWithoutFPGA>(setup.experiment, *setup.mapping, setup.mask,
|
|
*indexer, /*enable_fused_adaptive_gpu=*/true);
|
|
} catch (const std::exception &e) {
|
|
AnalysisError(std::string("analysis thread could not start, its frames are written unanalysed - ")
|
|
+ e.what());
|
|
}
|
|
|
|
while (auto held = queue.GetBlocking()) {
|
|
DataMessage &message = held->message;
|
|
std::optional<AzimuthalIntegrationProfile> profile;
|
|
const auto start_time = std::chrono::steady_clock::now();
|
|
profile.emplace(*setup.mapping);
|
|
try {
|
|
if (analysis)
|
|
analysis->Analyze(message, *profile, settings());
|
|
} catch (const JFJochException &e) {
|
|
// The frame is still written, without its analysis.
|
|
logger.ErrorException(e);
|
|
} catch (const std::exception &e) {
|
|
AnalysisError(fmt::format("image {} written unanalysed - {}", message.number, e.what()));
|
|
}
|
|
message.processing_time_s = std::chrono::duration<float>(std::chrono::steady_clock::now()
|
|
- start_time).count();
|
|
try {
|
|
sink(message, *profile);
|
|
if (merge)
|
|
merge_queue.Put(message.number);
|
|
if (status) {
|
|
int64_t analysed;
|
|
{
|
|
std::unique_lock ul(live_mutex);
|
|
analysed = ++live.images_analysed;
|
|
if (message.indexing_result.value_or(false))
|
|
++live.images_indexed;
|
|
}
|
|
status->Progress(analysed, experiment.GetImageNum());
|
|
Publish();
|
|
}
|
|
} catch (const std::exception &e) {
|
|
// The thread goes on with the next frame, so that the queue is drained.
|
|
AnalysisError(fmt::format("image {} lost - {}", message.number, e.what()));
|
|
}
|
|
}
|
|
}
|
|
|
|
std::optional<UnitCell> 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();
|
|
}
|
|
|
|
float StillsGPUAnalysis::GetDetectorDistance_mm() const {
|
|
return setup.experiment.GetDetectorDistance_mm();
|
|
}
|
|
|
|
const std::optional<StillsMergeOutput> &StillsGPUAnalysis::GetMerged() const {
|
|
return merged;
|
|
}
|
|
|
|
const DiffractionExperiment &StillsGPUAnalysis::GetMergeExperiment() const {
|
|
return merge_experiment;
|
|
}
|
|
|
|
void StillsGPUAnalysis::MergeThread() {
|
|
constexpr int64_t PREPARE_EVERY = 200; // frames between two Prepare calls after a checkpoint
|
|
const int64_t images = experiment.GetImageNum();
|
|
// The cell merged in, as rugnux merges it: the consensus of the frames indexed by the first merge (the
|
|
// collection's cell only picks the frames that are merged).
|
|
bool cell_set = false;
|
|
const auto have_cell = [&] {
|
|
if (!cell_set)
|
|
if (const auto c = indexer->GetConsensusUnitCell()) {
|
|
merge_experiment.SetUnitCell(*c);
|
|
merge->Cell(*c);
|
|
cell_set = true;
|
|
}
|
|
return cell_set;
|
|
};
|
|
size_t next = 0;
|
|
try {
|
|
const auto add = [&](int64_t n, IntegrationOutcome &&o) {
|
|
merge->Add(n, std::move(o));
|
|
if (next < CHECKPOINTS.size()
|
|
&& merge->Added() >= std::llround(CHECKPOINTS[next] * static_cast<double>(images))) {
|
|
if (have_cell())
|
|
PublishMerge(merge->Checkpoint(MERGE_THREADS),
|
|
fmt::format("{:.0f}% of the images (provisional)", 100.0 * CHECKPOINTS[next]));
|
|
++next;
|
|
} else if (next > 0 && merge->Added() % PREPARE_EVERY == 0)
|
|
merge->Prepare(MERGE_THREADS);
|
|
};
|
|
for (int64_t n = merge_queue.GetBlocking(); n >= 0; n = merge_queue.GetBlocking())
|
|
add(n, indexer->TakeIntegrationOutcome(n));
|
|
if (have_cell()) {
|
|
merged = merge->Final(static_cast<int>(AvailableCpus()));
|
|
logger.Info("Stills merge after the last image: {:.2f} s", merged->seconds);
|
|
}
|
|
} catch (const std::exception &e) {
|
|
logger.Warning("Stills merge failed: {}", e.what());
|
|
merged.reset();
|
|
}
|
|
}
|
|
|
|
void StillsGPUAnalysis::PublishMerge(const StillsMergeOutput &out, const std::string &title) {
|
|
const auto &o = out.merge.statistics.overall;
|
|
const std::string text = fmt::format("Merge of {}: {} crystals, {:.2f} A, CC1/2 {:.3f}, R_split {:.3f}, "
|
|
"completeness {:.1f}%, <I/sigma> {:.1f} ({:.1f} s)", title, out.images,
|
|
o.d_min, o.cc_half, o.r_split,
|
|
o.possible_unique_reflections > 0
|
|
? 100.0 * o.unique_reflections / o.possible_unique_reflections : 0.0,
|
|
o.mean_i_over_sigma, out.seconds);
|
|
logger.Info("{}", text);
|
|
if (status)
|
|
status->Note(text);
|
|
}
|