mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 68 Wiki Activity
Files
56512414aa0ec1b19cccacc3fcbbe91f920e610f
Jungfraujoch/image_analysis/scale_merge/RotationScaleMergeGPU.cu
T
jungfrauandClaude Opus 5 56512414aa Make the partials order total, and hoist 1/sigma out of the IRLS loop 
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>
2026-08-15 18:56:26 -04:00

1156 lines
61 KiB
Plaintext
Raw Blame History

// 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,
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;
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.
// `perm` (null for the partials, whose arrays are already frame-contiguous) maps a position in the
// frame's [lo,hi) range to the obs index, so the same kernel scales the fulls (emit-ordered) through a
// frame-grouping permutation without physically reordering the fulls arrays.
__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 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) {
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[perm ? perm[i] : 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) {
const int a = perm ? perm[i] : i;
if (!sco_ok[a]) continue;
const double coeff = sco_coeff[a];
const double w = inv_sigma[a];
const double w2 = w * w;
num += w2 * coeff * double(I[a]);
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) {
const int a = perm ? perm[i] : i;
if (!sco_ok[a]) continue;
const double coeff = sco_coeff[a];
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);
num += rw * w2 * coeff * double(I[a]);
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; }
}
// 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); }
}
}
}
// SmoothG corr adjust: corr[i] *= ratio[frame[i]] for frames flagged apply, in double then stored
// as float - matching CPU SmoothG's `corr = float(corr * (g/g_smooth))`. Grid-stride, resident corr.
__global__ void SmoothCorrKernel(int n_obs, const int32_t *__restrict__ frame,
const uint8_t *__restrict__ apply, const double *__restrict__ ratio,
float *__restrict__ corr) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
if (!apply[frame[i]]) continue;
const float c = corr[i];
if (isfinite(c)) corr[i] = float(double(c) * ratio[frame[i]]);
}
}
// 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;
}
}
// std::max / std::min return (a<b)?b:a and (b<a)?b:a - reproduce that exactly (NOT fmax/fmin, which
// differ on NaN) so the combine matches the CPU path bit-for-bit on the same inputs.
__device__ __forceinline__ double Dmax(double a, double b) { return (a < b) ? b : a; }
__device__ __forceinline__ double Dmin(double a, double b) { return (b < a) ? b : a; }
__device__ __forceinline__ float Fmax(float a, float b) { return (a < b) ? b : a; }
constexpr float COMBINE_MAX_FRAME_GAP = 2.0f; // == RotationScaleMerge::MAX_FRAME_GAP
// A partial is usable for the combine iff its corr and (I, sigma) are finite with corr>0, sigma>0.
__device__ __forceinline__ bool CombineUsable(int i, const float *I, const float *sigma,
const float *corr) {
const float c = corr[i];
if (!(c > 0.0f) || !isfinite(c)) return false;
return isfinite(I[i]) && isfinite(sigma[i]) && sigma[i] > 0.0f;
}
// All device pointers + scalars the combine kernels need, passed by value.
struct CombineParams {
int n_runs;
double min_partiality, capture_uncertainty_coeff, min_captured_fraction;
const float *I, *sigma, *corr, *partiality, *bkg, *var_bkg, *image_number, *d, *px, *py;
const int32_t *frame;
const uint8_t *on_ice;
const int32_t *perm, *rr_start, *rr_count, *rr_h, *rr_k, *rr_l, *rr_group;
int32_t *rr_nevents, *rr_nusable; // count pass outputs
const int32_t *rr_offset; // emit pass: per-run base offset into the fulls arrays
int32_t *f_h, *f_k, *f_l, *f_frame, *f_group;
float *f_I, *f_sigma, *f_d, *f_img, *f_px, *f_py, *f_var_bkg, *f_var_per_I;
uint8_t *f_on_ice;
};
// One thread's raw-hkl run: split its usable partials (already in image_number order within the run)
// into rocking events, and for each event pool background, seed F and run the 3-iter de-biased Poisson
// reweight - the exact objective of RotationScaleMerge::Combine::process_rawrun. When Emit, write the
// resulting full at rr_offset[r] + (event index); otherwise just count the emitted events. Both modes
// run the identical accept test, so the count pass predicts the emit pass exactly.
template <bool Emit>
__device__ void CombineRawRun(int r, const CombineParams &p) {
const int lo = p.rr_start[r];
const int hi = lo + p.rr_count[r];
const int group = p.rr_group[r];
int n_emit = 0;
int cursor = lo;
while (cursor < hi) {
while (cursor < hi && !CombineUsable(p.perm[cursor], p.I, p.sigma, p.corr)) ++cursor;
if (cursor >= hi) break;
const int ev_start = cursor; // first usable position of the event
int ev_end = cursor; // last usable position (inclusive), extended below
float last_img = p.image_number[p.perm[cursor]];
int probe = cursor + 1;
while (probe < hi) {
while (probe < hi && !CombineUsable(p.perm[probe], p.I, p.sigma, p.corr)) ++probe;
if (probe >= hi) break;
const float img = p.image_number[p.perm[probe]];
if (img - last_img > COMBINE_MAX_FRAME_GAP) break;
last_img = img;
ev_end = probe;
++probe;
}
cursor = ev_end + 1;
// Pass A: pooled background = mean of the event members' finite backgrounds.
double pooled_bkg = 0.0;
int n_pool = 0;
for (int m = ev_start; m <= ev_end; ++m) {
const int i = p.perm[m];
if (!CombineUsable(i, p.I, p.sigma, p.corr)) continue;
const float b = p.bkg[i];
if (isfinite(b)) { pooled_bkg += b; ++n_pool; }
}
pooled_bkg = n_pool > 0 ? pooled_bkg / n_pool : 0.0;
auto pooled_I = [&](int i) -> double {
const double n_bkg = Dmax(0.0, double(p.sigma[i]) * p.sigma[i] - p.I[i])
/ Fmax(p.bkg[i], 1.0f);
return double(p.I[i]) + n_bkg * (double(p.bkg[i]) - pooled_bkg);
};
// Pass B: seed F (inverse-variance mean of pooled_I*corr), plus peak / d / on_ice / partiality.
double sum_w = 0.0, sum_wI = 0.0, sum_partiality = 0.0;
float d = NAN;
const int first = p.perm[ev_start];
int peak_outcome = p.frame[first];
float peak_frame = p.image_number[first];
float peak_px = p.px[first], peak_py = p.py[first];
float peak_partiality = -1.0f;
const bool on_ice = p.on_ice[first];
for (int m = ev_start; m <= ev_end; ++m) {
const int i = p.perm[m];
if (!CombineUsable(i, p.I, p.sigma, p.corr)) continue;
const double sigma_corr = double(p.sigma[i]) * p.corr[i];
const double w = 1.0 / (sigma_corr * sigma_corr);
sum_w += w;
sum_wI += w * pooled_I(i) * p.corr[i];
sum_partiality += p.partiality[i];
if (p.partiality[i] > peak_partiality) {
peak_partiality = p.partiality[i];
peak_outcome = p.frame[i];
peak_frame = p.image_number[i];
peak_px = p.px[i]; peak_py = p.py[i];
}
if (!isfinite(d) && isfinite(p.d[i]) && p.d[i] > 0.0f) d = p.d[i];
}
double F = sum_wI / sum_w;
// Pass C: 3 de-biased Poisson reweights (variance = bkg part + corr*max(0,F)), plus the
// full's var(I) = var_bkg + var_per_I*I for the merge (see the host combine).
double sum_wb = 0.0, sum_cwb = 0.0;
for (int iter = 0; iter < 3; ++iter) {
sum_w = 0.0; sum_wI = 0.0; sum_wb = 0.0; sum_cwb = 0.0;
for (int m = ev_start; m <= ev_end; ++m) {
const int i = p.perm[m];
if (!CombineUsable(i, p.I, p.sigma, p.corr)) continue;
const double corr = p.corr[i];
const double I_corr = pooled_I(i) * corr;
const double sigma_corr = double(p.sigma[i]) * corr;
// The integrator's own non-signal variance (see the host combine).
const double bkg_var = corr * corr * (double) p.var_bkg[i];
const double a_var = bkg_var > 0.0 ? bkg_var : sigma_corr * sigma_corr;
double var = a_var + corr * Dmax(0.0, F);
const double w = 1.0 / var;
sum_w += w;
sum_wI += w * I_corr;
sum_wb += 1.0 / a_var;
sum_cwb += corr / (a_var * a_var);
}
F = sum_wI / sum_w;
}
const double var_bkg_full = 1.0 / sum_wb;
if (sum_w <= 0.0 || sum_partiality < p.min_partiality
|| sum_partiality < p.min_captured_fraction)
continue;
double sigma_full = 1.0 / sqrt(sum_w);
if (p.capture_uncertainty_coeff > 0.0) {
const double frac = Dmin(1.0, sum_partiality);
const double extra = p.capture_uncertainty_coeff * (1.0 - frac) * Dmax(0.0, F);
sigma_full = sqrt(sigma_full * sigma_full + extra * extra);
}
if (Emit) {
const int o = p.rr_offset[r] + n_emit;
p.f_h[o] = p.rr_h[r]; p.f_k[o] = p.rr_k[r]; p.f_l[o] = p.rr_l[r];
p.f_I[o] = float(F);
p.f_sigma[o] = float(sigma_full);
p.f_var_bkg[o] = float(var_bkg_full);
p.f_var_per_I[o] = float(var_bkg_full * var_bkg_full * sum_cwb);
p.f_d[o] = d;
p.f_img[o] = peak_frame;
p.f_px[o] = peak_px; p.f_py[o] = peak_py;
p.f_frame[o] = peak_outcome;
p.f_on_ice[o] = on_ice ? 1 : 0;
p.f_group[o] = group;
}
++n_emit;
}
if (!Emit) {
p.rr_nevents[r] = n_emit;
int n_usable = 0;
for (int m = lo; m < hi; ++m)
if (CombineUsable(p.perm[m], p.I, p.sigma, p.corr)) ++n_usable;
p.rr_nusable[r] = n_usable;
}
}
template <bool Emit>
__global__ void CombineKernel(CombineParams p) {
for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < p.n_runs; r += gridDim.x * blockDim.x)
CombineRawRun<Emit>(r, p);
}
template <typename T>
__global__ void FillKernel(T *p, int n, T v) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) p[i] = v;
}
// ===== error-model + merge reductions over the resident, scaled fulls (mirror MergeAndStats) =====
// Deterministic CC1/2 half from the frame index (splitmix64), matching Merge.cpp / RSM HalfForImage.
__device__ __forceinline__ int HalfForImageDev(long long image_id) {
unsigned long long z = (unsigned long long)image_id + 0x9e3779b97f4a7c15ULL;
z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
z = z ^ (z >> 31);
return (int)(z & 1ULL);
}
struct MergeParams {
int n_groups;
double min_partiality, error_model_a, error_model_b, reject_nsigma;
int for_search, error_model_active, reject_outliers;
const float *I, *sigma, *corr, *partiality, *d, *reject_median, *var_bkg, *var_per_I;
const int32_t *group, *frame;
const uint8_t *on_ice, *frame_cell_ok;
const int32_t *gperm, *gstart, *gcount;
const double *em_mean, *merged_I;
// outputs
double *sw, *swI, *em_mean_out;
int32_t *cnt;
double *s2, *I2, *dev2;
uint8_t *valid;
double *a_swI, *a_sw, *a_swIh0, *a_swIh1, *a_swh0, *a_swh1, *a_d;
int32_t *a_nh0, *a_nh1, *a_rejected;
double *r_absdev, *r_sumI;
int32_t *r_n, *r_nusable;
uint8_t *rejected_obs; // per-full: set by MergeAccum (outlier-rejected), read by MergeRmeas
};
// A full passes the merge / error-model filter (mirrors MergeAndStats::usable_merge).
__device__ __forceinline__ bool MergeUsable(int i, const MergeParams &p) {
const int g = p.group[i];
if (g < 0) return false;
if (!p.frame_cell_ok[p.frame[i]]) return false;
const float c = p.corr[i];
if (!(c > 0.0f) || !isfinite(c)) return false;
if (p.for_search && p.on_ice[i]) return false;
if (p.partiality[i] < p.min_partiality) return false;
const float I_corr = p.I[i] * c, sigma_corr = p.sigma[i] * c;
return isfinite(I_corr) && isfinite(sigma_corr) && sigma_corr > 0.0f;
}
// The looser R_meas filter (no ice / for_search - Mask = cell only).
__device__ __forceinline__ bool RmeasUsable(int i, const MergeParams &p) {
const int g = p.group[i];
if (g < 0) return false;
if (!p.frame_cell_ok[p.frame[i]]) return false;
const float c = p.corr[i];
if (!(c > 0.0f) || !isfinite(c)) return false;
if (p.partiality[i] < p.min_partiality) return false;
const float I_corr = p.I[i] * c, sigma_corr = p.sigma[i] * c;
return isfinite(I_corr) && isfinite(sigma_corr) && sigma_corr > 0.0f;
}
// One thread per group: inverse-variance sums over the group's usable fulls + the group mean (>=2).
__global__ void MergeEmStatsKernel(MergeParams p) {
for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
const int lo = p.gstart[g], hi = lo + p.gcount[g];
double sw = 0.0, swI = 0.0;
int cnt = 0;
for (int q = lo; q < hi; ++q) {
const int i = p.gperm[q];
if (!MergeUsable(i, p)) continue;
const double sigma_corr = double(p.sigma[i]) * p.corr[i];
const double w = 1.0 / (sigma_corr * sigma_corr);
sw += w; swI += w * (double(p.I[i]) * p.corr[i]); ++cnt;
}
p.sw[g] = sw; p.swI[g] = swI; p.cnt[g] = cnt;
p.em_mean_out[g] = (cnt >= 2 && sw > 0.0) ? swI / sw : NAN;
}
}
// One thread per full: the leverage-corrected error-model sample, or valid=0 if dropped.
__global__ void MergeSamplesKernel(int n_obs, MergeParams p) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
p.valid[i] = 0;
if (!MergeUsable(i, p)) continue;
const int g = p.group[i];
if (p.cnt[g] < 2) continue;
const double mean = p.em_mean[g];
if (!isfinite(mean)) continue;
const double sigma_corr = double(p.sigma[i]) * p.corr[i];
const double s2 = sigma_corr * sigma_corr;
const double factor = 1.0 - (1.0 / s2) / p.sw[g];
if (factor < 0.05) continue;
const double resid = double(p.I[i]) * p.corr[i] - mean;
p.s2[i] = s2; p.I2[i] = mean * mean; p.dev2[i] = resid * resid / factor; p.valid[i] = 1;
}
}
// One thread per group: the merge accumulators (inv-var sums + deterministic half-sets), with the
// error-model-corrected sigma. Mirrors MergeAndStats' merge loop (reject path stays on the host).
__global__ void MergeAccumKernel(MergeParams p) {
for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
const int lo = p.gstart[g], hi = lo + p.gcount[g];
double swI = 0, sw = 0, swIh0 = 0, swIh1 = 0, swh0 = 0, swh1 = 0, dd = NAN;
int nh0 = 0, nh1 = 0, rejected = 0;
const float rmed = p.reject_median[g];
for (int q = lo; q < hi; ++q) {
const int i = p.gperm[q];
if (!MergeUsable(i, p)) continue;
const float I_corr = p.I[i] * p.corr[i];
float sigma_corr = p.sigma[i] * p.corr[i];
if (p.error_model_active) {
const double I_for_b = isfinite(p.em_mean[g]) ? p.em_mean[g] : double(I_corr);
const double bi = p.error_model_b * I_for_b;
// The full's variance rebuilt at the reflection's expected intensity, so the merge
// weight cannot know this full's own fluctuation (see the host corrected_sigma).
const double c = p.corr[i];
double a_var = double(sigma_corr) * sigma_corr;
const double base = c * c * p.var_bkg[i] + c * p.var_per_I[i] * Dmax(0.0, I_for_b);
if (base > 0.0) a_var = base;
const double v = p.error_model_a * a_var + bi * bi;
if (v > 0.0) sigma_corr = float(sqrt(v));
}
if (p.reject_outliers && p.error_model_active && isfinite(rmed)
&& fabsf(I_corr - rmed) > p.reject_nsigma * sigma_corr) {
++rejected; p.rejected_obs[i] = 1; continue;
}
const double w = 1.0 / (double(sigma_corr) * sigma_corr);
const double wI = w * I_corr;
const int half = HalfForImageDev(p.frame[i]);
swI += wI; sw += w;
if (half) { swIh1 += wI; swh1 += w; ++nh1; } else { swIh0 += wI; swh0 += w; ++nh0; }
if (!isfinite(dd) && isfinite(p.d[i]) && p.d[i] > 0.0f) dd = p.d[i];
}
p.a_swI[g] = swI; p.a_sw[g] = sw; p.a_swIh0[g] = swIh0; p.a_swIh1[g] = swIh1;
p.a_swh0[g] = swh0; p.a_swh1[g] = swh1; p.a_nh0[g] = nh0; p.a_nh1[g] = nh1; p.a_d[g] = dd;
p.a_rejected[g] = rejected;
}
}
// One thread per group: R_meas accumulators (sum|I_corr - merged_I|, sum_I, n) + the count of
// observations this looser walk accepted. Mirrors MergeAndStats' R_meas re-walk (cell-only filter).
// That count is NOT the per-shell total_observations - it is wider than the merge, so it would
// over-report multiplicity; the host uses it only to skip empty groups.
__global__ void MergeRmeasKernel(MergeParams p) {
for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
const int lo = p.gstart[g], hi = lo + p.gcount[g];
const double mI = p.merged_I[g];
const bool have = isfinite(mI);
double absdev = 0, sumI = 0;
int n = 0, nusable = 0;
for (int q = lo; q < hi; ++q) {
const int i = p.gperm[q];
if (!RmeasUsable(i, p)) continue;
if (p.rejected_obs[i]) continue; // outlier-rejected in the merge -> also out of R_meas (XDS convention)
if (!isfinite(p.d[i]) || !(p.d[i] > 0.0f)) continue; // host counts only fulls with a shell
++nusable;
if (have) {
const double I_corr = double(p.I[i]) * p.corr[i];
absdev += fabs(I_corr - mI); sumI += I_corr; ++n;
}
}
p.r_absdev[g] = absdev; p.r_sumI[g] = sumI; p.r_n[g] = n; p.r_nusable[g] = nusable;
}
}
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 {
int device = 0; // the GPU this instance's buffers live on
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<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
CudaDevicePtr<int64_t> cc_n;
CudaDevicePtr<uint8_t> smooth_apply; // per-frame smooth-G apply flag + ratio, length n_frames
CudaDevicePtr<double> smooth_ratio;
// merge / error-model reductions over the resident fulls (reuse the fulls group CSR f_gperm/...)
CudaDevicePtr<uint8_t> frame_cell_ok;
CudaDevicePtr<double> m_sw, m_swI, m_em_mean; // per group (n_groups)
CudaDevicePtr<int32_t> m_cnt;
CudaDevicePtr<double> m_s2, m_I2, m_dev2; // per full (n_fulls)
CudaDevicePtr<uint8_t> m_valid;
CudaDevicePtr<uint8_t> m_rejected; // per-full outlier-rejected flag (MergeAccum -> MergeRmeas)
CudaDevicePtr<double> a_swI, a_sw, a_swIh0, a_swIh1, a_swh0, a_swh1, a_d; // merge accum per group
CudaDevicePtr<int32_t> a_nh0, a_nh1, a_rejected;
CudaDevicePtr<float> reject_median;
CudaDevicePtr<double> merged_I, r_absdev, r_sumI; // R_meas per group (merged_I uploaded)
CudaDevicePtr<int32_t> r_n, r_nusable;
int merge_for_search = 0; // filter context for one MergeAndStats call
double merge_min_part = 0.0;
// combine: extra per-obs inputs + the one-time raw-hkl run layout
CudaDevicePtr<float> bkg, var_bkg, image_number, d_obs, px_obs, py_obs;
int n_runs = 0, n_perm = 0;
CudaDevicePtr<int32_t> perm, rr_start, rr_count, rr_h, rr_k, rr_l, rr_group;
CudaDevicePtr<int32_t> rr_nevents, rr_nusable, rr_offset;
// combine: resident fulls SoA (rebuilt each Combine)
int n_fulls = 0;
CudaDevicePtr<int32_t> f_h, f_k, f_l, f_frame, f_group;
CudaDevicePtr<float> f_I, f_sigma, f_d, f_img, f_px, f_py, f_var_bkg, f_var_per_I;
CudaDevicePtr<uint8_t> f_on_ice;
// scale-fulls (Unity model, kept resident): all-ones partiality/rlp/zeta so the shared scaling kernels
// 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;
};
// Set the device this instance's memory lives on for the duration of a call, and put the caller's
// back afterwards. CUDA's current device is per-thread, so without this a single set_gpu() in the
// constructor silently re-pins the calling thread for the rest of its life - and, worse, the
// destructor would free several gigabytes against whatever device happened to be current then.
// CudaDevicePtr records no device of its own, so every entry point needs this.
namespace {
struct DeviceGuard {
int prev = 0;
bool active = false;
explicit DeviceGuard(int device, bool enable) : active(enable) {
if (!active)
return;
cudaGetDevice(&prev);
if (prev != device)
set_gpu(device);
}
~DeviceGuard() {
if (active)
set_gpu(prev);
}
};
} // namespace
RotationScaleMergeGPU::RotationScaleMergeGPU() : impl_(std::make_unique<Impl>()) {
if (get_gpu_count() > 0) {
// One instance, one device. It stays on device 0 for now - the merge is a single object and
// nothing else runs beside it - but it is recorded rather than assumed, so the guard below
// can put the caller's device back instead of leaving the thread moved.
impl_->device = 0;
DeviceGuard guard(impl_->device, true);
impl_->available = true;
}
}
RotationScaleMergeGPU::~RotationScaleMergeGPU() {
DeviceGuard guard(impl_->device, impl_->available);
impl_.reset();
}
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) {
DeviceGuard guard(impl_->device, impl_->available);
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.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));
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,
int n_group_perm, const int32_t *group_start,
const int32_t *group_count) {
DeviceGuard guard(impl_->device, impl_->available);
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) {
DeviceGuard guard(impl_->device, impl_->available);
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*/) {
DeviceGuard guard(impl_->device, impl_->available);
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.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.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());
}
CudaCheck(cudaGetLastError(), "kernel launch");
CudaCheck(cudaDeviceSynchronize(), "scale sync");
}
void RotationScaleMergeGPU::GetCorr(float *corr_out) const {
DeviceGuard guard(impl_->device, impl_->available);
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 {
DeviceGuard guard(impl_->device, impl_->available);
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");
}
void RotationScaleMergeGPU::SetFrameCellOk(const uint8_t *frame_cell_ok) {
DeviceGuard guard(impl_->device, impl_->available);
Upload(impl_->frame_cell_ok, frame_cell_ok, impl_->n_frames);
}
// The per-group inv-var mean (em_mean) + the per-full leverage-corrected error-model samples over the
// resident+scaled fulls. Stashes the filter context for the later MergeAccum/MergeRmeas calls.
void RotationScaleMergeGPU::MergeEmSamples(bool for_search, double min_partiality,
double *em_mean_out, int32_t *cnt_out, double *s2_out,
double *I2_out, double *dev2_out, uint8_t *valid_out) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int ng = d.n_groups, nf = d.n_fulls;
d.merge_for_search = for_search ? 1 : 0; d.merge_min_part = min_partiality;
d.m_sw = CudaDevicePtr<double>(std::max(1, ng)); d.m_swI = CudaDevicePtr<double>(std::max(1, ng));
d.m_em_mean = CudaDevicePtr<double>(std::max(1, ng)); d.m_cnt = CudaDevicePtr<int32_t>(std::max(1, ng));
d.m_s2 = CudaDevicePtr<double>(std::max(1, nf)); d.m_I2 = CudaDevicePtr<double>(std::max(1, nf));
d.m_dev2 = CudaDevicePtr<double>(std::max(1, nf)); d.m_valid = CudaDevicePtr<uint8_t>(std::max(1, nf));
MergeParams p{};
p.n_groups = ng; p.min_partiality = min_partiality;
p.for_search = d.merge_for_search;
p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get(); p.partiality = d.f_partiality.get();
p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get();
p.on_ice = d.f_on_ice.get(); p.frame_cell_ok = d.frame_cell_ok.get();
p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();
p.em_mean = d.m_em_mean.get(); p.sw = d.m_sw.get(); p.swI = d.m_swI.get();
p.em_mean_out = d.m_em_mean.get(); p.cnt = d.m_cnt.get();
p.s2 = d.m_s2.get(); p.I2 = d.m_I2.get(); p.dev2 = d.m_dev2.get(); p.valid = d.m_valid.get();
const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK);
MergeEmStatsKernel<<<grp_blocks, BLK>>>(p);
MergeSamplesKernel<<<obs_blocks, BLK>>>(nf, p);
CudaCheck(cudaGetLastError(), "merge em/samples launch");
CudaCheck(cudaDeviceSynchronize(), "merge em/samples sync");
CudaCheck(cudaMemcpy(em_mean_out, d.m_em_mean.get(), size_t(ng) * sizeof(double),
cudaMemcpyDeviceToHost), "dl em_mean");
CudaCheck(cudaMemcpy(cnt_out, d.m_cnt.get(), size_t(ng) * sizeof(int32_t),
cudaMemcpyDeviceToHost), "dl cnt");
if (nf > 0) {
CudaCheck(cudaMemcpy(s2_out, d.m_s2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost), "dl s2");
CudaCheck(cudaMemcpy(I2_out, d.m_I2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost), "dl I2");
CudaCheck(cudaMemcpy(dev2_out, d.m_dev2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost), "dl dev2");
CudaCheck(cudaMemcpy(valid_out, d.m_valid.get(), size_t(nf) * sizeof(uint8_t), cudaMemcpyDeviceToHost), "dl valid");
}
}
void RotationScaleMergeGPU::MergeAccum(double error_model_a, double error_model_b, bool error_model_active,
bool reject_outliers, double reject_nsigma, const float *reject_median,
double *swI, double *sw, double *swIh0, double *swIh1,
double *swh0, double *swh1, int32_t *nh0, int32_t *nh1, double *d_out,
int32_t *rejected, uint8_t *rejected_obs) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int ng = d.n_groups;
d.a_swI = CudaDevicePtr<double>(std::max(1, ng)); d.a_sw = CudaDevicePtr<double>(std::max(1, ng));
d.a_swIh0 = CudaDevicePtr<double>(std::max(1, ng)); d.a_swIh1 = CudaDevicePtr<double>(std::max(1, ng));
d.a_swh0 = CudaDevicePtr<double>(std::max(1, ng)); d.a_swh1 = CudaDevicePtr<double>(std::max(1, ng));
d.a_nh0 = CudaDevicePtr<int32_t>(std::max(1, ng)); d.a_nh1 = CudaDevicePtr<int32_t>(std::max(1, ng));
d.a_d = CudaDevicePtr<double>(std::max(1, ng)); d.a_rejected = CudaDevicePtr<int32_t>(std::max(1, ng));
d.m_rejected = CudaDevicePtr<uint8_t>(std::max(1, d.n_fulls));
CudaCheck(cudaMemset(d.m_rejected.get(), 0, size_t(std::max(1, d.n_fulls)) * sizeof(uint8_t)),
"zero m_rejected");
Upload(d.reject_median, reject_median, ng);
MergeParams p{};
p.n_groups = ng; p.min_partiality = d.merge_min_part;
p.for_search = d.merge_for_search;
p.error_model_a = error_model_a; p.error_model_b = error_model_b;
p.error_model_active = error_model_active ? 1 : 0;
p.reject_outliers = reject_outliers ? 1 : 0; p.reject_nsigma = reject_nsigma;
p.reject_median = d.reject_median.get();
p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get(); p.partiality = d.f_partiality.get();
p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get();
p.on_ice = d.f_on_ice.get(); p.frame_cell_ok = d.frame_cell_ok.get();
p.var_bkg = d.f_var_bkg.get(); p.var_per_I = d.f_var_per_I.get();
p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();
p.em_mean = d.m_em_mean.get();
p.a_swI = d.a_swI.get(); p.a_sw = d.a_sw.get(); p.a_swIh0 = d.a_swIh0.get(); p.a_swIh1 = d.a_swIh1.get();
p.a_swh0 = d.a_swh0.get(); p.a_swh1 = d.a_swh1.get(); p.a_nh0 = d.a_nh0.get(); p.a_nh1 = d.a_nh1.get();
p.a_d = d.a_d.get(); p.a_rejected = d.a_rejected.get();
p.rejected_obs = d.m_rejected.get();
const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
MergeAccumKernel<<<grp_blocks, BLK>>>(p);
CudaCheck(cudaGetLastError(), "merge accum launch");
CudaCheck(cudaDeviceSynchronize(), "merge accum sync");
auto dl = [&](void *h, const CudaDevicePtr<double> &s) {
CudaCheck(cudaMemcpy(h, s.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl accum"); };
dl(swI, d.a_swI); dl(sw, d.a_sw); dl(swIh0, d.a_swIh0); dl(swIh1, d.a_swIh1);
dl(swh0, d.a_swh0); dl(swh1, d.a_swh1); dl(d_out, d.a_d);
CudaCheck(cudaMemcpy(nh0, d.a_nh0.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl nh0");
CudaCheck(cudaMemcpy(nh1, d.a_nh1.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl nh1");
CudaCheck(cudaMemcpy(rejected, d.a_rejected.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl rej");
CudaCheck(cudaMemcpy(rejected_obs, d.m_rejected.get(), size_t(d.n_fulls) * sizeof(uint8_t), cudaMemcpyDeviceToHost),
"dl rejected_obs");
}
void RotationScaleMergeGPU::MergeRmeas(const double *merged_I, double *absdev, double *sumI,
int32_t *n, int32_t *nusable) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int ng = d.n_groups;
Upload(d.merged_I, merged_I, ng);
d.r_absdev = CudaDevicePtr<double>(std::max(1, ng)); d.r_sumI = CudaDevicePtr<double>(std::max(1, ng));
d.r_n = CudaDevicePtr<int32_t>(std::max(1, ng)); d.r_nusable = CudaDevicePtr<int32_t>(std::max(1, ng));
MergeParams p{};
p.n_groups = ng; p.min_partiality = d.merge_min_part;
p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get(); p.partiality = d.f_partiality.get();
p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get(); p.frame_cell_ok = d.frame_cell_ok.get();
p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();
p.merged_I = d.merged_I.get();
p.r_absdev = d.r_absdev.get(); p.r_sumI = d.r_sumI.get(); p.r_n = d.r_n.get(); p.r_nusable = d.r_nusable.get();
p.rejected_obs = d.m_rejected.get();
const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
MergeRmeasKernel<<<grp_blocks, BLK>>>(p);
CudaCheck(cudaGetLastError(), "merge rmeas launch");
CudaCheck(cudaDeviceSynchronize(), "merge rmeas sync");
CudaCheck(cudaMemcpy(absdev, d.r_absdev.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl absdev");
CudaCheck(cudaMemcpy(sumI, d.r_sumI.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl sumI");
CudaCheck(cudaMemcpy(n, d.r_n.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl rn");
CudaCheck(cudaMemcpy(nusable, d.r_nusable.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl rnusable");
}
void RotationScaleMergeGPU::SmoothCorr(const uint8_t *apply, const double *ratio) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
Upload(d.smooth_apply, apply, d.n_frames);
Upload(d.smooth_ratio, ratio, d.n_frames);
const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK);
SmoothCorrKernel<<<blocks, BLK>>>(d.n_obs, d.frame.get(), d.smooth_apply.get(),
d.smooth_ratio.get(), d.corr.get());
CudaCheck(cudaGetLastError(), "smooth corr launch");
CudaCheck(cudaDeviceSynchronize(), "smooth corr sync");
}
void RotationScaleMergeGPU::ComputePartialCC(double min_partiality, double *cc_out, int64_t *cc_n_out) {
DeviceGuard guard(impl_->device, impl_->available);
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 *var_bkg,
const float *image_number, const float *d,
const float *px, const float *py) {
DeviceGuard guard(impl_->device, impl_->available);
auto &dd = *impl_;
Upload(dd.bkg, bkg, dd.n_obs);
Upload(dd.var_bkg, var_bkg, dd.n_obs);
Upload(dd.image_number, image_number, dd.n_obs);
Upload(dd.d_obs, d, dd.n_obs);
Upload(dd.px_obs, px, dd.n_obs);
Upload(dd.py_obs, py, dd.n_obs);
}
void RotationScaleMergeGPU::SetRawRuns(int n_runs, int n_perm, const int32_t *perm,
const int32_t *rr_start, const int32_t *rr_count,
const int32_t *rr_h, const int32_t *rr_k, const int32_t *rr_l) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
d.n_runs = n_runs;
d.n_perm = n_perm;
Upload(d.perm, perm, n_perm);
Upload(d.rr_start, rr_start, n_runs);
Upload(d.rr_count, rr_count, n_runs);
Upload(d.rr_h, rr_h, n_runs);
Upload(d.rr_k, rr_k, n_runs);
Upload(d.rr_l, rr_l, n_runs);
d.rr_group = CudaDevicePtr<int32_t>(std::max(1, n_runs));
d.rr_nevents = CudaDevicePtr<int32_t>(std::max(1, n_runs));
d.rr_nusable = CudaDevicePtr<int32_t>(std::max(1, n_runs));
d.rr_offset = CudaDevicePtr<int32_t>(std::max(1, n_runs));
}
int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_partiality,
double capture_uncertainty_coeff, double min_captured_fraction) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
CudaCheck(cudaMemcpy(d.rr_group.get(), rawrun_group, size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyHostToDevice), "upload rr_group");
CombineParams p{};
p.n_runs = d.n_runs;
p.min_partiality = min_partiality;
p.capture_uncertainty_coeff = capture_uncertainty_coeff;
p.min_captured_fraction = min_captured_fraction;
p.I = d.I.get(); p.sigma = d.sigma.get(); p.corr = d.corr.get(); p.partiality = d.partiality.get();
p.bkg = d.bkg.get(); p.var_bkg = d.var_bkg.get(); p.image_number = d.image_number.get(); p.d = d.d_obs.get();
p.px = d.px_obs.get(); p.py = d.py_obs.get();
p.frame = d.frame.get(); p.on_ice = d.on_ice.get();
p.perm = d.perm.get(); p.rr_start = d.rr_start.get(); p.rr_count = d.rr_count.get();
p.rr_h = d.rr_h.get(); p.rr_k = d.rr_k.get(); p.rr_l = d.rr_l.get(); p.rr_group = d.rr_group.get();
p.rr_nevents = d.rr_nevents.get(); p.rr_nusable = d.rr_nusable.get();
const int blocks = std::min(65535, (d.n_runs + BLK - 1) / BLK);
// Count pass: how many fulls each run emits.
CombineKernel<false><<<blocks, BLK>>>(p);
CudaCheck(cudaGetLastError(), "combine count launch");
// Exclusive prefix sum on the host (deterministic) -> per-run output offset + total fulls.
std::vector<int32_t> nevents(d.n_runs);
CudaCheck(cudaMemcpy(nevents.data(), d.rr_nevents.get(), size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyDeviceToHost), "download nevents");
std::vector<int32_t> offset(d.n_runs);
int64_t acc = 0;
for (int r = 0; r < d.n_runs; ++r) { offset[r] = static_cast<int32_t>(acc); acc += nevents[r]; }
d.n_fulls = static_cast<int>(acc);
// Allocate the fulls SoA and emit.
const int nf = std::max(1, d.n_fulls);
d.f_h = CudaDevicePtr<int32_t>(nf); d.f_k = CudaDevicePtr<int32_t>(nf); d.f_l = CudaDevicePtr<int32_t>(nf);
d.f_frame = CudaDevicePtr<int32_t>(nf); d.f_group = CudaDevicePtr<int32_t>(nf);
d.f_I = CudaDevicePtr<float>(nf); d.f_sigma = CudaDevicePtr<float>(nf);
d.f_d = CudaDevicePtr<float>(nf); d.f_img = CudaDevicePtr<float>(nf);
d.f_px = CudaDevicePtr<float>(nf); d.f_py = CudaDevicePtr<float>(nf);
d.f_var_bkg = CudaDevicePtr<float>(nf); d.f_var_per_I = CudaDevicePtr<float>(nf);
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");
p.rr_offset = d.rr_offset.get();
p.f_h = d.f_h.get(); p.f_k = d.f_k.get(); p.f_l = d.f_l.get();
p.f_frame = d.f_frame.get(); p.f_group = d.f_group.get();
p.f_I = d.f_I.get(); p.f_sigma = d.f_sigma.get(); p.f_d = d.f_d.get(); p.f_img = d.f_img.get();
p.f_px = d.f_px.get(); p.f_py = d.f_py.get();
p.f_var_bkg = d.f_var_bkg.get(); p.f_var_per_I = d.f_var_per_I.get();
p.f_on_ice = d.f_on_ice.get();
if (d.n_fulls > 0) {
CombineKernel<true><<<blocks, BLK>>>(p);
CudaCheck(cudaGetLastError(), "combine emit launch");
}
CudaCheck(cudaDeviceSynchronize(), "combine sync");
return d.n_fulls;
}
void RotationScaleMergeGPU::GetFulls(int32_t *h, int32_t *k, int32_t *l, float *I, float *sigma, float *d,
float *image_number, int32_t *frame, uint8_t *on_ice,
int32_t *group) const {
DeviceGuard guard(impl_->device, impl_->available);
const auto &dd = *impl_;
const size_t n = static_cast<size_t>(dd.n_fulls);
if (n == 0) return;
auto dl = [&](void *dst, const void *src, size_t bytes) {
CudaCheck(cudaMemcpy(dst, src, bytes, cudaMemcpyDeviceToHost), "download fulls");
};
dl(h, dd.f_h.get(), n * sizeof(int32_t)); dl(k, dd.f_k.get(), n * sizeof(int32_t));
dl(l, dd.f_l.get(), n * sizeof(int32_t)); dl(frame, dd.f_frame.get(), n * sizeof(int32_t));
dl(group, dd.f_group.get(), n * sizeof(int32_t));
dl(I, dd.f_I.get(), n * sizeof(float)); dl(sigma, dd.f_sigma.get(), n * sizeof(float));
dl(d, dd.f_d.get(), n * sizeof(float)); dl(image_number, dd.f_img.get(), n * sizeof(float));
dl(on_ice, dd.f_on_ice.get(), n * sizeof(uint8_t));
}
void RotationScaleMergeGPU::GetFullsKeys(int32_t *frame, int32_t *group) const {
DeviceGuard guard(impl_->device, impl_->available);
const auto &d = *impl_;
if (d.n_fulls == 0) return;
const size_t bytes = size_t(d.n_fulls) * sizeof(int32_t);
CudaCheck(cudaMemcpy(frame, d.f_frame.get(), bytes, cudaMemcpyDeviceToHost), "download f_frame");
CudaCheck(cudaMemcpy(group, d.f_group.get(), bytes, cudaMemcpyDeviceToHost), "download f_group");
}
void RotationScaleMergeGPU::SetFullsFrameCSR(const int32_t *frame_perm, int n_perm,
const int32_t *frame_start, const int32_t *frame_count) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
Upload(d.f_frame_perm, frame_perm, n_perm);
Upload(d.f_frame_start, frame_start, d.n_frames);
Upload(d.f_frame_count, frame_count, d.n_frames);
}
void RotationScaleMergeGPU::SetFullsGroups(const int32_t *gperm, int n_gperm,
const int32_t *gstart, const int32_t *gcount) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
Upload(d.f_gperm, gperm, n_gperm);
Upload(d.f_gstart, gstart, d.n_groups);
Upload(d.f_gcount, gcount, d.n_groups);
}
void RotationScaleMergeGPU::ScaleFulls(int iters, double robust_k, double min_partiality) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int nf = d.n_fulls;
if (nf == 0) return;
const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK);
const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK);
// Unity model: partiality/rlp/zeta = 1 so coeff = mean; corr starts at 1.
FillKernel<<<obs_blocks, BLK>>>(d.f_corr.get(), nf, 1.0f);
FillKernel<<<obs_blocks, BLK>>>(d.f_partiality.get(), nf, 1.0f);
FillKernel<<<obs_blocks, BLK>>>(d.f_rlp.get(), nf, 1.0f);
FillKernel<<<obs_blocks, BLK>>>(d.f_zeta.get(), nf, 1.0f);
// g/scaled reset once (a frame, once fitted, stays fitted - same as ScalePartials).
CudaCheck(cudaMemset(d.scaled.get(), 0, size_t(d.n_frames) * sizeof(uint8_t)), "memset f scaled");
CudaCheck(cudaMemset(d.g.get(), 0, size_t(d.n_frames) * sizeof(double)), "memset f g");
for (int it = 0; it < iters; ++it) {
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());
// 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_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());
}
CudaCheck(cudaGetLastError(), "scale fulls launch");
CudaCheck(cudaDeviceSynchronize(), "scale fulls sync");
}
void RotationScaleMergeGPU::GetFullsCorr(float *corr) const {
DeviceGuard guard(impl_->device, impl_->available);
const auto &d = *impl_;
if (d.n_fulls == 0) return;
CudaCheck(cudaMemcpy(corr, d.f_corr.get(), size_t(d.n_fulls) * sizeof(float),
cudaMemcpyDeviceToHost), "download f_corr");
}
void RotationScaleMergeGPU::GetFullsPxPy(float *px, float *py) const {
DeviceGuard guard(impl_->device, impl_->available);
const auto &d = *impl_;
if (d.n_fulls == 0) return;
const size_t bytes = size_t(d.n_fulls) * sizeof(float);
CudaCheck(cudaMemcpy(px, d.f_px.get(), bytes, cudaMemcpyDeviceToHost), "download f_px");
CudaCheck(cudaMemcpy(py, d.f_py.get(), bytes, cudaMemcpyDeviceToHost), "download f_py");
}
void RotationScaleMergeGPU::GetFullsVariance(float *var_bkg, float *var_per_I) const {
DeviceGuard guard(impl_->device, impl_->available);
const auto &d = *impl_;
if (d.n_fulls == 0) return;
const size_t bytes = size_t(d.n_fulls) * sizeof(float);
CudaCheck(cudaMemcpy(var_bkg, d.f_var_bkg.get(), bytes, cudaMemcpyDeviceToHost), "download f_var_bkg");
CudaCheck(cudaMemcpy(var_per_I, d.f_var_per_I.get(), bytes, cudaMemcpyDeviceToHost),
"download f_var_per_I");
}
void RotationScaleMergeGPU::SetFullsCorr(const float *corr) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
if (d.n_fulls == 0) return;
CudaCheck(cudaMemcpy(d.f_corr.get(), corr, size_t(d.n_fulls) * sizeof(float),
cudaMemcpyHostToDevice), "upload f_corr");
}
Reference in New Issue View Git Blame Copy Permalink