diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index 0e45bd9ad..184e671ed 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -1633,8 +1634,11 @@ struct SlotPlace { // puts the long trials. A slot takes its engine's own part and a lane of its own where both fit, else its engine // alone on a lane it shares (one evaluation at a time: lanes side by side keep the card busier, more slots overlap // more host work); the first one on another card also the part the engines there share and a scale pool. On the -// CPU, the main one and those beside it are no more than the threads. -std::vector OwnSlotPlaces(const Context &main, int wanted, const RefinementSettings &settings) { +// CPU, the main one and those beside it are no more than the threads. `busiest`: what the main engine and the slots +// hold of the card that holds the most of them (0 on the CPU). +std::vector OwnSlotPlaces(const Context &main, int wanted, const RefinementSettings &settings, + size_t &busiest) { + busiest = 0; #ifdef JFJOCH_USE_CUDA if (main.engine->OnGPU()) { const int n = get_gpu_count(); @@ -1656,6 +1660,7 @@ std::vector OwnSlotPlaces(const Context &main, int wanted, const Refi places.push_back({d, false}); } } + busiest = *std::max_element(used.begin(), used.end()); return places; } #endif @@ -2237,10 +2242,20 @@ double RefinementDMin(double max_inv_d2) { // One refinement of a model, with the solvent build in it where settings.solvent asks for it. // `twin_laws`: against the data as twinned by them (one domain each beside the model's own), empty for untwinned; // `null_first`: the first of them the null (TwinNull) replaces the mates of. +// What a refinement holds of the card that holds the most of it (OwnSlotPlaces), said once (RefineModel). +struct CardHeld { + std::promise bytes; + std::once_flag said; + void Say(size_t b) { + std::call_once(said, [&] { bytes.set_value(b); }); + } +}; + RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector &merged, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double wavelength_A, Logger &logger, const RefinementSettings &settings, - const std::vector &twin_laws = {}, size_t null_first = 0) { + const std::vector &twin_laws = {}, size_t null_first = 0, + CardHeld *held = nullptr) { const auto start = std::chrono::steady_clock::now(); RefinementResult result; Context ctx; @@ -2485,8 +2500,10 @@ RefinementResult RefineOnce(const gemmi::Structure &placed, const std::vector 0) wanted = std::min(wanted, static_cast(settings.max_trials_at_once)); - const std::vector devices = settings.restrained ? OwnSlotPlaces(ctx, wanted - 1, settings) - : std::vector(); + size_t busiest = 0; + const std::vector devices = OwnSlotPlaces(ctx, settings.restrained ? wanted - 1 : 0, settings, busiest); + if (held != nullptr) + held->Say(busiest); const size_t slot_threads = std::max(1, settings.nthreads / (devices.size() + 1)); ScalePools pools; #ifdef JFJOCH_USE_CUDA @@ -3297,15 +3314,31 @@ double FinalRWork(const RefinementResult &r) { // standard errors, each added domain at a fraction of TWIN_DOMAIN_MIN at least. void DecideTwin(const gemmi::Structure &placed, const std::vector &merged, const gemmi::UnitCell &cell, const gemmi::SpaceGroup &sg, double wavelength_A, Logger &logger, - const RefinementSettings &settings, RefinementResult &result) { + const RefinementSettings &settings, RefinementResult &result, const std::shared_future &untwinned) { const RefinementStage stage("twin"); - result.untwinned_r_work = FinalRWork(result); - result.untwinned_r_free = FinalRFree(result); const std::vector laws = RefinementTwinLaws(cell, sg); const std::vector domains = RefinementTwinDomains(cell, sg); - result.twin_laws = static_cast(laws.size()); + // The twinned arm reads only the placed model and the data, so it runs beside the untwinned refinement + // (RefineModel): `result` is read and written only once `untwinned` says that refinement is in it, and + // nothing is decided where it failed or was cut short. + std::string split_law; + int split_pairs = 0; + auto take_untwinned = [&]() { + untwinned.wait(); + if (!result.ok || result.cancelled) + return false; + result.untwinned_r_work = FinalRWork(result); + result.untwinned_r_free = FinalRFree(result); + result.twin_laws = static_cast(laws.size()); + if (!split_law.empty()) { + result.twin_law = split_law; + result.twin_split_pairs = split_pairs; + } + return true; + }; if (laws.empty()) { - result.twin_decision = "NO_TWIN_LAW"; + if (take_untwinned()) + result.twin_decision = "NO_TWIN_LAW"; return; } // the laws whose twin pairs the free set keeps together @@ -3318,15 +3351,16 @@ void DecideTwin(const gemmi::Structure &placed, const std::vector 1) { RefinementSettings score = settings; @@ -3342,8 +3376,11 @@ void DecideTwin(const gemmi::Structure &placed, const std::vector &data = moved > 0 ? closed : merged; - RefinementResult result = RefineOnce(placed, data, cell, sg, wavelength_A, logger, settings); // Each arm, untwinned and twinned, is a whole refinement - the solvent built and the displacement model chosen // against its own target - and the decision is taken on where they ended. A model only scored is scored both // ways and decided by the same rule, so the R-free it reports is not an untwinned one by default. - if (settings.twin && result.ok && !result.cancelled) - DecideTwin(placed, data, cell, sg, wavelength_A, logger, settings, result); + // The arms run side by side, the twinned one on the next card where there is more than one, once the untwinned + // one has placed its engines and slots - and where twice what it holds of its busiest card does not fit the + // budget, after it, so that the twinned arm's engines are made as they would be alone. Its log is kept and + // follows the untwinned one's, as when they ran one after the other. + RefinementResult result; + if (!settings.twin) + result = RefineOnce(placed, data, cell, sg, wavelength_A, logger, settings); + else { + Logger twin_logger = Logger::Buffered(); + CardHeld held; + std::future held_bytes = held.bytes.get_future(); + std::promise untwinned_done; + const std::shared_future untwinned = untwinned_done.get_future().share(); + size_t budget = 0; +#ifdef JFJOCH_USE_CUDA + const int twin_card = get_gpu_count() > 1 ? (get_gpu() + 1) % get_gpu_count() : get_gpu(); + if (get_gpu_count() > 0) { + const size_t usable = get_gpu_total_memory(twin_card) / 5 * 4; + budget = usable > settings.gpu_resident_bytes ? usable - settings.gpu_resident_bytes : 0; + } +#endif + std::future twin_arm = std::async(std::launch::async, [&]() { +#ifdef JFJOCH_USE_CUDA + const ScopedGpuPin pin(twin_card); +#endif + if (2 * held_bytes.get() > budget) + untwinned.wait(); + DecideTwin(placed, data, cell, sg, wavelength_A, twin_logger, settings, result, untwinned); + }); + try { + result = RefineOnce(placed, data, cell, sg, wavelength_A, logger, settings, {}, 0, &held); + } catch (...) { + result = RefinementResult(); // not ok: the twinned arm decides nothing + held.Say(0); + untwinned_done.set_value(); + twin_arm.wait(); + throw; + } + held.Say(0); // where it stopped before placing its slots + untwinned_done.set_value(); + twin_arm.get(); + twin_logger.ReplayInto(logger); + } result.twin_free_closed = moved; result.map.reset(); return result;