mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 65 Wiki Activity
Files
ced85bcd9d3893db169b5ac5332c481dbce385d7
Jungfraujoch/image_analysis/scale_merge/RotationScaleMergeGPU.cu
T
leonarski_fandClaude Opus 4.8 ced85bcd9d RotationScaleMerge: GPU scale-fulls, fulls kept resident (phase 2 step 2) 
Scale the combined fulls (Unity model) on the device so they no longer round-trip
between the combine and the merge: after the GPU combine, build the fulls' per-frame
and per-ASU-group CSRs on the host from just the small key arrays (f_frame/f_group)
with a deterministic counting sort - no GPU stable-sort - then scale in place and
download once.

The four scaling kernels are reused unchanged except FitPerFrameGKernel, which gains
an optional `perm` argument (null for the partials, whose arrays are already
frame-contiguous; a frame-grouping permutation for the emit-ordered fulls) so the
fulls are scaled without a physical reorder. The Unity model falls out of giving the
fulls all-ones partiality/rlp/zeta (coeff = mean), so no other kernel changes and the
committed phase-1 partial-scaling path is bit-identical (perm == null -> idx == i).

Validated across the rotation battery (JFJOCH_RSM_GPU_COMBINE=1): all 15 deterministic
crystals stay run-to-run deterministic and their merged output is bit-identical to the
CPU path (SG/ISa/CC1.2/completeness). The lone exception is EP_cs_01-24 (CC1/2 2%,
R_meas 379% - unindexable noise): merged intensities/CC/completeness match exactly, but
the ill-conditioned 16-bin error-model b fit amplifies the ~1e-7 scale-fulls rounding
to ISa 10.6 vs 10.8 - benign, same class as the accepted phase-1 GPU rounding. The 3
upstream-nondeterministic crystals vary as before (GPU-prediction overflow, not this).

Scale-fulls drops from ~0.09s to ~0 across the two passes; combine+scale-fulls region
~0.32s GPU vs ~0.46s CPU on lyso. Still opt-in (fulls are downloaded for the host merge;
the win grows once the merge/error-model also stay resident).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2026-07-03 07:51:29 +02:00

