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

Make the partials order total, and hoist 1/sigma out of the IRLS loop

Browse Source 
The sort that orders every observation by (h,k,l,image_number) was not a total
order: two observations can genuinely share all four. The predictor emits BOTH
intersections of a reflection's rotation circle with the Ewald sphere, and near
the blind region - where zeta is smallest - the two are close enough in angle
that both are accepted on the same frame. Which of them came first was then
whatever the sort happened to produce.

That was observable. The combine takes on_ice from the FIRST member of a
rocking event, so the order decided whether a full was flagged as ice at all,
and its per-event sums are floating point, so it moved intensities in their
last bits. The observation's own index is now the final key, which orders them
by arrival - and, more usefully, makes the order unique, so it no longer
depends on which algorithm sorted it.

sigma never changes once it is uploaded, so 1/sigma is the same in all thirty
IRLS iterations of all three scaling iterations of all five scaling passes. It
was being recomputed every time: a 64-bit reciprocal is a hardware estimate
plus five refinement steps, and the profile put the three divisions in that
loop at 21 of its 31 double-precision instructions. It is computed once now, in
the pass that already streams every observation. The CPU has always hoisted it;
this is the GPU catching up. Same expression on the same operand, so the value
is what the loop used to compute, bit for bit.

Also: PrepScaleObsKernel is not a grid-stride loop, but the scale-fulls path
capped its grid at 65535 blocks like the grid-stride kernels around it. Above
16.8 million fulls that silently left the tail of sco_coeff/sco_ok stale. No
dataset here reaches it; the cap is simply wrong for that kernel.

And the AoS-to-SoA staging that feeds the GPU - the widest pass in Ingest,
reading an 80-byte struct and writing fourteen arrays out of it - ran on one
thread.

