diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index e04cb0658..a271fe370 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -483,12 +483,12 @@ void Rugnux::CollectSpeculativeProbe() { speculative_probe_.reset(); } -std::vector Rugnux::ExperimentKey(bool rotation_angles) const { - std::vector k{experiment_.GetBeamX_pxl(), experiment_.GetBeamY_pxl(), - experiment_.GetDetectorDistance_mm(), experiment_.GetPoniRot1_rad(), - experiment_.GetPoniRot2_rad(), experiment_.GetWavelength_A(), - static_cast(experiment_.GetMaxSpotCount())}; - if (const auto g = experiment_.GetGoniometer()) { +std::vector Rugnux::ExperimentKey(const DiffractionExperiment &e, bool rotation_angles) const { + std::vector k{e.GetBeamX_pxl(), e.GetBeamY_pxl(), + e.GetDetectorDistance_mm(), e.GetPoniRot1_rad(), + e.GetPoniRot2_rad(), e.GetWavelength_A(), + static_cast(e.GetMaxSpotCount())}; + if (const auto g = e.GetGoniometer()) { const Coord a = g->GetAxis(); const auto helical = g->GetHelicalStep().value_or(Coord(-1, -1, -1)); k.insert(k.end(), {1.0, a.x, a.y, a.z, helical.x, helical.y, helical.z}); @@ -497,11 +497,11 @@ std::vector Rugnux::ExperimentKey(bool rotation_angles) const { g->GetScreeningWedge().value_or(-1.0f)}); } else k.push_back(0.0); - if (const auto uc = experiment_.GetUnitCell()) + if (const auto uc = e.GetUnitCell()) k.insert(k.end(), {1.0, uc->a, uc->b, uc->c, uc->alpha, uc->beta, uc->gamma}); else k.push_back(0.0); - const auto &x = experiment_.GetIndexingSettings(); + const auto &x = e.GetIndexingSettings(); k.insert(k.end(), {x.GetTolerance(), x.GetFFT_MinUnitCell_A(), x.GetFFT_MaxUnitCell_A(), static_cast(x.GetFFT_NumVectors()), x.GetFFT_HighResolution_A(), x.GetFFT_MinAngle_deg(), x.GetFFT_MaxAngle_deg(), x.GetUnitCellDistTolerance(), @@ -514,23 +514,24 @@ std::vector Rugnux::ExperimentKey(bool rotation_angles) const { k.insert(k.end(), {f.signal_to_noise_threshold, static_cast(f.photon_count_threshold), static_cast(f.min_pix_per_spot.value_or(-1)), static_cast(f.max_pix_per_spot), f.high_resolution_limit.value_or(-1.0f)}); - k.insert(k.end(), {static_cast(experiment_.GetSpaceGroupOrP1().number), - static_cast(experiment_.GetImageNum())}); + k.insert(k.end(), {static_cast(e.GetSpaceGroupOrP1().number), + static_cast(e.GetImageNum())}); return k; } -std::vector Rugnux::SpotFindingKey(const AzimuthalIntegrationMapping &mapping) const { +std::vector Rugnux::SpotFindingKey(const DiffractionExperiment &x, + const AzimuthalIntegrationMapping &mapping) const { // Without the rotation angles: a frame's spots are found the same under any of them, and the angle // each spot carries is stamped again where a list is taken from the store (prefetch_spots). The // probes of the rotation-scale walk differ only there, so they find the spots of a frame once. - std::vector k = ExperimentKey(/*rotation_angles=*/false); + std::vector k = ExperimentKey(x, /*rotation_angles=*/false); const auto &f = config_.spot_finding; k.insert(k.end(), {static_cast(f.enable), f.low_resolution_limit.value_or(-1.0f), f.cutoff_spot_count_low_res, f.high_res_gap_Q_recipA.value_or(-1.0f), f.ice_ring_width_Q_recipA, static_cast(f.adaptive_threshold), f.false_pixels_per_frame, static_cast(f.measured_ring_q_recipA.size())}); k.insert(k.end(), f.measured_ring_q_recipA.begin(), f.measured_ring_q_recipA.end()); - k.insert(k.end(), {static_cast(experiment_.IsDetectIceRings()), + k.insert(k.end(), {static_cast(x.IsDetectIceRings()), static_cast(pixel_mask_.GetBinaryMaskChecksum()), static_cast(mapping.GetPixelToBinChecksum()), static_cast(mapping.GetCorrectionsChecksum())}); diff --git a/rugnux/Rugnux.h b/rugnux/Rugnux.h index 97b91a4de..355e1c0b3 100644 --- a/rugnux/Rugnux.h +++ b/rugnux/Rugnux.h @@ -787,7 +787,12 @@ class Rugnux { // not evaluate it again. Shared with the copies the run makes of itself. std::shared_ptr azint_geometry_ = std::make_shared(); - [[nodiscard]] std::vector SpotFindingKey(const AzimuthalIntegrationMapping &mapping) const; + [[nodiscard]] std::vector SpotFindingKey(const AzimuthalIntegrationMapping &mapping) const { + return SpotFindingKey(experiment_, mapping); + } + // The same, for spots found at another experiment (a trial beam centre) against its mapping. + [[nodiscard]] std::vector SpotFindingKey(const DiffractionExperiment &x, + const AzimuthalIntegrationMapping &mapping) const; void AddFirstPassMemo(const FirstPassMemo &memo); // The geometry walk's next indexing probe, started on a copy of the run as soon as the canonical // pass has post-refined - it depends on nothing that pass does afterwards - so that it runs beside @@ -809,7 +814,11 @@ class Rugnux { [[nodiscard]] std::vector FirstPassInputKey() const; // The part of FirstPassInputKey that is the experiment and the spot-finding settings; without the // goniometer's angles where rotation_angles is false (SpotFindingKey). - [[nodiscard]] std::vector ExperimentKey(bool rotation_angles = true) const; + [[nodiscard]] std::vector ExperimentKey(bool rotation_angles = true) const { + return ExperimentKey(experiment_, rotation_angles); + } + // The same, of another experiment beside the run's settings. + [[nodiscard]] std::vector ExperimentKey(const DiffractionExperiment &x, bool rotation_angles = true) const; // The file's detector distance, from before the first pass of the rotation two-pass: a pass whose // distance is still this one asks the post-refinement to test "the header distance is right" // (PostRefineSettings::distance_at_header); a pass that has walked off it does not. diff --git a/rugnux/RugnuxFirstPass.cpp b/rugnux/RugnuxFirstPass.cpp index eb257aa03..599a59494 100644 --- a/rugnux/RugnuxFirstPass.cpp +++ b/rugnux/RugnuxFirstPass.cpp @@ -88,6 +88,105 @@ using namespace rugnux_internal; namespace { + // The spots of one first-pass frame. A frame's spots are a pure function of that frame and the + // settings - the engine carries nothing from one image to the next - so it does not matter which + // engine, or which thread, finds them. + std::vector FindFirstPassSpots(JFJochReader &reader, const SpotFindingSettings &settings, + const JFJochReaderDataset &dataset, + MXAnalysisWithoutFPGA &analysis, + AzimuthalIntegrationProfile &profile, + JFJochReaderRawImage &img, int ordinal, int image_idx, + Logger &logger) { + std::vector spots; + try { + if (reader.ReadRawImage(image_idx, img)) { + DataMessage m{}; + m.number = ordinal; + m.original_number = image_idx; + auto first_pass = settings; + first_pass.indexing = false; + first_pass.quick_integration = false; + m.image = img.image; + if (dataset.efficiency.size() > image_idx) + m.image_collection_efficiency = dataset.efficiency[image_idx]; + analysis.Analyze(m, profile, first_pass); + spots = std::move(m.spots); + } + } catch (const std::exception &e) { + if (IsFatalResourceError(e)) throw; + logger.Warning("First-pass spot read failed for image {}: {}", image_idx, e.what()); + } + return spots; + } + + // First-pass spot finding on workers of its own, all on one card (`device`), for work that runs + // beside the first pass. Each worker keeps its engine from one call to the next - building one + // costs far more than a frame - and gives it back on that card. + class SpotFindingWorkers { + public: + SpotFindingWorkers(JFJochReader &reader, const DiffractionExperiment &x, + const AzimuthalIntegrationMapping &mapping, const PixelMask &mask, + const SpotFindingSettings &settings, const JFJochReaderDataset &dataset, + int start_image, int stride, size_t nworkers, int device) + : reader(reader), x(x), mapping(mapping), mask(mask), settings(settings), dataset(dataset), + start_image(start_image), stride(stride), device(device), engines(std::max(nworkers, 1)) {} + + ~SpotFindingWorkers() { + OnWorkers(engines.size(), [&](size_t t) { engines[t] = Engine{}; }); + } + + // The spots of each ordinal, in the order given. + std::vector> Find(const std::vector &ordinals) { + std::vector> found(ordinals.size()); + std::atomic next{0}; + OnWorkers(std::min(engines.size(), ordinals.size()), [&](size_t t) { + Engine &e = engines[t]; + if (!e.analysis) { + e.indexer = std::make_unique(x, nullptr, /*retain_outcomes=*/false); + e.analysis = std::make_unique(x, mapping, mask, *e.indexer, + /*enable_fused_adaptive_gpu=*/true); + e.profile = std::make_unique(mapping); + } + for (size_t i = next.fetch_add(1); i < ordinals.size(); i = next.fetch_add(1)) + found[i] = FindFirstPassSpots(reader, settings, dataset, *e.analysis, *e.profile, + e.raw_image, ordinals[i], start_image + ordinals[i] * stride, + logger); + }); + return found; + } + + private: + struct Engine { + // Before the analysis engine, which page-locks these bytes for its uploads. + JFJochReaderRawImage raw_image; + std::unique_ptr indexer; + std::unique_ptr analysis; + std::unique_ptr profile; + }; + + // fn(t) on n threads of its own, on the card. + template void OnWorkers(size_t n, Fn fn) { + std::vector> futures; + for (size_t t = 0; t < n; t++) + futures.emplace_back(std::async(std::launch::async, [this, &fn, t] { + pin_gpu(device); + fn(t); + })); + for (auto &f : futures) + f.get(); + } + + JFJochReader &reader; + const DiffractionExperiment &x; + const AzimuthalIntegrationMapping &mapping; + const PixelMask &mask; + const SpotFindingSettings settings; + const JFJochReaderDataset &dataset; + const int start_image, stride, device; + std::vector engines; + Logger logger{"Rugnux"}; + }; + // Field by field, bitwise, for the spot memo's verification. bool SameSpots(const std::vector &a, const std::vector &b) { if (a.size() != b.size()) @@ -248,27 +347,8 @@ bool Rugnux::FirstPassRotationIndexing(PipelineLocals &p) { // not matter which engine, or which thread, finds them. const auto find_spots = [&](MXAnalysisWithoutFPGA &analysis, AzimuthalIntegrationProfile &profile, JFJochReaderRawImage &img, int ordinal) { - const int image_idx = start_image + ordinal * config_.stride; - std::vector spots; - try { - if (reader_.ReadRawImage(image_idx, img)) { - DataMessage m{}; - m.number = ordinal; - m.original_number = image_idx; - auto first_pass = config_.spot_finding; - first_pass.indexing = false; - first_pass.quick_integration = false; - m.image = img.image; - if (dataset->efficiency.size() > image_idx) - m.image_collection_efficiency = dataset->efficiency[image_idx]; - analysis.Analyze(m, profile, first_pass); - spots = std::move(m.spots); - } - } catch (const std::exception &e) { - if (IsFatalResourceError(e)) throw; - logger.Warning("First-pass spot read failed for image {}: {}", image_idx, e.what()); - } - return spots; + return FindFirstPassSpots(reader_, config_.spot_finding, *dataset, analysis, profile, img, + ordinal, start_image + ordinal * config_.stride, logger); }; // Find the spots of every ordinal in the list that is not cached yet, on several workers, and @@ -1200,7 +1280,97 @@ bool Rugnux::FirstPassRotationIndexing(PipelineLocals &p) { std::vector> futures; // last, so it is waited for before the rest goes }; std::optional short_spec; + // The beam-centre check below runs a second first pass at the centre measured from the + // background, and it reads the first one's answer only after that pass is over. So the spots at + // the measured centre are found, and both schemes fed and indexed, here beside the file centre's + // pass - on a copy of the experiment at that centre, with workers, a mapping and a spot list of + // their own - and taken at the check's place only if what they read is still what the run has + // there (the experiment with its rotation angles, SpotFindingKey at the measured centre, the spot + // budget); otherwise that pass is run there as before. The schemes are scored there in either + // case, on the run's own indexer. + // + // Only with a second GPU, whose card its spot finding takes: on one card the two passes share the + // card and the reader, and measured, running them side by side saved nothing. There the check + // runs its pass after the first, as it always did. + struct CentreSpeculation { + std::vector experiment_key; + std::vector key; + size_t spots_per_image = 0; + std::unique_ptr experiment; + std::unique_ptr mask; + std::unique_ptr mapping; + std::map> spots; + SchemeIndexers ris; + std::future done; // last, so it is waited for before the rest goes + }; + std::optional centre_spec; + const auto start_centre_speculation = [&] { + if (get_gpu_count() < 2) + return; + JoinBeamCenterCapture(); + if (cancelled_ || !config_.beam_center_check || !background_center_) + return; + auto &cs = centre_spec.emplace(); + cs.experiment = std::make_unique(experiment_); + cs.experiment->BeamX_pxl(background_center_->beam_x_pxl).BeamY_pxl(background_center_->beam_y_pxl); + cs.mask = std::make_unique(pixel_mask_); + cs.mapping = std::make_unique(*cs.experiment, *cs.mask, *azint_geometry_); + cs.experiment_key = ExperimentKey(*cs.experiment); + cs.key = SpotFindingKey(*cs.experiment, *cs.mapping); + cs.spots_per_image = first_pass_spots_per_image; + { + // What an earlier pass already found at this centre under these settings. + std::lock_guard lock(first_pass_spots_->m); + for (const auto &[k, spots] : first_pass_spots_->entries) + if (k == cs.key) + cs.spots = spots; + } + // At this pass's angles, as prefetch_spots takes them. + const DiffractionGeometry spot_geometry = cs.experiment->GetDiffractionGeometry(); + for (auto &[ordinal, spots] : cs.spots) + StampSpotAngles(spots, spot_geometry); + cs.done = std::async(std::launch::async, [this, &cs, &schemes, &validation, &indexer_pool, &run_indexing, + settings = config_.spot_finding, dataset, start_image, + nworkers = engines.size()] { + SpotFindingWorkers workers(reader_, *cs.experiment, *cs.mapping, *cs.mask, settings, *dataset, + start_image, config_.stride, nworkers, /*device=*/1); + const auto find = [&](const std::vector &ordinals) { + std::vector wanted; + for (const int o : ordinals) + if (!cs.spots.contains(o)) + wanted.push_back(o); + auto found = workers.Find(wanted); + for (size_t i = 0; i < wanted.size(); i++) + cs.spots.emplace(wanted[i], std::move(found[i])); + }; + // As feed_scheme, at the measured centre. + const auto gonio = cs.experiment->GetGoniometer(); + for (const auto &[name, ordinals] : schemes) { + auto ri = std::make_unique(*cs.experiment, *indexer_pool); + ri->MaxSpotsPerImage(cs.spots_per_image); + ri->VerifyMemo(config_.verify_first_pass_memo); + for (size_t i = 0; i < ordinals.size(); i++) { + if (cancelled_ || ri->AccumulationFull()) + break; + if (!cs.spots.contains(ordinals[i])) + find({ordinals.begin() + i, + ordinals.begin() + std::min(ordinals.size(), i + SPOT_PREFETCH_CHUNK)}); + std::optional angle; + if (gonio) + angle = gonio->GetAngle_deg(static_cast(ordinals[i])) + gonio->GetWedge_deg() / 2.0f; + ri->ProcessImage(ordinals[i], cs.spots.at(ordinals[i]), angle); + } + cs.ris.push_back(std::move(ri)); + } + find(validation); + std::vector rp; + for (const auto &ri : cs.ris) + rp.push_back(ri.get()); + run_indexing(rp); + }); + }; FirstPass best = pick_best(*indexer_pool, *indexer, [&] { + start_centre_speculation(); if (cancelled_ || experiment_.GetUnitCell().has_value()) return; auto &sp = short_spec.emplace(); @@ -1322,7 +1492,30 @@ bool Rugnux::FirstPassRotationIndexing(PipelineLocals &p) { { const int majority = static_cast(validation.size()) / 2; try_beam_center(measured_x, measured_y); - FirstPass alt = pick_best(*indexer_pool, *indexer); + FirstPass alt; + if (centre_spec) + centre_spec->done.get(); + if (centre_spec && centre_spec->experiment_key == ExperimentKey() + && centre_spec->key == SpotFindingKey(*spot_mapping) + && centre_spec->spots_per_image == first_pass_spots_per_image) { + // What pick_best would have found and fed here: hand the spots to the run's cache and + // its memo as prefetch_spots does, and score the schemes on the run's own indexer. + spot_cache = centre_spec->spots; + { + std::lock_guard lock(first_pass_spots_->m); + auto &entries = first_pass_spots_->entries; + if (std::none_of(entries.begin(), entries.end(), + [&](const auto &e) { return e.first == centre_spec->key; })) { + constexpr size_t MAX_SPOT_KEYS = 8; + if (entries.size() == MAX_SPOT_KEYS) + entries.erase(entries.begin()); + entries.emplace_back(centre_spec->key, centre_spec->spots); + } + } + alt = score_schemes(centre_spec->ris, *indexer); + } else + alt = pick_best(*indexer_pool, *indexer); + centre_spec.reset(); const bool header_indexes = best.result.has_value() && best.score > majority; bool measured_indexes = alt.result.has_value() && alt.score > majority; if (!header_indexes) {