diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 7efc3d271..7d64bf9e9 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -672,11 +672,11 @@ void RotationScaleMerge::SortFullsByFrame() { } } -void RotationScaleMerge::FinalizePerFrameScale(int n_groups, const std::vector &partial_group_mean, - const std::vector &frame_scaled) { +void RotationScaleMerge::ComputePerFrameCC(const std::vector &partial_group_mean, + std::vector &cc, std::vector &cc_n) const { // Per-frame CC vs the merged reference (CalculateGlobalCC), computed once now (not every iteration). - std::vector cc(n_frames, NAN); - std::vector 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 0.0 && vy > 0.0) { cc[f] = cov / std::sqrt(vx * vy); cc_n[f] = static_cast(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 &cc, const std::vector &cc_n, + const std::vector &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 cc; + std::vector cc_n; + bool cc_on_gpu = false; +#ifdef JFJOCH_USE_CUDA + if (gpu_active_) { + std::vector 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 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); diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 6c338d716..d7abbf213 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -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 &partial_group_mean, + // Per-frame CC vs the partial merge reference (CPU; the GPU equivalent is gpu_->ComputePartialCC). + void ComputePerFrameCC(const std::vector &partial_group_mean, + std::vector &cc, std::vector &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 &cc, const std::vector &cc_n, const std::vector &frame_scaled); // Error model + merge + statistics over the fulls (the last stage). n_groups is the fulls group count. diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 1ca7182a9..b1828e9ae 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -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 scaled; CudaDevicePtr sco_coeff; CudaDevicePtr sco_ok; + CudaDevicePtr cc; // per-frame CC (diagnostic), length n_frames + CudaDevicePtr cc_n; // combine: extra per-obs inputs + the one-time raw-hkl run layout CudaDevicePtr bkg, image_number, d_obs; @@ -425,6 +474,8 @@ void RotationScaleMergeGPU::SetPartials(int n_obs, int n_frames, d.scaled = CudaDevicePtr(n_frames); d.sco_coeff = CudaDevicePtr(n_obs); d.sco_ok = CudaDevicePtr(n_obs); + d.cc = CudaDevicePtr(std::max(1, n_frames)); + d.cc_n = CudaDevicePtr(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<<>>(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, 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); diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.h b/image_analysis/scale_merge/RotationScaleMergeGPU.h index aba5c95c2..d9d146f05 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.h +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.h @@ -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