RotationScaleMerge: the resident ingest uploads the reflections once

The device ingest staged every reflection to the card twice - once for what the sort and the
cut read, once again for the observations the cut kept - and the second staging was a full
copy of the reflections into pinned memory and over the bus. Where the whole set of reflections
and the ingest's own peak beside it fit a quarter of the card (decided on the card's total
memory, so the answer depends on the card and not on the moment), the first upload now goes
into one resident copy, and the observations are taken from it in a single extract; the second
staging stays as the fallback where the copy does not fit, and where the card reads host
memory in place. The extract kernel also hands over the predicted positions, which the host
had gathered itself from a downloaded index array, and the two pinned staging buffers are kept
across the ingest instead of being allocated per staging. Same observations, same order, same
bits.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
This commit is contained in:
leonarski_fandClaude Fable 5.1 committed 2026-10-26 20:58:10 +01:00
1 parent 00f8c8ba4b
commit bbcae8ae7f
3 files changed
+64 -34

No files matched your search

@@ -733,8 +733,9 @@ void RotationScaleMerge::IngestHost() {
#ifdef JFJOCH_USE_CUDA
// The resident path's ingest. The reflections go to the card 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); once the scaling is known to fit, they go to the
// card a second time for the observations the cut kept (IngestObservations) - the observations, the runs
// (RotationScaleMergeGPU::IngestUpload .. IngestSort); once the scaling is known to fit, the observations
// the cut kept are taken from them on the card, or from a second upload where the first could not stay
// (IngestObservations) - 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() {
@@ -1330,9 +1330,16 @@ struct RotationScaleMergeGPU::Impl {
CudaDevicePtr<int32_t> in_new_idx;
int in_n_keep = 0;
CudaDevicePtr<float> in_dphi;
// The reflections of the ingest's frames, through two pinned buffers of whole frames, each filled on
// all threads while the other is on its way to the card; take(cnt, base, src) takes cnt of them,
// the first being reflection base, apart on the card.
// Every reflection as uploaded, kept on the card from the upload to the extract where the copy fits
// beside the ingest's peak (IngestUpload), so that the reflections cross the bus once; empty where
// it does not, and the extract stages them a second time.
CudaDevicePtr<Reflection> in_refl;
// The reflections of the ingest's frames, through two pinned buffers of whole frames (kept from the
// upload to the extract, which may stage them again), each filled on all threads while the other
// is on its way to the card; take(cnt, base, src) takes cnt of them, the first being reflection
// base, apart on the card.
CudaHostPtr<Reflection> in_stage[2];
int in_stage_cap = 0;
template <typename Take>
void StageReflections(const Reflection *const *refl, const int32_t *count, size_t nthreads, Take take);
RotationScaleMergeGPU::IngestLayout in_layout;
@@ -2512,7 +2519,7 @@ namespace {
struct IngestArrays {
ObsRecord *obs;
float *zeta, *corr, *bkg, *var_bkg, *image_number, *d, *dphi;
float *zeta, *corr, *bkg, *var_bkg, *image_number, *d, *dphi, *px, *py;
uint8_t *flags;
float d_max_scale; // coarser than this is kept out of the scale fit (OBS_LOW_RES_FLAG)
const int32_t *new_idx; // the cut's place of each reflection (an exclusive prefix over the kept)
@@ -2541,6 +2548,8 @@ namespace {
a.image_number[i] = r.image_number;
a.d[i] = r.d;
a.dphi[i] = r.delta_phi_deg;
a.px[i] = r.predicted_x;
a.py[i] = r.predicted_y;
a.flags[i] = (r.on_ice_ring ? OBS_ON_ICE_FLAG : 0) | (r.clipped ? OBS_CLIPPED_FLAG : 0)
| (r.overloaded ? OBS_OVERLOADED_FLAG : 0) | (r.d > a.d_max_scale ? OBS_LOW_RES_FLAG : 0);
}
@@ -2842,13 +2851,18 @@ void RotationScaleMergeGPU::Impl::StageReflections(const Reflection *const *refl
for (int o = 0; o < n_frames; ++o)
max_count = std::max(max_count, count[o]);
// Where the card reads host memory directly (gpu_direct_host_read()), the kernel takes the pinned
// buffer apart in place, without the upload.
// buffer apart in place, without the upload. Where the reflections stay resident (in_refl), each
// buffer goes to its place in the copy instead of a staging buffer of its own on the card.
constexpr size_t STAGE_BYTES = 64 << 20;
const bool direct = gpu_direct_host_read();
const int cap = std::max(static_cast<int>(STAGE_BYTES / sizeof(Reflection)), max_count);
CudaHostPtr<Reflection> stage[2] = {CudaHostPtr<Reflection>(cap), CudaHostPtr<Reflection>(cap)};
if (cap > in_stage_cap) {
in_stage[0] = CudaHostPtr<Reflection>(cap);
in_stage[1] = CudaHostPtr<Reflection>(cap);
in_stage_cap = cap;
}
CudaDevicePtr<Reflection> staged[2];
if (!direct) {
if (!direct && !in_refl) {
staged[0] = Alloc<Reflection>(cap);
staged[1] = Alloc<Reflection>(cap);
}
@@ -2863,7 +2877,7 @@ void RotationScaleMergeGPU::Impl::StageReflections(const Reflection *const *refl
const int cnt = in_src_start[f1] - base;
if (in_flight[b])
CudaCheck(cudaEventSynchronize(copied[b]), "staging buffer");
Reflection *dst = stage[b].get();
Reflection *dst = in_stage[b].get();
ParallelFor(f1 - f0, nthreads, [&](int fo) {
const int o = f0 + fo;
std::copy(refl[o], refl[o] + count[o], dst + (in_src_start[o] - base));
@@ -2871,10 +2885,11 @@ void RotationScaleMergeGPU::Impl::StageReflections(const Reflection *const *refl
if (cnt > 0) {
const Reflection *src = dst;
if (!direct) {
CudaCheck(cudaMemcpyAsync(staged[b].get(), dst, size_t(cnt) * sizeof(Reflection),
cudaMemcpyHostToDevice, s()), "upload reflections");
Reflection *to = in_refl ? in_refl.get() + base : staged[b].get();
CudaCheck(cudaMemcpyAsync(to, dst, size_t(cnt) * sizeof(Reflection), cudaMemcpyHostToDevice, s()),
"upload reflections");
CudaCheck(cudaEventRecord(copied[b], s()), "upload reflections");
src = staged[b].get();
src = to;
}
in_flight[b] = true;
take(cnt, base, src);
@@ -2904,9 +2919,18 @@ void RotationScaleMergeGPU::IngestUpload(int n_frames, const Reflection *const *
ThrowTooLargeForGPU(n);
// See AllocObservations: these are the first of the scaling's buffers.
cuda_trim_allocation_pool();
const int na = std::max(1, n);
// The reflections stay on the card until the observations are taken from them where that copy and
// the ingest's own peak beside it (67 bytes per reflection, see RotationScaleMerge::DeviceBytes)
// fit a quarter of the card - decided on the card's total, so the answer does not depend on the
// moment, and with room for a second engine's ingest beside this one. Not where the card reads
// the host's memory in place: there the copy would be the upload it saves.
constexpr size_t INGEST_PEAK_BYTES_PER_REFLECTION = 67;
d.in_refl = {};
if (!gpu_direct_host_read() && size_t(n) * (sizeof(Reflection) + INGEST_PEAK_BYTES_PER_REFLECTION) <= total_bytes / 4)
d.in_refl = d.Alloc<Reflection>(na);
// Only what the sort and the cut read; the observations are taken from the reflections once the cut
// has decided which they are (IngestObservations), so the two are never on the card together.
const int na = std::max(1, n);
d.in_img = d.Alloc<float>(na); d.in_d = d.Alloc<float>(na); d.in_bkg = d.Alloc<float>(na);
d.in_key = d.Alloc<uint64_t>(na);
const IngestKeys a{d.in_img.get(), d.in_d.get(), d.in_bkg.get(), d.in_key.get()};
@@ -3125,30 +3149,34 @@ void RotationScaleMergeGPU::IngestObservations(const Reflection *const *refl, co
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int n_keep = d.in_n_keep;
const int nka = std::max(1, n_keep);
d.AllocObservations(n_keep);
d.in_dphi = d.Alloc<float>(std::max(1, n_keep));
d.in_dphi = d.Alloc<float>(nka);
// The predicted positions, which only the fulls' and only on the host read, go there straight away.
CudaDevicePtr<float> dev_px = d.Alloc<float>(nka), dev_py = d.Alloc<float>(nka);
const int n = d.in_src_start[d.n_frames];
const 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.in_dphi.get(), d.flags.get(), d_max_scale, d.in_new_idx.get(),
d.in_src_start[d.n_frames], n_keep};
d.StageReflections(refl, count, nthreads, [&](int cnt, int base, const Reflection *src) {
d.d_obs.get(), d.in_dphi.get(), dev_px.get(), dev_py.get(), d.flags.get(), d_max_scale,
d.in_new_idx.get(), n, n_keep};
const auto extract = [&](int cnt, int base, const Reflection *src) {
IngestExtractKernel<<<(cnt + BLK - 1) / BLK, BLK, 0, d.s()>>>(cnt, base, src, a);
CudaCheck(cudaGetLastError(), "IngestExtractKernel launch");
});
// The predicted positions, which only the fulls' and only on the host read, stay there.
const int n = d.in_src_start[d.n_frames];
std::vector<int32_t> new_idx(n);
if (n > 0)
CopyAndWait(new_idx.data(), d.in_new_idx.get(), size_t(n) * sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(),
"download places");
};
if (d.in_refl) {
if (n > 0)
extract(n, 0, d.in_refl.get());
CudaCheck(cudaStreamSynchronize(d.s()), "extract reflections");
d.in_refl = {};
} else {
d.StageReflections(refl, count, nthreads, extract);
}
d.in_new_idx = {};
ParallelFor(d.n_frames, nthreads, [&](int o) {
for (int j = 0; j < count[o]; ++j) {
const int src = d.in_src_start[o] + j, i = new_idx[src];
if ((src + 1 < n ? new_idx[src + 1] : n_keep) == i) continue;
px[i] = refl[o][j].predicted_x;
py[i] = refl[o][j].predicted_y;
}
});
d.in_stage[0] = {}; d.in_stage[1] = {}; d.in_stage_cap = 0;
if (n_keep > 0) {
CopyAndWait(px, dev_px.get(), size_t(n_keep) * sizeof(float), cudaMemcpyDeviceToHost, d.s(), "download positions");
CopyAndWait(py, dev_py.get(), size_t(n_keep) * sizeof(float), cudaMemcpyDeviceToHost, d.s(), "download positions");
}
dev_px = {}; dev_py = {};
// The kept observations are in frame order, as the reflections were (IngestSort laid out the CSR).
CudaDevicePtr<double> flux;
@@ -61,7 +61,8 @@ public:
// 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 are sorted and cut on the card, and cross the bus
// twice - once for what the sort and the cut read, once for the observations the cut kept.
// once where their copy fits beside the ingest (IngestUpload), else twice - once for what the sort
// and the cut read, once for the observations the cut kept.
// The reflections of n_frames frames, frame o's count[o] of them at refl[o], uploaded in frame order;
// what the sort, the cut and the background below read is kept per reflection.