leonarski_fandClaude Opus 4.8
dd25de461d
First stage of moving the rotation scale/merge onto the GPU. The per-frame partial-scaling loop (inverse-variance group-mean reduction -> robust per-frame IRLS G -> corr update, x scaling_iter) now runs in RotationScaleMergeGPU (.cu) when a GPU is present; the CPU loops remain the fallback. The host keeps the one-time raw-hkl sort and the per-space-group gemmi ASU keying, and hands the GPU a group-ordered permutation + CSR so the per-group reduction is a DETERMINISTIC segmented reduction (one thread per group, fixed order, no atomics) - preserving the run-to-run determinism just won on the CPU path (a float atomicAdd reduction would have re-introduced jitter). Reduction is one-thread-per-group (groups average tens of obs, so a block-per-group wastes threads); the IRLS is one block per frame with a deterministic shared-memory reduction. Validated: bit-identical to the CPU path and deterministic run-to-run on lyso/cytC/Ins_H/pding (P41212 ISa 7.8 CC1/2 99.7%, etc.). The scaling kernels are ~7x faster than the CPU compute (~36 ms for 3 iters vs ~0.28 s); end-to-end scale/merge ~2.0 -> ~1.5 s. The remaining gap to the <1 s target is the per-pass host round-trip (corr down/upload for the CPU combine + per-SG group-CSR rebuild); phase 2 keeps the data resident by moving the 3D combine and the merge/error-model onto the GPU too, so nothing round-trips. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
290 lines
14 KiB
Plaintext
290 lines
14 KiB
Plaintext
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "RotationScaleMergeGPU.h"
|
|
|
|
#include <algorithm>
|
|
#include <cmath>
|
|
#include <string>
|
|
#include <cuda_runtime.h>
|
|
|
|
#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 <typename T>
|
|
void Upload(CudaDevicePtr<T> &dst, const T *src, int n) {
|
|
dst = CudaDevicePtr<T>(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<float> I, sigma, rlp, partiality, zeta;
|
|
CudaDevicePtr<uint8_t> on_ice;
|
|
CudaDevicePtr<int32_t> frame;
|
|
CudaDevicePtr<float> corr; // mutable, resident across iterations
|
|
CudaDevicePtr<int32_t> frame_start, frame_count;
|
|
// per space group
|
|
CudaDevicePtr<int32_t> group, group_perm, group_start, group_count;
|
|
// scratch
|
|
CudaDevicePtr<double> group_mean, g;
|
|
CudaDevicePtr<uint8_t> scaled;
|
|
CudaDevicePtr<float> sco_coeff;
|
|
CudaDevicePtr<uint8_t> sco_ok;
|
|
};
|
|
|
|
RotationScaleMergeGPU::RotationScaleMergeGPU() : impl_(std::make_unique<Impl>()) {
|
|
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<double>(n_frames);
|
|
d.scaled = CudaDevicePtr<uint8_t>(n_frames);
|
|
d.sco_coeff = CudaDevicePtr<float>(n_obs);
|
|
d.sco_ok = CudaDevicePtr<uint8_t>(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<double>(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<<<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());
|
|
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());
|
|
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(), 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());
|
|
}
|
|
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");
|
|
}
|