jungfrauandClaude Opus 5
16639e9de2
Every accumulator in this file that could be an integer already is one, and the spot finder's reduce_rings_shared says why: a preprocessed pixel is an exact int32, integer addition is associative, and a threshold that moves in its last bits between runs flips every pixel sitting on it. Four accumulators here were still floats, and they reach the intensity rather than a diagnostic. * The radial background curve. s_radv and rad_sum sum int32 pixel values, so int64 is not an approximation of the old sum, it IS the old sum - and the curve is subtracted from every reflection's background. * The learned profile grid and its second moments, which are sums of (px - bkg) / I over every strong reflection of the frame. There is no exact integer form, so these are fixed point at 2^20: a quantum of 1e-6 of one I-normalised pixel, far below the Poisson noise of the pixel it came from, and some five orders of headroom inside a signed 64-bit accumulator. * The normalisation total in build_profiles, which divides every cell of the profile - 128 lanes on one address, in arrival order. Now summed as integers, exactly, from the grid it normalises. * The fit's own reductions, s_num and s_den among them, which ARE the fitted intensity. These stay float, so fixed point would be a real precision trade over an unbounded range; instead each warp leaves its total in a slot of its own and every thread adds the slots up by warp index. WARP_ATOMIC_ADD is order-independent for the integer accumulators it was written for and not for these, which is what block_sum is for. With the prediction ordering of the previous commit, a run is now reproducible: the same command on the same images writes byte-identical .hkl and .mtz, at -N 1 and at -N 48, on a 16 Mpx rotation set and on a large-cell one. Before, all four differed. The battery is unchanged where it was ever stable: space group identical on all 24 crystals, reflection count on 20, R_meas on 22. The two that move are the two the battery has always seen move between runs of an unchanged binary - which is the point, since they stop moving now. Total 8m02s against 7m55s, inside the noise of per-crystal times quantised to a second. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_011n8riB6X59oRjkrSHzNPAU
943 lines
50 KiB
Plaintext
943 lines
50 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; // least readable/clean profile fraction that is kept (XDS MINPK)
|
|
int partial_ok; // keep a signal disk with unreadable pixels: anything but a box sum
|
|
float peak_frac; // most of the profile's peak value an unreadable pixel may carry
|
|
};
|
|
|
|
__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));
|
|
}
|
|
}
|
|
|
|
// Sum one value across the warp so a single lane does the shared-memory atomic.
|
|
//
|
|
// Every block-wide accumulation here has all 128 lanes targeting one address, and a shared float or
|
|
// 64-bit atomicAdd has no native instruction on Turing OR Ada - it compiles to a compare-and-swap
|
|
// retry loop, so those 128 lanes serialise into 128 retries. Reducing within the warp first leaves
|
|
// four atomics per block instead of 128. Every call site below sits after a loop that all threads
|
|
// reach, so the full-warp mask is the right one.
|
|
//
|
|
// The integer sums are unchanged by this: addition is associative. The float and double ones change
|
|
// in their last bits, and become MORE reproducible - a fixed shuffle tree replaces whatever order
|
|
// the atomics happened to arrive in.
|
|
template <typename T>
|
|
__device__ __forceinline__ T warp_sum(T v) {
|
|
#pragma unroll
|
|
for (int off = 16; off > 0; off >>= 1)
|
|
v += __shfl_down_sync(0xffffffffu, v, off);
|
|
return v;
|
|
}
|
|
|
|
#define WARP_ATOMIC_ADD(dst, val) do { \
|
|
const auto _w = warp_sum(val); \
|
|
if ((threadIdx.x & 31u) == 0u) atomicAdd(&(dst), _w); \
|
|
} while (0)
|
|
|
|
// Sum one float across the whole block, in a fixed order. warp_sum gives each warp its own total
|
|
// with a fixed shuffle tree; the warps then leave those totals in slots of their own and every
|
|
// thread adds the slots up by warp index. WARP_ATOMIC_ADD is the right tool for an INTEGER
|
|
// accumulator, where addition is associative, but on a float it adds the same handful of numbers in
|
|
// whatever order the warps happened to arrive - and s_num / s_den below are the numerator and
|
|
// denominator of the fitted intensity, so that order reached the answer. Every caller is reached by
|
|
// the whole block, as the WARP_ATOMIC_ADDs it replaces already were, and each concurrent quantity
|
|
// needs slots of its own.
|
|
constexpr int BRAGG_WARPS_MAX = 8; // the kernels launch 128 threads; room to 256
|
|
__device__ __forceinline__ float block_sum(float v, float *slots) {
|
|
const float w = warp_sum(v);
|
|
if ((threadIdx.x & 31u) == 0u) slots[threadIdx.x >> 5] = w;
|
|
__syncthreads();
|
|
float s = 0.0f;
|
|
const int nwarps = (blockDim.x + 31) >> 5;
|
|
for (int j = 0; j < nwarps; ++j) s += slots[j];
|
|
return s;
|
|
}
|
|
|
|
// Fixed-point scale for the learned profile and its moments. Those accumulators are sums of
|
|
// v = (px - bkg) / I over every strong reflection of the frame, and a float atomicAdd adds them in
|
|
// whatever order the blocks happen to arrive - so the profile, and every intensity fitted through
|
|
// it, moved between two runs of the same command on the same image. In fixed point the sum is
|
|
// integer addition, which is associative, exactly as the ring statistics and the box sums above
|
|
// already are. 2^20 puts the quantum at 1e-6 of one I-normalised pixel, far below the Poisson noise
|
|
// of the pixel it came from, and leaves the widest sum this kernel can build - every pixel of every
|
|
// reflection, times the r1^2 moment arm - some five orders inside a signed 64-bit accumulator.
|
|
constexpr double PROFILE_FIXED = 1048576.0; // 2^20
|
|
|
|
// Signed fixed-point value carried as two's complement, so the accumulators can use the unsigned
|
|
// 64-bit atomicAdd (the only 64-bit one CUDA offers) - the same convention as s_Isum above.
|
|
__device__ __forceinline__ unsigned long long profile_fixed(float v) {
|
|
return (unsigned long long) (long long) llrintf(v * (float) PROFILE_FIXED);
|
|
}
|
|
__device__ __forceinline__ float profile_float(unsigned long long v) {
|
|
return (float) ((double) (long long) v / PROFILE_FIXED);
|
|
}
|
|
|
|
// --- 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,
|
|
unsigned long long *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__ unsigned long long s_bkgsum;
|
|
__shared__ int s_accept, s_full;
|
|
__shared__ double s_bkg, s_thr;
|
|
__shared__ unsigned long long s_clipsum;
|
|
__shared__ long long s_thr_i;
|
|
__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.
|
|
// The radial sums are INTEGERS, for the reason reduce_rings_shared gives in the spot finder: a
|
|
// preprocessed pixel is an exact int32, so the sum is exact in 64 bits and integer addition is
|
|
// associative - the radial background curve is then the same whatever order the blocks arrive in.
|
|
// With float accumulators it moved in its last bits between runs of the same command, and it is
|
|
// subtracted from every reflection's background.
|
|
extern __shared__ unsigned long long s_rad[];
|
|
unsigned long long *s_radv = s_rad; // signed value carried as two's complement
|
|
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;
|
|
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;
|
|
long long l_bkg = 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);
|
|
// The window is the bounding box of the outer ellipse, so nearly half of it is neither the
|
|
// signal disk nor the background ring. Reading the pixel only once it is known to be wanted
|
|
// keeps those slots from fetching a cache line for nothing.
|
|
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;
|
|
const int32_t px = img[y * p.W + x];
|
|
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;
|
|
const int32_t px = img[y * p.W + x];
|
|
if (!valid(px)) continue;
|
|
l_bkg += px; ++l_nb;
|
|
}
|
|
}
|
|
WARP_ATOMIC_ADD(s_Isum, (unsigned long long) l_Isum);
|
|
WARP_ATOMIC_ADD(s_Ix, (unsigned long long) l_Ix);
|
|
WARP_ATOMIC_ADD(s_Iy, (unsigned long long) l_Iy);
|
|
WARP_ATOMIC_ADD(s_ninner, l_ni);
|
|
WARP_ATOMIC_ADD(s_ninner_valid, l_niv);
|
|
WARP_ATOMIC_ADD(s_nbkg, l_nb);
|
|
WARP_ATOMIC_ADD(s_bkgsum, (unsigned long long) l_bkg);
|
|
WARP_ATOMIC_ADD(s_ndisk, l_nd);
|
|
WARP_ATOMIC_ADD(s_nown, l_no);
|
|
__syncthreads();
|
|
|
|
for (int t = threadIdx.x; t < p.rad_w; t += blockDim.x) { s_radv[t] = 0; s_radn[t] = 0; }
|
|
if (threadIdx.x == 0) s_radbase = (int) lroundf(s_r0) - p.rad_w / 2;
|
|
if (threadIdx.x == 0) {
|
|
// A hole in the signal disk no longer discards the reflection in the profile modes - the fit
|
|
// renormalises to the pixels it can read and Pass B cuts on how much of the profile survived
|
|
// (XDS's MINPK). A box sum has no profile to renormalise with. See the CPU engine.
|
|
s_full = (s_ninner_valid == s_ninner) ? 1 : 0;
|
|
s_accept = ((s_full || p.partial_ok) && s_nbkg > 5) ? 1 : 0;
|
|
s_bkg = s_accept ? ((double) (long long) s_bkgsum / (double) s_nbkg) : 0.0;
|
|
s_thr = s_bkg + (double) p.bkg_clip_nsigma * sqrt(fmax(s_bkg, 1.0));
|
|
// The pixel is an integer, so comparing it against the floor of the threshold accepts
|
|
// exactly the same set - and does it with an integer compare instead of a widening
|
|
// conversion and a double comparison, per ring pixel.
|
|
s_thr_i = (long long) floor(s_thr);
|
|
s_clipsum = 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) {
|
|
long long c_l = 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 ((long long) px <= s_thr_i) {
|
|
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], (unsigned long long) (long long) px); // shared, not global
|
|
atomicAdd(&s_radn[idx], 1);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
WARP_ATOMIC_ADD(s_clipsum, (unsigned long long) c_l); WARP_ATOMIC_ADD(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]); // two's complement, summed as unsigned
|
|
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 = (double) (long long) s_clipsum / (double) s_clipn; n_bkg_used = s_clipn; }
|
|
const long long Isum = (long long) s_Isum;
|
|
// The sum is over the pixels actually READ, so that is the count the background is subtracted
|
|
// with; with nothing missing it is the whole disk, exactly as before.
|
|
const double I = (double) Isum - (double) s_ninner_valid * 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_valid * (double) s_ninner_valid * 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;
|
|
// A disk with a hole gives a centroid pulled away from it; see the CPU engine.
|
|
hasobs = s_full ? 1 : 0;
|
|
}
|
|
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_valid * bkg + var_bkg_term);
|
|
isum_o[i] = (float) Isum;
|
|
ninner_o[i] = s_ninner_valid;
|
|
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;
|
|
// The profile and its resolution shells are learned from COMPLETE reflections only; see the CPU
|
|
// engine.
|
|
strong_o[i] = (s_full && sigma > 0.0 && I / sigma >= STRONG_I_OVER_SIGMA) ? 1 : 0;
|
|
|
|
const float d = dd[i];
|
|
if (s_full && 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(unsigned long long *shell_grid, unsigned long long *global_grid,
|
|
unsigned long long *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;
|
|
for (int k = blockIdx.x * blockDim.x + threadIdx.x; k < GG; k += blockDim.x * gridDim.x)
|
|
global_grid[k] = 0;
|
|
if (blockIdx.x == 0 && threadIdx.x < MOM_STRIDE * (N_SHELL + 1)) mom[threadIdx.x] = 0;
|
|
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,
|
|
unsigned long long *shell_grid, unsigned long long *global_grid,
|
|
unsigned long long *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];
|
|
// The shell belongs to the reflection, not the pixel, so one thread works it out for the whole
|
|
// block: it is two software double-precision divisions, and all 128 threads were doing both.
|
|
__shared__ int s_sh;
|
|
if (threadIdx.x == 0) s_sh = compute_shell(dd[i], invd2mm);
|
|
__syncthreads();
|
|
const int sh = s_sh;
|
|
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;
|
|
unsigned long long *sg = shell_grid + (size_t) sh * p.GG;
|
|
__shared__ unsigned long long s_rad, s_tan, s_w;
|
|
if (threadIdx.x == 0) { s_rad = 0; s_tan = 0; s_w = 0; }
|
|
__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;
|
|
const unsigned long long vq = profile_fixed(v);
|
|
atomicAdd(&sg[k], vq);
|
|
atomicAdd(&global_grid[k], vq);
|
|
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 warp instead of one per thread. A thread's own l_* sum is over a fixed set of k, so
|
|
// it is the same every run; quantising here makes everything above it associative as well.
|
|
WARP_ATOMIC_ADD(s_rad, profile_fixed(l_rad));
|
|
WARP_ATOMIC_ADD(s_tan, profile_fixed(l_tan));
|
|
WARP_ATOMIC_ADD(s_w, profile_fixed(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 unsigned long long *shell_grid, const unsigned long long *global_grid,
|
|
const unsigned long long *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 unsigned long long *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 = profile_float(mom[MOM_STRIDE * b + 2]);
|
|
sigma2_r[b] = w > 0.0f ? fmaxf(0.25f, profile_float(mom[MOM_STRIDE * b + 0]) / w) : 1.0f;
|
|
sigma2_t[b] = w > 0.0f ? fmaxf(0.25f, profile_float(mom[MOM_STRIDE * b + 1]) / w) : 1.0f;
|
|
}
|
|
if (!p.empirical) return;
|
|
|
|
// The normalisation total is summed in fixed point as well. It divides every cell of the profile,
|
|
// so a float atomicAdd here - 128 lanes on one address, in arrival order - reached every intensity
|
|
// the profile fits, which is most of what was left moving between two runs of the same command.
|
|
__shared__ unsigned long long s_sum_q;
|
|
if (threadIdx.x == 0) s_sum_q = 0;
|
|
__syncthreads();
|
|
unsigned long long l_sum = 0;
|
|
for (int k = threadIdx.x; k < p.GG; k += blockDim.x) {
|
|
const long long gq = (long long) grid[k];
|
|
const unsigned long long g = gq > 0 ? (unsigned long long) gq : 0ull; // a profile has to be non-negative
|
|
P[k] = profile_float(g);
|
|
l_sum += g;
|
|
}
|
|
WARP_ATOMIC_ADD(s_sum_q, l_sum);
|
|
__syncthreads();
|
|
const float s_sum = profile_float(s_sum_q);
|
|
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 unsigned long long *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 ? (float) (long long) 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). The
|
|
// reweighting is the Kabsch/Otwinowski iteration: Kabsch, Acta Cryst D66, 133-144 (2010);
|
|
// Otwinowski & Minor, Methods Enzymol 276, 307-326 (1997).
|
|
// 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_I;
|
|
__shared__ float s_slot[5][BRAGG_WARPS_MAX]; // per-warp partials, see block_sum
|
|
// Max reductions. The profile is non-negative, and for non-negative floats the IEEE bit pattern
|
|
// orders exactly as the value does, so an integer atomicMax on that pattern is an exact float max.
|
|
__shared__ int s_ppeak_i, s_plost_i;
|
|
__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];
|
|
// As in learn_profile: one thread per block, not one per thread. Two double-precision divisions
|
|
// on a card whose double throughput is a sixty-fourth of its single is worth doing once.
|
|
__shared__ int s_sh;
|
|
if (threadIdx.x == 0) s_sh = compute_shell(dd[i], invd2mm);
|
|
__syncthreads();
|
|
const int sh = s_sh;
|
|
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; }
|
|
__syncthreads();
|
|
const float fx = px_x[i] - cx, fy = px_y[i] - cy;
|
|
// The two widths are the same for every cell of this reflection, so their reciprocals are
|
|
// taken once rather than per cell.
|
|
const float inv_2s2r = 0.5f / s2r, inv_2s2t = 0.5f / s2t;
|
|
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 * inv_2s2r - tn * tn * inv_2s2t);
|
|
Pbuf[k] = g; l_gs += g;
|
|
}
|
|
const float gsum = block_sum(l_gs, s_slot[0]);
|
|
// Normalising by a reciprocal rather than dividing per cell: the divisor is the same for
|
|
// every cell of the reflection, and a division here is thirteen instructions of a
|
|
// twenty-instruction loop.
|
|
const float inv_gs = __frcp_rn(gsum);
|
|
for (int k = threadIdx.x; k < Gf * Gf; k += blockDim.x) Pbuf[k] *= inv_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 the fit can actually see: s_pvalid over the whole grid against
|
|
// s_pgrid, the mass that falls on the detector at all (XDS's MINPK), s_pown the same over
|
|
// neighbour-owned pixels (what Reject cuts on), and s_mread / s_mall over the r1 disk alone,
|
|
// which is what the summation seed the runaway guard compares against actually summed. See the
|
|
// CPU engine.
|
|
if (threadIdx.x == 0) { s_ppeak_i = 0; s_plost_i = 0; }
|
|
__syncthreads();
|
|
float l_pgrid = 0.0f, l_pvalid = 0.0f, l_pown = 0.0f, l_mall = 0.0f, l_mread = 0.0f;
|
|
float l_ppeak = 0.0f, l_plost = 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;
|
|
const bool in_disk = (float) (dx * dx + dy * dy) < p.r1_sq;
|
|
l_pgrid += Pp;
|
|
l_ppeak = fmaxf(l_ppeak, Pp);
|
|
if (in_disk) l_mall += Pp;
|
|
if (!valid(img[y * p.W + x])) {
|
|
l_plost = fmaxf(l_plost, Pp);
|
|
continue;
|
|
}
|
|
l_pvalid += Pp;
|
|
const bool own = !p.overlap || 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 (in_disk && (own || p.overlap != 2)) l_mread += Pp;
|
|
}
|
|
atomicMax(&s_ppeak_i, __float_as_int(l_ppeak)); atomicMax(&s_plost_i, __float_as_int(l_plost));
|
|
const float s_pgrid = block_sum(l_pgrid, s_slot[0]);
|
|
const float s_pvalid = block_sum(l_pvalid, s_slot[1]);
|
|
const float s_pown = block_sum(l_pown, s_slot[2]);
|
|
const float s_mall = block_sum(l_mall, s_slot[3]);
|
|
const float s_mread = block_sum(l_mread, s_slot[4]);
|
|
if (s_pvalid < p.minpk * s_pgrid) {
|
|
if (threadIdx.x == 0) ok_o[i] = 0;
|
|
return;
|
|
}
|
|
// A hole in the profile's PEAK is a different defect from a hole in its wings. See the CPU engine.
|
|
if (__int_as_float(s_plost_i) > p.peak_frac * __int_as_float(s_ppeak_i)) {
|
|
if (threadIdx.x == 0) ok_o[i] = 0;
|
|
return;
|
|
}
|
|
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();
|
|
|
|
float s_num = 0.0f, s_den = 0.0f, s_wsum = 0.0f;
|
|
for (int iter = 0; iter < 4; ++iter) {
|
|
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);
|
|
// One reciprocal for the three weights. Written as three divisions the compiler emits
|
|
// the whole correctly-rounded sequence three times over - the same estimate, the same
|
|
// refinement, the same check - for a weight that only has to be a weight.
|
|
const float iv = __frcp_rn(v);
|
|
l_num += Pp * ((float) px - bkg) * iv;
|
|
l_den += Pp * Pp * iv;
|
|
l_wsum += Pp * iv;
|
|
}
|
|
s_num = block_sum(l_num, s_slot[0]);
|
|
s_den = block_sum(l_den, s_slot[1]);
|
|
s_wsum = block_sum(l_wsum, s_slot[2]);
|
|
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]);
|
|
// The signal part to remove is I itself, not max(0, I) - see the CPU engine.
|
|
float var_bkg = fmaxf(0.0f, 1.0f / s_den - 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. Pixels missing from both - excluded to a
|
|
// neighbour, or unreadable - scale the full-profile intensity back to the disk the seed
|
|
// read; nothing dropped gives exactly 1.
|
|
const float gs = s_mall > 0.0f ? s_mread / 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(unsigned long long) + 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<unsigned long long>(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,
|
|
.partial_ok = mode != IntegratorMode::BoxSum ? 1 : 0,
|
|
.peak_frac = static_cast<float>(MINPK_MAX_MISSING_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(unsigned long long) * 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);
|
|
}
|