Files
Jungfraujoch/image_analysis/bragg_integration/BraggIntegrationEngineGPU.cu
T
leonarski_fandClaude Opus 5 6f7b136ec2 Bragg integration: a shared signal pixel belongs to the nearer reflection
Nothing kept a neighbour's flux out of a reflection's own signal disk. The union mask
keeps neighbour cores out of the BACKGROUND ring, but the r1 disk was read whole, so on
a dense pattern a crowded reflection measures part of its neighbour as its own.

Ownership is decided once per image into a per-pixel (quantised distance, reflection)
key written with an atomic minimum, so the nearest predicted centre wins whatever order
the writes arrive in and the lowest index breaks a tie. `--overlap exclude`, now the
default, drops the pixels a nearer neighbour owns from the profile fit. A profile fit is
the amplitude of a normalised profile, so leaving pixels out renormalises the estimator
by construction and the reflection stays unbiased rather than being discarded; the
summation-fallback guard is scaled back to the disk the box-sum seed actually read, so
it still compares like with like. `--overlap reject` is the XDS MINPK alternative - drop
the reflection when less than `--overlap-minpk` of its expected profile is cleanly its
own. A box sum has no profile to renormalise with, so `exclude` is a no-op there and
only `reject` acts on it.

Widening the split - keeping a pixel only where no other centre is within its distance
PLUS a margin - was built and measured, and it is worse monotonically: the residual bias
of the pixels that were kept grows from +0.072 to +0.209 in ln intensity at 0 to 3 px of
margin. What the margin removes is the reflection's own profile, not the neighbour's
tail, so the plain nearest-centre split is the rule.

Measured on the full 38-crystal rotation battery against the same binary with the
treatment off: ISa better 15 / worse 8, summed shortfall against XDS 39.7 -> 28.1. Three
of the losses are the two-pass loop taking its other branch - their median mosaicity
moves between the two known attractors - rather than the change under test; excluding
those it is better 15 / worse 5 and the shortfall goes 31.3 -> 14.4. The two crowded
crystals gain 38% and 52% of their ISa, one of them passing XDS. High-shell CC1/2 over
the 35 crystals that neither flipped branch nor carry a collapsed error model is better
7 / worse 7. Space groups unchanged at 35/38. The owner map is built only when a
treatment is asked for and costs 1.1% of the battery's wall clock - 23% on a genuinely
crowded crystal, nothing where no two predictions touch.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-11 05:02:31 +02:00

