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

rugnux: ingest the scaling engine's host half beside the post-refinement

Browse Source 
RotationScaleMerge::Ingest is now IngestHost (read the outcomes, sort, build every host array,
smooth the geometry and recompute the partiality) followed by IngestDevice (allocate the GPU's
arrays and stage them). The host half reads the outcomes and nothing else, as the post-refinement
and the quality guard's merge do, so on the canonical pass it runs on its own thread beside both;
the device half runs where the ingest used to, after the guard has handed the GPU back, so the two
engines never hold the card at once. The device feed now stages the recomputed partiality in the
same slices as every other field, straight from the ingest array, instead of in a pass of its own
after the rest; the values on the device are the ones they were.

8a1a, no other load, interleaved runs: the first merge (P1) ends at 10.9 s after the image pass
against 11.3 with only the guard overlapped and 12.7 on the rc175 base. p.mtz byte-identical on the
gate sets (myob, thau, cytc, kdp, cytc -N 4 / -N 8, --prepass-fraction 1 on four sets, 8a1a,
8tyy, CPU build myob); --model, stills, calibration, --mode scale (GPU and CUDA_VISIBLE_DEVICES="")
outputs identical; --no-merge / --mode azint / --write-process-h5 _process.h5 differ only in the
datasets that differ base-vs-base (azint sums, bkgEstimate, profiling times, dates, git).

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
This commit is contained in:
leonarski_fandClaude Opus 5.5 committed 2026-10-07 21:31:29 +02:00
1 parent 19b3499a1a
commit 9cc8172da5
3 files changed
+119 -86
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

