diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index df4f18746..990692ac4 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -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(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(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(frame_count[o])); - const int bufsize = std::max(1, std::max(STAGE_SLICE, max_frame_obs)); - std::unique_ptr fbuf[NFLOAT]; - for (auto &b : fbuf) b.reset(new float[bufsize]); - const std::unique_ptr frm(new int32_t[bufsize]); - const std::unique_ptr 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( - 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 refl(n_frames); + std::vector count(n_frames); + for (int o = 0; o < n_frames; ++o) { + refl[o] = partials_out[o].reflections.data(); + count[o] = static_cast(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 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(layout.n_reflections); + ingest_runs = static_cast(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 h(n), k(n), l(n); + const std::unique_ptr 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 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(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().swap(perm); } #endif @@ -683,7 +697,13 @@ size_t RotationScaleMerge::DeviceBytes() const { const auto events = static_cast(resident_ingest ? rocking_events_at_ingest : n_partials_obs); const size_t bytes = BYTES_PER_OBS * static_cast(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::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(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(rawrun_start.size()), static_cast(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().swap(ingest_keep); - std::vector().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 &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 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(partials_out[o].reflections.size()); } } - const std::vector &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 &ke if (x.GetBandwidthFWHM().value_or(0.0f) > 0.0f) ingest_d.reset(new float[n_keep]); std::atomic 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( - 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 &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 &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 &rc = sg.recip[o]; - const Coord p0 = rc[0] * static_cast(r.h) + rc[1] * static_cast(r.k) - + rc[2] * static_cast(r.l); - const float phi_near = r.delta_phi_deg * static_cast(PI) / 180.0f; + const Coord p0 = rc[0] * static_cast(h) + rc[1] * static_cast(k) + + rc[2] * static_cast(l); + const float phi_near = delta_phi_deg * static_cast(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(PI); } diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 5ad12d9fb..42443704b 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -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 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 &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 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 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 ingest_keep; - std::vector 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 frame_start, frame_count; // CSR ranges of `partials` per frame std::vector frame_cell_ok; // per-frame cell-consistency mask (1 = kept) std::vector 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 &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 diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 793da8bc2..8892bd770 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -9,15 +9,19 @@ #include #include #include +#include #include #include +#include #include +#include #include #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 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 zeta; CudaDevicePtr on_ice, clipped; CudaDevicePtr frame; @@ -1279,6 +1289,13 @@ struct RotationScaleMergeGPU::Impl { CudaDevicePtr s_perm[3], s_seg_start[3]; CudaDevicePtr s_A, s_sw, s_swI, s_tcross, s_tref2, s_cross, s_ref2; CudaDevicePtr 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 in_src_start; + CudaDevicePtr in_key; + CudaDevicePtr 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(n); - d.zeta = d.Alloc(n); d.corr = d.Alloc(n); - d.bkg = d.Alloc(n); d.var_bkg = d.Alloc(n); - d.image_number = d.Alloc(n); d.d_obs = d.Alloc(n); - d.px_obs = d.Alloc(n); d.py_obs = d.Alloc(n); - d.frame = d.Alloc(n); - d.on_ice = d.Alloc(n); - d.clipped = d.Alloc(n); - d.g = d.Alloc(n_frames); - d.scaled = d.Alloc(n_frames); - d.info = d.Alloc(n_frames); - d.cc = d.Alloc(std::max(1, n_frames)); - d.cc_n = d.Alloc(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(n); + zeta = Alloc(n); corr = Alloc(n); + bkg = Alloc(n); var_bkg = Alloc(n); + image_number = Alloc(n); d_obs = Alloc(n); + px_obs = Alloc(n); py_obs = Alloc(n); + frame = Alloc(n); + on_ice = Alloc(n); + clipped = Alloc(n); +} + +void RotationScaleMergeGPU::Impl::AllocFrames(const int32_t *fstart, const int32_t *fcount) { + frame_start = Alloc(std::max(1, n_frames)); + frame_count = Alloc(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(std::max(1, n_frames)); + scaled = Alloc(std::max(1, n_frames)); + info = Alloc(std::max(1, n_frames)); + cc = Alloc(std::max(1, n_frames)); + cc_n = Alloc(std::max(1, n_frames)); +} + +void RotationScaleMergeGPU::Impl::AllocRunScratch() { + rr_group = Alloc(std::max(1, n_runs)); + rr_nevents = Alloc(std::max(1, n_runs)); + rr_noverloaded = Alloc(std::max(1, n_runs)); + rr_offset = Alloc(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(std::max(1, d.n_perm)); + d.rr_start = d.Alloc(nr); d.rr_count = d.Alloc(nr); + d.rr_h = d.Alloc(nr); d.rr_k = d.Alloc(nr); d.rr_l = d.Alloc(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(std::max(1, n_runs)); - d.rr_nevents = d.Alloc(std::max(1, n_runs)); - d.rr_noverloaded = d.Alloc(std::max(1, n_runs)); - d.rr_offset = d.Alloc(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(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(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 + __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(std::max(1, n)); + d.in_dphi = d.Alloc(std::max(1, n)); + + CudaDevicePtr src_start; + d.Upload(src_start, d.in_src_start.data(), n_frames + 1); + IngestFrameKernel<<>>(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(STAGE_BYTES / sizeof(Reflection)), max_count); + CudaHostPtr stage[2] = {CudaHostPtr(cap), CudaHostPtr(cap)}; + CudaDevicePtr staged[2] = {d.Alloc(cap), d.Alloc(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 start; + d.Upload(start, d.in_src_start.data(), d.n_frames + 1); + CudaDevicePtr mean = d.Alloc(std::max(1, d.n_frames)); + IngestMeanBkgKernel<<>>(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 d_min, + std::optional 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 flux; + d.Upload(flux, frame_flux, d.n_frames); + IngestFluxKernel<<>>(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 temp = d.Alloc(std::max(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 by_image = d.Alloc(na); + { + CudaDevicePtr ikey = d.Alloc(na), ikey_sorted = d.Alloc(na); + CudaDevicePtr idx = d.Alloc(na); + IngestImageKeyKernel<<>>(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 perm = d.Alloc(na); + CudaDevicePtr key_sorted = d.Alloc(na); + { + CudaDevicePtr key = d.Alloc(na); + IngestGatherKeyKernel<<>>(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(); + + // The raw-hkl runs of the sorted order. + int n_runs = 0; + CudaDevicePtr run_key = d.Alloc(na); + CudaDevicePtr run_count = d.Alloc(na); + if (n > 0) { + CudaDevicePtr n_runs_dev = d.Alloc(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 temp = d.Alloc(std::max(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(); + 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 key_exact = d.Alloc(nra); + CudaDevicePtr count_exact = d.Alloc(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 temp = d.Alloc(std::max(1, bytes)); + CudaCheck(cub::DeviceScan::ExclusiveSum(temp.get(), bytes, in, out, count, d.s()), "scan"); + CudaCheck(cudaStreamSynchronize(d.s()), "scan"); + }; + CudaDevicePtr run_start = d.Alloc(nra); + exclusive_sum(run_count.get(), run_start.get(), n_runs); + CudaDevicePtr run_d = d.Alloc(nra); + CudaDevicePtr keep_run = d.Alloc(nra); + const bool limited = d_min.has_value() || d_max.has_value(); + IngestRunKernel<<>>(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 keep = d.Alloc(na); + if (limited) { + CudaCheck(cudaMemsetAsync(keep.get(), 0, size_t(na) * sizeof(int32_t), d.s()), "memset keep"); + IngestKeepKernel<<>>(n_runs, run_start.get(), run_count.get(), perm.get(), + keep_run.get(), keep.get()); + CudaCheck(cudaGetLastError(), "IngestKeepKernel launch"); + } else { + IngestFillKernel<<>>(n, keep.get()); + IngestFillKernel<<>>(n_runs, keep_run.get()); + CudaCheck(cudaGetLastError(), "IngestFillKernel launch"); + } + CudaDevicePtr new_idx = d.Alloc(na); + exclusive_sum(keep.get(), new_idx.get(), n); + CudaDevicePtr run_idx = d.Alloc(nra), run_at = d.Alloc(nra); + exclusive_sum(keep_run.get(), run_idx.get(), n_runs); + { + CudaDevicePtr kept_count = d.Alloc(nra); + IngestKeptCountKernel<<>>(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(nka); + d.rr_start = d.Alloc(nkra); d.rr_count = d.Alloc(nkra); + d.rr_h = d.Alloc(nkra); d.rr_k = d.Alloc(nkra); d.rr_l = d.Alloc(nkra); + CudaDevicePtr rr_d = d.Alloc(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<<>>(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 start, fcount = d.Alloc(std::max(1, d.n_frames)); + d.Upload(start, d.in_src_start.data(), d.n_frames + 1); + IngestFrameCountKernel<<>>(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; + CudaDevicePtr out = d.Alloc(nka); + IngestCompactKernel<<>>(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 key(n); + std::vector 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(); + d.in_dphi = CudaDevicePtr(); + d.in_src_start.clear(); + d.in_layout = RotationScaleMergeGPU::IngestLayout{}; +} diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.h b/image_analysis/scale_merge/RotationScaleMergeGPU.h index 330466d80..a4fd45eda 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.h +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.h @@ -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 d_min, std::optional 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 frame_start, frame_count; + std::vector run_start, run_count, run_h, run_k, run_l; + std::vector 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 diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index 3ee74f930..4431125b4 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -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. diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 1ff1356f4..354c85b6c 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -155,6 +155,7 @@ ADD_EXECUTABLE(jfjoch_test ModelScalingTest.cpp ModelScaleGPUTest.cpp CorrectionSurfaceGPUTest.cpp + RotationIngestGPUTest.cpp TwinningAnalysisTest.cpp TranslationalNCSTest.cpp RfreeFlagsTest.cpp diff --git a/tests/RotationIngestGPUTest.cpp b/tests/RotationIngestGPUTest.cpp new file mode 100644 index 000000000..4e84ed4f8 --- /dev/null +++ b/tests/RotationIngestGPUTest.cpp @@ -0,0 +1,182 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include +#include "../common/CUDAWrapper.h" + +#ifdef JFJOCH_USE_CUDA + +#include +#include +#include +#include +#include + +#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> MakeFrames(int n_frames, uint32_t seed) { + std::mt19937 rng(seed); + std::uniform_int_distribution idx(-5, 5), count(0, 900), pick(0, 19); + std::uniform_real_distribution u(0.0f, 1.0f); + std::vector> 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(f) + 0.25f * static_cast(pick(rng) % 3); + r.d = 1.0f + 0.1f * static_cast(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 ptr(n_frames); + std::vector 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(frames[f].size()); + src_start[f + 1] = src_start[f] + count[f]; + } + const int n = src_start[n_frames]; + std::vector 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>{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 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 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 key(n); + for (int i = 0; i < n; ++i) + key[i] = {static_cast(all[i]->h), static_cast(all[i]->k), static_cast(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 keep(n, limited ? 0 : 1); + struct Run { int start, count; int32_t h, k, l; float d; }; + std::vector 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 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 h(n_keep), k(n_keep), l(n_keep), perm(n_keep); + std::vector 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(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