diff --git a/broker/JFJochServices.cpp b/broker/JFJochServices.cpp index 31dfaee5f..b38d4e5dc 100644 --- a/broker/JFJochServices.cpp +++ b/broker/JFJochServices.cpp @@ -4,9 +4,12 @@ #include "JFJochServices.h" #include "../common/JFJochException.h" +#include "../image_analysis/structure_refinement/ModelValidation.h" #include "../detector_control/SLSDetectorWrapper.h" #include "../detector_control/DectrisDetectorWrapper.h" +#include + JFJochServices::JFJochServices(Logger &in_logger) : logger(in_logger) {} void JFJochServices::Preflight(const DiffractionExperiment& experiment) { @@ -392,6 +395,8 @@ void JFJochServices::SetRugnuxModel(std::shared_ptr model, in s.model = std::move(model); s.model_size_bytes = s.model ? size_bytes : 0; SetRugnuxLive(s); + // Before the collection it is for starts, so the run's validation is not the process's first. + WarmUpModelValidation(s.model, std::thread::hardware_concurrency(), logger); } LiveCollection *JFJochServices::GetLiveCollection() { diff --git a/image_analysis/geom_refinement/BeamCenterFFT.cpp b/image_analysis/geom_refinement/BeamCenterFFT.cpp index ece2d6fb1..53c4468ac 100644 --- a/image_analysis/geom_refinement/BeamCenterFFT.cpp +++ b/image_analysis/geom_refinement/BeamCenterFFT.cpp @@ -4,7 +4,9 @@ #include "BeamCenterFFT.h" #include +#include #include +#include #include #include @@ -30,26 +32,58 @@ struct Prepared { double variance = 0.0; // of a over the valid pixels, after the whole preparation }; +// Order-preserving: unsigned comparison of the keys orders the floats as < does (-0 below +0). +uint32_t SortKey(float v) { + const uint32_t u = std::bit_cast(v); + return (u & 0x80000000u) ? ~u : (u | 0x80000000u); +} + +// For each k, the value nth_element would leave at position k of the finite values of x - without +// copying or reordering the image. Counting the finite values by the upper 16 bits of their key +// finds the bucket holding the k-th; counting that bucket by the lower 16 bits finds its key, +// which is the value. Two passes over x for all k together. +std::vector KthFinite(const std::vector &x, std::vector k) { + std::vector upper(1 << 16, 0); + for (const float v : x) + if (std::isfinite(v)) + upper[SortKey(v) >> 16]++; + std::vector bucket(k.size(), 0); + for (size_t i = 0; i < k.size(); i++) + while (k[i] >= upper[bucket[i]]) + k[i] -= upper[bucket[i]++]; + + std::vector> lower(k.size(), std::vector(1 << 16, 0)); + for (const float v : x) + if (std::isfinite(v)) + for (size_t i = 0; i < k.size(); i++) + if (SortKey(v) >> 16 == bucket[i]) + lower[i][SortKey(v) & 0xFFFF]++; + std::vector out; + for (size_t i = 0; i < k.size(); i++) { + uint32_t b = 0; + while (k[i] >= lower[i][b]) + k[i] -= lower[i][b++]; + const uint32_t key = (bucket[i] << 16) | b; + out.push_back(std::bit_cast((key & 0x80000000u) ? (key & 0x7FFFFFFFu) : ~key)); + } + return out; +} + Prepared PrepareImage(const std::vector &mean) { Prepared out; const size_t n = mean.size(); out.a.assign(n, 0.0f); out.m.assign(n, 0.0f); - std::vector valid; - valid.reserve(n); - for (float v : mean) - if (std::isfinite(v)) - valid.push_back(v); - if (valid.size() < 2) + const size_t n_valid = std::count_if(mean.begin(), mean.end(), [](float v) { return std::isfinite(v); }); + if (n_valid < 2) return out; - auto kth = [&valid](size_t k) { - std::nth_element(valid.begin(), valid.begin() + k, valid.end()); - return valid[k]; - }; - const float clip_hi = kth(static_cast(0.999 * (valid.size() - 1))); - const float median = kth(valid.size() / 2); + // Where these fall among equal values a -0 can come out for a +0, which the clipping below + // turns into the same image. + const std::vector q = KthFinite(mean, {static_cast(0.999 * (n_valid - 1)), n_valid / 2}); + const float clip_hi = q[0]; + const float median = q[1]; double sum = 0.0; size_t count = 0; diff --git a/image_analysis/structure_refinement/ModelValidation.cpp b/image_analysis/structure_refinement/ModelValidation.cpp index 0ad7cd61a..db415c610 100644 --- a/image_analysis/structure_refinement/ModelValidation.cpp +++ b/image_analysis/structure_refinement/ModelValidation.cpp @@ -10,6 +10,8 @@ #include #include #include +#include +#include #include #include #include @@ -1937,6 +1939,19 @@ std::vector OtherCardsTotalMemory(int device) { return out; } +namespace { + +// The warm-up last started (WarmUpModelValidation); invalid before the first. +std::mutex warm_up_m; +std::shared_future warm_up; + +std::shared_future LastWarmUp() { + std::lock_guard lock(warm_up_m); + return warm_up; +} + +} // namespace + ModelValidationResult ValidateAgainstModel(const std::vector &merged, const UnitCell &cell, const RugnuxModel &model, @@ -1949,6 +1964,8 @@ ModelValidationResult ValidateAgainstModel(const std::vector & const ModelValidationSchedule &schedule, const gemmi::Op *twin_law, const gemmi::Structure *placed) { + if (const auto w = LastWarmUp(); w.valid()) + w.wait(); #ifdef JFJOCH_USE_CUDA // A CUDA failure ends the validation, not the run, and nothing is moved to the CPU part way through: // the validation reports why it did not finish, and a validation that did not finish decides nothing, @@ -1972,6 +1989,47 @@ ModelValidationResult ValidateAgainstModel(const std::vector & #endif } +void WarmUpModelValidation(std::shared_ptr model, size_t nthreads, Logger logger) { +#ifdef JFJOCH_USE_CUDA + if (!model || get_gpu_count() == 0) + return; + std::lock_guard lock(warm_up_m); + std::shared_future previous = warm_up; + warm_up = std::async(std::launch::async, [model, nthreads, logger, previous]() mutable { + if (previous.valid()) + previous.wait(); + const auto start = std::chrono::steady_clock::now(); + // Held for as long as it runs: work beside a run waits for it where the card is too small for both. + GPUWorkBeside beside; + Logger quiet = Logger::Buffered(); + const double d_min = model->structure.resolution > 0.0 ? model->structure.resolution : 2.0; + std::vector fake = ModelReferenceIntensities(*model, std::nullopt, nullptr, d_min, quiet); + for (size_t i = 0; i < fake.size(); i++) { + fake[i].F = std::sqrt(fake[i].I); + fake[i].rfree_flag = i % 20 == 0; + } + const gemmi::UnitCell &c = model->structure.cell; + const UnitCell cell{static_cast(c.a), static_cast(c.b), static_cast(c.c), + static_cast(c.alpha), static_cast(c.beta), static_cast(c.gamma)}; + std::string failure; + try { + failure = Validate(fake, cell, *model, quiet, nullptr, true, nthreads, 0.0, {}, {}, nullptr, true, + nullptr).failure_reason; + } catch (const std::exception &e) { + cuda_clear_error(); + failure = e.what(); + } + logger.Info("Model validation warm-up on the GPU: {:.2f} s{}", + std::chrono::duration(std::chrono::steady_clock::now() - start).count(), + failure.empty() ? "" : " (did not finish: " + failure + ")"); + }).share(); +#else + (void) model; + (void) nthreads; + (void) logger; +#endif +} + namespace { // The reindexing operator as it reads on Miller indices ("k,h,-l" rather than "y,x,-z"). diff --git a/image_analysis/structure_refinement/ModelValidation.h b/image_analysis/structure_refinement/ModelValidation.h index b9dba2fe3..0eebb653b 100644 --- a/image_analysis/structure_refinement/ModelValidation.h +++ b/image_analysis/structure_refinement/ModelValidation.h @@ -408,3 +408,12 @@ std::vector ModelReferenceIntensities(const RugnuxModel &model const gemmi::SpaceGroup *space_group, double d_min, Logger &logger); + +// Run one validation of `model` against its own calculated intensities (to the model's stated +// resolution, 2 A where it states none) on a background thread, and throw it away. A process's first +// validation pays for what the GPU loads on first use - kernels, cuFFT plans and their sizes, on every +// card the null's replicates are spread over - and this pays it ahead, where nobody waits: the broker +// calls it when a model is uploaded, before the collection starts. ValidateAgainstModel waits for a +// warm-up still running, so a real validation never shares the cards with one and computes exactly +// what it would have without it. Its time goes to `logger` when it ends. No-op without a GPU. +void WarmUpModelValidation(std::shared_ptr model, size_t nthreads, Logger logger); diff --git a/receiver/JFJochReceiverRugnux.cpp b/receiver/JFJochReceiverRugnux.cpp index 17b707d42..98ffeb60f 100644 --- a/receiver/JFJochReceiverRugnux.cpp +++ b/receiver/JFJochReceiverRugnux.cpp @@ -351,6 +351,9 @@ void JFJochReceiverRugnux::ProbeThread() { setup.config.end_image = static_cast(frames); setup.config.rotation_postrefine_geometry = false; setup.config.run_scaling = false; + // The file's centre only: the check captures the centre on the GPUs and runs a second first pass, + // both beside the main run's pre-scan, for a centre the probe does not report. + setup.config.beam_center_check = false; ProcessResult r; { Rugnux probe_run(reader, setup.experiment, *dataset->pixel_mask, setup.config); diff --git a/rugnux/RugnuxFirstPass.cpp b/rugnux/RugnuxFirstPass.cpp index e6a2d7519..8228bc826 100644 --- a/rugnux/RugnuxFirstPass.cpp +++ b/rugnux/RugnuxFirstPass.cpp @@ -1099,13 +1099,15 @@ auto Rugnux::FirstPassRun::ScoreSchemes(const SchemeIndexers &ris, IndexAndRefin } // Feed the schemes, index them concurrently, and keep the best (ScoreSchemes). -// alongside: started once the schemes are fed and before they are indexed, to run beside them. +// alongside: started once the schemes are fed and their indexing has begun, to run beside it - it +// may wait (StartCentreSpeculation waits for the beam-centre capture), and the indexing goes on. auto Rugnux::FirstPassRun::PickBest(IndexerThreadPool &pool, IndexAndRefine &idx, const std::function &alongside) -> FirstPass { const SchemeIndexers ris = FeedSchemes(pool); + auto indexing = std::async(std::launch::async, [&ris] { RunIndexing(ris); }); if (alongside) alongside(); - RunIndexing(ris); + indexing.get(); return ScoreSchemes(ris, idx); } @@ -1558,9 +1560,10 @@ void Rugnux::FirstPassRun::StartShortAxisSpeculation() { } auto Rugnux::FirstPassRun::IndexSchemes() -> FirstPass { + // The short axis first: the measured centre's pass waits for the beam-centre capture. return PickBest(*indexer_pool, *indexer, [&] { - StartCentreSpeculation(); StartShortAxisSpeculation(); + StartCentreSpeculation(); }); } diff --git a/tests/ModelValidationTest.cpp b/tests/ModelValidationTest.cpp index 58b8a4ba9..ddea1bd65 100644 --- a/tests/ModelValidationTest.cpp +++ b/tests/ModelValidationTest.cpp @@ -863,6 +863,41 @@ TEST_CASE("ModelValidation_CCModelFollowsTheSignalByShell", "[ModelValidation]") std::filesystem::remove(path); } +// A warm-up (WarmUpModelValidation) only loads what the GPU loads on first use: the validation after it, +// which waits for it, gives what it gives without one. +TEST_CASE("ModelValidation_WarmUpChangesNothing", "[ModelValidation]") { + Logger logger("ModelValidation_WarmUpChangesNothing"); + + const auto path = WriteTemp("warm_up_test.pdb", ClusterPdb().c_str()); + auto obs = ModelReferenceIntensities(ReadModel(path), {}, {}, 2.5, logger); + REQUIRE(obs.size() > 1000); + std::mt19937 rng(20261019); + std::normal_distribution normal(0.0, 1.0); + for (size_t i = 0; i < obs.size(); i++) { + obs[i].I = static_cast(obs[i].I * (1.0 + 0.1 * normal(rng))); + obs[i].F = std::sqrt(std::max(0.0f, obs[i].I)); + obs[i].rfree_flag = (i % 20) == 0; + } + const UnitCell cell{.a = 30, .b = 34, .c = 38, .alpha = 90, .beta = 90, .gamma = 90}; + const auto model = std::make_shared(ReadModel(path)); + const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name("P 21 21 21"); + + const auto cold = ValidateAgainstModel(obs, cell, *model, logger, sg, false, 4); + WarmUpModelValidation(model, 4, logger); + const auto warm = ValidateAgainstModel(obs, cell, *model, logger, sg, false, 4); + + REQUIRE(cold.ok); + REQUIRE(warm.ok); + CHECK(warm.r_work == cold.r_work); + CHECK(warm.r_free == cold.r_free); + CHECK(warm.rigid_body_angle_deg == cold.rigid_body_angle_deg); + CHECK(warm.null_r_work_mean == cold.null_r_work_mean); + CHECK(warm.null_r_work_sd == cold.null_r_work_sd); + CHECK(warm.r_work_sigma == cold.r_work_sigma); + + std::filesystem::remove(path); +} + // The model path's structure factors come from FFTW; they must be gemmi's own transform to float // precision, in the same layout, so prepare_asu_data() reads the same reflections from either.