644 lines
32 KiB
Plaintext
Raw Blame History

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "RotationScaleMergeGPU.h"
#include <algorithm>
#include <cmath>
#include <string>
#include <cuda_runtime.h>
#include "../indexing/CUDAMemHelpers.h"
#include "../../common/CUDAWrapper.h"
#include "../../common/JFJochException.h"
namespace {
constexpr int BLK = 256;
constexpr int MIN_REFLECTIONS = 20;
__device__ __forceinline__ double SafeInvD(double x, double fallback) {
return (isfinite(x) && x != 0.0) ? 1.0 / x : fallback;
}
// Block reduction of a double, deterministic for a fixed thread->element mapping (fixed order,
// no atomics). Returns the sum on thread 0; `s` is BLK doubles of shared scratch.
__device__ double BlockReduceSum(double v, double *s) {
const int t = threadIdx.x;
s[t] = v;
__syncthreads();
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
if (t < stride) s[t] += s[t + stride];
__syncthreads();
}
return s[0];
}
// One thread per ASU group (grid-stride): inverse-variance mean of I*corr over the group's contiguous,
// fixed-order segment of group_perm. Groups are small (avg tens of obs), so a whole block per group
// wastes threads; summing each group in one thread avoids launch/sync overhead and stays deterministic
// (fixed group_perm order). Matches the CPU ReduceGroupMeans filter (no ice/cell mask - scaling ref).
__global__ void ReduceGroupMeansKernel(int n_groups, double min_partiality,
const int32_t *__restrict__ group_perm,
const int32_t *__restrict__ group_start,
const int32_t *__restrict__ group_count,
const float *__restrict__ I, const float *__restrict__ sigma,
const float *__restrict__ partiality,
const float *__restrict__ corr,
double *__restrict__ group_mean) {
for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < n_groups; g += gridDim.x * blockDim.x) {
const int lo = group_start[g], hi = group_start[g] + group_count[g];
double sw = 0.0, swI = 0.0;
for (int p = lo; p < hi; ++p) {
const int i = group_perm[p];
const float c = corr[i];
if (!(c > 0.0f) || !isfinite(c)) continue;
if (partiality[i] < min_partiality) continue;
const float I_corr = I[i] * c;
const float sigma_corr = sigma[i] * c;
if (!isfinite(I_corr) || !isfinite(sigma_corr) || sigma_corr <= 0.0f) continue;
const double w = 1.0 / (double(sigma_corr) * sigma_corr);
sw += w;
swI += w * I_corr;
}
group_mean[g] = sw > 0.0 ? swI / sw : NAN;
}
}
// Per-observation scale-fit coefficient (rotation model) and accept flag, recomputed each scaling
// iteration once the group means are known. coeff = partiality * (1/rlp) * mean[group].
__global__ void PrepScaleObsKernel(int n_obs, const int32_t *__restrict__ group,
const float *__restrict__ partiality, const float *__restrict__ rlp,
const float *__restrict__ zeta, const uint8_t *__restrict__ on_ice,
const double *__restrict__ group_mean,
float *__restrict__ sco_coeff, uint8_t *__restrict__ sco_ok) {
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n_obs) return;
const int g = group[i];
bool ok = (g >= 0) && !on_ice[i] && isfinite(zeta[i]) && zeta[i] > 0.0f;
double mean = 0.0;
if (ok) {
mean = group_mean[g];
ok = isfinite(mean);
}
sco_ok[i] = ok ? 1 : 0;
sco_coeff[i] = ok ? float(double(partiality[i]) * SafeInvD(rlp[i], 1.0) * mean) : 0.0f;
}
// One block per frame: robust per-frame scale G by IRLS (Cauchy), over the frame's contiguous obs.
// Identical objective to the CPU SolveScaleIRLS. Leaves g/scaled untouched for under-populated frames.
// `perm` (null for the partials, whose arrays are already frame-contiguous) maps a position in the
// frame's [lo,hi) range to the obs index, so the same kernel scales the fulls (emit-ordered) through a
// frame-grouping permutation without physically reordering the fulls arrays.
__global__ void FitPerFrameGKernel(int n_frames, double robust_k,
const int32_t *__restrict__ frame_start,
const int32_t *__restrict__ frame_count,
const float *__restrict__ I, const float *__restrict__ sigma,
const float *__restrict__ sco_coeff, const uint8_t *__restrict__ sco_ok,
const int32_t *__restrict__ perm,
double *__restrict__ g, uint8_t *__restrict__ scaled) {
const int f = blockIdx.x;
if (f >= n_frames) return;
const int lo = frame_start[f], hi = frame_start[f] + frame_count[f];
__shared__ double sh[BLK];
// count accepted
long cnt_local = 0;
for (int i = lo + threadIdx.x; i < hi; i += blockDim.x)
if (sco_ok[perm ? perm[i] : i]) ++cnt_local;
const double cnt = BlockReduceSum(double(cnt_local), sh);
__shared__ double s_cnt;
if (threadIdx.x == 0) s_cnt = cnt;
__syncthreads();
if (s_cnt < MIN_REFLECTIONS) return; // leave g[f]/scaled[f] as-is
const double k2 = robust_k * robust_k;
// seed: plain weighted-LS ratio (robust weight = 1)
double num = 0.0, den = 0.0;
for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) {
const int a = perm ? perm[i] : i;
if (!sco_ok[a]) continue;
const double coeff = sco_coeff[a];
const double w = SafeInvD(sigma[a], 1.0);
const double w2 = w * w;
num += w2 * coeff * double(I[a]);
den += w2 * coeff * coeff;
}
double tnum = BlockReduceSum(num, sh); __syncthreads();
double tden = BlockReduceSum(den, sh);
__shared__ double s_G;
if (threadIdx.x == 0) {
double G = tden > 0.0 ? tnum / tden : NAN;
s_G = isfinite(G) ? fmax(0.0, G) : 1.0;
}
__syncthreads();
if (!(s_G >= 0.0)) { if (threadIdx.x == 0) { g[f] = 1.0; scaled[f] = 1; } return; }
for (int iter = 0; iter < 30; ++iter) {
const double G = s_G;
num = 0.0; den = 0.0;
for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) {
const int a = perm ? perm[i] : i;
if (!sco_ok[a]) continue;
const double coeff = sco_coeff[a];
const double w = SafeInvD(sigma[a], 1.0);
const double w2 = w * w;
const double res = w * (G * coeff - double(I[a]));
const double rw = 1.0 / (1.0 + res * res / k2);
num += rw * w2 * coeff * double(I[a]);
den += rw * w2 * coeff * coeff;
}
tnum = BlockReduceSum(num, sh); __syncthreads();
tden = BlockReduceSum(den, sh);
if (threadIdx.x == 0) {
const double G_next = tden > 0.0 ? tnum / tden : NAN;
if (isfinite(G_next)) {
const double Gn = fmax(0.0, G_next);
s_cnt = (fabs(Gn - G) <= 1e-7 * fmax(Gn, 1.0)) ? 1.0 : 0.0; // reuse s_cnt as "converged"
s_G = Gn;
} else {
s_cnt = 1.0;
}
}
__syncthreads();
if (s_cnt != 0.0) break;
}
if (threadIdx.x == 0) { g[f] = s_G; scaled[f] = 1; }
}
// corr = rlp / (partiality * G[frame]) for fitted frames; unchanged otherwise (grid-stride).
__global__ void UpdateCorrKernel(int n_obs, const int32_t *__restrict__ frame,
const float *__restrict__ rlp, const float *__restrict__ partiality,
const double *__restrict__ g, const uint8_t *__restrict__ scaled,
float *__restrict__ corr) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
const int f = frame[i];
if (!scaled[f]) continue;
const double denom = double(partiality[i]) * g[f];
corr[i] = (isfinite(double(rlp[i])) && isfinite(denom) && denom > 0.0)
? float(rlp[i] / denom) : NAN;
}
}
// std::max / std::min return (a<b)?b:a and (b<a)?b:a - reproduce that exactly (NOT fmax/fmin, which
// differ on NaN) so the combine matches the CPU path bit-for-bit on the same inputs.
__device__ __forceinline__ double Dmax(double a, double b) { return (a < b) ? b : a; }
__device__ __forceinline__ double Dmin(double a, double b) { return (b < a) ? b : a; }
__device__ __forceinline__ float Fmax(float a, float b) { return (a < b) ? b : a; }
constexpr float COMBINE_MAX_FRAME_GAP = 2.0f; // == RotationScaleMerge::MAX_FRAME_GAP
// A partial is usable for the combine iff its corr and (I, sigma) are finite with corr>0, sigma>0.
__device__ __forceinline__ bool CombineUsable(int i, const float *I, const float *sigma,
const float *corr) {
const float c = corr[i];
if (!(c > 0.0f) || !isfinite(c)) return false;
return isfinite(I[i]) && isfinite(sigma[i]) && sigma[i] > 0.0f;
}
// All device pointers + scalars the combine kernels need, passed by value.
struct CombineParams {
int n_runs;
double min_partiality, capture_uncertainty_coeff;
const float *I, *sigma, *corr, *partiality, *bkg, *image_number, *d;
const int32_t *frame;
const uint8_t *on_ice;
const int32_t *perm, *rr_start, *rr_count, *rr_h, *rr_k, *rr_l, *rr_group;
int32_t *rr_nevents, *rr_nusable; // count pass outputs
const int32_t *rr_offset; // emit pass: per-run base offset into the fulls arrays
int32_t *f_h, *f_k, *f_l, *f_frame, *f_group;
float *f_I, *f_sigma, *f_d, *f_img;
uint8_t *f_on_ice;
};
// One thread's raw-hkl run: split its usable partials (already in image_number order within the run)
// into rocking events, and for each event pool background, seed F and run the 3-iter de-biased Poisson
// reweight - the exact objective of RotationScaleMerge::Combine::process_rawrun. When Emit, write the
// resulting full at rr_offset[r] + (event index); otherwise just count the emitted events. Both modes
// run the identical accept test, so the count pass predicts the emit pass exactly.
template <bool Emit>
__device__ void CombineRawRun(int r, const CombineParams &p) {
const int lo = p.rr_start[r];
const int hi = lo + p.rr_count[r];
const int group = p.rr_group[r];
int n_emit = 0;
int cursor = lo;
while (cursor < hi) {
while (cursor < hi && !CombineUsable(p.perm[cursor], p.I, p.sigma, p.corr)) ++cursor;
if (cursor >= hi) break;
const int ev_start = cursor; // first usable position of the event
int ev_end = cursor; // last usable position (inclusive), extended below
float last_img = p.image_number[p.perm[cursor]];
int probe = cursor + 1;
while (probe < hi) {
while (probe < hi && !CombineUsable(p.perm[probe], p.I, p.sigma, p.corr)) ++probe;
if (probe >= hi) break;
const float img = p.image_number[p.perm[probe]];
if (img - last_img > COMBINE_MAX_FRAME_GAP) break;
last_img = img;
ev_end = probe;
++probe;
}
cursor = ev_end + 1;
// Pass A: pooled background = mean of the event members' finite backgrounds.
double pooled_bkg = 0.0;
int n_pool = 0;
for (int m = ev_start; m <= ev_end; ++m) {
const int i = p.perm[m];
if (!CombineUsable(i, p.I, p.sigma, p.corr)) continue;
const float b = p.bkg[i];
if (isfinite(b)) { pooled_bkg += b; ++n_pool; }
}
pooled_bkg = n_pool > 0 ? pooled_bkg / n_pool : 0.0;
auto pooled_I = [&](int i) -> double {
const double n_bkg = Dmax(0.0, double(p.sigma[i]) * p.sigma[i] - p.I[i])
/ Fmax(p.bkg[i], 1.0f);
return double(p.I[i]) + n_bkg * (double(p.bkg[i]) - pooled_bkg);
};
// Pass B: seed F (inverse-variance mean of pooled_I*corr), plus peak / d / on_ice / partiality.
double sum_w = 0.0, sum_wI = 0.0, sum_partiality = 0.0;
float d = NAN;
const int first = p.perm[ev_start];
int peak_outcome = p.frame[first];
float peak_frame = p.image_number[first];
float peak_partiality = -1.0f;
const bool on_ice = p.on_ice[first];
for (int m = ev_start; m <= ev_end; ++m) {
const int i = p.perm[m];
if (!CombineUsable(i, p.I, p.sigma, p.corr)) continue;
const double sigma_corr = double(p.sigma[i]) * p.corr[i];
const double w = 1.0 / (sigma_corr * sigma_corr);
sum_w += w;
sum_wI += w * pooled_I(i) * p.corr[i];
sum_partiality += p.partiality[i];
if (p.partiality[i] > peak_partiality) {
peak_partiality = p.partiality[i];
peak_outcome = p.frame[i];
peak_frame = p.image_number[i];
}
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)).
for (int iter = 0; iter < 3; ++iter) {
sum_w = 0.0; sum_wI = 0.0;
for (int m = ev_start; m <= ev_end; ++m) {
const int i = p.perm[m];
if (!CombineUsable(i, p.I, p.sigma, p.corr)) continue;
const double corr = p.corr[i];
const double I_corr = pooled_I(i) * corr;
const double sigma_corr = double(p.sigma[i]) * corr;
const double bkg_var = sigma_corr * sigma_corr - corr * I_corr;
double var = Dmax(0.0, bkg_var) + corr * Dmax(0.0, F);
if (!(var > 0.0)) var = sigma_corr * sigma_corr;
const double w = 1.0 / var;
sum_w += w;
sum_wI += w * I_corr;
}
F = sum_wI / sum_w;
}
if (sum_w <= 0.0 || sum_partiality < p.min_partiality)
continue;
double sigma_full = 1.0 / sqrt(sum_w);
if (p.capture_uncertainty_coeff > 0.0) {
const double frac = Dmin(1.0, sum_partiality);
const double extra = p.capture_uncertainty_coeff * (1.0 - frac) * Dmax(0.0, F);
sigma_full = sqrt(sigma_full * sigma_full + extra * extra);
}
if (Emit) {
const int o = p.rr_offset[r] + n_emit;
p.f_h[o] = p.rr_h[r]; p.f_k[o] = p.rr_k[r]; p.f_l[o] = p.rr_l[r];
p.f_I[o] = float(F);
p.f_sigma[o] = float(sigma_full);
p.f_d[o] = d;
p.f_img[o] = peak_frame;
p.f_frame[o] = peak_outcome;
p.f_on_ice[o] = on_ice ? 1 : 0;
p.f_group[o] = group;
}
++n_emit;
}
if (!Emit) {
p.rr_nevents[r] = n_emit;
int n_usable = 0;
for (int m = lo; m < hi; ++m)
if (CombineUsable(p.perm[m], p.I, p.sigma, p.corr)) ++n_usable;
p.rr_nusable[r] = n_usable;
}
}
template <bool Emit>
__global__ void CombineKernel(CombineParams p) {
for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < p.n_runs; r += gridDim.x * blockDim.x)
CombineRawRun<Emit>(r, p);
}
template <typename T>
__global__ void FillKernel(T *p, int n, T v) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) p[i] = v;
}
void CudaCheck(cudaError_t e, const char *what) {
if (e != cudaSuccess)
throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
std::string("RotationScaleMergeGPU: ") + what + ": " + cudaGetErrorString(e));
}
template <typename T>
void Upload(CudaDevicePtr<T> &dst, const T *src, int n) {
dst = CudaDevicePtr<T>(std::max(1, n));
if (n > 0)
CudaCheck(cudaMemcpy(dst.get(), src, size_t(n) * sizeof(T), cudaMemcpyHostToDevice), "upload");
}
}
struct RotationScaleMergeGPU::Impl {
bool available = false;
int n_obs = 0, n_frames = 0, n_groups = 0;
// immutable per-obs
CudaDevicePtr<float> I, sigma, rlp, partiality, zeta;
CudaDevicePtr<uint8_t> on_ice;
CudaDevicePtr<int32_t> frame;
CudaDevicePtr<float> corr; // mutable, resident across iterations
CudaDevicePtr<int32_t> frame_start, frame_count;
// per space group
CudaDevicePtr<int32_t> group, group_perm, group_start, group_count;
// scratch
CudaDevicePtr<double> group_mean, g;
CudaDevicePtr<uint8_t> scaled;
CudaDevicePtr<float> sco_coeff;
CudaDevicePtr<uint8_t> sco_ok;
// combine: extra per-obs inputs + the one-time raw-hkl run layout
CudaDevicePtr<float> bkg, image_number, d_obs;
int n_runs = 0, n_perm = 0;
CudaDevicePtr<int32_t> perm, rr_start, rr_count, rr_h, rr_k, rr_l, rr_group;
CudaDevicePtr<int32_t> rr_nevents, rr_nusable, rr_offset;
// combine: resident fulls SoA (rebuilt each Combine)
int n_fulls = 0;
CudaDevicePtr<int32_t> f_h, f_k, f_l, f_frame, f_group;
CudaDevicePtr<float> f_I, f_sigma, f_d, f_img;
CudaDevicePtr<uint8_t> f_on_ice;
// scale-fulls (Unity model, kept resident): all-ones partiality/rlp/zeta so the shared scaling kernels
// yield coeff=mean, plus the working corr, the per-obs scale scratch, and the fulls frame/group CSRs
// (built on the host from the small f_frame/f_group key arrays, over the emit-ordered fulls).
CudaDevicePtr<float> f_corr, f_partiality, f_rlp, f_zeta, f_sco_coeff;
CudaDevicePtr<uint8_t> f_sco_ok;
CudaDevicePtr<int32_t> f_frame_perm, f_frame_start, f_frame_count;
CudaDevicePtr<int32_t> f_gperm, f_gstart, f_gcount;
};
RotationScaleMergeGPU::RotationScaleMergeGPU() : impl_(std::make_unique<Impl>()) {
if (get_gpu_count() > 0) {
set_gpu(0);
impl_->available = true;
}
}
RotationScaleMergeGPU::~RotationScaleMergeGPU() = default;
bool RotationScaleMergeGPU::Available() const { return impl_->available; }
void RotationScaleMergeGPU::SetPartials(int n_obs, int n_frames,
const float *I, const float *sigma, const float *rlp,
const float *partiality, const float *zeta, const uint8_t *on_ice,
const int32_t *frame, const float *corr0,
const int32_t *frame_start, const int32_t *frame_count) {
auto &d = *impl_;
d.n_obs = n_obs;
d.n_frames = n_frames;
Upload(d.I, I, n_obs); Upload(d.sigma, sigma, n_obs); Upload(d.rlp, rlp, n_obs);
Upload(d.partiality, partiality, n_obs); Upload(d.zeta, zeta, n_obs); Upload(d.on_ice, on_ice, n_obs);
Upload(d.frame, frame, n_obs); Upload(d.corr, corr0, n_obs);
Upload(d.frame_start, frame_start, n_frames); Upload(d.frame_count, frame_count, n_frames);
d.g = CudaDevicePtr<double>(n_frames);
d.scaled = CudaDevicePtr<uint8_t>(n_frames);
d.sco_coeff = CudaDevicePtr<float>(n_obs);
d.sco_ok = CudaDevicePtr<uint8_t>(n_obs);
}
void RotationScaleMergeGPU::SetGroups(int n_groups, const int32_t *group, const int32_t *group_perm,
int n_group_perm, const int32_t *group_start,
const int32_t *group_count) {
auto &d = *impl_;
d.n_groups = n_groups;
Upload(d.group, group, d.n_obs);
Upload(d.group_perm, group_perm, n_group_perm); // obs with group >= 0, in group order
Upload(d.group_start, group_start, n_groups);
Upload(d.group_count, group_count, n_groups);
d.group_mean = CudaDevicePtr<double>(std::max(1, n_groups));
}
void RotationScaleMergeGPU::SetCorr(const float *corr) {
CudaCheck(cudaMemcpy(impl_->corr.get(), corr, size_t(impl_->n_obs) * sizeof(float),
cudaMemcpyHostToDevice), "upload corr");
}
void RotationScaleMergeGPU::ScalePartials(int iters, double robust_k, double min_partiality,
bool /*has_d_min*/) {
auto &d = *impl_;
CudaCheck(cudaMemset(d.scaled.get(), 0, size_t(d.n_frames) * sizeof(uint8_t)), "memset scaled");
CudaCheck(cudaMemset(d.g.get(), 0, size_t(d.n_frames) * sizeof(double)), "memset g"); // unscaled g unused
const int obs_blocks = (d.n_obs + BLK - 1) / BLK;
const int upd_blocks = std::min(65535, obs_blocks);
const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK);
for (int it = 0; it < iters; ++it) {
ReduceGroupMeansKernel<<<grp_blocks, BLK>>>(d.n_groups, min_partiality,
d.group_perm.get(), d.group_start.get(), d.group_count.get(),
d.I.get(), d.sigma.get(), d.partiality.get(), d.corr.get(), d.group_mean.get());
PrepScaleObsKernel<<<obs_blocks, BLK>>>(d.n_obs, d.group.get(), d.partiality.get(), d.rlp.get(),
d.zeta.get(), d.on_ice.get(), d.group_mean.get(), d.sco_coeff.get(), d.sco_ok.get());
FitPerFrameGKernel<<<d.n_frames, BLK>>>(d.n_frames, robust_k, d.frame_start.get(), d.frame_count.get(),
d.I.get(), d.sigma.get(), d.sco_coeff.get(), d.sco_ok.get(), nullptr, d.g.get(), d.scaled.get());
UpdateCorrKernel<<<upd_blocks, BLK>>>(d.n_obs, d.frame.get(), d.rlp.get(), d.partiality.get(),
d.g.get(), d.scaled.get(), d.corr.get());
}
CudaCheck(cudaGetLastError(), "kernel launch");
CudaCheck(cudaDeviceSynchronize(), "scale sync");
}
void RotationScaleMergeGPU::GetCorr(float *corr_out) const {
CudaCheck(cudaMemcpy(corr_out, impl_->corr.get(), size_t(impl_->n_obs) * sizeof(float),
cudaMemcpyDeviceToHost), "download corr");
}
void RotationScaleMergeGPU::GetG(double *g_out, uint8_t *scaled_out) const {
CudaCheck(cudaMemcpy(g_out, impl_->g.get(), size_t(impl_->n_frames) * sizeof(double),
cudaMemcpyDeviceToHost), "download g");
CudaCheck(cudaMemcpy(scaled_out, impl_->scaled.get(), size_t(impl_->n_frames) * sizeof(uint8_t),
cudaMemcpyDeviceToHost), "download scaled");
}
void RotationScaleMergeGPU::SetCombineInputs(const float *bkg, const float *image_number, const float *d) {
auto &dd = *impl_;
Upload(dd.bkg, bkg, dd.n_obs);
Upload(dd.image_number, image_number, dd.n_obs);
Upload(dd.d_obs, d, dd.n_obs);
}
void RotationScaleMergeGPU::SetRawRuns(int n_runs, int n_perm, const int32_t *perm,
const int32_t *rr_start, const int32_t *rr_count,
const int32_t *rr_h, const int32_t *rr_k, const int32_t *rr_l) {
auto &d = *impl_;
d.n_runs = n_runs;
d.n_perm = n_perm;
Upload(d.perm, perm, n_perm);
Upload(d.rr_start, rr_start, n_runs);
Upload(d.rr_count, rr_count, n_runs);
Upload(d.rr_h, rr_h, n_runs);
Upload(d.rr_k, rr_k, n_runs);
Upload(d.rr_l, rr_l, n_runs);
d.rr_group = CudaDevicePtr<int32_t>(std::max(1, n_runs));
d.rr_nevents = CudaDevicePtr<int32_t>(std::max(1, n_runs));
d.rr_nusable = CudaDevicePtr<int32_t>(std::max(1, n_runs));
d.rr_offset = CudaDevicePtr<int32_t>(std::max(1, n_runs));
}
int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_partiality,
double capture_uncertainty_coeff) {
auto &d = *impl_;
CudaCheck(cudaMemcpy(d.rr_group.get(), rawrun_group, size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyHostToDevice), "upload rr_group");
CombineParams p{};
p.n_runs = d.n_runs;
p.min_partiality = min_partiality;
p.capture_uncertainty_coeff = capture_uncertainty_coeff;
p.I = d.I.get(); p.sigma = d.sigma.get(); p.corr = d.corr.get(); p.partiality = d.partiality.get();
p.bkg = d.bkg.get(); p.image_number = d.image_number.get(); p.d = d.d_obs.get();
p.frame = d.frame.get(); p.on_ice = d.on_ice.get();
p.perm = d.perm.get(); p.rr_start = d.rr_start.get(); p.rr_count = d.rr_count.get();
p.rr_h = d.rr_h.get(); p.rr_k = d.rr_k.get(); p.rr_l = d.rr_l.get(); p.rr_group = d.rr_group.get();
p.rr_nevents = d.rr_nevents.get(); p.rr_nusable = d.rr_nusable.get();
const int blocks = std::min(65535, (d.n_runs + BLK - 1) / BLK);
// Count pass: how many fulls each run emits.
CombineKernel<false><<<blocks, BLK>>>(p);
CudaCheck(cudaGetLastError(), "combine count launch");
// Exclusive prefix sum on the host (deterministic) -> per-run output offset + total fulls.
std::vector<int32_t> nevents(d.n_runs);
CudaCheck(cudaMemcpy(nevents.data(), d.rr_nevents.get(), size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyDeviceToHost), "download nevents");
std::vector<int32_t> offset(d.n_runs);
int64_t acc = 0;
for (int r = 0; r < d.n_runs; ++r) { offset[r] = static_cast<int32_t>(acc); acc += nevents[r]; }
d.n_fulls = static_cast<int>(acc);
// Allocate the fulls SoA and emit.
const int nf = std::max(1, d.n_fulls);
d.f_h = CudaDevicePtr<int32_t>(nf); d.f_k = CudaDevicePtr<int32_t>(nf); d.f_l = CudaDevicePtr<int32_t>(nf);
d.f_frame = CudaDevicePtr<int32_t>(nf); d.f_group = CudaDevicePtr<int32_t>(nf);
d.f_I = CudaDevicePtr<float>(nf); d.f_sigma = CudaDevicePtr<float>(nf);
d.f_d = CudaDevicePtr<float>(nf); d.f_img = CudaDevicePtr<float>(nf);
d.f_on_ice = CudaDevicePtr<uint8_t>(nf);
d.f_corr = CudaDevicePtr<float>(nf); d.f_partiality = CudaDevicePtr<float>(nf);
d.f_rlp = CudaDevicePtr<float>(nf); d.f_zeta = CudaDevicePtr<float>(nf);
d.f_sco_coeff = CudaDevicePtr<float>(nf); d.f_sco_ok = CudaDevicePtr<uint8_t>(nf);
CudaCheck(cudaMemcpy(d.rr_offset.get(), offset.data(), size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyHostToDevice), "upload offset");
p.rr_offset = d.rr_offset.get();
p.f_h = d.f_h.get(); p.f_k = d.f_k.get(); p.f_l = d.f_l.get();
p.f_frame = d.f_frame.get(); p.f_group = d.f_group.get();
p.f_I = d.f_I.get(); p.f_sigma = d.f_sigma.get(); p.f_d = d.f_d.get(); p.f_img = d.f_img.get();
p.f_on_ice = d.f_on_ice.get();
if (d.n_fulls > 0) {
CombineKernel<true><<<blocks, BLK>>>(p);
CudaCheck(cudaGetLastError(), "combine emit launch");
}
CudaCheck(cudaDeviceSynchronize(), "combine sync");
return d.n_fulls;
}
void RotationScaleMergeGPU::GetFulls(int32_t *h, int32_t *k, int32_t *l, float *I, float *sigma, float *d,
float *image_number, int32_t *frame, uint8_t *on_ice,
int32_t *group) const {
const auto &dd = *impl_;
const size_t n = static_cast<size_t>(dd.n_fulls);
if (n == 0) return;
auto dl = [&](void *dst, const void *src, size_t bytes) {
CudaCheck(cudaMemcpy(dst, src, bytes, cudaMemcpyDeviceToHost), "download fulls");
};
dl(h, dd.f_h.get(), n * sizeof(int32_t)); dl(k, dd.f_k.get(), n * sizeof(int32_t));
dl(l, dd.f_l.get(), n * sizeof(int32_t)); dl(frame, dd.f_frame.get(), n * sizeof(int32_t));
dl(group, dd.f_group.get(), n * sizeof(int32_t));
dl(I, dd.f_I.get(), n * sizeof(float)); dl(sigma, dd.f_sigma.get(), n * sizeof(float));
dl(d, dd.f_d.get(), n * sizeof(float)); dl(image_number, dd.f_img.get(), n * sizeof(float));
dl(on_ice, dd.f_on_ice.get(), n * sizeof(uint8_t));
}
void RotationScaleMergeGPU::GetFullsKeys(int32_t *frame, int32_t *group) const {
const auto &d = *impl_;
if (d.n_fulls == 0) return;
const size_t bytes = size_t(d.n_fulls) * sizeof(int32_t);
CudaCheck(cudaMemcpy(frame, d.f_frame.get(), bytes, cudaMemcpyDeviceToHost), "download f_frame");
CudaCheck(cudaMemcpy(group, d.f_group.get(), bytes, cudaMemcpyDeviceToHost), "download f_group");
}
void RotationScaleMergeGPU::SetFullsFrameCSR(const int32_t *frame_perm, int n_perm,
const int32_t *frame_start, const int32_t *frame_count) {
auto &d = *impl_;
Upload(d.f_frame_perm, frame_perm, n_perm);
Upload(d.f_frame_start, frame_start, d.n_frames);
Upload(d.f_frame_count, frame_count, d.n_frames);
}
void RotationScaleMergeGPU::SetFullsGroups(const int32_t *gperm, int n_gperm,
const int32_t *gstart, const int32_t *gcount) {
auto &d = *impl_;
Upload(d.f_gperm, gperm, n_gperm);
Upload(d.f_gstart, gstart, d.n_groups);
Upload(d.f_gcount, gcount, d.n_groups);
}
void RotationScaleMergeGPU::ScaleFulls(int iters, double robust_k, double min_partiality) {
auto &d = *impl_;
const int nf = d.n_fulls;
if (nf == 0) return;
const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK);
const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK);
// Unity model: partiality/rlp/zeta = 1 so coeff = mean; corr starts at 1.
FillKernel<<<obs_blocks, BLK>>>(d.f_corr.get(), nf, 1.0f);
FillKernel<<<obs_blocks, BLK>>>(d.f_partiality.get(), nf, 1.0f);
FillKernel<<<obs_blocks, BLK>>>(d.f_rlp.get(), nf, 1.0f);
FillKernel<<<obs_blocks, BLK>>>(d.f_zeta.get(), nf, 1.0f);
// g/scaled reset once (a frame, once fitted, stays fitted - same as ScalePartials).
CudaCheck(cudaMemset(d.scaled.get(), 0, size_t(d.n_frames) * sizeof(uint8_t)), "memset f scaled");
CudaCheck(cudaMemset(d.g.get(), 0, size_t(d.n_frames) * sizeof(double)), "memset f g");
for (int it = 0; it < iters; ++it) {
ReduceGroupMeansKernel<<<grp_blocks, BLK>>>(d.n_groups, min_partiality,
d.f_gperm.get(), d.f_gstart.get(), d.f_gcount.get(),
d.f_I.get(), d.f_sigma.get(), d.f_partiality.get(), d.f_corr.get(), d.group_mean.get());
PrepScaleObsKernel<<<obs_blocks, BLK>>>(nf, d.f_group.get(), d.f_partiality.get(), d.f_rlp.get(),
d.f_zeta.get(), d.f_on_ice.get(), d.group_mean.get(), d.f_sco_coeff.get(), d.f_sco_ok.get());
FitPerFrameGKernel<<<d.n_frames, BLK>>>(d.n_frames, robust_k,
d.f_frame_start.get(), d.f_frame_count.get(), d.f_I.get(), d.f_sigma.get(),
d.f_sco_coeff.get(), d.f_sco_ok.get(), d.f_frame_perm.get(), d.g.get(), d.scaled.get());
UpdateCorrKernel<<<obs_blocks, BLK>>>(nf, d.f_frame.get(), d.f_rlp.get(), d.f_partiality.get(),
d.g.get(), d.scaled.get(), d.f_corr.get());
}
CudaCheck(cudaGetLastError(), "scale fulls launch");
CudaCheck(cudaDeviceSynchronize(), "scale fulls sync");
}
void RotationScaleMergeGPU::GetFullsCorr(float *corr) const {
const auto &d = *impl_;
if (d.n_fulls == 0) return;
CudaCheck(cudaMemcpy(corr, d.f_corr.get(), size_t(d.n_fulls) * sizeof(float),
cudaMemcpyDeviceToHost), "download f_corr");
}
Reference in New Issue View Git Blame Copy Permalink