Full 24-crystal battery: same space group on all 24, none failed, one crystal
moved R_meas by 0.8 points with CC unchanged (it moves by that much between
runs of an identical binary). 15m32s -> 13m35s.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
jungfrau
2026-08-15 18:56:26 -04:00
co-authored by Claude Opus 5
parent e11c2a2b20
commit 56512414aa
2 changed files with 50 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/scale_merge/RotationScaleMerge.cpp
+26 -54
View File
@@ -1,6 +1,7 @@
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "../../common/ParallelFor.h"
#include "RotationScaleMerge.h"
#include <algorithm>
@@ -138,48 +139,6 @@ namespace {
return G;
}
// Run fn(i) for i in [0, n) over `nthreads` workers pulling from a shared atomic counter - the same
// self-load-balancing pattern the rest of the codebase uses (heavy frames don't stall light ones).
// Work-stealing per-item parallel: one atomic fetch per item. Use ONLY when the per-item work is
// heavy and uneven (e.g. per-frame fits) - the atomic amortises. For millions of tiny uniform items
// use ParallelChunks instead; a per-item atomic there is pure contention.
template <typename Fn>
void ParallelFor(int n, size_t nthreads, Fn fn) {
if (n <= 0) return;
if (nthreads <= 1 || n == 1) {
for (int i = 0; i < n; ++i) fn(i);
return;
}
const size_t local = std::min(nthreads, static_cast<size_t>(n));
std::atomic<int> next = 0;
std::vector<std::future<void>> futures;
futures.reserve(local);
for (size_t t = 0; t < local; ++t)
futures.emplace_back(std::async(std::launch::async, [&] {
for (int i = next.fetch_add(1); i < n; i = next.fetch_add(1))
fn(i);
}));
for (auto &f : futures) f.get();
}
// Chunked parallel: each worker gets one contiguous [lo, hi) range, no per-item synchronisation.
// Right for millions of cheap uniform items (the CPU stand-in for a flat CUDA grid-stride kernel).
template <typename Fn>
void ParallelChunks(int n, size_t nthreads, Fn fn) {
if (n <= 0) return;
const int nt = static_cast<int>(std::max<size_t>(1, std::min(nthreads, static_cast<size_t>(n))));
if (nt == 1) { fn(0, n); return; }
const int chunk = (n + nt - 1) / nt;
std::vector<std::future<void>> futures;
futures.reserve(nt);
for (int t = 0; t < nt; ++t) {
const int lo = t * chunk, hi = std::min(n, lo + chunk);
if (lo >= hi) break;
futures.emplace_back(std::async(std::launch::async, [&fn, lo, hi] { fn(lo, hi); }));
}
for (auto &f : futures) f.get();
}
double median_of(std::vector<double> &v) {
std::nth_element(v.begin(), v.begin() + v.size() / 2, v.end());
return v[v.size() / 2];
@@ -319,10 +278,17 @@ void RotationScaleMerge::Ingest() {
// Sort ONCE by (raw h,k,l, image_number) and split into raw-hkl runs. This is the one expensive sort;
// both the 3D combine (event split) and the per-space-group ASU grouping reuse this order.
// Sorting an index array whose comparator dereferences the 72-byte Obs is a cache miss per
// comparison over a multi-GB array, so sort a packed copy of the key fields instead. Started from
// the identity order with the same comparisons in the same order, std::sort takes exactly the same
// branches and produces exactly the same permutation.
// Sorting an index array whose comparator dereferences the 80-byte Obs is a cache miss per
// comparison over a multi-GB array, so sort a packed copy of the key fields instead.
//
// The observation's own index is the last key, which makes the order TOTAL. Two observations can
// genuinely share (h,k,l) and image_number: the predictor emits both intersections of a
// reflection's rotation circle with the Ewald sphere, and near the blind region the two are close
// enough in angle that both land on the same frame. Without a final key their relative order was
// whatever the sort happened to produce, and the combine reads it - it takes on_ice from the first
// member of an event, and its per-event sums are floating point. Ordering them by arrival makes
// the result independent of the sort algorithm, which is what lets the order be reproduced by a
// faster one.
struct SortKey { int32_t h, k, l; float image_number; int32_t idx; };
perm.resize(partials.size());
rawrun_start.clear(); rawrun_count.clear();
@@ -337,7 +303,8 @@ void RotationScaleMerge::Ingest() {
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;
return a.image_number < b.image_number;
if (a.image_number != b.image_number) return a.image_number < b.image_number;
return a.idx < b.idx;
});
for (size_t i = 0; i < keys.size(); ++i) perm[i] = keys[i].idx;
for (int i = 0; i < static_cast<int>(keys.size()); ) {
@@ -377,13 +344,18 @@ void RotationScaleMerge::Ingest() {
std::vector<uint8_t> onice(n);
std::vector<int32_t> frm(n);
std::vector<float> vbkg(n);
for (int i = 0; i < n; ++i) {
const auto &o = partials[i];
I[i] = o.I; sigma[i] = o.sigma; rlp[i] = o.rlp; part[i] = o.partiality;
zeta[i] = o.zeta; onice[i] = o.on_ice; frm[i] = o.frame; corr[i] = o.corr;
bkg[i] = o.bkg; vbkg[i] = o.var_bkg; img[i] = o.image_number; dd[i] = o.d;
px[i] = o.px; py[i] = o.py;
}
// Each observation writes its own slot in each of the fourteen arrays, so this splits
// straight over the range. It is the widest pass in Ingest - it reads the whole 80-byte
// struct and writes 57 bytes of it back out - and it ran on one thread.
ParallelChunks(n, nthreads, [&](int lo, int hi) {
for (int i = lo; i < hi; ++i) {
const auto &o = partials[i];
I[i] = o.I; sigma[i] = o.sigma; rlp[i] = o.rlp; part[i] = o.partiality;
zeta[i] = o.zeta; onice[i] = o.on_ice; frm[i] = o.frame; corr[i] = o.corr;
bkg[i] = o.bkg; vbkg[i] = o.var_bkg; img[i] = o.image_number; dd[i] = o.d;
px[i] = o.px; py[i] = o.py;
}
});
gpu_->SetPartials(n, n_frames, I.data(), sigma.data(), rlp.data(), part.data(), zeta.data(),
onice.data(), frm.data(), corr.data(), frame_start.data(), frame_count.data());
gpu_->SetCombineInputs(bkg.data(), vbkg.data(), img.data(), dd.data(), px.data(), py.data());
image_analysis/scale_merge/RotationScaleMergeGPU.cu
+24 -8
View File
@@ -70,9 +70,16 @@ namespace {
const float *__restrict__ partiality, const float *__restrict__ rlp,
const float *__restrict__ zeta, const uint8_t *__restrict__ on_ice,
const double *__restrict__ group_mean,
const float *__restrict__ sigma, double *__restrict__ inv_sigma,
float *__restrict__ sco_coeff, uint8_t *__restrict__ sco_ok) {
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n_obs) return;
// sigma never changes once it is uploaded, so its reciprocal is the same in every one of the
// thirty IRLS iterations and in every scaling pass. It was being recomputed - a full 64-bit
// reciprocal, five refinement steps - for every observation of every iteration. The CPU has
// always hoisted it (ScaleObs::weight); this is the GPU catching up. Same expression on the
// same operand, so the value is the one the loop used to compute, bit for bit.
inv_sigma[i] = SafeInvD(sigma[i], 1.0);
const int g = group[i];
bool ok = (g >= 0) && !on_ice[i] && isfinite(zeta[i]) && zeta[i] > 0.0f;
double mean = 0.0;
@@ -92,7 +99,7 @@ namespace {
__global__ void FitPerFrameGKernel(int n_frames, double robust_k,
const int32_t *__restrict__ frame_start,
const int32_t *__restrict__ frame_count,
const float *__restrict__ I, const float *__restrict__ sigma,
const float *__restrict__ I, const double *__restrict__ inv_sigma,
const float *__restrict__ sco_coeff, const uint8_t *__restrict__ sco_ok,
const int32_t *__restrict__ perm,
double *__restrict__ g, uint8_t *__restrict__ scaled) {
@@ -119,7 +126,7 @@ namespace {
const int a = perm ? perm[i] : i;
if (!sco_ok[a]) continue;
const double coeff = sco_coeff[a];
const double w = SafeInvD(sigma[a], 1.0);
const double w = inv_sigma[a];
const double w2 = w * w;
num += w2 * coeff * double(I[a]);
den += w2 * coeff * coeff;
@@ -141,7 +148,7 @@ namespace {
const int a = perm ? perm[i] : i;
if (!sco_ok[a]) continue;
const double coeff = sco_coeff[a];
const double w = SafeInvD(sigma[a], 1.0);
const double w = inv_sigma[a];
const double w2 = w * w;
const double res = w * (G * coeff - double(I[a]));
const double rw = 1.0 / (1.0 + res * res / k2);
@@ -610,6 +617,7 @@ struct RotationScaleMergeGPU::Impl {
// scratch
CudaDevicePtr<double> group_mean, g;
CudaDevicePtr<uint8_t> scaled;
CudaDevicePtr<double> inv_sigma; // 1/sigma, hoisted out of the IRLS loop (sigma never changes)
CudaDevicePtr<float> sco_coeff;
CudaDevicePtr<uint8_t> sco_ok;
CudaDevicePtr<double> cc; // per-frame CC (diagnostic), length n_frames
@@ -645,6 +653,7 @@ struct RotationScaleMergeGPU::Impl {
// yield coeff=mean, plus the working corr, the per-obs scale scratch, and the fulls frame/group CSRs
// (built on the host from the small f_frame/f_group key arrays, over the emit-ordered fulls).
CudaDevicePtr<float> f_corr, f_partiality, f_rlp, f_zeta, f_sco_coeff;
CudaDevicePtr<double> f_inv_sigma;
CudaDevicePtr<uint8_t> f_sco_ok;
CudaDevicePtr<int32_t> f_frame_perm, f_frame_start, f_frame_count;
CudaDevicePtr<int32_t> f_gperm, f_gstart, f_gcount;
@@ -706,6 +715,7 @@ void RotationScaleMergeGPU::SetPartials(int n_obs, int n_frames,
Upload(d.frame_start, frame_start, n_frames); Upload(d.frame_count, frame_count, n_frames);
d.g = CudaDevicePtr<double>(n_frames);
d.scaled = CudaDevicePtr<uint8_t>(n_frames);
d.inv_sigma = CudaDevicePtr<double>(n_obs);
d.sco_coeff = CudaDevicePtr<float>(n_obs);
d.sco_ok = CudaDevicePtr<uint8_t>(n_obs);
d.cc = CudaDevicePtr<double>(std::max(1, n_frames));
@@ -745,9 +755,10 @@ void RotationScaleMergeGPU::ScalePartials(int iters, double robust_k, double min
d.group_perm.get(), d.group_start.get(), d.group_count.get(),
d.I.get(), d.sigma.get(), d.partiality.get(), d.corr.get(), d.group_mean.get());
PrepScaleObsKernel<<<obs_blocks, BLK>>>(d.n_obs, d.group.get(), d.partiality.get(), d.rlp.get(),
d.zeta.get(), d.on_ice.get(), d.group_mean.get(), d.sco_coeff.get(), d.sco_ok.get());
d.zeta.get(), d.on_ice.get(), d.group_mean.get(), d.sigma.get(), d.inv_sigma.get(),
d.sco_coeff.get(), d.sco_ok.get());
FitPerFrameGKernel<<<d.n_frames, BLK>>>(d.n_frames, robust_k, d.frame_start.get(), d.frame_count.get(),
d.I.get(), d.sigma.get(), d.sco_coeff.get(), d.sco_ok.get(), nullptr, d.g.get(), d.scaled.get());
d.I.get(), d.inv_sigma.get(), d.sco_coeff.get(), d.sco_ok.get(), nullptr, d.g.get(), d.scaled.get());
UpdateCorrKernel<<<upd_blocks, BLK>>>(d.n_obs, d.frame.get(), d.rlp.get(), d.partiality.get(),
d.g.get(), d.scaled.get(), d.corr.get());
}
@@ -1006,6 +1017,7 @@ int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_parti
d.f_on_ice = CudaDevicePtr<uint8_t>(nf);
d.f_corr = CudaDevicePtr<float>(nf); d.f_partiality = CudaDevicePtr<float>(nf);
d.f_rlp = CudaDevicePtr<float>(nf); d.f_zeta = CudaDevicePtr<float>(nf);
d.f_inv_sigma = CudaDevicePtr<double>(nf);
d.f_sco_coeff = CudaDevicePtr<float>(nf); d.f_sco_ok = CudaDevicePtr<uint8_t>(nf);
CudaCheck(cudaMemcpy(d.rr_offset.get(), offset.data(), size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyHostToDevice), "upload offset");
@@ -1091,10 +1103,14 @@ void RotationScaleMergeGPU::ScaleFulls(int iters, double robust_k, double min_pa
ReduceGroupMeansKernel<<<grp_blocks, BLK>>>(d.n_groups, min_partiality,
d.f_gperm.get(), d.f_gstart.get(), d.f_gcount.get(),
d.f_I.get(), d.f_sigma.get(), d.f_partiality.get(), d.f_corr.get(), d.group_mean.get());
PrepScaleObsKernel<<<obs_blocks, BLK>>>(nf, d.f_group.get(), d.f_partiality.get(), d.f_rlp.get(),
d.f_zeta.get(), d.f_on_ice.get(), d.group_mean.get(), d.f_sco_coeff.get(), d.f_sco_ok.get());
// Not grid-stride, so its grid has to cover every full - unlike the grid-stride kernels
// below, which the 65535 cap is there for. Capped, it would silently leave the tail of
// sco_coeff/sco_ok stale above 16.8M fulls.
PrepScaleObsKernel<<<(nf + BLK - 1) / BLK, BLK>>>(nf, d.f_group.get(), d.f_partiality.get(),
d.f_rlp.get(), d.f_zeta.get(), d.f_on_ice.get(), d.group_mean.get(),
d.f_sigma.get(), d.f_inv_sigma.get(), d.f_sco_coeff.get(), d.f_sco_ok.get());
FitPerFrameGKernel<<<d.n_frames, BLK>>>(d.n_frames, robust_k,
d.f_frame_start.get(), d.f_frame_count.get(), d.f_I.get(), d.f_sigma.get(),
d.f_frame_start.get(), d.f_frame_count.get(), d.f_I.get(), d.f_inv_sigma.get(),
d.f_sco_coeff.get(), d.f_sco_ok.get(), d.f_frame_perm.get(), d.g.get(), d.scaled.get());
UpdateCorrKernel<<<obs_blocks, BLK>>>(nf, d.f_frame.get(), d.f_rlp.get(), d.f_partiality.get(),
d.g.get(), d.scaled.get(), d.f_corr.get());