diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index b6b0b6d78..a04a80182 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -101,6 +102,8 @@ struct Context { std::unique_ptr scale_gpu; #endif size_t nthreads = 1; + int device = -1; // the engine's GPU + size_t engine_bytes = 0; // and what it holds there XrayRowModel rows; std::vector> fcalc, fmask; double k_sol = 0, b_sol = 0; @@ -611,6 +614,8 @@ void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logg size_t bytes = 0; ctx.engine = MakeRefinementEngineGPU(device, ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, ctx.Model(), ctx.b_ceiling, budget, &bytes, settings.nthreads); + ctx.device = device; + ctx.engine_bytes = bytes; if (ctx.engine) ctx.scale_gpu = std::make_unique(device, ctx.refl.Size(), 1); else @@ -630,6 +635,187 @@ void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logg ctx.engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, settings.nthreads); } +// The riding hydrogens placed on their parents as they are, and the offsets and B they ride with. +void PlaceRiders(Context &ctx, RefinementTopology &topo) { + const RefinementStage stage("riders"); + PlaceRidingHydrogens(topo); + ctx.rider_offset.assign(ctx.atoms.size(), gemmi::Position()); + for (size_t i = 0; i < ctx.atoms.size(); i++) + if (ctx.rider_parent[i] >= 0) { + ctx.rider_offset[i] = ctx.atoms[i]->pos - ctx.atoms[ctx.rider_parent[i]]->pos; + ctx.atoms[i]->b_iso = ctx.atoms[ctx.rider_parent[i]]->b_iso; + } +} + +// The geometry and the repulsions of the model as it is. +RestraintSet RestraintsNow(Context &ctx, const RefinementTopology &topo) { + const RefinementStage stage("repulsions"); + RestraintSet set = topo.geometry; + set.Append(BuildRepulsions(topo, ctx.Model(), ctx.cell, *ctx.sg, REPULSION_MARGIN, ctx.nthreads)); + set.Finalize(ctx.atoms.size()); + return set; +} + +// Where a weight trial runs: a model with its topology and its engine. +struct TrialSlot { + Context *ctx; + RefinementTopology *topo; +}; + +// A slot of its own beside the main one. +struct OwnSlot { + Context ctx; + RefinementTopology topo; +}; + +// A copy of the model as it was before the topology was made (the same topology comes out of it, hydrogens +// included), with an engine of the main one's kind - a GPU's on `device`. Null where the card cannot make one. +std::unique_ptr MakeOwnSlot(const gemmi::Structure &before_topology, const Context &main, int device, + size_t nthreads, const RefinementSettings &settings) { + auto slot = std::make_unique(); + Context &ctx = slot->ctx; + ctx.st = before_topology; + ctx.cell = main.cell; + ctx.sg = main.sg; + ctx.d_min = main.d_min; + ctx.refl = main.refl; + ctx.nthreads = nthreads; + ctx.b_ceiling = main.b_ceiling; + slot->topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens); + ctx.atoms = AtomList(ctx.Model()); + ctx.rider_parent = slot->topo.rider_parent; +#ifdef JFJOCH_USE_CUDA + if (main.engine->OnGPU()) { + try { + ctx.engine = MakeRefinementEngineGPU(device, ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, ctx.Model(), + ctx.b_ceiling, SIZE_MAX, nullptr, nthreads); + if (!ctx.engine) + return nullptr; + ctx.scale_gpu = std::make_unique(device, ctx.refl.Size(), 1); + } catch (const JFJochException &) { + return nullptr; + } + ctx.device = device; + } +#else + (void) device; +#endif + if (!ctx.engine) + ctx.engine = MakeRefinementEngineCPU(ctx.cell, *ctx.sg, ctx.d_min, ctx.refl, nthreads); + ctx.engine->SetRestraints(1, slot->topo.adp); + return slot; +} + +// The devices of the slots beside the main one, `wanted` of them at most: the GPUs in turn from the main +// engine's, each taking engines while they fit in its budget. -1 for each one on the CPU. +std::vector OwnSlotDevices(const Context &main, int wanted, const RefinementSettings &settings) { +#ifdef JFJOCH_USE_CUDA + if (main.engine->OnGPU()) { + const int n = get_gpu_count(); + std::vector used(n, 0); + used[main.device] = main.engine_bytes; + std::vector devices; + for (int k = 1; k <= wanted; k++) { + const int d = (main.device + k) % n; + const size_t usable = get_gpu_total_memory(d) / 5 * 4; + const size_t budget = usable > settings.gpu_resident_bytes ? usable - settings.gpu_resident_bytes : 0; + if (used[d] + main.engine_bytes <= budget) { + used[d] += main.engine_bytes; + devices.push_back(d); + } + } + return devices; + } +#else + (void) settings; +#endif + return std::vector(std::max(wanted, 0), -1); +} + +// One weight trial: a full refinement from pos0 and b0, macro-cycles of the positions at X-ray weight w and, +// where adp_w is not 0, the B factors at X-ray weight adp_w. +struct TrialOutcome { + RefinementWeightTrial trial; + std::vector pos; + std::vector b; + std::string line; // what the log says of it +}; + +TrialOutcome RunTrial(TrialSlot slot, const std::vector &refinable, const std::vector &pos0, + const std::vector &b0, double w, double adp_w, const RefinementSettings &settings) { + Context &ctx = *slot.ctx; + RefinementTopology &topo = *slot.topo; + const RefinementStage stage(adp_w > 0 ? "rung.b" : "rung.xyz"); + ctx.SetPositions(pos0); + ctx.SetB(b0); + for (int cycle = 0; cycle < settings.macro_cycles; cycle++) { + PlaceRiders(ctx, topo); + ctx.Refit(); + { + const RestraintSet set = RestraintsNow(ctx, topo); + const RefinementStage s1("set_restraints"); + ctx.engine->SetRestraints(0, set); + } + Minimise(ctx, AtomTarget::Mode::Xyz, refinable, w, 0, settings.xyz_iterations); + if (adp_w > 0) + Minimise(ctx, AtomTarget::Mode::B, refinable, adp_w, 1, settings.b_iterations); + } + PlaceRiders(ctx, topo); + TrialOutcome out; + RefinementWeightTrial &t = out.trial; + t.weight = w; + t.adp_weight = adp_w; + const RFactors rf = ctx.Refit(); + t.r_work = rf.r_work; + t.r_free = rf.r_free; + { + const RefinementStage s1("measure"); + t.geometry = MeasureGeometry(topo.geometry, + BuildRepulsions(topo, ctx.Model(), ctx.cell, *ctx.sg, 0.0, ctx.nthreads), ctx.XyzB()); + } + t.within_bounds = t.geometry.rmsd_bonds <= MAX_RMSD_BONDS && t.geometry.rmsd_angles <= MAX_RMSD_ANGLES; + out.line = fmt::format("Refinement, X-ray weight {:.3g}, B {}: R-work {:.4f}, R-free {:.4f}, rmsd bonds {:.4f} A, " + "angles {:.2f} deg{}", + w, adp_w > 0 ? fmt::format("refined at weight {:.3g}", adp_w) : std::string("held"), + rf.r_work, rf.r_free, t.geometry.rmsd_bonds, t.geometry.rmsd_angles, + t.within_bounds ? "" : " (geometry out of bounds)"); + out.pos = ctx.Positions(); + out.b = ctx.B(); + return out; +} + +// The trials (w, adp_w) side by side: trial i in slot i % slots, each slot's trials in turn on a thread of its +// own (on its GPU), returned in the order given. Every trial starts from pos0 and b0 and takes nothing from the +// one before it in its slot, so the outcome is the same for any number of slots. +std::vector RunTrials(const std::vector &trial_slots, + const std::vector> &trials, + const std::vector &refinable, const std::vector &pos0, + const std::vector &b0, const RefinementSettings &settings) { + std::vector out(trials.size()); + std::vector errors(trial_slots.size()); + auto run_slot = [&](size_t s) { + try { +#ifdef JFJOCH_USE_CUDA + const ScopedGpuPin pin(trial_slots[s].ctx->device >= 0 ? trial_slots[s].ctx->device : get_gpu()); +#endif + for (size_t i = s; i < trials.size(); i += trial_slots.size()) + out[i] = RunTrial(trial_slots[s], refinable, pos0, b0, trials[i].first, trials[i].second, settings); + } catch (...) { + errors[s] = std::current_exception(); + } + }; + std::vector threads; + for (size_t s = 1; s < std::min(trial_slots.size(), trials.size()); s++) + threads.emplace_back(run_slot, s); + run_slot(0); + for (std::thread &t : threads) + t.join(); + for (const std::exception_ptr &e : errors) + if (e) + std::rethrow_exception(e); + return out; +} + } // namespace RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector &merged, @@ -689,23 +875,14 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector topo_stage("topology"); RefinementTopology topo = BuildRefinementTopology(ctx.st, settings.riding_hydrogens); topo_stage.reset(); // the hydrogens it added: the atoms and the engine made again for them ctx.atoms = AtomList(ctx.Model()); ctx.rider_parent = topo.rider_parent; - auto place_riders = [&]() { - const RefinementStage stage("riders"); - PlaceRidingHydrogens(topo); - ctx.rider_offset.assign(ctx.atoms.size(), gemmi::Position()); - for (size_t i = 0; i < ctx.atoms.size(); i++) - if (ctx.rider_parent[i] >= 0) { - ctx.rider_offset[i] = ctx.atoms[i]->pos - ctx.atoms[ctx.rider_parent[i]]->pos; - ctx.atoms[i]->b_iso = ctx.atoms[ctx.rider_parent[i]]->b_iso; - } - }; - place_riders(); + PlaceRiders(ctx, topo); if (static_cast(ctx.atoms.size()) != result.atoms) { MakeEngine(ctx, result.on_gpu, settings, logger); result.on_gpu = ctx.engine->OnGPU(); @@ -725,16 +902,9 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vectorSetRestraints(1, topo.adp); { - const RestraintSet set = restraints_now(); + const RestraintSet set = RestraintsNow(ctx, topo); ctx.engine->SetRestraints(0, set); RestraintSet repulsions = BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0, ctx.nthreads); result.geometry_start = MeasureGeometry(topo.geometry, repulsions, ctx.XyzB()); @@ -750,67 +920,75 @@ RefinementResult RefineModel(const gemmi::Structure &placed, const std::vector pos0 = ctx.Positions(); const std::vector b0 = ctx.B(); + + // The trials run side by side, as many at a time as the longer ladder has rungs and the cards take engines: + // on the main model and engine, and on copies with engines of their own; each with a share of the threads. + std::optional slot_stage("slots"); + const int wanted = static_cast(std::max(settings.weight_ladder.size(), settings.adp_ladder.size())); + const std::vector devices = OwnSlotDevices(ctx, wanted - 1, settings); + const size_t slot_threads = std::max(1, settings.nthreads / (devices.size() + 1)); + std::vector> own(devices.size()); + { + std::vector threads; + for (size_t k = 0; k < devices.size(); k++) + threads.emplace_back([&, k]() { +#ifdef JFJOCH_USE_CUDA + const ScopedGpuPin pin(devices[k] >= 0 ? devices[k] : get_gpu()); +#endif + own[k] = MakeOwnSlot(before_topology, ctx, devices[k], slot_threads, settings); + }); + for (std::thread &t : threads) + t.join(); + } + const size_t main_threads = ctx.nthreads; + ctx.nthreads = slot_threads; + std::vector trial_slots{{&ctx, &topo}}; + for (const std::unique_ptr &o : own) + if (o) + trial_slots.push_back({&o->ctx, &o->topo}); + slot_stage.reset(); + logger.Debug("Refinement: {} weight trials at a time", trial_slots.size()); + int best = -1; std::vector best_pos; std::vector best_b; - // One full refinement from the start: macro-cycles of the positions at X-ray weight w and, where adp_w is - // not 0, the B factors at X-ray weight adp_w; kept where it beats the best so far. The rungs come in order - // of trust - tighter geometry, then fewer parameters - and a later one displaces an earlier only by more - // than the tie in R-free. - auto trial = [&](double w, double adp_w) { - const RefinementStage stage(adp_w > 0 ? "rung.b" : "rung.xyz"); - ctx.SetPositions(pos0); - ctx.SetB(b0); - for (int cycle = 0; cycle < settings.macro_cycles; cycle++) { - place_riders(); - ctx.Refit(); - { - const RestraintSet set = restraints_now(); - const RefinementStage s1("set_restraints"); - ctx.engine->SetRestraints(0, set); + // A trial is kept where it beats the best so far. The rungs come in order of trust - tighter geometry, + // then fewer parameters - and a later one displaces an earlier only by more than the tie in R-free. + auto choose = [&](std::vector &outcomes) { + for (TrialOutcome &o : outcomes) { + logger.Info("{}", o.line); + result.trials.push_back(o.trial); + const RefinementWeightTrial &t = result.trials.back(); + const int k = static_cast(result.trials.size()) - 1; + const bool better = best < 0 ? true + : t.within_bounds && (!result.trials[best].within_bounds || + t.r_free < result.trials[best].r_free - R_FREE_TIE); + if (better) { + best = k; + best_pos = std::move(o.pos); + best_b = std::move(o.b); } - Minimise(ctx, AtomTarget::Mode::Xyz, refinable, w, 0, settings.xyz_iterations); - if (adp_w > 0) - Minimise(ctx, AtomTarget::Mode::B, refinable, adp_w, 1, settings.b_iterations); - } - place_riders(); - RefinementWeightTrial t; - t.weight = w; - t.adp_weight = adp_w; - const RFactors rf = ctx.Refit(); - t.r_work = rf.r_work; - t.r_free = rf.r_free; - { - const RefinementStage s1("measure"); - t.geometry = MeasureGeometry(topo.geometry, BuildRepulsions(topo, ctx.Model(), cell, sg, 0.0, ctx.nthreads), ctx.XyzB()); - } - t.within_bounds = t.geometry.rmsd_bonds <= MAX_RMSD_BONDS && t.geometry.rmsd_angles <= MAX_RMSD_ANGLES; - logger.Info("Refinement, X-ray weight {:.3g}, B {}: R-work {:.4f}, R-free {:.4f}, rmsd bonds {:.4f} A, " - "angles {:.2f} deg{}", - w, adp_w > 0 ? fmt::format("refined at weight {:.3g}", adp_w) : std::string("held"), rf.r_work, - rf.r_free, t.geometry.rmsd_bonds, t.geometry.rmsd_angles, - t.within_bounds ? "" : " (geometry out of bounds)"); - result.trials.push_back(t); - const int k = static_cast(result.trials.size()) - 1; - const bool better = best < 0 ? true - : t.within_bounds && (!result.trials[best].within_bounds || - t.r_free < result.trials[best].r_free - R_FREE_TIE); - if (better) { - best = k; - best_pos = ctx.Positions(); - best_b = ctx.B(); } }; // the positions' weight, with the B factors as they are + std::vector> xyz_trials; for (double ladder : settings.weight_ladder) - trial(w0 * ladder, 0.0); + xyz_trials.emplace_back(w0 * ladder, 0.0); + std::vector outcomes = RunTrials(trial_slots, xyz_trials, refinable, pos0, b0, settings); + choose(outcomes); // then the B factors refined too, at that weight, against their own weight const double w = result.trials[best].weight; + std::vector> b_trials; for (double adp : settings.adp_ladder) - trial(w, w * adp); + b_trials.emplace_back(w, w * adp); + outcomes = RunTrials(trial_slots, b_trials, refinable, pos0, b0, settings); + choose(outcomes); + for (size_t s = 1; s < trial_slots.size(); s++) + ctx.evaluations += trial_slots[s].ctx->evaluations; + ctx.nthreads = main_threads; ctx.SetPositions(best_pos); ctx.SetB(best_b); - place_riders(); + PlaceRiders(ctx, topo); const RefinementWeightTrial &chosen = result.trials[best]; result.restrained = true; result.weight = chosen.weight;