mx/Jungfraujoch
3
2
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases 81 Wiki Activity
Files
9cf987cd4e7d6663a64ef22300d5e2ee36d473ec
Jungfraujoch/image_analysis/scale_merge/RotationScaleMergeGPU.cu
T
leonarski_fandClaude Opus 5.5 7d656c13de rugnux merge engine: the immutable fields the scaling gathers in one 16-byte record 
The group reduction reaches every observation through the group permutation, and so fetched one
32-byte sector per field for I, sigma, partiality and corr. I, sigma, partiality and prescaling_corr
now sit together in one aligned record, so it fetches two sectors per observation; the grouping
kernels' finiteness test and the combine read the same record. The per-observation fit coefficient,
its flag and 1/sigma are no longer stored: the per-frame fit forms them as it reads each observation
(the same expressions, so the same bits), which removes PrepScaleObsKernel and 13 bytes per
observation of device memory. DeviceBytes' per-observation figure drops by those 13 bytes.

Every sum keeps its order: p.hkl md5 identical to the base on cytc, 8a1a and 8qaw.
RTX 5080, nsys, merge-engine kernels: 8a1a 2459 -> 1924 ms, 8qaw 3624 -> 3025 ms, cytc 501 -> 476 ms
(ReduceGroupMeans 1.5-1.6x faster, PrepScaleObs gone, FitPerFrameG absorbs it).

Measured and not taken: storing all per-observation arrays in group order makes the reduction 3-4x
faster on its own, but the per-frame fit then has to gather, and moving 15 arrays at every grouping
costs as much as the reduction saves (net 0 on the 5080); a record for the combine-only fields
did not speed the combine.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3
2026-10-12 11:21:49 +02:00

