Files
Jungfraujoch/image_analysis/bragg_integration/BraggIntegrationEngineGPU.h
T
leonarski_fandClaude Opus 5 641f890a40 Build the frame's constants once, and page-lock what the integration engine copies
The per-image geometry refinement is the largest stage of the image loop, and a
third of it was arithmetic on numbers that never change.

The residual derives the detector angles' sines and cosines, the goniometer's
back-rotation - a three-argument hypot, a sine, a cosine and a division - and the
reciprocal basis of the cell on every evaluation. On the rotation path the detector
angles and the axis are held fixed and stored as plain doubles, so all of it is
constant, not merely constant per block: there is one frame per image and one cell.
Three solves an image, fifty iterations a solve and a thousand spots make it tens of
thousands of repetitions of the same result. The frame's constants are now built
once and handed in. The body they feed is the same body, split out rather than
copied, so no expression is reassociated - in particular the reciprocal vector is
still formed as the basis times the inverse volume, with the volume not folded into
the basis.

The spot confidence weights depend only on each spot's resolution and intensity,
which no solver touches, and were recomputed identically for each of the three
passes. They are computed once. The sort behind them ordered indices through a
projection that chased a random eighty-byte-strided element per comparison; it now
sorts a packed resolution and index, which makes the same comparisons in the same
sequence and therefore the same permutation. The spot list itself was copied per
image through an initializer list whose elements are const; it is passed as a view.

The integration engine was the last one in the loop copying through pageable host
memory - three transfers in and eight out per image, twenty-six bytes a reflection,
while every other engine already page-locks its staging. A driver copy from pageable
memory stages through its own pinned buffer on the calling thread, which is why an
asynchronous copy was averaging a hundred and thirteen microseconds. Page-locked, the
same seventeen thousand calls cost four hundred and thirty-two milliseconds instead
of one and a half seconds, and the wait moves to the synchronisation point where it
belongs.

Two smaller ones: the reflections were copied into the per-image message for a
process file that a merging run does not write, so the copy is made where a writer
exists; and the intensity statistics and the Wilson estimate walked the same
eighty-byte array twice to read twelve bytes, which is now one pass with each
accumulation in its own order.

Every reflection file is byte-identical on four crystals; the process file's
reflections match dataset for dataset, and its azimuthal arrays differ no more
between this build and the last than the last differs from itself.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EGpGdgmJ8MyY9pCGWjktyi
2026-08-25 01:08:22 +02:00

84 lines
4.8 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#pragma once
#include <cstdint>
#include <memory>
#include <vector>
#include "BraggIntegrationEngine.h"
#include "../indexing/CUDAMemHelpers.h"
// CUDA engine: reproduces BraggIntegrationEngineCPU up to floating-point precision. Each stage is a
// kernel with one CUDA block per reflection cooperating over the small window via shared-memory
// reductions (the natural mapping for thousands of independent, tiny per-spot integrations).
//
// Pipeline (profile modes): reset -> mark_mask -> boxsum -> learn_profile -> build_profiles -> fit
// (the resolution shell is computed inline, so there is no separate shell pass). BoxSum mode stops
// after boxsum (that pass is the BraggIntegrate2D box integrator and the seed of the profile fit).
// The preprocessed image already lives on the device (ImagePreprocessorBufferGPU::getGPUBuffer());
// only the per-frame predicted centres are uploaded.
class BraggIntegrationEngineGPU : public BraggIntegrationEngine {
std::shared_ptr<CudaStream> stream;
int threads;
size_t fit_shared_bytes;
int rad_w = 0; // radial-background window of boxsum, in bins of one pixel
size_t boxsum_shared_bytes = 0;
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<float> d_px_x, d_px_y, d_d;
CudaDevicePtr<int> d_cx, d_cy;
CudaDevicePtr<float> d_I, d_sigma, d_bkg, d_bkg_var, d_var_bkg, d_obs_x, d_obs_y;
CudaDevicePtr<float> d_isum; // box-sum raw sum, for the radial correction
CudaDevicePtr<int> d_ninner, d_rbin, d_kbin;
CudaDevicePtr<uint8_t> d_ok, d_strong, d_has_obs;
// --- radial background curvature correction (see BraggIntegrationEngine) ---
int n_rad = 0; // radial bins, 0 when the correction is off
CudaDevicePtr<unsigned long long> d_rad_sum; // integer pixel sums, see boxsum
CudaDevicePtr<float> d_k_diff;
CudaDevicePtr<int> d_rad_cnt;
// --- fixed-size device arrays ---
// The learning/fit math is single precision: FP64 is heavily throttled on consumer GPUs and the
// extraction is Poisson-noise limited, so float reproduces the double CPU path to ~1e-4.
CudaDevicePtr<uint8_t> d_mask; // per-pixel inner-stencil reflection mask
// Per-pixel (distance, reflection) key naming the nearest predicted centre; allocated only when
// an overlap treatment is on, so the default path costs no extra device memory.
CudaDevicePtr<uint32_t> d_owner;
// Fixed-point (see PROFILE_FIXED): a float atomicAdd here made the profile depend on the order
// the blocks arrived in, and with it every intensity fitted through it.
CudaDevicePtr<unsigned long long> d_shell_grid, d_global_grid; // learned profile accumulators (N_SHELL*GG, GG)
CudaDevicePtr<float> d_shell_P, d_global_P; // normalised profiles (empirical mode)
CudaDevicePtr<unsigned long long> d_mom; // learned 2nd moments, 3 per shell + global
CudaDevicePtr<float> d_sigma2_r, d_sigma2_t; // radial/tangential widths, N_SHELL + global
CudaDevicePtr<int> d_shell_n, d_global_n;
CudaDevicePtr<unsigned long long> d_invd2; // [min,max] inv-d^2 as monotonic bit patterns
// --- host staging (copied back once per frame) ---
// Pinned, like every other engine's staging: a copy out of pageable memory does not return until the
// driver has staged it through a bounce buffer, so eight of them in a row are eight serialised
// round-trips rather than eight queued transfers.
CudaHostPtr<float> h_px_x, h_px_y, h_d;
CudaHostPtr<float> h_I, h_sigma, h_bkg, h_var_bkg, h_obs_x, h_obs_y;
CudaHostPtr<uint8_t> h_ok, h_has_obs;
void EnsureCapacity(size_t n);
public:
BraggIntegrationEngineGPU(const DiffractionExperiment &experiment, std::shared_ptr<CudaStream> stream);
std::vector<Reflection> Run(const ImagePreprocessorBuffer &image,
const std::vector<Reflection> &predicted, size_t npredicted,
int64_t image_number) override;
};