From 0109676bdfbdb1dfed9bd92f195e8395122975f4 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Mon, 24 Aug 2026 21:49:57 +0200 Subject: [PATCH] Keep the constants out of the dual numbers, and clear only the boxes Two costs in the per-image loop, each measured before it was touched. Geometry refinement is the largest item in that loop - about half to two thirds of its processor time on the datasets where the loop matters, and all of it on the host. It runs three solves per image, and each one spends four fifths of itself inside the solver at barely two iterations: the cost is not convergence, it is what every residual evaluation does. The residual carried the blocks it does not refine as dual numbers, so each evaluation recomputed the two detector rotations, the whole orthogonalisation matrix, three cross products and the cell volume - all of them constant for the image - through the derivative machinery, several million times per run. Split the observed and predicted sides so the un-refined blocks pass as plain doubles, evaluate the cell side once when the functor is built, and let the rotator take a point whose type differs from the angle's. A dual number times a double is a dual number times a dual number whose derivatives are zero, so the arithmetic is the same one with the zeros removed. Integration cleared the owner and mask images for the whole frame before every image. On a large detector that is more than three hundred megabytes of writes to reset pixels of which about one in twenty-five is ever marked, and it cost most of what the integration kernels themselves cost. The marking kernel gained an unmarking mode - one kernel, so the two cannot drift apart - and the engine clears whichever way is cheaper for the frame in front of it, with a flag to force the full clear the first time and after anything threw. The size test is not decoration: without it, clearing box by box is slower than the memset on a small detector with many predictions, which is what the measurement said before it was added. Faster on thirteen of thirteen matched pairs: refinement by a quarter to a third, whole-run wall by one to eight per cent depending on how much of the run is the loop. The two changes pay in opposite regimes - refinement where the loop is processor-bound, the clear where the detector is large enough for the card to be the constraint. Every reflection file over seven datasets is byte-identical, and the solver did not merely land in the same place: it took the same path, agreeing digit for digit on iteration, residual and Jacobian evaluation counts. A new test runs two mismatched frames through one engine and compares against a fresh one, which is what a mark left behind would break. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_016NNnL26LAvruQ9eLUUWvrJ --- docs/CHANGELOG.md | 1 + .../BraggIntegrationEngineGPU.cu | 55 +++- .../BraggIntegrationEngineGPU.h | 7 + image_analysis/geom_refinement/XtalResidual.h | 256 +++++++++++------- tests/BraggIntegrationEngineGPUTest.cpp | 55 ++++ 5 files changed, 263 insertions(+), 111 deletions(-) diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index 9f77680e..46199d68 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -19,6 +19,7 @@ This is an UNSTABLE release. It includes many experimental features, as well as * rugnux: an image integrated in pyFAI through the `.poni` file written by `--mode calibration` now comes out with the correct azimuth. Radial integration is unchanged. * rugnux: the `.poni` file declares pyFAI's `orientation`, which needs pyFAI 2024.01 or newer. * The build resolves a single Eigen for the whole project, and refuses to configure if Ceres picks up a different one; a build that mixed two Eigen versions was undefined behaviour and crashed at -O2. +* rugnux: per-image geometry refinement and integration are faster, with identical output. * rugnux: scaling and merging are faster, with identical output. * The per-image resolution estimate now predicts the resolution the merged data reach, rather than reporting the highest-resolution spot found; rugnux reports the run's value as `SPOT_RESOLUTION_ESTIMATE` in its report. * rugnux: fixing the space group with `-S` no longer prevents the lattice from being found; the group is applied to scaling and merging rather than to the indexing search. diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineGPU.cu b/image_analysis/bragg_integration/BraggIntegrationEngineGPU.cu index 6659d8c3..7f9a8ccc 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineGPU.cu +++ b/image_analysis/bragg_integration/BraggIntegrationEngineGPU.cu @@ -48,8 +48,12 @@ constexpr int MOM_STRIDE = 3; // 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. --- +// Mark the per-pixel reflection mask and the overlap owner, or - with clear set - put the very same +// pixels back the way they were. The two directions share one kernel because they have to visit +// exactly the same set: what is written here is what has to be unwritten, and a second kernel that +// recomputed the boxes for itself would be free to drift from this one. __global__ void mark_mask(const float *px_x, const float *px_y, uint8_t *mask, uint32_t *owner, - BraggGpuParams p, int n) { + BraggGpuParams p, int n, int clear) { const int i = blockIdx.x; if (i >= n) return; const float cx = px_x[i], cy = px_y[i]; @@ -63,9 +67,15 @@ __global__ void mark_mask(const float *px_x, const float *px_y, uint8_t *mask, u 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)); + if (d.inner < p.r2_sq) mask[y * p.W + x] = clear ? (uint8_t) 0 : (uint8_t) 1; + if (p.overlap && d.signal < p.claim_sq) { + // Clearing writes one fixed value, so the blocks that share a pixel cannot disagree and + // no atomic is needed; marking keeps the nearest centre and does. + if (clear) + owner[y * p.W + x] = BRAGG_OWNER_NONE; + else + atomicMin(&owner[y * p.W + x], BraggOwnerKey(sqrtf(d.signal), p.inv_claim, i)); + } } } @@ -761,6 +771,11 @@ BraggIntegrationEngineGPU::BraggIntegrationEngineGPU(const DiffractionExperiment // 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(std::ceil(r3 + BraggStencilGrow_px(static_cast(r_max), stencil))) + 1) + 1; + + // The pixels one reflection can mark: the box mark_mask walks, at the widest aperture on the + // detector. Run() weighs that against the frame to decide how to clear the mask afterwards. + const int mark_half = static_cast(std::ceil(r2 + BraggStencilGrow_px(static_cast(r_max), stencil))) + 1; + mask_box_px = static_cast(2 * mark_half + 1) * static_cast(2 * mark_half + 1); boxsum_shared_bytes = static_cast(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 @@ -876,16 +891,30 @@ std::vector BraggIntegrationEngineGPU::Run(const ImagePreprocessorBu // 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 + // The mask and the owner map can be taken back either by clearing the whole frame or by revisiting + // the boxes that were marked, and which is cheaper depends on the detector: on 18 Mpx the marks are + // a twentieth of the frame and clearing all of it costs 109 us of card time per image - three + // quarters of what the two integration kernels themselves cost - while on 2.5 Mpx with twelve + // thousand predictions the boxes cover more than the frame does. So compare the two. The factor of + // two is the measured penalty for scattering the writes over boxes instead of streaming them: + // 0.73 TB/s against 1.2 TB/s for the memset. + const bool clear_by_box = static_cast(n) * mask_box_px * 2 < npixel; + + // Whoever wrote the buffers last left them clean, so nothing has to be cleared here; `dirty` says + // the previous call did not get that far - or chose the frame-wide clear, or is the first call - + // and the frame has to be taken back wholesale before this image marks it. + if (dirty) { + 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<<>>(d_px_x, d_px_y, d_mask, d_owner, p, n); + dirty = true; + mark_mask<<>>(d_px_x, d_px_y, d_mask, d_owner, p, n, 0); boxsum<<>>(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, @@ -913,6 +942,11 @@ std::vector BraggIntegrationEngineGPU::Run(const ImagePreprocessorBu d_I, d_sigma, d_var_bkg, d_ok, p, n); } + // Nothing below reads the mask or the owner map, so put them back now, over the boxes that were + // marked rather than over the frame - where that is the cheaper of the two. + if (clear_by_box) + mark_mask<<>>(d_px_x, d_px_y, d_mask, d_owner, p, n, 1); + 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)); @@ -924,6 +958,9 @@ std::vector BraggIntegrationEngineGPU::Run(const ImagePreprocessorBu 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)); + // The clearing pass has run, so the buffers are clean again for the next image. + if (clear_by_box) + dirty = false; for (size_t i = 0; i < npredicted; ++i) { if (!h_ok[i]) continue; diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineGPU.h b/image_analysis/bragg_integration/BraggIntegrationEngineGPU.h index 8cdf3dd1..659be40a 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineGPU.h +++ b/image_analysis/bragg_integration/BraggIntegrationEngineGPU.h @@ -28,6 +28,13 @@ class BraggIntegrationEngineGPU : public BraggIntegrationEngine { size_t capacity = 0; // per-reflection device/host arrays hold at least this many reflections + // Whether d_mask / d_owner may still carry the marks of an earlier image. Run() clears what it + // marked before it returns when that is cheaper than clearing the frame, and this is then false in + // the steady state; it is true before the first call, after one that threw part-way through, and + // whenever the marks covered enough of the frame that clearing all of it was the cheaper choice. + bool dirty = true; + size_t mask_box_px = 0; // pixels one reflection's mark_mask box covers, at the widest aperture + // --- per-reflection device arrays (grown by EnsureCapacity) --- CudaDevicePtr d_px_x, d_px_y, d_d; CudaDevicePtr d_cx, d_cy; diff --git a/image_analysis/geom_refinement/XtalResidual.h b/image_analysis/geom_refinement/XtalResidual.h index 72c28145..1a7ed046 100644 --- a/image_analysis/geom_refinement/XtalResidual.h +++ b/image_analysis/geom_refinement/XtalResidual.h @@ -40,11 +40,16 @@ struct AngleAxisRotator { w[2] = aa[2] * theta_inverse; } - void Rotate(const T pt[3], T result[3]) const { + // The point does not have to share the rotation's type. Where one side of the product is a + // constant of the problem and the other carries the derivatives, the constant side stays a plain + // double: multiplying a Jet by a double is the same arithmetic on the value and on every + // derivative as multiplying it by a Jet whose derivatives are all zero, minus the zeros. + template + void Rotate(const P pt[3], R result[3]) const { if (at_zero) { // At zero the rotation is I + hat(angle_axis); the Taylor form keeps the derivatives // meaningful where the angle-axis vanishes. - const T w_cross_pt[3] = {aa[1] * pt[2] - aa[2] * pt[1], + const R w_cross_pt[3] = {aa[1] * pt[2] - aa[2] * pt[1], aa[2] * pt[0] - aa[0] * pt[2], aa[0] * pt[1] - aa[1] * pt[0]}; result[0] = pt[0] + w_cross_pt[0]; @@ -53,10 +58,10 @@ struct AngleAxisRotator { return; } - const T w_cross_pt[3] = {w[1] * pt[2] - w[2] * pt[1], + const R w_cross_pt[3] = {w[1] * pt[2] - w[2] * pt[1], w[2] * pt[0] - w[0] * pt[2], w[0] * pt[1] - w[1] * pt[0]}; - const T tmp = (w[0] * pt[0] + w[1] * pt[1] + w[2] * pt[2]) * (T(1.0) - costheta); + const R tmp = (w[0] * pt[0] + w[1] * pt[1] + w[2] * pt[2]) * (T(1.0) - costheta); result[0] = pt[0] * costheta + w_cross_pt[0] * sintheta + w[0] * tmp; result[1] = pt[1] * costheta + w_cross_pt[1] * sintheta + w[1] * tmp; @@ -100,40 +105,41 @@ struct XtalResidual { "Lambda cannot be close to zero"); } - template - bool operator()(const T *const beam, - const T *const distance_mm, - const T *const detector_rot, - const T *const rotation_axis, - const T *const p0, - const T *const p1, - const T *const p2, - T *residual) const { + // The observed reciprocal vector: the spot's detector position taken through the current beam, + // distance and tilt and then back-rotated into the reference crystal frame. Split out of + // operator() so that a caller holding some of those blocks fixed can hand them over as plain + // doubles - C is the type of the blocks that are NOT refined, D that of the distance - and the + // arithmetic that touches only those then stays out of the dual numbers. With C = D = T this is + // the code operator() used to run inline, term for term. + template + void ObservedRecip(const T *const beam, const D &distance_mm, + const C *const detector_rot, const C *const rotation_axis, + T *recip_obs) const { // PyFAI convention: poni_rot = Rz(-rot3) * Rx(-rot2) * Ry(+rot1). // detector_rot[0] = rot1, detector_rot[1] = rot2 are refined; rot3 is fixed // (e.g. from a PONI import) and baked in here as a constant so that a non-zero // rot3 is not silently dropped during refinement. - const T rot1 = detector_rot[0]; - const T rot2 = detector_rot[1]; + const C rot1 = detector_rot[0]; + const C rot2 = detector_rot[1]; // Ry(+rot1): rotation around Y-axis - const T c1 = ceres::cos(rot1); - const T s1 = ceres::sin(rot1); + const C c1 = ceres::cos(rot1); + const C s1 = ceres::sin(rot1); // Rx(-rot2): rotation around X-axis with inverted sign (PyFAI left-handed) - const T c2 = ceres::cos(rot2); - const T s2 = ceres::sin(rot2); + const C c2 = ceres::cos(rot2); + const C s2 = ceres::sin(rot2); // Rz(-rot3): rotation around Z (beam); constant, identity when rot3 == 0. Its sine and cosine // are taken in the constructor - they do not depend on any parameter, so recomputing them per // evaluation is two libm calls per residual for nothing. - const T c3 = T(cos_rot3); - const T s3 = T(sin_rot3); + const double c3 = cos_rot3; + const double s3 = sin_rot3; // Detector coordinates in mm const T det_x = (T(obs_x) - beam[0]) * T(pixel_size); const T det_y = (T(obs_y) - beam[1]) * T(pixel_size); - const T det_z = T(distance_mm[0]); + const D &det_z = distance_mm; // Apply Ry(rot1) first: rotate around Y const T t1_x = c1 * det_x + s1 * det_z; @@ -161,25 +167,29 @@ struct XtalResidual { // Apply goniometer "back-to-start" rotation: // brings observed reciprocal from image orientation into reference crystal frame - const T aa_back[3] = { - T(angle_rad) * rotation_axis[0], - T(angle_rad) * rotation_axis[1], - T(angle_rad) * rotation_axis[2] + const C aa_back[3] = { + C(angle_rad) * rotation_axis[0], + C(angle_rad) * rotation_axis[1], + C(angle_rad) * rotation_axis[2] }; - T recip_obs[3]; - ceres::AngleAxisRotatePoint(aa_back, recip_raw, recip_obs); - - const Eigen::Matrix e_obs_recip(recip_obs[0], recip_obs[1], recip_obs[2]); + const AngleAxisRotator rot_back(aa_back); + rot_back.Rotate(recip_raw, recip_obs); + } + // The predicted reciprocal vector h a* + k b* + l c* in the UNROTATED crystal frame, i.e. what is + // left of the prediction once the orientation is taken out of it. It depends on the cell alone, so + // a caller that holds the cell fixed evaluates this once in double and never again. + template + void PredictedRecipUnrot(const C *const p1, const C *const p2, C *recip_unrot_out) const { // Build unit cell lengths and B (convention: columns are a, b, c prior to global rotation) - Eigen::Matrix e_uc_len = Eigen::Matrix::Zero(); - Eigen::Matrix B = Eigen::Matrix::Identity(); + Eigen::Matrix e_uc_len = Eigen::Matrix::Zero(); + Eigen::Matrix B = Eigen::Matrix::Identity(); if (symmetry == gemmi::CrystalSystem::Hexagonal) { e_uc_len << p1[0], p1[0], p1[2]; - B(0, 1) = T(-0.5); // cos(120) - B(1, 1) = T(sqrt(3.0) / 2.0); // sin(120) + B(0, 1) = C(-0.5); // cos(120) + B(1, 1) = C(sqrt(3.0) / 2.0); // sin(120) } else if (symmetry == gemmi::CrystalSystem::Orthorhombic) { e_uc_len << p1[0], p1[1], p1[2]; } else if (symmetry == gemmi::CrystalSystem::Tetragonal) { @@ -193,40 +203,41 @@ struct XtalResidual { B(2, 2) = ceres::sin(p2[0]); } else { // Triclinic: p1 = (a,b,c), p2 = (alpha, beta, gamma) in radians - const T ca = ceres::cos(p2[0]); - const T cb = ceres::cos(p2[1]); - const T cg = ceres::cos(p2[2]); - const T sg = ceres::sin(p2[2]); + const C ca = ceres::cos(p2[0]); + const C cb = ceres::cos(p2[1]); + const C cg = ceres::cos(p2[2]); + const C sg = ceres::sin(p2[2]); e_uc_len << p1[0], p1[1], p1[2]; - B(0, 0) = T(1); - B(1, 0) = T(0); - B(2, 0) = T(0); + B(0, 0) = C(1); + B(1, 0) = C(0); + B(2, 0) = C(0); B(0, 1) = cg; B(1, 1) = sg; - B(2, 1) = T(0); + B(2, 1) = C(0); // c vector components: - const T cx = cb; - const T cy = (ca - cb * cg) / sg; - const T v = T(1) - cx * cx - cy * cy; - const T cz = (v >= T(0)) ? ceres::sqrt(v) : T(0); + const C cx = cb; + const C cy = (ca - cb * cg) / sg; + const C v = C(1) - cx * cx - cy * cy; + const C cz = (v >= C(0)) ? ceres::sqrt(v) : C(0); B(0, 2) = cx; B(1, 2) = cy; B(2, 2) = cz; } - // Build unrotated direct lattice columns: (B * D), then rotate them by p0. - // This avoids AngleAxisToRotationMatrix + matrix multiplications. - const T L0 = e_uc_len[0]; - const T L1 = e_uc_len[1]; - const T L2 = e_uc_len[2]; + // Build the unrotated direct lattice columns: (B * D). The caller turns the reciprocal vector + // they give by the orientation afterwards, which avoids AngleAxisToRotationMatrix and the + // matrix multiplications - see the note on the cross products below. + const C L0 = e_uc_len[0]; + const C L1 = e_uc_len[1]; + const C L2 = e_uc_len[2]; - T col0_unrot[3] = {B(0, 0) * L0, B(1, 0) * L0, B(2, 0) * L0}; - T col1_unrot[3] = {B(0, 1) * L1, B(1, 1) * L1, B(2, 1) * L1}; - T col2_unrot[3] = {B(0, 2) * L2, B(1, 2) * L2, B(2, 2) * L2}; + C col0_unrot[3] = {B(0, 0) * L0, B(1, 0) * L0, B(2, 0) * L0}; + C col1_unrot[3] = {B(0, 1) * L1, B(1, 1) * L1, B(2, 1) * L1}; + C col2_unrot[3] = {B(0, 2) * L2, B(1, 2) * L2, B(2, 2) * L2}; // Build the reciprocal vector in the UNROTATED frame and turn it once at the end, rather than // rotating all three direct columns first. A rotation commutes with the cross product and leaves @@ -234,30 +245,49 @@ struct XtalResidual { // basis of the rotated cell is the rotation of the unrotated one, and h a* + k b* + l c* is the // rotation of the same combination built from the unrotated columns. Three rotations become one, // and the cross products now run on columns that are mostly zero for every system but triclinic. - const Eigen::Matrix a_unrot(col0_unrot[0], col0_unrot[1], col0_unrot[2]); - const Eigen::Matrix b_unrot(col1_unrot[0], col1_unrot[1], col1_unrot[2]); - const Eigen::Matrix c_unrot(col2_unrot[0], col2_unrot[1], col2_unrot[2]); + const Eigen::Matrix a_unrot(col0_unrot[0], col0_unrot[1], col0_unrot[2]); + const Eigen::Matrix b_unrot(col1_unrot[0], col1_unrot[1], col1_unrot[2]); + const Eigen::Matrix c_unrot(col2_unrot[0], col2_unrot[1], col2_unrot[2]); - const Eigen::Matrix bxc = b_unrot.cross(c_unrot); - const Eigen::Matrix cxa = c_unrot.cross(a_unrot); - const Eigen::Matrix axb = a_unrot.cross(b_unrot); + const Eigen::Matrix bxc = b_unrot.cross(c_unrot); + const Eigen::Matrix cxa = c_unrot.cross(a_unrot); + const Eigen::Matrix axb = a_unrot.cross(b_unrot); - const T invV = T(1) / a_unrot.dot(bxc); + const C invV = C(1) / a_unrot.dot(bxc); - const T h = T(exp_h); - const T k = T(exp_k); - const T l = T(exp_l); + const C h = C(exp_h); + const C k = C(exp_k); + const C l = C(exp_l); - const Eigen::Matrix recip_unrot = (bxc * h + cxa * k + axb * l) * invV; + const Eigen::Matrix recip_unrot = (bxc * h + cxa * k + axb * l) * invV; + + recip_unrot_out[0] = recip_unrot[0]; + recip_unrot_out[1] = recip_unrot[1]; + recip_unrot_out[2] = recip_unrot[2]; + } + + template + bool operator()(const T *const beam, + const T *const distance_mm, + const T *const detector_rot, + const T *const rotation_axis, + const T *const p0, + const T *const p1, + const T *const p2, + T *residual) const { + T recip_obs[3]; + ObservedRecip(beam, distance_mm[0], detector_rot, rotation_axis, recip_obs); + + T recip_unrot[3]; + PredictedRecipUnrot(p1, p2, recip_unrot); - T recip_in[3] = {recip_unrot[0], recip_unrot[1], recip_unrot[2]}; T e_pred_recip[3]; const AngleAxisRotator rot_p0(p0); - rot_p0.Rotate(recip_in, e_pred_recip); + rot_p0.Rotate(recip_unrot, e_pred_recip); - residual[0] = e_obs_recip[0] - e_pred_recip[0]; - residual[1] = e_obs_recip[1] - e_pred_recip[1]; - residual[2] = e_obs_recip[2] - e_pred_recip[2]; + residual[0] = recip_obs[0] - e_pred_recip[0]; + residual[1] = recip_obs[1] - e_pred_recip[1]; + residual[2] = recip_obs[2] - e_pred_recip[2]; return true; } @@ -277,8 +307,10 @@ struct XtalResidual { // lengths, cell angles - baked in as constants, so that only beam(2) and distance(1) remain parameter // blocks. That is the rotation post-refinement's detector step, which deliberately holds the crystal at // the value its cell/axis step settled on: the seven-block form differentiates 17 parameters in order to -// use 3, this one runs on Jet. It forwards to XtalResidual with the very same values, so the -// residual and the beam/distance columns of its Jacobian are unchanged. +// use 3, this one runs on Jet. It calls the same two halves of XtalResidual with the very +// same values, so the residual and the beam/distance columns of its Jacobian are unchanged - and since +// the whole crystal is a constant here, the predicted side of the residual is worked out once in the +// constructor and never enters a dual number at all. struct XtalResidualBeamDistance { XtalResidualBeamDistance(const XtalResidual &residual, const double *detector_rot, @@ -288,28 +320,29 @@ struct XtalResidualBeamDistance { const double *uc_angle) : residual(residual), detector_rot{detector_rot[0], detector_rot[1]}, - rotation_axis{rotation_axis[0], rotation_axis[1], rotation_axis[2]}, - orientation{orientation[0], orientation[1], orientation[2]}, - uc_len{uc_len[0], uc_len[1], uc_len[2]}, - uc_angle{uc_angle[0], uc_angle[1], uc_angle[2]} { + rotation_axis{rotation_axis[0], rotation_axis[1], rotation_axis[2]} { + double recip_unrot[3]; + this->residual.PredictedRecipUnrot(uc_len, uc_angle, recip_unrot); + const AngleAxisRotator rot_p0(orientation); + rot_p0.Rotate(recip_unrot, e_pred_recip); } template bool operator()(const T *const beam, const T *const distance, T *residual_out) const { - const T rot[2] = {T(detector_rot[0]), T(detector_rot[1])}; - const T axis[3] = {T(rotation_axis[0]), T(rotation_axis[1]), T(rotation_axis[2])}; - const T p0[3] = {T(orientation[0]), T(orientation[1]), T(orientation[2])}; - const T p1[3] = {T(uc_len[0]), T(uc_len[1]), T(uc_len[2])}; - const T p2[3] = {T(uc_angle[0]), T(uc_angle[1]), T(uc_angle[2])}; - return residual(beam, distance, rot, axis, p0, p1, p2, residual_out); + T recip_obs[3]; + residual.ObservedRecip(beam, distance[0], detector_rot, rotation_axis, recip_obs); + residual_out[0] = recip_obs[0] - e_pred_recip[0]; + residual_out[1] = recip_obs[1] - e_pred_recip[1]; + residual_out[2] = recip_obs[2] - e_pred_recip[2]; + return true; } const XtalResidual residual; const double detector_rot[2]; const double rotation_axis[3]; - const double orientation[3]; - const double uc_len[3]; - const double uc_angle[3]; + // The crystal is held whole here, so the predicted reciprocal vector is a constant of the + // residual - cell, orientation and all - and is worked out once, in the constructor. + double e_pred_recip[3]; }; // Same residual with only the detector distance baked in, leaving the other six blocks free: 16 @@ -318,8 +351,8 @@ struct XtalResidualBeamDistance { // degenerate - so the block was only ever declared in order to be frozen. Dropping it is worth more // than one parameter in seventeen suggests: Jet keeps its derivatives in an Eigen vector, and // 17 doubles is one lane past the four-wide boundary, so it costs five packet operations and cannot be -// over-aligned where 16 costs four and can. The distance is forwarded unchanged, so the residual and -// every remaining Jacobian column are the same. +// over-aligned where 16 costs four and can. The distance is passed through as a plain double, so the +// residual and every remaining Jacobian column are the same. struct XtalResidualFixedDistance { XtalResidualFixedDistance(const XtalResidual &residual, double distance_mm) : residual(residual), distance_mm(distance_mm) { @@ -328,8 +361,20 @@ struct XtalResidualFixedDistance { template bool operator()(const T *const beam, const T *const detector_rot, const T *const rotation_axis, const T *const p0, const T *const p1, const T *const p2, T *residual_out) const { - const T distance[1] = {T(distance_mm)}; - return residual(beam, distance, detector_rot, rotation_axis, p0, p1, p2, residual_out); + T recip_obs[3]; + residual.ObservedRecip(beam, distance_mm, detector_rot, rotation_axis, recip_obs); + + T recip_unrot[3]; + residual.PredictedRecipUnrot(p1, p2, recip_unrot); + + T e_pred_recip[3]; + const AngleAxisRotator rot_p0(p0); + rot_p0.Rotate(recip_unrot, e_pred_recip); + + residual_out[0] = recip_obs[0] - e_pred_recip[0]; + residual_out[1] = recip_obs[1] - e_pred_recip[1]; + residual_out[2] = recip_obs[2] - e_pred_recip[2]; + return true; } const XtalResidual residual; @@ -340,8 +385,9 @@ struct XtalResidualFixedDistance { // that only beam(2) and orientation(3) remain parameter blocks. Ceres sizes its autodiff dual numbers // from the DECLARED blocks, not from which of them the caller then holds constant, so the seven-block // form above differentiates 17 parameters even when 5 are free; this one runs on Jet. It -// forwards to XtalResidual with the very same values, so the residual and the beam/orientation columns -// of its Jacobian are unchanged. +// calls the same two halves of XtalResidual with the very same values, so the residual and the +// beam/orientation columns of its Jacobian are unchanged - and the cell half, which no free parameter +// touches, is worked out once in the constructor rather than on every dual number. struct XtalResidualBeamOrientation { XtalResidualBeamOrientation(const XtalResidual &residual, double distance_mm, @@ -352,25 +398,31 @@ struct XtalResidualBeamOrientation { : residual(residual), distance_mm(distance_mm), detector_rot{detector_rot[0], detector_rot[1]}, - rotation_axis{rotation_axis[0], rotation_axis[1], rotation_axis[2]}, - uc_len{uc_len[0], uc_len[1], uc_len[2]}, - uc_angle{uc_angle[0], uc_angle[1], uc_angle[2]} { + rotation_axis{rotation_axis[0], rotation_axis[1], rotation_axis[2]} { + this->residual.PredictedRecipUnrot(uc_len, uc_angle, recip_unrot); } template bool operator()(const T *const beam, const T *const p0, T *residual_out) const { - const T distance[1] = {T(distance_mm)}; - const T rot[2] = {T(detector_rot[0]), T(detector_rot[1])}; - const T axis[3] = {T(rotation_axis[0]), T(rotation_axis[1]), T(rotation_axis[2])}; - const T p1[3] = {T(uc_len[0]), T(uc_len[1]), T(uc_len[2])}; - const T p2[3] = {T(uc_angle[0]), T(uc_angle[1]), T(uc_angle[2])}; - return residual(beam, distance, rot, axis, p0, p1, p2, residual_out); + T recip_obs[3]; + residual.ObservedRecip(beam, distance_mm, detector_rot, rotation_axis, recip_obs); + + T e_pred_recip[3]; + const AngleAxisRotator rot_p0(p0); + rot_p0.Rotate(recip_unrot, e_pred_recip); + + residual_out[0] = recip_obs[0] - e_pred_recip[0]; + residual_out[1] = recip_obs[1] - e_pred_recip[1]; + residual_out[2] = recip_obs[2] - e_pred_recip[2]; + return true; } const XtalResidual residual; const double distance_mm; const double detector_rot[2]; const double rotation_axis[3]; - const double uc_len[3]; - const double uc_angle[3]; + // The cell is held fixed here, so h a* + k b* + l c* in the unrotated frame - the B matrix, the + // three cross products and the cell volume - is a constant of the residual, worked out once here + // rather than on every dual number. + double recip_unrot[3]; }; diff --git a/tests/BraggIntegrationEngineGPUTest.cpp b/tests/BraggIntegrationEngineGPUTest.cpp index b0a0f899..ff6f56e6 100644 --- a/tests/BraggIntegrationEngineGPUTest.cpp +++ b/tests/BraggIntegrationEngineGPUTest.cpp @@ -305,6 +305,61 @@ TEST_CASE("BraggIntegrationEngineGPU_MatchesCPU") { } } +// The mask and the owner map are cleared by the run that marked them rather than at the start of the +// next one, so a reused engine has to give the same answer as a fresh one. A first frame whose spots +// are somewhere else entirely is what would show a leftover mark: a stale mask pixel is read as a +// neighbour's signal and dropped from the background ring, a stale owner steals a pixel outright. +TEST_CASE("BraggIntegrationEngineGPU_ReusedEngineMatchesFresh") { + if (get_gpu_count() == 0) { + WARN("No CUDA GPU present. Skipping BraggIntegrationEngineGPU_ReusedEngineMatchesFresh"); + return; + } + + for (OverlapMode ovl : {OverlapMode::Off, OverlapMode::Exclude}) { + const DiffractionExperiment experiment = + MakeExperiment(IntegratorMode::ProfileGaussian, std::nullopt, 4.0f, false, DetJF(2), + 0.0f, 0.0f, 0.0f, 0.0f, ovl); + const size_t width = experiment.GetXPixelsNum(); + const size_t height = experiment.GetYPixelsNum(); + const size_t npixel = experiment.GetPixelsNum(); + + // Two frames whose spot grids do not line up, so the first frame's marks fall on the second + // frame's background rings rather than back onto its own disks. + const Scene first = BuildScene(width, height, 47); + const Scene second = BuildScene(width, height, 60); + REQUIRE(first.predicted.size() > 60); + REQUIRE(second.predicted.size() > 60); + + auto integrate = [&](BraggIntegrationEngineGPU &engine, const Scene &scene, + const std::shared_ptr &stream) { + ImagePreprocessorBufferGPU img(npixel); + for (size_t i = 0; i < npixel; ++i) img[i] = scene.image[i]; + REQUIRE(cudaMemcpyAsync(img.getGPUBuffer(), img.getBuffer().data(), + npixel * sizeof(int32_t), cudaMemcpyHostToDevice, *stream) == cudaSuccess); + return engine.Run(img, scene.predicted, scene.predicted.size(), 7); + }; + + auto stream_fresh = std::make_shared(); + BraggIntegrationEngineGPU fresh(experiment, stream_fresh); + const auto out_fresh = integrate(fresh, second, stream_fresh); + + auto stream_reused = std::make_shared(); + BraggIntegrationEngineGPU reused(experiment, stream_reused); + integrate(reused, first, stream_reused); + const auto out_reused = integrate(reused, second, stream_reused); + + INFO("overlap mode " << static_cast(ovl)); + REQUIRE(out_reused.size() == out_fresh.size()); + for (size_t i = 0; i < out_fresh.size(); ++i) { + INFO("reflection " << i); + CHECK(out_reused[i].h == out_fresh[i].h); + CHECK(out_reused[i].I == out_fresh[i].I); + CHECK(out_reused[i].sigma == out_fresh[i].sigma); + CHECK(out_reused[i].bkg == out_fresh[i].bkg); + } + } +} + // Hidden ([.]) benchmark: the raison d'etre of the GPU port is < 2 ms/frame (vs ~142 ms on the CPU // for ProfileIntegrate2D). Run explicitly with: ./jfjoch_test "[bragg_bench]" TEST_CASE("BraggIntegrationEngineGPU_Benchmark", "[.][bragg_bench]") {