leonarski_fandClaude Opus 4.8
dd25de461d
First stage of moving the rotation scale/merge onto the GPU. The per-frame partial-scaling loop (inverse-variance group-mean reduction -> robust per-frame IRLS G -> corr update, x scaling_iter) now runs in RotationScaleMergeGPU (.cu) when a GPU is present; the CPU loops remain the fallback. The host keeps the one-time raw-hkl sort and the per-space-group gemmi ASU keying, and hands the GPU a group-ordered permutation + CSR so the per-group reduction is a DETERMINISTIC segmented reduction (one thread per group, fixed order, no atomics) - preserving the run-to-run determinism just won on the CPU path (a float atomicAdd reduction would have re-introduced jitter). Reduction is one-thread-per-group (groups average tens of obs, so a block-per-group wastes threads); the IRLS is one block per frame with a deterministic shared-memory reduction. Validated: bit-identical to the CPU path and deterministic run-to-run on lyso/cytC/Ins_H/pding (P41212 ISa 7.8 CC1/2 99.7%, etc.). The scaling kernels are ~7x faster than the CPU compute (~36 ms for 3 iters vs ~0.28 s); end-to-end scale/merge ~2.0 -> ~1.5 s. The remaining gap to the <1 s target is the per-pass host round-trip (corr down/upload for the CPU combine + per-SG group-CSR rebuild); phase 2 keeps the data resident by moving the 3D combine and the merge/error-model onto the GPU too, so nothing round-trips. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
61 lines
3.1 KiB
C++
61 lines
3.1 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 <optional>
|
|
#include <vector>
|
|
|
|
// GPU engine for the RotationScaleMerge hot loops. The class keeps the per-observation data resident on
|
|
// the device as a structure-of-arrays (coalesced) and runs the scaling loop there. The host keeps the
|
|
// one-time raw-hkl sort and the per-space-group ASU keying (gemmi); it hands the GPU the dense group ids
|
|
// plus a group-ordered permutation so the per-group reduction is a deterministic segmented reduction
|
|
// (one block per group, fixed order, no atomics) - matching the run-to-run determinism of the CPU path.
|
|
//
|
|
// Only compiled when CUDA is available; the header is safe to include unconditionally (the impl behind
|
|
// the pimpl is null without CUDA, and RotationScaleMerge falls back to the CPU loops).
|
|
class RotationScaleMergeGPU {
|
|
public:
|
|
RotationScaleMergeGPU();
|
|
~RotationScaleMergeGPU();
|
|
RotationScaleMergeGPU(const RotationScaleMergeGPU &) = delete;
|
|
RotationScaleMergeGPU &operator=(const RotationScaleMergeGPU &) = delete;
|
|
|
|
// True if a GPU was found and the engine is usable.
|
|
[[nodiscard]] bool Available() const;
|
|
|
|
// Upload the immutable per-observation fields (once). Arrays are length n_obs unless noted; the
|
|
// frame CSR (frame_start/frame_count) is length n_frames and indexes the obs arrays in frame order.
|
|
void SetPartials(int n_obs, int n_frames,
|
|
const float *I, const float *sigma, const float *rlp, const float *partiality,
|
|
const float *zeta, const uint8_t *on_ice, const int32_t *frame,
|
|
const float *corr0,
|
|
const int32_t *frame_start, const int32_t *frame_count);
|
|
|
|
// Per space group: the dense ASU-group id per obs, and a group-ordered permutation of the obs whose
|
|
// group >= 0 (group_perm), with its CSR (group_start/group_count, length n_groups) - so each group's
|
|
// observations are a contiguous, fixed-order segment for the reduction.
|
|
void SetGroups(int n_groups, const int32_t *group,
|
|
const int32_t *group_perm, int n_group_perm,
|
|
const int32_t *group_start, const int32_t *group_count);
|
|
|
|
// Re-upload the working corr (length n_obs) before a scaling pass (the host mutates it via smooth-G
|
|
// between passes). SetPartials uploads the initial corr; this refreshes it.
|
|
void SetCorr(const float *corr);
|
|
|
|
// Run `iters` of {reduce group means -> per-frame robust IRLS G -> update corr} on the device,
|
|
// in place on the resident corr. Rotation model (partiality folded via the stored partiality).
|
|
void ScalePartials(int iters, double robust_k, double min_partiality, bool has_d_min);
|
|
|
|
// Copy the updated corr back to the host (length n_obs), and the fitted per-frame G (length
|
|
// n_frames, as double) plus the per-frame "was fitted" flag.
|
|
void GetCorr(float *corr_out) const;
|
|
void GetG(double *g_out, uint8_t *scaled_out) const;
|
|
|
|
private:
|
|
struct Impl;
|
|
std::unique_ptr<Impl> impl_;
|
|
};
|