mx/Jungfraujoch
3
2
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases 81 Wiki Activity
Files
76adf3f0fa66a6799174fd9bf0ed76cef9ff73f2
Jungfraujoch/image_analysis/scale_merge/RotationScaleMergeGPU.cu
T
leonarski_fandClaude Opus 5.5 bfd715b64d rugnux --no-gpu: run on the CPU only, without touching CUDA, on a machine with a GPU 
hide_gpus() in CUDAWrapper makes get_gpu_count() answer 0 without calling
cudaGetDeviceCount, so every GPU decision in the process (indexer Auto ->
FFTW, preprocessing/spot finding/decode, pre-scan, scaling engine, model
validation, MR) takes the path it takes on a machine without a GPU, and no
CUDA context is created. rugnux reads --no-gpu in main, ahead of the parser,
because set_gpu_blocking_sync() already asks CUDA for the devices. The
banner says 'CPU only (--no-gpu)'. The 'run it on the CPU' advice in the
out-of-GPU-memory messages and in the install docs now names --no-gpu
instead of CUDA_VISIBLE_DEVICES=.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
2026-10-14 10:13:23 +02:00

3171 lines
174 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 <type_traits>
#include <vector>
#include <cuda_runtime.h>
#include <math_constants.h>
#include <cub/device/device_radix_sort.cuh>
#include <cub/device/device_run_length_encode.cuh>
#include <cub/device/device_scan.cuh>
#include "../indexing/CUDAMemHelpers.h"
#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"
#include "../../common/ParallelFor.h"
#include "../../common/Reflection.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::SynchronousAfterWorkBeside;
__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 (rugnux --no-gpu ..., 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 (rugnux --no-gpu ..., 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 (rugnux --no-gpu ...)",
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.
//
// Plain Synchronous, not ALLOC: this function waits for the work beside itself, with its own limit.
template <typename T>
CudaDevicePtr<T> Alloc(size_t n) const {
try {
return CudaDevicePtr<T>(n, CudaAlloc::Synchronous);
} 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, CudaAlloc::Synchronous);
} 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);
// The per-observation arrays for n observations, the per-frame ones for n_frames and the frame CSR
// (SetPartialsLayout, the resident ingest and CopyObservationsFrom all lay them out the same way).
void AllocObservations(int n);
// (frame_start null: the CSR is left to the caller to fill.)
void AllocFrames(const int32_t *frame_start, const int32_t *frame_count);
void AllocRunScratch();
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
// the resident ingest (IngestUpload .. IngestFinish): each frame's first reflection
// (host, length n_frames + 1), and per observation its raw hkl as the sort key and its delta_phi as
// integrated, until the smoothing has used them
std::vector<int32_t> in_src_start;
CudaDevicePtr<uint64_t> in_key;
CudaDevicePtr<float> in_dphi;
RotationScaleMergeGPU::IngestLayout in_layout;
};
// 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.AllocFrames(frame_start, frame_count);
d.AllocObservations(n_obs);
}
void RotationScaleMergeGPU::Impl::AllocObservations(int n_in) {
const int n = std::max(1, n_in);
obs = Alloc<ObsRecord>(n);
zeta = Alloc<float>(n); corr = Alloc<float>(n);
bkg = Alloc<float>(n); var_bkg = Alloc<float>(n);
image_number = Alloc<float>(n); d_obs = Alloc<float>(n);
px_obs = Alloc<float>(n); py_obs = Alloc<float>(n);
frame = Alloc<int32_t>(n);
on_ice = Alloc<uint8_t>(n);
clipped = Alloc<uint8_t>(n);
}
void RotationScaleMergeGPU::Impl::AllocFrames(const int32_t *fstart, const int32_t *fcount) {
frame_start = Alloc<int32_t>(std::max(1, n_frames));
frame_count = Alloc<int32_t>(std::max(1, n_frames));
if (n_frames > 0 && fstart) {
CopyAndWait(frame_start.get(), fstart, size_t(n_frames) * sizeof(int32_t), cudaMemcpyHostToDevice, s(),
"frame start");
CopyAndWait(frame_count.get(), fcount, size_t(n_frames) * sizeof(int32_t), cudaMemcpyHostToDevice, s(),
"frame count");
}
g = Alloc<double>(std::max(1, n_frames));
scaled = Alloc<uint8_t>(std::max(1, n_frames));
info = Alloc<double>(std::max(1, n_frames));
cc = Alloc<double>(std::max(1, n_frames));
cc_n = Alloc<int64_t>(std::max(1, n_frames));
}
void RotationScaleMergeGPU::Impl::AllocRunScratch() {
rr_group = Alloc<int32_t>(std::max(1, n_runs));
rr_nevents = Alloc<int32_t>(std::max(1, n_runs));
rr_noverloaded = Alloc<int32_t>(std::max(1, n_runs));
rr_offset = Alloc<int32_t>(std::max(1, n_runs));
}
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_;
d.n_obs = f.n_obs;
d.n_frames = f.n_frames;
d.n_runs = f.n_runs;
d.n_perm = f.n_perm;
d.AllocFrames(nullptr, nullptr);
d.AllocObservations(d.n_obs);
const auto copy = [&](auto &dst, const auto &src, size_t n) {
CudaCheck(cudaMemcpyPeerAsync(dst.get(), d.device, src.get(), f.device,
n * sizeof(*src.get()), d.s()), "copy observations");
};
copy(d.frame_start, f.frame_start, d.n_frames); copy(d.frame_count, f.frame_count, d.n_frames);
const size_t n = d.n_obs;
copy(d.obs, f.obs, n); copy(d.zeta, f.zeta, n); copy(d.corr, f.corr, n);
copy(d.bkg, f.bkg, n); copy(d.var_bkg, f.var_bkg, n); copy(d.image_number, f.image_number, n);
copy(d.d_obs, f.d_obs, n); copy(d.px_obs, f.px_obs, n); copy(d.py_obs, f.py_obs, n);
copy(d.frame, f.frame, n); copy(d.on_ice, f.on_ice, n); copy(d.clipped, f.clipped, n);
const size_t nr = std::max(1, d.n_runs);
d.perm = d.Alloc<int32_t>(std::max(1, d.n_perm));
d.rr_start = d.Alloc<int32_t>(nr); d.rr_count = d.Alloc<int32_t>(nr);
d.rr_h = d.Alloc<int32_t>(nr); d.rr_k = d.Alloc<int32_t>(nr); d.rr_l = d.Alloc<int32_t>(nr);
copy(d.perm, f.perm, d.n_perm);
copy(d.rr_start, f.rr_start, d.n_runs); copy(d.rr_count, f.rr_count, d.n_runs);
copy(d.rr_h, f.rr_h, d.n_runs); copy(d.rr_k, f.rr_k, d.n_runs); copy(d.rr_l, f.rr_l, d.n_runs);
CudaCheck(cudaStreamSynchronize(d.s()), "copy observations");
d.AllocRunScratch();
}
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.AllocRunScratch();
}
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::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");
}
// --- the resident ingest ---
namespace {
constexpr int REFL_WORDS = sizeof(Reflection) / sizeof(uint32_t);
static_assert(sizeof(Reflection) % sizeof(uint32_t) == 0 && alignof(Reflection) == alignof(uint32_t),
"the ingest reads a reflection as whole 32-bit words");
// The raw hkl as one sort key: each index as the sixteen bits the host's SortKey holds, offset so
// that the unsigned order of the key is the signed order of (h, k, l).
__device__ __forceinline__ uint64_t RawHklKey(int32_t h, int32_t k, int32_t l) {
const auto u = [](int32_t v) { return uint64_t(uint16_t(int16_t(v)) ^ 0x8000u); };
return (u(h) << 32) | (u(k) << 16) | u(l);
}
__host__ __device__ __forceinline__ int32_t RawHklIndex(uint64_t key, int shift) {
return int16_t(uint16_t(((key >> shift) & 0xffffu) ^ 0x8000u));
}
// A float as an unsigned key in the order of the float comparison: -0 is +0, as the comparison
// sees it (a NaN never reaches here - the host comparator would not be an order with one).
__device__ __forceinline__ uint32_t FloatOrderKey(float v) {
if (v == 0.0f) v = 0.0f;
const uint32_t u = __float_as_uint(v);
return (u & 0x80000000u) ? ~u : (u | 0x80000000u);
}
struct IngestArrays {
ObsRecord *obs;
float *zeta, *corr, *bkg, *var_bkg, *image_number, *d, *px, *py, *dphi;
int32_t *frame;
uint8_t *on_ice, *clipped;
uint64_t *key;
};
// One block per BLK reflections: the block reads its reflections as whole words (coalesced) into
// shared memory and each thread takes its own apart. The deterministic correction is the product of
// its three factors, in the host's order; the frame's flux is divided out once it is known.
__global__ void IngestExtractKernel(int n, int base, const Reflection *__restrict__ refl, IngestArrays a) {
__shared__ uint32_t buf[BLK * REFL_WORDS];
const int first = blockIdx.x * BLK;
const int cnt = min(BLK, n - first);
const uint32_t *src = reinterpret_cast<const uint32_t *>(refl + first);
for (int w = threadIdx.x; w < cnt * REFL_WORDS; w += BLK)
buf[w] = src[w];
__syncthreads();
if (int(threadIdx.x) >= cnt)
return;
const Reflection &r = *reinterpret_cast<const Reflection *>(buf + threadIdx.x * REFL_WORDS);
const int i = base + first + threadIdx.x;
a.obs[i] = ObsRecord{r.I, r.sigma, r.partiality,
__fmul_rn(__fmul_rn(r.prescaling_corr, r.qe_corr), r.flight_corr)};
a.zeta[i] = r.zeta;
a.corr[i] = r.image_scale_corr;
a.bkg[i] = r.bkg;
a.var_bkg[i] = r.var_bkg;
a.image_number[i] = r.image_number;
a.d[i] = r.d;
a.px[i] = r.predicted_x;
a.py[i] = r.predicted_y;
a.dphi[i] = r.delta_phi_deg;
a.on_ice[i] = r.on_ice_ring ? 1 : 0;
a.clipped[i] = (r.clipped ? 1 : 0) | (r.overloaded ? 2 : 0); // OBS_CLIPPED | OBS_OVERLOADED
a.key[i] = RawHklKey(r.h, r.k, r.l);
}
__global__ void IngestFrameKernel(int n_frames, const int32_t *__restrict__ start, int32_t *__restrict__ frame) {
for (int o = blockIdx.x; o < n_frames; o += gridDim.x)
for (int i = start[o] + threadIdx.x; i < start[o + 1]; i += blockDim.x)
frame[i] = o;
}
// One thread per frame, its reflections in order: the host's sum to the bit.
__global__ void IngestMeanBkgKernel(int n_frames, const int32_t *__restrict__ start,
const float *__restrict__ bkg, double *__restrict__ mean) {
for (int o = blockIdx.x * blockDim.x + threadIdx.x; o < n_frames; o += gridDim.x * blockDim.x) {
double sum = 0.0;
int n = 0;
for (int i = start[o]; i < start[o + 1]; ++i) {
const float b = bkg[i];
if (isfinite(b) && b > 0.0f) { sum = __dadd_rn(sum, double(b)); ++n; }
}
mean[o] = n > 0 ? __ddiv_rn(sum, double(n)) : CUDART_NAN;
}
}
__global__ void IngestFluxKernel(int n, const int32_t *__restrict__ frame, const double *__restrict__ flux,
ObsRecord *__restrict__ obs) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
obs[i].prescaling_corr = __double2float_rn(__ddiv_rn(double(obs[i].prescaling_corr), flux[frame[i]]));
}
__global__ void IngestImageKeyKernel(int n, const float *__restrict__ image_number,
uint32_t *__restrict__ key, int32_t *__restrict__ idx) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {
key[i] = FloatOrderKey(image_number[i]);
idx[i] = i;
}
}
__global__ void IngestGatherKeyKernel(int n, const int32_t *__restrict__ idx, const uint64_t *__restrict__ key,
uint64_t *__restrict__ out) {
for (int j = blockIdx.x * blockDim.x + threadIdx.x; j < n; j += gridDim.x * blockDim.x)
out[j] = key[idx[j]];
}
// Per run: its resolution (that of its first member with a usable one, as the host's run scan takes
// it) and whether the resolution range keeps it.
__global__ void IngestRunKernel(int n_runs, const int32_t *__restrict__ start, const int32_t *__restrict__ count,
const int32_t *__restrict__ perm, const float *__restrict__ d_obs,
bool has_d_min, double d_min, bool has_d_max, double d_max,
float *__restrict__ run_d, int32_t *__restrict__ keep_run) {
for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x) {
float d = CUDART_NAN_F;
for (int p = start[r]; p < start[r] + count[r]; ++p) {
const float dp = d_obs[perm[p]];
if (isfinite(dp) && dp > 0.0f) { d = dp; break; }
}
run_d[r] = d;
keep_run[r] = (isfinite(d) && d > 0.0f && !(has_d_min && d < d_min) && !(has_d_max && d > d_max)) ? 1 : 0;
}
}
__global__ void IngestKeepKernel(int n_runs, const int32_t *__restrict__ start, const int32_t *__restrict__ count,
const int32_t *__restrict__ perm, const int32_t *__restrict__ keep_run,
int32_t *__restrict__ keep) {
for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x)
if (keep_run[r])
for (int p = start[r]; p < start[r] + count[r]; ++p)
keep[perm[p]] = 1; // distinct runs own disjoint observations
}
__global__ void IngestFrameCountKernel(int n_frames, const int32_t *__restrict__ start,
const int32_t *__restrict__ keep, int32_t *__restrict__ count) {
for (int o = blockIdx.x * blockDim.x + threadIdx.x; o < n_frames; o += gridDim.x * blockDim.x) {
int c = 0;
for (int i = start[o]; i < start[o + 1]; ++i) c += keep[i];
count[o] = c;
}
}
// The surviving runs, in order, each owning its sorted members renumbered to the observations they
// became; their raw hkl unpacked from the key.
struct RunCut {
int n_runs;
const int32_t *start, *count, *keep_run, *run_idx, *run_at, *perm, *new_idx;
const uint64_t *key;
const float *d;
int32_t *out_start, *out_count, *out_h, *out_k, *out_l, *out_perm;
float *out_d;
};
__global__ void IngestCutRunsKernel(RunCut c) {
for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < c.n_runs; r += gridDim.x * blockDim.x) {
if (!c.keep_run[r]) continue;
const int rr = c.run_idx[r], at = c.run_at[r];
c.out_start[rr] = at;
c.out_count[rr] = c.count[r];
c.out_h[rr] = RawHklIndex(c.key[r], 32);
c.out_k[rr] = RawHklIndex(c.key[r], 16);
c.out_l[rr] = RawHklIndex(c.key[r], 0);
c.out_d[rr] = c.d[r];
for (int p = 0; p < c.count[r]; ++p)
c.out_perm[at + p] = c.new_idx[c.perm[c.start[r] + p]];
}
}
__global__ void IngestFillKernel(int n, int32_t *__restrict__ v) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
v[i] = 1;
}
__global__ void IngestKeptCountKernel(int n_runs, const int32_t *__restrict__ keep_run,
const int32_t *__restrict__ count, int32_t *__restrict__ kept) {
for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x)
kept[r] = keep_run[r] ? count[r] : 0;
}
template <typename T>
__global__ void IngestCompactKernel(int n, const int32_t *__restrict__ keep, const int32_t *__restrict__ new_idx,
const T *__restrict__ in, T *__restrict__ out) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
if (keep[i]) out[new_idx[i]] = in[i];
}
int Blocks(int n) { return std::max(1, std::min(65535, (n + BLK - 1) / BLK)); }
// --- the fulls' CSRs on the device (BuildFullsCSRs) ---
// A full's group as a sort key: the fulls in no group sort after every group.
__global__ void FullGroupKeyKernel(int n, const int32_t *__restrict__ group, int32_t *__restrict__ key) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
key[i] = group[i] >= 0 ? group[i] : INT32_MAX;
}
// Counts per bin (integers, so in any order); a negative bin counts nowhere.
__global__ void CountBinsKernel(int n, const int32_t *__restrict__ bin, int32_t *__restrict__ count) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
if (bin[i] >= 0) atomicAdd(&count[bin[i]], 1);
}
// --- the ingest's geometry smoothing on the card (RotationScaleMerge::IngestOnDevice) ---
// Explicit roundings throughout, so that no compiler contracts a product into a sum differently on
// another card: the same bits on every GPU, and within the last bits of the host's own loops.
struct Vec3 { float x, y, z; };
__device__ __forceinline__ Vec3 VAdd(Vec3 a, Vec3 b) { return {__fadd_rn(a.x, b.x), __fadd_rn(a.y, b.y), __fadd_rn(a.z, b.z)}; }
__device__ __forceinline__ Vec3 VScale(Vec3 a, float s) { return {__fmul_rn(a.x, s), __fmul_rn(a.y, s), __fmul_rn(a.z, s)}; }
__device__ __forceinline__ float VDot(Vec3 a, Vec3 b) {
return __fadd_rn(__fadd_rn(__fmul_rn(a.x, b.x), __fmul_rn(a.y, b.y)), __fmul_rn(a.z, b.z));
}
struct DeltaPhiGeometry {
const float *recip; // per frame: a*, b*, c* (9 floats)
const uint8_t *ok; // per frame
Vec3 S0, m1, m2, m3;
float m2_S0, m3_S0;
};
// RotationScaleMerge::SmoothedDeltaPhi / SolvePhi.
__global__ void IngestDeltaPhiKernel(int n, const int32_t *__restrict__ frame, const uint64_t *__restrict__ key,
DeltaPhiGeometry g, float *__restrict__ dphi,
unsigned long long *__restrict__ n_resolved) {
constexpr float PI_F = 3.14159265358979323846f;
unsigned long long resolved = 0;
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {
const int o = frame[i];
if (!g.ok[o]) continue;
const float *rc = g.recip + 9 * o;
const Vec3 a{rc[0], rc[1], rc[2]}, b{rc[3], rc[4], rc[5]}, c{rc[6], rc[7], rc[8]};
const float h = float(RawHklIndex(key[i], 32)), k = float(RawHklIndex(key[i], 16)),
l = float(RawHklIndex(key[i], 0));
const Vec3 p0 = VAdd(VAdd(VScale(a, h), VScale(b, k)), VScale(c, l));
const float phi_near = __fdiv_rn(__fmul_rn(dphi[i], PI_F), 180.0f);
const float p0_sq = VDot(p0, p0);
if (!(p0_sq > 0.0f)) continue;
const float p0_m1 = VDot(p0, g.m1), p0_m2 = VDot(p0, g.m2), p0_m3 = VDot(p0, g.m3);
const float rho_sq = __fsub_rn(p0_sq, __fmul_rn(p0_m2, p0_m2));
const float p_m3 = __fdiv_rn(__fsub_rn(__fdiv_rn(-p0_sq, 2.0f), __fmul_rn(p0_m2, g.m2_S0)), g.m3_S0);
if (!(rho_sq > __fmul_rn(p_m3, p_m3))) continue;
const float root = __fsqrt_rn(__fsub_rn(rho_sq, __fmul_rn(p_m3, p_m3)));
float best = 0.0f;
bool have = false;
for (const float p_m1 : {root, -root}) {
const float cosphi = __fdiv_rn(__fadd_rn(__fmul_rn(p_m1, p0_m1), __fmul_rn(p_m3, p0_m3)), rho_sq);
const float sinphi = __fdiv_rn(__fsub_rn(__fmul_rn(p_m1, p0_m3), __fmul_rn(p_m3, p0_m1)), rho_sq);
const float phi = -atan2f(sinphi, cosphi);
if (!have || fabsf(__fsub_rn(phi, phi_near)) < fabsf(__fsub_rn(best, phi_near))) { best = phi; have = true; }
}
if (!isfinite(best)) continue;
dphi[i] = __fdiv_rn(__fmul_rn(best, 180.0f), PI_F);
++resolved;
}
atomicAdd(n_resolved, resolved);
}
// One exact-Bragg angle per rocking event (RotationScaleMerge::SmoothMosaicityAndPartiality), one
// thread per raw-hkl run.
__global__ void IngestRelayKernel(int n_runs, const int32_t *__restrict__ start, const int32_t *__restrict__ count,
const int32_t *__restrict__ perm, const float *__restrict__ img,
float *__restrict__ dphi, double increment_deg, float max_frame_gap,
unsigned long long *__restrict__ n_events,
unsigned long long *__restrict__ n_partials) {
unsigned long long ev = 0, pa = 0;
for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x) {
const int lo = start[r], hi = lo + count[r];
for (int i = lo; i < hi; ) {
int j = i + 1;
while (j < hi && __fsub_rn(img[perm[j]], img[perm[j - 1]]) <= max_frame_gap)
++j;
if (j - i >= 2) {
const float f0 = img[perm[i]];
double sum = 0.0;
int m_n = 0;
for (int m = i; m < j; ++m) {
const int p = perm[m];
if (!isfinite(dphi[p])) continue;
sum = __dadd_rn(sum, __dadd_rn(double(dphi[p]),
__dmul_rn(increment_deg, double(__fsub_rn(img[p], f0)))));
++m_n;
}
if (m_n >= 2) {
const double c = __ddiv_rn(sum, double(m_n));
for (int m = i; m < j; ++m) {
const int p = perm[m];
dphi[p] = __double2float_rn(__dsub_rn(c, __dmul_rn(increment_deg,
double(__fsub_rn(img[p], f0)))));
}
++ev;
pa += j - i;
}
}
i = j;
}
}
atomicAdd(n_events, ev);
atomicAdd(n_partials, pa);
}
// The frames each rocking event spans over the observations usable as integrated
// (RotationScaleMerge::RockingEventFramesOver), counted into a histogram: integers, so the order the
// runs are counted in does not matter.
__global__ void IngestRockingKernel(int n_runs, const int32_t *__restrict__ start, const int32_t *__restrict__ count,
const int32_t *__restrict__ perm, const float *__restrict__ img,
const float *__restrict__ corr, const ObsRecord *__restrict__ obs,
float max_frame_gap, int n_frames, int32_t *__restrict__ hist) {
for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x) {
bool open = false;
float first = 0.0f, last = 0.0f;
for (int q = start[r]; q < start[r] + count[r]; ++q) {
const int i = perm[q];
const float c = corr[i];
const ObsRecord o = obs[i];
if (!(isfinite(c) && c > 0.0f && isfinite(o.I) && isfinite(o.sigma) && o.sigma > 0.0f))
continue;
const float im = img[i];
if (open && __fsub_rn(im, last) <= max_frame_gap) {
last = im;
continue;
}
if (open)
atomicAdd(&hist[min(int(lroundf(last) - lroundf(first)) + 1, n_frames + 1)], 1);
first = last = im;
open = true;
}
if (open)
atomicAdd(&hist[min(int(lroundf(last) - lroundf(first)) + 1, n_frames + 1)], 1);
}
}
// The partiality recomputed from the frame's smoothed mosaicity, broadened by the energy bandwidth
// (RotationScaleMerge::SmoothMosaicityAndPartiality), and corr brought to it.
struct PartialityParams {
const float *mos; // per frame
double wedge, bandwidth_sigma, half_wavelength;
};
__global__ void IngestPartialityKernel(int n, const int32_t *__restrict__ frame, const float *__restrict__ zeta,
const float *__restrict__ d, const float *__restrict__ dphi,
PartialityParams q, ObsRecord *__restrict__ obs, float *__restrict__ corr) {
constexpr double PI_D = 3.14159265358979323846;
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {
const float z = zeta[i];
float mos = q.mos[frame[i]];
if (q.bandwidth_sigma > 0.0 && d[i] > 0.0f) {
const double sin_theta = __ddiv_rn(q.half_wavelength, double(d[i]));
const double sigma_bw = __ddiv_rn(__dmul_rn(__dmul_rn(__ddiv_rn(180.0, PI_D), q.bandwidth_sigma), sin_theta),
__dsqrt_rn(__dsub_rn(1.0, __dmul_rn(sin_theta, sin_theta))));
mos = __double2float_rn(__dsqrt_rn(__dadd_rn(double(__fmul_rn(mos, mos)), __dmul_rn(sigma_bw, sigma_bw))));
}
if (!(isfinite(mos) && mos > 1e-6f && isfinite(z) && z > 0.0f && isfinite(dphi[i])))
continue;
const double half_wedge = __ddiv_rn(q.wedge, 2.0);
const double c1 = __ddiv_rn(double(z), __dsqrt_rn(2.0));
const double arg_plus = __ddiv_rn(__dmul_rn(__dadd_rn(double(dphi[i]), half_wedge), c1), double(mos));
const double arg_minus = __ddiv_rn(__dmul_rn(__dsub_rn(double(dphi[i]), half_wedge), c1), double(mos));
const float p_new = __double2float_rn(__ddiv_rn(__dsub_rn(erf(arg_plus), erf(arg_minus)), 2.0));
const float p_old = obs[i].partiality;
if (p_old > 0.0f && p_new > 0.0f)
corr[i] = __fmul_rn(corr[i], __fdiv_rn(p_old, p_new));
obs[i].partiality = p_new;
}
}
}
void RotationScaleMergeGPU::IngestUpload(int n_frames, const Reflection *const *refl, const int32_t *count,
size_t nthreads) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
d.n_frames = n_frames;
d.in_src_start.assign(n_frames + 1, 0);
int max_count = 0;
for (int o = 0; o < n_frames; ++o) {
d.in_src_start[o + 1] = d.in_src_start[o] + count[o];
max_count = std::max(max_count, count[o]);
}
const int n = d.in_src_start[n_frames];
d.n_obs = n;
size_t free_bytes = 0, total_bytes = 0;
CudaCheck(cudaMemGetInfo(&free_bytes, &total_bytes), "memory info");
if (size_t(n) * OBS_BYTES > total_bytes)
ThrowTooLargeForGPU(n);
d.AllocObservations(n);
d.in_key = d.Alloc<uint64_t>(std::max(1, n));
d.in_dphi = d.Alloc<float>(std::max(1, n));
CudaDevicePtr<int32_t> src_start;
d.Upload(src_start, d.in_src_start.data(), n_frames + 1);
IngestFrameKernel<<<std::max(1, std::min(n_frames, 65535)), BLK, 0, d.s()>>>(n_frames, src_start.get(),
d.frame.get());
CudaCheck(cudaGetLastError(), "IngestFrameKernel launch");
// Through two pinned buffers of whole frames, each filled on all threads while the other is on
// its way to the card and taken apart there.
constexpr size_t STAGE_BYTES = 64 << 20;
const int cap = std::max(static_cast<int>(STAGE_BYTES / sizeof(Reflection)), max_count);
CudaHostPtr<Reflection> stage[2] = {CudaHostPtr<Reflection>(cap), CudaHostPtr<Reflection>(cap)};
CudaDevicePtr<Reflection> staged[2] = {d.Alloc<Reflection>(cap), d.Alloc<Reflection>(cap)};
CudaEvent copied[2];
bool in_flight[2] = {false, false};
IngestArrays a{d.obs.get(), d.zeta.get(), d.corr.get(), d.bkg.get(), d.var_bkg.get(), d.image_number.get(),
d.d_obs.get(), d.px_obs.get(), d.py_obs.get(), d.in_dphi.get(), d.frame.get(),
d.on_ice.get(), d.clipped.get(), d.in_key.get()};
int b = 0;
for (int f0 = 0; f0 < n_frames; ) {
int f1 = f0 + 1;
while (f1 < n_frames && d.in_src_start[f1 + 1] - d.in_src_start[f0] <= cap)
++f1;
const int base = d.in_src_start[f0];
const int cnt = d.in_src_start[f1] - base;
if (in_flight[b])
CudaCheck(cudaEventSynchronize(copied[b]), "staging buffer");
Reflection *dst = stage[b].get();
ParallelFor(f1 - f0, nthreads, [&](int fo) {
const int o = f0 + fo;
std::copy(refl[o], refl[o] + count[o], dst + (d.in_src_start[o] - base));
});
if (cnt > 0) {
CudaCheck(cudaMemcpyAsync(staged[b].get(), dst, size_t(cnt) * sizeof(Reflection),
cudaMemcpyHostToDevice, d.s()), "upload reflections");
CudaCheck(cudaEventRecord(copied[b], d.s()), "upload reflections");
in_flight[b] = true;
IngestExtractKernel<<<(cnt + BLK - 1) / BLK, BLK, 0, d.s()>>>(cnt, base, staged[b].get(), a);
CudaCheck(cudaGetLastError(), "IngestExtractKernel launch");
}
b ^= 1;
f0 = f1;
}
CudaCheck(cudaStreamSynchronize(d.s()), "upload reflections");
}
void RotationScaleMergeGPU::IngestMeanBackground(double *mean_bkg) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
CudaDevicePtr<int32_t> start;
d.Upload(start, d.in_src_start.data(), d.n_frames + 1);
CudaDevicePtr<double> mean = d.Alloc<double>(std::max(1, d.n_frames));
IngestMeanBkgKernel<<<Blocks(d.n_frames), BLK, 0, d.s()>>>(d.n_frames, start.get(), d.bkg.get(), mean.get());
CudaCheck(cudaGetLastError(), "IngestMeanBkgKernel launch");
if (d.n_frames > 0)
CopyAndWait(mean_bkg, mean.get(), size_t(d.n_frames) * sizeof(double), cudaMemcpyDeviceToHost, d.s(),
"download mean background");
}
int RotationScaleMergeGPU::IngestSort(const double *frame_flux, std::optional<double> d_min,
std::optional<double> d_max) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int n = d.n_obs;
const int na = std::max(1, n);
{
CudaDevicePtr<double> flux;
d.Upload(flux, frame_flux, d.n_frames);
IngestFluxKernel<<<Blocks(n), BLK, 0, d.s()>>>(n, d.frame.get(), flux.get(), d.obs.get());
CudaCheck(cudaGetLastError(), "IngestFluxKernel launch");
CudaCheck(cudaStreamSynchronize(d.s()), "flux");
}
// (h, k, l, image_number, index) is a total order, so two stable radix passes - by the frame position
// with the index as the value, then by the raw hkl - give the one permutation std::sort gives on the
// host, whatever either algorithm does on the way.
const auto sort_pairs = [&](auto *key_in, auto *key_out, const int32_t *val_in, int32_t *val_out, int end_bit) {
size_t bytes = 0;
CudaCheck(cub::DeviceRadixSort::SortPairs(nullptr, bytes, key_in, key_out, val_in, val_out, n, 0, end_bit,
d.s()), "sort temp size");
CudaDevicePtr<uint8_t> temp = d.Alloc<uint8_t>(std::max<size_t>(1, bytes));
CudaCheck(cub::DeviceRadixSort::SortPairs(temp.get(), bytes, key_in, key_out, val_in, val_out, n, 0, end_bit,
d.s()), "sort");
CudaCheck(cudaStreamSynchronize(d.s()), "sort");
};
CudaDevicePtr<int32_t> by_image = d.Alloc<int32_t>(na);
{
CudaDevicePtr<uint32_t> ikey = d.Alloc<uint32_t>(na), ikey_sorted = d.Alloc<uint32_t>(na);
CudaDevicePtr<int32_t> idx = d.Alloc<int32_t>(na);
IngestImageKeyKernel<<<Blocks(n), BLK, 0, d.s()>>>(n, d.image_number.get(), ikey.get(), idx.get());
CudaCheck(cudaGetLastError(), "IngestImageKeyKernel launch");
if (n > 0) sort_pairs(ikey.get(), ikey_sorted.get(), idx.get(), by_image.get(), 32);
}
CudaDevicePtr<int32_t> perm = d.Alloc<int32_t>(na);
CudaDevicePtr<uint64_t> key_sorted = d.Alloc<uint64_t>(na);
{
CudaDevicePtr<uint64_t> key = d.Alloc<uint64_t>(na);
IngestGatherKeyKernel<<<Blocks(n), BLK, 0, d.s()>>>(n, by_image.get(), d.in_key.get(), key.get());
CudaCheck(cudaGetLastError(), "IngestGatherKeyKernel launch");
if (n > 0) sort_pairs(key.get(), key_sorted.get(), by_image.get(), perm.get(), 48);
}
by_image = CudaDevicePtr<int32_t>();
// The raw-hkl runs of the sorted order.
int n_runs = 0;
CudaDevicePtr<uint64_t> run_key = d.Alloc<uint64_t>(na);
CudaDevicePtr<int32_t> run_count = d.Alloc<int32_t>(na);
if (n > 0) {
CudaDevicePtr<int32_t> n_runs_dev = d.Alloc<int32_t>(1);
size_t bytes = 0;
CudaCheck(cub::DeviceRunLengthEncode::Encode(nullptr, bytes, key_sorted.get(), run_key.get(),
run_count.get(), n_runs_dev.get(), n, d.s()), "runs temp size");
CudaDevicePtr<uint8_t> temp = d.Alloc<uint8_t>(std::max<size_t>(1, bytes));
CudaCheck(cub::DeviceRunLengthEncode::Encode(temp.get(), bytes, key_sorted.get(), run_key.get(),
run_count.get(), n_runs_dev.get(), n, d.s()), "runs");
CopyAndWait(&n_runs, n_runs_dev.get(), sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(), "run count");
}
key_sorted = CudaDevicePtr<uint64_t>();
const int nra = std::max(1, n_runs);
{
// Cut down to the runs there are: the encode had to be given room for one per reflection.
CudaDevicePtr<uint64_t> key_exact = d.Alloc<uint64_t>(nra);
CudaDevicePtr<int32_t> count_exact = d.Alloc<int32_t>(nra);
if (n_runs > 0) {
CopyAndWait(key_exact.get(), run_key.get(), size_t(n_runs) * sizeof(uint64_t), cudaMemcpyDeviceToDevice,
d.s(), "runs");
CopyAndWait(count_exact.get(), run_count.get(), size_t(n_runs) * sizeof(int32_t), cudaMemcpyDeviceToDevice,
d.s(), "runs");
}
run_key = std::move(key_exact);
run_count = std::move(count_exact);
}
const auto exclusive_sum = [&](const int32_t *in, int32_t *out, int count) {
if (count == 0) return;
size_t bytes = 0;
CudaCheck(cub::DeviceScan::ExclusiveSum(nullptr, bytes, in, out, count, d.s()), "scan temp size");
CudaDevicePtr<uint8_t> temp = d.Alloc<uint8_t>(std::max<size_t>(1, bytes));
CudaCheck(cub::DeviceScan::ExclusiveSum(temp.get(), bytes, in, out, count, d.s()), "scan");
CudaCheck(cudaStreamSynchronize(d.s()), "scan");
};
CudaDevicePtr<int32_t> run_start = d.Alloc<int32_t>(nra);
exclusive_sum(run_count.get(), run_start.get(), n_runs);
CudaDevicePtr<float> run_d = d.Alloc<float>(nra);
CudaDevicePtr<int32_t> keep_run = d.Alloc<int32_t>(nra);
const bool limited = d_min.has_value() || d_max.has_value();
IngestRunKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(n_runs, run_start.get(), run_count.get(), perm.get(),
d.d_obs.get(), d_min.has_value(), d_min.value_or(0.0),
d_max.has_value(), d_max.value_or(0.0),
run_d.get(), keep_run.get());
CudaCheck(cudaGetLastError(), "IngestRunKernel launch");
// Which observations survive (all of them without a range, as the host keeps them), and where each
// run and each observation lands.
CudaDevicePtr<int32_t> keep = d.Alloc<int32_t>(na);
if (limited) {
CudaCheck(cudaMemsetAsync(keep.get(), 0, size_t(na) * sizeof(int32_t), d.s()), "memset keep");
IngestKeepKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(n_runs, run_start.get(), run_count.get(), perm.get(),
keep_run.get(), keep.get());
CudaCheck(cudaGetLastError(), "IngestKeepKernel launch");
} else {
IngestFillKernel<<<Blocks(n), BLK, 0, d.s()>>>(n, keep.get());
IngestFillKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(n_runs, keep_run.get());
CudaCheck(cudaGetLastError(), "IngestFillKernel launch");
}
CudaDevicePtr<int32_t> new_idx = d.Alloc<int32_t>(na);
exclusive_sum(keep.get(), new_idx.get(), n);
CudaDevicePtr<int32_t> run_idx = d.Alloc<int32_t>(nra), run_at = d.Alloc<int32_t>(nra);
exclusive_sum(keep_run.get(), run_idx.get(), n_runs);
{
CudaDevicePtr<int32_t> kept_count = d.Alloc<int32_t>(nra);
IngestKeptCountKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(n_runs, keep_run.get(), run_count.get(),
kept_count.get());
CudaCheck(cudaGetLastError(), "IngestKeptCountKernel launch");
exclusive_sum(kept_count.get(), run_at.get(), n_runs);
}
int32_t last_keep = 0, last_new_idx = 0, last_keep_run = 0, last_run_idx = 0;
if (n > 0) {
CopyAndWait(&last_keep, keep.get() + n - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(), "count");
CopyAndWait(&last_new_idx, new_idx.get() + n - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(), "count");
}
if (n_runs > 0) {
CopyAndWait(&last_keep_run, keep_run.get() + n_runs - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(),
"count");
CopyAndWait(&last_run_idx, run_idx.get() + n_runs - 1, sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(),
"count");
}
const int n_keep = last_new_idx + last_keep;
const int n_keep_run = last_run_idx + last_keep_run;
d.n_runs = n_keep_run;
d.n_perm = n_keep;
const int nka = std::max(1, n_keep), nkra = std::max(1, n_keep_run);
d.perm = d.Alloc<int32_t>(nka);
d.rr_start = d.Alloc<int32_t>(nkra); d.rr_count = d.Alloc<int32_t>(nkra);
d.rr_h = d.Alloc<int32_t>(nkra); d.rr_k = d.Alloc<int32_t>(nkra); d.rr_l = d.Alloc<int32_t>(nkra);
CudaDevicePtr<float> rr_d = d.Alloc<float>(nkra);
RunCut c{n_runs, run_start.get(), run_count.get(), keep_run.get(), run_idx.get(), run_at.get(), perm.get(),
new_idx.get(), run_key.get(), run_d.get(),
d.rr_start.get(), d.rr_count.get(), d.rr_h.get(), d.rr_k.get(), d.rr_l.get(), d.perm.get(), rr_d.get()};
IngestCutRunsKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(c);
CudaCheck(cudaGetLastError(), "IngestCutRunsKernel launch");
// The frame CSR of what survived.
auto &L = d.in_layout;
L = RotationScaleMergeGPU::IngestLayout{};
L.n_reflections = n;
L.n_runs_all = n_runs;
L.frame_count.assign(d.n_frames, 0);
{
CudaDevicePtr<int32_t> start, fcount = d.Alloc<int32_t>(std::max(1, d.n_frames));
d.Upload(start, d.in_src_start.data(), d.n_frames + 1);
IngestFrameCountKernel<<<Blocks(d.n_frames), BLK, 0, d.s()>>>(d.n_frames, start.get(), keep.get(),
fcount.get());
CudaCheck(cudaGetLastError(), "IngestFrameCountKernel launch");
if (d.n_frames > 0)
CopyAndWait(L.frame_count.data(), fcount.get(), size_t(d.n_frames) * sizeof(int32_t),
cudaMemcpyDeviceToHost, d.s(), "download frame count");
}
L.frame_start.assign(d.n_frames, 0);
for (int o = 0, at = 0; o < d.n_frames; ++o) { L.frame_start[o] = at; at += L.frame_count[o]; }
// The observations themselves, cut field by field, so that only one field is ever held twice.
if (n_keep < n) {
const auto cut = [&](auto &field) {
using T = std::remove_reference_t<decltype(*field.get())>;
CudaDevicePtr<T> out = d.Alloc<T>(nka);
IngestCompactKernel<T><<<Blocks(n), BLK, 0, d.s()>>>(n, keep.get(), new_idx.get(), field.get(), out.get());
CudaCheck(cudaGetLastError(), "IngestCompactKernel launch");
CudaCheck(cudaStreamSynchronize(d.s()), "cut");
field = std::move(out);
};
cut(d.obs); cut(d.zeta); cut(d.corr); cut(d.bkg); cut(d.var_bkg); cut(d.image_number); cut(d.d_obs);
cut(d.px_obs); cut(d.py_obs); cut(d.frame); cut(d.on_ice); cut(d.clipped); cut(d.in_key); cut(d.in_dphi);
}
d.n_obs = n_keep;
d.AllocFrames(L.frame_start.data(), L.frame_count.data());
d.AllocRunScratch();
L.run_start.resize(n_keep_run); L.run_count.resize(n_keep_run);
L.run_h.resize(n_keep_run); L.run_k.resize(n_keep_run); L.run_l.resize(n_keep_run); L.run_d.resize(n_keep_run);
const auto get = [&](auto &v, const auto &src) {
if (!v.empty())
CopyAndWait(v.data(), src.get(), v.size() * sizeof(v[0]), cudaMemcpyDeviceToHost, d.s(), "download runs");
};
get(L.run_start, d.rr_start); get(L.run_count, d.rr_count);
get(L.run_h, d.rr_h); get(L.run_k, d.rr_k); get(L.run_l, d.rr_l); get(L.run_d, rr_d);
return n_keep;
}
void RotationScaleMergeGPU::IngestGetLayout(IngestLayout &out) const {
out = impl_->in_layout;
}
int64_t RotationScaleMergeGPU::IngestDeltaPhi(const DeltaPhiFrame &g) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
CudaDevicePtr<float> recip;
CudaDevicePtr<uint8_t> ok;
d.Upload(recip, g.recip.data(), static_cast<int>(g.recip.size()));
d.Upload(ok, g.ok.data(), static_cast<int>(g.ok.size()));
CudaDevicePtr<unsigned long long> n_resolved = d.Alloc<unsigned long long>(1);
CudaCheck(cudaMemsetAsync(n_resolved.get(), 0, sizeof(unsigned long long), d.s()), "memset");
const auto v = [](const float *a) { return Vec3{a[0], a[1], a[2]}; };
const DeltaPhiGeometry dg{recip.get(), ok.get(), v(g.S0), v(g.m1), v(g.m2), v(g.m3), g.m2_S0, g.m3_S0};
IngestDeltaPhiKernel<<<Blocks(d.n_obs), BLK, 0, d.s()>>>(d.n_obs, d.frame.get(), d.in_key.get(), dg,
d.in_dphi.get(), n_resolved.get());
CudaCheck(cudaGetLastError(), "IngestDeltaPhiKernel launch");
unsigned long long out = 0;
CopyAndWait(&out, n_resolved.get(), sizeof(out), cudaMemcpyDeviceToHost, d.s(), "resolved count");
return static_cast<int64_t>(out);
}
void RotationScaleMergeGPU::IngestRelayEvents(double increment_deg, float max_frame_gap, int64_t &n_events,
int64_t &n_partials) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
CudaDevicePtr<unsigned long long> counts = d.Alloc<unsigned long long>(2);
CudaCheck(cudaMemsetAsync(counts.get(), 0, 2 * sizeof(unsigned long long), d.s()), "memset");
IngestRelayKernel<<<Blocks(d.n_runs), BLK, 0, d.s()>>>(d.n_runs, d.rr_start.get(), d.rr_count.get(), d.perm.get(),
d.image_number.get(), d.in_dphi.get(), increment_deg,
max_frame_gap, counts.get(), counts.get() + 1);
CudaCheck(cudaGetLastError(), "IngestRelayKernel launch");
unsigned long long out[2] = {0, 0};
CopyAndWait(out, counts.get(), sizeof(out), cudaMemcpyDeviceToHost, d.s(), "relay counts");
n_events = static_cast<int64_t>(out[0]);
n_partials = static_cast<int64_t>(out[1]);
}
void RotationScaleMergeGPU::IngestRockingHistogram(float max_frame_gap, int32_t *hist) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int n_bins = d.n_frames + 2;
CudaDevicePtr<int32_t> h = d.Alloc<int32_t>(n_bins);
CudaCheck(cudaMemsetAsync(h.get(), 0, size_t(n_bins) * sizeof(int32_t), d.s()), "memset");
IngestRockingKernel<<<Blocks(d.n_runs), BLK, 0, d.s()>>>(d.n_runs, d.rr_start.get(), d.rr_count.get(), d.perm.get(),
d.image_number.get(), d.corr.get(), d.obs.get(),
max_frame_gap, d.n_frames, h.get());
CudaCheck(cudaGetLastError(), "IngestRockingKernel launch");
CopyAndWait(hist, h.get(), size_t(n_bins) * sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(), "rocking histogram");
}
void RotationScaleMergeGPU::IngestPartiality(const float *mosaicity, double wedge_deg, double bandwidth_sigma,
double half_wavelength_A) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
CudaDevicePtr<float> mos;
d.Upload(mos, mosaicity, d.n_frames);
const PartialityParams q{mos.get(), wedge_deg, bandwidth_sigma, half_wavelength_A};
IngestPartialityKernel<<<Blocks(d.n_obs), BLK, 0, d.s()>>>(d.n_obs, d.frame.get(), d.zeta.get(), d.d_obs.get(),
d.in_dphi.get(), q, d.obs.get(), d.corr.get());
CudaCheck(cudaGetLastError(), "IngestPartialityKernel launch");
CudaCheck(cudaStreamSynchronize(d.s()), "partiality");
}
void RotationScaleMergeGPU::IngestFinish(float *corr) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
if (d.n_obs > 0)
CopyAndWait(corr, d.corr.get(), size_t(d.n_obs) * sizeof(float), cudaMemcpyDeviceToHost, d.s(),
"download corr");
d.in_key = CudaDevicePtr<uint64_t>();
d.in_dphi = CudaDevicePtr<float>();
d.in_src_start.clear();
d.in_layout = RotationScaleMergeGPU::IngestLayout{};
}
void RotationScaleMergeGPU::BuildFullsCSRs(int32_t *frame_count, std::vector<int32_t> &gstart,
std::vector<int32_t> &gperm) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int nf = d.n_fulls, nfr = d.n_frames, ng = d.n_groups;
const int nfa = std::max(1, nf);
const auto stable_sort = [&](const int32_t *key_in, int32_t *key_out, const int32_t *val_in, int32_t *val_out) {
size_t bytes = 0;
CudaCheck(cub::DeviceRadixSort::SortPairs(nullptr, bytes, key_in, key_out, val_in, val_out, nf, 0, 32, d.s()),
"sort temp size");
CudaDevicePtr<uint8_t> temp = d.Alloc<uint8_t>(std::max<size_t>(1, bytes));
CudaCheck(cub::DeviceRadixSort::SortPairs(temp.get(), bytes, key_in, key_out, val_in, val_out, nf, 0, 32, d.s()),
"sort");
CudaCheck(cudaStreamSynchronize(d.s()), "sort");
};
// The counts come back for the starts, a prefix the host forms in order.
const auto count = [&](const int32_t *bin, int n_bins, CudaDevicePtr<int32_t> &out, std::vector<int32_t> &host) {
out = d.Alloc<int32_t>(std::max(1, n_bins));
CudaCheck(cudaMemsetAsync(out.get(), 0, size_t(std::max(1, n_bins)) * sizeof(int32_t), d.s()), "memset counts");
CountBinsKernel<<<Blocks(nf), BLK, 0, d.s()>>>(nf, bin, out.get());
CudaCheck(cudaGetLastError(), "CountBinsKernel launch");
host.resize(n_bins);
if (n_bins > 0)
CopyAndWait(host.data(), out.get(), size_t(n_bins) * sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(),
"download counts");
};
const auto prefix = [](const std::vector<int32_t> &c) {
std::vector<int32_t> start(c.size() + 1, 0);
for (size_t i = 0; i < c.size(); ++i) start[i + 1] = start[i] + c[i];
return start;
};
CudaDevicePtr<int32_t> idx = d.Alloc<int32_t>(nfa), key = d.Alloc<int32_t>(nfa), sorted = d.Alloc<int32_t>(nfa);
if (nf > 0) IotaKernel<<<(nf + BLK - 1) / BLK, BLK, 0, d.s()>>>(nf, idx.get());
CudaCheck(cudaGetLastError(), "IotaKernel launch");
std::vector<int32_t> fcount;
count(d.f_frame.get(), nfr, d.f_frame_count, fcount);
std::copy(fcount.begin(), fcount.end(), frame_count);
const std::vector<int32_t> fstart = prefix(fcount);
d.Upload(d.f_frame_start, fstart.data(), nfr);
d.f_frame_perm = d.Alloc<int32_t>(nfa);
if (nf > 0) stable_sort(d.f_frame.get(), key.get(), idx.get(), d.f_frame_perm.get());
std::vector<int32_t> gcount;
count(d.f_group.get(), ng, d.f_gcount, gcount);
gstart = prefix(gcount);
d.Upload(d.f_gstart, gstart.data(), ng);
const int n_grouped = gstart.back();
FullGroupKeyKernel<<<Blocks(nf), BLK, 0, d.s()>>>(nf, d.f_group.get(), key.get());
CudaCheck(cudaGetLastError(), "FullGroupKeyKernel launch");
CudaDevicePtr<int32_t> by_group = d.Alloc<int32_t>(nfa);
if (nf > 0) stable_sort(key.get(), sorted.get(), idx.get(), by_group.get());
d.f_gperm = d.Alloc<int32_t>(std::max(1, n_grouped));
gperm.resize(n_grouped);
if (n_grouped > 0) {
CopyAndWait(d.f_gperm.get(), by_group.get(), size_t(n_grouped) * sizeof(int32_t), cudaMemcpyDeviceToDevice,
d.s(), "group perm");
CopyAndWait(gperm.data(), by_group.get(), size_t(n_grouped) * sizeof(int32_t), cudaMemcpyDeviceToHost, d.s(),
"download group perm");
}
}
Reference in New Issue View Git Blame Copy Permalink