796 lines
41 KiB
Plaintext

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include "BraggIntegrationEngineGPU.h"
using namespace bragg_engine;
namespace {
inline void cuda_err(cudaError_t val) {
if (val != cudaSuccess)
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
}
// Fixed scalars passed by value to every kernel (mirrors BraggIntegrationEngine's members).
struct BraggGpuParams {
int W, H;
float r1_sq, r2, r2_sq, r3, r3_sq;
int R, G, GG;
float bkg_clip_nsigma; // high-side background sigma-clip multiplier (0 = no clip)
int empirical; // ProfileEmpirical vs ProfileGaussian
int use_ellipse;
float bw_sigma;
float c_radial;
float F_px;
float beam_x, beam_y;
float bkg_trim; // idea 1: symmetric trimmed-mean background fraction (0 = plain ring mean)
BraggStencilParams stencil; // per-reflection signal/background geometry (BraggStencil.h)
int rad_w; // radial-background window held in shared memory, in bins of one pixel
int n_kern; // radial-background kernels in the table
int overlap; // 0 = off, 1 = reject, 2 = exclude (OverlapMode)
int boxsum_reject; // BoxSum mode under Reject: drop on the disk-AREA fraction (see the CPU engine)
int exclude_px; // drop a neighbour's pixels from the disk: Exclude, and not a box sum
float claim_sq, inv_claim; // how far a reflection claims pixels in the owner map
float minpk; // Reject: least clean profile fraction that is kept
};
__device__ inline bool valid(int32_t v) { return v != INT32_MIN && v != INT32_MAX; }
// Learned second moments, (sum v*rad^2, sum v*tan^2, sum v) per shell plus one global slot at N_SHELL.
constexpr int MOM_STRIDE = 3;
// --- Mark the r2 signal region of every predicted reflection (race-free: all writes are 1). The
// region is the INNER stencil ellipse in the neighbour's own frame; see the CPU engine.
// The same sweep fills the owner map when an overlap treatment is on: a pixel two signal regions
// share belongs to the nearer centre, which an atomicMin over (distance, index) settles without
// any ordering or sorting (BraggStencil.h). The claim disk sits inside the ellipse this loop
// already walks, so ownership costs one atomic on the pixels already visited. ---
__global__ void mark_mask(const float *px_x, const float *px_y, uint8_t *mask, uint32_t *owner,
BraggGpuParams p, int n) {
const int i = blockIdx.x;
if (i >= n) return;
const float cx = px_x[i], cy = px_y[i];
const BraggStencil st = MakeBraggStencil(cx, cy, p.stencil);
const int x0 = max(0, (int) floorf(cx - st.ex_in - 1.0f));
const int x1 = min(p.W - 1, (int) ceilf(cx + st.ex_in + 1.0f));
const int y0 = max(0, (int) floorf(cy - st.ey_in - 1.0f));
const int y1 = min(p.H - 1, (int) ceilf(cy + st.ey_in + 1.0f));
const int bw = x1 - x0 + 1, bh = y1 - y0 + 1;
if (bw <= 0 || bh <= 0) return;
for (int t = threadIdx.x; t < bw * bh; t += blockDim.x) {
const int x = x0 + t % bw, y = y0 + t / bw;
const BraggStencilDist d = BraggStencilDistances(st, (float) x - cx, (float) y - cy);
if (d.inner < p.r2_sq) mask[y * p.W + x] = 1;
if (p.overlap && d.signal < p.claim_sq)
atomicMin(&owner[y * p.W + x], BraggOwnerKey(sqrtf(d.signal), p.inv_claim, i));
}
}
// --- Pass A box-sum: rough I / background / centroid / strong flag, one block per reflection. ---
__global__ void boxsum(const float *px_x, const float *px_y, const float *dd,
const int32_t *img, const uint8_t *mask, const uint32_t *owner,
BraggGpuParams p, int n,
int *cx_o, int *cy_o, float *I_o, float *sigma_o, float *bkg_o,
float *bkgvar_o, float *varbkg_o, float *obsx_o, float *obsy_o, uint8_t *ok_o, uint8_t *strong_o,
uint8_t *hasobs_o, unsigned long long *invd2mm,
float *isum_o, int *ninner_o, int *rbin_o, int *kbin_o,
float *rad_sum, int *rad_cnt, int n_rad) {
const int i = blockIdx.x;
if (i >= n) return;
__shared__ unsigned long long s_Isum, s_Ix, s_Iy;
__shared__ int s_ninner, s_ninner_valid, s_nbkg, s_ndisk, s_nown;
__shared__ double s_bkgsum;
__shared__ int s_accept;
__shared__ double s_bkg, s_thr, s_clipsum;
__shared__ int s_clipn;
__shared__ float s_r0;
// The stencil spans r0 +- the radial semi-axis of the outer ellipse, so the window has to be
// sized from the widest aperture on the detector - which the host does, into p.rad_w. A fixed
// 32-bin window silently dropped the outer bins as soon as anything was elongated.
extern __shared__ float s_rad[];
float *s_radv = s_rad;
int *s_radn = (int *) (s_rad + p.rad_w);
__shared__ int s_radbase;
if (threadIdx.x == 0) {
const float rx = px_x[i] - p.beam_x, ry = px_y[i] - p.beam_y;
s_r0 = sqrtf(rx * rx + ry * ry);
s_Isum = 0; s_Ix = 0; s_Iy = 0;
s_ninner = 0; s_ninner_valid = 0; s_nbkg = 0; s_bkgsum = 0.0;
s_ndisk = 0; s_nown = 0;
}
__syncthreads();
const float cx = px_x[i], cy = px_y[i];
// Every thread builds the same stencil from the same inputs; it is a few flops against a whole
// bounding box of pixel reads, so it costs less than staging it through shared memory.
const BraggStencil st = MakeBraggStencil(cx, cy, p.stencil);
const int x0 = max(0, (int) floorf(cx - st.ex_out - 1.0f));
const int x1 = min(p.W - 1, (int) ceilf(cx + st.ex_out + 1.0f));
const int y0 = max(0, (int) floorf(cy - st.ey_out - 1.0f));
const int y1 = min(p.H - 1, (int) ceilf(cy + st.ey_out + 1.0f));
const int bw = x1 - x0 + 1, bh = y1 - y0 + 1;
const int area = (bw > 0 && bh > 0) ? bw * bh : 0;
long long l_Isum = 0, l_Ix = 0, l_Iy = 0;
int l_ni = 0, l_niv = 0, l_nb = 0, l_nd = 0, l_no = 0;
double l_bkg = 0.0;
for (int t = threadIdx.x; t < area; t += blockDim.x) {
const int x = x0 + t % bw, y = y0 + t / bw;
const BraggStencilDist d = BraggStencilDistances(st, (float) x - cx, (float) y - cy);
const int32_t px = img[y * p.W + x];
if (d.signal < p.r1_sq) {
// A pixel a nearer neighbour owns carries that neighbour's flux; see the CPU engine.
++l_nd;
if (p.overlap) {
if (BraggOwnedBy(owner[y * p.W + x], i)) ++l_no;
else if (p.exclude_px) continue;
}
++l_ni;
if (valid(px)) { l_Isum += px; l_Ix += (long long) x * px; l_Iy += (long long) y * px; ++l_niv; }
} else if (d.inner >= p.r2_sq && d.outer < p.r3_sq) {
if (mask[y * p.W + x]) continue;
if (!valid(px)) continue;
l_bkg += (double) px; ++l_nb;
}
}
atomicAdd(&s_Isum, (unsigned long long) l_Isum);
atomicAdd(&s_Ix, (unsigned long long) l_Ix);
atomicAdd(&s_Iy, (unsigned long long) l_Iy);
atomicAdd(&s_ninner, l_ni);
atomicAdd(&s_ninner_valid, l_niv);
atomicAdd(&s_nbkg, l_nb);
atomicAdd(&s_bkgsum, l_bkg);
atomicAdd(&s_ndisk, l_nd);
atomicAdd(&s_nown, l_no);
__syncthreads();
for (int t = threadIdx.x; t < p.rad_w; t += blockDim.x) { s_radv[t] = 0.0f; s_radn[t] = 0; }
if (threadIdx.x == 0) s_radbase = (int) lroundf(s_r0) - p.rad_w / 2;
if (threadIdx.x == 0) {
s_accept = (s_ninner_valid == s_ninner && s_nbkg > 5) ? 1 : 0;
s_bkg = s_accept ? (s_bkgsum / (double) s_nbkg) : 0.0;
s_thr = s_bkg + (double) p.bkg_clip_nsigma * sqrt(fmax(s_bkg, 1.0));
s_clipsum = 0.0; s_clipn = 0;
}
__syncthreads();
// Trimmed-mean background (idea 1): collect the annulus into shared memory, bitonic-sort the block, and
// average the middle (1 - 2*bkg_trim) fraction - robust to the high-side contamination that biases the
// plain ring mean. Falls back to the flat mean when the ring exceeds the shared buffer.
__shared__ int s_bvals[BKG_TRIM_MAX];
__shared__ int s_bn;
if (threadIdx.x == 0) s_bn = 0;
__syncthreads();
const bool do_trim = s_accept && p.bkg_trim > 0.0f && s_nbkg > 5 && s_nbkg <= BKG_TRIM_MAX;
if (do_trim) {
for (int t = threadIdx.x; t < area; t += blockDim.x) {
const int x = x0 + t % bw, y = y0 + t / bw;
const BraggStencilDist d = BraggStencilDistances(st, (float) x - cx, (float) y - cy);
if (!(d.inner >= p.r2_sq && d.outer < p.r3_sq)) continue;
if (mask[y * p.W + x]) continue;
const int32_t px = img[y * p.W + x];
if (!valid(px)) continue;
const int slot = atomicAdd(&s_bn, 1);
if (slot < BKG_TRIM_MAX) s_bvals[slot] = px;
}
__syncthreads();
const int nb = min(s_bn, BKG_TRIM_MAX);
int n2 = 1;
while (n2 < nb) n2 <<= 1;
for (int t = threadIdx.x + nb; t < n2; t += blockDim.x) s_bvals[t] = INT32_MAX; // pad to power of 2
__syncthreads();
for (int k = 2; k <= n2; k <<= 1) // ascending bitonic sort of s_bvals[0..n2)
for (int j = k >> 1; j > 0; j >>= 1) {
for (int idx = threadIdx.x; idx < n2; idx += blockDim.x) {
const int ixj = idx ^ j;
if (ixj > idx) {
const bool up = ((idx & k) == 0);
const int a = s_bvals[idx], b = s_bvals[ixj];
if ((up && a > b) || (!up && a < b)) { s_bvals[idx] = b; s_bvals[ixj] = a; }
}
}
__syncthreads();
}
if (threadIdx.x == 0) {
const int lo = (int) (nb * p.bkg_trim), hi = nb - lo;
if (hi > lo) {
double s = 0.0;
for (int t = lo; t < hi; ++t) s += (double) s_bvals[t];
s_bkg = s / (double) (hi - lo);
}
}
__syncthreads();
}
// Second ring pass for the high-side sigma-clip (re-reads the annulus; avoids storing bkg values).
if (s_accept && p.bkg_clip_nsigma > 0.0f && !do_trim) {
double c_l = 0.0; int cn_l = 0;
for (int t = threadIdx.x; t < area; t += blockDim.x) {
const int x = x0 + t % bw, y = y0 + t / bw;
const BraggStencilDist d = BraggStencilDistances(st, (float) x - cx, (float) y - cy);
if (!(d.inner >= p.r2_sq && d.outer < p.r3_sq)) continue;
if (mask[y * p.W + x]) continue;
const int32_t px = img[y * p.W + x];
if (!valid(px)) continue;
if ((double) px <= s_thr) {
c_l += px; ++cn_l;
if (n_rad > 0) {
// The radial curve is binned on the TRUE detector radius, so this is the
// unshrunk radial projection - no per-pixel sqrt.
const int idx = (int) lroundf(s_r0 + d.rad) - s_radbase;
if (idx >= 0 && idx < p.rad_w) {
atomicAdd(&s_radv[idx], (float) px); // shared, not global
atomicAdd(&s_radn[idx], 1);
}
}
}
}
atomicAdd(&s_clipsum, c_l); atomicAdd(&s_clipn, cn_l);
}
__syncthreads();
if (n_rad > 0) {
for (int t = threadIdx.x; t < p.rad_w; t += blockDim.x) {
if (s_radn[t] == 0) continue;
const int b = min(max(s_radbase + t, 0), n_rad - 1);
atomicAdd(&rad_sum[b], s_radv[t]);
atomicAdd(&rad_cnt[b], s_radn[t]);
}
__syncthreads();
}
if (threadIdx.x != 0) return;
if (!s_accept) { ok_o[i] = 0; strong_o[i] = 0; hasobs_o[i] = 0; return; }
// A box sum has no profile to renormalise a disk it has taken pixels out of, so dropping the
// reflection is the only overlap treatment it has (see the CPU engine).
if (p.boxsum_reject && (float) s_nown < p.minpk * (float) s_ndisk) {
ok_o[i] = 0; strong_o[i] = 0; hasobs_o[i] = 0; return;
}
double bkg = s_bkg;
int n_bkg_used = s_nbkg; // pixels behind the FINAL background value (trim/clip shrink it)
if (do_trim) {
const int nb = min(s_bn, BKG_TRIM_MAX);
const int lo = (int) (nb * p.bkg_trim), hi = nb - lo;
if (hi > lo) n_bkg_used = hi - lo;
}
if (p.bkg_clip_nsigma > 0.0f && s_clipn > 5) { bkg = s_clipsum / (double) s_clipn; n_bkg_used = s_clipn; }
const long long Isum = (long long) s_Isum;
const double I = (double) Isum - (double) s_ninner * bkg;
// See the CPU engine: bkg is estimated from n_bkg_used pixels and subtracted n_inner times, so
// var(I) = Isum + n_inner^2 * bkg/n_bkg_used. Both engines must agree.
const double bkg_var = bkg / (double) n_bkg_used;
const double var_bkg_term = (double) s_ninner * (double) s_ninner * bkg_var;
double sigma = 1.0;
uint8_t hasobs = 0; double ox = 0.0, oy = 0.0;
if (Isum > 0) {
sigma = fmax(sigma, sqrt((double) Isum + var_bkg_term));
ox = (double) (long long) s_Ix / (double) Isum;
oy = (double) (long long) s_Iy / (double) Isum;
hasobs = 1;
}
cx_o[i] = (int) lroundf(cx);
cy_o[i] = (int) lroundf(cy);
I_o[i] = (float) I; sigma_o[i] = (float) sigma; bkg_o[i] = (float) bkg;
bkgvar_o[i] = (float) bkg_var;
varbkg_o[i] = (float) ((double) s_ninner * bkg + var_bkg_term);
isum_o[i] = (float) Isum;
ninner_o[i] = s_ninner;
rbin_o[i] = min(max((int) lroundf(s_r0), 0), n_rad > 0 ? n_rad - 1 : 0);
kbin_o[i] = BraggStencilKernelIndex(st, p.n_kern);
obsx_o[i] = (float) ox; obsy_o[i] = (float) oy; hasobs_o[i] = hasobs;
ok_o[i] = 1;
strong_o[i] = (sigma > 0.0 && I / sigma >= STRONG_I_OVER_SIGMA) ? 1 : 0;
const float d = dd[i];
if (d > 0.0f) {
// Positive doubles keep IEEE bit-pattern ordering, so atomicMin/Max on the ull view works.
const unsigned long long b = (unsigned long long) __double_as_longlong(1.0 / ((double) d * d));
atomicMin(&invd2mm[0], b);
atomicMax(&invd2mm[1], b);
}
}
// Resolution shell of one reflection from the global inv-d^2 range (mirrors CPU shell_of). Computed
// inline in learn_profile and fit so no separate shell array/kernel is needed.
__device__ inline int compute_shell(float d, const unsigned long long *invd2mm) {
const unsigned long long mn = invd2mm[0], mx = invd2mm[1];
if (!(d > 0.0f) || mx <= mn) return 0;
const double invd2 = 1.0 / ((double) d * d);
const double dmn = __longlong_as_double((long long) mn), dmx = __longlong_as_double((long long) mx);
const int s = (int) ((invd2 - dmn) / (dmx - dmn) * N_SHELL);
return s < 0 ? 0 : (s >= N_SHELL ? N_SHELL - 1 : s);
}
// --- Zero the profile accumulators and seed the inv-d^2 range, in one launch (replaces a handful
// of small cudaMemsetAsync calls, which matter when kernel-launch latency is high). ---
__global__ void reset(float *shell_grid, float *global_grid, float *mom, int *shell_n, int *global_n,
unsigned long long *invd2mm, int GG) {
for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < N_SHELL * GG; k += blockDim.x * gridDim.x)
shell_grid[k] = 0.0f;
for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < GG; k += blockDim.x * gridDim.x)
global_grid[k] = 0.0f;
if (blockIdx.x == 0 && threadIdx.x < MOM_STRIDE * (N_SHELL + 1)) mom[threadIdx.x] = 0.0f;
if (blockIdx.x == 0 && threadIdx.x < N_SHELL) shell_n[threadIdx.x] = 0;
if (blockIdx.x == 0 && threadIdx.x == 0) {
*global_n = 0;
invd2mm[0] = ~0ull; // min seed
invd2mm[1] = 0ull; // max seed
}
}
// --- Learn the profile: each strong spot adds its bkg-subtracted, I-normalised grid to its shell
// (and the global grid), and its second moments to the same shell's moment slot. The grid stays
// in the DETECTOR frame (that is where the empirical profile is applied); the moments are taken in
// the spot's own radial/tangential frame, because a detector-frame stack is azimuthally averaged
// and cannot tell a radially smeared spot from a tangentially wide one (see the CPU engine).
// One block per reflection. ---
__global__ void learn_profile(const int32_t *img, const uint32_t *owner,
const float *px_x, const float *px_y,
const int *cx_a, const int *cy_a, const float *dd,
const unsigned long long *invd2mm,
const float *I_a, const float *bkg_a, const uint8_t *ok_a, const uint8_t *strong_a,
float *shell_grid, float *global_grid, float *mom, int *shell_n, int *global_n,
BraggGpuParams p, int n) {
const int i = blockIdx.x;
if (i >= n || !ok_a[i] || !strong_a[i]) return;
const float I = I_a[i];
if (!(I > 0.0f)) return;
const int cx = cx_a[i], cy = cy_a[i], sh = compute_shell(dd[i], invd2mm);
const float bkg = bkg_a[i];
const float rx = px_x[i] - p.beam_x, ry = px_y[i] - p.beam_y;
const float Rpx = sqrtf(rx * rx + ry * ry);
const float ux = Rpx > 1e-6f ? rx / Rpx : 1.0f, uy = Rpx > 1e-6f ? ry / Rpx : 0.0f;
float *sg = shell_grid + (size_t) sh * p.GG;
__shared__ float s_rad, s_tan, s_w;
if (threadIdx.x == 0) { s_rad = 0.0f; s_tan = 0.0f; s_w = 0.0f; }
__syncthreads();
float l_rad = 0.0f, l_tan = 0.0f, l_w = 0.0f;
for (int k = threadIdx.x; k < p.GG; k += blockDim.x) {
const int dx = k % p.G - p.R, dy = k / p.G - p.R;
const int x = cx + dx, y = cy + dy;
if (x < 0 || y < 0 || x >= p.W || y >= p.H) continue;
const int32_t px = img[y * p.W + x];
if (!valid(px)) continue;
if (p.overlap == 2 && !BraggOwnedBy(owner[y * p.W + x], i)) continue;
const float v = ((float) px - bkg) / I;
atomicAdd(&sg[k], v);
atomicAdd(&global_grid[k], v);
if ((float) (dx * dx + dy * dy) >= p.r1_sq) continue;
const float rad = dx * ux + dy * uy, tn = -dx * uy + dy * ux;
l_rad += v * rad * rad; l_tan += v * tn * tn; l_w += v;
}
// Reduce in shared memory first: every block adds to the same handful of global words, so one
// atomic per block instead of one per thread.
atomicAdd(&s_rad, l_rad); atomicAdd(&s_tan, l_tan); atomicAdd(&s_w, l_w);
__syncthreads();
if (threadIdx.x == 0) {
atomicAdd(&mom[MOM_STRIDE * sh + 0], s_rad);
atomicAdd(&mom[MOM_STRIDE * sh + 1], s_tan);
atomicAdd(&mom[MOM_STRIDE * sh + 2], s_w);
atomicAdd(&mom[MOM_STRIDE * N_SHELL + 0], s_rad);
atomicAdd(&mom[MOM_STRIDE * N_SHELL + 1], s_tan);
atomicAdd(&mom[MOM_STRIDE * N_SHELL + 2], s_w);
atomicAdd(&shell_n[sh], 1);
atomicAdd(global_n, 1);
}
}
// --- Turn the learned moments into radial/tangential variances and, for empirical, normalise the
// learned grid into a profile. One block per grid: blocks [0,N_SHELL) are the shells, block
// N_SHELL is the global one. ---
__global__ void build_profiles(const float *shell_grid, const float *global_grid, const float *mom,
const int *shell_n, const int *global_n,
float *shell_P, float *global_P, float *sigma2_r, float *sigma2_t,
BraggGpuParams p) {
const int b = blockIdx.x;
const float *grid; int nstrong; float *P;
if (b < N_SHELL) { grid = shell_grid + (size_t) b * p.GG; nstrong = shell_n[b]; P = shell_P + (size_t) b * p.GG; }
else { grid = global_grid; nstrong = *global_n; P = global_P; }
if (threadIdx.x == 0) {
const float w = mom[MOM_STRIDE * b + 2];
sigma2_r[b] = w > 0.0f ? fmaxf(0.25f, mom[MOM_STRIDE * b + 0] / w) : 1.0f;
sigma2_t[b] = w > 0.0f ? fmaxf(0.25f, mom[MOM_STRIDE * b + 1] / w) : 1.0f;
}
if (!p.empirical) return;
__shared__ float s_sum;
if (threadIdx.x == 0) s_sum = 0.0f;
__syncthreads();
float l_sum = 0.0f;
for (int k = threadIdx.x; k < p.GG; k += blockDim.x) {
const float g = fmaxf(0.0f, grid[k]); // a profile has to be non-negative
P[k] = g;
l_sum += g;
}
atomicAdd(&s_sum, l_sum);
__syncthreads();
const bool normalise = nstrong > 0 && s_sum > 0.0f;
for (int k = threadIdx.x; k < p.GG; k += blockDim.x) P[k] = normalise ? P[k] / s_sum : 0.0f;
}
// --- Radial background curvature correction: one thread per reflection, no pixel reads.
// The disk and the annulus are concentric, so a background linear in position cancels between
// them; what is left is the curvature of the radial background. k_diff (annulus minus disk
// histogram over radial offset) turns that into one short dot product. An empty radial bin
// contributes the reflection's own background, so an empty neighbourhood gives exactly zero
// correction because the kernel weights sum to zero. The kernel depends on how far this
// reflection's ring was elongated, so k_diff is a table and kbin_a picks the row.
// Mirrors BraggIntegrationEngineCPU. ---
__global__ void radial_correct(const float *rad_sum, const int *rad_cnt, int n_rad,
const float *k_diff, int k_len, int k_off,
const float *isum_a, const int *ninner_a, const int *rbin_a,
const int *kbin_a,
const uint8_t *ok_a, float *bkg_o, float *I_o, int n) {
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n || !ok_a[i]) return;
const float bkg = bkg_o[i];
const float *kern = k_diff + (size_t) kbin_a[i] * k_len;
float corr = 0.0f;
for (int k = 0; k < k_len; ++k) {
int b = rbin_a[i] + k - k_off;
b = min(max(b, 0), n_rad - 1);
const float v = rad_cnt[b] > 0 ? rad_sum[b] / (float) rad_cnt[b] : bkg;
corr += kern[k] * v;
}
const float bkg_new = bkg - corr;
bkg_o[i] = bkg_new;
I_o[i] = isum_a[i] - (float) ninner_a[i] * bkg_new;
}
// --- Pass B Kabsch profile fit: I = sum P(c-B)/v over sum P^2/v, v = B + max(I,0)P (iterate).
// One block per reflection; the (possibly elongated) profile is built in shared memory. ---
__global__ void fit(const int32_t *img, const uint32_t *owner, const float *px_x, const float *px_y,
const int *cx_a, const int *cy_a, const float *dd, const unsigned long long *invd2mm,
const float *I_seed, const float *sigma_seed, const float *bkg_a,
const float *bkgvar_a, const float *varbkg_seed, const uint8_t *ok_a,
const float *shell_P, const float *global_P,
const float *sigma2_r, const float *sigma2_t, const int *shell_n,
float *I_o, float *sigma_o, float *varbkg_o, uint8_t *ok_o, BraggGpuParams p, int n) {
const int i = blockIdx.x;
if (i >= n) return;
extern __shared__ float Pbuf[];
__shared__ float s_gs, s_num, s_den, s_I, s_wsum;
__shared__ float s_pown, s_mall, s_mown;
__shared__ int s_Rf, s_Gf;
if (!ok_a[i]) { if (threadIdx.x == 0) ok_o[i] = 0; return; }
const int cx = cx_a[i], cy = cy_a[i];
const int sh = compute_shell(dd[i], invd2mm);
const bool use_shell = shell_n[sh] >= MIN_STRONG_PER_SHELL; // else fall back to the global profile
const float bkg = bkg_a[i];
if (p.empirical) {
const float *Psrc = use_shell ? (shell_P + (size_t) sh * p.GG) : global_P;
if (threadIdx.x == 0) { s_Rf = p.R; s_Gf = p.G; }
__syncthreads();
for (int k = threadIdx.x; k < p.GG; k += blockDim.x) Pbuf[k] = Psrc[k];
__syncthreads();
} else {
const int si = use_shell ? sh : N_SHELL; // N_SHELL is the global slot
const float s2t = sigma2_t[si];
const float rx = px_x[i] - p.beam_x, ry = px_y[i] - p.beam_y;
const float Rpx = sqrtf(rx * rx + ry * ry);
const float tan2t = Rpx / p.F_px;
float s2r = s2t, ux = 1.0f, uy = 0.0f;
bool elong = false;
if (p.use_ellipse) {
// Measured radial excess over the tangential width, floored by the analytic bandwidth +
// parallax/capture term (see the CPU engine).
const float sbw = p.bw_sigma * Rpx;
const float radial_extra = fmaxf(sigma2_r[si] - s2t, sbw * sbw + p.c_radial * tan2t * tan2t);
if (Rpx > 1e-6f && radial_extra > 0.25f) { ux = rx / Rpx; uy = ry / Rpx; s2r = s2t + radial_extra; elong = true; }
}
const int Rf = elong ? min(3 * p.R, (int) ceilf(p.r2 + 2.0f * sqrtf(s2r))) : p.R;
const int Gf = 2 * Rf + 1;
if (threadIdx.x == 0) { s_Rf = Rf; s_Gf = Gf; s_gs = 0.0f; }
__syncthreads();
const float fx = px_x[i] - cx, fy = px_y[i] - cy;
float l_gs = 0.0f;
for (int k = threadIdx.x; k < Gf * Gf; k += blockDim.x) {
const float ex = (k % Gf - Rf) - fx, ey = (k / Gf - Rf) - fy;
const float rad = ex * ux + ey * uy, tn = -ex * uy + ey * ux;
const float g = expf(-rad * rad / (2.0f * s2r) - tn * tn / (2.0f * s2t));
Pbuf[k] = g; l_gs += g;
}
atomicAdd(&s_gs, l_gs);
__syncthreads();
const float gs = s_gs;
for (int k = threadIdx.x; k < Gf * Gf; k += blockDim.x) Pbuf[k] /= gs;
__syncthreads();
}
const int Rf = s_Rf, Gf = s_Gf, GfGf = Gf * Gf;
const float B = fmaxf(bkg, (float) PIXEL_VARIANCE_FLOOR);
// How much of the expected profile is cleanly this reflection's own: s_pown over the whole grid
// (XDS's MINPK quantity, what Reject cuts on) and s_mown / s_mall over the r1 disk alone, which
// is what the summation seed the runaway guard compares against actually saw. See the CPU engine.
if (p.overlap) {
if (threadIdx.x == 0) { s_pown = 0.0f; s_mall = 0.0f; s_mown = 0.0f; }
__syncthreads();
float l_pown = 0.0f, l_mall = 0.0f, l_mown = 0.0f;
for (int k = threadIdx.x; k < GfGf; k += blockDim.x) {
const float Pp = Pbuf[k];
if (Pp <= 0.0f) continue;
const int dx = k % Gf - Rf, dy = k / Gf - Rf;
const int x = cx + dx, y = cy + dy;
if (x < 0 || y < 0 || x >= p.W || y >= p.H) continue;
if (!valid(img[y * p.W + x])) continue;
const bool own = BraggOwnedBy(owner[y * p.W + x], i);
// Zeroing the profile here is how Exclude drops the pixel: the fit skips any cell with
// P <= 0 already, and P is not renormalised, so the fitted amplitude comes out on the
// scale of the WHOLE profile - the renormalisation is the estimator's own doing.
if (!own && p.overlap == 2) Pbuf[k] = 0.0f;
if (own) l_pown += Pp;
if ((float) (dx * dx + dy * dy) < p.r1_sq) {
l_mall += Pp;
if (own) l_mown += Pp;
}
}
atomicAdd(&s_pown, l_pown); atomicAdd(&s_mall, l_mall); atomicAdd(&s_mown, l_mown);
__syncthreads();
if (p.overlap == 1 && s_pown < p.minpk) {
if (threadIdx.x == 0) ok_o[i] = 0;
return;
}
}
if (threadIdx.x == 0) s_I = I_seed[i];
__syncthreads();
for (int iter = 0; iter < 4; ++iter) {
if (threadIdx.x == 0) { s_num = 0.0f; s_den = 0.0f; s_wsum = 0.0f; }
__syncthreads();
const float Ihere = s_I;
float l_num = 0.0f, l_den = 0.0f, l_wsum = 0.0f;
for (int k = threadIdx.x; k < GfGf; k += blockDim.x) {
const float Pp = Pbuf[k];
if (Pp <= 0.0f) continue;
const int x = cx + (k % Gf - Rf), y = cy + (k / Gf - Rf);
if (x < 0 || y < 0 || x >= p.W || y >= p.H) continue;
const int32_t px = img[y * p.W + x];
if (!valid(px)) continue;
const float v = fmaxf(B + Ihere * Pp, (float) WEIGHT_VARIANCE_MIN_FRACTION * B);
l_num += Pp * ((float) px - bkg) / v;
l_den += Pp * Pp / v;
l_wsum += Pp / v;
}
atomicAdd(&s_num, l_num); atomicAdd(&s_den, l_den); atomicAdd(&s_wsum, l_wsum);
__syncthreads();
if (threadIdx.x == 0 && s_den > 0.0f) s_I = s_num / s_den;
__syncthreads();
}
if (threadIdx.x == 0) {
if (s_den > 0.0f) {
// 1/s_den takes the background as exact; dI/dbkg = -s_wsum/s_den adds the background
// estimate's own error (see the CPU engine).
const float wr = s_wsum / s_den;
float I = s_I, sigma = sqrtf(1.0f / s_den + wr * wr * bkgvar_a[i]);
float var_bkg = fmaxf(0.0f, 1.0f / s_den - fmaxf(0.0f, I) + wr * wr * bkgvar_a[i]);
// Guard against profile-fit runaways (see the CPU engine): fall back to the summation seed
// when the profile result diverges from it. Exclude has taken pixels out of both, so the
// full-profile intensity is scaled back to the disk the seed read; nothing excluded gives
// exactly 1.
const float gs = (p.overlap == 2 && s_mall > 0.0f) ? s_mown / s_mall : 1.0f;
if (fabsf(I * gs - I_seed[i]) > (float) PROFILE_SUMMATION_MAX_NSIGMA * sigma_seed[i]) {
I = I_seed[i];
sigma = sigma_seed[i];
var_bkg = varbkg_seed[i];
}
I_o[i] = I; sigma_o[i] = sigma; varbkg_o[i] = var_bkg; ok_o[i] = 1;
} else ok_o[i] = 0;
}
}
} // namespace
BraggIntegrationEngineGPU::BraggIntegrationEngineGPU(const DiffractionExperiment &experiment,
std::shared_ptr<CudaStream> stream)
: BraggIntegrationEngine(experiment),
stream(std::move(stream)),
d_mask(npixel),
d_shell_grid(static_cast<size_t>(bragg_engine::N_SHELL) * GG),
d_global_grid(GG),
d_shell_P(static_cast<size_t>(bragg_engine::N_SHELL) * GG),
d_global_P(GG),
d_mom(MOM_STRIDE * (bragg_engine::N_SHELL + 1)),
d_sigma2_r(bragg_engine::N_SHELL + 1),
d_sigma2_t(bragg_engine::N_SHELL + 1),
d_shell_n(bragg_engine::N_SHELL),
d_global_n(1),
d_invd2(2) {
threads = 128;
// Fit profile grid: R for empirical / box, up to 3R (radially elongated) for the Gaussian.
const int max_Rf = empirical ? R : 3 * R;
const int max_Gf = 2 * max_Rf + 1;
fit_shared_bytes = static_cast<size_t>(max_Gf) * max_Gf * sizeof(float);
// Shared radial window of the background curve in boxsum. The stencil spans r0 +- the radial
// semi-axis of the outer ellipse, widest at the far corner of the detector, and the window
// covers twice that so a reflection anywhere on the detector fits.
rad_w = 2 * (static_cast<int>(std::ceil(r3 + BraggStencilGrow_px(static_cast<float>(r_max), stencil))) + 1) + 1;
boxsum_shared_bytes = static_cast<size_t>(rad_w) * (sizeof(float) + sizeof(int));
// The current device, not device 0: workers are pinned round-robin across GPUs, so device 0's
// shared-memory size can belong to a different card than the one these kernels launch on.
int device = 0;
cuda_err(cudaGetDevice(&device));
cudaDeviceProp prop{};
cuda_err(cudaGetDeviceProperties(&prop, device));
if (fit_shared_bytes > prop.sharedMemPerBlock)
throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
"BraggIntegrationEngineGPU: profile grid exceeds shared memory (r2 too large)");
// boxsum's dynamic window sits alongside its static shared arrays (the trimmed-mean buffer
// dominates them), so leave room for those rather than budgeting the whole block.
if (boxsum_shared_bytes + sizeof(int) * BKG_TRIM_MAX + 256 > prop.sharedMemPerBlock)
throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
"BraggIntegrationEngineGPU: background ring exceeds shared memory");
// Radial background curve: one bin per pixel of distance from the beam, out to the far corner.
// Allocated whenever the correction COULD run, so the auto mode can turn it on for an individual
// image; whether it runs for a given image is decided in Run() from the per-image bkg_radial.
if (overlap != OverlapMode::Off)
d_owner = CudaDevicePtr<uint32_t>(npixel);
if (bkg_radial || bkg_radial_auto) {
n_rad = static_cast<int>(std::ceil(r_max)) + 2;
d_rad_sum = CudaDevicePtr<float>(n_rad);
d_rad_cnt = CudaDevicePtr<int>(n_rad);
d_k_diff = CudaDevicePtr<float>(k_diff.size());
cuda_err(cudaMemcpy(d_k_diff, k_diff.data(), sizeof(float) * k_diff.size(),
cudaMemcpyHostToDevice));
}
}
void BraggIntegrationEngineGPU::EnsureCapacity(size_t n) {
if (n <= capacity)
return;
// Grow with slack. cudaMalloc/cudaFree take a device-wide lock in the driver, so growing to
// exactly n makes every image that sets a new reflection-count record stall all other workers.
const size_t new_capacity = std::max(n, capacity + capacity / 2);
d_px_x = CudaDevicePtr<float>(new_capacity);
d_px_y = CudaDevicePtr<float>(new_capacity);
d_d = CudaDevicePtr<float>(new_capacity);
d_cx = CudaDevicePtr<int>(new_capacity);
d_cy = CudaDevicePtr<int>(new_capacity);
d_I = CudaDevicePtr<float>(new_capacity);
d_sigma = CudaDevicePtr<float>(new_capacity);
d_bkg = CudaDevicePtr<float>(new_capacity);
d_bkg_var = CudaDevicePtr<float>(new_capacity);
d_var_bkg = CudaDevicePtr<float>(new_capacity);
d_isum = CudaDevicePtr<float>(new_capacity);
d_ninner = CudaDevicePtr<int>(new_capacity);
d_rbin = CudaDevicePtr<int>(new_capacity);
d_kbin = CudaDevicePtr<int>(new_capacity);
d_obs_x = CudaDevicePtr<float>(new_capacity);
d_obs_y = CudaDevicePtr<float>(new_capacity);
d_ok = CudaDevicePtr<uint8_t>(new_capacity);
d_strong = CudaDevicePtr<uint8_t>(new_capacity);
d_has_obs = CudaDevicePtr<uint8_t>(new_capacity);
h_px_x.resize(new_capacity); h_px_y.resize(new_capacity); h_d.resize(new_capacity);
h_I.resize(new_capacity); h_sigma.resize(new_capacity); h_bkg.resize(new_capacity);
h_var_bkg.resize(new_capacity);
h_obs_x.resize(new_capacity); h_obs_y.resize(new_capacity);
h_ok.resize(new_capacity); h_has_obs.resize(new_capacity);
capacity = new_capacity;
}
std::vector<Reflection> BraggIntegrationEngineGPU::Run(const ImagePreprocessorBuffer &image,
const std::vector<Reflection> &predicted,
size_t npredicted, int64_t image_number) {
std::vector<BraggFitResult> results(npredicted);
if (image.size() != npixel || npredicted == 0)
return Finalize(predicted, npredicted, results, image_number);
const int32_t *img = image.getGPUBuffer();
if (img == nullptr)
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
"BraggIntegrationEngineGPU: image buffer is not on the GPU");
EnsureCapacity(npredicted);
const int n = static_cast<int>(npredicted);
for (size_t i = 0; i < npredicted; ++i) {
h_px_x[i] = predicted[i].predicted_x;
h_px_y[i] = predicted[i].predicted_y;
h_d[i] = predicted[i].d;
}
cuda_err(cudaMemcpyAsync(d_px_x, h_px_x.data(), sizeof(float) * npredicted, cudaMemcpyHostToDevice, *stream));
cuda_err(cudaMemcpyAsync(d_px_y, h_px_y.data(), sizeof(float) * npredicted, cudaMemcpyHostToDevice, *stream));
cuda_err(cudaMemcpyAsync(d_d, h_d.data(), sizeof(float) * npredicted, cudaMemcpyHostToDevice, *stream));
BraggGpuParams p{
.W = static_cast<int>(xpixel), .H = static_cast<int>(ypixel),
.r1_sq = r1_sq, .r2 = r2, .r2_sq = r2_sq, .r3 = r3, .r3_sq = r3_sq,
.R = R, .G = G, .GG = GG,
.bkg_clip_nsigma = mode != IntegratorMode::BoxSum ? bkg_clip_nsigma : 0.0f,
.empirical = empirical ? 1 : 0,
.use_ellipse = use_ellipse ? 1 : 0,
.bw_sigma = static_cast<float>(bw_sigma), .c_radial = static_cast<float>(c_radial),
.F_px = static_cast<float>(F_px),
.beam_x = beam_x, .beam_y = beam_y,
.bkg_trim = bkg_trim, // effective trim fraction (0 for stills), set by the base ctor from settings
.stencil = stencil, .rad_w = rad_w, .n_kern = n_kern,
.overlap = overlap == OverlapMode::Reject ? 1 : (overlap == OverlapMode::Exclude ? 2 : 0),
.boxsum_reject = (overlap == OverlapMode::Reject && mode == IntegratorMode::BoxSum) ? 1 : 0,
.exclude_px = (overlap == OverlapMode::Exclude && mode != IntegratorMode::BoxSum) ? 1 : 0,
.claim_sq = claim * claim, .inv_claim = inv_claim, .minpk = overlap_min_peak,
};
// Whether the radial correction runs for THIS image. n_rad only says the buffers exist - under
// the auto mode they are allocated for every image and used for the ones the ice score selects.
const int rad_n = bkg_radial ? n_rad : 0;
// Pass A: reset accumulators, mask, then box-sum.
cuda_err(cudaMemsetAsync(d_mask, 0, npixel, *stream));
if (p.overlap)
cuda_err(cudaMemsetAsync(d_owner, 0xff, sizeof(uint32_t) * npixel, *stream)); // BRAGG_OWNER_NONE
if (rad_n > 0) {
cuda_err(cudaMemsetAsync(d_rad_sum, 0, sizeof(float) * n_rad, *stream));
cuda_err(cudaMemsetAsync(d_rad_cnt, 0, sizeof(int) * n_rad, *stream));
}
reset<<<32, 256, 0, *stream>>>(d_shell_grid, d_global_grid, d_mom, d_shell_n, d_global_n, d_invd2, GG);
mark_mask<<<n, threads, 0, *stream>>>(d_px_x, d_px_y, d_mask, d_owner, p, n);
boxsum<<<n, threads, boxsum_shared_bytes, *stream>>>(d_px_x, d_px_y, d_d, img, d_mask, d_owner, p, n,
d_cx, d_cy, d_I, d_sigma, d_bkg, d_bkg_var, d_var_bkg, d_obs_x, d_obs_y,
d_ok, d_strong, d_has_obs, d_invd2,
d_isum, d_ninner, d_rbin, d_kbin,
d_rad_sum, d_rad_cnt, rad_n);
// Correct the flat annulus background for the curvature of the radial background before anything
// downstream (profile fit, variance) reads it.
if (rad_n > 0)
radial_correct<<<(n + threads - 1) / threads, threads, 0, *stream>>>(
d_rad_sum, d_rad_cnt, rad_n, d_k_diff, k_len, k_off,
d_isum, d_ninner, d_rbin, d_kbin, d_ok, d_bkg, d_I, n);
if (mode != IntegratorMode::BoxSum) {
// Pass B: learn (shell computed inline) -> build -> fit.
learn_profile<<<n, threads, 0, *stream>>>(img, d_owner, d_px_x, d_px_y, d_cx, d_cy, d_d, d_invd2,
d_I, d_bkg, d_ok, d_strong,
d_shell_grid, d_global_grid, d_mom, d_shell_n, d_global_n, p, n);
build_profiles<<<bragg_engine::N_SHELL + 1, threads, 0, *stream>>>(
d_shell_grid, d_global_grid, d_mom, d_shell_n, d_global_n,
d_shell_P, d_global_P, d_sigma2_r, d_sigma2_t, p);
fit<<<n, threads, fit_shared_bytes, *stream>>>(img, d_owner, d_px_x, d_px_y, d_cx, d_cy, d_d, d_invd2,
d_I, d_sigma, d_bkg, d_bkg_var, d_var_bkg, d_ok, d_shell_P, d_global_P,
d_sigma2_r, d_sigma2_t, d_shell_n,
d_I, d_sigma, d_var_bkg, d_ok, p, n);
}
cuda_err(cudaMemcpyAsync(h_I.data(), d_I, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
cuda_err(cudaMemcpyAsync(h_sigma.data(), d_sigma, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
cuda_err(cudaMemcpyAsync(h_bkg.data(), d_bkg, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
cuda_err(cudaMemcpyAsync(h_var_bkg.data(), d_var_bkg, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
cuda_err(cudaMemcpyAsync(h_ok.data(), d_ok, sizeof(uint8_t) * npredicted, cudaMemcpyDeviceToHost, *stream));
// Pass A always fills the box-sum centroid (observed spot position), so copy it back in all modes -
// post-refinement uses it as the observed position (beam-centre / distance).
cuda_err(cudaMemcpyAsync(h_obs_x.data(), d_obs_x, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
cuda_err(cudaMemcpyAsync(h_obs_y.data(), d_obs_y, sizeof(float) * npredicted, cudaMemcpyDeviceToHost, *stream));
cuda_err(cudaMemcpyAsync(h_has_obs.data(), d_has_obs, sizeof(uint8_t) * npredicted, cudaMemcpyDeviceToHost, *stream));
cuda_err(cudaStreamSynchronize(*stream));
for (size_t i = 0; i < npredicted; ++i) {
if (!h_ok[i]) continue;
results[i].I = h_I[i];
results[i].sigma = h_sigma[i];
results[i].bkg = h_bkg[i];
results[i].var_bkg = h_var_bkg[i];
results[i].ok = true;
if (h_has_obs[i]) {
results[i].observed_x = h_obs_x[i];
results[i].observed_y = h_obs_y[i];
results[i].has_observed = true;
}
}
return Finalize(predicted, npredicted, results, image_number);
}