diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index f8a8ff5c9..b4711ecf3 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -77,6 +77,7 @@ ADD_LIBRARY(JFJochCommon STATIC DetectorModuleGeometry.cpp DetectorModuleGeometry.h DetectorSetup.h DetectorSetup.cpp ZeroCopyReturnValue.h Histogram.h DiffractionGeometry.h CUDAWrapper.cpp CUDAWrapper.h + HostMemory.cpp HostMemory.h ThreadAffinity.cpp ThreadAffinity.h ADUHistogram.cpp ADUHistogram.h RawToConvertedGeometryCore.h diff --git a/common/HostMemory.cpp b/common/HostMemory.cpp new file mode 100644 index 000000000..54c4a8b66 --- /dev/null +++ b/common/HostMemory.cpp @@ -0,0 +1,77 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "HostMemory.h" + +#if defined(_WIN32) +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +#elif defined(__APPLE__) +#include +#include +#include +#else +#include +#endif + +size_t get_host_total_memory() { +#if defined(_WIN32) + MEMORYSTATUSEX status{}; + status.dwLength = sizeof(status); + if (!GlobalMemoryStatusEx(&status)) + return 0; + return static_cast(status.ullTotalPhys); +#elif defined(__APPLE__) + uint64_t bytes = 0; + size_t len = sizeof(bytes); + if (sysctlbyname("hw.memsize", &bytes, &len, nullptr, 0) != 0) + return 0; + return static_cast(bytes); +#else + const long pages = sysconf(_SC_PHYS_PAGES), page = sysconf(_SC_PAGE_SIZE); + if (pages <= 0 || page <= 0) + return 0; + return static_cast(pages) * static_cast(page); +#endif +} + +size_t DefaultHostMemoryBudget() { + return static_cast(HOST_MEMORY_BUDGET_FRACTION * static_cast(get_host_total_memory())); +} + +HostWorkerPool &HostWorkerPool::Instance() { + static HostWorkerPool pool; + return pool; +} + +HostWorkerPool::HostWorkerPool() { + SetSize(0); +} + +void HostWorkerPool::SetSize(size_t bytes) { + std::lock_guard lock(m); + size = bytes > 0 ? bytes + : static_cast(HOST_WORKER_POOL_FRACTION * static_cast(DefaultHostMemoryBudget())); + cv.notify_all(); +} + +size_t HostWorkerPool::Size() { + std::lock_guard lock(m); + return size; +} + +HostWorkerPool::Reservation::Reservation(HostWorkerPool &pool, size_t bytes) : pool(pool), bytes(bytes) { + std::unique_lock lock(pool.m); + pool.cv.wait(lock, [&] { return pool.reserved == 0 || pool.reserved + bytes <= pool.size; }); + pool.reserved += bytes; +} + +HostWorkerPool::Reservation::~Reservation() { + { + std::lock_guard lock(pool.m); + pool.reserved -= bytes; + } + pool.cv.notify_all(); +} diff --git a/common/HostMemory.h b/common/HostMemory.h new file mode 100644 index 000000000..d016ddb58 --- /dev/null +++ b/common/HostMemory.h @@ -0,0 +1,60 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include + +// The machine's physical memory in bytes - its total, not what is free at the moment. 0 where it +// cannot be read. +size_t get_host_total_memory(); + +// The host memory a run plans against: this fraction of the machine's physical memory, decided from +// the total, never from what happens to be free (as what runs on a GPU is decided on the card's total). +// The plan decides only what runs at the same time - two merges side by side or one after the other, +// how many large per-worker buffers are held at once - and never what is computed, so no result +// depends on it. +constexpr double HOST_MEMORY_BUDGET_FRACTION = 0.75; +size_t DefaultHostMemoryBudget(); + +// The share of the budget that the large per-worker buffers of parallel steps - a translation function's +// grids, a null replicate's model maps - may hold at once. The rest is what the run holds beside them: +// its reflections, its merge, the detector-sized arrays. +constexpr double HOST_WORKER_POOL_FRACTION = 1.0 / 3.0; + +// One pool for those buffers, process-wide. A worker reserves its bytes before it makes them and hands +// them back when it is done; one that would take the pool past its size waits until others have handed +// theirs back - except that a reservation made while nothing is reserved is always let through, however +// large. So how many such workers run at once is decided by their bytes, not by the thread count, and +// however the steps holding them nest or run side by side; what each worker computes, and where its +// result goes, does not change. A worker must not reserve again while it holds a reservation, and none +// may reserve on ParallelFor's pool: a holder's own parallel passes need those threads. Work spread over +// that pool is bounded by its count instead (Size() / bytes each). +class HostWorkerPool { +public: + static HostWorkerPool &Instance(); + + // The pool's size in bytes; 0 sets it back to HOST_WORKER_POOL_FRACTION of DefaultHostMemoryBudget(). + void SetSize(size_t bytes); + [[nodiscard]] size_t Size(); + + class Reservation { + public: + Reservation(HostWorkerPool &pool, size_t bytes); + ~Reservation(); + Reservation(const Reservation &) = delete; + Reservation &operator=(const Reservation &) = delete; + private: + HostWorkerPool &pool; + size_t bytes; + }; + +private: + HostWorkerPool(); + std::mutex m; + std::condition_variable cv; + size_t size = 0; + size_t reserved = 0; +}; diff --git a/image_analysis/geom_refinement/BeamCenterFFTCPU.cpp b/image_analysis/geom_refinement/BeamCenterFFTCPU.cpp index 659c5c0ca..1496752f5 100644 --- a/image_analysis/geom_refinement/BeamCenterFFTCPU.cpp +++ b/image_analysis/geom_refinement/BeamCenterFFTCPU.cpp @@ -83,34 +83,39 @@ public: BeamCenterConvSurfaces2D BeamCenterFFTCPU::PointSurfaces(const std::vector &a, const std::vector &m, int64_t h, int64_t w) { - // The three forward transforms are independent of each other, and so are the four inverse ones, - // so each group runs at the same time, one workspace per transform. A workspace is built exactly - // as a single one was - its own buffers and its own plan for them - so every transform is the one - // it was, bit for bit; the four inverse ones reuse the three forward workspaces and one more. + // Two transforms at a time, each on a workspace of its own. A workspace is padded to twice the + // detector each way, so on a large detector the four workspaces and three spectra that the seven + // transforms used to hold at once were gigabytes. In order: the forward transforms of the image + // and the mask; the two inverse ones that read the image's spectrum, which is then handed back; + // the squared image's forward transform beside the mask's own inverse; and the last inverse. + // A workspace is built exactly as a single one was - its own buffers and its own plan for them - + // so every transform is the one it was, bit for bit. const int64_t ny = BeamCenterFFTPadSize(2 * h), nx = BeamCenterFFTPadSize(2 * w); - std::vector> fft; - for (int i = 0; i < 4; i++) - fft.push_back(std::make_unique(ny, nx)); + Rfft2 ws0(ny, nx), ws1(ny, nx); + + std::vector> A, M; + { + auto fA = std::async(std::launch::async, [&] { return ws0.Forward(a, h, w); }); + M = ws1.Forward(m, h, w); + A = fA.get(); + } + + BeamCenterConvSurfaces2D out; + { + auto fC = std::async(std::launch::async, [&] { return ws0.InverseProduct(A, A, 2 * h, 2 * w); }); + out.S = ws1.InverseProduct(A, M, 2 * h, 2 * w); + out.C = fC.get(); + } + std::vector>().swap(A); std::vector a2(a.size()); for (size_t i = 0; i < a2.size(); i++) a2[i] = a[i] * a[i]; - auto fA = std::async(std::launch::async, [&] { return fft[0]->Forward(a, h, w); }); - auto fM = std::async(std::launch::async, [&] { return fft[1]->Forward(m, h, w); }); - auto fA2 = std::async(std::launch::async, [&] { return fft[2]->Forward(a2, h, w); }); - const auto A = fA.get(); - const auto M = fM.get(); - const auto A2 = fA2.get(); - - BeamCenterConvSurfaces2D out; - auto fC = std::async(std::launch::async, [&] { return fft[0]->InverseProduct(A, A, 2 * h, 2 * w); }); - auto fS = std::async(std::launch::async, [&] { return fft[1]->InverseProduct(A, M, 2 * h, 2 * w); }); - auto fQ = std::async(std::launch::async, [&] { return fft[2]->InverseProduct(A2, M, 2 * h, 2 * w); }); - auto fD = std::async(std::launch::async, [&] { return fft[3]->InverseProduct(M, M, 2 * h, 2 * w); }); - out.C = fC.get(); - out.S = fS.get(); - out.Q = fQ.get(); - out.D = fD.get(); + { + auto fD = std::async(std::launch::async, [&] { return ws1.InverseProduct(M, M, 2 * h, 2 * w); }); + const auto A2 = ws0.Forward(a2, h, w); + out.Q = ws0.InverseProduct(A2, M, 2 * h, 2 * w); + out.D = fD.get(); + } return out; } - diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 86dc97067..8610dfc36 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -596,10 +596,8 @@ void RotationScaleMerge::IngestHost() { SmoothMosaicityAndPartiality(); ParallelChunks(static_cast(partials.size()), ThreadsForWork(partials.size(), nthreads), [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { - partials[i].delta_phi = ingest_delta_phi[i]; + for (int i = lo; i < hi; ++i) partials[i].partiality = ingest_partiality[i]; - } }); // The corr every pass starts from (see the header) was filled by the build, straight from the @@ -781,6 +779,43 @@ size_t RotationScaleMerge::DeviceBytes() const { return std::max(bytes + bytes / 5, ingest); } +namespace { + // The host bytes the ingest holds per partial from then on: the corr every Run starts from and, with + // the GPU resident, the predicted position the fulls are given; on the host, the partial itself, the + // raw-hkl order, the finiteness mask and the ASU grouping's three per-partial arrays. Per raw-hkl run: + // its start, count, h, k, l, d and group. + constexpr size_t HOST_INGEST_BYTES_PER_PARTIAL_GPU = 3 * sizeof(float); + constexpr size_t HOST_INGEST_BYTES_PER_RUN = 7 * sizeof(int32_t); + // What one merge holds on the host at its peak, per rocking event (a full each at most: the full, its + // download staging and the merge's per-full fields), per distinct hkl (an ASU group and a merged + // reflection each at most: the merge's per-group accumulators, R_meas sums and their staging, the ISa + // scatter, the merged reflection), and per partial where the engine scales on the host (the scaling + // loop's copies of the partials and the memo of their corr). + constexpr size_t HOST_MERGE_BYTES_PER_EVENT = 190, HOST_MERGE_BYTES_PER_HKL = 420; + constexpr size_t HOST_MERGE_BYTES_PER_PARTIAL_CPU = 65; +} + +size_t RotationScaleMerge::HostIngestBytes() const { + constexpr size_t PER_PARTIAL_CPU = sizeof(PartialObs) + sizeof(float) + sizeof(uint8_t) + 4 * sizeof(int32_t); + const size_t per_partial = OnGPU() ? HOST_INGEST_BYTES_PER_PARTIAL_GPU : PER_PARTIAL_CPU; + return per_partial * static_cast(n_partials_obs) + HOST_INGEST_BYTES_PER_RUN * rawrun_start.size(); +} + +size_t RotationScaleMerge::HostMergeBytes() const { + // The device path's ingest counted the rocking events; without it every partial may be one. + const auto events = static_cast(rocking_events_at_ingest > 0 ? rocking_events_at_ingest : n_partials_obs); + return HOST_MERGE_BYTES_PER_EVENT * events + HOST_MERGE_BYTES_PER_HKL * rawrun_start.size() + + (OnGPU() ? 0 : HOST_MERGE_BYTES_PER_PARTIAL_CPU * static_cast(n_partials_obs)); +} + +size_t RotationScaleMerge::HostBytesBound(size_t n_reflections, bool on_gpu) { + constexpr size_t PER_PARTIAL_CPU = sizeof(PartialObs) + sizeof(float) + sizeof(uint8_t) + 4 * sizeof(int32_t) + + HOST_MERGE_BYTES_PER_PARTIAL_CPU; + const size_t per_reflection = (on_gpu ? HOST_INGEST_BYTES_PER_PARTIAL_GPU : PER_PARTIAL_CPU) + + HOST_INGEST_BYTES_PER_RUN + HOST_MERGE_BYTES_PER_EVENT + HOST_MERGE_BYTES_PER_HKL; + return per_reflection * n_reflections; +} + std::unique_ptr RotationScaleMerge::CloneIngest(int device, Logger &log) const { TimingMark mark(*logger, "rsm ingest copy"); std::unique_ptr c(new RotationScaleMerge(*this)); @@ -831,12 +866,13 @@ void RotationScaleMerge::IngestDevice(const RotationScaleMerge *observations_fro const int cnt = std::min(n - at, slice); ParallelChunks(cnt, nthreads, [&](int lo, int hi) { for (int i = lo; i < hi; ++i) { - const auto &o = partials[at + i]; - fbuf[0][i] = o.I; fbuf[1][i] = o.sigma; fbuf[2][i] = o.prescaling_corr; - fbuf[3][i] = o.partiality; fbuf[4][i] = o.zeta; fbuf[5][i] = o.corr; - fbuf[6][i] = o.bkg; fbuf[7][i] = o.var_bkg; fbuf[8][i] = o.image_number; - fbuf[9][i] = o.d; - flags[i] = static_cast((o.on_ice ? 1 : 0) | (o.clipped << 1)); // SetObsFlags + const PartialObs &o = partials[at + i]; + const Reflection &r = *o.r; + fbuf[0][i] = r.I; fbuf[1][i] = r.sigma; fbuf[2][i] = o.prescaling_corr; + fbuf[3][i] = o.partiality; fbuf[4][i] = r.zeta; fbuf[5][i] = o.corr; + fbuf[6][i] = r.bkg; fbuf[7][i] = r.var_bkg; fbuf[8][i] = r.image_number; + fbuf[9][i] = r.d; + flags[i] = static_cast((r.on_ice_ring ? 1 : 0) | (ClippedFlags(r) << 1)); // SetObsFlags } }); constexpr F fields[NFLOAT] = {F::I, F::Sigma, F::PrescalingCorr, F::Partiality, @@ -849,7 +885,7 @@ void RotationScaleMerge::IngestDevice(const RotationScaleMerge *observations_fro auto pos = std::make_shared(); pos->px.resize(n); pos->py.resize(n); ParallelChunks(n, nthreads, [&](int lo, int hi) { - for (int i = lo; i < hi; ++i) { pos->px[i] = partials[i].px; pos->py[i] = partials[i].py; } + for (int i = lo; i < hi; ++i) { pos->px[i] = partials[i].r->predicted_x; pos->py[i] = partials[i].r->predicted_y; } }); obs_position = std::move(pos); gpu_->SetRawRuns(static_cast(rawrun_start.size()), static_cast(perm.size()), perm.data(), @@ -1000,25 +1036,18 @@ void RotationScaleMerge::BuildInRangeObservations(std::unique_ptr &ke for (const auto &r : partials_out[o].reflections) { if (!keep[src]) { ++src; continue; } ++src; - Obs &obs = partials[at]; - obs.h = r.h; obs.k = r.k; obs.l = r.l; - obs.I = r.I; obs.sigma = r.sigma; obs.d = r.d; + PartialObs &obs = partials[at]; + obs.r = &r; // The whole deterministic correction - Lorentz x polarization x sensor efficiency x air path - is the // numerator of corr, so the two halves are multiplied back together here. The incident flux is // divided out at the same time so the fits below see only what is left; it is the one term in // the product that is measured from the data rather than computed from the geometry. obs.prescaling_corr = static_cast(r.prescaling_corr * r.qe_corr * r.flight_corr / frame_flux[o]); - obs.partiality = r.partiality; obs.zeta = r.zeta; obs.delta_phi = r.delta_phi_deg; obs.bkg = r.bkg; obs.var_bkg = r.var_bkg; - obs.var_per_I = 0.0f; obs.capture = 0.0f; - obs.px = r.predicted_x; obs.py = r.predicted_y; - obs.image_number = r.image_number; + obs.partiality = r.partiality; obs.frame = o; - obs.on_ice = r.on_ice_ring ? 1 : 0; - obs.clipped = (r.clipped ? OBS_CLIPPED : 0) | (r.overloaded ? OBS_OVERLOADED : 0); obs.corr = r.image_scale_corr; - obs.group = -1; - finite_ok[at] = (std::isfinite(obs.I) && std::isfinite(obs.prescaling_corr) && obs.prescaling_corr != 0.0f - && std::isfinite(obs.sigma) && obs.sigma > 0.0f) ? 1 : 0; + finite_ok[at] = (std::isfinite(r.I) && std::isfinite(obs.prescaling_corr) && obs.prescaling_corr != 0.0f + && std::isfinite(r.sigma) && r.sigma > 0.0f) ? 1 : 0; ingest_delta_phi[at] = SmoothedDeltaPhi(sg, o, r.h, r.k, r.l, r.delta_phi_deg, resolved); ingest_partiality[at] = r.partiality; ingest_image_number[at] = r.image_number; @@ -1562,8 +1591,8 @@ void RotationScaleMerge::ReducePartialGroupMeans(int n_groups, std::vectorI, o.r->sigma, min_partiality, I_corr); if (w <= 0.0) continue; sw += w; swI += w * I_corr; @@ -2813,11 +2842,11 @@ int RotationScaleMerge::RockingEventFrames(int64_t *n_events) const { } return RockingEventFramesOver( [&](int i) { - const Obs &o = partials[i]; - return std::isfinite(o.corr) && o.corr > 0.0f && std::isfinite(o.I) - && std::isfinite(o.sigma) && o.sigma > 0.0f; + const PartialObs &o = partials[i]; + return std::isfinite(o.corr) && o.corr > 0.0f && std::isfinite(o.r->I) + && std::isfinite(o.r->sigma) && o.r->sigma > 0.0f; }, - [&](int i) { return partials[i].image_number; }, n_events); + [&](int i) { return partials[i].r->image_number; }, n_events); } template @@ -4389,15 +4418,23 @@ void RotationScaleMerge::Combine() { // A raw hkl's partials lie a frame apart in memory, and the work on one rocking event is long // enough that the CPU does not reach the next runs' loads on its own - the combine waited on a // cache miss for nearly every partial. So each run asks for the partials of the run - // COMBINE_PREFETCH_RUNS ahead. + // COMBINE_PREFETCH_RUNS ahead, and for the reflections of the run half as far ahead, whose partials + // - which say where the reflections are - were asked for that many runs before. constexpr int COMBINE_PREFETCH_RUNS = 8; auto prefetch_run = [&](int r) { - if (r >= n_run) return; - for (int p = rawrun_start[r]; p < rawrun_start[r] + rawrun_count[r]; ++p) { - const char *o = reinterpret_cast(&partials[perm[p]]); - Prefetch(o); - Prefetch(o + sizeof(Obs) - 1); - } + if (r < n_run) + for (int p = rawrun_start[r]; p < rawrun_start[r] + rawrun_count[r]; ++p) { + const char *o = reinterpret_cast(&partials[perm[p]]); + Prefetch(o); + Prefetch(o + sizeof(PartialObs) - 1); + } + const int rr = r - COMBINE_PREFETCH_RUNS / 2; + if (rr < n_run) + for (int p = rawrun_start[rr]; p < rawrun_start[rr] + rawrun_count[rr]; ++p) { + const char *o = reinterpret_cast(partials[perm[p]].r); + Prefetch(o); + Prefetch(o + sizeof(Reflection) - 1); + } }; // Combine one raw-hkl run into fulls, appended to `out`. Independent per run (the events of one hkl @@ -4409,9 +4446,9 @@ void RotationScaleMerge::Combine() { std::vector ev; // usable perm-indices of this raw hkl, in image-number order ev.reserve(hi - lo); for (int p = lo; p < hi; ++p) { - const auto &o = partials[perm[p]]; + const PartialObs &o = partials[perm[p]]; if (!std::isfinite(o.corr) || o.corr <= 0.0f) continue; - if (!std::isfinite(o.I) || !std::isfinite(o.sigma) || o.sigma <= 0.0f) continue; + if (!std::isfinite(o.r->I) || !std::isfinite(o.r->sigma) || o.r->sigma <= 0.0f) continue; ev.push_back(perm[p]); } const int group = rawrun_group[r]; @@ -4419,9 +4456,9 @@ void RotationScaleMerge::Combine() { size_t i = 0; while (i < ev.size()) { size_t kk = i + 1; - float last_frame = partials[ev[i]].image_number; + float last_frame = partials[ev[i]].r->image_number; while (kk < ev.size()) { - const float frame = partials[ev[kk]].image_number; + const float frame = partials[ev[kk]].r->image_number; if (frame - last_frame > max_frame_gap) break; last_frame = frame; ++kk; @@ -4430,11 +4467,11 @@ void RotationScaleMerge::Combine() { double pooled_bkg = 0.0; int n_pool = 0; for (size_t m = i; m < kk; ++m) { - const float b = partials[ev[m]].bkg; + const float b = partials[ev[m]].r->bkg; if (std::isfinite(b)) { pooled_bkg += b; ++n_pool; } } pooled_bkg = n_pool > 0 ? pooled_bkg / n_pool : 0.0; - auto pooled_I = [&](const Obs &r2) { + auto pooled_I = [&](const Reflection &r2) { const double n_bkg = std::max(0.0, static_cast(r2.sigma) * r2.sigma - r2.I) / std::max(r2.bkg, 1.0f); return static_cast(r2.I) + n_bkg * (static_cast(r2.bkg) - pooled_bkg); @@ -4443,12 +4480,12 @@ void RotationScaleMerge::Combine() { double sum_w = 0.0, sum_wI = 0.0, sum_partiality = 0.0; float d = NAN; int peak_outcome = partials[ev[i]].frame; - float peak_frame = partials[ev[i]].image_number; - float peak_px = partials[ev[i]].px, peak_py = partials[ev[i]].py; + float peak_frame = partials[ev[i]].r->image_number; + float peak_px = partials[ev[i]].r->predicted_x, peak_py = partials[ev[i]].r->predicted_y; float peak_partiality = -1.0f; - const bool on_ice = partials[ev[i]].on_ice; + const bool on_ice = partials[ev[i]].r->on_ice_ring; uint8_t flags = 0; - for (size_t m = i; m < kk; ++m) flags |= partials[ev[m]].clipped; + for (size_t m = i; m < kk; ++m) flags |= ClippedFlags(*partials[ev[m]].r); // A saturated pixel anywhere in the rocking curve means its brightest part was not measured: // the event is not a measurement of the reflection, and summing what is left - or scaling // the flanks up by the partiality model - reads it 2-3x low. Dropped whole, as XDS drops an @@ -4460,17 +4497,18 @@ void RotationScaleMerge::Combine() { } const bool clipped = flags & OBS_CLIPPED; for (size_t m = i; m < kk; ++m) { - const auto &r2 = partials[ev[m]]; - const double sigma_corr = static_cast(r2.sigma) * r2.corr; + const PartialObs &p2 = partials[ev[m]]; + const Reflection &r2 = *p2.r; + const double sigma_corr = static_cast(r2.sigma) * p2.corr; const double w = 1.0 / (sigma_corr * sigma_corr); sum_w += w; - sum_wI += w * pooled_I(r2) * r2.corr; - sum_partiality += r2.partiality; - if (r2.partiality > peak_partiality) { - peak_partiality = r2.partiality; - peak_outcome = r2.frame; + sum_wI += w * pooled_I(r2) * p2.corr; + sum_partiality += p2.partiality; + if (p2.partiality > peak_partiality) { + peak_partiality = p2.partiality; + peak_outcome = p2.frame; peak_frame = r2.image_number; - peak_px = r2.px; peak_py = r2.py; + peak_px = r2.predicted_x; peak_py = r2.predicted_y; } if (!std::isfinite(d) && std::isfinite(r2.d) && r2.d > 0.0f) d = r2.d; } @@ -4482,8 +4520,8 @@ void RotationScaleMerge::Combine() { for (int iter = 0; iter < 3; ++iter) { sum_w = 0.0; sum_wI = 0.0; sum_wb = 0.0; sum_cwb = 0.0; for (size_t m = i; m < kk; ++m) { - const auto &r2 = partials[ev[m]]; - const double corr = r2.corr; + const Reflection &r2 = *partials[ev[m]].r; + const double corr = partials[ev[m]].corr; const double I_corr = pooled_I(r2) * corr; const double sigma_corr = static_cast(r2.sigma) * corr; // The non-signal variance as the integrator measured it, which is where it belongs: @@ -4661,16 +4699,17 @@ void RotationScaleMerge::ComputePerFrameCC(const std::vector &partial_gr size_t n = 0; const int lo = frame_start[f], hi = frame_start[f] + frame_count[f]; for (int i = lo; i < hi; ++i) { - const auto &o = partials[i]; + const PartialObs &o = partials[i]; + const Reflection &r = *o.r; const int g = partial_group[i]; - if (o.on_ice) continue; + if (r.on_ice_ring) continue; if (g < 0) continue; if (o.partiality < min_partiality) continue; - if (!std::isfinite(o.I) || !std::isfinite(o.corr) || o.corr <= 0.0f) continue; - if (!std::isfinite(o.sigma) || o.sigma <= 0.0f) continue; + if (!std::isfinite(r.I) || !std::isfinite(o.corr) || o.corr <= 0.0f) continue; + if (!std::isfinite(r.sigma) || r.sigma <= 0.0f) continue; const double mean = partial_group_mean[g]; if (!std::isfinite(mean)) continue; - const double img = static_cast(o.I) * o.corr; + const double img = static_cast(r.I) * o.corr; sx += img; sy += mean; sx2 += img * img; sy2 += mean * mean; sxy += img * mean; ++n; } @@ -6435,7 +6474,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool #endif // The CPU partial scaling below writes every partial's corr when it is done, from a copy it starts // off corr_ingested itself, and nothing between here and there reads corr - so there the restore - // is not made, which saves a pass over the eighty-byte records to set four bytes of each. + // is not made, which saves a pass over the partials to set four bytes of each. if (!restored_on_gpu && !scale_partials) ParallelChunks(static_cast(partials.size()), ThreadsForWork(partials.size(), nthreads), [&](int lo, int hi) { @@ -6568,8 +6607,8 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool FillVector scaling_corr(n_partials); ParallelChunks(n_partials, partial_threads, [&](int lo, int hi) { for (int q = lo; q < hi; ++q) { - const Obs &o = partials[scaling_index[q]]; - scaling_refs[q] = {o.frame, o.I, o.sigma, o.partiality, o.prescaling_corr}; + const PartialObs &o = partials[scaling_index[q]]; + scaling_refs[q] = {o.frame, o.r->I, o.r->sigma, o.partiality, o.prescaling_corr}; scaling_corr[q] = corr_ingested[scaling_index[q]]; // the restore (see the top) } }); @@ -6580,9 +6619,9 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool // fitted on them is fitted on noise that the reference does not answer to. The Rotation // model needs zeta > 0. auto fitted_obs = [&](int i) { - const Obs &o = partials[i]; - return partial_group[i] >= 0 && !o.on_ice && o.partiality >= min_partiality - && std::isfinite(o.zeta) && o.zeta > 0.0f; + const PartialObs &o = partials[i]; + return partial_group[i] >= 0 && !o.r->on_ice_ring && o.partiality >= min_partiality + && std::isfinite(o.r->zeta) && o.r->zeta > 0.0f; }; const int n_fit_frames = static_cast(frame_start.size()); std::vector fit_start(n_fit_frames + 1, 0); @@ -6597,8 +6636,8 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool int32_t k = fit_start[f]; for (int i = frame_start[f]; i < frame_start[f] + frame_count[f]; ++i) { if (!fitted_obs(i)) continue; - const Obs &o = partials[i]; - scaling_fit[k++] = {partial_group[i], o.I, o.sigma, o.partiality, o.prescaling_corr}; + const PartialObs &o = partials[i]; + scaling_fit[k++] = {partial_group[i], o.r->I, o.r->sigma, o.partiality, o.prescaling_corr}; } }); @@ -6679,7 +6718,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool ParallelChunks(static_cast(partials.size()), ThreadsForWork(partials.size(), nthreads), [&](int lo, int hi) { for (int i = lo; i < hi; ++i) { - Obs &o = partials[i]; + PartialObs &o = partials[i]; if (apply[o.frame] && std::isfinite(o.corr)) o.corr = static_cast(o.corr * ratio[o.frame]); } @@ -6738,8 +6777,8 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool [&](int lo, int hi) { int64_t n = 0; for (int i = lo; i < hi; ++i) { - Obs &o = partials[i]; - if (!(std::isfinite(o.zeta) && o.zeta >= search_min_zeta)) { + PartialObs &o = partials[i]; + if (!(std::isfinite(o.r->zeta) && o.r->zeta >= search_min_zeta)) { if (std::isfinite(o.corr) && o.corr > 0.0f) ++n; o.corr = 0.0f; } @@ -6873,7 +6912,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool // to grow they are freed first: growing in place would copy the previous chain's records into // the new allocation only for the gather to overwrite them. if (static_cast(nf) > fulls.capacity()) { - fulls = {}; + Free(fulls); fulls_staging = {}; } fulls.resize(nf); @@ -7122,5 +7161,14 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool // its CC1/2 is the answer without a second merge. A caller that did not ask keeps the NaN. if (measure_cc_before_corrections) r.cc_half_before_corrections = cc_half_before_corrections.value_or(r.statistics.overall.cc_half); + // The fulls and the merge's working arrays are rebuilt by every arm before they are read, so they go + // with it: kept, they sat beside the next arm's combine and through everything the caller did between + // merges - on a large cell in P1, the fulls and about 330 bytes per ASU group, gigabytes. Faulting them + // in again costs each arm a fraction of a second. + ReleaseMergeScratch(); + Free(fulls); + Free(fulls_staging); + Free(fulls_group_start); + Free(fulls_group_perm); return r; } diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index b1aa6d658..fb3ad61cb 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -128,7 +128,10 @@ public: // experiment: read live (its space group is changed by the caller between Run() calls). // partial_outcomes: the per-frame partials; the final per-frame scale (G, CC, mosaicity) is written - // back onto them so the offline per-image scaling table is still exported. + // back onto them so the offline per-image scaling table is still exported. Where the engine scales on + // the host (ReadsReflections), it reads the reflections themselves from the ingest on, so they + // must stay where they are, as they are, for as long as it merges: not freed, moved or changed, + // but for their h, k and l, which it does not read after the ingest. // reference_cell: the consensus cell (for the completeness count and the cell-consistency mask). RotationScaleMerge(const DiffractionExperiment &experiment, std::vector &partial_outcomes, @@ -183,27 +186,31 @@ public: // 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); } + // Hands an array (or a struct of them) back to the allocator. `v = {}` would not: on a std::vector it + // is the initializer-list assignment, which empties the vector and keeps its memory. + template static void Free(T &v) { v = T(); } + // Hand back the merge's working arrays - MergeAndStats' accumulators, R_meas sums, ISa scatter, // error-model pools and per-full fields. Each is resized and refilled before it is read, so the // next merge starts from them as an engine's first merge does, and gives the same answer. On a // large cell in P1 they are gigabytes, kept between merges only to save faulting them in again; // this is for a gap between merges where something else sits on the run's memory peak. void ReleaseMergeScratch() { - merge_fields = {}; - merge_acc = {}; - merge_accum = {}; - merge_merged_I = {}; - merge_rejected = {}; - merge_half_set = {}; - merge_obs_key = {}; - merge_wobs = {}; - merge_rmeas = {}; - merge_rmeas_staging = {}; - asymptote_scatter = {}; - em_samples = {}; - em_cut_pool = {}; - em_fit_pool = {}; - em_chi2 = {}; + Free(merge_fields); + Free(merge_acc); + Free(merge_accum); + Free(merge_merged_I); + Free(merge_rejected); + Free(merge_half_set); + Free(merge_obs_key); + Free(merge_wobs); + Free(merge_rmeas); + Free(merge_rmeas_staging); + Free(asymptote_scatter); + Free(em_samples); + Free(em_cut_pool); + Free(em_fit_pool); + Free(em_chi2); } // For an engine that will not merge again: everything on the host but the per-frame state @@ -211,21 +218,21 @@ public: // last_cc_n, last_frame_scaled). Its card memory goes with the engine. void ReleaseAllButPerFrameScale() { ReleaseMergeScratch(); - partials = {}; - fulls = {}; - fulls_staging = {}; - fulls_group_start = {}; - fulls_group_perm = {}; - perm = {}; - rawrun_start = {}; - rawrun_count = {}; - rawrun_h = {}; - rawrun_k = {}; - rawrun_l = {}; - rawrun_d = {}; - rawrun_group = {}; - finite_ok = {}; - corr_ingested = {}; + Free(partials); + Free(fulls); + Free(fulls_staging); + Free(fulls_group_start); + Free(fulls_group_perm); + Free(perm); + Free(rawrun_start); + Free(rawrun_count); + Free(rawrun_h); + Free(rawrun_k); + Free(rawrun_l); + Free(rawrun_d); + Free(rawrun_group); + Free(finite_ok); + Free(corr_ingested); obs_position.reset(); } @@ -247,6 +254,15 @@ public: [[nodiscard]] size_t DeviceBytes() const; // The device scaling's estimated peak (see the .cpp). Call after the ingest. [[nodiscard]] size_t GPUScalingPeakBytes() const; + + // The host memory an engine holds, for a caller deciding what may run beside it within the host budget + // (HostMemory.h); it decides nothing in the engine. HostIngestBytes: what the ingest holds from then on, + // between merges - call after IngestDevice(). HostMergeBytes: what one merge holds on top of it at its + // peak, from the counts the ingest knows (see the .cpp). HostBytesBound: the two together, before an + // ingest of `n_reflections`, at most. + [[nodiscard]] size_t HostIngestBytes() const; + [[nodiscard]] size_t HostMergeBytes() const; + [[nodiscard]] static size_t HostBytesBound(size_t n_reflections, bool on_gpu); // Above this fraction of the card's memory the scaling runs on the CPU. static constexpr double CPU_SCALING_ABOVE_CARD_FRACTION = 0.8; @@ -259,6 +275,9 @@ public: [[nodiscard]] int GetGpuDevice() const { return gpu_device; } // Whether this engine scales on the GPU (decided by the ingest). [[nodiscard]] bool OnGPU() const; + // Whether the engine still reads the source reflections in later Runs (see the constructor): where it + // scales on the host, or keeps its partials for the observation dump. + [[nodiscard]] bool ReadsReflections() const { return !partials.empty(); } // Whether Run() also hands back the merge's fulls unmerged (Result::scaled_fulls). Off by default: // it is a copy as long as the fulls, wanted only by a caller that writes them. @@ -296,10 +315,10 @@ private: static constexpr uint8_t OBS_CLIPPED = 1; static constexpr uint8_t OBS_OVERLOADED = 2; - // One integrated observation - a per-frame partial during scaling/combine, or a combined full during - // scale-fulls/merge. Flat (not nested per image); a POD so the arrays translate straight to CUDA. - // No member initialisers, so that `partials` and `fulls` (FillVector) are first written by the - // parallel passes that fill them: whoever makes one sets every field. + // One combined full during scale-fulls/merge (a per-frame partial is a PartialObs, below). Flat (not + // nested per image); a POD so the arrays translate straight to CUDA. No member initialisers - here or in + // PartialObs - so that `fulls` and `partials` (FillVector) are first written by the parallel passes + // that fill them: whoever makes one sets every field. struct Obs { int32_t h, k, l; float I, sigma, d, prescaling_corr, partiality, zeta, delta_phi, bkg, var_bkg; @@ -321,6 +340,24 @@ private: int32_t group; // dense ASU-group id for the current space group; <0 = never mergeable }; + // One per-frame partial as the host holds it - without a GPU, or for the observation dump: where its + // reflection is, and the four values the ingest works out for it or the scaling changes. Every other + // field is the reflection's own, read from it: the run keeps the reflections for as long as the engine + // merges, and nothing changes in them after the images are integrated but the Miller indices of a + // reindexing, which only the ingest reads (from the reflections, into the raw-hkl runs). Twenty-four + // bytes where a copy of the fields was eighty-four - on a long axis, a hundred million partials. + struct PartialObs { + const Reflection *r; + int32_t frame; // index of the outcome whose per-frame scale G applies to this obs + float prescaling_corr; // the reflection's whole deterministic correction over the frame's flux + float partiality; // recomputed from the frame-order-smoothed mosaicity (Ingest) + float corr; // image_scale_corr (working; updated by scaling) + }; + // A reflection's Obs::clipped. + static uint8_t ClippedFlags(const Reflection &r) { + return (r.clipped ? OBS_CLIPPED : 0) | (r.overloaded ? OBS_OVERLOADED : 0); + } + // One leverage-corrected error-model sample per usable full: its counting variance at the group's // mean, that mean squared, its squared deviation from the mean - and the resolution it sits at, // because the fit is re-run over the samples that survive the automatic resolution cutoff. See @@ -414,7 +451,7 @@ private: // built only when the device pipeline is NOT resident (no GPU, or an observation dump was asked // for); on the resident path (resident_ingest) the device holds the observations and the host // keeps only the ingest arrays below, until the end of Ingest. - FillVector partials; // all per-frame partials, grouped by frame + FillVector partials; // all per-frame partials, grouped by frame bool resident_ingest = false; // What the ingest-time geometry smoothing (SmoothMosaicityAndPartiality) needs, one array per field // in `partials` order: the two fields it rewrites, the two the rocking-event walks read in raw-hkl @@ -696,13 +733,13 @@ private: std::vector partial_group_start, partial_group_count; NotCopied partial_group_perm; // Without a GPU: each partial's ASU group (-1 = in none), in observation order - where the CPU - // passes read a partial's group from. A partial's own Obs::group is not stamped: written into the - // eighty-byte records at random, once per grouping, it cost more than the grouping itself. + // passes read a partial's group from. It is not kept on the partial: written into the records at + // random, once per grouping, it cost more than the grouping itself. NotCopied partial_group; // The CPU scaling loops' working copy of the partials, or of the fulls. A loop passes over them // several times per iteration and reads a handful of fields each time, so it works on these rather - // than on the eighty-byte Obs. ScalingRef holds what the group means and the corr update read, in the CSR's + // than on the partials and their reflections. ScalingRef holds what the group means and the corr update read, in the CSR's // group order with the partials in no group after them, and corr is kept beside it in the same // order - so neither pass gathers. ScalingFit holds, frame by frame, the observations // FitScalingFrames fits: the fit's filters, which do not change over the loop, applied once, diff --git a/image_analysis/structure_refinement/ModelGrid.cpp b/image_analysis/structure_refinement/ModelGrid.cpp index 2dde37fe2..c10ea95fd 100644 --- a/image_analysis/structure_refinement/ModelGrid.cpp +++ b/image_analysis/structure_refinement/ModelGrid.cpp @@ -317,3 +317,10 @@ void PutMaskOnGrid(gemmi::Grid &grid, const gemmi::Model &model, const st if (shrink_has_offsets) masker.shrink(grid); } + +size_t ModelGridPoints(const gemmi::UnitCell &cell, const gemmi::SpaceGroup *sg, double d_min, double rate) { + const double spacing = d_min / (2 * rate); + const std::array m = gemmi::good_grid_size({{cell.a / spacing, cell.b / spacing, cell.c / spacing}}, + gemmi::GridSizeRounding::Up, sg); + return static_cast(m[0]) * static_cast(m[1]) * static_cast(m[2]); +} diff --git a/image_analysis/structure_refinement/ModelGrid.h b/image_analysis/structure_refinement/ModelGrid.h index 4998f2226..96f647292 100644 --- a/image_analysis/structure_refinement/ModelGrid.h +++ b/image_analysis/structure_refinement/ModelGrid.h @@ -27,3 +27,12 @@ void PutModelDensityOnGrid(gemmi::DensityCalculator, float> & // `grid` sized, with its cell and group set. void PutMaskOnGrid(gemmi::Grid &grid, const gemmi::Model &model, const std::vector &orbit_leaders, size_t nthreads); + +// The number of points of the grid a model's structure factors to d_min are computed on at DensityCalculator's +// `rate` (a spacing of d_min / (2 rate)) in this cell and space group: the grid's size, without making it. +size_t ModelGridPoints(const gemmi::UnitCell &cell, const gemmi::SpaceGroup *sg, double d_min, double rate = 1.5); + +// What computing a model's structure factors on the CPU holds at its peak, per point of that grid: the density +// and the bulk-solvent mask grids, and for each of them its transform's input copy, spectrum and coefficients, +// with the two made at once (ModelValidation's compute_model_factors). +constexpr size_t MODEL_FACTORS_BYTES_PER_GRID_POINT = 32; diff --git a/image_analysis/structure_refinement/ModelValidation.cpp b/image_analysis/structure_refinement/ModelValidation.cpp index 5556c7532..65794338a 100644 --- a/image_analysis/structure_refinement/ModelValidation.cpp +++ b/image_analysis/structure_refinement/ModelValidation.cpp @@ -35,6 +35,7 @@ #include "../../common/CorrelationCoefficient.h" #include "../../common/JFJochMath.h" // PI (M_PI is not standard, and MSVC does not define it) +#include "../../common/HostMemory.h" #include "../../common/Logger.h" #include "../../common/ParallelFor.h" #include "../scale_merge/ReindexAmbiguity.h" // ReindexReflections @@ -1247,7 +1248,17 @@ ModelValidationResult Validate(const std::vector &merged, // the order the concurrent replicates finish in. std::vector null_r_work(NULL_REPLICATES), null_margin(NULL_REPLICATES); std::atomic placed_redraws{0}; + // A replicate whose structure factors are made on the CPU holds a gigabyte or more of grids on a large + // cell, so it reserves them against the host's worker pool (HostWorkerPool) first: as many replicates + // run at once as it holds, whatever the thread count, and each computes what it computes either way. + size_t replicate_bytes = MODEL_FACTORS_BYTES_PER_GRID_POINT + * ModelGridPoints(null_model.cell, null_model.find_spacegroup(), null_d_min); +#ifdef JFJOCH_USE_CUDA + if ((sf_gpu != nullptr && null_d_min == sf_gpu->DMin()) || (sf_null != nullptr && null_d_min == sf_null->DMin())) + replicate_bytes = 0; +#endif auto run_replicate = [&](int i) { + const HostWorkerPool::Reservation reserved(HostWorkerPool::Instance(), replicate_bytes); std::mt19937 rng(NULL_SEED + 1 + i); gemmi::Mat33 rot = null_rotation[i]; for (int redraw = 0;; redraw++) { diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index 9f7adaf6c..442339577 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/image_analysis/structure_refinement/MolecularReplacement.cpp @@ -33,6 +33,8 @@ #include "../../common/JFJochMath.h" // PI #include "../../common/Logger.h" +#include "../../common/CUDAWrapper.h" // get_gpu_count +#include "../../common/HostMemory.h" #include "../../common/ParallelFor.h" #include "ModelFFT.h" #include "ModelGrid.h" @@ -45,7 +47,6 @@ #include "RigidBodyGPUEngine.h" #include "ModelStructureFactorsGPU.h" #include "ModelScaleGPU.h" -#include "../../common/CUDAWrapper.h" #include "../../common/JFJochException.h" #include "../indexing/CUDAMemHelpers.h" // cuda_trim_allocation_pool #endif @@ -766,11 +767,18 @@ double null_sigma(const gemmi::Structure &placed, int copies, const gemmi::Mat33 where = translate(turned); } std::vector r_work(NULL_REPLICATES); + // Without a GPU each replicate's structure factors are made on the CPU, a gigabyte or more of grids on a + // large cell, so it reserves them against the host's worker pool (HostWorkerPool) first: as many run at + // once as it holds - with the nulls of other candidates beside it - and each computes the same either way. + const size_t replicate_bytes = get_gpu_count() == 0 + ? MODEL_FACTORS_BYTES_PER_GRID_POINT * ModelGridPoints(placed.cell, placed.find_spacegroup(), null_d_min) + : 0; std::atomic next{0}; std::vector> workers; for (size_t t = 0; t < std::min(std::max(nthreads, 1), NULL_REPLICATES); t++) workers.push_back(std::async(std::launch::async, [&] { for (int i = next++; i < NULL_REPLICATES; i = next++) { + const HostWorkerPool::Reservation reserved(HostWorkerPool::Instance(), replicate_bytes); gemmi::Structure rep = placed; std::vector &chains = rep.models[0].chains; for (int k = 0; k < copies; k++) { @@ -1561,7 +1569,16 @@ std::vector TranslationFunction(const MolecularTransform &ft, c const std::vector> *known, size_t nthreads) { const TranslationGather gather = make_gather(setup, known != nullptr); std::vector out(rotations.size()); - ParallelFor(static_cast(rotations.size()), nthreads, [&](int i) { + // A rotation's buffers - the grid's coefficients and the function, its terms' factors - are hundreds of + // megabytes on a large cell, so as many rotations run at once as the host's worker pool holds + // (HostWorkerPool), at most nthreads. Counted here rather than reserved in each: the rotations run on + // ParallelFor's pool, whose workers must not wait on one another's memory. + const size_t nu = setup.nu, nv = setup.nv, nw = setup.nw; + const size_t bytes = nu * nv * (nw / 2 + 1) * sizeof(std::complex) + nu * nv * nw * sizeof(float) + + setup.Terms() * (setup.ops * sizeof(std::complex) + sizeof(double)); + const size_t workers = std::clamp(HostWorkerPool::Instance().Size() / std::max(bytes, 1), 1, + std::max(nthreads, 1)); + ParallelFor(static_cast(rotations.size()), workers, [&](int i) { out[i] = translation_function(ft, setup, gather, rotations[i], known); }); return out; @@ -2793,8 +2810,15 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec to_refine.push_back(&f); if (!to_refine.empty()) { std::vector> running; + // Without a GPU each placement's structure factors and rigid body are made on the CPU, a + // gigabyte or more of grids on a large cell, so each reserves them against the host's worker + // pool first (HostWorkerPool), as a null replicate does. + const size_t placement_bytes = get_gpu_count() == 0 + ? MODEL_FACTORS_BYTES_PER_GRID_POINT * ModelGridPoints(cell, &sg, d_min) + : 0; for (First *f : to_refine) running.push_back(std::async(std::launch::async, [&, f] { + const HostWorkerPool::Reservation reserved(HostWorkerPool::Instance(), placement_bytes); f->score = refine_placement(*f->placed, sg, *merge.merged, d_min, engines, nthreads, &scorers, &f->moved); f->clash = ClashFraction(*f->placed); diff --git a/image_analysis/structure_refinement/MolecularReplacement.h b/image_analysis/structure_refinement/MolecularReplacement.h index 31b6094c3..1c659212c 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.h +++ b/image_analysis/structure_refinement/MolecularReplacement.h @@ -181,7 +181,8 @@ struct TranslationPeak { TranslationPeak TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup, const gemmi::Mat33 &rotation, const std::vector> *known = nullptr); -// The same at each of `rotations`, `nthreads` of them at once, the grid's gather worked out once for all. +// The same at each of `rotations`, `nthreads` of them at once - fewer where their buffers do not all fit the host's +// worker pool (HostWorkerPool) - the grid's gather worked out once for all. std::vector TranslationFunction(const MolecularTransform &ft, const TranslationSetup &setup, const std::vector &rotations, const std::vector> *known, size_t nthreads); diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index 5a97c214a..37f7c6484 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -32,6 +32,7 @@ #include "../reader/JFJochHDF5Reader.h" #include "../reader/JFJochStartMessageReader.h" +#include "../common/HostMemory.h" #include "../common/JFJochMath.h" #include "../common/ParallelFor.h" #include "../common/Logger.h" @@ -317,6 +318,20 @@ Rugnux::Rugnux(JFJochReader &reader, DiffractionExperiment experiment, LogHeaderBeamCenter(experiment_, pixel_mask_); + // What the run plans its host memory against, decided once from the machine's total (HostMemory.h): + // work that would not fit it side by side runs one part after the other, and the large per-worker + // buffers of the model's steps are held a share of it at a time. Only when things run changes. + const size_t machine = get_host_total_memory(); + host_budget_ = config_.host_memory_budget_bytes > 0 ? config_.host_memory_budget_bytes : DefaultHostMemoryBudget(); + HostWorkerPool::Instance().SetSize( + std::max(1, static_cast(HOST_WORKER_POOL_FRACTION * static_cast(host_budget_)))); + Logger("Rugnux").Info("Host memory budget {:.1f} GB ({}): what does not fit it side by side runs in series", + static_cast(host_budget_) / 1e9, + config_.host_memory_budget_bytes > 0 + ? std::string("set by the caller") + : fmt::format("{:.0f}% of the machine's {:.1f} GB", 100.0 * HOST_MEMORY_BUDGET_FRACTION, + static_cast(machine) / 1e9)); + // Every decision on a rotation run - per-frame scaling, the correction surfaces, the outlier test, // the error model, the CC1/2 half-sets and the resolution cut, the frame ledger, the space-group // search - is taken with Friedel mates merged, -A or not. Splitting them halves the multiplicity diff --git a/rugnux/Rugnux.h b/rugnux/Rugnux.h index 6d434c98e..2eea2d4cb 100644 --- a/rugnux/Rugnux.h +++ b/rugnux/Rugnux.h @@ -73,6 +73,11 @@ struct ProcessConfig { // The per-image loop is capped in that case (see RunPipeline); an explicit -N never is. bool nthreads_auto = false; + // The host memory the run plans against, in bytes; 0 (the default) takes the machine's + // (DefaultHostMemoryBudget). It decides only what runs side by side - never a number - so a caller + // setting it (the tests do) changes the run's time and its memory, and nothing it writes. + size_t host_memory_budget_bytes = 0; + // Output prefix for the _process.h5 (and scaled reflections). Empty => process without writing. std::string output_prefix; // What the reflection files record about where they came from (WriteReflections). @@ -800,6 +805,8 @@ class Rugnux { DiffractionExperiment experiment_; PixelMask pixel_mask_; ProcessConfig config_; + // The host memory budget (ProcessConfig::host_memory_budget_bytes, else the machine's), decided once. + size_t host_budget_ = 0; // The run's cancel flag, shared with the copies the run makes of itself to run passes side by // side (see RunAllPasses), so cancelling the run cancels those too. class CancelFlag { diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index 6ec378d7e..604d5f2cb 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -839,14 +839,17 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // the scaling engine is finished - nothing after this point reads it on that pass // (the search and the reference re-merge run only on the canonical one) - so it is // freed BEFORE the fit: the engine holds about twice the partial payload again, - // and with it gone the fit no longer sits on the pass's memory peak. - if (prepass_postrefine_obs) + // and with it gone the fit no longer sits on the pass's memory peak. An engine that + // read the reflections in its merges left the gather to here, and from here nothing + // on the pass reads them either: each image's is handed back as it is gathered. + const bool release = geometry_prepass && rsm && config_.reference_data.empty(); + if (prepass_postrefine_obs || release) rsm.reset(); auto pr_obs = prepass_postrefine_obs ? std::move(*prepass_postrefine_obs) : gather ? GatherPostRefineObservations(indexer->GetIntegrationOutcome(), config_.nthreads, - /*release_reflections=*/false, logger, + /*release_reflections=*/release, logger, max_images) : PostRefineObservations{}; prepass_postrefine_obs.reset(); @@ -945,13 +948,41 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { /*full_stats=*/false, /*measure_cc_before_corrections=*/true); return ScaleMergeResult{std::move(r.merged), std::move(r.statistics), r.cc_half_before_corrections}; }; + // What the host holds from here on (HostMemory.h): the reflections, and beside them the engines below. + // The counts are known now, before anything is ingested, so each concurrency below is decided on them: + // where two pieces of work would not fit the host budget together, one waits for the other. None of + // them reads anything another writes while it runs (see each), so only the order changes. + size_t n_reflections = 0, n_reflections_guard = 0; + { + const auto &outcomes = indexer->GetIntegrationOutcome(); + for (size_t o = 0; o < outcomes.size(); ++o) { + n_reflections += outcomes[o].reflections.size(); + if (static_cast(o) < prepass_images) + n_reflections_guard += outcomes[o].reflections.size(); + } + } + const size_t reflection_bytes = n_reflections * sizeof(Reflection); + const auto host_gb = [](size_t bytes) { return static_cast(bytes) / 1e9; }; std::future guard_ahead; - if (merge_first_part && !(postrefine_probe_ && !geometry_prepass && postrefine_probe_only_)) - guard_ahead = std::async(std::launch::async, [&] { - if (guard_on_second_gpu) - pin_gpu(1); // the card and its NUMA node, as an image-loop worker has them - return merge_guard(); - }); + if (merge_first_part && !(postrefine_probe_ && !geometry_prepass && postrefine_probe_only_)) { + // The guard's merge runs beside the main engine's ingest. Where an engine of the guard's frames and + // one of all the frames would not fit the budget together - counted at the most an engine of that + // many reflections can hold - the guard merges first, on its own, and the ingest starts after it. + const bool on_gpu = get_gpu_count() > 0; + const size_t beside = reflection_bytes + RotationScaleMerge::HostBytesBound(n_reflections_guard, on_gpu) + + RotationScaleMerge::HostBytesBound(n_reflections, on_gpu); + if (beside <= host_budget_ || guard_on_second_gpu) + guard_ahead = std::async(std::launch::async, [&] { + if (guard_on_second_gpu) + pin_gpu(1); // the card and its NUMA node, as an image-loop worker has them + return merge_guard(); + }); + else { + logger.Info("Host memory: the quality guard's merge and this pass's ingest need up to {:.1f} GB side by " + "side, over the {:.1f} GB budget - the guard merges first", host_gb(beside), host_gb(host_budget_)); + guard_merge = merge_guard(); + } + } // The scaling engine's ingest reads the outcomes and nothing else either, so its host half runs // beside the post-refinement too; with one GPU the device half waits for the guard to hand the // card back. @@ -1093,6 +1124,19 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { std::unique_ptr rsm_all; Logger all_arm_log = Logger::Buffered(); std::future all_arm_ahead; + const auto start_all_arm = [&rsm_all, &all_arm_log] { + return std::async(std::launch::async, [&rsm_all, &all_arm_log] { + pin_gpu(rsm_all->GetGpuDevice()); // the card and its NUMA node, as the guard's merge has them + TimingMark mark(all_arm_log, "merge P1, all observations"); + auto r = rsm_all->Run(/*for_search=*/true, /*full_stats=*/true, + /*measure_cc_before_corrections=*/false); + // This engine merges once: what is left for it to do is the write-back of its + // per-frame scales, so its observations and its merge's working arrays go now, + // rather than through the search the others are about to run. + rsm_all->ReleaseAllButPerFrameScale(); + return r; + }); + }; if (is_rotation) { if (experiment_.GetRefineRotationWedgeInScaling() || rot_ss.GetRotationWedgeForScaling().has_value()) @@ -1147,7 +1191,8 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { } // With one card the guard's merge has the card to itself first. if (merge_first_part && !guard_on_second_gpu) { - guard_merge = guard_ahead.valid() ? guard_ahead.get() : merge_guard(); + if (!guard_merge) + guard_merge = guard_ahead.valid() ? guard_ahead.get() : merge_guard(); guard_log.ReplayInto(logger); } rsm->IngestDevice(); @@ -1178,27 +1223,43 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { release_gather(); engine_log.ReplayInto(logger); rsm->SetLogger(logger); + // The canonical pass's post-refinement gathers (from the moment the engine has ingested) and fits + // beside the merges that follow. Where its gather and the selection it ranks would not fit the host + // budget beside the reflections, this engine and its first merge, it is done before that merge. + const size_t host_engine = reflection_bytes + rsm->HostIngestBytes(); + const size_t host_merge = rsm->HostMergeBytes(); + constexpr size_t POSTREFINE_RANKED_BYTES = 16; // PostRefineRotationGeometry's ranking, per partial + const size_t host_postrefine = n_reflections * (sizeof(PostRefinePartial) + POSTREFINE_RANKED_BYTES); + bool postrefine_beside = p.postrefine_ahead.valid(); + if (postrefine_beside && host_engine + host_merge + host_postrefine > host_budget_) { + logger.Info("Host memory: the post-refinement ({:.1f} GB) does not fit beside the merge ({:.1f} GB, " + "{:.1f} GB held) within the {:.1f} GB budget - it is done before the merge", + host_gb(host_postrefine), host_gb(host_merge), host_gb(host_engine), host_gb(host_budget_)); + p.postrefine_ahead.wait(); + postrefine_beside = false; + } + bool all_arm_beside = false; if (all_arm) { rsm_all = rsm->CloneIngest(all_arm_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); + // ... and beside this engine's merges only where two merges fit the host budget as well, + // with the post-refinement where it is still running; otherwise it merges where the search + // first reads it, after this engine's merges, on its own. + const size_t host_beside = host_engine + 2 * host_merge + (postrefine_beside ? host_postrefine : 0); + all_arm_beside = host_beside <= host_budget_; + if (!all_arm_beside) + logger.Info("Host memory: two P1 merges side by side need about {:.1f} GB, over the {:.1f} GB " + "budget - the all-observation merge is made after this engine's", + host_gb(host_beside), host_gb(host_budget_)); } if (rsm_all) { // Its observations are copied from this engine's card arrays, not staged from the // reflections a second time. rsm_all->IngestDevice(&*rsm); - all_arm_ahead = std::async(std::launch::async, [&rsm_all, &all_arm_log] { - pin_gpu(rsm_all->GetGpuDevice()); // the card and its NUMA node, as the guard's merge has them - TimingMark mark(all_arm_log, "merge P1, all observations"); - auto r = rsm_all->Run(/*for_search=*/true, /*full_stats=*/true, - /*measure_cc_before_corrections=*/false); - // This engine merges once: what is left for it to do is the write-back of its - // per-frame scales, so its observations and its merge's working arrays go now, - // rather than through the search the others are about to run. - rsm_all->ReleaseAllButPerFrameScale(); - return r; - }); + if (all_arm_beside) + all_arm_ahead = start_all_arm(); } } @@ -1209,8 +1270,10 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // per-frame fields - so take the post-refine's view NOW and hand each image's vector back as // it is read: on a fine-sliced long axis the vectors are gigabytes, held to the end of the // pass for those 32 bytes a partial. A run with a reference MTZ keeps them: resolving the - // indexing ambiguity re-reads them (and re-ingests). - if (geometry_prepass && rsm && config_.reference_data.empty()) + // indexing ambiguity re-reads them (and re-ingests). So does an engine that scales on the host, + // which reads them in its merges (RotationScaleMerge::ReadsReflections): there the view is taken + // once the engine is done with them (prepare_post_refine). + if (geometry_prepass && rsm && config_.reference_data.empty() && !rsm->ReadsReflections()) prepass_postrefine_obs = GatherPostRefineObservations( indexer->GetIntegrationOutcome(), config_.nthreads, /*release_reflections=*/true, logger); @@ -1640,6 +1703,8 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { // 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 (rsm_all && !all_arm_ahead.valid()) + all_arm_ahead = start_all_arm(); // not beside the merges above (the host budget): now if (all_arm_ahead.valid()) { TimingMark wait_mark(logger, "wait for the all-observation merge"); try { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c6dd46391..e5fb8f6bc 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -102,6 +102,7 @@ ADD_EXECUTABLE(jfjoch_test HotPixelFinderTest.cpp BandwidthEstimateTest.cpp ParallelForTest.cpp + HostMemoryTest.cpp RugnuxTest.cpp ResultReportTest.cpp DiagnosticOutputTest.cpp diff --git a/tests/HostMemoryTest.cpp b/tests/HostMemoryTest.cpp new file mode 100644 index 000000000..e58d17ebd --- /dev/null +++ b/tests/HostMemoryTest.cpp @@ -0,0 +1,48 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include + +#include +#include +#include +#include +#include +#include + +#include "../common/HostMemory.h" + +TEST_CASE("HostMemory_BudgetFromTheTotal", "[HostMemory]") { + const size_t total = get_host_total_memory(); + REQUIRE(total > 0); + CHECK(DefaultHostMemoryBudget() == static_cast(HOST_MEMORY_BUDGET_FRACTION * static_cast(total))); +} + +// As many reservations at once as the pool holds, and one at a time where a single one is larger than it. +TEST_CASE("HostWorkerPool_BoundsWorkersByBytes", "[HostMemory]") { + HostWorkerPool &pool = HostWorkerPool::Instance(); + for (const auto &[size, bytes, at_most] : std::vector>{ + {100, 40, 2}, {100, 100, 1}, {100, 1000, 1}, {1000, 10, 8}}) { + CAPTURE(size, bytes); + pool.SetSize(size); + std::atomic running{0}, most{0}, done{0}; + std::vector workers; + for (int i = 0; i < 8; i++) + workers.emplace_back([&] { + const HostWorkerPool::Reservation reserved(pool, bytes); + const int now = ++running; + int seen = most.load(); + while (now > seen && !most.compare_exchange_weak(seen, now)) {} + std::this_thread::sleep_for(std::chrono::milliseconds(20)); + --running; + ++done; + }); + for (auto &w : workers) + w.join(); + CHECK(done == 8); + CHECK(most >= 1); + CHECK(most <= at_most); + } + pool.SetSize(0); + CHECK(pool.Size() == static_cast(HOST_WORKER_POOL_FRACTION * static_cast(DefaultHostMemoryBudget()))); +} diff --git a/tests/RugnuxLargeTest.cpp b/tests/RugnuxLargeTest.cpp index 19ba12005..38a6de228 100644 --- a/tests/RugnuxLargeTest.cpp +++ b/tests/RugnuxLargeTest.cpp @@ -9,11 +9,15 @@ #include #include +#include +#include #include +#include #include #include "TestData.h" #include "../common/DiffractionExperiment.h" +#include "../common/HostMemory.h" #include "../common/IndexingSettings.h" #include "../reader/JFJochHDF5Reader.h" #include "../rugnux/Rugnux.h" @@ -121,3 +125,63 @@ TEST_CASE("Rugnux_FirstPassMemoMatchesRecomputation", "[large]") { reader.Close(); } + +// The host memory budget decides only what runs side by side - the quality guard's merge beside the +// ingest, the post-refinement beside the merge chain, the two search merges beside each other, how many +// large per-worker buffers are held at once - and never what is computed. So a run planned against a +// budget nothing fits in (everything in series) writes what a run with room for everything writes, to +// the byte. +TEST_CASE("Rugnux_HostBudgetChangesOnlyTheSchedule", "[large]") { + const auto master = jfjoch_test::LargeDataFile("rotation_master.h5"); + if (!master) + SKIP("rotation_master.h5 not available (git-lfs data not pulled)"); + + RegisterHDF5Filter(); + const auto slurp = [](const std::string &path) { + std::ifstream f(path, std::ios::binary); + return std::string(std::istreambuf_iterator(f), std::istreambuf_iterator()); + }; + const auto run = [&](size_t budget, const std::string &prefix) { + JFJochHDF5Reader reader; + REQUIRE_NOTHROW(reader.ReadFile(*master)); + auto dataset = reader.GetDataset(); + REQUIRE(dataset); + DiffractionExperiment experiment(dataset->experiment); + IndexingSettings indexing; + indexing.Algorithm(IndexingAlgorithmEnum::Auto); + indexing.RotationIndexing(true); + experiment.ImportIndexingSettings(indexing); + + ProcessConfig config; + config.mode = ProcessMode::FullAnalysis; + config.nthreads = default_threads(); + config.spot_finding = DiffractionExperiment::DefaultDataProcessingSettings(); + config.spot_finding.indexing = true; + config.rotation_indexing = true; + config.two_pass_rotation = true; + config.rotation_postrefine_geometry = true; + config.run_scaling = true; + config.write_process_h5 = false; + config.output_prefix = prefix; + config.host_memory_budget_bytes = budget; + + Rugnux process(reader, experiment, *dataset->pixel_mask, config); + ProcessResult result; + REQUIRE_NOTHROW(result = process.Run()); + reader.Close(); + return result; + }; + const std::filesystem::path dir = std::filesystem::temp_directory_path() / "rugnux_host_budget_test"; + std::filesystem::create_directories(dir); + const ProcessResult roomy = run(size_t(1) << 40, (dir / "roomy").string()); + const ProcessResult tight = run(size_t(1) << 20, (dir / "tight").string()); + + REQUIRE(roomy.has_merge_statistics); + CHECK(tight.merge_statistics_text == roomy.merge_statistics_text); + CHECK(tight.merged_i_sigma == roomy.merged_i_sigma); + const std::string hkl = slurp((dir / "roomy.hkl").string()); + CHECK_FALSE(hkl.empty()); + CHECK(slurp((dir / "tight.hkl").string()) == hkl); + std::filesystem::remove_all(dir); + HostWorkerPool::Instance().SetSize(0); // process-wide: back to the machine's for the tests after this one +}