diff --git a/image_analysis/IndexAndRefine.cpp b/image_analysis/IndexAndRefine.cpp index 6d48a4b82..b21475948 100644 --- a/image_analysis/IndexAndRefine.cpp +++ b/image_analysis/IndexAndRefine.cpp @@ -586,18 +586,24 @@ void IndexAndRefine::QuickPredictAndIntegrate(DataMessage &msg, // caller's callback binds it to the right image (GPU-resident buffer, host buffer, or the assembled // FPGA image read straight on the CPU). auto integration_start_time = std::chrono::steady_clock::now(); - i_outcome.reflections = integrate(prediction.GetReflections(), nrefl, msg.number); - msg.integrated_reflections = i_outcome.reflections.size(); + const std::vector integrated = integrate(prediction.GetReflections(), nrefl, msg.number); + msg.integrated_reflections = integrated.size(); auto integration_end_time = std::chrono::steady_clock::now(); msg.integration_time_s = std::chrono::duration(integration_end_time - integration_start_time).count(); - CalcISigmaAndWilsonBFactor(msg, i_outcome.reflections); + CalcISigmaAndWilsonBFactor(msg, integrated); + + // A retained outcome is kept in the pass's arena (see ReflectionArena), not beside the worker's + // other allocations. + i_outcome.reflections = ReflectionVector(integrated.begin(), integrated.end(), + ArenaAllocator(retain_outcomes_ ? &reflection_arena + : nullptr)); ScaleImage(msg, i_outcome); // Copy reflections to outgoing message if (keep_reflections_in_message_) - msg.reflections = i_outcome.reflections; + msg.reflections.assign(i_outcome.reflections.begin(), i_outcome.reflections.end()); // Persist the per-image result for the whole-run scaling/merge pass, unless the caller opted out // (viewer interactive use only needs the current image, returned above via msg). diff --git a/image_analysis/IndexAndRefine.h b/image_analysis/IndexAndRefine.h index 0fb0d1fdd..56904f78e 100644 --- a/image_analysis/IndexAndRefine.h +++ b/image_analysis/IndexAndRefine.h @@ -62,6 +62,8 @@ class IndexAndRefine { }; mutable std::mutex reflections_mutex; + // Holds the retained outcomes' reflections; declared before them, so it outlives them. + ReflectionArena reflection_arena; std::vector integration_outcome; std::vector mosaicity; // Optional per-frame mosaicity used for Bragg prediction, indexed by image number. When set (the diff --git a/image_analysis/IntegrationOutcome.h b/image_analysis/IntegrationOutcome.h index c0d8855fa..a16f41957 100644 --- a/image_analysis/IntegrationOutcome.h +++ b/image_analysis/IntegrationOutcome.h @@ -8,11 +8,16 @@ #include "../common/Reflection.h" #include "../common/CrystalLattice.h" #include "../common/DiffractionGeometry.h" +#include "ReflectionArena.h" + +// A frame's integrated reflections. The whole-run passes keep them in a ReflectionArena (see there); +// anywhere else the allocator is plain new/delete. +using ReflectionVector = std::vector>; struct IntegrationOutcome { DiffractionGeometry geom; CrystalLattice latt; - std::vector reflections; + ReflectionVector reflections; std::optional mosaicity_deg; std::optional image_scale_cc; std::optional image_scale_cc_n; diff --git a/image_analysis/ReflectionArena.h b/image_analysis/ReflectionArena.h new file mode 100644 index 000000000..193dc71b7 --- /dev/null +++ b/image_analysis/ReflectionArena.h @@ -0,0 +1,113 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include +#include +#include +#include + +// Large blocks that a whole pass's per-image reflection vectors are carved from. +// +// A rotation pass keeps every frame's integrated reflections until its scaling is done - thousands of +// vectors of a few megabytes each, gigabytes together, allocated by the image workers in the +// allocator's per-thread arenas. When the pass hands them back, that memory mostly stays in those +// arenas, as holes between whatever else the workers allocated, and the next pass's workers, being +// new threads, do not reuse it: on a fine-sliced long axis gigabytes of freed reflections were carried +// to the end of the run. Here a block is allocated in one piece (large enough that it is its own +// mapping) and returned in one piece once the last vector in it is gone, so a pass's reflections +// leave nothing behind. +// +// A vector is carved from the current block by bumping a pointer; the space of a vector freed early +// is only reclaimed with its whole block. That suits the use here: the vectors are written once, +// when the frame is integrated, and handed back together. +class ReflectionArena { +public: + ReflectionArena() = default; + ReflectionArena(const ReflectionArena &) = delete; + ReflectionArena &operator=(const ReflectionArena &) = delete; + ~ReflectionArena() { + for (auto &b : blocks) + ::operator delete(b.data); + } + + void *Allocate(size_t bytes) { + bytes = (bytes + kAlign - 1) / kAlign * kAlign; + std::unique_lock ul(m); + if (blocks.empty() || blocks.back().size - blocks.back().used < bytes) { + const size_t size = std::max(kBlockBytes, bytes); + blocks.push_back(Block{static_cast(::operator new(size)), size, 0, 0}); + } + Block &b = blocks.back(); + void *p = b.data + b.used; + b.used += bytes; + b.live++; + return p; + } + + void Deallocate(void *p) { + std::unique_lock ul(m); + for (size_t i = 0; i < blocks.size(); ++i) { + Block &b = blocks[i]; + if (static_cast(p) < b.data || static_cast(p) >= b.data + b.size) + continue; + if (--b.live == 0) { + if (i + 1 == blocks.size()) { + b.used = 0; // the block being filled: keep it for the next vector + } else { + ::operator delete(b.data); + blocks.erase(blocks.begin() + static_cast(i)); + } + } + return; + } + } + +private: + // 64 MiB: above the largest size glibc ever serves from its heaps (32 MiB), so every block is a + // mapping of its own and goes back to the system when freed. A frame's reflections are a few + // megabytes, so the tail a block cannot fit is a few percent of it. + static constexpr size_t kBlockBytes = size_t(64) << 20; + static constexpr size_t kAlign = alignof(std::max_align_t); + struct Block { + char *data; + size_t size, used, live; + }; + std::mutex m; + std::vector blocks; +}; + +// A std::allocator that carves from a ReflectionArena, or is plain new/delete without one. A copy of +// a container is made on the heap (select_on_container_copy_construction), so only the containers the +// arena's owner fills itself live in the arena; a move carries the arena along with the storage. +template +struct ArenaAllocator { + using value_type = T; + using propagate_on_container_move_assignment = std::true_type; + using propagate_on_container_swap = std::true_type; + + ReflectionArena *arena = nullptr; + + ArenaAllocator() = default; + explicit ArenaAllocator(ReflectionArena *a) : arena(a) {} + template ArenaAllocator(const ArenaAllocator &o) : arena(o.arena) {} + + T *allocate(size_t n) { + if (arena) + return static_cast(arena->Allocate(n * sizeof(T))); + return static_cast(::operator new(n * sizeof(T))); + } + void deallocate(T *p, size_t) { + if (arena) + arena->Deallocate(p); + else + ::operator delete(p); + } + ArenaAllocator select_on_container_copy_construction() const { return ArenaAllocator(); } + + template bool operator==(const ArenaAllocator &o) const { return arena == o.arena; } + template bool operator!=(const ArenaAllocator &o) const { return arena != o.arena; } +}; diff --git a/image_analysis/geom_refinement/PostRefine.cpp b/image_analysis/geom_refinement/PostRefine.cpp index 65a6e905c..4ad21536e 100644 --- a/image_analysis/geom_refinement/PostRefine.cpp +++ b/image_analysis/geom_refinement/PostRefine.cpp @@ -225,7 +225,7 @@ PostRefineObservations GatherPostRefineObservations(std::vector().swap(outcomes[o].reflections); + ReflectionVector().swap(outcomes[o].reflections); } }); logger.Info("Post-refine: {} partials gathered", n_pts); diff --git a/image_analysis/scale_merge/Merge.cpp b/image_analysis/scale_merge/Merge.cpp index 2945ce2f7..dea9f923f 100644 --- a/image_analysis/scale_merge/Merge.cpp +++ b/image_analysis/scale_merge/Merge.cpp @@ -414,7 +414,7 @@ struct ShellAccum { CorrelationCoefficient cc_ref; }; -std::pair ImageReferenceCC(const std::vector &reflections, +std::pair ImageReferenceCC(std::span reflections, const std::map &reference, const HKLKeyGenerator &generator, std::optional d_min_limit, diff --git a/image_analysis/scale_merge/Merge.h b/image_analysis/scale_merge/Merge.h index 025dfe7cb..cf7b6df80 100644 --- a/image_analysis/scale_merge/Merge.h +++ b/image_analysis/scale_merge/Merge.h @@ -6,6 +6,7 @@ #include #include #include +#include #include #include @@ -308,7 +309,7 @@ std::vector MergeAll(const DiffractionExperiment &x, // This is the per-image image_scale_cc: ScaleOnTheFly sets it, and StillsPartialityRefine recomputes it // after refining the partiality model, so the reported CC always describes the corrections that will be // merged - which matters because --min-image-cc drops images by it. -std::pair ImageReferenceCC(const std::vector &reflections, +std::pair ImageReferenceCC(std::span reflections, const std::map &reference, const HKLKeyGenerator &generator, std::optional d_min_limit, diff --git a/image_analysis/scale_merge/ReindexAmbiguity.cpp b/image_analysis/scale_merge/ReindexAmbiguity.cpp index 672bf4bfb..d776097ea 100644 --- a/image_analysis/scale_merge/ReindexAmbiguity.cpp +++ b/image_analysis/scale_merge/ReindexAmbiguity.cpp @@ -164,7 +164,7 @@ bool ReindexAmbiguityResolver::Accept(const Reflection &r) const { return AcceptReflection(r, s.GetHighResolutionLimit_A(), s.GetLowResolutionLimit_A()); } -double ReindexAmbiguityResolver::ReferenceCC(const std::vector &reflections, +double ReindexAmbiguityResolver::ReferenceCC(std::span reflections, const gemmi::Op &op) const { double sx = 0, sy = 0, sxx = 0, syy = 0, sxy = 0; size_t n = 0; @@ -195,7 +195,7 @@ double ReindexAmbiguityResolver::ReferenceCC(const std::vector &refl // Serial stills index each crystal in one of the merohedrally-equivalent hands at random; pick, for this // image alone, the reindexing whose intensities correlate best with the external reference and apply it. -void ReindexAmbiguityResolver::Resolve(std::vector &reflections) const { +void ReindexAmbiguityResolver::Resolve(std::span reflections) const { if (ops.empty()) return; diff --git a/image_analysis/scale_merge/ReindexAmbiguity.h b/image_analysis/scale_merge/ReindexAmbiguity.h index 4b8f6f32f..76c8b6bc0 100644 --- a/image_analysis/scale_merge/ReindexAmbiguity.h +++ b/image_analysis/scale_merge/ReindexAmbiguity.h @@ -5,6 +5,7 @@ #include #include +#include #include #include "gemmi/symmetry.hpp" @@ -82,7 +83,7 @@ class ReindexAmbiguityResolver { // Pearson correlation of this image's partiality/Lorentz-corrected intensities against the reference, // matched by ASU index under `op`. Scale-invariant (the unknown per-image scale cancels). NaN when // too few reflections match. - double ReferenceCC(const std::vector &reflections, const gemmi::Op &op) const; + double ReferenceCC(std::span reflections, const gemmi::Op &op) const; public: // The reference must be an external, correctly-handed dataset (not the data's own running merge). @@ -90,5 +91,5 @@ public: // Reindex this image's reflections in place into the hand best-correlated with the reference (done // once per image, for good). No-op when there is no ambiguity. - void Resolve(std::vector &reflections) const; + void Resolve(std::span reflections) const; }; diff --git a/reader/HDF5MetadataSource.cpp b/reader/HDF5MetadataSource.cpp index e65edee4e..7c7de1169 100644 --- a/reader/HDF5MetadataSource.cpp +++ b/reader/HDF5MetadataSource.cpp @@ -303,9 +303,10 @@ CrystalLattice ApplyReindex(const CrystalLattice &latt, const std::optional bool ReadReflectionsFromGroup(HDF5Object &file, const std::string &image_group_name, - std::vector &reflections, + ReflectionContainer &reflections, const std::optional> &reindex) { if (!file.Exists("/entry/reflections") || !file.Exists(image_group_name)) return false; diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index 557c79afe..992916202 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -6016,7 +6016,7 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b } else { for (auto &io : indexer->GetIntegrationOutcome()) { io.latt = io.latt.Multiply(reduce); - std::vector kept; + ReflectionVector kept; kept.reserve(io.reflections.size()); for (auto r : io.reflections) if (reduce_hkl(r))