2379 lines
131 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 "OutlierBand.h"
#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdio>
#include <string>
#include <vector>
#include <cuda_runtime.h>
#include <cub/device/device_radix_sort.cuh>
#include <cub/device/device_scan.cuh>
#include "../indexing/CUDAMemHelpers.h"
#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"
#include "../../common/ParallelFor.h"
namespace {
constexpr int BLK = 256;
constexpr int MIN_REFLECTIONS = 20;
// Every kernel and copy here is queued on the instance's own stream (Impl::stream), not on the
// legacy NULL stream: the merge and the image analysis of a probe pass beside it would otherwise
// wait for each other's work at every launch and every synchronisation. As in BeamCenterFFTGPU no buffer comes from the pool. A
// pooled buffer is freed with cudaFreeAsync on the thread's allocation stream, which is not
// ordered after this one. Each entry point below waits for its own work before it returns, so no
// free here has yet overtaken a read, but that holds only by that convention, and not at all for a
// free while unwinding from a failed call; nor can compute-sanitizer --track-stream-ordered-races
// see it, and it reported every reassigned merge buffer as a use-after-free. cudaFree
// synchronises the device first.
constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous;
__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 four immutable fields of an observation that the scaling, the combine and the grouping read
// together, in one 16-byte record: the kernels that reach an observation through a permutation
// (the group reduction above all) then fetch one sector for it instead of one per field.
struct alignas(16) ObsRecord { float I, sigma, partiality, prescaling_corr; }; // floats 0, 1, 2, 3
// 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 ObsRecord *__restrict__ obs,
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 ObsRecord o = obs[i];
float I_corr;
const double w = ReferenceWeight(o.I, o.sigma, o.partiality, 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;
}
}
// The scale-fit coefficient of an observation (rotation model), recomputed each scaling iteration
// once the group means are known: coeff = partiality * (1/prescaling_corr) * mean[group], or false
// where the observation is not one the reference is built from (the host FitPerFrameG says why).
__device__ __forceinline__ bool ScaleCoeff(int i, double min_partiality, const int32_t *__restrict__ group,
const ObsRecord &o, const float *__restrict__ zeta,
const uint8_t *__restrict__ on_ice,
const double *__restrict__ group_mean, float &coeff) {
const int g = group[i];
if (!((g >= 0) && !on_ice[i] && isfinite(zeta[i]) && zeta[i] > 0.0f && o.partiality >= min_partiality))
return false;
const double mean = group_mean[g];
if (!isfinite(mean)) return false;
coeff = float(double(o.partiality) * SafeInvD(o.prescaling_corr, 1.0) * mean);
return true;
}
// 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 frames with fewer than min_obs usable observations. `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. info[f] receives the
// fit's information on G, sum w^2 c^2 (0 for a frame left alone). The weight is 1/sigma, as the
// CPU's ScaleObs::weight.
__global__ void FitPerFrameGKernel(int n_frames,
const int32_t *__restrict__ frame_start,
const int32_t *__restrict__ frame_count,
const int32_t *__restrict__ perm, double min_partiality,
const int32_t *__restrict__ group, const ObsRecord *__restrict__ obs,
const float *__restrict__ zeta, const uint8_t *__restrict__ on_ice,
const double *__restrict__ group_mean, long min_obs,
double *__restrict__ g, uint8_t *__restrict__ scaled,
double *__restrict__ info) {
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];
if (threadIdx.x == 0) info[f] = 0.0;
// The sums are formed in the same pass as the count; a frame with too few observations
// discards them.
long cnt_local = 0;
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;
const ObsRecord o = obs[a];
float coeff_f;
if (!ScaleCoeff(a, min_partiality, group, o, zeta, on_ice, group_mean, coeff_f)) continue;
++cnt_local;
const double coeff = coeff_f;
const double w = SafeInvD(o.sigma, 1.0);
const double w2 = w * w;
num += w2 * coeff * double(o.I);
den += w2 * coeff * coeff;
}
const double cnt = BlockReduceSum(double(cnt_local), sh);
__shared__ double s_cnt;
if (threadIdx.x == 0) s_cnt = cnt;
__syncthreads();
if (s_cnt < min_obs) return; // leave g[f]/scaled[f] as-is
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;
info[f] = tden;
}
}
// 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 ObsRecord *__restrict__ obs, 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;
const ObsRecord o = obs[i];
if (o.partiality < min_partiality) continue;
const float c = corr[i];
if (!isfinite(o.I) || !isfinite(c) || !(c > 0.0f)) continue;
if (!isfinite(o.sigma) || !(o.sigma > 0.0f)) continue;
const double mean = group_mean[g];
if (!isfinite(mean)) continue;
const double img = double(o.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 ObsRecord *__restrict__ obs,
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 ObsRecord o = obs[i];
const double denom = double(o.partiality) * g[f];
corr[i] = (isfinite(double(o.prescaling_corr)) && isfinite(denom) && denom > 0.0)
? float(o.prescaling_corr / 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 ObsRecord *obs, const float *corr) {
const float c = corr[i];
if (!(c > 0.0f) || !isfinite(c)) return false;
return isfinite(obs[i].I) && isfinite(obs[i].sigma) && obs[i].sigma > 0.0f;
}
// RotationScaleMerge::OBS_CLIPPED / OBS_OVERLOADED, the flag bits a partial's `clipped` byte carries.
constexpr uint8_t OBS_CLIPPED_FLAG = 1;
constexpr uint8_t OBS_OVERLOADED_FLAG = 2;
// 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 ObsRecord *__restrict__ obs;
const float *__restrict__ corr, *__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
int32_t *rr_noverloaded; // count pass output: events dropped for a saturated pixel
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, *f_capture;
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, n_overloaded = 0;
int cursor = lo;
while (cursor < hi) {
while (cursor < hi && !CombineUsable(p.perm[cursor], p.obs, 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.obs, 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.obs, 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.obs[i].sigma) * p.obs[i].sigma - p.obs[i].I)
/ Fmax(p.bkg[i], 1.0f);
return double(p.obs[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];
uint8_t flags = 0;
for (int m = ev_start; m <= ev_end; ++m) flags |= p.clipped[p.perm[m]];
// An overloaded event is not a measurement of the reflection; see the host combine.
if (flags & OBS_OVERLOADED_FLAG) { ++n_overloaded; continue; }
const bool clipped = flags & OBS_CLIPPED_FLAG;
for (int m = ev_start; m <= ev_end; ++m) {
const int i = p.perm[m];
if (!CombineUsable(i, p.obs, p.corr)) continue;
const double sigma_corr = double(p.obs[i].sigma) * 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.obs[i].partiality;
if (p.obs[i].partiality > peak_partiality) {
peak_partiality = p.obs[i].partiality;
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.obs, p.corr)) continue;
const double corr = p.corr[i];
const double I_corr = pooled_I(i) * corr;
const double sigma_corr = double(p.obs[i].sigma) * 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);
const double capture = p.capture_uncertainty_coeff * (1.0 - Dmin(1.0, sum_partiality));
if (p.capture_uncertainty_coeff > 0.0) {
const double extra = capture * 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_capture[o] = float(capture);
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;
p.rr_noverloaded[r] = n_overloaded;
}
}
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);
}
// ===== the per-space-group grouping, built on the device from the raw-hkl runs =====
// The host's finiteness test (finite_ok in RotationScaleMerge::Ingest), on the same uploaded values.
__device__ __forceinline__ bool ObsFinite(int i, const ObsRecord *obs) {
const ObsRecord o = obs[i];
return isfinite(o.I) && isfinite(o.prescaling_corr) && o.prescaling_corr != 0.0f && isfinite(o.sigma)
&& o.sigma > 0.0f;
}
struct GroupParams {
int n_runs, n_groups;
const int32_t *perm, *rr_start, *rr_count, *rr_group;
const int32_t *sorted_run, *group_first; // group g = runs sorted_run[group_first[g] .. [g+1])
const ObsRecord *obs;
int32_t *group, *group_start, *group_count, *group_perm;
};
// Every observation's group: its raw-hkl run's, or -1 when the run has none or the observation
// failed the finiteness test. Runs own disjoint observations.
__global__ void StampGroupsKernel(GroupParams p) {
for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < p.n_runs; r += gridDim.x * blockDim.x) {
const int g = p.rr_group[r];
for (int q = p.rr_start[r]; q < p.rr_start[r] + p.rr_count[r]; ++q) {
const int i = p.perm[q];
p.group[i] = (g >= 0 && ObsFinite(i, p.obs)) ? g : -1;
}
}
}
__global__ void CountGroupsKernel(GroupParams p) {
for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
int c = 0;
for (int j = p.group_first[g]; j < p.group_first[g + 1]; ++j) {
const int r = p.sorted_run[j];
for (int q = p.rr_start[r]; q < p.rr_start[r] + p.rr_count[r]; ++q)
c += ObsFinite(p.perm[q], p.obs) ? 1 : 0;
}
p.group_count[g] = c;
}
}
// Each group gathers its observations and puts them in ascending index order (an in-place heap
// sort: a group is tens of observations, rarely thousands) - the order the host's stable counting
// sort over the observations produced, which is unique because the indices are distinct.
__global__ void FillGroupsKernel(GroupParams p) {
for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
int32_t *a = p.group_perm + p.group_start[g];
int n = 0;
for (int j = p.group_first[g]; j < p.group_first[g + 1]; ++j) {
const int r = p.sorted_run[j];
for (int q = p.rr_start[r]; q < p.rr_start[r] + p.rr_count[r]; ++q) {
const int i = p.perm[q];
if (ObsFinite(i, p.obs)) a[n++] = i;
}
}
auto sift = [a](int root, int end) {
while (2 * root + 1 < end) {
int child = 2 * root + 1;
if (child + 1 < end && a[child] < a[child + 1]) ++child;
if (a[root] >= a[child]) return;
const int32_t t = a[root]; a[root] = a[child]; a[child] = t;
root = child;
}
};
for (int k = n / 2 - 1; k >= 0; --k) sift(k, n);
for (int end = n - 1; end > 0; --end) {
const int32_t t = a[0]; a[0] = a[end]; a[end] = t;
sift(0, end);
}
}
}
// One field of a slice of the records, from a contiguous array: rec[i].field = v[i], the field being
// the record's float number `field`.
__global__ void SetRecordFieldKernel(int n, const float *__restrict__ v, ObsRecord *__restrict__ rec, int field) {
float *f = reinterpret_cast<float *>(rec);
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
f[size_t(i) * 4 + field] = v[i];
}
// The fulls' records for their Unity scaling: partiality and prescaling_corr 1.
__global__ void FullsRecordKernel(int n, const float *__restrict__ I, const float *__restrict__ sigma,
ObsRecord *__restrict__ obs) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
obs[i] = {I[i], sigma[i], 1.0f, 1.0f};
}
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;
}
// ===== correction-surface terms built on the device from the resident fulls =====
// ApplyCellSurface's host build, step for step: which fulls give a term, the terms in fulls order,
// the group CSR over them, each subset's per-(block, cell) segments and the shells' weights. Every
// order the host makes is a stable counting sort, which a stable radix sort of the indices
// reproduces; the only arithmetic (1/d^2, the shell weights) is rounded as the host rounds it.
// Per full: 1 if it gives a term (all, even frame, odd frame) - the host's `usable(o) && cell >= 0`.
// A full's partiality is 1 (the combine sets it so on both paths).
__global__ void SurfaceKeepKernel(int n, const int32_t *group, const float *corr, const uint8_t *on_ice,
const int32_t *frame, const int32_t *cell, double min_partiality,
int32_t *keep, int32_t *keep_even, int32_t *keep_odd) {
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
const float c = corr[i];
const bool k = group[i] >= 0 && c > 0.0f && isfinite(c) && 1.0f >= min_partiality && !on_ice[i]
&& cell[i] >= 0;
const bool odd = frame[i] & 1;
keep[i] = k; keep_even[i] = k && !odd; keep_odd[i] = k && odd;
}
// The term of each kept full at its place (the scans of the keep flags), its frame parity, and its
// place in the even / odd lists.
__global__ void SurfaceWriteTermKernel(int n, const int32_t *keep_at, const int32_t *even_at,
const int32_t *odd_at, const int32_t *last_flags,
const float *I, const float *sigma, const float *corr, const float *d,
const int32_t *cell, const int32_t *group, const int32_t *frame,
RotationScaleMergeGPU::SurfaceTerm *term, uint8_t *parity,
int32_t *idx_even, int32_t *idx_odd, uint32_t *group_key) {
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n) return;
// A full is kept when its exclusive-scan position steps up to the next one (the last full by
// its own flag).
const bool kept = i + 1 < n ? keep_at[i + 1] > keep_at[i] : last_flags[0] != 0;
if (!kept) return;
const int w = keep_at[i];
term[w] = {I[i], sigma[i], corr[i], d[i], cell[i], group[i]};
group_key[w] = static_cast<uint32_t>(group[i]);
const uint8_t par = static_cast<uint8_t>(frame[i] & 1);
parity[w] = par;
if (par) idx_odd[odd_at[i]] = w;
else idx_even[even_at[i]] = w;
}
__global__ void IotaKernel(int n, int32_t *v) {
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) v[i] = i;
}
// Counting into bins: integer adds, so the counts do not depend on their order.
__global__ void CountKeysKernel(int n, const uint32_t *key, int32_t *count) {
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i < n) atomicAdd(&count[key[i]], 1);
}
// The block-and-cell key of position k of a subset of n terms: block b holds positions
// [n b / nb, n (b + 1) / nb), as the host's ParallelFor over blocks cuts them.
__global__ void SurfaceSubsetKeyKernel(int n, int nb, int ncell, const int32_t *sel,
const RotationScaleMergeGPU::SurfaceTerm *term, uint32_t *key) {
const int k = blockIdx.x * blockDim.x + threadIdx.x;
if (k >= n) return;
auto lo = [n, nb](int b) { return static_cast<int64_t>(n) * b / nb; };
int b = static_cast<int>(static_cast<int64_t>(k) * nb / n);
while (b + 1 < nb && lo(b + 1) <= k) ++b;
while (b > 0 && lo(b) > k) --b;
key[k] = static_cast<uint32_t>(b * ncell + term[sel[k]].cell);
}
// 1/d^2 of each term, as the host's shells take it.
__global__ void SurfaceS2Kernel(int n, const RotationScaleMergeGPU::SurfaceTerm *term, float *s2) {
const int k = blockIdx.x * blockDim.x + threadIdx.x;
if (k >= n) return;
const float d = term[k].d;
s2[k] = d > 0.0f ? __fdiv_rn(1.0f, __fmul_rn(d, d)) : 0.0f;
}
// The shell of each term (how many edges are at or below its s^2 - the host's upper_bound) and
// the (cell, shell) bin its weight goes to.
__global__ void SurfaceShellKernel(int n, const float *s2, const float *edge, int n_edge, int nshell,
const RotationScaleMergeGPU::SurfaceTerm *term, uint8_t *shell,
uint32_t *bin) {
const int k = blockIdx.x * blockDim.x + threadIdx.x;
if (k >= n) return;
int sh = 0;
while (sh < n_edge && edge[sh] <= s2[k]) ++sh;
shell[k] = static_cast<uint8_t>(sh);
bin[k] = static_cast<uint32_t>(term[k].cell * nshell + sh);
}
// The shell of each group: that of its last term in fulls order (the host overwrites it term by term).
__global__ void SurfaceGroupShellKernel(int n_groups, const int32_t *gstart, const int32_t *gperm,
const uint8_t *shell, int32_t *group_shell) {
const int g = blockIdx.x * blockDim.x + threadIdx.x;
if (g >= n_groups) return;
group_shell[g] = gstart[g + 1] > gstart[g] ? shell[gperm[gstart[g + 1] - 1]] : 0;
}
// One thread per (cell, shell) bin: the statistical weight of its terms, in fulls order.
__global__ void SurfaceShellWeightKernel(int n_bins, const int32_t *bin_start, const int32_t *order,
const RotationScaleMergeGPU::SurfaceTerm *term, double *shw) {
const int b = blockIdx.x * blockDim.x + threadIdx.x;
if (b >= n_bins) return;
double w = 0.0;
for (int q = bin_start[b]; q < bin_start[b + 1]; ++q) {
const auto &t = term[order[q]];
const double sc = __dmul_rn(double(t.sigma), double(t.corr));
if (!(sc > 0.0)) continue;
w = __dadd_rn(w, 1.0 / __dmul_rn(sc, sc));
}
shw[b] = w;
}
// ===== 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;
double reject_up, reject_down; // the multiplicative part of the outlier band (OutlierBand.h)
int for_search, error_model_active, reject_outliers;
const float *I, *sigma, *corr, *partiality, *d, *reject_median, *var_bkg, *var_per_I, *capture;
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;
int32_t *sample_at; // per full: 1 where it gives a sample, then (scanned) where it goes
ErrorModelSample *sample; // the samples, compacted in fulls order
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;
}
}
// The host's counting_variance (RotationScaleMerge::MergeAndStats), to the bit: each product
// rounded on its own, as the host build rounds it, and summed in the host's order.
__device__ __forceinline__ double CountingVarianceExact(float corr, float var_bkg, float var_per_I, float capture,
double I_for_b, double own_s2) {
const double I_exp = 0.0 < I_for_b ? I_for_b : 0.0;
const double c = corr, cap = capture;
double base = __dmul_rn(__dmul_rn(c, c), double(var_bkg));
base = __dadd_rn(base, __dmul_rn(__dmul_rn(c, double(var_per_I)), I_exp));
base = __dadd_rn(base, __dmul_rn(__dmul_rn(__dmul_rn(cap, cap), I_exp), I_exp));
return base > 0.0 ? base : own_s2;
}
// The host's model_sigma, to the bit (see CountingVarianceExact).
__device__ __forceinline__ float ModelSigmaExact(int i, const MergeParams &p, double I_for_b, float sigma_corr) {
if (!p.error_model_active) return sigma_corr;
const double a_var = CountingVarianceExact(p.corr[i], p.var_bkg[i], p.var_per_I[i], p.capture[i], I_for_b,
__dmul_rn(double(sigma_corr), double(sigma_corr)));
const double bI = __dmul_rn(p.error_model_b, I_for_b);
const double v = __dadd_rn(__dmul_rn(p.error_model_a, a_var), __dmul_rn(bI, bI));
return v > 0.0 ? float(sqrt(v)) : sigma_corr;
}
// One thread per group: what the pooled outlier test's widening is measured from (the host's
// reject_var_add block) - the group's d (its first usable full's), the excess dI^2 - var(dI) of its
// two hands under the error model's sigmas (NaN where a hand is missing), and the d range of the
// fulls that entered. Summed over the group's fulls in index order, as the host sums them.
__global__ void BijvoetExcessKernel(MergeParams p, float *group_d, double *excess, float *d_lo, float *d_hi) {
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[2] = {0.0, 0.0}, swI[2] = {0.0, 0.0};
float gd = NAN, dmin = INFINITY, dmax = 0.0f;
bool first = true;
for (int q = lo; q < hi; ++q) {
const int i = p.gperm[q];
if (!MergeUsable(i, p)) continue;
if (first) { gd = p.d[i]; first = false; }
if (!p.has_hands[g]) continue;
const float I_corr = p.I[i] * p.corr[i];
const double I_for_b = isfinite(p.em_mean[g]) ? p.em_mean[g] : double(I_corr);
const float sc = ModelSigmaExact(i, p, I_for_b, p.sigma[i] * p.corr[i]);
if (!(sc > 0.0f) || !isfinite(sc)) continue;
const double w = 1.0 / __dmul_rn(double(sc), double(sc));
const int h = p.hand[i];
sw[h] = __dadd_rn(sw[h], w);
swI[h] = __dadd_rn(swI[h], __dmul_rn(w, double(I_corr)));
const float d = p.d[i];
if (isfinite(d) && d > 0.0f) { dmin = fminf(dmin, d); dmax = fmaxf(dmax, d); }
}
group_d[g] = gd; d_lo[g] = dmin; d_hi[g] = dmax;
excess[g] = NAN;
if (sw[0] > 0.0 && sw[1] > 0.0 && isfinite(gd)) {
const double di = swI[0] / sw[0] - swI[1] / sw[1];
excess[g] = __dmul_rn(di, di) - (1.0 / sw[0] + 1.0 / sw[1]);
}
}
}
// One thread per group: the sums the strong-reflection ISa asymptote is read from (MergeAndStats),
// over the group's usable fulls in index order, with the host's rounding.
__global__ void GroupScatterKernel(MergeParams p, double *sum, double *sum_sq, double *sum_var, int32_t *n) {
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 s = 0.0, s2 = 0.0, sv = 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 I_corr = __dmul_rn(double(p.I[i]), double(p.corr[i]));
const double sigma_corr = __dmul_rn(double(p.sigma[i]), double(p.corr[i]));
s = __dadd_rn(s, I_corr);
s2 = __dadd_rn(s2, __dmul_rn(I_corr, I_corr));
sv = __dadd_rn(sv, CountingVarianceExact(p.corr[i], p.var_bkg[i], p.var_per_I[i], p.capture[i],
p.em_mean[g], __dmul_rn(sigma_corr, sigma_corr)));
++cnt;
}
sum[g] = s; sum_sq[g] = s2; sum_var[g] = sv; n[g] = cnt;
}
}
// The host's ObservationKey (RotationScaleMerge.cpp): integer arithmetic, so the same on both.
__device__ __forceinline__ uint64_t ObservationKeyDevice(int32_t h, int32_t k, int32_t l, int32_t frame) {
uint64_t z = uint64_t(uint32_t(h)) * 0x9E3779B97F4A7C15ULL
^ uint64_t(uint32_t(k)) * 0xC2B2AE3D27D4EB4FULL
^ uint64_t(uint32_t(l)) * 0x165667B19E3779F9ULL
^ uint64_t(uint32_t(frame)) * 0x27D4EB2F165667C5ULL;
z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
return z ^ (z >> 31);
}
// One thread per full: its observation key, and whether the merge can use it.
__global__ void ObservationKeyKernel(int n_obs, MergeParams p, const int32_t *h, const int32_t *k,
const int32_t *l, uint64_t *key, uint8_t *usable) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
key[i] = ObservationKeyDevice(h[i], k[i], l[i], p.frame[i]);
usable[i] = MergeUsable(i, p) ? 1 : 0;
}
}
// One thread per group: the CC1/2 half of each usable full, the parity of its key's rank among the
// group's usable fulls (the host's AssignHalvesOfMembers). Unusable fulls keep half 0.
__global__ void HalvesKernel(MergeParams p, const uint64_t *key, const uint8_t *usable, uint8_t *half) {
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];
for (int a = lo; a < hi; ++a) {
const int i = p.gperm[a];
if (!usable[i]) continue;
int rank = 0;
for (int b = lo; b < hi; ++b) {
const int j = p.gperm[b];
if (usable[j] && key[j] < key[i]) ++rank;
}
half[i] = uint8_t(rank & 1);
}
}
}
// The leverage-corrected error-model sample of full i; false where the full gives none.
__device__ __forceinline__ bool EmSample(int i, const MergeParams &p, ErrorModelSample &out) {
if (!MergeUsable(i, p)) return false;
const int g = p.group[i];
if (p.cnt[g] < 2) return false;
// 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)) return false;
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) return false;
const double resid = double(p.I[i]) * p.corr[i] - mean;
out.s2 = CountingVarianceExact(p.corr[i], p.var_bkg[i], p.var_per_I[i], p.capture[i], p.em_mean[g], s2);
out.I2 = mean * mean;
out.dev2 = resid * resid / factor;
out.d = p.d[i];
return true;
}
// One thread per full: whether it gives an error-model sample (the scan turns it into where).
__global__ void MergeSampleFlagKernel(int n_obs, MergeParams p) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
ErrorModelSample s;
p.sample_at[i] = EmSample(i, p, s) ? 1 : 0;
}
}
// One thread per full: its sample, at its place among the samples in fulls order.
__global__ void MergeSampleWriteKernel(int n_obs, MergeParams p) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
ErrorModelSample s;
if (EmSample(i, p, s)) p.sample[p.sample_at[i]] = s;
}
}
// 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 I_exp = Dmax(0.0, I_for_b);
const double cap = p.capture[i];
const double base = c * c * p.var_bkg[i] + c * p.var_per_I[i] * I_exp + cap * cap * I_exp * I_exp;
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 I_b = isfinite(p.em_mean[g]) ? p.em_mean[g] : double(I_corr);
if (OutsideOutlierBand(I_corr, rmed, p.reject_nsigma, double(sigma_corr) * sigma_corr + v_add,
I_b, p.error_model_b * I_b, p.reject_up, p.reject_down)) {
++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;
}
}
// --- correction-surface fit (ApplyCellSurface) ---
// The host fit is the reference here: these kernels form the same sums from the same terms in the
// same order, and every rounding is spelled out (__dmul_rn / __dadd_rn round each step on its own,
// fma rounds once) to be the one the host build makes - nvcc would otherwise contract a multiply
// and an add wherever it sees them, and GCC at -march=x86-64-v3 does so only in some of them. The
// products are the host's (I * corr) * a and (sigma * corr) * a.
using SurfaceTerm = RotationScaleMergeGPU::SurfaceTerm;
__device__ __forceinline__ double SurfaceIs(const SurfaceTerm &t, double a) {
return __dmul_rn(__dmul_rn(double(t.I), double(t.corr)), a);
}
__device__ __forceinline__ double SurfaceSigma(const SurfaceTerm &t, double a) {
return __dmul_rn(__dmul_rn(double(t.sigma), double(t.corr)), a);
}
// One thread per ASU group: the group's inverse-variance sums over its terms of frame parity `parity`
// (< 0 = all), in fulls order - the host's `reference`. The host builds that loop twice, and only
// the parity-filtered copy fuses the multiply-add of swI; the unfiltered one rounds the product.
__global__ void SurfaceReferenceKernel(int n_groups, int parity, const int32_t *__restrict__ gperm,
const int32_t *__restrict__ gstart,
const SurfaceTerm *__restrict__ term,
const uint8_t *__restrict__ term_parity,
const double *__restrict__ A,
double *__restrict__ sw, double *__restrict__ swI) {
for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < n_groups; g += gridDim.x * blockDim.x) {
double s_w = 0.0, s_wI = 0.0;
for (int k = gstart[g]; k < gstart[g + 1]; ++k) {
const int i = gperm[k];
if (parity >= 0 && term_parity[i] != parity) continue;
const SurfaceTerm t = term[i];
const double a = A[t.cell];
const double Is = SurfaceIs(t, a), sc = SurfaceSigma(t, a);
const double w = 1.0 / __dmul_rn(sc, sc);
s_w = __dadd_rn(s_w, w);
s_wI = parity >= 0 ? fma(Is, w, s_wI) : __dadd_rn(s_wI, __dmul_rn(Is, w));
}
sw[g] = s_w; swI[g] = s_wI;
}
}
// The fit's sums of one round. Each term's contribution is formed on its own thread (one per term,
// in segment order), and only the two multiply-adds that sum them are left to the thread of each
// (reduction block, cell) - the per-block accumulators of the host fit, one slot at a time, walking
// its segment in term order. A term the host skips is marked by a NaN Iref (a kept one is finite).
__global__ void SurfaceFitTermKernel(int n, const int32_t *__restrict__ perm,
const SurfaceTerm *__restrict__ term,
const double *__restrict__ A,
const double *__restrict__ sw, const double *__restrict__ swI,
double *__restrict__ w_Is, double *__restrict__ w_Iref,
double *__restrict__ Iref_out) {
for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < n; k += gridDim.x * blockDim.x) {
const SurfaceTerm t = term[perm[k]];
Iref_out[k] = NAN;
if (sw[t.group] <= 0.0) continue;
const double Iref = swI[t.group] / sw[t.group], a = A[t.cell];
const double Is = SurfaceIs(t, a), sc = SurfaceSigma(t, a);
if (!isfinite(Iref) || Iref <= 0.0 || !(sc > 0.0)) continue;
const double w = 1.0 / __dmul_rn(sc, sc);
w_Is[k] = __dmul_rn(w, Is);
w_Iref[k] = __dmul_rn(w, Iref);
Iref_out[k] = Iref;
}
}
__global__ void SurfaceFitSegmentKernel(int n_seg, const int32_t *__restrict__ seg_start,
const double *__restrict__ w_Is, const double *__restrict__ w_Iref,
const double *__restrict__ Iref,
double *__restrict__ tcross, double *__restrict__ tref2) {
for (int s = blockIdx.x * blockDim.x + threadIdx.x; s < n_seg; s += gridDim.x * blockDim.x) {
double xcross = 0.0, xref2 = 0.0;
for (int k = seg_start[s]; k < seg_start[s + 1]; ++k) {
if (isnan(Iref[k])) continue;
xcross = fma(w_Is[k], Iref[k], xcross);
xref2 = fma(w_Iref[k], Iref[k], xref2);
}
tcross[s] = xcross; tref2[s] = xref2;
}
}
// One thread per cell: the blocks' sums added up in block order, as the host adds its slots.
__global__ void SurfaceFitCellKernel(int n_blocks, int ncell, const double *__restrict__ tcross,
const double *__restrict__ tref2,
double *__restrict__ cross, double *__restrict__ ref2) {
for (int c = blockIdx.x * blockDim.x + threadIdx.x; c < ncell; c += gridDim.x * blockDim.x) {
double sc = 0.0, sr = 0.0;
for (int b = 0; b < n_blocks; ++b) {
sc = __dadd_rn(sc, tcross[b * ncell + c]);
sr = __dadd_rn(sr, tref2[b * ncell + c]);
}
cross[c] = sc; ref2[c] = sr;
}
}
void CudaCheck(cudaError_t e, const char *what);
// A copy on the instance's stream, waited for - what cudaMemcpy on the NULL stream was, without also
// waiting for every other stream on the card.
void CopyAndWait(void *dst, const void *src, size_t bytes, cudaMemcpyKind kind, cudaStream_t s,
const char *what) {
CudaCheck(cudaMemcpyAsync(dst, src, bytes, kind, s), what);
CudaCheck(cudaStreamSynchronize(s), what);
}
void CudaCheck(cudaError_t e, const char *what) {
if (e != cudaSuccess)
throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
std::string("RotationScaleMergeGPU: ") + what + ": " + cudaGetErrorString(e));
}
// Device bytes per observation in the arrays SetPartialsLayout allocates: twelve floats (the record's
// I, sigma, partiality and prescaling_corr; zeta, corr, bkg, var_bkg, image_number, d, px, py), frame
// (4 bytes), on_ice and clipped (1 each). The merge's other buffers come on top of these.
constexpr size_t OBS_BYTES = 12 * 4 + 4 + 2 * 1;
// Every buffer here is sized by the data set, and the scaling does not fall back to the CPU
// mid-run: the two paths differ in the last bits, so the result would depend on the card. A
// buffer that does not fit ends the run, and says why and what to do instead.
[[noreturn]] void ThrowTooLargeForGPU(int64_t n_obs) {
size_t free_bytes = 0, total_bytes = 0;
cudaMemGetInfo(&free_bytes, &total_bytes);
cudaGetLastError();
char msg[768];
snprintf(msg, sizeof(msg),
"Scaling %lld partial observations needs more GPU memory than this card has: %.1f GB for "
"their per-observation arrays alone (%zu bytes each), before the merge's other buffers, "
"on a %.1f GB card. This data set is too large for GPU scaling on this card - run it on "
"the CPU (CUDA_VISIBLE_DEVICES= rugnux ..., or a CPU-only build) on a machine with plenty of "
"host memory: the CPU path keeps all of it in RAM, which for a set this large can be "
"tens of GB",
static_cast<long long>(n_obs), double(n_obs) * OBS_BYTES / 1e9, OBS_BYTES,
double(total_bytes) / 1e9);
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, msg);
}
// A buffer that did not fit although the per-observation arrays fit the card, with nothing else in
// this process running on the GPU any more: what holds the rest is this run or another program.
[[noreturn]] void ThrowOutOfGPUMemory(int64_t n_obs, size_t bytes) {
size_t free_bytes = 0, total_bytes = 0;
cudaMemGetInfo(&free_bytes, &total_bytes);
cudaGetLastError();
char msg[768];
snprintf(msg, sizeof(msg),
"Scaling %lld partial observations ran out of GPU memory: a buffer of %.1f MB did not fit, "
"with %.1f GB of the %.1f GB card free (the per-observation arrays take %.1f GB). Either "
"another program is using this GPU, or this data set is too large for GPU scaling on this "
"card - run it on the CPU (CUDA_VISIBLE_DEVICES= rugnux ..., or a CPU-only build) on a "
"machine with plenty of host memory",
static_cast<long long>(n_obs), double(bytes) / 1e6, double(free_bytes) / 1e9,
double(total_bytes) / 1e9, double(n_obs) * OBS_BYTES / 1e9);
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, msg);
}
// How long a merge waits for GPU work running beside it (GPUWorkBeside) to give its memory back.
// That work is a probe pass of seconds, so this long means it is stuck, not slow.
constexpr auto GPU_BUSY_TIMEOUT = std::chrono::minutes(10);
[[noreturn]] void ThrowGPUBusy(int64_t n_obs) {
char msg[512];
snprintf(msg, sizeof(msg),
"GPU busy: scaling %lld partial observations waited %lld minutes for GPU work running "
"beside it in this process to give back its memory, and it did not. The data set fits "
"this card; run it again, or on the CPU (CUDA_VISIBLE_DEVICES= rugnux ...)",
static_cast<long long>(n_obs),
static_cast<long long>(std::chrono::duration_cast<std::chrono::minutes>(GPU_BUSY_TIMEOUT).count()));
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, msg);
}
}
struct RotationScaleMergeGPU::Impl {
// First, so it goes last: the buffers below are freed before the stream their work ran on.
std::unique_ptr<CudaStream> stream;
cudaStream_t s() const { return stream->get(); }
int device = 0; // the GPU this instance's buffers live on
bool available = false;
int n_obs = 0, n_frames = 0, n_groups = 0;
// The arrays fit the card (SetPartialsLayout checks that first), but not necessarily beside other
// GPU work of this process: a probe pass started on a copy of the run holds several GB of engines
// for a few seconds while this merge runs. So a buffer that does not fit waits for that work to
// end and is asked for again - the same buffer, so what the merge computes does not depend on the
// wait. It is asked for again also when nothing is running beside it any more: work that ended
// between the failure and the check has just given its memory back.
//
// The failed allocation is handled here, so its error must not stay behind as this thread's last
// error: the next kernel launch's cudaGetLastError() would report it over work that went fine.
template <typename T>
CudaDevicePtr<T> Alloc(size_t n) const {
try {
return CudaDevicePtr<T>(n, ALLOC);
} catch (const JFJochException &) {
cuda_throw_if_context_lost(); // not memory: waiting would only hide the real error
cuda_clear_error();
}
if (!wait_for_gpu_work_beside(GPU_BUSY_TIMEOUT))
ThrowGPUBusy(n_obs);
try {
return CudaDevicePtr<T>(n, ALLOC);
} catch (const JFJochException &) {
ThrowOutOfGPUMemory(n_obs, n * sizeof(T));
}
}
template <typename T>
void Upload(CudaDevicePtr<T> &dst, const T *src, int n) const {
dst = Alloc<T>(std::max(1, n));
if (n > 0)
CopyAndWait(dst.get(), src, size_t(n) * sizeof(T), cudaMemcpyHostToDevice, s(), "upload");
}
// immutable per-obs
CudaDevicePtr<ObsRecord> obs;
void UploadRecordChunk(int field, int offset, int count, const float *v);
CudaDevicePtr<float> 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, info;
CudaDevicePtr<uint8_t> scaled;
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<int32_t> m_sample_at; // per full (n_fulls)
CudaDevicePtr<ErrorModelSample> m_sample; // the compacted error-model samples
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_noverloaded, 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, f_capture;
CudaDevicePtr<uint8_t> f_on_ice, f_clipped;
// scale-fulls (Unity model, kept resident): all-ones partiality/zeta and records with partiality and
// prescaling_corr 1, so the shared scaling kernels yield coeff=mean, plus the working corr 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_zeta;
CudaDevicePtr<ObsRecord> f_obs;
CudaDevicePtr<int32_t> f_frame_perm, f_frame_start, f_frame_count;
CudaDevicePtr<int32_t> f_gperm, f_gstart, f_gcount;
// correction-surface fit (one ApplyCellSurface call at a time): its terms and their group CSR, the
// three subsets' per-(block, cell) segments, the surface and the sums of the round
int s_ncell = 0, s_n_groups = 0;
int s_n_blocks[3] = {0, 0, 0};
int s_n_sel[3] = {0, 0, 0};
size_t s_slots = 0; // capacity of s_tcross / s_tref2
CudaDevicePtr<RotationScaleMergeGPU::SurfaceTerm> s_term;
CudaDevicePtr<uint8_t> s_parity;
CudaDevicePtr<int32_t> s_gperm, s_gstart;
CudaDevicePtr<int32_t> s_perm[3], s_seg_start[3];
CudaDevicePtr<double> s_A, s_sw, s_swI, s_tcross, s_tref2, s_cross, s_ref2;
CudaDevicePtr<double> s_w_Is, s_w_Iref, s_Iref; // per term of a subset, in segment order
};
// 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)
cudaSetDevice(prev); // not set_gpu: a destructor must not throw
}
};
} // namespace
RotationScaleMergeGPU::RotationScaleMergeGPU(int device) : impl_(std::make_unique<Impl>()) {
if (get_gpu_count() > 0) {
// One instance, one device, recorded so the guard below can put the caller's device back
// instead of leaving the thread moved.
impl_->device = device;
DeviceGuard guard(impl_->device, true);
impl_->stream = std::make_unique<CudaStream>();
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;
// Known now, before anything is allocated: arrays that could not fit even an empty card fail here,
// not after gigabytes of them have been.
size_t free_bytes = 0, total_bytes = 0;
CudaCheck(cudaMemGetInfo(&free_bytes, &total_bytes), "memory info");
if (size_t(std::max(0, n_obs)) * OBS_BYTES > total_bytes)
ThrowTooLargeForGPU(n_obs);
d.Upload(d.frame_start, frame_start, n_frames); d.Upload(d.frame_count, frame_count, n_frames);
const int n = std::max(1, n_obs);
d.obs = d.Alloc<ObsRecord>(n);
d.zeta = d.Alloc<float>(n); d.corr = d.Alloc<float>(n);
d.bkg = d.Alloc<float>(n); d.var_bkg = d.Alloc<float>(n);
d.image_number = d.Alloc<float>(n); d.d_obs = d.Alloc<float>(n);
d.px_obs = d.Alloc<float>(n); d.py_obs = d.Alloc<float>(n);
d.frame = d.Alloc<int32_t>(n);
d.on_ice = d.Alloc<uint8_t>(n);
d.clipped = d.Alloc<uint8_t>(n);
d.g = d.Alloc<double>(n_frames);
d.scaled = d.Alloc<uint8_t>(n_frames);
d.info = d.Alloc<double>(n_frames);
d.cc = d.Alloc<double>(std::max(1, n_frames));
d.cc_n = d.Alloc<int64_t>(std::max(1, n_frames));
}
namespace {
template <typename T>
void UploadChunk(CudaDevicePtr<T> &dst, int offset, int count, const T *v, cudaStream_t s) {
if (count > 0)
CopyAndWait(dst.get() + offset, v, size_t(count) * sizeof(T), cudaMemcpyHostToDevice, s,
"upload chunk");
}
}
// A slice of one field of the records (`field` its float number), staged contiguously and then put
// in place on the device: a strided copy from pageable memory goes element by element.
void RotationScaleMergeGPU::Impl::UploadRecordChunk(int field, int offset, int count, const float *v) {
if (count == 0)
return;
auto staged = Alloc<float>(count);
CopyAndWait(staged.get(), v, size_t(count) * sizeof(float), cudaMemcpyHostToDevice, s(), "upload record chunk");
SetRecordFieldKernel<<<std::min(65535, (count + BLK - 1) / BLK), BLK, 0, s()>>>(count, staged.get(),
obs.get() + offset, field);
CudaCheck(cudaGetLastError(), "record field launch");
CudaCheck(cudaStreamSynchronize(s()), "upload record 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: d.UploadRecordChunk(0, offset, count, v); break;
case ObsField::Sigma: d.UploadRecordChunk(1, offset, count, v); break;
case ObsField::Partiality: d.UploadRecordChunk(2, offset, count, v); break;
case ObsField::PrescalingCorr: d.UploadRecordChunk(3, offset, count, v); break;
case ObsField::Zeta: UploadChunk(d.zeta, offset, count, v, impl_->s()); break;
case ObsField::Corr0: UploadChunk(d.corr, offset, count, v, impl_->s()); break;
case ObsField::Bkg: UploadChunk(d.bkg, offset, count, v, impl_->s()); break;
case ObsField::VarBkg: UploadChunk(d.var_bkg, offset, count, v, impl_->s()); break;
case ObsField::ImageNumber: UploadChunk(d.image_number, offset, count, v, impl_->s()); break;
case ObsField::D: UploadChunk(d.d_obs, offset, count, v, impl_->s()); break;
case ObsField::Px: UploadChunk(d.px_obs, offset, count, v, impl_->s()); break;
case ObsField::Py: UploadChunk(d.py_obs, offset, count, v, impl_->s()); break;
}
}
void RotationScaleMergeGPU::CopyObservationsFrom(const RotationScaleMergeGPU &from) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const auto &f = *from.impl_;
const auto copy = [&](auto &dst, const auto &src) {
CudaCheck(cudaMemcpyPeerAsync(dst.get(), d.device, src.get(), f.device,
size_t(d.n_obs) * sizeof(*src.get()), d.s()), "copy observations");
};
copy(d.obs, f.obs); copy(d.zeta, f.zeta); copy(d.corr, f.corr);
copy(d.bkg, f.bkg); copy(d.var_bkg, f.var_bkg); copy(d.image_number, f.image_number);
copy(d.d_obs, f.d_obs); copy(d.px_obs, f.px_obs); copy(d.py_obs, f.py_obs);
copy(d.frame, f.frame); copy(d.on_ice, f.on_ice); copy(d.clipped, f.clipped);
CudaCheck(cudaStreamSynchronize(d.s()), "copy observations");
}
void RotationScaleMergeGPU::SetObsFrame(int offset, int count, const int32_t *frame) {
DeviceGuard guard(impl_->device, impl_->available);
UploadChunk(impl_->frame, offset, count, frame, impl_->s());
}
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, impl_->s());
}
void RotationScaleMergeGPU::SetObsClipped(int offset, int count, const uint8_t *clipped) {
DeviceGuard guard(impl_->device, impl_->available);
UploadChunk(impl_->clipped, offset, count, clipped, impl_->s());
}
void RotationScaleMergeGPU::BuildGroups(int n_groups, const int32_t *rawrun_group, int n_sorted,
const int32_t *sorted_run, const int32_t *group_first) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
d.n_groups = n_groups;
CopyAndWait(d.rr_group.get(), rawrun_group, size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyHostToDevice, impl_->s(), "upload rr_group");
CudaDevicePtr<int32_t> sorted, first;
d.Upload(sorted, sorted_run, n_sorted);
d.Upload(first, group_first, n_groups + 1);
if (d.group.get() == nullptr) d.group = d.Alloc<int32_t>(std::max(1, d.n_obs));
d.group_count = d.Alloc<int32_t>(std::max(1, n_groups));
d.group_start = d.Alloc<int32_t>(std::max(1, n_groups));
GroupParams p{};
p.n_runs = d.n_runs; p.n_groups = n_groups;
p.perm = d.perm.get(); p.rr_start = d.rr_start.get(); p.rr_count = d.rr_count.get(); p.rr_group = d.rr_group.get();
p.sorted_run = sorted.get(); p.group_first = first.get();
p.obs = d.obs.get();
p.group = d.group.get(); p.group_start = d.group_start.get(); p.group_count = d.group_count.get();
const int run_blocks = std::min(65535, (d.n_runs + BLK - 1) / BLK);
const int grp_blocks = std::min(65535, (n_groups + BLK - 1) / BLK);
StampGroupsKernel<<<std::max(1, run_blocks), BLK, 0, impl_->s()>>>(p);
CudaCheck(cudaGetLastError(), "StampGroupsKernel launch");
CountGroupsKernel<<<std::max(1, grp_blocks), BLK, 0, impl_->s()>>>(p);
CudaCheck(cudaGetLastError(), "CountGroupsKernel launch");
// The starts are a prefix over the counts: a few megabytes there and back, summed in order.
std::vector<int32_t> count(n_groups), start(n_groups);
if (n_groups > 0)
CopyAndWait(count.data(), d.group_count.get(), size_t(n_groups) * sizeof(int32_t),
cudaMemcpyDeviceToHost, impl_->s(), "download group counts");
int acc = 0;
for (int g = 0; g < n_groups; ++g) { start[g] = acc; acc += count[g]; }
if (n_groups > 0)
CopyAndWait(d.group_start.get(), start.data(), size_t(n_groups) * sizeof(int32_t),
cudaMemcpyHostToDevice, impl_->s(), "upload group starts");
d.group_perm = d.Alloc<int32_t>(std::max(1, acc));
p.group_perm = d.group_perm.get();
FillGroupsKernel<<<std::max(1, grp_blocks), BLK, 0, impl_->s()>>>(p);
CudaCheck(cudaGetLastError(), "FillGroupsKernel launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "groups sync");
d.group_mean = d.Alloc<double>(std::max(1, n_groups));
}
void RotationScaleMergeGPU::SetCorr(const float *corr) {
DeviceGuard guard(impl_->device, impl_->available);
CopyAndWait(impl_->corr.get(), corr, size_t(impl_->n_obs) * sizeof(float),
cudaMemcpyHostToDevice, impl_->s(), "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(cudaMemsetAsync(d.scaled.get(), 0, size_t(d.n_frames) * sizeof(uint8_t), impl_->s()), "memset scaled");
CudaCheck(cudaMemsetAsync(d.g.get(), 0, size_t(d.n_frames) * sizeof(double), impl_->s()), "memset g"); // unscaled g unused
const int upd_blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK);
const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK);
for (int it = 0; it < iters; ++it) {
ReduceGroupMeansKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(d.n_groups, min_partiality,
d.group_perm.get(), d.group_start.get(), d.group_count.get(),
d.obs.get(), d.corr.get(), d.group_mean.get());
CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch");
FitPerFrameGKernel<<<d.n_frames, BLK, 0, impl_->s()>>>(d.n_frames, d.frame_start.get(), d.frame_count.get(),
nullptr, min_partiality, d.group.get(), d.obs.get(), d.zeta.get(), d.on_ice.get(), d.group_mean.get(),
long(MIN_REFLECTIONS), d.g.get(), d.scaled.get(), d.info.get());
CudaCheck(cudaGetLastError(), "FitPerFrameGKernel launch");
UpdateCorrKernel<<<upd_blocks, BLK, 0, impl_->s()>>>(d.n_obs, d.frame.get(), d.obs.get(),
d.g.get(), d.scaled.get(), d.corr.get());
}
CudaCheck(cudaGetLastError(), "kernel launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "scale sync");
}
void RotationScaleMergeGPU::GetCorr(float *corr_out) const {
DeviceGuard guard(impl_->device, impl_->available);
CopyAndWait(corr_out, impl_->corr.get(), size_t(impl_->n_obs) * sizeof(float),
cudaMemcpyDeviceToHost, impl_->s(), "download corr");
}
void RotationScaleMergeGPU::GetInfo(double *info_out) const {
DeviceGuard guard(impl_->device, impl_->available);
CopyAndWait(info_out, impl_->info.get(), size_t(impl_->n_frames) * sizeof(double),
cudaMemcpyDeviceToHost, impl_->s(), "download info");
}
void RotationScaleMergeGPU::GetG(double *g_out, uint8_t *scaled_out) const {
DeviceGuard guard(impl_->device, impl_->available);
CopyAndWait(g_out, impl_->g.get(), size_t(impl_->n_frames) * sizeof(double),
cudaMemcpyDeviceToHost, impl_->s(), "download g");
CopyAndWait(scaled_out, impl_->scaled.get(), size_t(impl_->n_frames) * sizeof(uint8_t),
cudaMemcpyDeviceToHost, impl_->s(), "download scaled");
}
void RotationScaleMergeGPU::SetFrameCellOk(const uint8_t *frame_cell_ok) {
DeviceGuard guard(impl_->device, impl_->available);
impl_->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,
std::vector<ErrorModelSample> &samples) {
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 = d.Alloc<double>(std::max(1, ng)); d.m_swI = d.Alloc<double>(std::max(1, ng));
d.m_em_mean = d.Alloc<double>(std::max(1, ng)); d.m_cnt = d.Alloc<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 = d.Alloc<double>(std::max(1, 2 * ng)); d.m_swIh = d.Alloc<double>(std::max(1, 2 * ng));
d.m_cnth = d.Alloc<int32_t>(std::max(1, 2 * ng));
if (nf > 0) d.Upload(d.m_hand, hand, nf);
if (ng > 0) d.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_sample_at = d.Alloc<int32_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.var_bkg = d.f_var_bkg.get(); p.var_per_I = d.f_var_per_I.get(); p.capture = d.f_capture.get();
p.sample_at = d.m_sample_at.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, 0, impl_->s()>>>(p);
MergeSampleFlagKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(nf, p);
CudaCheck(cudaGetLastError(), "merge em/samples launch");
// Where each sample goes: the fulls that give one, in order, as the host's compaction lists them.
int n_samples = 0;
if (nf > 0) {
int32_t last_flag = 0;
CopyAndWait(&last_flag, d.m_sample_at.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost,
impl_->s(), "dl last sample flag");
size_t bytes = 0;
CudaCheck(cub::DeviceScan::ExclusiveSum(nullptr, bytes, d.m_sample_at.get(), d.m_sample_at.get(), nf,
impl_->s()), "sample scan size");
auto scratch = d.Alloc<uint8_t>(std::max<size_t>(bytes, 1));
CudaCheck(cub::DeviceScan::ExclusiveSum(scratch.get(), bytes, d.m_sample_at.get(), d.m_sample_at.get(), nf,
impl_->s()), "sample scan");
CopyAndWait(&n_samples, d.m_sample_at.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost,
impl_->s(), "dl sample count");
n_samples += last_flag;
}
d.m_sample = d.Alloc<ErrorModelSample>(std::max(1, n_samples));
p.sample = d.m_sample.get();
MergeSampleWriteKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(nf, p);
CudaCheck(cudaGetLastError(), "merge samples write launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "merge em/samples sync");
CopyAndWait(em_mean_out, d.m_em_mean.get(), size_t(ng) * sizeof(double),
cudaMemcpyDeviceToHost, impl_->s(), "dl em_mean");
CopyAndWait(cnt_out, d.m_cnt.get(), size_t(ng) * sizeof(int32_t),
cudaMemcpyDeviceToHost, impl_->s(), "dl cnt");
samples.resize(n_samples);
if (n_samples > 0)
CopyAndWait(samples.data(), d.m_sample.get(), size_t(n_samples) * sizeof(ErrorModelSample),
cudaMemcpyDeviceToHost, impl_->s(), "dl samples");
d.m_sample = CudaDevicePtr<ErrorModelSample>();
}
void RotationScaleMergeGPU::MergeAccum(double error_model_a, double error_model_b, bool error_model_active,
bool reject_outliers, double reject_nsigma, double reject_mult_sd,
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 = d.Alloc<double>(std::max(1, ng)); d.a_sw = d.Alloc<double>(std::max(1, ng));
d.a_swIh0 = d.Alloc<double>(std::max(1, ng)); d.a_swIh1 = d.Alloc<double>(std::max(1, ng));
d.a_swh0 = d.Alloc<double>(std::max(1, ng)); d.a_swh1 = d.Alloc<double>(std::max(1, ng));
d.a_swht0 = d.Alloc<double>(std::max(1, ng)); d.a_swht1 = d.Alloc<double>(std::max(1, ng));
d.a_nh0 = d.Alloc<int32_t>(std::max(1, ng)); d.a_nh1 = d.Alloc<int32_t>(std::max(1, ng));
d.a_d = d.Alloc<double>(std::max(1, ng)); d.a_rejected = d.Alloc<int32_t>(std::max(1, ng));
d.a_on_ice = d.Alloc<uint8_t>(std::max(1, ng));
d.Upload(d.m_rejected, rejected_obs, d.n_fulls);
d.Upload(d.reject_median, reject_median, ng);
if (reject_var_add && ng > 0) d.Upload(d.reject_var_add, reject_var_add, ng);
else d.reject_var_add = CudaDevicePtr<float>();
d.Upload(d.m_half, half, d.n_fulls);
d.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_up = std::expm1(reject_nsigma * reject_mult_sd);
p.reject_down = -std::expm1(-reject_nsigma * reject_mult_sd);
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.capture = d.f_capture.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, 0, impl_->s()>>>(p);
CudaCheck(cudaGetLastError(), "merge accum launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "merge accum sync");
CopyAndWait(rejected_obs, d.m_rejected.get(), size_t(d.n_fulls) * sizeof(uint8_t), cudaMemcpyDeviceToHost, impl_->s(),
"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) {
CopyAndWait(h, s.get() + g0, size_t(n) * sizeof(*s.get()), cudaMemcpyDeviceToHost, impl_->s(),
"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;
d.Upload(d.merged_I, merged_I, ng);
d.r_absdev = d.Alloc<double>(std::max(1, ng)); d.r_sumI = d.Alloc<double>(std::max(1, ng));
d.r_wabsdev = d.Alloc<double>(std::max(1, ng)); d.r_wsumI = d.Alloc<double>(std::max(1, ng));
d.r_sumv = d.Alloc<double>(std::max(1, ng)); d.r_sumv2 = d.Alloc<double>(std::max(1, ng));
d.r_n = d.Alloc<int32_t>(std::max(1, ng)); d.r_nusable = d.Alloc<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.capture = d.f_capture.get();
p.rejected_obs = d.m_rejected.get();
const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
MergeRmeasKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(p);
CudaCheck(cudaGetLastError(), "merge rmeas launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "merge rmeas sync");
CopyAndWait(absdev, d.r_absdev.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl absdev");
CopyAndWait(sumI, d.r_sumI.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sumI");
CopyAndWait(wabsdev, d.r_wabsdev.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl wabsdev");
CopyAndWait(wsumI, d.r_wsumI.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl wsumI");
CopyAndWait(sumv, d.r_sumv.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sumv");
CopyAndWait(sumv2, d.r_sumv2.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sumv2");
CopyAndWait(n, d.r_n.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(), "dl rn");
CopyAndWait(nusable, d.r_nusable.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(), "dl rnusable");
}
void RotationScaleMergeGPU::MergeBijvoetExcess(double error_model_a, double error_model_b, float *group_d,
double *excess, float &d_lo, float &d_hi) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int ng = d.n_groups;
auto g_d = d.Alloc<float>(std::max(1, ng)), g_lo = d.Alloc<float>(std::max(1, ng)),
g_hi = d.Alloc<float>(std::max(1, ng));
auto g_excess = d.Alloc<double>(std::max(1, ng));
MergeParams p{};
p.n_groups = ng; p.min_partiality = d.merge_min_part; 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.hand = d.m_hand.get(); p.has_hands = d.m_has_hands.get();
p.error_model_a = error_model_a; p.error_model_b = error_model_b; p.error_model_active = 1;
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.capture = d.f_capture.get();
const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
BijvoetExcessKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(p, g_d.get(), g_excess.get(), g_lo.get(), g_hi.get());
CudaCheck(cudaGetLastError(), "Bijvoet excess launch");
std::vector<float> lo(ng), hi(ng);
CopyAndWait(group_d, g_d.get(), size_t(ng) * sizeof(float), cudaMemcpyDeviceToHost, impl_->s(), "dl group d");
CopyAndWait(excess, g_excess.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(),
"dl excess");
CopyAndWait(lo.data(), g_lo.get(), size_t(ng) * sizeof(float), cudaMemcpyDeviceToHost, impl_->s(), "dl d lo");
CopyAndWait(hi.data(), g_hi.get(), size_t(ng) * sizeof(float), cudaMemcpyDeviceToHost, impl_->s(), "dl d hi");
d_lo = INFINITY; d_hi = 0.0f;
for (int g = 0; g < ng; ++g) { d_lo = std::min(d_lo, lo[g]); d_hi = std::max(d_hi, hi[g]); }
}
void RotationScaleMergeGPU::MergeGroupScatter(double *sum, double *sum_sq, double *sum_var, int32_t *n) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int ng = d.n_groups;
auto g_sum = d.Alloc<double>(std::max(1, ng)), g_sum_sq = d.Alloc<double>(std::max(1, ng)),
g_sum_var = d.Alloc<double>(std::max(1, ng));
auto g_n = d.Alloc<int32_t>(std::max(1, ng));
MergeParams p{};
p.n_groups = ng; p.min_partiality = d.merge_min_part; 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.var_bkg = d.f_var_bkg.get(); p.var_per_I = d.f_var_per_I.get();
p.capture = d.f_capture.get();
const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
GroupScatterKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(p, g_sum.get(), g_sum_sq.get(), g_sum_var.get(), g_n.get());
CudaCheck(cudaGetLastError(), "group scatter launch");
CopyAndWait(sum, g_sum.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sum");
CopyAndWait(sum_sq, g_sum_sq.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sum_sq");
CopyAndWait(sum_var, g_sum_var.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(),
"dl sum_var");
CopyAndWait(n, g_n.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(), "dl n");
}
void RotationScaleMergeGPU::MergeHalves(uint64_t *obs_key, uint8_t *half) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int ng = d.n_groups, nf = d.n_fulls;
if (nf == 0) return;
auto key = d.Alloc<uint64_t>(nf);
auto usable = d.Alloc<uint8_t>(nf), h = d.Alloc<uint8_t>(nf);
MergeParams p{};
p.n_groups = ng; p.min_partiality = d.merge_min_part; 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();
const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK);
const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
ObservationKeyKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(nf, p, d.f_h.get(), d.f_k.get(), d.f_l.get(), key.get(),
usable.get());
CudaCheck(cudaGetLastError(), "observation key launch");
CudaCheck(cudaMemsetAsync(h.get(), 0, size_t(nf), impl_->s()), "memset halves");
HalvesKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(p, key.get(), usable.get(), h.get());
CudaCheck(cudaGetLastError(), "halves launch");
CopyAndWait(obs_key, key.get(), size_t(nf) * sizeof(uint64_t), cudaMemcpyDeviceToHost, impl_->s(), "dl keys");
CopyAndWait(half, h.get(), size_t(nf), cudaMemcpyDeviceToHost, impl_->s(), "dl halves");
}
void RotationScaleMergeGPU::SmoothCorr(const uint8_t *apply, const double *ratio) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
d.Upload(d.smooth_apply, apply, d.n_frames);
d.Upload(d.smooth_ratio, ratio, d.n_frames);
const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK);
SmoothCorrKernel<<<blocks, BLK, 0, impl_->s()>>>(d.n_obs, d.frame.get(), d.smooth_apply.get(),
d.smooth_ratio.get(), d.corr.get());
CudaCheck(cudaGetLastError(), "smooth corr launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "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;
d.Upload(d.smooth_apply, apply, d.n_frames);
d.Upload(d.smooth_ratio, ratio, d.n_frames);
const int blocks = std::min(65535, (d.n_fulls + BLK - 1) / BLK);
SmoothCorrKernel<<<blocks, BLK, 0, impl_->s()>>>(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(cudaStreamSynchronize(impl_->s()), "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 = d.Alloc<unsigned long long>(1);
CudaCheck(cudaMemsetAsync(dropped.get(), 0, sizeof(unsigned long long), impl_->s()), "zero zeta drop count");
const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK);
FilterZetaKernel<<<blocks, BLK, 0, impl_->s()>>>(d.n_obs, min_zeta, d.zeta.get(), d.corr.get(), dropped.get());
CudaCheck(cudaGetLastError(), "zeta filter launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "zeta filter sync");
unsigned long long n = 0;
CopyAndWait(&n, dropped.get(), sizeof(unsigned long long), cudaMemcpyDeviceToHost, impl_->s(),
"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_;
d.Upload(d.filter_reject, reject, d.n_frames);
const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK);
FilterFrameKernel<<<blocks, BLK, 0, impl_->s()>>>(d.n_obs, d.frame.get(), d.filter_reject.get(), d.corr.get());
CudaCheck(cudaGetLastError(), "frame filter launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "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, 0, impl_->s()>>>(d.n_groups, min_partiality,
d.group_perm.get(), d.group_start.get(), d.group_count.get(),
d.obs.get(), d.corr.get(), d.group_mean.get());
CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch");
PerFrameCCKernel<<<d.n_frames, BLK, 0, impl_->s()>>>(d.n_frames, min_partiality,
d.frame_start.get(), d.frame_count.get(), d.obs.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(cudaStreamSynchronize(impl_->s()), "partial CC sync");
CopyAndWait(cc_out, d.cc.get(), size_t(d.n_frames) * sizeof(double),
cudaMemcpyDeviceToHost, impl_->s(), "download cc");
CopyAndWait(cc_n_out, d.cc_n.get(), size_t(d.n_frames) * sizeof(int64_t),
cudaMemcpyDeviceToHost, impl_->s(), "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;
d.Upload(d.perm, perm, n_perm);
d.Upload(d.rr_start, rr_start, n_runs);
d.Upload(d.rr_count, rr_count, n_runs);
d.Upload(d.rr_h, rr_h, n_runs);
d.Upload(d.rr_k, rr_k, n_runs);
d.Upload(d.rr_l, rr_l, n_runs);
d.rr_group = d.Alloc<int32_t>(std::max(1, n_runs));
d.rr_nevents = d.Alloc<int32_t>(std::max(1, n_runs));
d.rr_noverloaded = d.Alloc<int32_t>(std::max(1, n_runs));
d.rr_offset = d.Alloc<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, int64_t &n_overloaded) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
CopyAndWait(d.rr_group.get(), rawrun_group, size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyHostToDevice, impl_->s(), "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.obs = d.obs.get(); p.corr = d.corr.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();
p.rr_noverloaded = d.rr_noverloaded.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, 0, impl_->s()>>>(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);
CopyAndWait(nevents.data(), d.rr_nevents.get(), size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyDeviceToHost, impl_->s(), "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);
std::vector<int32_t> noverloaded(d.n_runs);
CopyAndWait(noverloaded.data(), d.rr_noverloaded.get(), size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyDeviceToHost, impl_->s(), "download noverloaded");
n_overloaded = 0;
for (int r = 0; r < d.n_runs; ++r) n_overloaded += noverloaded[r];
// Allocate the fulls SoA and emit.
const int nf = std::max(1, d.n_fulls);
d.f_h = d.Alloc<int32_t>(nf); d.f_k = d.Alloc<int32_t>(nf); d.f_l = d.Alloc<int32_t>(nf);
d.f_frame = d.Alloc<int32_t>(nf); d.f_group = d.Alloc<int32_t>(nf);
d.f_I = d.Alloc<float>(nf); d.f_sigma = d.Alloc<float>(nf);
d.f_d = d.Alloc<float>(nf); d.f_img = d.Alloc<float>(nf);
d.f_px = d.Alloc<float>(nf); d.f_py = d.Alloc<float>(nf);
d.f_var_bkg = d.Alloc<float>(nf); d.f_var_per_I = d.Alloc<float>(nf); d.f_capture = d.Alloc<float>(nf);
d.f_on_ice = d.Alloc<uint8_t>(nf); d.f_clipped = d.Alloc<uint8_t>(nf);
d.f_corr = d.Alloc<float>(nf); d.f_partiality = d.Alloc<float>(nf);
d.f_zeta = d.Alloc<float>(nf);
d.f_obs = d.Alloc<ObsRecord>(nf);
CopyAndWait(d.rr_offset.get(), offset.data(), size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyHostToDevice, impl_->s(), "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_capture = d.f_capture.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, 0, impl_->s()>>>(p);
CudaCheck(cudaGetLastError(), "combine emit launch");
}
CudaCheck(cudaStreamSynchronize(impl_->s()), "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) {
CopyAndWait(dst, src, bytes, cudaMemcpyDeviceToHost, impl_->s(), "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);
CopyAndWait(frame, d.f_frame.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "download f_frame");
CopyAndWait(group, d.f_group.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "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_;
d.Upload(d.f_frame_perm, frame_perm, n_perm);
d.Upload(d.f_frame_start, frame_start, d.n_frames);
d.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_;
d.Upload(d.f_gperm, gperm, n_gperm);
d.Upload(d.f_gstart, gstart, d.n_groups);
d.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, 0, impl_->s()>>>(d.f_corr.get(), nf, 1.0f);
CudaCheck(cudaGetLastError(), "FillKernel launch");
FillKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(d.f_partiality.get(), nf, 1.0f);
CudaCheck(cudaGetLastError(), "FillKernel launch");
FullsRecordKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(nf, d.f_I.get(), d.f_sigma.get(), d.f_obs.get());
CudaCheck(cudaGetLastError(), "FullsRecordKernel launch");
FillKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(d.f_zeta.get(), nf, 1.0f);
CudaCheck(cudaGetLastError(), "FillKernel launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "reset fulls scale sync");
}
void RotationScaleMergeGPU::FitFullsScale(double min_partiality) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int nf = d.n_fulls;
if (nf == 0) return;
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(cudaMemsetAsync(d.scaled.get(), 0, size_t(d.n_frames) * sizeof(uint8_t), impl_->s()), "memset f scaled");
CudaCheck(cudaMemsetAsync(d.g.get(), 0, size_t(d.n_frames) * sizeof(double), impl_->s()), "memset f g");
ReduceGroupMeansKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(d.n_groups, min_partiality,
d.f_gperm.get(), d.f_gstart.get(), d.f_gcount.get(),
d.f_obs.get(), d.f_corr.get(), d.group_mean.get());
CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch");
FitPerFrameGKernel<<<d.n_frames, BLK, 0, impl_->s()>>>(d.n_frames, d.f_frame_start.get(), d.f_frame_count.get(),
d.f_frame_perm.get(), min_partiality, d.f_group.get(), d.f_obs.get(), d.f_zeta.get(), d.f_on_ice.get(),
d.group_mean.get(), 1L, d.g.get(), d.scaled.get(), d.info.get());
CudaCheck(cudaGetLastError(), "scale fulls launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "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;
CopyAndWait(corr, d.f_corr.get(), size_t(d.n_fulls) * sizeof(float),
cudaMemcpyDeviceToHost, impl_->s(), "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);
CopyAndWait(px, d.f_px.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "download f_px");
CopyAndWait(py, d.f_py.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "download f_py");
}
void RotationScaleMergeGPU::GetFullsVariance(float *var_bkg, float *var_per_I, float *capture) 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);
CopyAndWait(var_bkg, d.f_var_bkg.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "download f_var_bkg");
CopyAndWait(var_per_I, d.f_var_per_I.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(),
"download f_var_per_I");
CopyAndWait(capture, d.f_capture.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "download f_capture");
}
void RotationScaleMergeGPU::SetFullsCorr(const float *corr) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
if (d.n_fulls == 0) return;
CopyAndWait(d.f_corr.get(), corr, size_t(d.n_fulls) * sizeof(float),
cudaMemcpyHostToDevice, impl_->s(), "upload f_corr");
}
void RotationScaleMergeGPU::SurfaceSetTerms(int n_terms, const SurfaceTerm *term, const uint8_t *parity,
int n_groups, const int32_t *gperm, const int32_t *gstart,
int ncell) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
d.s_ncell = ncell;
d.s_n_groups = n_groups;
d.Upload(d.s_term, term, n_terms);
d.Upload(d.s_parity, parity, n_terms);
d.Upload(d.s_gperm, gperm, n_terms);
d.Upload(d.s_gstart, gstart, n_groups + 1);
d.s_A = d.Alloc<double>(std::max(1, ncell));
d.s_cross = d.Alloc<double>(std::max(1, ncell));
d.s_ref2 = d.Alloc<double>(std::max(1, ncell));
d.s_sw = d.Alloc<double>(std::max(1, n_groups));
d.s_swI = d.Alloc<double>(std::max(1, n_groups));
d.s_w_Is = d.Alloc<double>(std::max(1, n_terms));
d.s_w_Iref = d.Alloc<double>(std::max(1, n_terms));
d.s_Iref = d.Alloc<double>(std::max(1, n_terms));
for (int &nb : d.s_n_blocks) nb = 0;
}
void RotationScaleMergeGPU::SurfaceSetSubset(int subset, int n_blocks, const int32_t *perm,
const int32_t *seg_start) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int n_seg = n_blocks * d.s_ncell;
d.s_n_sel[subset] = n_blocks > 0 ? seg_start[n_seg] : 0;
d.Upload(d.s_perm[subset], perm, d.s_n_sel[subset]);
d.Upload(d.s_seg_start[subset], seg_start, n_seg + 1);
d.s_n_blocks[subset] = n_blocks;
if (size_t(n_seg) > d.s_slots) {
d.s_tcross = d.Alloc<double>(n_seg);
d.s_tref2 = d.Alloc<double>(n_seg);
d.s_slots = n_seg;
}
}
namespace {
// The positions, in order, of a stable radix sort of `value` by `key` (keys below 2^bits).
void StableSortByKey(cudaStream_t s, int n, const uint32_t *key, const int32_t *value, uint32_t *key_out,
int32_t *value_out, int bits, const char *what) {
size_t bytes = 0;
CudaCheck(cub::DeviceRadixSort::SortPairs(nullptr, bytes, key, key_out, value, value_out, n, 0, bits, s), what);
CudaDevicePtr<uint8_t> scratch(std::max<size_t>(bytes, 1), ALLOC);
CudaCheck(cub::DeviceRadixSort::SortPairs(scratch.get(), bytes, key, key_out, value, value_out, n, 0, bits, s),
what);
}
// start[b] = how many keys are below b, for b in [0, n_bins] (a counting sort's bin starts).
void BinStarts(cudaStream_t s, int n, const uint32_t *key, int n_bins, int32_t *start, const char *what) {
CudaCheck(cudaMemsetAsync(start, 0, size_t(n_bins + 1) * sizeof(int32_t), s), what);
if (n > 0) CountKeysKernel<<<(n + BLK - 1) / BLK, BLK, 0, s>>>(n, key, start);
CudaCheck(cudaGetLastError(), what);
size_t bytes = 0;
CudaCheck(cub::DeviceScan::ExclusiveSum(nullptr, bytes, start, start, n_bins + 1, s), what);
CudaDevicePtr<uint8_t> scratch(std::max<size_t>(bytes, 1), ALLOC);
CudaCheck(cub::DeviceScan::ExclusiveSum(scratch.get(), bytes, start, start, n_bins + 1, s), what);
}
int BitsFor(int64_t n_keys) {
int bits = 1;
while ((int64_t(1) << bits) < n_keys) ++bits;
return bits;
}
}
int RotationScaleMergeGPU::SurfaceBuildTerms(const int32_t *cell, const float *corr, double min_partiality,
int ncell, int n_groups, int block) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const cudaStream_t st = impl_->s();
const int nf = d.n_fulls;
d.s_ncell = ncell;
d.s_n_groups = n_groups;
for (int &nb : d.s_n_blocks) nb = 0;
for (int &ns : d.s_n_sel) ns = 0;
// Which fulls give a term, and where each goes in the three lists.
int n_terms = 0, n_even = 0, n_odd = 0;
CudaDevicePtr<int32_t> idx_even, idx_odd;
CudaDevicePtr<uint32_t> group_key;
{
CudaDevicePtr<int32_t> d_cell, keep, even, odd;
CudaDevicePtr<float> d_corr;
d.Upload(d_cell, cell, nf);
d.Upload(d_corr, corr, nf);
keep = d.Alloc<int32_t>(std::max(1, nf)); even = d.Alloc<int32_t>(std::max(1, nf));
odd = d.Alloc<int32_t>(std::max(1, nf));
if (nf > 0) {
SurfaceKeepKernel<<<(nf + BLK - 1) / BLK, BLK, 0, st>>>(nf, d.f_group.get(), d_corr.get(), d.f_on_ice.get(),
d.f_frame.get(), d_cell.get(), min_partiality, keep.get(), even.get(), odd.get());
CudaCheck(cudaGetLastError(), "surface keep launch");
}
// The last full's flags, before the scans turn the flags into places.
int32_t last[3] = {0, 0, 0};
auto d_last = d.Alloc<int32_t>(3);
if (nf > 0) {
CopyAndWait(&last[0], keep.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl keep");
CopyAndWait(&last[1], even.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl keep");
CopyAndWait(&last[2], odd.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl keep");
size_t bytes = 0;
CudaCheck(cub::DeviceScan::ExclusiveSum(nullptr, bytes, keep.get(), keep.get(), nf, st), "surface scan");
auto scratch = d.Alloc<uint8_t>(std::max<size_t>(bytes, 1));
for (int32_t *v : {keep.get(), even.get(), odd.get()})
CudaCheck(cub::DeviceScan::ExclusiveSum(scratch.get(), bytes, v, v, nf, st), "surface scan");
int32_t at[3];
CopyAndWait(&at[0], keep.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl count");
CopyAndWait(&at[1], even.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl count");
CopyAndWait(&at[2], odd.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, st, "dl count");
n_terms = at[0] + last[0]; n_even = at[1] + last[1]; n_odd = at[2] + last[2];
}
CopyAndWait(d_last.get(), last, sizeof(last), cudaMemcpyHostToDevice, st, "ul last flags");
d.s_term = d.Alloc<SurfaceTerm>(std::max(1, n_terms));
d.s_parity = d.Alloc<uint8_t>(std::max(1, n_terms));
group_key = d.Alloc<uint32_t>(std::max(1, n_terms));
idx_even = d.Alloc<int32_t>(std::max(1, n_even));
idx_odd = d.Alloc<int32_t>(std::max(1, n_odd));
if (nf > 0) {
SurfaceWriteTermKernel<<<(nf + BLK - 1) / BLK, BLK, 0, st>>>(nf, keep.get(), even.get(), odd.get(),
d_last.get(), d.f_I.get(), d.f_sigma.get(), d_corr.get(), d.f_d.get(), d_cell.get(), d.f_group.get(),
d.f_frame.get(), d.s_term.get(), d.s_parity.get(), idx_even.get(), idx_odd.get(), group_key.get());
CudaCheck(cudaGetLastError(), "surface term launch");
}
CudaCheck(cudaStreamSynchronize(st), "surface terms sync");
}
// The group CSR over the terms: the terms of group g at gperm[gstart[g] .. gstart[g + 1]), in
// fulls order.
auto iota = d.Alloc<int32_t>(std::max(1, n_terms));
if (n_terms > 0) IotaKernel<<<(n_terms + BLK - 1) / BLK, BLK, 0, st>>>(n_terms, iota.get());
CudaCheck(cudaGetLastError(), "iota launch");
d.s_gperm = d.Alloc<int32_t>(std::max(1, n_terms));
d.s_gstart = d.Alloc<int32_t>(n_groups + 1);
{
auto key_sorted = d.Alloc<uint32_t>(std::max(1, n_terms));
if (n_terms > 0)
StableSortByKey(st, n_terms, group_key.get(), iota.get(), key_sorted.get(), d.s_gperm.get(),
BitsFor(n_groups), "surface group sort");
BinStarts(st, n_terms, group_key.get(), n_groups, d.s_gstart.get(), "surface group starts");
}
group_key = CudaDevicePtr<uint32_t>();
// Each subset cut into the host's reduction blocks, every block's terms ordered by cell.
const int32_t *sel_of[3] = {idx_even.get(), idx_odd.get(), iota.get()};
const int n_of[3] = {n_even, n_odd, n_terms};
for (int subset = 0; subset < 3; ++subset) {
const int n = n_of[subset];
const int nb = ReductionBlocks(n, block);
const int n_seg = nb * ncell;
d.s_n_sel[subset] = n;
d.s_n_blocks[subset] = nb;
d.s_perm[subset] = d.Alloc<int32_t>(std::max(1, n));
d.s_seg_start[subset] = d.Alloc<int32_t>(n_seg + 1);
if (nb == 0) continue;
auto key = d.Alloc<uint32_t>(n), key_sorted = d.Alloc<uint32_t>(n);
SurfaceSubsetKeyKernel<<<(n + BLK - 1) / BLK, BLK, 0, st>>>(n, nb, ncell, sel_of[subset], d.s_term.get(),
key.get());
CudaCheck(cudaGetLastError(), "surface subset key launch");
StableSortByKey(st, n, key.get(), sel_of[subset], key_sorted.get(), d.s_perm[subset].get(), BitsFor(n_seg),
"surface subset sort");
BinStarts(st, n, key.get(), n_seg, d.s_seg_start[subset].get(), "surface subset starts");
if (size_t(n_seg) > d.s_slots) {
d.s_tcross = d.Alloc<double>(n_seg);
d.s_tref2 = d.Alloc<double>(n_seg);
d.s_slots = n_seg;
}
}
CudaCheck(cudaStreamSynchronize(st), "surface subsets sync");
d.s_A = d.Alloc<double>(std::max(1, ncell));
d.s_cross = d.Alloc<double>(std::max(1, ncell));
d.s_ref2 = d.Alloc<double>(std::max(1, ncell));
d.s_sw = d.Alloc<double>(std::max(1, n_groups));
d.s_swI = d.Alloc<double>(std::max(1, n_groups));
d.s_w_Is = d.Alloc<double>(std::max(1, n_terms));
d.s_w_Iref = d.Alloc<double>(std::max(1, n_terms));
d.s_Iref = d.Alloc<double>(std::max(1, n_terms));
return n_terms;
}
void RotationScaleMergeGPU::SurfaceRelease() {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
d.s_slots = 0;
d.s_term = {}; d.s_parity = {}; d.s_gperm = {}; d.s_gstart = {};
for (int subset = 0; subset < 3; ++subset) {
d.s_perm[subset] = {}; d.s_seg_start[subset] = {};
d.s_n_blocks[subset] = d.s_n_sel[subset] = 0;
}
d.s_A = {}; d.s_sw = {}; d.s_swI = {}; d.s_tcross = {}; d.s_tref2 = {}; d.s_cross = {}; d.s_ref2 = {};
d.s_w_Is = {}; d.s_w_Iref = {}; d.s_Iref = {};
}
void RotationScaleMergeGPU::SurfaceShellWeights(int n_terms, int nshell, int32_t *group_shell, double *shw) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const cudaStream_t st = impl_->s();
const int ncell = d.s_ncell, n_groups = d.s_n_groups;
const int n = n_terms;
// The edges: order statistics of s^2, read off a sort.
auto s2 = d.Alloc<float>(std::max(1, n)), s2_sorted = d.Alloc<float>(std::max(1, n));
SurfaceS2Kernel<<<(n + BLK - 1) / BLK, BLK, 0, st>>>(n, d.s_term.get(), s2.get());
CudaCheck(cudaGetLastError(), "surface s2 launch");
{
size_t bytes = 0;
CudaCheck(cub::DeviceRadixSort::SortKeys(nullptr, bytes, s2.get(), s2_sorted.get(), n, 0, 32, st), "s2 sort");
auto scratch = d.Alloc<uint8_t>(std::max<size_t>(bytes, 1));
CudaCheck(cub::DeviceRadixSort::SortKeys(scratch.get(), bytes, s2.get(), s2_sorted.get(), n, 0, 32, st),
"s2 sort");
}
std::vector<float> edge(nshell - 1);
for (int i = 1; i < nshell; ++i)
CopyAndWait(&edge[i - 1], s2_sorted.get() + size_t(n) * size_t(i) / size_t(nshell), sizeof(float),
cudaMemcpyDeviceToHost, st, "dl edge");
s2_sorted = CudaDevicePtr<float>();
CudaDevicePtr<float> d_edge;
d.Upload(d_edge, edge.data(), nshell - 1);
auto shell = d.Alloc<uint8_t>(std::max(1, n));
auto bin = d.Alloc<uint32_t>(std::max(1, n)), bin_sorted = d.Alloc<uint32_t>(std::max(1, n));
SurfaceShellKernel<<<(n + BLK - 1) / BLK, BLK, 0, st>>>(n, s2.get(), d_edge.get(), nshell - 1, nshell,
d.s_term.get(), shell.get(), bin.get());
CudaCheck(cudaGetLastError(), "surface shell launch");
auto g_shell = d.Alloc<int32_t>(std::max(1, n_groups));
if (n_groups > 0)
SurfaceGroupShellKernel<<<(n_groups + BLK - 1) / BLK, BLK, 0, st>>>(n_groups, d.s_gstart.get(),
d.s_gperm.get(), shell.get(), g_shell.get());
CudaCheck(cudaGetLastError(), "surface group shell launch");
// The weights: the terms of each (cell, shell) bin in fulls order, summed one bin per thread.
const int n_bins = ncell * nshell;
auto iota = d.Alloc<int32_t>(std::max(1, n)), order = d.Alloc<int32_t>(std::max(1, n));
auto bin_start = d.Alloc<int32_t>(n_bins + 1);
auto d_shw = d.Alloc<double>(std::max(1, n_bins));
IotaKernel<<<(n + BLK - 1) / BLK, BLK, 0, st>>>(n, iota.get());
CudaCheck(cudaGetLastError(), "iota launch");
StableSortByKey(st, n, bin.get(), iota.get(), bin_sorted.get(), order.get(), BitsFor(n_bins), "surface bin sort");
BinStarts(st, n, bin.get(), n_bins, bin_start.get(), "surface bin starts");
SurfaceShellWeightKernel<<<(n_bins + BLK - 1) / BLK, BLK, 0, st>>>(n_bins, bin_start.get(), order.get(),
d.s_term.get(), d_shw.get());
CudaCheck(cudaGetLastError(), "surface shell weight launch");
CopyAndWait(shw, d_shw.get(), size_t(n_bins) * sizeof(double), cudaMemcpyDeviceToHost, st, "dl shw");
CopyAndWait(group_shell, g_shell.get(), size_t(n_groups) * sizeof(int32_t), cudaMemcpyDeviceToHost, st,
"dl group shell");
}
void RotationScaleMergeGPU::SurfaceReference(int parity, const double *A) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
CopyAndWait(d.s_A.get(), A, size_t(d.s_ncell) * sizeof(double), cudaMemcpyHostToDevice, impl_->s(),
"upload surface");
const int grp_blocks = std::min(65535, (d.s_n_groups + BLK - 1) / BLK);
if (grp_blocks > 0)
SurfaceReferenceKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(d.s_n_groups, parity, d.s_gperm.get(),
d.s_gstart.get(), d.s_term.get(), d.s_parity.get(), d.s_A.get(), d.s_sw.get(), d.s_swI.get());
CudaCheck(cudaGetLastError(), "surface reference launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "surface reference sync");
}
void RotationScaleMergeGPU::SurfaceGetReference(double *sw, double *swI) const {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const size_t bytes = size_t(d.s_n_groups) * sizeof(double);
CopyAndWait(sw, d.s_sw.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "dl surface sw");
CopyAndWait(swI, d.s_swI.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "dl surface swI");
}
void RotationScaleMergeGPU::SurfaceFitSums(int subset, double *cross, double *ref2) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int ncell = d.s_ncell, nb = d.s_n_blocks[subset], n_seg = nb * ncell;
if (nb == 0) {
std::fill(cross, cross + ncell, 0.0);
std::fill(ref2, ref2 + ncell, 0.0);
return;
}
SurfaceFitTermKernel<<<std::min(65535, (d.s_n_sel[subset] + BLK - 1) / BLK), BLK, 0, impl_->s()>>>(
d.s_n_sel[subset], d.s_perm[subset].get(), d.s_term.get(), d.s_A.get(), d.s_sw.get(), d.s_swI.get(),
d.s_w_Is.get(), d.s_w_Iref.get(), d.s_Iref.get());
CudaCheck(cudaGetLastError(), "surface fit term launch");
SurfaceFitSegmentKernel<<<std::min(65535, (n_seg + BLK - 1) / BLK), BLK, 0, impl_->s()>>>(n_seg,
d.s_seg_start[subset].get(), d.s_w_Is.get(), d.s_w_Iref.get(), d.s_Iref.get(),
d.s_tcross.get(), d.s_tref2.get());
CudaCheck(cudaGetLastError(), "surface fit segment launch");
SurfaceFitCellKernel<<<std::min(65535, (ncell + BLK - 1) / BLK), BLK, 0, impl_->s()>>>(nb, ncell,
d.s_tcross.get(), d.s_tref2.get(), d.s_cross.get(), d.s_ref2.get());
CudaCheck(cudaGetLastError(), "surface cell sum launch");
CopyAndWait(cross, d.s_cross.get(), size_t(ncell) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(),
"dl surface cross");
CopyAndWait(ref2, d.s_ref2.get(), size_t(ncell) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(),
"dl surface ref2");
}
Reference in New Issue View Git Blame Copy Permalink