mx/Jungfraujoch
3
2 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 89 Wiki Activity

scale/merge: the resident ingest never materialises the full Obs record

Browse Source 
The previous commit freed the host Obs array once the device held every field;
this one stops building it. With the GPU pipeline resident (and no observation
dump) the build lays down a forty-byte record - hkl, frame, rocking position,
resolution, zeta, delta_phi, partiality, and the rocking-event usability bit -
holding exactly what the host still touches between the build and the release:
the geometry smoothing and the partiality recompute, the corr snapshot, and
the rocking-event median. Every immutable per-obs field is staged to the
device straight from the source reflections, in frame-aligned bounded slices,
and the recomputed partiality follows after the smoothing. The smoothing loops
and the rocking-event walk run on either record through the same code
(templated on the record type), so the arithmetic is exactly the one the full
record ran; the CPU and dump paths keep the full Obs as before.

On a rotation crystal with a ~640 A axis and 95.7 M first-pass partials the
full record was 7.7 GB, built and held beside the 8.4 GB of source reflections
it copied from; the narrow one is 3.8 GB and is handed back at the end of
Ingest. Measured there, on top of the previous commit's release: peak RSS
20.3-22.3 GB -> 18.3 GB, and the phases that used to carry the peak drop far
more - the second pass's post-refine probe and space-group search run at
13-15 GB where they ran at 18-20, and the first pass's post-Ingest plateau at
12-16 where it ran at 20-21. What remains on top is no longer the scale/merge:
the new peak is the first-pass ingest sitting on the image loop's own
high-water mark (~17 GB), which this code never touches.

