diff --git a/image_analysis/scale_merge/CMakeLists.txt b/image_analysis/scale_merge/CMakeLists.txt index 80594da7..cd020430 100644 --- a/image_analysis/scale_merge/CMakeLists.txt +++ b/image_analysis/scale_merge/CMakeLists.txt @@ -15,4 +15,9 @@ ADD_LIBRARY(JFJochScaleMerge HKLKey.h ScalingResult.h ScalingResult.cpp) -TARGET_LINK_LIBRARIES(JFJochScaleMerge Ceres::ceres Eigen3::Eigen JFJochCommon) \ No newline at end of file +TARGET_LINK_LIBRARIES(JFJochScaleMerge Ceres::ceres Eigen3::Eigen JFJochCommon) + +IF (JFJOCH_CUDA_AVAILABLE) + TARGET_SOURCES(JFJochScaleMerge PRIVATE ../indexing/CUDAMemHelpers.h + RotationScaleMergeGPU.cu RotationScaleMergeGPU.h) +ENDIF() \ No newline at end of file diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 1d881e70..5eec365e 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -238,6 +238,27 @@ void RotationScaleMerge::Ingest() { total, n_frames, rawrun_start.size()); SmoothMosaicityAndPartiality(); + +#ifdef JFJOCH_USE_CUDA + // Bring the partial-scaling loop onto the GPU when one is present. Upload the immutable per-obs + // fields once (corr lives on the device, refreshed each pass); the CPU keeps the sort/keying/combine. + gpu_ = std::make_unique(); + gpu_active_ = gpu_->Available(); + if (gpu_active_) { + const int n = static_cast(partials.size()); + std::vector I(n), sigma(n), rlp(n), part(n), zeta(n), corr(n); + std::vector onice(n); + std::vector frm(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; + } + 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()); + logger.Info("RotationScaleMerge: GPU partial-scaling active"); + } +#endif } void RotationScaleMerge::SmoothMosaicityAndPartiality() { @@ -340,6 +361,28 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) { } } }); + +#ifdef JFJOCH_USE_CUDA + // Group-ordered permutation (obs bucketed by ASU group, obs-index order) + its CSR, so the GPU + // reduction is a deterministic segmented reduction (fixed order, no atomics). Built per space group. + if (gpu_active_) { + const int n = static_cast(partials.size()); + std::vector group_ids(n), gcount(n_groups, 0); + for (int i = 0; i < n; ++i) { + group_ids[i] = partials[i].group; + if (partials[i].group >= 0) ++gcount[partials[i].group]; + } + std::vector gstart(n_groups, 0); + int acc = 0; + for (int g = 0; g < n_groups; ++g) { gstart[g] = acc; acc += gcount[g]; } + std::vector gperm(acc), gfill = gstart; + for (int i = 0; i < n; ++i) { + const int g = partials[i].group; + if (g >= 0) gperm[gfill[g]++] = i; + } + gpu_->SetGroups(n_groups, group_ids.data(), gperm.data(), acc, gstart.data(), gcount.data()); + } +#endif return n_groups; } @@ -1000,10 +1043,28 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search, const int n_groups = ComputeAsuGroups(keygen); // one ASU grouping, shared by partials and fulls lap("group hkl"); std::vector partial_mean; - for (int it = 0; it < scaling_iter; ++it) { - ReduceGroupMeans(partials, n_groups, false, {}, partial_mean); - FitPerFrameG(partials, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial); - UpdateCorr(partials, g_partial, frame_scaled_scratch); + bool scaled_on_gpu = false; +#ifdef JFJOCH_USE_CUDA + if (gpu_active_) { + // Refresh corr on the device (smooth-G mutated it on the host between passes), run the whole + // scaling loop on the GPU, then read corr + per-frame G/scaled back. + std::vector corr(partials.size()); + for (size_t i = 0; i < partials.size(); ++i) corr[i] = partials[i].corr; + gpu_->SetCorr(corr.data()); + gpu_->ScalePartials(scaling_iter, SCALE_ROBUST_K, min_partiality, d_min_limit.has_value()); + gpu_->GetCorr(corr.data()); + frame_scaled_scratch.assign(n_frames, 0); + gpu_->GetG(g_partial.data(), frame_scaled_scratch.data()); + for (size_t i = 0; i < partials.size(); ++i) partials[i].corr = corr[i]; + scaled_on_gpu = true; + } +#endif + if (!scaled_on_gpu) { + for (int it = 0; it < scaling_iter; ++it) { + ReduceGroupMeans(partials, n_groups, false, {}, partial_mean); + FitPerFrameG(partials, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial); + UpdateCorr(partials, g_partial, frame_scaled_scratch); + } } lap("scale partials"); const std::vector partial_scaled = frame_scaled_scratch; diff --git a/image_analysis/scale_merge/RotationScaleMerge.h b/image_analysis/scale_merge/RotationScaleMerge.h index 8fafb1d0..d858a51a 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.h +++ b/image_analysis/scale_merge/RotationScaleMerge.h @@ -14,6 +14,10 @@ #include "../IntegrationOutcome.h" #include "Merge.h" // MergedReflection, MergeStatistics +#ifdef JFJOCH_USE_CUDA +#include +#include "RotationScaleMergeGPU.h" +#endif // Dedicated, allocate-once scale+combine+merge for rotation data (the -P rot3d path). // @@ -124,6 +128,13 @@ private: // Working per-group arrays (sized to the current group count; reused). std::vector group_h, group_k, group_l; +#ifdef JFJOCH_USE_CUDA + // GPU engine for the partial-scaling loop (segmented reduce + per-frame IRLS + corr update). Null / + // inactive when no GPU; the CPU loops are the fallback. Built in Ingest. + std::unique_ptr gpu_; + bool gpu_active_ = false; +#endif + // --- helpers (each a flat pass; see the .cpp) --- // Compute the dense ASU-group id for the current space group by grouping the (pre-sorted) raw-hkl // runs by their ASU key - one gemmi ASU reduction per distinct raw hkl, not per observation. Fills diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu new file mode 100644 index 00000000..b0e0eb9b --- /dev/null +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -0,0 +1,289 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "RotationScaleMergeGPU.h" + +#include +#include +#include +#include + +#include "../indexing/CUDAMemHelpers.h" +#include "../../common/CUDAWrapper.h" +#include "../../common/JFJochException.h" + +namespace { + constexpr int BLK = 256; + constexpr int MIN_REFLECTIONS = 20; + + __device__ __forceinline__ double SafeInvD(double x, double fallback) { + return (isfinite(x) && x != 0.0) ? 1.0 / x : fallback; + } + + // Block reduction of a double, deterministic for a fixed thread->element mapping (fixed order, + // no atomics). Returns the sum on thread 0; `s` is BLK doubles of shared scratch. + __device__ double BlockReduceSum(double v, double *s) { + const int t = threadIdx.x; + s[t] = v; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (t < stride) s[t] += s[t + stride]; + __syncthreads(); + } + return s[0]; + } + + // One thread per ASU group (grid-stride): inverse-variance mean of I*corr over the group's contiguous, + // fixed-order segment of group_perm. Groups are small (avg tens of obs), so a whole block per group + // wastes threads; summing each group in one thread avoids launch/sync overhead and stays deterministic + // (fixed group_perm order). Matches the CPU ReduceGroupMeans filter (no ice/cell mask - scaling ref). + __global__ void ReduceGroupMeansKernel(int n_groups, double min_partiality, + const int32_t *__restrict__ group_perm, + const int32_t *__restrict__ group_start, + const int32_t *__restrict__ group_count, + const float *__restrict__ I, const float *__restrict__ sigma, + const float *__restrict__ partiality, + const float *__restrict__ corr, + double *__restrict__ group_mean) { + for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < n_groups; g += gridDim.x * blockDim.x) { + const int lo = group_start[g], hi = group_start[g] + group_count[g]; + double sw = 0.0, swI = 0.0; + for (int p = lo; p < hi; ++p) { + const int i = group_perm[p]; + const float c = corr[i]; + if (!(c > 0.0f) || !isfinite(c)) continue; + if (partiality[i] < min_partiality) continue; + const float I_corr = I[i] * c; + const float sigma_corr = sigma[i] * c; + if (!isfinite(I_corr) || !isfinite(sigma_corr) || sigma_corr <= 0.0f) continue; + const double w = 1.0 / (double(sigma_corr) * sigma_corr); + sw += w; + swI += w * I_corr; + } + group_mean[g] = sw > 0.0 ? swI / sw : NAN; + } + } + + // Per-observation scale-fit coefficient (rotation model) and accept flag, recomputed each scaling + // iteration once the group means are known. coeff = partiality * (1/rlp) * mean[group]. + __global__ void PrepScaleObsKernel(int n_obs, const int32_t *__restrict__ group, + const float *__restrict__ partiality, const float *__restrict__ rlp, + const float *__restrict__ zeta, const uint8_t *__restrict__ on_ice, + const double *__restrict__ group_mean, + float *__restrict__ sco_coeff, uint8_t *__restrict__ sco_ok) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n_obs) return; + const int g = group[i]; + bool ok = (g >= 0) && !on_ice[i] && isfinite(zeta[i]) && zeta[i] > 0.0f; + double mean = 0.0; + if (ok) { + mean = group_mean[g]; + ok = isfinite(mean); + } + sco_ok[i] = ok ? 1 : 0; + sco_coeff[i] = ok ? float(double(partiality[i]) * SafeInvD(rlp[i], 1.0) * mean) : 0.0f; + } + + // One block per frame: robust per-frame scale G by IRLS (Cauchy), over the frame's contiguous obs. + // Identical objective to the CPU SolveScaleIRLS. Leaves g/scaled untouched for under-populated frames. + __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__ sco_coeff, const uint8_t *__restrict__ sco_ok, + double *__restrict__ g, uint8_t *__restrict__ scaled) { + 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]; + + // count accepted + long cnt_local = 0; + for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) + if (sco_ok[i]) ++cnt_local; + const double cnt = BlockReduceSum(double(cnt_local), sh); + __shared__ double s_cnt; + if (threadIdx.x == 0) s_cnt = cnt; + __syncthreads(); + if (s_cnt < MIN_REFLECTIONS) return; // leave g[f]/scaled[f] as-is + + const double k2 = robust_k * robust_k; + + // seed: plain weighted-LS ratio (robust weight = 1) + double num = 0.0, den = 0.0; + for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) { + if (!sco_ok[i]) continue; + const double coeff = sco_coeff[i]; + const double w = SafeInvD(sigma[i], 1.0); + const double w2 = w * w; + num += w2 * coeff * double(I[i]); + den += w2 * coeff * coeff; + } + double tnum = BlockReduceSum(num, sh); __syncthreads(); + double tden = BlockReduceSum(den, sh); + __shared__ double s_G; + if (threadIdx.x == 0) { + double G = tden > 0.0 ? tnum / tden : NAN; + s_G = isfinite(G) ? fmax(0.0, G) : 1.0; + } + __syncthreads(); + if (!(s_G >= 0.0)) { if (threadIdx.x == 0) { g[f] = 1.0; scaled[f] = 1; } return; } + + for (int iter = 0; iter < 30; ++iter) { + const double G = s_G; + num = 0.0; den = 0.0; + for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) { + if (!sco_ok[i]) continue; + const double coeff = sco_coeff[i]; + const double w = SafeInvD(sigma[i], 1.0); + const double w2 = w * w; + const double res = w * (G * coeff - double(I[i])); + const double rw = 1.0 / (1.0 + res * res / k2); + num += rw * w2 * coeff * double(I[i]); + den += rw * w2 * coeff * coeff; + } + tnum = BlockReduceSum(num, sh); __syncthreads(); + tden = BlockReduceSum(den, sh); + if (threadIdx.x == 0) { + const double G_next = tden > 0.0 ? tnum / tden : NAN; + if (isfinite(G_next)) { + const double Gn = fmax(0.0, G_next); + s_cnt = (fabs(Gn - G) <= 1e-7 * fmax(Gn, 1.0)) ? 1.0 : 0.0; // reuse s_cnt as "converged" + s_G = Gn; + } else { + s_cnt = 1.0; + } + } + __syncthreads(); + if (s_cnt != 0.0) break; + } + if (threadIdx.x == 0) { g[f] = s_G; scaled[f] = 1; } + } + + // 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, + const double *__restrict__ g, const uint8_t *__restrict__ scaled, + float *__restrict__ corr) { + for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) { + const int f = frame[i]; + if (!scaled[f]) continue; + const double denom = double(partiality[i]) * g[f]; + corr[i] = (isfinite(double(rlp[i])) && isfinite(denom) && denom > 0.0) + ? float(rlp[i] / denom) : NAN; + } + } + + void CudaCheck(cudaError_t e, const char *what) { + if (e != cudaSuccess) + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, + std::string("RotationScaleMergeGPU: ") + what + ": " + cudaGetErrorString(e)); + } + + template + void Upload(CudaDevicePtr &dst, const T *src, int n) { + dst = CudaDevicePtr(std::max(1, n)); + if (n > 0) + CudaCheck(cudaMemcpy(dst.get(), src, size_t(n) * sizeof(T), cudaMemcpyHostToDevice), "upload"); + } +} + +struct RotationScaleMergeGPU::Impl { + bool available = false; + int n_obs = 0, n_frames = 0, n_groups = 0; + + // immutable per-obs + CudaDevicePtr I, sigma, rlp, partiality, zeta; + CudaDevicePtr on_ice; + CudaDevicePtr frame; + CudaDevicePtr corr; // mutable, resident across iterations + CudaDevicePtr frame_start, frame_count; + // per space group + CudaDevicePtr group, group_perm, group_start, group_count; + // scratch + CudaDevicePtr group_mean, g; + CudaDevicePtr scaled; + CudaDevicePtr sco_coeff; + CudaDevicePtr sco_ok; +}; + +RotationScaleMergeGPU::RotationScaleMergeGPU() : impl_(std::make_unique()) { + if (get_gpu_count() > 0) { + set_gpu(0); + impl_->available = true; + } +} + +RotationScaleMergeGPU::~RotationScaleMergeGPU() = default; + +bool RotationScaleMergeGPU::Available() const { return impl_->available; } + +void RotationScaleMergeGPU::SetPartials(int n_obs, int n_frames, + const float *I, const float *sigma, const float *rlp, + const float *partiality, const float *zeta, const uint8_t *on_ice, + const int32_t *frame, const float *corr0, + const int32_t *frame_start, const int32_t *frame_count) { + auto &d = *impl_; + d.n_obs = n_obs; + d.n_frames = n_frames; + Upload(d.I, I, n_obs); Upload(d.sigma, sigma, n_obs); Upload(d.rlp, rlp, n_obs); + Upload(d.partiality, partiality, n_obs); Upload(d.zeta, zeta, n_obs); Upload(d.on_ice, on_ice, n_obs); + Upload(d.frame, frame, n_obs); Upload(d.corr, corr0, n_obs); + Upload(d.frame_start, frame_start, n_frames); Upload(d.frame_count, frame_count, n_frames); + d.g = CudaDevicePtr(n_frames); + d.scaled = CudaDevicePtr(n_frames); + d.sco_coeff = CudaDevicePtr(n_obs); + d.sco_ok = CudaDevicePtr(n_obs); +} + +void RotationScaleMergeGPU::SetGroups(int n_groups, const int32_t *group, const int32_t *group_perm, + int n_group_perm, const int32_t *group_start, + const int32_t *group_count) { + auto &d = *impl_; + d.n_groups = n_groups; + Upload(d.group, group, d.n_obs); + Upload(d.group_perm, group_perm, n_group_perm); // obs with group >= 0, in group order + Upload(d.group_start, group_start, n_groups); + Upload(d.group_count, group_count, n_groups); + d.group_mean = CudaDevicePtr(std::max(1, n_groups)); +} + +void RotationScaleMergeGPU::SetCorr(const float *corr) { + CudaCheck(cudaMemcpy(impl_->corr.get(), corr, size_t(impl_->n_obs) * sizeof(float), + cudaMemcpyHostToDevice), "upload corr"); +} + +void RotationScaleMergeGPU::ScalePartials(int iters, double robust_k, double min_partiality, + bool /*has_d_min*/) { + auto &d = *impl_; + CudaCheck(cudaMemset(d.scaled.get(), 0, size_t(d.n_frames) * sizeof(uint8_t)), "memset scaled"); + CudaCheck(cudaMemset(d.g.get(), 0, size_t(d.n_frames) * sizeof(double)), "memset g"); // unscaled g unused + const int obs_blocks = (d.n_obs + BLK - 1) / BLK; + const int upd_blocks = std::min(65535, obs_blocks); + const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK); + for (int it = 0; it < iters; ++it) { + 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()); + PrepScaleObsKernel<<>>(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()); + FitPerFrameGKernel<<>>(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(), d.g.get(), d.scaled.get()); + UpdateCorrKernel<<>>(d.n_obs, d.frame.get(), d.rlp.get(), d.partiality.get(), + d.g.get(), d.scaled.get(), d.corr.get()); + } + CudaCheck(cudaGetLastError(), "kernel launch"); + CudaCheck(cudaDeviceSynchronize(), "scale sync"); +} + +void RotationScaleMergeGPU::GetCorr(float *corr_out) const { + CudaCheck(cudaMemcpy(corr_out, impl_->corr.get(), size_t(impl_->n_obs) * sizeof(float), + cudaMemcpyDeviceToHost), "download corr"); +} + +void RotationScaleMergeGPU::GetG(double *g_out, uint8_t *scaled_out) const { + CudaCheck(cudaMemcpy(g_out, impl_->g.get(), size_t(impl_->n_frames) * sizeof(double), + cudaMemcpyDeviceToHost), "download g"); + CudaCheck(cudaMemcpy(scaled_out, impl_->scaled.get(), size_t(impl_->n_frames) * sizeof(uint8_t), + cudaMemcpyDeviceToHost), "download scaled"); +} diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.h b/image_analysis/scale_merge/RotationScaleMergeGPU.h new file mode 100644 index 00000000..cb01e49f --- /dev/null +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.h @@ -0,0 +1,60 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include +#include + +// 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 +// plus a group-ordered permutation so the per-group reduction is a deterministic segmented reduction +// (one block per group, fixed order, no atomics) - matching the run-to-run determinism of the CPU path. +// +// Only compiled when CUDA is available; the header is safe to include unconditionally (the impl behind +// the pimpl is null without CUDA, and RotationScaleMerge falls back to the CPU loops). +class RotationScaleMergeGPU { +public: + RotationScaleMergeGPU(); + ~RotationScaleMergeGPU(); + RotationScaleMergeGPU(const RotationScaleMergeGPU &) = delete; + RotationScaleMergeGPU &operator=(const RotationScaleMergeGPU &) = delete; + + // True if a GPU was found and the engine is usable. + [[nodiscard]] bool Available() const; + + // Upload the immutable per-observation fields (once). Arrays are length n_obs unless noted; the + // frame CSR (frame_start/frame_count) is length n_frames and indexes the obs arrays in frame order. + void SetPartials(int n_obs, int n_frames, + const float *I, const float *sigma, const float *rlp, const float *partiality, + const float *zeta, const uint8_t *on_ice, const int32_t *frame, + const float *corr0, + const int32_t *frame_start, const int32_t *frame_count); + + // Per space group: the dense ASU-group id per obs, and a group-ordered permutation of the obs whose + // group >= 0 (group_perm), with its CSR (group_start/group_count, length n_groups) - so each group's + // observations are a contiguous, fixed-order segment for the reduction. + void SetGroups(int n_groups, const int32_t *group, + const int32_t *group_perm, int n_group_perm, + const int32_t *group_start, const int32_t *group_count); + + // Re-upload the working corr (length n_obs) before a scaling pass (the host mutates it via smooth-G + // between passes). SetPartials uploads the initial corr; this refreshes it. + void SetCorr(const float *corr); + + // Run `iters` of {reduce group means -> per-frame robust IRLS G -> update corr} on the device, + // in place on the resident corr. Rotation model (partiality folded via the stored partiality). + void ScalePartials(int iters, double robust_k, double min_partiality, bool has_d_min); + + // Copy the updated corr back to the host (length n_obs), and the fitted per-frame G (length + // n_frames, as double) plus the per-frame "was fitted" flag. + void GetCorr(float *corr_out) const; + void GetG(double *g_out, uint8_t *scaled_out) const; + +private: + struct Impl; + std::unique_ptr impl_; +};