leonarski_fandClaude Opus 5.5
5c5eb00904
The delta-CC1/2 disposition keeps a weak stretch in the merge (downgraded) wherever removing it would not raise CC1/2, and the merge then carries each of its observations at the small 1/sigma^2 its scaled-up counting error gives it. R_meas counted those observations at full weight, so the frames that add next to nothing to the intensities set the number. hq-pool battery: 8xtf kept 115 weak frames (scale 0.1-0.2 of the run's) with the same CC1/2, <I/sigma> and CC_model per shell as rc173 with 84 frames rejected, and R_meas was 1.3-1.5x higher in every shell; 9w3y (33 deg rejected -> 0, every model metric better) and lyso_x10sa_strong read the same way. The disposition itself is not segmentation-dependent: conviction is on the batch grid, not the ledger ranges, and those sets' dispositions changed because the corrected data changed. R_meas now weights each observation by its merge weight v = 1/sigma^2 under the error model (corrected_sigma on the host, ModelSigma on the GPU, with the error model of the last MergeAccum), normalised per reflection to Kish's effective count, so equal sigmas give the ordinary formula (WeightedRmeasTerms). Where the proportional term dominates - strong reflections - frames are weighted alike, as in the merge: 8xtf's lowest shell reads 15.0%, the rc173 run with 84 frames rejected 15.2%. The per-hand table is weighted the same way. R_MEAS_UNWEIGHTED / REFRES_R_MEAS_UNWEIGHTED keep the XDS/AIMLESS convention and are what to set beside XDS; the battery scorer records them. MULTIPLICITY stays a count. Effect (weighted / unweighted): lyso_x06da_ref 0.0454 / 0.0479, 8xtf 0.225 / 0.907, insu_I_x06da_ref REFRES 0.072 / 0.232. GPU and CPU paths agree to the last printed digit on lyso_x06da_ref. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01D1G8gJVAy6gp1K5Dz3NE5C
1345 lines
74 KiB
Plaintext
1345 lines
74 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];
|
|
}
|
|
|
|
// The observation's inverse-variance weight in the scaling reference (0 = not in it), and its
|
|
// scaled intensity. The same filter as the host's ReferenceWeight.
|
|
__device__ __forceinline__ double ReferenceWeight(float I, float sigma, float partiality, float corr,
|
|
double min_partiality, float &I_corr) {
|
|
if (!(corr > 0.0f) || !isfinite(corr)) return 0.0;
|
|
if (partiality < min_partiality) return 0.0;
|
|
const float Ic = I * corr;
|
|
const float sigma_corr = sigma * corr;
|
|
if (!isfinite(Ic) || !isfinite(sigma_corr) || sigma_corr <= 0.0f) return 0.0;
|
|
I_corr = Ic;
|
|
return 1.0 / (double(sigma_corr) * sigma_corr);
|
|
}
|
|
|
|
// 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];
|
|
float I_corr;
|
|
const double w = ReferenceWeight(I[i], sigma[i], partiality[i], corr[i], min_partiality, I_corr);
|
|
if (w <= 0.0) continue;
|
|
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/prescaling_corr) * mean[group].
|
|
__global__ void PrepScaleObsKernel(int n_obs, double min_partiality, const int32_t *__restrict__ group,
|
|
const float *__restrict__ partiality, const float *__restrict__ prescaling_corr,
|
|
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;
|
|
const int g = group[i];
|
|
// Only the observations the reference is built from (the host FitPerFrameG says why).
|
|
bool ok = (g >= 0) && !on_ice[i] && isfinite(zeta[i]) && zeta[i] > 0.0f && partiality[i] >= min_partiality;
|
|
double mean = 0.0;
|
|
if (ok) {
|
|
mean = group_mean[g];
|
|
ok = isfinite(mean);
|
|
}
|
|
// sigma never changes once it is uploaded, so its reciprocal is the same in every scaling
|
|
// pass; hoisted as the CPU does (ScaleObs::weight).
|
|
inv_sigma[i] = SafeInvD(sigma[i], 1.0);
|
|
sco_ok[i] = ok ? 1 : 0;
|
|
sco_coeff[i] = ok ? float(double(partiality[i]) * SafeInvD(prescaling_corr[i], 1.0) * mean) : 0.0f;
|
|
}
|
|
|
|
// One block per frame: the per-frame scale G as the weighted least-squares slope over the frame's
|
|
// contiguous obs - the same plain fit as the host SolveScale (which says why it is not robust).
|
|
// 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,
|
|
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];
|
|
|
|
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
|
|
|
|
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;
|
|
}
|
|
const double tnum = BlockReduceSum(num, sh); __syncthreads();
|
|
const double tden = BlockReduceSum(den, sh);
|
|
if (threadIdx.x == 0) {
|
|
const double G = tden > 0.0 ? tnum / tden : NAN;
|
|
g[f] = isfinite(G) ? fmax(0.0, G) : 1.0;
|
|
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]]);
|
|
}
|
|
}
|
|
|
|
// Zero corr where the rocking geometry is too tangential for the de-novo search, counting what that
|
|
// removed from the merge. zeta is compared in double, as the host does with a double threshold. The
|
|
// count is an integer sum, so the order the atomics land in cannot change it.
|
|
__global__ void FilterZetaKernel(int n_obs, double min_zeta, const float *__restrict__ zeta,
|
|
float *__restrict__ corr, unsigned long long *dropped) {
|
|
unsigned long long local = 0;
|
|
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
|
|
const float z = zeta[i];
|
|
if (isfinite(z) && double(z) >= min_zeta) continue;
|
|
const float c = corr[i];
|
|
if (isfinite(c) && c > 0.0f) ++local;
|
|
corr[i] = 0.0f;
|
|
}
|
|
if (local > 0) atomicAdd(dropped, local);
|
|
}
|
|
|
|
// Zero corr on the rejected frames (--min-image-cc), grid-stride over the resident corr.
|
|
__global__ void FilterFrameKernel(int n_obs, const int32_t *__restrict__ frame,
|
|
const uint8_t *__restrict__ reject, float *__restrict__ corr) {
|
|
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x)
|
|
if (reject[frame[i]]) corr[i] = 0.0f;
|
|
}
|
|
|
|
// corr = prescaling_corr / (partiality * G[frame]) for fitted frames; unchanged otherwise (grid-stride).
|
|
__global__ void UpdateCorrKernel(int n_obs, const int32_t *__restrict__ frame,
|
|
const float *__restrict__ prescaling_corr, 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(prescaling_corr[i])) && isfinite(denom) && denom > 0.0)
|
|
? float(prescaling_corr[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; }
|
|
|
|
// 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;
|
|
// RotationScaleMerge::max_frame_gap, passed in rather than recomputed here so that the host
|
|
// and the device compare the same float and the combine stays bit-identical on both paths.
|
|
float max_frame_gap;
|
|
const float *__restrict__ I, *__restrict__ sigma, *__restrict__ corr,
|
|
*__restrict__ partiality, *__restrict__ bkg, *__restrict__ var_bkg,
|
|
*__restrict__ image_number, *__restrict__ d, *__restrict__ px, *__restrict__ py;
|
|
const int32_t *__restrict__ frame;
|
|
const uint8_t *__restrict__ on_ice, *__restrict__ clipped;
|
|
const int32_t *__restrict__ perm, *__restrict__ rr_start, *__restrict__ rr_count,
|
|
*__restrict__ rr_h, *__restrict__ rr_k, *__restrict__ rr_l,
|
|
*__restrict__ rr_group;
|
|
int32_t *rr_nevents; // count pass output
|
|
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, *f_clipped;
|
|
};
|
|
|
|
// 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 > p.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];
|
|
bool clipped = false;
|
|
for (int m = ev_start; m <= ev_end; ++m) clipped = clipped || p.clipped[p.perm[m]];
|
|
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;
|
|
// a_var has no F in it, so these two are the same in every reweight and only
|
|
// the last round's values are ever read. Two thirds of them were two divisions
|
|
// each, thrown away.
|
|
if (iter == 2) {
|
|
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_clipped[o] = clipped ? 1 : 0;
|
|
p.f_group[o] = group;
|
|
}
|
|
++n_emit;
|
|
}
|
|
|
|
if (!Emit)
|
|
p.rr_nevents[r] = n_emit;
|
|
}
|
|
|
|
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) =====
|
|
|
|
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 float *reject_var_add; // per group: the shell's measured share of dI^2/4
|
|
const int32_t *group, *frame;
|
|
const uint8_t *on_ice, *frame_cell_ok, *half;
|
|
const uint8_t *hand, *has_hands; // per-full Bijvoet hand; per-group "two hands to split"
|
|
const int32_t *gperm, *gstart, *gcount;
|
|
const double *em_mean, *merged_I, *cc_factor;
|
|
// outputs
|
|
double *sw, *swI, *em_mean_out;
|
|
double *swh, *swIh; // per (group, hand), indexed 2*g + hand
|
|
int32_t *cnt, *cnth;
|
|
double *s2, *I2, *dev2;
|
|
uint8_t *valid;
|
|
double *a_swI, *a_sw, *a_swIh0, *a_swIh1, *a_swh0, *a_swh1, *a_swht0, *a_swht1, *a_d;
|
|
int32_t *a_nh0, *a_nh1, *a_rejected;
|
|
uint8_t *a_on_ice;
|
|
double *r_absdev, *r_sumI, *r_wabsdev, *r_wsumI, *r_sumv, *r_sumv2;
|
|
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, swh[2] = {0.0, 0.0}, swIh[2] = {0.0, 0.0};
|
|
int cnt = 0, cnth[2] = {0, 0};
|
|
const bool split = p.has_hands && p.has_hands[g];
|
|
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);
|
|
const double wI = w * (double(p.I[i]) * p.corr[i]);
|
|
sw += w; swI += wI; ++cnt;
|
|
if (split) {
|
|
const int hh = p.hand[i];
|
|
swh[hh] += w; swIh[hh] += wI; ++cnth[hh];
|
|
}
|
|
}
|
|
if (p.swh) {
|
|
p.swh[2 * g] = swh[0]; p.swh[2 * g + 1] = swh[1];
|
|
p.swIh[2 * g] = swIh[0]; p.swIh[2 * g + 1] = swIh[1];
|
|
p.cnth[2 * g] = cnth[0]; p.cnth[2 * g + 1] = cnth[1];
|
|
}
|
|
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;
|
|
// The hand's own mean where the hand has two of its own, the pooled pair where it has
|
|
// not - see the host obs_hand block. em_mean (the merge weights) is untouched.
|
|
const int sub = 2 * g + (p.hand ? p.hand[i] : 0);
|
|
const bool on_hand = p.swh && p.has_hands[g] && p.cnth[sub] >= 2 && p.swh[sub] > 0.0;
|
|
const double mean = on_hand ? p.swIh[sub] / p.swh[sub] : 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) / (on_hand ? p.swh[sub] : 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;
|
|
}
|
|
}
|
|
|
|
// The full's sigma under the error model with its variance at intensity I_for_b - the host model_sigma.
|
|
// At the reflection's expected intensity the merge weight cannot know this full's own fluctuation.
|
|
__device__ __forceinline__ float ModelSigma(int i, const MergeParams &p, double I_for_b, float sigma_raw) {
|
|
const double bi = p.error_model_b * I_for_b;
|
|
const double c = p.corr[i];
|
|
double a_var = double(sigma_raw) * sigma_raw;
|
|
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;
|
|
return v > 0.0 ? float(sqrt(v)) : sigma_raw;
|
|
}
|
|
|
|
// 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, swht0 = 0, swht1 = 0, dd = NAN;
|
|
int nh0 = 0, nh1 = 0, rejected = 0, onice = 0;
|
|
const float rmed = p.reject_median[g];
|
|
// The pooled cut, widened by the shell's measured Bijvoet variance - see the host loop.
|
|
const double v_add = p.reject_var_add ? p.reject_var_add[g] : 0.0;
|
|
for (int q = lo; q < hi; ++q) {
|
|
const int i = p.gperm[q];
|
|
if (!MergeUsable(i, p)) continue;
|
|
if (p.rejected_obs[i]) { ++rejected; continue; } // by the host's Wilson test
|
|
const float I_corr = p.I[i] * p.corr[i];
|
|
const float sigma_raw = p.sigma[i] * p.corr[i];
|
|
const float sigma_corr = p.error_model_active
|
|
? ModelSigma(i, p, isfinite(p.em_mean[g]) ? p.em_mean[g] : double(I_corr), sigma_raw)
|
|
: sigma_raw;
|
|
if (p.reject_outliers && p.error_model_active && isfinite(rmed)) {
|
|
const double tol = p.reject_nsigma * sqrt(double(sigma_corr) * sigma_corr + v_add);
|
|
if (fabs(double(I_corr) - rmed) > tol) {
|
|
++rejected; p.rejected_obs[i] = 1; continue;
|
|
}
|
|
}
|
|
const double w = 1.0 / (double(sigma_corr) * sigma_corr);
|
|
const double wI = w * I_corr;
|
|
// The half-sets are assigned once on the host (AssignHalvesByRank), not decided here:
|
|
// a rank within the reflection cannot be read off a running count, and computing it
|
|
// in one place is what makes this kernel and the host merge loop agree exactly.
|
|
const int half = p.half[i];
|
|
const double wt = w * p.cc_factor[p.frame[i]]; // see the host merge loop (cc_weight)
|
|
swI += wI; sw += w;
|
|
if (half) { swIh1 += wI; swh1 += w; swht1 += wt; ++nh1; }
|
|
else { swIh0 += wI; swh0 += w; swht0 += wt; ++nh0; }
|
|
if (p.on_ice[i]) onice = 1; // see the host merge loop
|
|
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_swht0[g] = swht0; p.a_swht1[g] = swht1;
|
|
p.a_nh0[g] = nh0; p.a_nh1[g] = nh1; p.a_d[g] = dd;
|
|
p.a_rejected[g] = rejected; p.a_on_ice[g] = uint8_t(onice);
|
|
}
|
|
}
|
|
|
|
// One thread per group: R_meas accumulators (sum|I_corr - merged_I|, sum I_corr, the same weighted by
|
|
// v = the observation's merge weight 1/sigma^2 under the error model (as MergeAccumKernel weights it),
|
|
// sum v, sum v^2, 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, wabsdev = 0, wsumI = 0, sumv = 0, sumv2 = 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];
|
|
const float sigma_raw = p.sigma[i] * p.corr[i];
|
|
const double sc = p.error_model_active
|
|
? ModelSigma(i, p, isfinite(p.em_mean[g]) ? p.em_mean[g] : I_corr, sigma_raw)
|
|
: sigma_raw;
|
|
const double v = 1.0 / (sc * sc);
|
|
const double dev = fabs(I_corr - mI);
|
|
absdev += dev; sumI += I_corr;
|
|
wabsdev += v * dev; wsumI += v * I_corr; sumv += v; sumv2 += v * v; ++n;
|
|
}
|
|
}
|
|
p.r_absdev[g] = absdev; p.r_sumI[g] = sumI; p.r_wabsdev[g] = wabsdev; p.r_wsumI[g] = wsumI;
|
|
p.r_sumv[g] = sumv; p.r_sumv2[g] = sumv2;
|
|
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, prescaling_corr, partiality, zeta;
|
|
CudaDevicePtr<uint8_t> on_ice, clipped;
|
|
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;
|
|
CudaDevicePtr<uint8_t> filter_reject; // per-frame --min-image-cc rejection flag, length n_frames
|
|
// 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_swh, m_swIh; // per (group, hand), 2 * n_groups
|
|
CudaDevicePtr<int32_t> m_cnth;
|
|
CudaDevicePtr<uint8_t> m_hand; // per full: Bijvoet hand (0 = I(+))
|
|
CudaDevicePtr<uint8_t> m_has_hands; // per group: acentric, and the merge pools the hands
|
|
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<uint8_t> m_half; // per-full CC1/2 half-set, assigned on the host
|
|
CudaDevicePtr<double> a_swI, a_sw, a_swIh0, a_swIh1, a_swh0, a_swh1, a_swht0, a_swht1, a_d; // merge accum per group
|
|
CudaDevicePtr<double> cc_factor; // per-frame CC1/2 weight factor (length n_frames)
|
|
CudaDevicePtr<int32_t> a_nh0, a_nh1, a_rejected;
|
|
CudaDevicePtr<uint8_t> a_on_ice; // per group: any contributing full on an ice ring
|
|
CudaDevicePtr<float> reject_median;
|
|
CudaDevicePtr<float> reject_var_add; // the pooled test's widening
|
|
CudaDevicePtr<double> merged_I, r_absdev, r_sumI, r_wabsdev, r_wsumI, r_sumv, r_sumv2; // 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;
|
|
double merge_em_a = 1.0, merge_em_b = 0.0; // the error model of the last MergeAccum, for MergeRmeas
|
|
int merge_em_active = 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_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, f_clipped;
|
|
// scale-fulls (Unity model, kept resident): all-ones partiality/prescaling_corr/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::SetPartialsLayout(int n_obs, int n_frames,
|
|
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.frame_start, frame_start, n_frames); Upload(d.frame_count, frame_count, n_frames);
|
|
const int n = std::max(1, n_obs);
|
|
d.I = CudaDevicePtr<float>(n); d.sigma = CudaDevicePtr<float>(n);
|
|
d.prescaling_corr = CudaDevicePtr<float>(n); d.partiality = CudaDevicePtr<float>(n);
|
|
d.zeta = CudaDevicePtr<float>(n); d.corr = CudaDevicePtr<float>(n);
|
|
d.bkg = CudaDevicePtr<float>(n); d.var_bkg = CudaDevicePtr<float>(n);
|
|
d.image_number = CudaDevicePtr<float>(n); d.d_obs = CudaDevicePtr<float>(n);
|
|
d.px_obs = CudaDevicePtr<float>(n); d.py_obs = CudaDevicePtr<float>(n);
|
|
d.frame = CudaDevicePtr<int32_t>(n);
|
|
d.on_ice = CudaDevicePtr<uint8_t>(n);
|
|
d.clipped = CudaDevicePtr<uint8_t>(n);
|
|
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));
|
|
}
|
|
|
|
namespace {
|
|
template <typename T>
|
|
void UploadChunk(CudaDevicePtr<T> &dst, int offset, int count, const T *v) {
|
|
if (count > 0)
|
|
CudaCheck(cudaMemcpy(dst.get() + offset, v, size_t(count) * sizeof(T),
|
|
cudaMemcpyHostToDevice), "upload chunk");
|
|
}
|
|
}
|
|
|
|
void RotationScaleMergeGPU::SetObsField(ObsField f, int offset, int count, const float *v) {
|
|
DeviceGuard guard(impl_->device, impl_->available);
|
|
auto &d = *impl_;
|
|
switch (f) {
|
|
case ObsField::I: UploadChunk(d.I, offset, count, v); break;
|
|
case ObsField::Sigma: UploadChunk(d.sigma, offset, count, v); break;
|
|
case ObsField::PrescalingCorr: UploadChunk(d.prescaling_corr, offset, count, v); break;
|
|
case ObsField::Partiality: UploadChunk(d.partiality, offset, count, v); break;
|
|
case ObsField::Zeta: UploadChunk(d.zeta, offset, count, v); break;
|
|
case ObsField::Corr0: UploadChunk(d.corr, offset, count, v); break;
|
|
case ObsField::Bkg: UploadChunk(d.bkg, offset, count, v); break;
|
|
case ObsField::VarBkg: UploadChunk(d.var_bkg, offset, count, v); break;
|
|
case ObsField::ImageNumber: UploadChunk(d.image_number, offset, count, v); break;
|
|
case ObsField::D: UploadChunk(d.d_obs, offset, count, v); break;
|
|
case ObsField::Px: UploadChunk(d.px_obs, offset, count, v); break;
|
|
case ObsField::Py: UploadChunk(d.py_obs, offset, count, v); break;
|
|
}
|
|
}
|
|
|
|
void RotationScaleMergeGPU::SetObsFrame(int offset, int count, const int32_t *frame) {
|
|
DeviceGuard guard(impl_->device, impl_->available);
|
|
UploadChunk(impl_->frame, offset, count, frame);
|
|
}
|
|
|
|
void RotationScaleMergeGPU::SetObsOnIce(int offset, int count, const uint8_t *on_ice) {
|
|
DeviceGuard guard(impl_->device, impl_->available);
|
|
UploadChunk(impl_->on_ice, offset, count, on_ice);
|
|
}
|
|
|
|
void RotationScaleMergeGPU::SetObsClipped(int offset, int count, const uint8_t *clipped) {
|
|
DeviceGuard guard(impl_->device, impl_->available);
|
|
UploadChunk(impl_->clipped, offset, count, clipped);
|
|
}
|
|
|
|
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 min_partiality, bool /*has_d_min*/) {
|
|
DeviceGuard guard(impl_->device, impl_->available);
|
|
auto &d = *impl_;
|
|
// Reset per call: the host keeps the G of a frame across calls (RunScalingLoop), so a frame this
|
|
// call did not fit must read as unfitted, not as fitted with the value of the call before.
|
|
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());
|
|
CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch");
|
|
PrepScaleObsKernel<<<obs_blocks, BLK>>>(d.n_obs, min_partiality, d.group.get(), d.partiality.get(), d.prescaling_corr.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());
|
|
CudaCheck(cudaGetLastError(), "PrepScaleObsKernel launch");
|
|
FitPerFrameGKernel<<<d.n_frames, BLK>>>(d.n_frames, 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());
|
|
CudaCheck(cudaGetLastError(), "FitPerFrameGKernel launch");
|
|
UpdateCorrKernel<<<upd_blocks, BLK>>>(d.n_obs, d.frame.get(), d.prescaling_corr.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,
|
|
const uint8_t *hand, const uint8_t *has_hands,
|
|
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));
|
|
// The error model is fitted on the Bijvoet hands (see the host obs_hand block). Absent - a merge
|
|
// that already separates them - the hand sums are not allocated and the kernels take the pooled
|
|
// group, which is what they did before.
|
|
if (hand && has_hands) {
|
|
d.m_swh = CudaDevicePtr<double>(std::max(1, 2 * ng)); d.m_swIh = CudaDevicePtr<double>(std::max(1, 2 * ng));
|
|
d.m_cnth = CudaDevicePtr<int32_t>(std::max(1, 2 * ng));
|
|
if (nf > 0) Upload(d.m_hand, hand, nf);
|
|
if (ng > 0) Upload(d.m_has_hands, has_hands, ng);
|
|
} else {
|
|
d.m_swh = CudaDevicePtr<double>(); d.m_swIh = CudaDevicePtr<double>();
|
|
d.m_cnth = CudaDevicePtr<int32_t>();
|
|
d.m_hand = CudaDevicePtr<uint8_t>(); d.m_has_hands = CudaDevicePtr<uint8_t>();
|
|
}
|
|
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.half = d.m_half.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.swh = d.m_swh.get(); p.swIh = d.m_swIh.get(); p.cnth = d.m_cnth.get();
|
|
p.hand = d.m_hand.get(); p.has_hands = d.m_has_hands.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,
|
|
const float *reject_var_add,
|
|
const uint8_t *half, const double *frame_cc_factor,
|
|
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_swht0 = CudaDevicePtr<double>(std::max(1, ng)); d.a_swht1 = 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.a_on_ice = CudaDevicePtr<uint8_t>(std::max(1, ng));
|
|
Upload(d.m_rejected, rejected_obs, d.n_fulls);
|
|
Upload(d.reject_median, reject_median, ng);
|
|
if (reject_var_add && ng > 0) Upload(d.reject_var_add, reject_var_add, ng);
|
|
else d.reject_var_add = CudaDevicePtr<float>();
|
|
Upload(d.m_half, half, d.n_fulls);
|
|
Upload(d.cc_factor, frame_cc_factor, d.n_frames);
|
|
|
|
d.merge_em_a = error_model_a; d.merge_em_b = error_model_b; d.merge_em_active = error_model_active ? 1 : 0;
|
|
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.reject_var_add = d.reject_var_add.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.half = d.m_half.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.cc_factor = d.cc_factor.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_swht0 = d.a_swht0.get(); p.a_swht1 = d.a_swht1.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.a_on_ice = d.a_on_ice.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");
|
|
CudaCheck(cudaMemcpy(rejected_obs, d.m_rejected.get(), size_t(d.n_fulls) * sizeof(uint8_t), cudaMemcpyDeviceToHost),
|
|
"dl rejected_obs");
|
|
}
|
|
|
|
void RotationScaleMergeGPU::MergeAccumRange(int g0, int n, double *swI, double *sw, double *swIh0,
|
|
double *swIh1, double *swh0, double *swh1, double *swh_typ0,
|
|
double *swh_typ1, int32_t *nh0, int32_t *nh1, double *d_out,
|
|
int32_t *rejected, uint8_t *on_ice_out) {
|
|
DeviceGuard guard(impl_->device, impl_->available);
|
|
auto &d = *impl_;
|
|
auto dl = [&](void *h, const auto &s) {
|
|
CudaCheck(cudaMemcpy(h, s.get() + g0, size_t(n) * sizeof(*s.get()), 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(swh_typ0, d.a_swht0); dl(swh_typ1, d.a_swht1);
|
|
dl(d_out, d.a_d);
|
|
dl(nh0, d.a_nh0); dl(nh1, d.a_nh1); dl(rejected, d.a_rejected); dl(on_ice_out, d.a_on_ice);
|
|
}
|
|
|
|
void RotationScaleMergeGPU::MergeRmeas(const double *merged_I, double *absdev, double *sumI, double *wabsdev,
|
|
double *wsumI, double *sumv, double *sumv2, 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_wabsdev = CudaDevicePtr<double>(std::max(1, ng)); d.r_wsumI = CudaDevicePtr<double>(std::max(1, ng));
|
|
d.r_sumv = CudaDevicePtr<double>(std::max(1, ng)); d.r_sumv2 = 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.r_wabsdev = d.r_wabsdev.get(); p.r_wsumI = d.r_wsumI.get();
|
|
p.r_sumv = d.r_sumv.get(); p.r_sumv2 = d.r_sumv2.get();
|
|
p.error_model_a = d.merge_em_a; p.error_model_b = d.merge_em_b; p.error_model_active = d.merge_em_active;
|
|
p.em_mean = d.m_em_mean.get(); p.var_bkg = d.f_var_bkg.get(); p.var_per_I = d.f_var_per_I.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(wabsdev, d.r_wabsdev.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl wabsdev");
|
|
CudaCheck(cudaMemcpy(wsumI, d.r_wsumI.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl wsumI");
|
|
CudaCheck(cudaMemcpy(sumv, d.r_sumv.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl sumv");
|
|
CudaCheck(cudaMemcpy(sumv2, d.r_sumv2.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl sumv2");
|
|
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::SmoothFullsCorr(const uint8_t *apply, const double *ratio) {
|
|
DeviceGuard guard(impl_->device, impl_->available);
|
|
auto &d = *impl_;
|
|
if (d.n_fulls == 0) return;
|
|
Upload(d.smooth_apply, apply, d.n_frames);
|
|
Upload(d.smooth_ratio, ratio, d.n_frames);
|
|
const int blocks = std::min(65535, (d.n_fulls + BLK - 1) / BLK);
|
|
SmoothCorrKernel<<<blocks, BLK>>>(d.n_fulls, d.f_frame.get(), d.smooth_apply.get(),
|
|
d.smooth_ratio.get(), d.f_corr.get());
|
|
CudaCheck(cudaGetLastError(), "smooth fulls corr launch");
|
|
CudaCheck(cudaDeviceSynchronize(), "smooth fulls corr sync");
|
|
}
|
|
|
|
int64_t RotationScaleMergeGPU::FilterCorrByZeta(double min_zeta) {
|
|
DeviceGuard guard(impl_->device, impl_->available);
|
|
auto &d = *impl_;
|
|
CudaDevicePtr<unsigned long long> dropped(1);
|
|
CudaCheck(cudaMemset(dropped.get(), 0, sizeof(unsigned long long)), "zero zeta drop count");
|
|
const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK);
|
|
FilterZetaKernel<<<blocks, BLK>>>(d.n_obs, min_zeta, d.zeta.get(), d.corr.get(), dropped.get());
|
|
CudaCheck(cudaGetLastError(), "zeta filter launch");
|
|
CudaCheck(cudaDeviceSynchronize(), "zeta filter sync");
|
|
unsigned long long n = 0;
|
|
CudaCheck(cudaMemcpy(&n, dropped.get(), sizeof(unsigned long long), cudaMemcpyDeviceToHost),
|
|
"dl zeta drop count");
|
|
return static_cast<int64_t>(n);
|
|
}
|
|
|
|
void RotationScaleMergeGPU::FilterCorrByFrame(const uint8_t *reject) {
|
|
DeviceGuard guard(impl_->device, impl_->available);
|
|
auto &d = *impl_;
|
|
Upload(d.filter_reject, reject, d.n_frames);
|
|
const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK);
|
|
FilterFrameKernel<<<blocks, BLK>>>(d.n_obs, d.frame.get(), d.filter_reject.get(), d.corr.get());
|
|
CudaCheck(cudaGetLastError(), "frame filter launch");
|
|
CudaCheck(cudaDeviceSynchronize(), "frame filter 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());
|
|
CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch");
|
|
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::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_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,
|
|
float max_frame_gap) {
|
|
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.max_frame_gap = max_frame_gap;
|
|
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.clipped = d.clipped.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();
|
|
|
|
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_clipped = 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(); p.f_clipped = d.f_clipped.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,
|
|
uint8_t *clipped, 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));
|
|
dl(clipped, dd.f_clipped.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::ResetFullsScale() {
|
|
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);
|
|
// Unity model: partiality/prescaling_corr/zeta = 1 so coeff = mean; corr starts at 1.
|
|
FillKernel<<<obs_blocks, BLK>>>(d.f_corr.get(), nf, 1.0f);
|
|
CudaCheck(cudaGetLastError(), "FillKernel launch");
|
|
FillKernel<<<obs_blocks, BLK>>>(d.f_partiality.get(), nf, 1.0f);
|
|
CudaCheck(cudaGetLastError(), "FillKernel launch");
|
|
FillKernel<<<obs_blocks, BLK>>>(d.f_rlp.get(), nf, 1.0f);
|
|
CudaCheck(cudaGetLastError(), "FillKernel launch");
|
|
FillKernel<<<obs_blocks, BLK>>>(d.f_zeta.get(), nf, 1.0f);
|
|
CudaCheck(cudaGetLastError(), "FillKernel launch");
|
|
CudaCheck(cudaDeviceSynchronize(), "reset fulls scale sync");
|
|
}
|
|
|
|
void RotationScaleMergeGPU::ScaleFulls(int iters, 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);
|
|
|
|
// Reset per call, as ScalePartials: the host keeps the G of a frame across calls.
|
|
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());
|
|
CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch");
|
|
// 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, min_partiality, 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());
|
|
CudaCheck(cudaGetLastError(), "PrepScaleObsKernel launch");
|
|
FitPerFrameGKernel<<<d.n_frames, BLK>>>(d.n_frames,
|
|
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());
|
|
CudaCheck(cudaGetLastError(), "FitPerFrameGKernel launch");
|
|
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");
|
|
}
|