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

rotation merge: the resident ingest sorts and cuts the reflections on the card

Browse Source 
With the GPU resident, the host ingest swept the source reflections four times (background
and h histogram, sort keys, the in-range build, the device staging), bucket-sorted the keys
and cut the raw-hkl runs: 5-6 s of host time per engine on a 78 M-reflection sweep, run
twice beside the quality guard. Now the frames' reflections are copied once, through two
pinned buffers, to the card, where the per-frame background is summed, the reflections
are ordered by two stable radix passes ((image, index), then raw hkl - the one order the
host's total comparator gives), cut into runs by run-length encoding, filtered by the
resolution range and compacted. The observations, the runs and their order are the host
ingest's to the bit; a new GPU test checks the layout against the host's sort and cut.
The geometry smoothing still runs on the host in this step, on arrays read back.

The all-observation arm's engine is cloned after the main engine's device ingest (its
DeviceBytes are known only then) and copies the runs on the card as well. DeviceBytes now
also covers the ingest's own peak (90 bytes per reflection, 64 per run before the cut).

Measured: p.hkl and MTZ data identical on cytc and 7qij.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3
This commit is contained in:
leonarski_f
2026-10-12 22:38:25 +02:00
co-authored by Claude Opus 5.5
parent e24fb88b53
commit f88ff0440d
7 changed files with 933 additions and 185 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/scale_merge/RotationScaleMerge.cpp
+123 -145
View File
@@ -399,11 +399,21 @@ void RotationScaleMerge::IngestHost() {
resident_ingest = gpu_active_ && observation_dump_path.empty();
#endif
size_t total = 0;
for (int o = 0; o < n_frames; ++o) total += partials_out[o].reflections.size();
frame_cell_ok.assign(n_frames, 1);
g_partial.assign(n_frames, 1.0);
const float dist_tol = x.GetIndexingSettings().GetUnitCellDistTolerance();
const float ang_tol = x.GetIndexingSettings().GetUnitCellAngleTolerance_deg();
if (resident_ingest) {
// The reflections are read on the card (IngestOnDevice); here only the per-frame cells.
if (reference_cell)
for (int o = 0; o < n_frames; ++o)
frame_cell_ok[o] = partials_out[o].latt.GetUnitCell().is_close(*reference_cell, dist_tol, ang_tol)
? 1 : 0;
return;
}
size_t total = 0;
for (int o = 0; o < n_frames; ++o) total += partials_out[o].reflections.size();
// Where each frame's slice of the key array starts. Every reflection of a frame becomes one key, so
// this is known before anything is read - which is what lets the sweep below hand a frame to each
@@ -417,9 +427,6 @@ void RotationScaleMerge::IngestHost() {
}
}
const float dist_tol = x.GetIndexingSettings().GetUnitCellDistTolerance();
const float ang_tol = x.GetIndexingSettings().GetUnitCellAngleTolerance_deg();
// The per-frame mean background - the incident-flux meter, see MeasureIncidentFlux - is
// accumulated here rather than in a pass of its own. One frame is still summed in its own order,
// which is what keeps the mean bit-exact.
@@ -582,94 +589,101 @@ void RotationScaleMerge::IngestHost() {
n_partials_obs, n_frames, rawrun_start.size());
SmoothMosaicityAndPartiality();
if (!resident_ingest) {
ParallelChunks(static_cast<int>(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];
partials[i].partiality = ingest_partiality[i];
}
});
}
ParallelChunks(static_cast<int>(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];
partials[i].partiality = ingest_partiality[i];
}
});
// The corr every pass starts from (see the header) was filled by the build, straight from the
// source reflections, and brought to the recomputed partiality by the smoothing.
if (!resident_ingest)
rocking_event_frames_at_start = -1;
#ifdef JFJOCH_USE_CUDA
// 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
rocking_event_frames_at_start = -1;
}
#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;
}
});
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;
// The resident path's ingest. The reflections go to the card once, frame by frame, and the card measures
// the per-frame background, sorts, cuts the raw-hkl runs and drops what the resolution range excludes
// (RotationScaleMergeGPU::IngestUpload .. IngestSort) - the observations, the runs and the order of each
// that BuildInRangeObservations lays down on the host, to the bit, and left where the device pipeline
// reads them. The geometry smoothing then runs here on what the card hands back.
void RotationScaleMerge::IngestOnDevice() {
std::vector<const Reflection *> refl(n_frames);
std::vector<int32_t> count(n_frames);
for (int o = 0; o < n_frames; ++o) {
refl[o] = partials_out[o].reflections.data();
count[o] = static_cast<int32_t>(partials_out[o].reflections.size());
}
TimingMark upload_mark(*logger, "rsm ingest upload");
gpu_->IngestUpload(n_frames, refl.data(), count.data(), nthreads);
upload_mark.Stop();
TimingMark sort_mark(*logger, "rsm ingest sort");
std::vector<double> mean_bkg(n_frames);
gpu_->IngestMeanBackground(mean_bkg.data());
MeasureIncidentFlux(mean_bkg);
const int n = gpu_->IngestSort(frame_flux.data(), d_min_limit, d_max_limit);
n_partials_obs = n;
RotationScaleMergeGPU::IngestLayout layout;
gpu_->IngestGetLayout(layout);
ingest_reflections = static_cast<size_t>(layout.n_reflections);
ingest_runs = static_cast<size_t>(layout.n_runs_all);
sort_mark.Stop();
TimingMark smoothing_mark(*logger, "rsm ingest smoothing");
frame_start = std::move(layout.frame_start);
frame_count = std::move(layout.frame_count);
rawrun_start = std::move(layout.run_start);
rawrun_count = std::move(layout.run_count);
rawrun_h = std::move(layout.run_h);
rawrun_k = std::move(layout.run_k);
rawrun_l = std::move(layout.run_l);
rawrun_d = std::move(layout.run_d);
ingest_delta_phi.reset(new float[n]);
ingest_partiality.reset(new float[n]);
ingest_image_number.reset(new float[n]);
ingest_zeta.reset(new float[n]);
ingest_rock_ok.reset(new uint8_t[n]);
corr_ingested.resize(n);
perm.resize(n);
std::vector<int32_t> h(n), k(n), l(n);
const std::unique_ptr<float[]> dphi(new float[n]), d(new float[n]), I(new float[n]), sigma(new float[n]);
gpu_->IngestGetObservations(h.data(), k.data(), l.data(), dphi.get(), ingest_partiality.get(),
ingest_image_number.get(), ingest_zeta.get(), d.get(), corr_ingested.data(),
I.get(), sigma.get(), perm.data());
if (x.GetBandwidthFWHM().value_or(0.0f) > 0.0f) {
ingest_d.reset(new float[n]);
std::copy(d.get(), d.get() + n, ingest_d.get());
}
const SmoothedGeometry sg = SmoothedFrameGeometry();
std::atomic<int64_t> n_resolved{0};
ParallelFor(n_frames, nthreads, [&](int o) {
int64_t resolved = 0;
for (int i = frame_start[o]; i < frame_start[o] + frame_count[o]; ++i) {
ingest_delta_phi[i] = SmoothedDeltaPhi(sg, o, h[i], k[i], l[i], dphi[i], resolved);
ingest_rock_ok[i] = (std::isfinite(corr_ingested[i]) && corr_ingested[i] > 0.0f && std::isfinite(I[i])
&& std::isfinite(sigma[i]) && sigma[i] > 0.0f) ? 1 : 0;
}
n_resolved += resolved;
});
if (sg.half > 0)
logger->Info("Smoothed per-frame geometry over +-{} frames (chosen by cross-validation); "
"recomputed delta_phi for {} of {} partials", sg.half, n_resolved.load(), n);
if (n < layout.n_reflections)
logger->Info("RotationScaleMerge: dropped {} of {} observations ({} of {} distinct hkl) outside the "
"requested resolution range - nothing downstream could have merged them",
layout.n_reflections - n, layout.n_reflections,
layout.n_runs_all - static_cast<int>(rawrun_start.size()), layout.n_runs_all);
rawrun_group.assign(rawrun_start.size(), -1);
logger->Info("RotationScaleMerge: ingested {} partial observations from {} frames ({} distinct hkl)",
n_partials_obs, n_frames, rawrun_start.size());
SmoothMosaicityAndPartiality();
// The ingest arrays' one host reader after this is RockingEventFrames - which on this path sees only
// ingest-time values (see the header) - so its answer is taken now.
rocking_event_frames_at_ingest = RockingEventFrames(&rocking_events_at_ingest);
gpu_->IngestSetPartialityAndCorr(ingest_partiality.get(), corr_ingested.data());
std::vector<int32_t>().swap(perm);
}
#endif
@@ -683,7 +697,13 @@ size_t RotationScaleMerge::DeviceBytes() const {
const auto events = static_cast<size_t>(resident_ingest ? rocking_events_at_ingest : n_partials_obs);
const size_t bytes = BYTES_PER_OBS * static_cast<size_t>(n_partials_obs) + BYTES_PER_EVENT * events
+ BYTES_PER_HKL * rawrun_start.size();
return bytes + bytes / 5;
// The resident ingest's own peak comes before any of that, over every reflection read and every run
// before the resolution cut: the observation arrays, the sort key and the integrated delta_phi of
// each (66 bytes) and 24 bytes of the sort's or the cut's working arrays beside them, plus sixteen
// arrays per run at the cut (IngestSort).
constexpr size_t INGEST_BYTES_PER_REFLECTION = 90, INGEST_BYTES_PER_RUN = 64;
const size_t ingest = INGEST_BYTES_PER_REFLECTION * ingest_reflections + INGEST_BYTES_PER_RUN * ingest_runs;
return std::max(bytes + bytes / 5, ingest);
}
std::unique_ptr<RotationScaleMerge> RotationScaleMerge::CloneIngest(int device, Logger &log) const {
@@ -703,11 +723,10 @@ void RotationScaleMerge::IngestDevice(const RotationScaleMerge *observations_fro
const ScopedGpuPin pin(gpu_device); // see IngestHost
TimingMark mark(*logger, "rsm ingest (device)");
#ifdef JFJOCH_USE_CUDA
if (gpu_active_ && observations_from) {
gpu_->SetPartialsLayout(n_partials_obs, n_frames, frame_start.data(), frame_count.data());
if (gpu_active_ && observations_from)
gpu_->CopyObservationsFrom(*observations_from->gpu_);
} else if (resident_ingest)
StageObservations();
else if (resident_ingest)
IngestOnDevice();
if (gpu_active_ && !resident_ingest && !observations_from) {
const int n = static_cast<int>(partials.size());
gpu_->SetPartialsLayout(n, n_frames, frame_start.data(), frame_count.data());
@@ -748,11 +767,11 @@ void RotationScaleMerge::IngestDevice(const RotationScaleMerge *observations_fro
gpu_->SetObsOnIce(at, cnt, onice.get());
gpu_->SetObsClipped(at, cnt, clip.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());
}
if (gpu_active_) {
gpu_->SetFrameCellOk(frame_cell_ok.data());
logger->Info("RotationScaleMerge: GPU scaling + combine + scale-fulls + merge active");
}
@@ -765,8 +784,6 @@ void RotationScaleMerge::IngestDevice(const RotationScaleMerge *observations_fro
ingest_partiality.reset();
ingest_image_number.reset();
ingest_rock_ok.reset();
std::vector<uint8_t>().swap(ingest_keep);
std::vector<int32_t>().swap(ingest_src_start);
}
// Build `partials` (and the per-frame CSR, the finiteness mask and `perm`) from the source reflections,
@@ -803,15 +820,14 @@ void RotationScaleMerge::BuildInRangeObservations(std::unique_ptr<SortKey[]> &ke
// Where each frame's reflections sit in the key numbering, and how many of them survive. Building in
// index order keeps each frame's observations contiguous and in the order they had, so the per-frame
// ranges still describe them and every per-frame sum keeps its sequence of roundings.
ingest_src_start.assign(n_frames, 0);
std::vector<int32_t> src_start(n_frames, 0);
{
int32_t at = 0;
for (int o = 0; o < n_frames; ++o) {
ingest_src_start[o] = at;
src_start[o] = at;
at += static_cast<int32_t>(partials_out[o].reflections.size());
}
}
const std::vector<int32_t> &src_start = ingest_src_start;
frame_count.assign(n_frames, 0);
ParallelFor(n_frames, nthreads, [&](int o) {
int32_t c = 0;
@@ -890,41 +906,6 @@ void RotationScaleMerge::BuildInRangeObservations(std::unique_ptr<SortKey[]> &ke
if (x.GetBandwidthFWHM().value_or(0.0f) > 0.0f)
ingest_d.reset(new float[n_keep]);
std::atomic<int64_t> n_resolved{0};
#ifdef JFJOCH_USE_CUDA
if (resident_ingest) {
// 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) {
int32_t src = src_start[o];
int32_t at = frame_start[o];
int64_t resolved = 0;
for (const auto &r : partials_out[o].reflections) {
if (!keep[src]) { ++src; continue; }
++src;
ingest_delta_phi[at] = SmoothedDeltaPhi(sg, o, r, resolved);
ingest_partiality[at] = r.partiality;
ingest_image_number[at] = r.image_number;
ingest_zeta[at] = r.zeta;
if (ingest_d) ingest_d[at] = r.d;
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;
ingest_rock_ok[at] = (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;
}
n_resolved += resolved;
});
} else {
#endif
partials.clear();
partials.resize(n_keep);
corr_ingested.resize(n_keep);
@@ -953,7 +934,7 @@ void RotationScaleMerge::BuildInRangeObservations(std::unique_ptr<SortKey[]> &ke
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;
ingest_delta_phi[at] = SmoothedDeltaPhi(sg, o, r, resolved);
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;
ingest_zeta[at] = r.zeta;
@@ -963,12 +944,8 @@ void RotationScaleMerge::BuildInRangeObservations(std::unique_ptr<SortKey[]> &ke
}
n_resolved += resolved;
});
#ifdef JFJOCH_USE_CUDA
}
#endif
n_partials_obs = n_keep;
ingest_keep = std::move(keep);
if (sg.half > 0)
logger->Info("Smoothed per-frame geometry over +-{} frames (chosen by cross-validation); "
"recomputed delta_phi for {} of {} partials", sg.half, n_resolved.load(), n_keep);
@@ -1154,16 +1131,17 @@ RotationScaleMerge::SmoothedGeometry RotationScaleMerge::SmoothedFrameGeometry()
// The delta_phi of reflection r of frame o, re-solved against the frame's smoothed lattice where there
// is one (see SmoothedFrameGeometry), as integrated otherwise. Counts the re-solved ones in n.
float RotationScaleMerge::SmoothedDeltaPhi(const SmoothedGeometry &sg, int o, const Reflection &r, int64_t &n) {
float RotationScaleMerge::SmoothedDeltaPhi(const SmoothedGeometry &sg, int o, int32_t h, int32_t k, int32_t l,
float delta_phi_deg, int64_t &n) {
if (sg.ok.empty() || !sg.ok[o])
return r.delta_phi_deg;
return delta_phi_deg;
const std::array<Coord, 3> &rc = sg.recip[o];
const Coord p0 = rc[0] * static_cast<float>(r.h) + rc[1] * static_cast<float>(r.k)
+ rc[2] * static_cast<float>(r.l);
const float phi_near = r.delta_phi_deg * static_cast<float>(PI) / 180.0f;
const Coord p0 = rc[0] * static_cast<float>(h) + rc[1] * static_cast<float>(k)
+ rc[2] * static_cast<float>(l);
const float phi_near = delta_phi_deg * static_cast<float>(PI) / 180.0f;
const auto phi = SolvePhi(p0, sg.S0, sg.m1, sg.m2, sg.m3, sg.m2_S0, sg.m3_S0, phi_near);
if (!phi || !std::isfinite(*phi))
return r.delta_phi_deg;
return delta_phi_deg;
++n;
return *phi * 180.0f / static_cast<float>(PI);
}
image_analysis/scale_merge/RotationScaleMerge.h
+14 -10
View File
@@ -120,7 +120,9 @@ public:
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.
// run the host half beside other work and the device half once the GPU has room. With the GPU
// resident (no observation dump) the host half only reads the per-frame cells: the reflections go
// to the card, and are sorted and cut there, in the device half.
void IngestHost();
// With `observations_from` - the engine this one was copied from (CloneIngest), its own
// IngestDevice done - the observations are copied from that engine's card arrays rather than
@@ -161,7 +163,7 @@ public:
// A second engine on this one's ingest, on `device` and logging to `log`, for a merge to run beside
// this engine's: the host half of the ingest is copied, and the copy's IngestDevice(this) is still
// to be called, after this engine's own. Call between IngestHost() and IngestDevice(). No merge depends on the merges run
// to be called. Call after this engine's IngestDevice(). No merge depends on the merges run
// before it (see RunArm), so a merge on the copy gives what the same merge gives on this engine.
// The copy does not write its per-frame scales back (WriteBackPerFrameScale does).
[[nodiscard]] std::unique_ptr<RotationScaleMerge> CloneIngest(int device, Logger &log) const;
@@ -172,7 +174,8 @@ public:
// events (the combine makes at most one full of each) and the distinct hkl (at most one ASU group
// each) - at what each costs at the engine's peak (the correction-surface terms, or the merge with
// its outlier test), plus a fifth. Measured on 11 data sets of 13-123 M observations, every
// engine's peak was at most 81% of this.
// engine's peak was at most 81% of this. Where the ingest on the card peaks higher, its peak.
// Call after IngestDevice().
[[nodiscard]] size_t DeviceBytes() const;
// Merge only the first n frames of the outcomes (by default all of them). Set before Ingest().
@@ -270,7 +273,8 @@ private:
float m2_S0 = 0.0f, m3_S0 = 0.0f;
};
SmoothedGeometry SmoothedFrameGeometry() const;
static float SmoothedDeltaPhi(const SmoothedGeometry &sg, int o, const Reflection &r, int64_t &n);
static float SmoothedDeltaPhi(const SmoothedGeometry &sg, int o, int32_t h, int32_t k, int32_t l,
float delta_phi_deg, int64_t &n);
const DiffractionExperiment &x;
std::vector<IntegrationOutcome> &partials_out; // written back at the end of scaling
@@ -342,7 +346,7 @@ private:
// record is built at all - on a fine-sliced long axis that record was gigabytes, held beside the
// source reflections it was copied from - and the smoothing does not sweep the source reflections
// again either. Filled by
// BuildInRangeObservations and handed back at the end of Ingest. Plain arrays, not sized vectors, for
// BuildInRangeObservations (IngestOnDevice on the resident path) and handed back at the end of Ingest. Plain arrays, not sized vectors, for
// the reason the staging buffers are (see Ingest): the build writes every element, and a sized vector
// would first zero all of them on one thread.
NotCopied<float[]> ingest_delta_phi, ingest_partiality, ingest_image_number, ingest_zeta, ingest_d;
@@ -350,10 +354,6 @@ private:
// that walk reads on the resident path, where the host corr is never refreshed (see
// partials_released). Resident path only.
NotCopied<uint8_t[]> ingest_rock_ok;
// Which source reflections the resolution range kept (indexed like the sort keys: frame by frame,
// from ingest_src_start), so that the i-th kept one of a frame can be found again.
std::vector<uint8_t> ingest_keep;
std::vector<int32_t> ingest_src_start;
// 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;
@@ -367,6 +367,8 @@ private:
bool partials_released = false;
int rocking_event_frames_at_ingest = 0;
int64_t rocking_events_at_ingest = 0;
// The reflections and raw-hkl runs the resident ingest read before the resolution cut (DeviceBytes).
size_t ingest_reflections = 0, ingest_runs = 0;
std::vector<int32_t> frame_start, frame_count; // CSR ranges of `partials` per frame
std::vector<uint8_t> frame_cell_ok; // per-frame cell-consistency mask (1 = kept)
std::vector<uint8_t> finite_ok; // per-obs AcceptReflection finiteness (immutable; 1 = kept)
@@ -592,7 +594,9 @@ private:
// Every kept partial's delta_phi is re-solved against `sg` as it is built (SmoothedDeltaPhi).
void BuildInRangeObservations(std::unique_ptr<SortKey[]> &keys, int n_obs, const SmoothedGeometry &sg);
#ifdef JFJOCH_USE_CUDA
void StageObservations();
// The resident path's ingest (see IngestDevice): the observations and the raw-hkl runs built on the
// card from the reflections, exactly as BuildInRangeObservations builds them on the host.
void IngestOnDevice();
#endif
// Compute the dense ASU-group id for the current space group by grouping the (pre-sorted) raw-hkl
image_analysis/scale_merge/RotationScaleMergeGPU.cu
+570 -26
View File
@@ -9,15 +9,19 @@
#include <cmath>
#include <cstdio>
#include <string>
#include <type_traits>
#include <vector>
#include <cuda_runtime.h>
#include <math_constants.h>
#include <cub/device/device_radix_sort.cuh>
#include <cub/device/device_run_length_encode.cuh>
#include <cub/device/device_scan.cuh>
#include "../indexing/CUDAMemHelpers.h"
#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"
#include "../../common/ParallelFor.h"
#include "../../common/Reflection.h"
namespace {
constexpr int BLK = 256;
@@ -1209,6 +1213,12 @@ struct RotationScaleMergeGPU::Impl {
// immutable per-obs
CudaDevicePtr<ObsRecord> obs;
void UploadRecordChunk(int field, int offset, int count, const float *v);
// The per-observation arrays for n observations, the per-frame ones for n_frames and the frame CSR
// (SetPartialsLayout, the resident ingest and CopyObservationsFrom all lay them out the same way).
void AllocObservations(int n);
// (frame_start null: the CSR is left to the caller to fill.)
void AllocFrames(const int32_t *frame_start, const int32_t *frame_count);
void AllocRunScratch();
CudaDevicePtr<float> zeta;
CudaDevicePtr<uint8_t> on_ice, clipped;
CudaDevicePtr<int32_t> frame;
@@ -1279,6 +1289,13 @@ struct RotationScaleMergeGPU::Impl {
CudaDevicePtr<int32_t> s_perm[3], s_seg_start[3];
CudaDevicePtr<double> s_A, s_sw, s_swI, s_tcross, s_tref2, s_cross, s_ref2;
CudaDevicePtr<double> s_w_Is, s_w_Iref, s_Iref; // per term of a subset, in segment order
// the resident ingest (IngestUpload .. IngestSetPartialityAndCorr): each frame's first reflection
// (host, length n_frames + 1), and per observation its raw hkl as the sort key and its delta_phi as
// integrated, until the smoothing has used them
std::vector<int32_t> in_src_start;
CudaDevicePtr<uint64_t> in_key;
CudaDevicePtr<float> in_dphi;
RotationScaleMergeGPU::IngestLayout in_layout;
};
// Set the device this instance's memory lives on for the duration of a call, and put the caller's
@@ -1334,21 +1351,43 @@ void RotationScaleMergeGPU::SetPartialsLayout(int n_obs, int n_frames,
CudaCheck(cudaMemGetInfo(&free_bytes, &total_bytes), "memory info");
if (size_t(std::max(0, n_obs)) * OBS_BYTES > total_bytes)
ThrowTooLargeForGPU(n_obs);
d.Upload(d.frame_start, frame_start, n_frames); d.Upload(d.frame_count, frame_count, n_frames);
const int n = std::max(1, n_obs);
d.obs = d.Alloc<ObsRecord>(n);
d.zeta = d.Alloc<float>(n); d.corr = d.Alloc<float>(n);
d.bkg = d.Alloc<float>(n); d.var_bkg = d.Alloc<float>(n);
d.image_number = d.Alloc<float>(n); d.d_obs = d.Alloc<float>(n);
d.px_obs = d.Alloc<float>(n); d.py_obs = d.Alloc<float>(n);
d.frame = d.Alloc<int32_t>(n);
d.on_ice = d.Alloc<uint8_t>(n);
d.clipped = d.Alloc<uint8_t>(n);
d.g = d.Alloc<double>(n_frames);
d.scaled = d.Alloc<uint8_t>(n_frames);
d.info = d.Alloc<double>(n_frames);
d.cc = d.Alloc<double>(std::max(1, n_frames));
d.cc_n = d.Alloc<int64_t>(std::max(1, n_frames));
d.AllocFrames(frame_start, frame_count);
d.AllocObservations(n_obs);
}
void RotationScaleMergeGPU::Impl::AllocObservations(int n_in) {
const int n = std::max(1, n_in);
obs = Alloc<ObsRecord>(n);
zeta = Alloc<float>(n); corr = Alloc<float>(n);
bkg = Alloc<float>(n); var_bkg = Alloc<float>(n);
image_number = Alloc<float>(n); d_obs = Alloc<float>(n);
px_obs = Alloc<float>(n); py_obs = Alloc<float>(n);
frame = Alloc<int32_t>(n);
on_ice = Alloc<uint8_t>(n);
clipped = Alloc<uint8_t>(n);
}
void RotationScaleMergeGPU::Impl::AllocFrames(const int32_t *fstart, const int32_t *fcount) {
frame_start = Alloc<int32_t>(std::max(1, n_frames));
frame_count = Alloc<int32_t>(std::max(1, n_frames));
if (n_frames > 0 && fstart) {
CopyAndWait(frame_start.get(), fstart, size_t(n_frames) * sizeof(int32_t), cudaMemcpyHostToDevice, s(),
"frame start");
CopyAndWait(frame_count.get(), fcount, size_t(n_frames) * sizeof(int32_t), cudaMemcpyHostToDevice, s(),
"frame count");
}
g = Alloc<double>(std::max(1, n_frames));
scaled = Alloc<uint8_t>(std::max(1, n_frames));
info = Alloc<double>(std::max(1, n_frames));
cc = Alloc<double>(std::max(1, n_frames));
cc_n = Alloc<int64_t>(std::max(1, n_frames));
}
void RotationScaleMergeGPU::Impl::AllocRunScratch() {
rr_group = Alloc<int32_t>(std::max(1, n_runs));
rr_nevents = Alloc<int32_t>(std::max(1, n_runs));
rr_noverloaded = Alloc<int32_t>(std::max(1, n_runs));
rr_offset = Alloc<int32_t>(std::max(1, n_runs));
}
namespace {
@@ -1397,15 +1436,31 @@ void RotationScaleMergeGPU::CopyObservationsFrom(const RotationScaleMergeGPU &fr
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const auto &f = *from.impl_;
const auto copy = [&](auto &dst, const auto &src) {
d.n_obs = f.n_obs;
d.n_frames = f.n_frames;
d.n_runs = f.n_runs;
d.n_perm = f.n_perm;
d.AllocFrames(nullptr, nullptr);
d.AllocObservations(d.n_obs);
const auto copy = [&](auto &dst, const auto &src, size_t n) {
CudaCheck(cudaMemcpyPeerAsync(dst.get(), d.device, src.get(), f.device,
size_t(d.n_obs) * sizeof(*src.get()), d.s()), "copy observations");
n * sizeof(*src.get()), d.s()), "copy observations");
};
copy(d.obs, f.obs); copy(d.zeta, f.zeta); copy(d.corr, f.corr);
copy(d.bkg, f.bkg); copy(d.var_bkg, f.var_bkg); copy(d.image_number, f.image_number);
copy(d.d_obs, f.d_obs); copy(d.px_obs, f.px_obs); copy(d.py_obs, f.py_obs);
copy(d.frame, f.frame); copy(d.on_ice, f.on_ice); copy(d.clipped, f.clipped);
copy(d.frame_start, f.frame_start, d.n_frames); copy(d.frame_count, f.frame_count, d.n_frames);
const size_t n = d.n_obs;
copy(d.obs, f.obs, n); copy(d.zeta, f.zeta, n); copy(d.corr, f.corr, n);
copy(d.bkg, f.bkg, n); copy(d.var_bkg, f.var_bkg, n); copy(d.image_number, f.image_number, n);
copy(d.d_obs, f.d_obs, n); copy(d.px_obs, f.px_obs, n); copy(d.py_obs, f.py_obs, n);
copy(d.frame, f.frame, n); copy(d.on_ice, f.on_ice, n); copy(d.clipped, f.clipped, n);
const size_t nr = std::max(1, d.n_runs);
d.perm = d.Alloc<int32_t>(std::max(1, d.n_perm));
d.rr_start = d.Alloc<int32_t>(nr); d.rr_count = d.Alloc<int32_t>(nr);
d.rr_h = d.Alloc<int32_t>(nr); d.rr_k = d.Alloc<int32_t>(nr); d.rr_l = d.Alloc<int32_t>(nr);
copy(d.perm, f.perm, d.n_perm);
copy(d.rr_start, f.rr_start, d.n_runs); copy(d.rr_count, f.rr_count, d.n_runs);
copy(d.rr_h, f.rr_h, d.n_runs); copy(d.rr_k, f.rr_k, d.n_runs); copy(d.rr_l, f.rr_l, d.n_runs);
CudaCheck(cudaStreamSynchronize(d.s()), "copy observations");
d.AllocRunScratch();
}
void RotationScaleMergeGPU::SetObsFrame(int offset, int count, const int32_t *frame) {
@@ -1884,11 +1939,7 @@ void RotationScaleMergeGPU::SetRawRuns(int n_runs, int n_perm, const int32_t *pe
d.Upload(d.rr_h, rr_h, n_runs);
d.Upload(d.rr_k, rr_k, n_runs);
d.Upload(d.rr_l, rr_l, n_runs);
d.rr_group = d.Alloc<int32_t>(std::max(1, n_runs));
d.rr_nevents = d.Alloc<int32_t>(std::max(1, n_runs));
d.rr_noverloaded = d.Alloc<int32_t>(std::max(1, n_runs));
d.rr_offset = d.Alloc<int32_t>(std::max(1, n_runs));
d.AllocRunScratch();
}
int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_partiality,
@@ -2376,3 +2427,496 @@ void RotationScaleMergeGPU::SurfaceFitSums(int subset, double *cross, double *re
CopyAndWait(ref2, d.s_ref2.get(), size_t(ncell) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(),
"dl surface ref2");
}
// --- the resident ingest ---
namespace {
constexpr int REFL_WORDS = sizeof(Reflection) / sizeof(uint32_t);
static_assert(sizeof(Reflection) % sizeof(uint32_t) == 0 && alignof(Reflection) == alignof(uint32_t),
"the ingest reads a reflection as whole 32-bit words");
// The raw hkl as one sort key: each index as the sixteen bits the host's SortKey holds, offset so
// that the unsigned order of the key is the signed order of (h, k, l).
__device__ __forceinline__ uint64_t RawHklKey(int32_t h, int32_t k, int32_t l) {
const auto u = [](int32_t v) { return uint64_t(uint16_t(int16_t(v)) ^ 0x8000u); };
return (u(h) << 32) | (u(k) << 16) | u(l);
}
__host__ __device__ __forceinline__ int32_t RawHklIndex(uint64_t key, int shift) {
return int16_t(uint16_t(((key >> shift) & 0xffffu) ^ 0x8000u));
}
// A float as an unsigned key in the order of the float comparison: -0 is +0, as the comparison
// sees it (a NaN never reaches here - the host comparator would not be an order with one).
__device__ __forceinline__ uint32_t FloatOrderKey(float v) {
if (v == 0.0f) v = 0.0f;
const uint32_t u = __float_as_uint(v);
return (u & 0x80000000u) ? ~u : (u | 0x80000000u);
}
struct IngestArrays {
ObsRecord *obs;
float *zeta, *corr, *bkg, *var_bkg, *image_number, *d, *px, *py, *dphi;
int32_t *frame;
uint8_t *on_ice, *clipped;
uint64_t *key;
};
// One block per BLK reflections: the block reads its reflections as whole words (coalesced) into
// shared memory and each thread takes its own apart. The deterministic correction is the product of
// its three factors, in the host's order; the frame's flux is divided out once it is known.
__global__ void IngestExtractKernel(int n, int base, const Reflection *__restrict__ refl, IngestArrays a) {
__shared__ uint32_t buf[BLK * REFL_WORDS];
const int first = blockIdx.x * BLK;
const int cnt = min(BLK, n - first);
const uint32_t *src = reinterpret_cast<const uint32_t *>(refl + first);
for (int w = threadIdx.x; w < cnt * REFL_WORDS; w += BLK)
buf[w] = src[w];
__syncthreads();
if (int(threadIdx.x) >= cnt)
return;
const Reflection &r = *reinterpret_cast<const Reflection *>(buf + threadIdx.x * REFL_WORDS);
const int i = base + first + threadIdx.x;
a.obs[i] = ObsRecord{r.I, r.sigma, r.partiality,
__fmul_rn(__fmul_rn(r.prescaling_corr, r.qe_corr), r.flight_corr)};
a.zeta[i] = r.zeta;
a.corr[i] = r.image_scale_corr;
a.bkg[i] = r.bkg;
a.var_bkg[i] = r.var_bkg;
a.image_number[i] = r.image_number;
a.d[i] = r.d;
a.px[i] = r.predicted_x;
a.py[i] = r.predicted_y;
a.dphi[i] = r.delta_phi_deg;
a.on_ice[i] = r.on_ice_ring ? 1 : 0;
a.clipped[i] = (r.clipped ? 1 : 0) | (r.overloaded ? 2 : 0); // OBS_CLIPPED | OBS_OVERLOADED
a.key[i] = RawHklKey(r.h, r.k, r.l);
}
__global__ void IngestFrameKernel(int n_frames, const int32_t *__restrict__ start, int32_t *__restrict__ frame) {
for (int o = blockIdx.x; o < n_frames; o += gridDim.x)
for (int i = start[o] + threadIdx.x; i < start[o + 1]; i += blockDim.x)
frame[i] = o;
}
// One thread per frame, its reflections in order: the host's sum to the bit.
__global__ void IngestMeanBkgKernel(int n_frames, const int32_t *__restrict__ start,
const float *__restrict__ bkg, double *__restrict__ mean) {
for (int o = blockIdx.x * blockDim.x + threadIdx.x; o < n_frames; o += gridDim.x * blockDim.x) {
double sum = 0.0;
int n = 0;
for (int i = start[o]; i < start[o + 1]; ++i) {
const float b = bkg[i];
if (isfinite(b) && b > 0.0f) { sum = __dadd_rn(sum, double(b)); ++n; }
}
mean[o] = n > 0 ? __ddiv_rn(sum, double(n)) : CUDART_NAN;
}
}
__global__ void IngestFluxKernel(int n, const int32_t *__restrict__ frame, const double *__restrict__ flux,
ObsRecord *__restrict__ obs) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
obs[i].prescaling_corr = __double2float_rn(__ddiv_rn(double(obs[i].prescaling_corr), flux[frame[i]]));
}
__global__ void IngestImageKeyKernel(int n, const float *__restrict__ image_number,
uint32_t *__restrict__ key, int32_t *__restrict__ idx) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {
key[i] = FloatOrderKey(image_number[i]);
idx[i] = i;
}
}
__global__ void IngestGatherKeyKernel(int n, const int32_t *__restrict__ idx, const uint64_t *__restrict__ key,
uint64_t *__restrict__ out) {
for (int j = blockIdx.x * blockDim.x + threadIdx.x; j < n; j += gridDim.x * blockDim.x)
out[j] = key[idx[j]];
}
// Per run: its resolution (that of its first member with a usable one, as the host's run scan takes
// it) and whether the resolution range keeps it.
__global__ void IngestRunKernel(int n_runs, const int32_t *__restrict__ start, const int32_t *__restrict__ count,
const int32_t *__restrict__ perm, const float *__restrict__ d_obs,
bool has_d_min, double d_min, bool has_d_max, double d_max,
float *__restrict__ run_d, int32_t *__restrict__ keep_run) {
for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x) {
float d = CUDART_NAN_F;
for (int p = start[r]; p < start[r] + count[r]; ++p) {
const float dp = d_obs[perm[p]];
if (isfinite(dp) && dp > 0.0f) { d = dp; break; }
}
run_d[r] = d;
keep_run[r] = (isfinite(d) && d > 0.0f && !(has_d_min && d < d_min) && !(has_d_max && d > d_max)) ? 1 : 0;
}
}
__global__ void IngestKeepKernel(int n_runs, const int32_t *__restrict__ start, const int32_t *__restrict__ count,
const int32_t *__restrict__ perm, const int32_t *__restrict__ keep_run,
int32_t *__restrict__ keep) {
for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x)
if (keep_run[r])
for (int p = start[r]; p < start[r] + count[r]; ++p)
keep[perm[p]] = 1; // distinct runs own disjoint observations
}
__global__ void IngestFrameCountKernel(int n_frames, const int32_t *__restrict__ start,
const int32_t *__restrict__ keep, int32_t *__restrict__ count) {
for (int o = blockIdx.x * blockDim.x + threadIdx.x; o < n_frames; o += gridDim.x * blockDim.x) {
int c = 0;
for (int i = start[o]; i < start[o + 1]; ++i) c += keep[i];
count[o] = c;
}
}
// The surviving runs, in order, each owning its sorted members renumbered to the observations they
// became; their raw hkl unpacked from the key.
struct RunCut {
int n_runs;
const int32_t *start, *count, *keep_run, *run_idx, *run_at, *perm, *new_idx;
const uint64_t *key;
const float *d;
int32_t *out_start, *out_count, *out_h, *out_k, *out_l, *out_perm;
float *out_d;
};
__global__ void IngestCutRunsKernel(RunCut c) {
for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < c.n_runs; r += gridDim.x * blockDim.x) {
if (!c.keep_run[r]) continue;
const int rr = c.run_idx[r], at = c.run_at[r];
c.out_start[rr] = at;
c.out_count[rr] = c.count[r];
c.out_h[rr] = RawHklIndex(c.key[r], 32);
c.out_k[rr] = RawHklIndex(c.key[r], 16);
c.out_l[rr] = RawHklIndex(c.key[r], 0);
c.out_d[rr] = c.d[r];
for (int p = 0; p < c.count[r]; ++p)
c.out_perm[at + p] = c.new_idx[c.perm[c.start[r] + p]];
}
}
__global__ void IngestFillKernel(int n, int32_t *__restrict__ v) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
v[i] = 1;
}
__global__ void IngestKeptCountKernel(int n_runs, const int32_t *__restrict__ keep_run,
const int32_t *__restrict__ count, int32_t *__restrict__ kept) {
for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x)
kept[r] = keep_run[r] ? count[r] : 0;
}
template <typename T>
__global__ void IngestCompactKernel(int n, const int32_t *__restrict__ keep, const int32_t *__restrict__ new_idx,
const T *__restrict__ in, T *__restrict__ out) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
if (keep[i]) out[new_idx[i]] = in[i];
}
int Blocks(int n) { return std::max(1, std::min(65535, (n + BLK - 1) / BLK)); }
}
void RotationScaleMergeGPU::IngestUpload(int n_frames, const Reflection *const *refl, const int32_t *count,
size_t nthreads) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
d.n_frames = n_frames;
d.in_src_start.assign(n_frames + 1, 0);
int max_count = 0;
for (int o = 0; o < n_frames; ++o) {
d.in_src_start[o + 1] = d.in_src_start[o] + count[o];
max_count = std::max(max_count, count[o]);
}
const int n = d.in_src_start[n_frames];
d.n_obs = n;
size_t free_bytes = 0, total_bytes = 0;
CudaCheck(cudaMemGetInfo(&free_bytes, &total_bytes), "memory info");
if (size_t(n) * OBS_BYTES > total_bytes)
ThrowTooLargeForGPU(n);
d.AllocObservations(n);
d.in_key = d.Alloc<uint64_t>(std::max(1, n));
d.in_dphi = d.Alloc<float>(std::max(1, n));
CudaDevicePtr<int32_t> src_start;
d.Upload(src_start, d.in_src_start.data(), n_frames + 1);
IngestFrameKernel<<<std::max(1, std::min(n_frames, 65535)), BLK, 0, d.s()>>>(n_frames, src_start.get(),
d.frame.get());
CudaCheck(cudaGetLastError(), "IngestFrameKernel launch");
// Through two pinned buffers of whole frames, each filled on all threads while the other is on
// its way to the card and taken apart there.
constexpr size_t STAGE_BYTES = 64 << 20;
const int cap = std::max(static_cast<int>(STAGE_BYTES / sizeof(Reflection)), max_count);
CudaHostPtr<Reflection> stage[2] = {CudaHostPtr<Reflection>(cap), CudaHostPtr<Reflection>(cap)};
CudaDevicePtr<Reflection> staged[2] = {d.Alloc<Reflection>(cap), d.Alloc<Reflection>(cap)};
CudaEvent copied[2];
bool in_flight[2] = {false, false};
IngestArrays a{d.obs.get(), d.zeta.get(), d.corr.get(), d.bkg.get(), d.var_bkg.get(), d.image_number.get(),
d.d_obs.get(), d.px_obs.get(), d.py_obs.get(), d.in_dphi.get(), d.frame.get(),
d.on_ice.get(), d.clipped.get(), d.in_key.get()};
int b = 0;
for (int f0 = 0; f0 < n_frames; ) {
int f1 = f0 + 1;
while (f1 < n_frames && d.in_src_start[f1 + 1] - d.in_src_start[f0] <= cap)
++f1;
const int base = d.in_src_start[f0];
const int cnt = d.in_src_start[f1] - base;
if (in_flight[b])
CudaCheck(cudaEventSynchronize(copied[b]), "staging buffer");
Reflection *dst = stage[b].get();
ParallelFor(f1 - f0, nthreads, [&](int fo) {
const int o = f0 + fo;
std::copy(refl[o], refl[o] + count[o], dst + (d.in_src_start[o] - base));
});
if (cnt > 0) {
CudaCheck(cudaMemcpyAsync(staged[b].get(), dst, size_t(cnt) * sizeof(Reflection),
cudaMemcpyHostToDevice, d.s()), "upload reflections");
CudaCheck(cudaEventRecord(copied[b], d.s()), "upload reflections");
in_flight[b] = true;
IngestExtractKernel<<<(cnt + BLK - 1) / BLK, BLK, 0, d.s()>>>(cnt, base, staged[b].get(), a);
CudaCheck(cudaGetLastError(), "IngestExtractKernel launch");
}
b ^= 1;
f0 = f1;
}
CudaCheck(cudaStreamSynchronize(d.s()), "upload reflections");
}
void RotationScaleMergeGPU::IngestMeanBackground(double *mean_bkg) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
CudaDevicePtr<int32_t> start;
d.Upload(start, d.in_src_start.data(), d.n_frames + 1);
CudaDevicePtr<double> mean = d.Alloc<double>(std::max(1, d.n_frames));
IngestMeanBkgKernel<<<Blocks(d.n_frames), BLK, 0, d.s()>>>(d.n_frames, start.get(), d.bkg.get(), mean.get());
CudaCheck(cudaGetLastError(), "IngestMeanBkgKernel launch");
if (d.n_frames > 0)
CopyAndWait(mean_bkg, mean.get(), size_t(d.n_frames) * sizeof(double), cudaMemcpyDeviceToHost, d.s(),
"download mean background");
}
int RotationScaleMergeGPU::IngestSort(const double *frame_flux, std::optional<double> d_min,
std::optional<double> d_max) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int n = d.n_obs;
const int na = std::max(1, n);
{
CudaDevicePtr<double> flux;
d.Upload(flux, frame_flux, d.n_frames);
IngestFluxKernel<<<Blocks(n), BLK, 0, d.s()>>>(n, d.frame.get(), flux.get(), d.obs.get());
CudaCheck(cudaGetLastError(), "IngestFluxKernel launch");
CudaCheck(cudaStreamSynchronize(d.s()), "flux");
}
// (h, k, l, image_number, index) is a total order, so two stable radix passes - by the frame position
// with the index as the value, then by the raw hkl - give the one permutation std::sort gives on the
// host, whatever either algorithm does on the way.
const auto sort_pairs = [&](auto *key_in, auto *key_out, const int32_t *val_in, int32_t *val_out, int end_bit) {
size_t bytes = 0;
CudaCheck(cub::DeviceRadixSort::SortPairs(nullptr, bytes, key_in, key_out, val_in, val_out, n, 0, end_bit,
d.s()), "sort temp size");
CudaDevicePtr<uint8_t> temp = d.Alloc<uint8_t>(std::max<size_t>(1, bytes));
CudaCheck(cub::DeviceRadixSort::SortPairs(temp.get(), bytes, key_in, key_out, val_in, val_out, n, 0, end_bit,
d.s()), "sort");
CudaCheck(cudaStreamSynchronize(d.s()), "sort");
};
CudaDevicePtr<int32_t> by_image = d.Alloc<int32_t>(na);
{
CudaDevicePtr<uint32_t> ikey = d.Alloc<uint32_t>(na), ikey_sorted = d.Alloc<uint32_t>(na);
CudaDevicePtr<int32_t> idx = d.Alloc<int32_t>(na);
IngestImageKeyKernel<<<Blocks(n), BLK, 0, d.s()>>>(n, d.image_number.get(), ikey.get(), idx.get());
CudaCheck(cudaGetLastError(), "IngestImageKeyKernel launch");
if (n > 0) sort_pairs(ikey.get(), ikey_sorted.get(), idx.get(), by_image.get(), 32);
}
CudaDevicePtr<int32_t> perm = d.Alloc<int32_t>(na);
CudaDevicePtr<uint64_t> key_sorted = d.Alloc<uint64_t>(na);
{
CudaDevicePtr<uint64_t> key = d.Alloc<uint64_t>(na);
IngestGatherKeyKernel<<<Blocks(n), BLK, 0, d.s()>>>(n, by_image.get(), d.in_key.get(), key.get());
CudaCheck(cudaGetLastError(), "IngestGatherKeyKernel launch");
if (n > 0) sort_pairs(key.get(), key_sorted.get(), by_image.get(), perm.get(), 48);
}
by_image = CudaDevicePtr<int32_t>();
// The raw-hkl runs of the sorted order.
int n_runs = 0;
CudaDevicePtr<uint64_t> run_key = d.Alloc<uint64_t>(na);
CudaDevicePtr<int32_t> run_count = d.Alloc<int32_t>(na);
if (n > 0) {
CudaDevicePtr<int32_t> n_runs_dev = d.Alloc<int32_t>(1);
size_t bytes = 0;
CudaCheck(cub::DeviceRunLengthEncode::Encode(nullptr, bytes, key_sorted.get(), run_key.get(),
run_count.get(), n_runs_dev.get(), n, d.s()), "runs temp size");
CudaDevicePtr<uint8_t> temp = d.Alloc<uint8_t>(std::max<size_t>(1, bytes));
CudaCheck(cub::DeviceRunLengthEncode::Encode(temp.get(), bytes, key_sorted.get(), run_key.get(),
run_count.get(), n_runs_dev.get(), n, d.s()), "runs");
CopyAndWait(&n_runs, n_runs_dev.get(), sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(), "run count");
}
key_sorted = CudaDevicePtr<uint64_t>();
const int nra = std::max(1, n_runs);
{
// Cut down to the runs there are: the encode had to be given room for one per reflection.
CudaDevicePtr<uint64_t> key_exact = d.Alloc<uint64_t>(nra);
CudaDevicePtr<int32_t> count_exact = d.Alloc<int32_t>(nra);
if (n_runs > 0) {
CopyAndWait(key_exact.get(), run_key.get(), size_t(n_runs) * sizeof(uint64_t), cudaMemcpyDeviceToDevice,
d.s(), "runs");
CopyAndWait(count_exact.get(), run_count.get(), size_t(n_runs) * sizeof(int32_t), cudaMemcpyDeviceToDevice,
d.s(), "runs");
}
run_key = std::move(key_exact);
run_count = std::move(count_exact);
}
const auto exclusive_sum = [&](const int32_t *in, int32_t *out, int count) {
if (count == 0) return;
size_t bytes = 0;
CudaCheck(cub::DeviceScan::ExclusiveSum(nullptr, bytes, in, out, count, d.s()), "scan temp size");
CudaDevicePtr<uint8_t> temp = d.Alloc<uint8_t>(std::max<size_t>(1, bytes));
CudaCheck(cub::DeviceScan::ExclusiveSum(temp.get(), bytes, in, out, count, d.s()), "scan");
CudaCheck(cudaStreamSynchronize(d.s()), "scan");
};
CudaDevicePtr<int32_t> run_start = d.Alloc<int32_t>(nra);
exclusive_sum(run_count.get(), run_start.get(), n_runs);
CudaDevicePtr<float> run_d = d.Alloc<float>(nra);
CudaDevicePtr<int32_t> keep_run = d.Alloc<int32_t>(nra);
const bool limited = d_min.has_value() || d_max.has_value();
IngestRunKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(n_runs, run_start.get(), run_count.get(), perm.get(),
d.d_obs.get(), d_min.has_value(), d_min.value_or(0.0),
d_max.has_value(), d_max.value_or(0.0),
run_d.get(), keep_run.get());
CudaCheck(cudaGetLastError(), "IngestRunKernel launch");
// Which observations survive (all of them without a range, as the host keeps them), and where each
// run and each observation lands.
CudaDevicePtr<int32_t> keep = d.Alloc<int32_t>(na);
if (limited) {
CudaCheck(cudaMemsetAsync(keep.get(), 0, size_t(na) * sizeof(int32_t), d.s()), "memset keep");
IngestKeepKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(n_runs, run_start.get(), run_count.get(), perm.get(),
keep_run.get(), keep.get());
CudaCheck(cudaGetLastError(), "IngestKeepKernel launch");
} else {
IngestFillKernel<<<Blocks(n), BLK, 0, d.s()>>>(n, keep.get());
IngestFillKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(n_runs, keep_run.get());
CudaCheck(cudaGetLastError(), "IngestFillKernel launch");
}
CudaDevicePtr<int32_t> new_idx = d.Alloc<int32_t>(na);
exclusive_sum(keep.get(), new_idx.get(), n);
CudaDevicePtr<int32_t> run_idx = d.Alloc<int32_t>(nra), run_at = d.Alloc<int32_t>(nra);
exclusive_sum(keep_run.get(), run_idx.get(), n_runs);
{
CudaDevicePtr<int32_t> kept_count = d.Alloc<int32_t>(nra);
IngestKeptCountKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(n_runs, keep_run.get(), run_count.get(),
kept_count.get());
CudaCheck(cudaGetLastError(), "IngestKeptCountKernel launch");
exclusive_sum(kept_count.get(), run_at.get(), n_runs);
}
int32_t last_keep = 0, last_new_idx = 0, last_keep_run = 0, last_run_idx = 0;
if (n > 0) {
CopyAndWait(&last_keep, keep.get() + n - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(), "count");
CopyAndWait(&last_new_idx, new_idx.get() + n - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(), "count");
}
if (n_runs > 0) {
CopyAndWait(&last_keep_run, keep_run.get() + n_runs - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(),
"count");
CopyAndWait(&last_run_idx, run_idx.get() + n_runs - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(),
"count");
}
const int n_keep = last_new_idx + last_keep;
const int n_keep_run = last_run_idx + last_keep_run;
d.n_runs = n_keep_run;
d.n_perm = n_keep;
const int nka = std::max(1, n_keep), nkra = std::max(1, n_keep_run);
d.perm = d.Alloc<int32_t>(nka);
d.rr_start = d.Alloc<int32_t>(nkra); d.rr_count = d.Alloc<int32_t>(nkra);
d.rr_h = d.Alloc<int32_t>(nkra); d.rr_k = d.Alloc<int32_t>(nkra); d.rr_l = d.Alloc<int32_t>(nkra);
CudaDevicePtr<float> rr_d = d.Alloc<float>(nkra);
RunCut c{n_runs, run_start.get(), run_count.get(), keep_run.get(), run_idx.get(), run_at.get(), perm.get(),
new_idx.get(), run_key.get(), run_d.get(),
d.rr_start.get(), d.rr_count.get(), d.rr_h.get(), d.rr_k.get(), d.rr_l.get(), d.perm.get(), rr_d.get()};
IngestCutRunsKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(c);
CudaCheck(cudaGetLastError(), "IngestCutRunsKernel launch");
// The frame CSR of what survived.
auto &L = d.in_layout;
L = RotationScaleMergeGPU::IngestLayout{};
L.n_reflections = n;
L.n_runs_all = n_runs;
L.frame_count.assign(d.n_frames, 0);
{
CudaDevicePtr<int32_t> start, fcount = d.Alloc<int32_t>(std::max(1, d.n_frames));
d.Upload(start, d.in_src_start.data(), d.n_frames + 1);
IngestFrameCountKernel<<<Blocks(d.n_frames), BLK, 0, d.s()>>>(d.n_frames, start.get(), keep.get(),
fcount.get());
CudaCheck(cudaGetLastError(), "IngestFrameCountKernel launch");
if (d.n_frames > 0)
CopyAndWait(L.frame_count.data(), fcount.get(), size_t(d.n_frames) * sizeof(int32_t),
cudaMemcpyDeviceToHost, d.s(), "download frame count");
}
L.frame_start.assign(d.n_frames, 0);
for (int o = 0, at = 0; o < d.n_frames; ++o) { L.frame_start[o] = at; at += L.frame_count[o]; }
// The observations themselves, cut field by field, so that only one field is ever held twice.
if (n_keep < n) {
const auto cut = [&](auto &field) {
using T = std::remove_reference_t<decltype(*field.get())>;
CudaDevicePtr<T> out = d.Alloc<T>(nka);
IngestCompactKernel<T><<<Blocks(n), BLK, 0, d.s()>>>(n, keep.get(), new_idx.get(), field.get(), out.get());
CudaCheck(cudaGetLastError(), "IngestCompactKernel launch");
CudaCheck(cudaStreamSynchronize(d.s()), "cut");
field = std::move(out);
};
cut(d.obs); cut(d.zeta); cut(d.corr); cut(d.bkg); cut(d.var_bkg); cut(d.image_number); cut(d.d_obs);
cut(d.px_obs); cut(d.py_obs); cut(d.frame); cut(d.on_ice); cut(d.clipped); cut(d.in_key); cut(d.in_dphi);
}
d.n_obs = n_keep;
d.AllocFrames(L.frame_start.data(), L.frame_count.data());
d.AllocRunScratch();
L.run_start.resize(n_keep_run); L.run_count.resize(n_keep_run);
L.run_h.resize(n_keep_run); L.run_k.resize(n_keep_run); L.run_l.resize(n_keep_run); L.run_d.resize(n_keep_run);
const auto get = [&](auto &v, const auto &src) {
if (!v.empty())
CopyAndWait(v.data(), src.get(), v.size() * sizeof(v[0]), cudaMemcpyDeviceToHost, d.s(), "download runs");
};
get(L.run_start, d.rr_start); get(L.run_count, d.rr_count);
get(L.run_h, d.rr_h); get(L.run_k, d.rr_k); get(L.run_l, d.rr_l); get(L.run_d, rr_d);
return n_keep;
}
void RotationScaleMergeGPU::IngestGetLayout(IngestLayout &out) const {
out = impl_->in_layout;
}
void RotationScaleMergeGPU::IngestGetObservations(int32_t *h, int32_t *k, int32_t *l, float *delta_phi,
float *partiality, float *image_number, float *zeta, float *d_out,
float *corr, float *I, float *sigma, int32_t *perm) const {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const size_t n = d.n_obs;
if (n == 0) return;
std::vector<uint64_t> key(n);
std::vector<ObsRecord> obs(n);
CopyAndWait(key.data(), d.in_key.get(), n * sizeof(uint64_t), cudaMemcpyDeviceToHost, d.s(), "download keys");
CopyAndWait(obs.data(), d.obs.get(), n * sizeof(ObsRecord), cudaMemcpyDeviceToHost, d.s(), "download records");
for (size_t i = 0; i < n; ++i) {
h[i] = RawHklIndex(key[i], 32); k[i] = RawHklIndex(key[i], 16); l[i] = RawHklIndex(key[i], 0);
partiality[i] = obs[i].partiality; I[i] = obs[i].I; sigma[i] = obs[i].sigma;
}
CopyAndWait(delta_phi, d.in_dphi.get(), n * sizeof(float), cudaMemcpyDeviceToHost, d.s(), "download dphi");
CopyAndWait(image_number, d.image_number.get(), n * sizeof(float), cudaMemcpyDeviceToHost, d.s(), "download img");
CopyAndWait(zeta, d.zeta.get(), n * sizeof(float), cudaMemcpyDeviceToHost, d.s(), "download zeta");
CopyAndWait(d_out, d.d_obs.get(), n * sizeof(float), cudaMemcpyDeviceToHost, d.s(), "download d");
CopyAndWait(corr, d.corr.get(), n * sizeof(float), cudaMemcpyDeviceToHost, d.s(), "download corr");
CopyAndWait(perm, d.perm.get(), size_t(d.n_perm) * sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(),
"download perm");
}
void RotationScaleMergeGPU::IngestSetPartialityAndCorr(const float *partiality, const float *corr) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
d.UploadRecordChunk(2, 0, d.n_obs, partiality);
if (d.n_obs > 0)
CopyAndWait(d.corr.get(), corr, size_t(d.n_obs) * sizeof(float), cudaMemcpyHostToDevice, d.s(), "upload corr");
d.in_key = CudaDevicePtr<uint64_t>();
d.in_dphi = CudaDevicePtr<float>();
d.in_src_start.clear();
d.in_layout = RotationScaleMergeGPU::IngestLayout{};
}
image_analysis/scale_merge/RotationScaleMergeGPU.h
+42 -3
View File
@@ -10,6 +10,8 @@
#include "ErrorModel.h"
struct Reflection;
// GPU engine for the RotationScaleMerge hot loops. The class keeps the per-observation data resident on
// the device as a structure-of-arrays (coalesced) and runs the scaling loop there. The host keeps the
// one-time raw-hkl sort and the per-space-group ASU keying (gemmi); it hands the GPU the dense group ids
@@ -49,11 +51,48 @@ public:
void SetObsOnIce(int offset, int count, const uint8_t *on_ice);
void SetObsClipped(int offset, int count, const uint8_t *clipped);
// Every per-observation field the uploads above set, copied from `from` (which has them all, on
// its own device or this one) instead of being staged again from the host. Call after
// SetPartialsLayout with the same layout, before anything changes the fields on `from`.
// Every per-observation field the uploads above set, the frame CSR and the raw-hkl runs, copied from
// `from` (which has them all, on its own device or this one) instead of being staged again from the
// host. Call once, in place of SetPartialsLayout, the uploads and SetRawRuns, before anything changes
// the fields on `from`.
void CopyObservationsFrom(const RotationScaleMergeGPU &from);
// --- the resident ingest (RotationScaleMerge::IngestDevice with no observation array on the host) ---
// The steps below build, from the reflections, exactly the observations and raw-hkl runs that
// RotationScaleMerge::BuildInRangeObservations builds on the host, and leave them where the uploads
// above and SetRawRuns would: the reflections cross the bus once, and are sorted and cut on the card.
// The reflections of n_frames frames, frame o's count[o] of them at refl[o], uploaded in frame order;
// every field the ingest and the merge read is kept per reflection.
void IngestUpload(int n_frames, const Reflection *const *refl, const int32_t *count, size_t nthreads);
// Each frame's mean background over its reflections with a finite, positive one, summed in their
// order (NaN where there is none): the incident-flux meter. Length n_frames.
void IngestMeanBackground(double *mean_bkg);
// Divide every reflection's deterministic correction by its frame's flux (length n_frames); order the
// reflections by (raw h,k,l, image_number, index) and cut that order into raw-hkl runs, each with the
// resolution of its first member that has a usable one; and drop the runs outside [d_min, d_max]
// (either may be absent) with their reflections. Returns how many observations are left.
int IngestSort(const double *frame_flux, std::optional<double> d_min, std::optional<double> d_max);
// The layout IngestSort left: the reflection and run counts before the resolution cut, the
// observations' frame CSR, and the runs that survived it.
struct IngestLayout {
int64_t n_reflections = 0;
int n_runs_all = 0;
std::vector<int32_t> frame_start, frame_count;
std::vector<int32_t> run_start, run_count, run_h, run_k, run_l;
std::vector<float> run_d;
};
void IngestGetLayout(IngestLayout &out) const;
// What the host's geometry smoothing reads per observation (length IngestSort's count each): the raw
// index, delta_phi and partiality as integrated, the frame position, zeta, d, image_scale_corr, I and
// sigma; and the runs' permutation of the observations.
void IngestGetObservations(int32_t *h, int32_t *k, int32_t *l, float *delta_phi, float *partiality,
float *image_number, float *zeta, float *d, float *corr, float *I,
float *sigma, int32_t *perm) const;
// The partiality and corr the smoothing left, in place of the integrated ones; the transient
// ingest arrays go.
void IngestSetPartialityAndCorr(const float *partiality, const float *corr);
// Per space group: the dense ASU-group id per obs, and a group-ordered permutation of the obs whose
// group >= 0, with its CSR - so each group's observations are a contiguous, fixed-order segment for
// the reduction. Built on the device from the raw-hkl runs uploaded by SetRawRuns: rawrun_group is
rugnux/RugnuxScaleMerge.cpp
+1 -1
View File
@@ -1074,6 +1074,7 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
// holds more, but a merge that does not fit beside it waits for it to end. On a 16 GB card that
// is the data sets up to about 50 M observations at high multiplicity; the large ones merge
// their arms one after the other.
rsm->IngestDevice();
const int32_t cards = get_gpu_count();
if (!geometry_prepass && !experiment_.GetGemmiSpaceGroup().has_value() && rsm->GetSearchMinZeta() > 0.0
&& cards > 0) {
@@ -1086,7 +1087,6 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
rsm_all->SetFrenchWilson(false);
}
}
rsm->IngestDevice();
if (rsm_all) {
// Its observations are copied from this engine's card arrays, not staged from the
// reflections a second time.
tests/CMakeLists.txt
+1
View File
@@ -155,6 +155,7 @@ ADD_EXECUTABLE(jfjoch_test
ModelScalingTest.cpp
ModelScaleGPUTest.cpp
CorrectionSurfaceGPUTest.cpp
RotationIngestGPUTest.cpp
TwinningAnalysisTest.cpp
TranslationalNCSTest.cpp
RfreeFlagsTest.cpp
tests/RotationIngestGPUTest.cpp
+182
View File
@@ -0,0 +1,182 @@
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include <catch2/catch_all.hpp>
#include "../common/CUDAWrapper.h"
#ifdef JFJOCH_USE_CUDA
#include <algorithm>
#include <cmath>
#include <numeric>
#include <random>
#include <vector>
#include "../common/Reflection.h"
#include "../image_analysis/scale_merge/RotationScaleMergeGPU.h"
namespace {
// Frames of reflections with every awkward case the host ingest orders: repeated hkl on one frame (both
// Ewald intersections), the same frame position twice, unusable resolutions and backgrounds, indices
// that fill the sixteen bits of the sort key.
std::vector<std::vector<Reflection>> MakeFrames(int n_frames, uint32_t seed) {
std::mt19937 rng(seed);
std::uniform_int_distribution<int> idx(-5, 5), count(0, 900), pick(0, 19);
std::uniform_real_distribution<float> u(0.0f, 1.0f);
std::vector<std::vector<Reflection>> frames(n_frames);
for (int f = 0; f < n_frames; ++f) {
const int n = (f == 7) ? 0 : count(rng);
for (int j = 0; j < n; ++j) {
Reflection r{};
r.h = idx(rng); r.k = idx(rng); r.l = idx(rng);
if (pick(rng) == 0) r.h = (j % 2) ? 32767 : -32768;
r.image_number = static_cast<float>(f) + 0.25f * static_cast<float>(pick(rng) % 3);
r.d = 1.0f + 0.1f * static_cast<float>(std::abs(r.h % 7) + std::abs(r.k) + std::abs(r.l));
if (pick(rng) == 1) r.d = NAN;
if (pick(rng) == 2) r.d = 0.0f;
r.I = 100.0f * u(rng);
r.sigma = 1.0f + u(rng);
r.bkg = (pick(rng) == 3) ? NAN : 10.0f * u(rng) - 1.0f;
r.var_bkg = u(rng);
r.prescaling_corr = 0.5f + u(rng);
r.qe_corr = 0.9f + 0.1f * u(rng);
r.flight_corr = 1.0f + 0.01f * u(rng);
r.partiality = u(rng);
r.zeta = u(rng);
r.image_scale_corr = u(rng);
r.delta_phi_deg = u(rng) - 0.5f;
r.predicted_x = 1000.0f * u(rng);
r.predicted_y = 1000.0f * u(rng);
r.on_ice_ring = pick(rng) == 4;
r.clipped = pick(rng) == 5;
r.overloaded = pick(rng) == 6;
frames[f].push_back(r);
}
}
return frames;
}
struct HostKey { int16_t h, k, l; float img; int32_t idx; float d; };
} // namespace
TEST_CASE("RotationIngestGPU_SortAndCutMatchHost", "[RotationScale][gpu]") {
if (get_gpu_count() == 0)
SKIP("no GPU");
const int n_frames = 60;
const auto frames = MakeFrames(n_frames, 17);
std::vector<const Reflection *> ptr(n_frames);
std::vector<int32_t> count(n_frames), src_start(n_frames + 1, 0);
for (int f = 0; f < n_frames; ++f) {
ptr[f] = frames[f].data();
count[f] = static_cast<int32_t>(frames[f].size());
src_start[f + 1] = src_start[f] + count[f];
}
const int n = src_start[n_frames];
std::vector<const Reflection *> all(n);
for (int f = 0; f < n_frames; ++f)
for (int j = 0; j < count[f]; ++j) all[src_start[f] + j] = &frames[f][j];
for (const auto &[d_min, d_max] : {std::pair<std::optional<double>, std::optional<double>>{std::nullopt, std::nullopt},
{1.4, 2.1}, {std::nullopt, 1.8}}) {
RotationScaleMergeGPU gpu(0);
gpu.IngestUpload(n_frames, ptr.data(), count.data(), 4);
// The flux meter, summed per frame in reflection order.
std::vector<double> mean(n_frames);
gpu.IngestMeanBackground(mean.data());
for (int f = 0; f < n_frames; ++f) {
double sum = 0.0;
int c = 0;
for (const auto &r : frames[f])
if (std::isfinite(r.bkg) && r.bkg > 0.0f) { sum += r.bkg; ++c; }
if (c > 0) CHECK(mean[f] == sum / c);
else CHECK(std::isnan(mean[f]));
}
std::vector<double> flux(n_frames);
for (int f = 0; f < n_frames; ++f) flux[f] = 0.9 + 0.01 * f;
const int n_keep = gpu.IngestSort(flux.data(), d_min, d_max);
// The host ingest: one total order, its raw-hkl runs, each with its first usable resolution, and
// the runs the range keeps with their observations.
std::vector<HostKey> key(n);
for (int i = 0; i < n; ++i)
key[i] = {static_cast<int16_t>(all[i]->h), static_cast<int16_t>(all[i]->k), static_cast<int16_t>(all[i]->l),
all[i]->image_number, i, all[i]->d};
std::sort(key.begin(), key.end(), [](const HostKey &a, const HostKey &b) {
if (a.h != b.h) return a.h < b.h;
if (a.k != b.k) return a.k < b.k;
if (a.l != b.l) return a.l < b.l;
if (a.img != b.img) return a.img < b.img;
return a.idx < b.idx;
});
const bool limited = d_min || d_max;
std::vector<uint8_t> keep(n, limited ? 0 : 1);
struct Run { int start, count; int32_t h, k, l; float d; };
std::vector<Run> runs;
int n_runs_all = 0;
for (int i = 0; i < n; ) {
int j = i;
float d = NAN;
while (j < n && key[j].h == key[i].h && key[j].k == key[i].k && key[j].l == key[i].l) {
if (!std::isfinite(d) && std::isfinite(key[j].d) && key[j].d > 0.0f) d = key[j].d;
++j;
}
++n_runs_all;
const bool in = !limited || (std::isfinite(d) && d > 0.0f && !(d_min && d < *d_min) && !(d_max && d > *d_max));
if (in) {
runs.push_back({i, j - i, key[i].h, key[i].k, key[i].l, d});
for (int p = i; p < j; ++p) keep[key[p].idx] = 1;
}
i = j;
}
std::vector<int32_t> new_idx(n, -1);
int at = 0;
for (int i = 0; i < n; ++i) if (keep[i]) new_idx[i] = at++;
REQUIRE(n_keep == at);
RotationScaleMergeGPU::IngestLayout L;
gpu.IngestGetLayout(L);
CHECK(L.n_reflections == n);
CHECK(L.n_runs_all == n_runs_all);
REQUIRE(L.run_start.size() == runs.size());
std::vector<int32_t> h(n_keep), k(n_keep), l(n_keep), perm(n_keep);
std::vector<float> dphi(n_keep), part(n_keep), img(n_keep), zeta(n_keep), d(n_keep), corr(n_keep),
I(n_keep), sigma(n_keep);
gpu.IngestGetObservations(h.data(), k.data(), l.data(), dphi.data(), part.data(), img.data(), zeta.data(),
d.data(), corr.data(), I.data(), sigma.data(), perm.data());
int at_run = 0;
for (size_t r = 0; r < runs.size(); ++r) {
CHECK(L.run_start[r] == at_run);
CHECK(L.run_count[r] == runs[r].count);
CHECK(L.run_h[r] == runs[r].h);
CHECK(L.run_k[r] == runs[r].k);
CHECK(L.run_l[r] == runs[r].l);
CHECK((L.run_d[r] == runs[r].d || (std::isnan(L.run_d[r]) && std::isnan(runs[r].d))));
for (int p = 0; p < runs[r].count; ++p)
CHECK(perm[at_run + p] == new_idx[key[runs[r].start + p].idx]);
at_run += runs[r].count;
}
for (int f = 0, s = 0; f < n_frames; ++f) {
int c = 0;
for (int j = 0; j < count[f]; ++j) c += keep[src_start[f] + j];
CHECK(L.frame_start[f] == s);
CHECK(L.frame_count[f] == c);
s += c;
}
for (int i = 0; i < n; ++i) {
if (!keep[i]) continue;
const int o = new_idx[i];
const Reflection &r = *all[i];
CHECK(h[o] == static_cast<int16_t>(r.h));
CHECK(I[o] == r.I);
CHECK(corr[o] == r.image_scale_corr);
CHECK(dphi[o] == r.delta_phi_deg);
CHECK(img[o] == r.image_number);
}
}
}
#endif