Jungfraujoch/docs/GPU_DATA_ANALYSIS.md

GPU data analysis

Jungfraujoch runs most of its crystallographic analysis on NVIDIA GPUs where one is present: the per-image loop of jfjoch_broker, jfjoch_viewer and rugnux, and most of rugnux's dataset-level steps - the pre-scan, scaling and merging, model validation, molecular replacement and refinement. This page lists the GPU engines (all of them in .cu files under image_analysis/), with what each computes. The methods themselves are described once, in the CPU/GPU data analysis reference and its siblings, which each entry links to.

Every engine has a CPU counterpart that computes the same thing, and the CPU-only build (JFJOCH_USE_CUDA=OFF) runs the whole pipeline on it, the ffbidx indexer included (an in-tree CPU port of the library's kernels). How closely the two agree is stated engine by engine below and in the source: where the arithmetic is integer, or the order of every sum is fixed to the host's, the result is the same to the bit; where the GPU uses float distances, cuFFT for FFTW, or a different (fixed) summation order, the two agree to rounding. Every engine is deterministic on its own - no floating-point atomics decide a result - so the same input gives the same output on every run of a GPU build.

Getting the image onto the card

• Bitshuffle+LZ4 decoder (BSLZ4DecoderGPU.cu) - decompresses the stored chunk on the device, so only the compressed bytes cross PCIe, and un-transposes it straight into the preprocessing step. The same pixels as the host decoder. The approach is Jon Wright's (ESRF). §0
• CBF byte-offset decoder (CBFDecoderGPU.cu) - decodes a miniCBF x-CBF_BYTE_OFFSET stream on the device into the same int32 pixels as the host decoder, refusing a stream that ends before the image does. §0
• marCCD frames (ImagePreprocessorGPU.cu) - the stored 16-bit pixels are uploaded as they are and widened to int32 on the card, inside preprocessing. §0
• Preprocessing (ImagePreprocessorGPU.cu, ImagePreprocessorBufferGPU.cu) - applies the pixel mask and flags error and saturated pixels, counting each, and keeps the preprocessed image on the device for every engine below; on request it gathers the pixels of the strong spots, or small squares around given positions, back to the host. Integer, so identical to the host. §1.6

Per image: profiles and spots

• Azimuthal integration (AzIntEngineGPU.cu) - bins every pixel into its resolution (and optionally azimuth) bin with the geometric corrections applied and sums them, as a radial profile. Fixed-point sums, so the profile does not depend on the order the threads add in. §2
• ROI integration (ROIIntegrationGPU.cu) - sums the pixels of user-defined detector regions of interest (box, circle, azimuthal), each pixel feeding any subset of them, on the EIGER/CPU path where the FPGA does not. Integer sums. FPGA: ROI integration
• Local-statistics spot finder (ImageSpotFinderGPU.cu) - flags strong pixels by their signal-to-noise against a local box. §3.1
• Adaptive spot finder (AdaptiveSpotFinderGPU.cu) - reduces the image into resolution rings once, and from that one pass gives both the azimuthal profile and each ring's robust background, which sets the self-calibrating detection threshold; the result is intersected with the local test above. §3.2
• Spot extractor (SpotExtractorGPU.cu) - joins the flagged pixels into connected components and sums each into a spot descriptor on the device, so only the spot list comes back. The same components, in the same order, as the host labelling (the sparse formulation of the ACTS/traccc project). §3.4

Pre-scan accumulators

These read a subset of the frames before processing starts, on every card at once, and are added together across cards; their sums are integers, so the split of the frames between the cards changes nothing.

• Beam-stop projection and mask (ShadowAccumulatorGPU.cu, ShadowMaskGPU.cu) - accumulates the per-pixel maximum, sum and count of the raw counts straight from the compressed frames, then draws the beam-stop shadow mask from them, step for step as the host does. §1.5
• Hot pixels (HotPixelsGPU.cu) - per frame, the ring-sector and ring order statistics, the levels made from them and the per-pixel counts of lit frames, from which persistently hot pixels are masked. Exact order statistics and integer sums: the host's result to the bit. §1.6
• Beam centre from the background (BeamCenterBackgroundGPU.cu) - bins the usable pixels about a trial centre and sums each cell, the two passes of the background-isotropy fit; the fit itself stays on the host. §1.4
• Beam centre by FFT (BeamCenterFFTGPU.cu) - the self-convolutions that score every candidate centre at once, their combination into one score surface, and the shortlist of its peaks, with only the shortlist brought back. Runs on FFTW on the host instead where it would need more than half of the card's total memory, decided before the card is touched and logged. §1.4

Indexing

• FFT indexer (FFTIndexerGPU.cu) - projects the spots' reciprocal-space vectors onto a dense set of directions, histograms each projection and transforms it with cuFFT, 4096 directions at a time, and picks each direction's most prominent lattice peak; the candidate vectors then go to the host's lattice reduction. cuFFT gives the same bits at any batch size, so the batching changes nothing but the memory held. A rotation run gives its indexers' memory back to the card once a pass has done its last indexing. §5.1, §5.2
• Fast-feedback indexer (ffbidx) - third-party, from the fast-feedback-indexer project (PSI); a known-cell indexer for serial stills, used as fetched. Its CPU counterpart is an in-tree port with the same samples and scores, agreeing to rounding. §4.3

Prediction and integration

• Still prediction (BraggPredictionGPU.cu) - the reflections of a lattice within the excitation-error cutoff - or the window that follows the run's rocking curve, centred on its radial scale - with their detector positions and per-reflection corrections. §8.2
• Rotation prediction (BraggPredictionRotGPU.cu) - the reflections crossing the Ewald sphere within an image's rotation range, with their rotation angle and partiality; they are ordered, and under a cap the best-recorded ones selected, on the device, in exactly the order the host gives. §8.3
• Bragg integration (BraggIntegrationEngineGPU.cu) - for every predicted reflection: the three-ring regions and neighbour ownership, box summation, the per-shell profile learned from the strong spots and rebuilt per reflection, the radial background-curvature correction, and the profile fit. The profile fit is in single precision on both engines, in the same arithmetic, so the two differ only in summation order. §9.2, §9.3

Scaling and merging (rotation)

One engine (RotationScaleMergeGPU.cu) keeps every observation resident on the card for the whole scaling run; each per-group sum is a segmented reduction in a fixed order, so it is deterministic. Whether the scaling runs there at all is decided before any observation is on the card, from the card's total memory: a set whose sort alone, or whose scaling (estimated from its observation and distinct-hkl counts once the sort has counted them), would take more than 80% of it scales on the CPU instead, with a warning in the log (RotationScaleMerge: ... scaling on the CPU instead). Judged on the total and not on what is free, the same input on the same card always takes the same path.

• Ingest - the integrated reflections uploaded frame by frame, twice: first only what the sort and the resolution cut need, then, once the scaling is known to fit, the observations the cut kept, written straight to their place (the predicted detector positions stay on the host); the frame flux applied; the sort into raw-hkl runs and the resolution cut; the recomputed rotation angles (delta-phi), the rocking events relaid, and the partiality recomputed from the smoothed mosaicity. The same sort and cut as the host, to the bit. §10.2, §10.6
• Per-frame scaling (ReduceGroupMeans, FitPerFrameG, UpdateCorr, SmoothCorr, the zeta and frame filters, PerFrameCC) - the reference intensity of every symmetry group, the per-frame scale fitted against it, its smoothing over the sweep, and the per-frame correlation behind --min-image-cc. §10.1, §10.3
• Combine - each reflection's partials summed into a full, with its capture fraction and variance; then the fulls' own per-frame scale and its smoothing. §10.6
• Correction surfaces - the decay, absorption and modulation surfaces' terms, their reference and their least-squares sums over the fulls, on held-out subsets. §10.6
• Merge - the inverse-variance merge (I(+)/I(-) and the mean), the half-sets for CC1/2, R_meas, the per-group scatter, the anomalous (Bijvoet) excess and the error-model samples. §10.4, §10.5
• Error model (ErrorModelGPU.cu) - the binned fit of the error model's a and b, the bin medians identical to the host's; a and b can differ from the host fit in their last bits (summation order). §10.5
• Wilson outliers (WilsonOutliersGPU.cu) - the per-observation Wilson-statistics rejection, the host's result to the bit. §13.3
• Pair sort (PairSortGPU.cu) - sorts the (deviation, weight) pairs of the outlier test's weighted median; the same sequence as std::sort. §13.3

Model validation

The model's structure factors and the rigid body run on the card where there is one, every part placed on it or on the CPU up front against one budget, 80% of the card's total memory (§14.8); the GPU and CPU agree to rounding (float distances, cuFFT for FFTW, the crystal's symmetry composed in reciprocal space instead of by symmetrizing the grid). On a machine with several cards, the rigid-body engines the first card has no room for go on the others, each against its own budget.

• Model density (ModelDensityGPU.cu) - the density of one copy of the model's atoms on a grid, as a gather over bricks of the grid, bit-identical from run to run. §14.1
• Bulk-solvent mask (ModelMaskGPU.cu) - the solvent mask from every symmetry image of every atom, with island removal and the shrink, as gemmi computes it. §14.2
• Structure factors and maps (ModelStructureFactorsGPU.cu) - F_calc and F_mask from the two grids by cuFFT, composed over the group's operators and read off at the data's reflections; and the inverse transform of map coefficients into maps. §14.1, §14.4
• Scale fit (ModelScaleGPU.cu) - the sums of the overall-scale, anisotropic-B and bulk-solvent fit over the reflections, in double and a fixed order, for many fits at once (the solvent grid, the null replicates); the Levenberg-Marquardt control stays on the host. Used where its engine fits the validation's budget, the CPU fit otherwise. §14.2
• Rigid body (RigidBodyGPU.cu) - the rigid-body target for a placement: the model moved, its F_calc and mask, the scale, the residuals and their Jacobian; used for the placement itself and for every replicate of the null. §14.8, §14.5

Molecular replacement

• Molecular transform (MolecularReplacementGPU.cu, MapToFPhiGPU) - the Fourier transform of the search model's density in a large box, by cuFFT, where its buffers fit the search's budget (80% of the card, the search holding nothing else yet), kept only as far out as it is read.
• Rotation function (MolecularReplacementGPU.cu) - for every trial orientation, the model's intensities summed over the copies, normalised by resolution bin and correlated with the data's normalised intensities.
• Translation function (MolecularReplacementGPU.cu) - for each orientation, the function over every position in the cell, beside the copies already placed where there are any, and its mean, spread and highest point.

What of the search goes on the card is planned once per merge on the same budget as model validation's, before the search starts (§14.10). On a machine with several cards, every other card holds a copy of the molecular transform where it fits, the rotation and translation functions are split over the cards by orientation, and the rigid-body engines are spread over them as in model validation; the result does not depend on how many cards there are.

The placements are then refined, scored and checked against their null with the rigid body and the structure factors above; the likelihood scoring runs on the host. The search is described in §14.10, its keys and decision rules in Molecular replacement.

Refinement

The restrained refinement of a placed model (§15) evaluates its target and its gradients on the card: one engine (RefinementGPU.cu, with the host half in RefinementGPU.cpp) keeps the reflections, the grids, the symmetry and the restraints resident, and per evaluation takes only the atoms' positions up and one value and one gradient per atom back. Everything that depends on the cell, the group and the reflections - the grids, the operators, the adjoint's lists - is worked out once on the host with gemmi when the engine is made. The GPU and CPU engines compute the same target to rounding (float distances and densities, cuFFT for FFTW, summation order); the GPU's result is the same from run to run, as no sum of a result is a floating-point atomic: every per-point, per-reflection and per-atom sum is a gather in a fixed order, and every total a cub::DeviceReduce tree, whose order for a given size is fixed.

One evaluation - target and gradient of every atom - is a density gather, a forward and an inverse transform and a gradient gather, whatever the number of atoms (§15.2):

• Density (ModelDensityGPU.cu, the gather of model validation) - the atoms' blurred density of one copy of the model on the structure-factor grid. The grid is cut into bricks of 8^3 points; each atom's box is counted against the bricks it touches, the counts scanned (cub::DeviceScan) into a list of (brick, atom) pairs, and the list sorted stably by brick (cub::DeviceRadixSort), so each brick holds its atoms in model order. One block per brick then stages its atoms in shared memory (128 at a time) and every point adds, in model order, each atom whose radius reaches it - no atomics, the grid the CPU walk's to the bit for the same atoms. The pair count is the one wait inside an evaluation (it sizes the sort). The atoms' density coefficients are uploaded again only where some atom's B changed; the positions go up through a pinned buffer.
• Forward transform - cuFFT r2c of the grid into its half-u spectrum, the plans made once per engine with an explicit work area.
• Composition and target (compose_target_kernel, one thread per reflection) - F_calc of the reflection composed from the spectrum over the group's operators (each operator's phase exp(2 pi i h.t) computed exactly for the translation's fraction, sums in double), times the centring count and the unblur; then, on the working reflections, F_model = k (F_calc + F_mask) with the macro-cycle's scale, mask and D, Σ held on the card, the likelihood term -ln L of the reflection, and the coefficient conj(dT/dF) the derivative map is made of. The terms are summed by cub::DeviceReduce.
• Adjoint of the composition (scatter_kernel, one thread per spectrum point) - the coefficients taken back onto the half-u spectrum as a gather, not a scatter: a CSR list made on the host at setup gives, for every spectrum point a (reflection, operator) term lands on, the terms in a fixed order, conjugated where they land as a Friedel mate and doubled onto the planes k_u = 0 and k_u = nu/2, which hold their own mates. Each point sums its terms in double and is written once.
• Inverse transform - cuFFT c2r of that spectrum into the derivative map, on the density's grid.
• Atom gradient (atom_gradient_kernel, one warp per atom) - each atom's dT/dx, dy, dz and dB: the derivative map summed against the derivatives of the atom's five-Gaussian density over the same box and radius as the gather, the warp's 32 lanes walking the box's points in steps of 32 without an integer division (the box index stepped on, the wrap into the grid by one subtraction), and the lanes' sums added in a fixed shuffle tree. The riding hydrogens' derivatives are passed to their parents on the host.

Beside the X-ray target, and in the same submission:

• Restraints (restraint_kernel, one thread per restraint) - bonds, angles, torsions, chirals, planes (up to 32 atoms), repulsions (with their symmetry image), metal links and B similarity (§15.5), each restraint writing its energy and its gradient into one slot per atom it involves; the energies summed by cub::DeviceReduce.
• Slot gather (slot_gather_kernel, one thread per atom) - each atom's restraint gradient, its slots summed in the order a CSR list gives. Two restraint sets are held (positions and B), set again when the repulsions and contacts are rebuilt.

EvaluateWithRestraints queues the X-ray evaluation, writes the restraints' positions into a pinned buffer while the card works on it, queues the restraints behind it and waits once; the target, the restraint energy and both gradients come back into pinned buffers. The minimiser (L-BFGS) stays on the host.

Once per macro-cycle (§15.3):

• F_calc and F_mask - the forward path above without the target, and the bulk-solvent mask drawn on the card (ModelMaskGPU.cu) on its own fine grid (d_min/3 or 0.4 A, whichever is finer, at most 2^26 points), with a grid, spectrum and cuFFT plan of its own where that grid is not the density's: one block per (atom, symmetry image) sets the points within the atom's REFMAC radius, the islands are removed by a union-find labelling (integer atomicCAS, rooted at each component's lowest index, so the partition does not depend on the order of the unions), and the shrink is two passes - the edge points marked, then each edge point's stencil written as solvent. Every write is an idempotent store, so the mask is the same bit for bit from run to run; F_mask is read off at the reflections (fmask_kernel).
• Robust scale fit (ModelScaleGPU.cu) - the k_sol/B_sol grid of fits, each fit's overall scale and anisotropic B by gemmi's Levenberg-Marquardt, many fits per launch: the sums over the reflections (the WSSR, J^T J and J^T r, in double) by a fixed launch shape of 48 blocks of 128 threads, a fixed shuffle tree within each warp and the blocks summed in order; the damping and the 7x7 solve stay on the host. The fine pass is robust: four rounds of residuals, weights 1/sqrt(1 + (r/c)^2) and the fit again, with c the median residual of the first round found on the card by a radix select (residual_median: the residuals are not negative, so their order is that of their bits, found a byte at a time in four passes of integer counts - the element std::nth_element puts at n/2). After a trial's first macro-cycle the fit starts from the fine pass around the previous winner, skipping the coarse one. Its stream has the card's highest priority, so that the fit's short, host-waited launches do not queue behind the long kernels of the other trials on the card.
• D and Σ, the likelihood's parameters, are fitted on the host (bin by bin, on all of the run's threads).

Trials, cards and waits. The weight trials of one set (§15.7) run side by side, each on a slot of its own - a copy of the model, its topology, a refinement engine and a scale-fit engine - as many slots as the largest set of trials has (six at the defaults), the ones beside the main one built on their own threads while the placement runs. The slots' engines go on the cards in turn from the main one's, each card taking engines while they fit its budget; on a machine with several cards the trials are spread over all of them. A slot whose engine cannot be made is left out rather than moved to the CPU, and the outcome does not depend on how many trials run at once or on which card. The engine waits on its stream through an event made with cudaEventBlockingSync (StreamWaiter), so a waiting trial sleeps instead of spinning a core - with six engines on one card the spinning took half of the refinement's CPU time.

Memory. RefinementGPUEngine::DeviceBytes counts the engine's reserve before anything is allocated: per atom (positions, boxes, densities, mask atoms, gradients), the grid, its spectrum and cuFFT's work area, the mask's own grid and spectrum, per reflection (indices, scales, F_o, the row model, F_calc, F_mask, coefficients, terms), the adjoint's lists, the gather's pair reserve - sized for every atom at the largest B a refinement may reach with the widest blur - and the restraints at about 30 records and 60 slots per atom. The engine goes on the card when that total, plus what is already resident there (a merge still at work), is within 80% of the card's total memory; the trial slots are counted against the same budget, card by card.

Falling back. Where the GPU cannot take the case (an atom's box wider than the cell, more than 192 mask operators), where the engine is over the budget, or where CUDA fails while the engine is made, the refinement runs on the CPU engine, with a line in the log (Refinement: on the CPU - ...). --no-gpu and a build without CUDA take the CPU engine as well. The refinement has no Metal engine: on macOS it runs on the CPU.

Indicative timings (RTX 5080, -N 8, the refinement alone): a 5,000-atom model (half of it riding hydrogens) at 1.65 A with 31k reflections refines in about 1.3 s; a 4,500-atom model at 1.09 A with 92k reflections in about 2 s; a 50,000-atom model at 3.6 A with 51k reflections in about 7 s, of which the kernels take 5 s - there the density gather and the atom gradient (about 0.9 and 0.8 ms per evaluation, over a thousand evaluations) are half or more of the kernel time, the scale fit 10-25% (more at high resolution) and the fine mask 15-18% at low resolution. Inside a full rugnux --model run of a 5,000-atom model the refinement took 1.5 s of 8 s. The CPU engine (-N 8) takes about a minute and a half for a set the card refines in about 2 s.

Apple GPUs (Metal)

A macOS build has no CUDA, and three engines run on the Apple GPU instead (none under rugnux --no-gpu; where the GPU fails, the CPU takes over with a warning):

• Bragg integration (BraggIntegrationEngineMetal.mm, .metal) - the CPU engine with its pixel work moved to the GPU: the neighbour mask, the ownership of shared pixels, the box summation's walk over every reflection, the background clip and the single-precision profile fit. Everything computed in double and the profile learning stay on the CPU, as Apple GPUs have no double precision. The fit's Gaussian uses an exp of our own (BraggExp.h) on both, as no GPU exp matches the CPU's. The same result as the CPU engine to the bit, which a macOS CI step checks on every push; rugnux --no-gpu keeps it off the GPU, and so does --background-trim, whose trimmed mean has no Metal kernel. §9.2
• Model density (ModelDensityMetal.mm, .metal) - the density of the model's atoms on a grid, for model validation and molecular replacement, as a gather over bricks of the grid as under CUDA; the same grid as the CPU's to the bit. §14.1
• Molecular-replacement rotation function (MolecularReplacementMetal.mm, .metal) - the rotation function's trial orientations scored from a copy of the molecular transform; equal to the CPU's to rounding. §14.10