diff --git a/common/CUDAWrapper.cpp b/common/CUDAWrapper.cpp index 19838a246..dc4d31102 100644 --- a/common/CUDAWrapper.cpp +++ b/common/CUDAWrapper.cpp @@ -61,6 +61,10 @@ std::vector get_gpu_names() { return {}; } +size_t get_gpu_total_memory(int32_t) { + return 0; +} + void set_gpu(int32_t dev_id) {} void pin_gpu() {} diff --git a/common/CUDAWrapper.cu b/common/CUDAWrapper.cu index f150cf101..434b7b5ee 100644 --- a/common/CUDAWrapper.cu +++ b/common/CUDAWrapper.cu @@ -49,6 +49,13 @@ std::vector get_gpu_names() { return names; } +size_t get_gpu_total_memory(int32_t dev_id) { + cudaDeviceProp prop{}; + if (cudaGetDeviceProperties(&prop, dev_id) != cudaSuccess) + return 0; + return prop.totalGlobalMem; +} + void set_gpu(int32_t dev_id) { auto dev_count = get_gpu_count(); diff --git a/common/CUDAWrapper.h b/common/CUDAWrapper.h index b7405aca6..83e02bca8 100644 --- a/common/CUDAWrapper.h +++ b/common/CUDAWrapper.h @@ -4,6 +4,7 @@ #pragma once #include +#include #include #include #include @@ -17,6 +18,11 @@ std::vector get_gpu_names(); // The same list collapsed for a person: "4x NVIDIA A100-SXM4-80GB", or several such groups separated // by ", " on a mixed machine. Empty when no GPU is visible. std::string get_gpu_description(); + +// The device's total memory in bytes - not what is free at the moment. 0 without CUDA, or when the +// device cannot be queried. +size_t get_gpu_total_memory(int32_t dev_id); + void set_gpu(int32_t dev_id); // Pin the calling thread to the next GPU in round-robin order, using a process-wide counter diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 144e85d5a..dcec6a2ba 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -706,6 +706,25 @@ void RotationScaleMerge::StageObservations() { } #endif +std::unique_ptr RotationScaleMerge::CloneIngest(int device, Logger &log) const { + TimingMark mark(*logger, "rsm ingest copy"); + std::unique_ptr c(new RotationScaleMerge(*this)); + c->logger = &log; + c->gpu_device = device; + c->write_back_per_frame_scale = false; + // Of the ingest arrays, IngestDevice reads only this one (StageObservations); the rest served + // IngestHost alone. + if (ingest_partiality) { + c->ingest_partiality.reset(new float[n_partials_obs]); + std::copy_n(ingest_partiality.get(), n_partials_obs, c->ingest_partiality.get()); + } +#ifdef JFJOCH_USE_CUDA + if (gpu_) + c->gpu_ = std::make_unique(device); +#endif + return c; +} + void RotationScaleMerge::IngestDevice() { const ScopedGpuPin pin(gpu_device); // see IngestHost TimingMark mark(*logger, "rsm ingest (device)"); @@ -6211,9 +6230,8 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, bool full_st // returned is the one run last, so that every per-pass state the merge leaves behind is that arm's. // // Where the prior is to scale the partials and that merge resolves its error model, the other arm - // cannot change the choice, so it is not made. The partials' arm then starts from G = 1, which is - // what the fulls-alone arm would have left it (it is the only state that arm leaves for the next), - // so the merge is the one the pair would have returned. + // cannot change the choice, so it is not made. Every arm starts from G = 1 (RunArm), so the merge + // is the one the pair would have returned. if (rocking_event_frames_at_start < 0) rocking_event_frames_at_start = RockingEventFrames(&rocking_events_at_start); const double events_per_frame = n_frames > 0 ? static_cast(rocking_events_at_start) / n_frames : 0.0; @@ -6221,7 +6239,6 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, bool full_st if (for_search && !scale_partials_choice) return RunArm(prior, for_search, full_stats, measure_cc_before_corrections); if (!scale_partials_choice && prior) { - std::fill(g_partial.begin(), g_partial.end(), 1.0); Result with = RunArm(true, for_search, full_stats, measure_cc_before_corrections); if (ErrorModelResolved(with)) { scale_partials_choice = true; @@ -6270,6 +6287,12 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool }); } + // ... and from G = 1, so that no merge starts from another's scales. A frame this merge does not + // fit would otherwise keep the G an earlier merge fitted it with, and the fulls' CC1/2 weight + // reads it (frame_scale_fulls): the merge would depend on which merges ran before it on this + // engine, and two merges could not run side by side on two engines with the same result. + std::fill(g_partial.begin(), g_partial.end(), 1.0); + TimingMark section_mark(*logger, "rsm scale partials and filter frames"); // --- 1. Per-frame partial scaling (Rotation model, per-image G only). --- const int n_groups = ComputeAsuGroups(keygen); // one ASU grouping, shared by partials and fulls @@ -6325,7 +6348,6 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool } #endif if (!scale_partials) { - std::fill(g_partial.begin(), g_partial.end(), 1.0); frame_scaled_scratch.assign(n_frames, 1); partial_loop.converged = true; } else if (!scaled_on_gpu) { @@ -6462,6 +6484,9 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool ReducePartialGroupMeans(n_groups, partial_mean); ComputePerFrameCC(partial_mean, cc, cc_n); } + last_cc = cc; + last_cc_n = cc_n; + last_frame_scaled = partial_scaled; if (write_back_per_frame_scale) FinalizePerFrameScale(cc, cc_n, partial_scaled); diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index b23610398..ae68d4560 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -151,6 +151,27 @@ public: // was written. void SetWriteBackPerFrameScale(bool on) { write_back_per_frame_scale = on; } + // Write back what the last Run() would have written with the write-back on: for a merge that ran + // beside others with it off, at the point where it would have run in series. + void WriteBackPerFrameScale() { FinalizePerFrameScale(last_cc, last_cc_n, last_frame_scaled); } + + // A second engine on this one's ingest, on `device` and logging to `log`, for a merge to run beside + // this engine's: the host half of the ingest is copied, and the copy's IngestDevice() is still to + // be called. Call between IngestHost() and IngestDevice(). No merge depends on the merges run + // before it (see RunArm), so a merge on the copy gives what the same merge gives on this engine. + // The copy does not write its per-frame scales back (WriteBackPerFrameScale does). + [[nodiscard]] std::unique_ptr CloneIngest(int device, Logger &log) const; + + // The partial observations the ingest kept. + [[nodiscard]] int64_t ObservationCount() const { return n_partials_obs; } + + // The card memory an engine may hold during a search merge, per observation the ingest kept: what + // a caller goes by when it decides whether a second engine fits beside one already on the card. + // Between what two engines merging side by side were measured at, ~150 bytes per observation each + // (20.6M observations, 4.2M fulls), and the layout's worst case of ~280 (71 bytes per observation, + // 83 per full, ~130 per ASU group, with a full and a group for every observation). + static constexpr size_t DEVICE_BYTES_PER_OBSERVATION = 256; + // Merge only the first n frames of the outcomes (by default all of them). Set before Ingest(). void SetFrameLimit(int n) { frame_limit = n; } @@ -180,6 +201,17 @@ public: [[nodiscard]] double GetSearchMinZeta() const { return search_min_zeta; } private: + // A unique_ptr that a copy of the engine starts without (see CloneIngest). + template struct NotCopied : std::unique_ptr { + using std::unique_ptr::unique_ptr; + using std::unique_ptr::operator=; + NotCopied() = default; + NotCopied(NotCopied &&) noexcept = default; + NotCopied &operator=(NotCopied &&) noexcept = default; + NotCopied(const NotCopied &) {} + }; + RotationScaleMerge(const RotationScaleMerge &) = default; // CloneIngest only + static constexpr uint8_t OBS_CLIPPED = 1; static constexpr uint8_t OBS_OVERLOADED = 2; @@ -296,11 +328,11 @@ private: // BuildInRangeObservations and handed back at the end of Ingest. Plain arrays, not sized vectors, for // the reason the staging buffers are (see Ingest): the build writes every element, and a sized vector // would first zero all of them on one thread. - std::unique_ptr ingest_delta_phi, ingest_partiality, ingest_image_number; + NotCopied ingest_delta_phi, ingest_partiality, ingest_image_number; // The rocking-event usability test (RockingEventFrames) on the ingest-time values - which is what // that walk reads on the resident path, where the host corr is never refreshed (see // partials_released). Resident path only. - std::unique_ptr ingest_rock_ok; + NotCopied ingest_rock_ok; // Which source reflections the resolution range kept (indexed like the sort keys: frame by frame, // from ingest_src_start), so that the i-th kept one of a frame can be found again. std::vector ingest_keep; @@ -329,7 +361,7 @@ private: std::vector frame_start, frame_count; // CSR ranges of `partials` per frame std::vector frame_cell_ok; // per-frame cell-consistency mask (1 = kept) std::vector finite_ok; // per-obs AcceptReflection finiteness (immutable; 1 = kept) - std::vector g_partial; // per-frame partial scale G (the RESIDUAL after the flux) + std::vector g_partial; // per-frame partial scale G (the RESIDUAL after the flux); 1 at the start of every merge std::vector frame_flux; // per-frame incident flux, run median = 1 (see the .cpp) // corr as Ingest built it, kept for the whole life of the object and put back at the start of every // Run. A pass that starts anywhere else is not the pass that would have run on its own, and the @@ -509,7 +541,7 @@ private: // GPU engine: the whole hot path (scaling, combine, scale-fulls, per-frame CC, smooth-G, merge + // error model) runs on the device, resident, when a GPU is present. Null / inactive otherwise, with // the CPU loops as the bit-parity fallback. Built in Ingest. - std::unique_ptr gpu_; + NotCopied gpu_; bool gpu_active_ = false; // The fulls of this pass live on the device (the GPU combine ran), so the correction surfaces can // build their terms there. Set by RunArm. @@ -519,6 +551,10 @@ private: std::vector surface_corr; #endif bool write_back_per_frame_scale = true; // see SetWriteBackPerFrameScale + // What the last Run wrote back, or would have (WriteBackPerFrameScale). + std::vector last_cc; + std::vector last_cc_n; + std::vector last_frame_scaled; int frame_limit = -1; // see SetFrameLimit int gpu_device = 0; // see SetGpuDevice bool export_scaled_fulls = false; // see SetExportScaledFulls @@ -557,7 +593,7 @@ private: // which ComputeAsuGroups builds for the GPU reduction anyway. ReducePartialGroupMeans sums each // group over it, on all threads and in the order the serial loop would. std::vector partial_group_start, partial_group_count; - std::unique_ptr partial_group_perm; + NotCopied partial_group_perm; // The CPU partial scaling loop's working copy of the partials. It passes over them several times // per iteration and reads a handful of fields each time, so it works on these rather than on the diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index 4dde9ed34..6067ad494 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -958,6 +958,15 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { 100.0 * *bkg_starved_fraction_); return true; } + // The all-observation arm of the space-group search (below) is the same ingest merged with the + // Lorentz filter off. No merge depends on the merges made before it (RotationScaleMerge::RunArm), + // so it is merged on a copy of the engine, started here and read where the search needs it: beside + // the P1 merge, the re-scaled search merge and the first arm. Started whenever the search will run; + // a search that turns out not to need it (no point group from the first arm) does not read it. + // Declared before the future, which waits for the merge when it goes, so the engine outlives it. + std::unique_ptr rsm_all; + Logger all_arm_log = Logger::Buffered(); + std::future all_arm_ahead; if (is_rotation) { if (experiment_.GetRefineRotationWedgeInScaling() || rot_ss.GetRotationWedgeForScaling().has_value()) @@ -985,6 +994,32 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { config_.nthreads, logger, config_.observation_dump_path); rsm->IngestHost(); } + // On a card of its own from three cards up (card 0 has this engine, card 1 the quality guard's + // and the speculative passes). Otherwise on the last card, beside the engine already there - + // this one, or the guard's, which merges at most as many observations - and only where two + // engines of this size fit the card's memory, judged on its total, not on what happens to be + // free: the answer must not depend on the moment. On one card that is the small and middling + // data sets; the large ones merge their arms one after the other. + const int32_t cards = get_gpu_count(); + if (!geometry_prepass && !experiment_.GetGemmiSpaceGroup().has_value() && rsm->GetSearchMinZeta() > 0.0 + && cards > 0) { + const int32_t device = std::min(cards - 1, 2); + const size_t engine_bytes = static_cast(rsm->ObservationCount()) + * RotationScaleMerge::DEVICE_BYTES_PER_OBSERVATION; + if (cards >= 3 || 2 * engine_bytes <= get_gpu_total_memory(device)) { + rsm_all = rsm->CloneIngest(device, all_arm_log); + rsm_all->SetSearchMinZeta(0.0); + rsm_all->SetExportScaledFulls(false); // as scale_and_merge sets them for a search merge + rsm_all->SetFrenchWilson(false); + all_arm_ahead = std::async(std::launch::async, [&rsm_all, &all_arm_log, device] { + pin_gpu(device); // the card and its NUMA node, as the guard's merge has them + TimingMark mark(all_arm_log, "merge P1, all observations"); + rsm_all->IngestDevice(); + return rsm_all->Run(/*for_search=*/true, /*full_stats=*/true, + /*measure_cc_before_corrections=*/false); + }); + } + } rsm->IngestDevice(); } @@ -1001,6 +1036,28 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { indexer->GetIntegrationOutcome(), config_.nthreads, /*release_reflections=*/true, logger); + // What a rotation merge leaves in the result, taken in the order the merges are made - which, for + // the all-observation arm merged beside the others, is where the search reads it. + const auto take_merge = [&](const std::string &label, bool for_search, RotationScaleMerge::Result r) + -> ScaleMergeResult { + result.error_model_isa = r.isa; + result.error_model_isa_asymptotic = r.isa_asymptotic; + result.error_model_isa_resolved = r.isa_resolved; + if (r.resolution_fit_A) result.resolution_fit_A = r.resolution_fit_A; + result.error_model_a = r.error_model_a; + result.error_model_b = r.error_model_b; + result.scaling_iterations_partials = r.scaling_iterations_partials; + result.scaling_iterations_fulls = r.scaling_iterations_fulls; + result.scaling_converged = r.scaling_converged; + if (for_search) + result.search_scaling_converged = r.scaling_converged; + Note(fmt::format("merge in {}: {} unique", label, r.merged.size())); + return ScaleMergeResult{std::move(r.merged), std::move(r.statistics), + r.cc_half_before_corrections, + std::move(r.reference_statistics), r.reference_isa, + std::move(r.scaled_fulls)}; + }; + // measure_cc: whether anything reads the result's cc_half_before_corrections, which costs a // merge of its own - only the first merge, which the two-pass quality guard reads. auto scale_and_merge = [&](const std::string &label, bool for_search, bool measure_cc = false) @@ -1017,24 +1074,9 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // Amplitudes are made once, on the merge that is written, where the anisotropy tensor // their prior uses is known (see below) - not on every merge the tail makes on its way. rsm->SetFrenchWilson(false); - auto r = rsm->Run(for_search, /*full_stats=*/!geometry_prepass, - /*measure_cc_before_corrections=*/measure_cc); - result.error_model_isa = r.isa; - result.error_model_isa_asymptotic = r.isa_asymptotic; - result.error_model_isa_resolved = r.isa_resolved; - if (r.resolution_fit_A) result.resolution_fit_A = r.resolution_fit_A; - result.error_model_a = r.error_model_a; - result.error_model_b = r.error_model_b; - result.scaling_iterations_partials = r.scaling_iterations_partials; - result.scaling_iterations_fulls = r.scaling_iterations_fulls; - result.scaling_converged = r.scaling_converged; - if (for_search) - result.search_scaling_converged = r.scaling_converged; - Note(fmt::format("merge in {}: {} unique", label, r.merged.size())); - return ScaleMergeResult{std::move(r.merged), std::move(r.statistics), - r.cc_half_before_corrections, - std::move(r.reference_statistics), r.reference_isa, - std::move(r.scaled_fulls)}; + return take_merge(label, for_search, + rsm->Run(for_search, /*full_stats=*/!geometry_prepass, + /*measure_cc_before_corrections=*/measure_cc)); } // Stills (rotation goes through RotationScaleMerge above): self-scale each image against the // running merge with ScaleOnTheFly (fixed partiality), then merge directly. This runs even @@ -1347,6 +1389,30 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // Set per merge, not once: the two arms below are scaled separately and their ISa differ by up // to a factor of six. 0 (b not measurable) leaves the cut as the caller set it. sg_opts.merge_isa = result.error_model_isa; + // The all-observation merge made beside everything since the ingest (see rsm_all), and its + // search started beside this one's: a search reads its merge and its options and nothing + // else. Whether the all-observation arm is wanted at all is known only from the first arm + // (below), so nothing of it reaches the run before then. A merge that failed beside the + // others is made again where it is wanted, in series, as it would have been without it. + std::optional all_merge; + std::future all_search_ahead; + if (all_arm_ahead.valid()) { + TimingMark wait_mark(logger, "wait for the all-observation merge"); + try { + all_merge = all_arm_ahead.get(); + } catch (const std::exception &e) { + logger.Warning("Space-group search: the all-observation merge beside the others failed ({}); " + "it is made again in series", e.what()); + } + wait_mark.Stop(); + } + if (all_merge) { + SearchSpaceGroupOptions all_opts = sg_opts; + all_opts.merge_isa = all_merge->isa; // as take_merge leaves it below + all_search_ahead = std::async(std::launch::async, [&merged = all_merge->merged, all_opts] { + return SearchSpaceGroup(merged, all_opts); + }); + } TimingMark first_arm_mark(logger, "SearchSpaceGroup (first arm)"); auto sg_search = SearchSpaceGroup(sm.merged, sg_opts); first_arm_mark.Stop(); @@ -1380,15 +1446,28 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // all-observation arm on both twin gates and promoted to R32 by the filtered one. So the // losing arm's refusal is logged and carried to the report below; the rule is unchanged. if (rsm && rsm->GetSearchMinZeta() > 0.0 && !sg_search.point_group_hm.empty()) { - const double zeta = rsm->GetSearchMinZeta(); - rsm->SetSearchMinZeta(0.0); - auto sm_all = scale_and_merge("P1, all observations", true); - rsm->SetSearchMinZeta(zeta); + ScaleMergeResult sm_all; + SearchSpaceGroupResult alt; merged_filtered_isa = sg_opts.merge_isa; // the filtered arm's, before it is replaced - sg_opts.merge_isa = result.error_model_isa; // this arm's own error model - TimingMark all_arm_mark(logger, "SearchSpaceGroup (all-observation arm)"); - auto alt = SearchSpaceGroup(sm_all.merged, sg_opts); - all_arm_mark.Stop(); + if (all_merge) { + TimingMark all_arm_mark(logger, "SearchSpaceGroup (all-observation arm, wait)"); + alt = all_search_ahead.get(); + all_arm_mark.Stop(); + phase("Scale/combine/merge (P1, all observations)"); + all_arm_log.ReplayInto(logger); + rsm_all->WriteBackPerFrameScale(); + sm_all = take_merge("P1, all observations", true, std::move(*all_merge)); + sg_opts.merge_isa = result.error_model_isa; // this arm's own error model + } else { + const double zeta = rsm->GetSearchMinZeta(); + rsm->SetSearchMinZeta(0.0); + sm_all = scale_and_merge("P1, all observations", true); + rsm->SetSearchMinZeta(zeta); + sg_opts.merge_isa = result.error_model_isa; // this arm's own error model + TimingMark all_arm_mark(logger, "SearchSpaceGroup (all-observation arm)"); + alt = SearchSpaceGroup(sm_all.merged, sg_opts); + all_arm_mark.Stop(); + } // The order of the point group each arm confirmed. Read from the search, not from the // space group it went on to pick: Stage B leaves best_space_group unset when no // candidate has enough absences to be eligible, which is common on the filtered arm @@ -1535,6 +1614,13 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { sm = std::move(sm_all); } + // The arm merged beside the search is done with - or was not needed - before anything below + // changes the experiment it reads, and its card memory goes back before the merges below. + if (all_search_ahead.valid()) + all_search_ahead.wait(); + all_merge.reset(); + rsm_all.reset(); + // Everything from here to the adopted group's merge: the lattice re-asks and the choices // between the candidates the arms left. TimingMark decisions_mark(logger, "space-group decisions");