mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 68 Wiki Activity

Take the double precision out of the box integrator's inner loops

Browse Source 
boxsum summed its ring background in double and compared each ring pixel
against a double threshold. The pixels are integers: a sum of at most a
thousand int32 values is exact in a 64-bit integer AND exact in a double, so
the two agree bit for bit, and comparing an integer against the floor of the
threshold accepts exactly the same pixels as comparing it against the threshold
itself. Both loops now do integer arithmetic.

That was 39% of the card's double-precision pipe on the development machine and
about three quarters of it on the production one, where the double rate is
unchanged from Turing while the single rate has doubled - so this is worth more
there than here.

Alongside it, three things in the combine kernel. rr_nusable was computed by a
whole extra walk over every observation and then never downloaded or read by
anything. sum_wb and sum_cwb have no F in them, so they are the same in all
three reweights and only the last round's values are ever used - two thirds of
them were two divisions each, discarded. And CombineParams was the one
parameter struct in the file without __restrict__, so the compiler could not
assume the observation arrays and the freshly allocated fulls arrays were
distinct.

Measured on a crystal with 66 million partial observations: boxsum 12.2 s ->
8.3 s, the combine kernel 8.0 s -> 7.6 s, whole crystal 1m17s -> 1m12s. Battery
15m32s -> 9m59s. Same space group on all 24 crystals, none failed.

Two things measured and NOT kept, recorded so they are not tried again: sorting
the raw-hkl runs by length so a warp holds runs of similar length - it trades
away the locality of neighbouring runs in the permutation and came out slower
(7.6 s -> 8.8 s); and page-locking the integrator's host staging arrays
individually - eleven separate registrations of small heap allocations overlap
on shared pages and the driver refuses them.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
jungfrau
2026-08-15 23:22:29 -04:00
co-authored by Claude Opus 5
parent 3aa239fce8
commit 09cb9e2be9
2 changed files with 38 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/bragg_integration/BraggIntegrationEngineGPU.cu
+18 -12
View File
@@ -107,9 +107,11 @@ __global__ void boxsum(const float *px_x, const float *px_y, const float *dd,
__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__ unsigned long long s_bkgsum;
__shared__ int s_accept, s_full;
__shared__ double s_bkg, s_thr, s_clipsum;
__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
@@ -123,7 +125,7 @@ __global__ void boxsum(const float *px_x, const float *px_y, const float *dd,
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_ninner = 0; s_ninner_valid = 0; s_nbkg = 0; s_bkgsum = 0;
s_ndisk = 0; s_nown = 0;
}
__syncthreads();
@@ -141,7 +143,7 @@ __global__ void boxsum(const float *px_x, const float *px_y, const float *dd,
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;
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);
@@ -162,7 +164,7 @@ __global__ void boxsum(const float *px_x, const float *px_y, const float *dd,
if (mask[y * p.W + x]) continue;
const int32_t px = img[y * p.W + x];
if (!valid(px)) continue;
l_bkg += (double) px; ++l_nb;
l_bkg += px; ++l_nb;
}
}
WARP_ATOMIC_ADD(s_Isum, (unsigned long long) l_Isum);
@@ -171,7 +173,7 @@ __global__ void boxsum(const float *px_x, const float *px_y, const float *dd,
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, l_bkg);
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();
@@ -184,9 +186,13 @@ __global__ void boxsum(const float *px_x, const float *px_y, const float *dd,
// (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 ? (s_bkgsum / (double) s_nbkg) : 0.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));
s_clipsum = 0.0; s_clipn = 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();
@@ -240,7 +246,7 @@ __global__ void boxsum(const float *px_x, const float *px_y, const float *dd,
// 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;
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);
@@ -248,7 +254,7 @@ __global__ void boxsum(const float *px_x, const float *px_y, const float *dd,
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) {
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
@@ -261,7 +267,7 @@ __global__ void boxsum(const float *px_x, const float *px_y, const float *dd,
}
}
}
WARP_ATOMIC_ADD(s_clipsum, c_l); WARP_ATOMIC_ADD(s_clipn, cn_l);
WARP_ATOMIC_ADD(s_clipsum, (unsigned long long) c_l); WARP_ATOMIC_ADD(s_clipn, cn_l);
}
__syncthreads();
@@ -290,7 +296,7 @@ __global__ void boxsum(const float *px_x, const float *px_y, const float *dd,
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; }
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.
image_analysis/scale_merge/RotationScaleMergeGPU.cu
+20 -16
View File
@@ -266,11 +266,15 @@ namespace {
struct CombineParams {
int n_runs;
double min_partiality, capture_uncertainty_coeff, min_captured_fraction;
const float *I, *sigma, *corr, *partiality, *bkg, *var_bkg, *image_number, *d, *px, *py;
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 float *__restrict__ I, *__restrict__ sigma, *__restrict__ corr,
*__restrict__ partiality, *__restrict__ bkg, *__restrict__ var_bkg,
*__restrict__ image_number, *__restrict__ d, *__restrict__ px, *__restrict__ py;
const int32_t *__restrict__ frame;
const uint8_t *__restrict__ on_ice;
const int32_t *__restrict__ perm, *__restrict__ rr_start, *__restrict__ rr_count,
*__restrict__ rr_h, *__restrict__ rr_k, *__restrict__ rr_l,
*__restrict__ rr_group;
int32_t *rr_nevents; // count pass output
const int32_t *rr_offset; // emit pass: per-run base offset into the fulls arrays
int32_t *f_h, *f_k, *f_l, *f_frame, *f_group;
float *f_I, *f_sigma, *f_d, *f_img, *f_px, *f_py, *f_var_bkg, *f_var_per_I;
@@ -371,8 +375,13 @@ namespace {
const double w = 1.0 / var;
sum_w += w;
sum_wI += w * I_corr;
sum_wb += 1.0 / a_var;
sum_cwb += corr / (a_var * a_var);
// a_var has no F in it, so these two are the same in every reweight and only
// the last round's values are ever read. Two thirds of them were two divisions
// each, thrown away.
if (iter == 2) {
sum_wb += 1.0 / a_var;
sum_cwb += corr / (a_var * a_var);
}
}
F = sum_wI / sum_w;
}
@@ -406,13 +415,8 @@ namespace {
++n_emit;
}
if (!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>
@@ -643,7 +647,7 @@ struct RotationScaleMergeGPU::Impl {
CudaDevicePtr<float> bkg, var_bkg, image_number, d_obs, px_obs, py_obs;
int n_runs = 0, n_perm = 0;
CudaDevicePtr<int32_t> perm, rr_start, rr_count, rr_h, rr_k, rr_l, rr_group;
CudaDevicePtr<int32_t> rr_nevents, rr_nusable, rr_offset;
CudaDevicePtr<int32_t> rr_nevents, 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;
@@ -965,9 +969,9 @@ void RotationScaleMergeGPU::SetRawRuns(int n_runs, int n_perm, const int32_t *pe
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));
}
@@ -989,7 +993,7 @@ int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_parti
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();
p.rr_nevents = d.rr_nevents.get();
const int blocks = std::min(65535, (d.n_runs + BLK - 1) / BLK);