Byte-identical output on that crystal and on two ordinary rotation sentinels:
merged, P1 and unmerged MTZ data blocks, .hkl, .cif, per-image table, and the
report modulo date/git/command line/wall time.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
leonarski_fandClaude Opus 5 committed 2026-09-16 14:37:13 +02:00
1 parent a1b63c3b5e
commit fd274d5bda
2 files changed
+200 -38
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/scale_merge/RotationScaleMerge.cpp
+174 -35
View File
@@ -300,6 +300,14 @@ void RotationScaleMerge::MeasureIncidentFlux(const std::vector<double> &mean_bkg
void RotationScaleMerge::Ingest() {
n_frames = static_cast<int>(partials_out.size());
partials_released = false;
resident_ingest = false;
#ifdef JFJOCH_USE_CUDA
// Probed before the observation array is built: which record the build lays down (the narrow
// resident one or the full Obs) depends on whether the device pipeline will be active.
gpu_ = std::make_unique<RotationScaleMergeGPU>();
gpu_active_ = gpu_->Available();
resident_ingest = gpu_active_ && observation_dump_path.empty();
#endif
size_t total = 0;
for (const auto &o : partials_out) total += o.reflections.size();
@@ -491,23 +499,40 @@ void RotationScaleMerge::Ingest() {
BuildInRangeObservations(keys);
rawrun_group.assign(rawrun_start.size(), -1);
logger.Info("RotationScaleMerge: ingested {} partial observations from {} frames ({} distinct hkl)",
partials.size(), n_frames, rawrun_start.size());
n_partials_obs, n_frames, rawrun_start.size());
SmoothMosaicityAndPartiality();
// The corr every pass starts from (see the header).
corr_ingested.resize(partials.size());
ParallelChunks(static_cast<int>(partials.size()), ThreadsForWork(partials.size(), nthreads),
[&](int lo, int hi) {
for (int i = lo; i < hi; ++i) corr_ingested[i] = partials[i].corr;
});
// The corr every pass starts from (see the header). The resident build filled it already,
// straight from the source reflections (and the smoothing between there and here never
// touches corr).
if (!resident_ingest) {
corr_ingested.resize(partials.size());
ParallelChunks(static_cast<int>(partials.size()), ThreadsForWork(partials.size(), nthreads),
[&](int lo, int hi) {
for (int i = lo; i < hi; ++i) corr_ingested[i] = partials[i].corr;
});
}
#ifdef JFJOCH_USE_CUDA
// Bring the partial-scaling loop onto the GPU when one is present. Upload the immutable per-obs
// fields once (corr lives on the device, refreshed each pass); the CPU keeps the sort/keying/combine.
gpu_ = std::make_unique<RotationScaleMergeGPU>();
gpu_active_ = gpu_->Available();
if (gpu_active_) {
// Feed the device (probed at the top of Ingest). On the resident path the layout and every
// immutable field went up inside BuildInRangeObservations, straight from the source
// reflections; what is staged here is the partiality, which the recompute above rewrote.
if (resident_ingest) {
using F = RotationScaleMergeGPU::ObsField;
constexpr int STAGE_SLICE = 8 << 20;
const int n = n_partials_obs;
const int slice = std::min(std::max(n, 1), STAGE_SLICE);
const std::unique_ptr<float[]> buf(new float[slice]);
for (int at = 0; at < n; at += slice) {
const int cnt = std::min(n - at, slice);
ParallelChunks(cnt, nthreads, [&](int lo, int hi) {
for (int i = lo; i < hi; ++i) buf[i] = partials_ingest[at + i].partiality;
});
gpu_->SetObsField(F::Partiality, at, cnt, buf.get());
}
}
if (gpu_active_ && !resident_ingest) {
const int n = static_cast<int>(partials.size());
gpu_->SetPartialsLayout(n, n_frames, frame_start.data(), frame_count.data());
// Staged in bounded slices: fourteen full-length arrays at once are more than half of the
@@ -546,19 +571,21 @@ void RotationScaleMerge::Ingest() {
gpu_->SetObsFrame(at, cnt, frm.get());
gpu_->SetObsOnIce(at, cnt, onice.get());
}
}
if (gpu_active_) {
gpu_->SetRawRuns(static_cast<int>(rawrun_start.size()), static_cast<int>(perm.size()), perm.data(),
rawrun_start.data(), rawrun_count.data(),
rawrun_h.data(), rawrun_k.data(), rawrun_l.data());
gpu_->SetFrameCellOk(frame_cell_ok.data());
logger.Info("RotationScaleMerge: GPU scaling + combine + scale-fulls + merge active");
}
// The device now holds every per-obs field the resident pipeline reads, so the host copy's one
// remaining reader is RockingEventFrames - which on this path sees only ingest-time values (see
// the header). Take its answer now and hand the array back. The dump path keeps it: the CPU
// combine is what writes the dump.
if (gpu_active_ && observation_dump_path.empty()) {
// The device now holds every per-obs field the resident pipeline reads, so the host record's
// one remaining reader is RockingEventFrames - which on this path sees only ingest-time values
// (see the header). Take its answer now and hand the record back. The dump path keeps the full
// Obs array: the CPU combine is what writes the dump.
if (resident_ingest) {
rocking_event_frames_at_ingest = RockingEventFrames();
std::vector<Obs>().swap(partials);
std::vector<ObsIngest>().swap(partials_ingest);
partials_released = true;
}
#endif
@@ -663,9 +690,99 @@ void RotationScaleMerge::BuildInRangeObservations(std::vector<SortKey> &keys) {
// The observations themselves, one frame per thread. finite_ok is the per-obs AcceptReflection
// finiteness (immutable) - it lets ComputeAsuGroups stamp the ASU-group id per obs from a flat
// 1-byte array instead of re-reading the fat Obs struct for every space group.
finite_ok.assign(n_keep, 0);
#ifdef JFJOCH_USE_CUDA
if (resident_ingest) {
// The resident path: build the narrow host record, then stage every immutable per-obs field
// to the device straight from the source reflections, in frame-aligned bounded slices. The
// 80-byte Obs is never materialised here - beside the source reflections it was another
// copy of the whole payload - and the slice buffers are bounded whatever the run.
partials_ingest.clear();
partials_ingest.resize(n_keep);
corr_ingested.resize(n_keep);
ParallelFor(n_frames, nthreads, [&](int o) {
int32_t src = src_start[o];
int32_t at = frame_start[o];
for (const auto &r : partials_out[o].reflections) {
if (!keep[src]) { ++src; continue; }
++src;
ObsIngest &obs = partials_ingest[at];
obs.h = r.h; obs.k = r.k; obs.l = r.l;
obs.frame = o;
obs.image_number = r.image_number;
obs.d = r.d;
obs.zeta = r.zeta;
obs.delta_phi = r.delta_phi_deg;
obs.partiality = r.partiality;
const float prescaling_corr = static_cast<float>(
r.prescaling_corr * r.qe_corr * r.flight_corr / frame_flux[o]);
corr_ingested[at] = r.image_scale_corr;
finite_ok[at] = (std::isfinite(r.I) && std::isfinite(prescaling_corr)
&& prescaling_corr != 0.0f
&& std::isfinite(r.sigma) && r.sigma > 0.0f) ? 1 : 0;
obs.rock_ok = (std::isfinite(r.image_scale_corr) && r.image_scale_corr > 0.0f
&& std::isfinite(r.I) && std::isfinite(r.sigma) && r.sigma > 0.0f)
? 1 : 0;
++at;
}
});
gpu_->SetPartialsLayout(n_keep, n_frames, frame_start.data(), frame_count.data());
using F = RotationScaleMergeGPU::ObsField;
constexpr int STAGE_SLICE = 8 << 20;
constexpr int NFLOAT = 11;
int max_frame_obs = 0;
for (int o = 0; o < n_frames; ++o) max_frame_obs = std::max(max_frame_obs, static_cast<int>(frame_count[o]));
const int bufsize = std::max(1, std::max(STAGE_SLICE, max_frame_obs));
std::unique_ptr<float[]> fbuf[NFLOAT];
for (auto &b : fbuf) b.reset(new float[bufsize]);
const std::unique_ptr<int32_t[]> frm(new int32_t[bufsize]);
const std::unique_ptr<uint8_t[]> onice(new uint8_t[bufsize]);
int f0 = 0;
while (f0 < n_frames) {
int f1 = f0 + 1;
while (f1 < n_frames && frame_start[f1] + frame_count[f1] - frame_start[f0] <= bufsize)
++f1;
const int32_t base = frame_start[f0];
const int cnt = frame_start[f1 - 1] + frame_count[f1 - 1] - base;
if (cnt > 0) {
ParallelFor(f1 - f0, nthreads, [&](int fo) {
const int o = f0 + fo;
int32_t src = src_start[o];
int at = frame_start[o] - base;
for (const auto &r : partials_out[o].reflections) {
if (!keep[src]) { ++src; continue; }
++src;
fbuf[0][at] = r.I;
fbuf[1][at] = r.sigma;
fbuf[2][at] = static_cast<float>(
r.prescaling_corr * r.qe_corr * r.flight_corr / frame_flux[o]);
fbuf[3][at] = r.zeta;
fbuf[4][at] = r.image_scale_corr;
fbuf[5][at] = r.bkg;
fbuf[6][at] = r.var_bkg;
fbuf[7][at] = r.image_number;
fbuf[8][at] = r.d;
fbuf[9][at] = r.predicted_x;
fbuf[10][at] = r.predicted_y;
frm[at] = o;
onice[at] = r.on_ice_ring ? 1 : 0;
++at;
}
});
constexpr F fields[NFLOAT] = {F::I, F::Sigma, F::PrescalingCorr, F::Zeta, F::Corr0,
F::Bkg, F::VarBkg, F::ImageNumber, F::D, F::Px, F::Py};
for (int k = 0; k < NFLOAT; ++k)
gpu_->SetObsField(fields[k], base, cnt, fbuf[k].get());
gpu_->SetObsFrame(base, cnt, frm.get());
gpu_->SetObsOnIce(base, cnt, onice.get());
}
f0 = f1;
}
} else {
#endif
partials.clear();
partials.resize(n_keep);
finite_ok.assign(n_keep, 0);
ParallelFor(n_frames, nthreads, [&](int o) {
int32_t src = src_start[o];
int32_t at = frame_start[o];
@@ -692,6 +809,9 @@ void RotationScaleMerge::BuildInRangeObservations(std::vector<SortKey> &keys) {
++at;
}
});
#ifdef JFJOCH_USE_CUDA
}
#endif
n_partials_obs = n_keep;
@@ -761,7 +881,7 @@ namespace {
// window would over-smooth a fast-slipping crystal and under-smooth a sparse one. The frame's own
// value is excluded when scoring, so a window is only preferred if it PREDICTS that frame better,
// which is exactly the noise-versus-drift trade.
void RotationScaleMerge::SmoothGeometry() {
template <class T> void RotationScaleMerge::SmoothGeometry(std::vector<T> &obs) {
const auto gon = x.GetGoniometer();
if (!gon || n_frames < 5)
return;
@@ -868,10 +988,10 @@ void RotationScaleMerge::SmoothGeometry() {
}
std::atomic<int64_t> changed{0};
ParallelChunks(static_cast<int>(partials.size()), nthreads, [&](int lo, int hi) {
ParallelChunks(static_cast<int>(obs.size()), nthreads, [&](int lo, int hi) {
int64_t local = 0;
for (int i = lo; i < hi; ++i) {
auto &o = partials[i];
auto &o = obs[i];
if (o.frame < 0 || o.frame >= n_frames || !ok[o.frame])
continue;
const std::array<Coord, 3> &r = frame_recip[o.frame];
@@ -889,11 +1009,16 @@ void RotationScaleMerge::SmoothGeometry() {
logger.Info("Smoothed per-frame geometry over +-{} frames (chosen by cross-validation); "
"recomputed delta_phi for {} of {} partials", best_half, changed.load(),
partials.size());
obs.size());
}
void RotationScaleMerge::SmoothMosaicityAndPartiality() {
SmoothGeometry();
if (resident_ingest) SmoothMosaicityAndPartialityOn(partials_ingest);
else SmoothMosaicityAndPartialityOn(partials);
}
template <class T> void RotationScaleMerge::SmoothMosaicityAndPartialityOn(std::vector<T> &obs) {
SmoothGeometry(obs);
// One rocking event, one exact-Bragg angle. Every partial's delta_phi is solved from its OWN
// frame's lattice - by the predictor, and again by SmoothGeometry above - so the per-frame
@@ -919,15 +1044,15 @@ void RotationScaleMerge::SmoothMosaicityAndPartiality() {
int i = lo;
while (i < hi) {
int j = i + 1;
while (j < hi && partials[perm[j]].image_number
- partials[perm[j - 1]].image_number <= max_frame_gap)
while (j < hi && obs[perm[j]].image_number
- obs[perm[j - 1]].image_number <= max_frame_gap)
++j;
if (j - i >= 2) {
const float f0 = partials[perm[i]].image_number;
const float f0 = obs[perm[i]].image_number;
double sum = 0.0;
int n = 0;
for (int m = i; m < j; ++m) {
const auto &o = partials[perm[m]];
const auto &o = obs[perm[m]];
if (!std::isfinite(o.delta_phi)) continue;
sum += o.delta_phi + increment_deg * (o.image_number - f0);
++n;
@@ -935,7 +1060,7 @@ void RotationScaleMerge::SmoothMosaicityAndPartiality() {
if (n >= 2) {
const double c = sum / n;
for (int m = i; m < j; ++m) {
auto &o = partials[perm[m]];
auto &o = obs[perm[m]];
o.delta_phi = static_cast<float>(c - increment_deg
* (o.image_number - f0));
}
@@ -1010,9 +1135,9 @@ void RotationScaleMerge::SmoothMosaicityAndPartiality() {
// partiality recomputed here matches the one integration was based on. sin(theta_B) = lambda/(2d).
const double bandwidth_sigma = x.GetBandwidthFWHM().value_or(0.0f) / 2.3548;
const double half_wavelength_A = x.GetWavelength_A() / 2.0;
ParallelChunks(static_cast<int>(partials.size()), nthreads, [&](int lo, int hi) {
ParallelChunks(static_cast<int>(obs.size()), nthreads, [&](int lo, int hi) {
for (int i = lo; i < hi; ++i) {
auto &o = partials[i];
auto &o = obs[i];
float mos = mos_smooth[o.frame];
if (bandwidth_sigma > 0.0 && o.d > 0.0f) {
const double sin_theta = half_wavelength_A / o.d;
@@ -2004,6 +2129,21 @@ namespace {
int RotationScaleMerge::RockingEventFrames() const {
if (partials_released)
return rocking_event_frames_at_ingest;
if (resident_ingest) // Ingest's own call, before the narrow record is handed back
return RockingEventFramesOver(
[&](int i) { return partials_ingest[i].rock_ok != 0; },
[&](int i) { return partials_ingest[i].image_number; });
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;
},
[&](int i) { return partials[i].image_number; });
}
template <class UsableFn, class ImgFn>
int RotationScaleMerge::RockingEventFramesOver(UsableFn usable, ImgFn img) const {
std::vector<int> hist(n_frames + 2, 0);
std::vector<float> im;
int n_event = 0;
@@ -2011,10 +2151,9 @@ int RotationScaleMerge::RockingEventFrames() const {
const int lo = rawrun_start[r], hi = lo + rawrun_count[r];
im.clear();
for (int p = lo; p < hi; ++p) {
const Obs &o = partials[perm[p]];
if (std::isfinite(o.corr) && o.corr > 0.0f && std::isfinite(o.I)
&& std::isfinite(o.sigma) && o.sigma > 0.0f)
im.push_back(o.image_number);
const int i = perm[p];
if (usable(i))
im.push_back(img(i));
}
for (size_t i = 0; i < im.size();) {
size_t k = i + 1;
image_analysis/scale_merge/RotationScaleMerge.h
+26 -3
View File
@@ -128,6 +128,22 @@ private:
// re-run over the samples that survive the automatic resolution cutoff. See MergeAndStats.
struct Sample { double s2, I2, dev2; float d; };
// The record the ingest keeps on the host when the device pipeline is resident (a GPU and no
// observation dump): every immutable per-obs field is uploaded straight from the source
// reflections inside BuildInRangeObservations, and the host keeps only what its own remaining
// stages touch - the geometry smoothing (delta_phi, then the recomputed partiality), plus the
// rocking-event walk's inputs. Forty bytes against the Obs's eighty, and on that path the fat
// record is never materialised at all: on a fine-sliced long axis it alone is many gigabytes,
// held beside the source reflections it was copied from.
struct ObsIngest {
int32_t h, k, l, frame;
float image_number, d, zeta, delta_phi, partiality;
// The rocking-event usability test (RockingEventFrames), evaluated on the ingest-time
// values - which is what that walk reads on the resident path, where the host corr is
// never refreshed (see partials_released).
uint8_t rock_ok;
};
// The narrow per-observation record the ingest sort orders: the raw hkl the runs are cut on, the
// frame position that breaks a tie inside one, the observation's own index (which makes the order
// total - see the .cpp), and the resolution the range test reads. Twenty-four bytes against the
@@ -190,8 +206,12 @@ private:
double resolution_cc_target = 0.30;
int report_shell_count = 9;
// Flat buffers, allocated once by Ingest() and reused across Run() calls.
// Flat buffers, allocated once by Ingest() and reused across Run() calls. Exactly one of the
// two observation arrays is built: the narrow resident record when the device pipeline is
// active and no observation dump was asked for (resident_ingest), the full Obs otherwise.
std::vector<Obs> partials; // all per-frame partials, grouped by frame
std::vector<ObsIngest> partials_ingest; // the resident-path record (same order)
bool resident_ingest = false;
// How many observations Ingest built (== partials.size() until any release below), so the
// sizes survive when the array itself does not.
int n_partials_obs = 0;
@@ -404,8 +424,9 @@ private:
// per-frame partials of one rocking event tile the curve consistently (they sum toward 1) before the
// 3D combine. Deterministic (frame order); replaces the old arrival-order mosaicity moving average
// that prediction applied. SG-independent, so done once in Ingest.
void SmoothGeometry();
void SmoothMosaicityAndPartiality();
template <class T> void SmoothGeometry(std::vector<T> &obs);
void SmoothMosaicityAndPartiality(); // dispatch: the resident narrow record or the full Obs
template <class T> void SmoothMosaicityAndPartialityOn(std::vector<T> &obs);
void Combine(); // partials -> fulls (CPU)
@@ -470,6 +491,8 @@ private:
// Walked on the partials, because the combine also runs on the GPU and a full keeps only the frame
// of its peak partial.
[[nodiscard]] int RockingEventFrames() const;
template <class UsableFn, class ImgFn>
[[nodiscard]] int RockingEventFramesOver(UsableFn usable, ImgFn img) const;
// Per-batch delta-CC1/2 on the corrected fulls: measure what keeping each batch of the sweep costs
// the merged intensities, convict the batches that cost significantly, slide the conviction's edges
// onto the frames that carry it, and turn the result into the disposition ledger. Fills