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

RotationScaleMerge: GPU post-smooth group means + per-frame CC

Browse Source 
The single-threaded ReduceGroupMeans over the 6.3M partials (~0.07s/2-pass) and the
per-frame diagnostic CC now run on the resident partials on the GPU: after SmoothG,
the smoothed corr is uploaded once (and left resident for the combine, dropping the
combine's redundant re-upload), then the post-smooth group means (reusing the scaling
reduce) and the per-frame Pearson CC (a new one-block-per-frame kernel) run there and
only the tiny per-frame cc/cc_n come back. FinalizePerFrameScale is split into
ComputePerFrameCC (host reference) + the writeback; the GPU path uses ComputePartialCC.

The per-frame CC is diagnostic only (the per-image scaling table), so the tree
reduction's ~ulp difference from the CPU is immaterial and it does not touch merged
intensities. smooth+CC region ~0.10s GPU vs ~0.15s CPU on lyso. Validated across the
battery: 15/15 deterministic crystals run-to-run deterministic and merged output
bit-identical to the CPU path (only EP_cs_01-24, unindexable noise, keeps its benign
error-model-b wobble). CPU fallbacks (JFJOCH_RSM_CPU_COMBINE / _NO_GPU) unchanged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
leonarski_f
2026-07-03 08:58:26 +02:00
co-authored by Claude Opus 4.8
parent 617184041f
commit 7d804eb799
4 changed files with 114 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/scale_merge/RotationScaleMerge.cpp
+32 -10
View File
@@ -672,11 +672,11 @@ void RotationScaleMerge::SortFullsByFrame() {
}
}
void RotationScaleMerge::FinalizePerFrameScale(int n_groups, const std::vector<double> &partial_group_mean,
const std::vector<uint8_t> &frame_scaled) {
void RotationScaleMerge::ComputePerFrameCC(const std::vector<double> &partial_group_mean,
std::vector<double> &cc, std::vector<int64_t> &cc_n) const {
// Per-frame CC vs the merged reference (CalculateGlobalCC), computed once now (not every iteration).
std::vector<double> cc(n_frames, NAN);
std::vector<int64_t> cc_n(n_frames, 0);
cc.assign(n_frames, NAN);
cc_n.assign(n_frames, 0);
ParallelFor(n_frames, nthreads, [&](int f) {
double sx = 0, sy = 0, sx2 = 0, sy2 = 0, sxy = 0;
size_t n = 0;
@@ -701,7 +701,12 @@ void RotationScaleMerge::FinalizePerFrameScale(int n_groups, const std::vector<d
const double vy = sy2 - sy * sy / nd;
if (vx > 0.0 && vy > 0.0) { cc[f] = cov / std::sqrt(vx * vy); cc_n[f] = static_cast<int64_t>(n); }
});
}
// Write the per-frame G / CC / mosaicity (from the given cc/cc_n) back onto the partials for the offline
// per-image scaling table. cc/cc_n are computed on the host (ComputePerFrameCC) or GPU (ComputePartialCC).
void RotationScaleMerge::FinalizePerFrameScale(const std::vector<double> &cc, const std::vector<int64_t> &cc_n,
const std::vector<uint8_t> &frame_scaled) {
for (int f = 0; f < n_frames; ++f) {
auto &o = partials_out[f];
if (frame_scaled[f]) {
@@ -1092,9 +1097,28 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search,
SmoothG(partials, g_partial, window);
}
// Per-frame CC + write G/CC/mosaicity back onto the partials (once).
ReduceGroupMeans(partials, n_groups, false, {}, partial_mean);
FinalizePerFrameScale(n_groups, partial_mean, partial_scaled);
// Per-frame CC + write G/CC/mosaicity back onto the partials (once). On the GPU the smoothed corr is
// uploaded here and stays resident for the combine; the post-smooth group means + per-frame CC run on
// the resident partials and only the tiny per-frame cc/cc_n come back.
std::vector<double> cc;
std::vector<int64_t> cc_n;
bool cc_on_gpu = false;
#ifdef JFJOCH_USE_CUDA
if (gpu_active_) {
std::vector<float> corr(partials.size());
for (size_t i = 0; i < partials.size(); ++i) corr[i] = partials[i].corr;
gpu_->SetCorr(corr.data()); // smoothed corr; also consumed by the GPU combine below
cc.resize(n_frames);
cc_n.resize(n_frames);
gpu_->ComputePartialCC(min_partiality, cc.data(), cc_n.data());
cc_on_gpu = true;
}
#endif
if (!cc_on_gpu) {
ReduceGroupMeans(partials, n_groups, false, {}, partial_mean);
ComputePerFrameCC(partial_mean, cc, cc_n);
}
FinalizePerFrameScale(cc, cc_n, partial_scaled);
// --- 3. 3D combine of per-frame partials into fulls (fulls inherit their ASU group here). ---
bool combined_on_gpu = false;
@@ -1105,9 +1129,7 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search,
// stable-sort), scale the fulls in place, and download only once. Mirrors Combine() + the Unity
// scale-fulls loop below. The diagnostic dump (serial, one writer) has no GPU path -> CPU fallback.
if (gpu_active_ && gpu_combine_ && observation_dump_path.empty()) {
std::vector<float> corr(partials.size()); // refresh the smoothed corr on the device
for (size_t i = 0; i < partials.size(); ++i) corr[i] = partials[i].corr;
gpu_->SetCorr(corr.data());
// The smoothed corr is already resident (uploaded for the per-frame CC just above).
const int nf = gpu_->Combine(rawrun_group.data(), min_partiality, capture_uncertainty_coeff);
g_full.assign(n_frames, 1.0);
image_analysis/scale_merge/RotationScaleMerge.h
+7 -3
View File
@@ -176,9 +176,13 @@ private:
// step slices). Shared by the CPU Combine tail and the GPU combine path.
void SortFullsByFrame();
// Per-frame CC vs the partial merge reference, then write G/CC/mosaicity back onto the partials
// (once, at the end of partial scaling) so the offline per-image scaling table is still exported.
void FinalizePerFrameScale(int n_groups, const std::vector<double> &partial_group_mean,
// Per-frame CC vs the partial merge reference (CPU; the GPU equivalent is gpu_->ComputePartialCC).
void ComputePerFrameCC(const std::vector<double> &partial_group_mean,
std::vector<double> &cc, std::vector<int64_t> &cc_n) const;
// Write G/CC/mosaicity back onto the partials (once, at the end of partial scaling) from the given
// per-frame cc/cc_n, so the offline per-image scaling table is still exported.
void FinalizePerFrameScale(const std::vector<double> &cc, const std::vector<int64_t> &cc_n,
const std::vector<uint8_t> &frame_scaled);
// Error model + merge + statistics over the fulls (the last stage). n_groups is the fulls group count.
image_analysis/scale_merge/RotationScaleMergeGPU.cu
+70
View File
@@ -166,6 +166,53 @@ namespace {
if (threadIdx.x == 0) { g[f] = s_G; scaled[f] = 1; }
}
// One block per frame: Pearson CC of (I*corr) vs the merged group mean over the frame's partials,
// == FinalizePerFrameScale's per-frame loop. Diagnostic only (per-image scaling table), so the tree
// reduction's ~ulp difference from the CPU is immaterial; deterministic run-to-run.
__global__ void PerFrameCCKernel(int n_frames, double min_partiality,
const int32_t *__restrict__ frame_start,
const int32_t *__restrict__ frame_count,
const float *__restrict__ I, const float *__restrict__ sigma,
const float *__restrict__ partiality, const float *__restrict__ corr,
const uint8_t *__restrict__ on_ice, const int32_t *__restrict__ group,
const double *__restrict__ group_mean,
double *__restrict__ cc_out, int64_t *__restrict__ cc_n_out) {
const int f = blockIdx.x;
if (f >= n_frames) return;
const int lo = frame_start[f], hi = frame_start[f] + frame_count[f];
__shared__ double sh[BLK];
double sx = 0, sy = 0, sx2 = 0, sy2 = 0, sxy = 0;
long nl = 0;
for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) {
if (on_ice[i]) continue;
const int g = group[i];
if (g < 0) continue;
if (partiality[i] < min_partiality) continue;
const float c = corr[i];
if (!isfinite(I[i]) || !isfinite(c) || !(c > 0.0f)) continue;
if (!isfinite(sigma[i]) || !(sigma[i] > 0.0f)) continue;
const double mean = group_mean[g];
if (!isfinite(mean)) continue;
const double img = double(I[i]) * c;
sx += img; sy += mean; sx2 += img * img; sy2 += mean * mean; sxy += img * mean; ++nl;
}
const double tsx = BlockReduceSum(sx, sh); __syncthreads();
const double tsy = BlockReduceSum(sy, sh); __syncthreads();
const double tsx2 = BlockReduceSum(sx2, sh); __syncthreads();
const double tsy2 = BlockReduceSum(sy2, sh); __syncthreads();
const double tsxy = BlockReduceSum(sxy, sh); __syncthreads();
const double tn = BlockReduceSum(double(nl), sh);
if (threadIdx.x == 0) {
cc_out[f] = NAN; cc_n_out[f] = 0;
if (tn >= MIN_REFLECTIONS) {
const double cov = tsxy - tsx * tsy / tn;
const double vx = tsx2 - tsx * tsx / tn;
const double vy = tsy2 - tsy * tsy / tn;
if (vx > 0.0 && vy > 0.0) { cc_out[f] = cov / sqrt(vx * vy); cc_n_out[f] = int64_t(tn); }
}
}
}
// corr = rlp / (partiality * G[frame]) for fitted frames; unchanged otherwise (grid-stride).
__global__ void UpdateCorrKernel(int n_obs, const int32_t *__restrict__ frame,
const float *__restrict__ rlp, const float *__restrict__ partiality,
@@ -378,6 +425,8 @@ struct RotationScaleMergeGPU::Impl {
CudaDevicePtr<uint8_t> scaled;
CudaDevicePtr<float> sco_coeff;
CudaDevicePtr<uint8_t> sco_ok;
CudaDevicePtr<double> cc; // per-frame CC (diagnostic), length n_frames
CudaDevicePtr<int64_t> cc_n;
// combine: extra per-obs inputs + the one-time raw-hkl run layout
CudaDevicePtr<float> bkg, image_number, d_obs;
@@ -425,6 +474,8 @@ void RotationScaleMergeGPU::SetPartials(int n_obs, int n_frames,
d.scaled = CudaDevicePtr<uint8_t>(n_frames);
d.sco_coeff = CudaDevicePtr<float>(n_obs);
d.sco_ok = CudaDevicePtr<uint8_t>(n_obs);
d.cc = CudaDevicePtr<double>(std::max(1, n_frames));
d.cc_n = CudaDevicePtr<int64_t>(std::max(1, n_frames));
}
void RotationScaleMergeGPU::SetGroups(int n_groups, const int32_t *group, const int32_t *group_perm,
@@ -479,6 +530,25 @@ void RotationScaleMergeGPU::GetG(double *g_out, uint8_t *scaled_out) const {
cudaMemcpyDeviceToHost), "download scaled");
}
void RotationScaleMergeGPU::ComputePartialCC(double min_partiality, double *cc_out, int64_t *cc_n_out) {
auto &d = *impl_;
const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK);
// Post-smooth group means (reuse the scaling reduce; reads the resident, smoothed corr), then the
// per-frame CC over the resident partials. Only the tiny per-frame cc/cc_n come back to the host.
ReduceGroupMeansKernel<<<grp_blocks, BLK>>>(d.n_groups, min_partiality,
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());
PerFrameCCKernel<<<d.n_frames, BLK>>>(d.n_frames, min_partiality,
d.frame_start.get(), d.frame_count.get(), d.I.get(), d.sigma.get(), d.partiality.get(),
d.corr.get(), d.on_ice.get(), d.group.get(), d.group_mean.get(), d.cc.get(), d.cc_n.get());
CudaCheck(cudaGetLastError(), "partial CC launch");
CudaCheck(cudaDeviceSynchronize(), "partial CC sync");
CudaCheck(cudaMemcpy(cc_out, d.cc.get(), size_t(d.n_frames) * sizeof(double),
cudaMemcpyDeviceToHost), "download cc");
CudaCheck(cudaMemcpy(cc_n_out, d.cc_n.get(), size_t(d.n_frames) * sizeof(int64_t),
cudaMemcpyDeviceToHost), "download cc_n");
}
void RotationScaleMergeGPU::SetCombineInputs(const float *bkg, const float *image_number, const float *d) {
auto &dd = *impl_;
Upload(dd.bkg, bkg, dd.n_obs);
image_analysis/scale_merge/RotationScaleMergeGPU.h
+5
View File
@@ -54,6 +54,11 @@ public:
void GetCorr(float *corr_out) const;
void GetG(double *g_out, uint8_t *scaled_out) const;
// Post-smooth per-frame diagnostic CC: recompute the group means from the resident (smoothed) corr
// and the Pearson CC of each frame's I*corr vs its group mean, downloading only the per-frame cc /
// cc_n (length n_frames). Mirrors ReduceGroupMeans(partials) + FinalizePerFrameScale's CC loop.
void ComputePartialCC(double min_partiality, double *cc_out, int64_t *cc_n_out);
// --- 3D combine (partials -> fulls), all on the device ---
// The per-obs fields the combine needs on top of the scaling inputs (image-local bkg, fractional