image_analysis/scale_merge/RotationScaleMerge.cpp
+88 -82
View File
@@ -371,6 +371,11 @@ void RotationScaleMerge::MeasureIncidentFlux(const std::vector<double> &mean_bkg
}
void RotationScaleMerge::Ingest() {
IngestHost();
IngestDevice();
}
void RotationScaleMerge::IngestHost() {
TimingMark mark(*logger, "rsm ingest");
scale_partials_choice.reset();
n_frames = static_cast<int>(partials_out.size());
@@ -617,23 +622,86 @@ void RotationScaleMerge::Ingest() {
});
#ifdef JFJOCH_USE_CUDA
// 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] = ingest_partiality[at + i];
// On the resident path the ingest arrays' one host reader after the device feed is
// RockingEventFrames - which on this path sees only ingest-time values (see the header) - so its
// answer is taken now, and IngestDevice hands the arrays back once the device holds the rest.
if (resident_ingest)
rocking_event_frames_at_ingest = RockingEventFrames(&rocking_events_at_ingest);
#endif
}
#ifdef JFJOCH_USE_CUDA
// The resident path's device feed: every per-obs field the device reads, staged straight from the
// source reflections in frame-aligned bounded slices - the partiality from the ingest array, where the
// recompute left it.
void RotationScaleMerge::StageObservations() {
gpu_->SetPartialsLayout(n_partials_obs, n_frames, frame_start.data(), frame_count.data());
using F = RotationScaleMergeGPU::ObsField;
// Two million observations a slice: the fifteen staging arrays are then about a hundred
// megabytes together, where eight million put four hundred on top of the ingest's peak, and a
// slice still spans tens of frames for the threads to share.
constexpr int STAGE_SLICE = 2 << 20;
constexpr int NFLOAT = 12;
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]), clip(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 = ingest_src_start[o];
int at = frame_start[o] - base;
for (const auto &r : partials_out[o].reflections) {
if (!ingest_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;
fbuf[11][at] = ingest_partiality[base + at];
frm[at] = o;
onice[at] = r.on_ice_ring ? 1 : 0;
clip[at] = (r.clipped ? OBS_CLIPPED : 0) | (r.overloaded ? OBS_OVERLOADED : 0);
++at;
}
});
gpu_->SetObsField(F::Partiality, at, cnt, buf.get());
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,
F::Partiality};
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());
gpu_->SetObsClipped(base, cnt, clip.get());
}
f0 = f1;
}
}
#endif
void RotationScaleMerge::IngestDevice() {
TimingMark mark(*logger, "rsm ingest (device)");
#ifdef JFJOCH_USE_CUDA
if (resident_ingest)
StageObservations();
if (gpu_active_ && !resident_ingest) {
const int n = static_cast<int>(partials.size());
gpu_->SetPartialsLayout(n, n_frames, frame_start.data(), frame_count.data());
@@ -682,14 +750,10 @@ void RotationScaleMerge::Ingest() {
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 ingest arrays'
// one remaining reader is RockingEventFrames - which on this path sees only ingest-time values
// (see the header). Take its answer now and hand them 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(&rocking_events_at_ingest);
// The device now holds every per-obs field the resident pipeline reads: the ingest arrays go.
// The dump path keeps the full Obs array: the CPU combine is what writes the dump.
if (resident_ingest)
partials_released = true;
}
#endif
ingest_delta_phi.reset();
ingest_partiality.reset();
@@ -812,10 +876,9 @@ void RotationScaleMerge::BuildInRangeObservations(std::unique_ptr<SortKey[]> &ke
ingest_image_number.reset(new float[n_keep]);
#ifdef JFJOCH_USE_CUDA
if (resident_ingest) {
// The resident path: fill the ingest arrays, 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.
// The resident path: fill the ingest arrays only. The 80-byte Obs is never materialised -
// beside the source reflections it would be another copy of the whole payload - and the
// device is fed straight from the source reflections by StageObservations.
ingest_rock_ok.reset(new uint8_t[n_keep]);
corr_ingested.resize(n_keep);
ParallelFor(n_frames, nthreads, [&](int o) {
@@ -840,63 +903,6 @@ void RotationScaleMerge::BuildInRangeObservations(std::unique_ptr<SortKey[]> &ke
}
});
gpu_->SetPartialsLayout(n_keep, n_frames, frame_start.data(), frame_count.data());
using F = RotationScaleMergeGPU::ObsField;
// Two million observations a slice: the thirteen staging arrays are then under a hundred
// megabytes together, where eight million put four hundred on top of the ingest's peak, and a
// slice still spans tens of frames for the threads to share.
constexpr int STAGE_SLICE = 2 << 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]), clip(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;
clip[at] = (r.clipped ? OBS_CLIPPED : 0) | (r.overloaded ? OBS_OVERLOADED : 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());
gpu_->SetObsClipped(base, cnt, clip.get());
}
f0 = f1;
}
} else {
#endif
partials.clear();
image_analysis/scale_merge/RotationScaleMerge.h
+8
View File
@@ -116,6 +116,11 @@ public:
// Copy the per-frame partials into the flat buffers. Call once before the first Run().
void Ingest();
// Ingest in its two halves, which is what Ingest() calls: the host half reads the outcomes and
// builds every host array, the device half allocates the GPU's arrays and fills them. A caller can
// run the host half beside other work and the device half once the GPU has room.
void IngestHost();
void IngestDevice();
// Scale (per-frame G) -> smooth G -> 3D combine -> scale fulls -> merge -> error model -> statistics
// for the space group currently set on the experiment, reusing the ingested buffers.
@@ -506,6 +511,9 @@ private:
// observation array is built, so the two never coexist - together they would be another third of
// the payload on top of it.
void BuildInRangeObservations(std::unique_ptr<SortKey[]> &keys, int n_obs);
#ifdef JFJOCH_USE_CUDA
void StageObservations();
#endif
// Compute the dense ASU-group id for the current space group by grouping the (pre-sorted) raw-hkl
// runs by their ASU key - one gemmi ASU reduction per distinct raw hkl, not per observation. Fills
rugnux/RugnuxScaleMerge.cpp
+23 -4
View File
@@ -724,6 +724,18 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
std::future<ScaleMergeResult> guard_ahead;
if (merge_first_part && !(postrefine_probe_ && !geometry_prepass && postrefine_probe_only_))
guard_ahead = std::async(std::launch::async, merge_guard);
// The scaling engine's ingest reads the outcomes and nothing else either, so its host half runs
// beside the post-refinement too; the device half waits for the guard to hand the GPU back.
Logger engine_log = Logger::Buffered();
std::future<void> engine_ahead;
if (is_rotation && postrefine_probe_ && !geometry_prepass && !postrefine_probe_only_
&& !experiment_.GetRefineRotationWedgeInScaling() && !rot_ss.GetRotationWedgeForScaling().has_value())
engine_ahead = std::async(std::launch::async, [&] {
rsm.emplace(experiment_, indexer->GetIntegrationOutcome(), result.consensus_cell,
static_cast<int>(config_.scaling_iter),
config_.nthreads, engine_log, config_.observation_dump_path);
rsm->IngestHost();
});
if (postrefine_probe_ && !geometry_prepass && is_rotation) {
post_refine_geometry();
if (speculate_geometry_probe_ && !cancelled_ && prepass_detector_geometry_ && result.post_refine
@@ -774,10 +786,17 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
// A reference MTZ is allowed for rotation: it fixes the space group / cell (on the CLI) and
// resolves the indexing ambiguity (below), but is NOT used to scale - the rotation merge stays
// self-consistent.
rsm.emplace(experiment_, indexer->GetIntegrationOutcome(), result.consensus_cell,
static_cast<int>(config_.scaling_iter),
config_.nthreads, logger, config_.observation_dump_path);
rsm->Ingest();
if (engine_ahead.valid()) {
engine_ahead.get();
engine_log.ReplayInto(logger);
rsm->SetLogger(logger);
} else {
rsm.emplace(experiment_, indexer->GetIntegrationOutcome(), result.consensus_cell,
static_cast<int>(config_.scaling_iter),
config_.nthreads, logger, config_.observation_dump_path);
rsm->IngestHost();
}
rsm->IngestDevice();
}
// The geometry pre-pass reads the per-image reflections exactly twice more: they were just