RotationScaleMerge: GPU partial-scaling loop (CUDA port, phase 1)
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>
This commit is contained in:
Claude Opus 4.8
@@ -15,4 +15,9 @@ ADD_LIBRARY(JFJochScaleMerge
|
||||
HKLKey.h
|
||||
ScalingResult.h
|
||||
ScalingResult.cpp)
|
||||
TARGET_LINK_LIBRARIES(JFJochScaleMerge Ceres::ceres Eigen3::Eigen JFJochCommon)
|
||||
TARGET_LINK_LIBRARIES(JFJochScaleMerge Ceres::ceres Eigen3::Eigen JFJochCommon)
|
||||
|
||||
IF (JFJOCH_CUDA_AVAILABLE)
|
||||
TARGET_SOURCES(JFJochScaleMerge PRIVATE ../indexing/CUDAMemHelpers.h
|
||||
RotationScaleMergeGPU.cu RotationScaleMergeGPU.h)
|
||||
ENDIF()
|
||||
@@ -238,6 +238,27 @@ void RotationScaleMerge::Ingest() {
|
||||
total, n_frames, rawrun_start.size());
|
||||
|
||||
SmoothMosaicityAndPartiality();
|
||||
|
||||
#ifdef JFJOCH_USE_CUDA
|
||||
// Bring the partial-scaling loop onto the GPU when one is present. Upload the immutable per-obs
|
||||
// fields once (corr lives on the device, refreshed each pass); the CPU keeps the sort/keying/combine.
|
||||
gpu_ = std::make_unique<RotationScaleMergeGPU>();
|
||||
gpu_active_ = gpu_->Available();
|
||||
if (gpu_active_) {
|
||||
const int n = static_cast<int>(partials.size());
|
||||
std::vector<float> I(n), sigma(n), rlp(n), part(n), zeta(n), corr(n);
|
||||
std::vector<uint8_t> onice(n);
|
||||
std::vector<int32_t> frm(n);
|
||||
for (int i = 0; i < n; ++i) {
|
||||
const auto &o = partials[i];
|
||||
I[i] = o.I; sigma[i] = o.sigma; rlp[i] = o.rlp; part[i] = o.partiality;
|
||||
zeta[i] = o.zeta; onice[i] = o.on_ice; frm[i] = o.frame; corr[i] = o.corr;
|
||||
}
|
||||
gpu_->SetPartials(n, n_frames, I.data(), sigma.data(), rlp.data(), part.data(), zeta.data(),
|
||||
onice.data(), frm.data(), corr.data(), frame_start.data(), frame_count.data());
|
||||
logger.Info("RotationScaleMerge: GPU partial-scaling active");
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
void RotationScaleMerge::SmoothMosaicityAndPartiality() {
|
||||
@@ -340,6 +361,28 @@ int RotationScaleMerge::ComputeAsuGroups(const HKLKeyGenerator &keygen) {
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
#ifdef JFJOCH_USE_CUDA
|
||||
// Group-ordered permutation (obs bucketed by ASU group, obs-index order) + its CSR, so the GPU
|
||||
// reduction is a deterministic segmented reduction (fixed order, no atomics). Built per space group.
|
||||
if (gpu_active_) {
|
||||
const int n = static_cast<int>(partials.size());
|
||||
std::vector<int32_t> group_ids(n), gcount(n_groups, 0);
|
||||
for (int i = 0; i < n; ++i) {
|
||||
group_ids[i] = partials[i].group;
|
||||
if (partials[i].group >= 0) ++gcount[partials[i].group];
|
||||
}
|
||||
std::vector<int32_t> gstart(n_groups, 0);
|
||||
int acc = 0;
|
||||
for (int g = 0; g < n_groups; ++g) { gstart[g] = acc; acc += gcount[g]; }
|
||||
std::vector<int32_t> gperm(acc), gfill = gstart;
|
||||
for (int i = 0; i < n; ++i) {
|
||||
const int g = partials[i].group;
|
||||
if (g >= 0) gperm[gfill[g]++] = i;
|
||||
}
|
||||
gpu_->SetGroups(n_groups, group_ids.data(), gperm.data(), acc, gstart.data(), gcount.data());
|
||||
}
|
||||
#endif
|
||||
return n_groups;
|
||||
}
|
||||
|
||||
@@ -1000,10 +1043,28 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search,
|
||||
const int n_groups = ComputeAsuGroups(keygen); // one ASU grouping, shared by partials and fulls
|
||||
lap("group hkl");
|
||||
std::vector<double> partial_mean;
|
||||
for (int it = 0; it < scaling_iter; ++it) {
|
||||
ReduceGroupMeans(partials, n_groups, false, {}, partial_mean);
|
||||
FitPerFrameG(partials, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial);
|
||||
UpdateCorr(partials, g_partial, frame_scaled_scratch);
|
||||
bool scaled_on_gpu = false;
|
||||
#ifdef JFJOCH_USE_CUDA
|
||||
if (gpu_active_) {
|
||||
// Refresh corr on the device (smooth-G mutated it on the host between passes), run the whole
|
||||
// scaling loop on the GPU, then read corr + per-frame G/scaled back.
|
||||
std::vector<float> corr(partials.size());
|
||||
for (size_t i = 0; i < partials.size(); ++i) corr[i] = partials[i].corr;
|
||||
gpu_->SetCorr(corr.data());
|
||||
gpu_->ScalePartials(scaling_iter, SCALE_ROBUST_K, min_partiality, d_min_limit.has_value());
|
||||
gpu_->GetCorr(corr.data());
|
||||
frame_scaled_scratch.assign(n_frames, 0);
|
||||
gpu_->GetG(g_partial.data(), frame_scaled_scratch.data());
|
||||
for (size_t i = 0; i < partials.size(); ++i) partials[i].corr = corr[i];
|
||||
scaled_on_gpu = true;
|
||||
}
|
||||
#endif
|
||||
if (!scaled_on_gpu) {
|
||||
for (int it = 0; it < scaling_iter; ++it) {
|
||||
ReduceGroupMeans(partials, n_groups, false, {}, partial_mean);
|
||||
FitPerFrameG(partials, frame_start, frame_count, partial_mean, /*unity=*/false, g_partial);
|
||||
UpdateCorr(partials, g_partial, frame_scaled_scratch);
|
||||
}
|
||||
}
|
||||
lap("scale partials");
|
||||
const std::vector<uint8_t> partial_scaled = frame_scaled_scratch;
|
||||
|
||||
@@ -14,6 +14,10 @@
|
||||
#include "../IntegrationOutcome.h"
|
||||
|
||||
#include "Merge.h" // MergedReflection, MergeStatistics
|
||||
#ifdef JFJOCH_USE_CUDA
|
||||
#include <memory>
|
||||
#include "RotationScaleMergeGPU.h"
|
||||
#endif
|
||||
|
||||
// Dedicated, allocate-once scale+combine+merge for rotation data (the -P rot3d path).
|
||||
//
|
||||
@@ -124,6 +128,13 @@ private:
|
||||
// Working per-group arrays (sized to the current group count; reused).
|
||||
std::vector<int32_t> group_h, group_k, group_l;
|
||||
|
||||
#ifdef JFJOCH_USE_CUDA
|
||||
// GPU engine for the partial-scaling loop (segmented reduce + per-frame IRLS + corr update). Null /
|
||||
// inactive when no GPU; the CPU loops are the fallback. Built in Ingest.
|
||||
std::unique_ptr<RotationScaleMergeGPU> gpu_;
|
||||
bool gpu_active_ = false;
|
||||
#endif
|
||||
|
||||
// --- helpers (each a flat pass; see the .cpp) ---
|
||||
// Compute the dense ASU-group id for the current space group by grouping the (pre-sorted) raw-hkl
|
||||
// runs by their ASU key - one gemmi ASU reduction per distinct raw hkl, not per observation. Fills
|
||||
|
||||
@@ -0,0 +1,289 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#include "RotationScaleMergeGPU.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <string>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include "../indexing/CUDAMemHelpers.h"
|
||||
#include "../../common/CUDAWrapper.h"
|
||||
#include "../../common/JFJochException.h"
|
||||
|
||||
namespace {
|
||||
constexpr int BLK = 256;
|
||||
constexpr int MIN_REFLECTIONS = 20;
|
||||
|
||||
__device__ __forceinline__ double SafeInvD(double x, double fallback) {
|
||||
return (isfinite(x) && x != 0.0) ? 1.0 / x : fallback;
|
||||
}
|
||||
|
||||
// Block reduction of a double, deterministic for a fixed thread->element mapping (fixed order,
|
||||
// no atomics). Returns the sum on thread 0; `s` is BLK doubles of shared scratch.
|
||||
__device__ double BlockReduceSum(double v, double *s) {
|
||||
const int t = threadIdx.x;
|
||||
s[t] = v;
|
||||
__syncthreads();
|
||||
for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
|
||||
if (t < stride) s[t] += s[t + stride];
|
||||
__syncthreads();
|
||||
}
|
||||
return s[0];
|
||||
}
|
||||
|
||||
// One thread per ASU group (grid-stride): inverse-variance mean of I*corr over the group's contiguous,
|
||||
// fixed-order segment of group_perm. Groups are small (avg tens of obs), so a whole block per group
|
||||
// wastes threads; summing each group in one thread avoids launch/sync overhead and stays deterministic
|
||||
// (fixed group_perm order). Matches the CPU ReduceGroupMeans filter (no ice/cell mask - scaling ref).
|
||||
__global__ void ReduceGroupMeansKernel(int n_groups, double min_partiality,
|
||||
const int32_t *__restrict__ group_perm,
|
||||
const int32_t *__restrict__ group_start,
|
||||
const int32_t *__restrict__ group_count,
|
||||
const float *__restrict__ I, const float *__restrict__ sigma,
|
||||
const float *__restrict__ partiality,
|
||||
const float *__restrict__ corr,
|
||||
double *__restrict__ group_mean) {
|
||||
for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < n_groups; g += gridDim.x * blockDim.x) {
|
||||
const int lo = group_start[g], hi = group_start[g] + group_count[g];
|
||||
double sw = 0.0, swI = 0.0;
|
||||
for (int p = lo; p < hi; ++p) {
|
||||
const int i = group_perm[p];
|
||||
const float c = corr[i];
|
||||
if (!(c > 0.0f) || !isfinite(c)) continue;
|
||||
if (partiality[i] < min_partiality) continue;
|
||||
const float I_corr = I[i] * c;
|
||||
const float sigma_corr = sigma[i] * c;
|
||||
if (!isfinite(I_corr) || !isfinite(sigma_corr) || sigma_corr <= 0.0f) continue;
|
||||
const double w = 1.0 / (double(sigma_corr) * sigma_corr);
|
||||
sw += w;
|
||||
swI += w * I_corr;
|
||||
}
|
||||
group_mean[g] = sw > 0.0 ? swI / sw : NAN;
|
||||
}
|
||||
}
|
||||
|
||||
// Per-observation scale-fit coefficient (rotation model) and accept flag, recomputed each scaling
|
||||
// iteration once the group means are known. coeff = partiality * (1/rlp) * mean[group].
|
||||
__global__ void PrepScaleObsKernel(int n_obs, const int32_t *__restrict__ group,
|
||||
const float *__restrict__ partiality, const float *__restrict__ rlp,
|
||||
const float *__restrict__ zeta, const uint8_t *__restrict__ on_ice,
|
||||
const double *__restrict__ group_mean,
|
||||
float *__restrict__ sco_coeff, uint8_t *__restrict__ sco_ok) {
|
||||
const int i = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (i >= n_obs) return;
|
||||
const int g = group[i];
|
||||
bool ok = (g >= 0) && !on_ice[i] && isfinite(zeta[i]) && zeta[i] > 0.0f;
|
||||
double mean = 0.0;
|
||||
if (ok) {
|
||||
mean = group_mean[g];
|
||||
ok = isfinite(mean);
|
||||
}
|
||||
sco_ok[i] = ok ? 1 : 0;
|
||||
sco_coeff[i] = ok ? float(double(partiality[i]) * SafeInvD(rlp[i], 1.0) * mean) : 0.0f;
|
||||
}
|
||||
|
||||
// One block per frame: robust per-frame scale G by IRLS (Cauchy), over the frame's contiguous obs.
|
||||
// Identical objective to the CPU SolveScaleIRLS. Leaves g/scaled untouched for under-populated frames.
|
||||
__global__ void FitPerFrameGKernel(int n_frames, double robust_k,
|
||||
const int32_t *__restrict__ frame_start,
|
||||
const int32_t *__restrict__ frame_count,
|
||||
const float *__restrict__ I, const float *__restrict__ sigma,
|
||||
const float *__restrict__ sco_coeff, const uint8_t *__restrict__ sco_ok,
|
||||
double *__restrict__ g, uint8_t *__restrict__ scaled) {
|
||||
const int f = blockIdx.x;
|
||||
if (f >= n_frames) return;
|
||||
const int lo = frame_start[f], hi = frame_start[f] + frame_count[f];
|
||||
__shared__ double sh[BLK];
|
||||
|
||||
// count accepted
|
||||
long cnt_local = 0;
|
||||
for (int i = lo + threadIdx.x; i < hi; i += blockDim.x)
|
||||
if (sco_ok[i]) ++cnt_local;
|
||||
const double cnt = BlockReduceSum(double(cnt_local), sh);
|
||||
__shared__ double s_cnt;
|
||||
if (threadIdx.x == 0) s_cnt = cnt;
|
||||
__syncthreads();
|
||||
if (s_cnt < MIN_REFLECTIONS) return; // leave g[f]/scaled[f] as-is
|
||||
|
||||
const double k2 = robust_k * robust_k;
|
||||
|
||||
// seed: plain weighted-LS ratio (robust weight = 1)
|
||||
double num = 0.0, den = 0.0;
|
||||
for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) {
|
||||
if (!sco_ok[i]) continue;
|
||||
const double coeff = sco_coeff[i];
|
||||
const double w = SafeInvD(sigma[i], 1.0);
|
||||
const double w2 = w * w;
|
||||
num += w2 * coeff * double(I[i]);
|
||||
den += w2 * coeff * coeff;
|
||||
}
|
||||
double tnum = BlockReduceSum(num, sh); __syncthreads();
|
||||
double tden = BlockReduceSum(den, sh);
|
||||
__shared__ double s_G;
|
||||
if (threadIdx.x == 0) {
|
||||
double G = tden > 0.0 ? tnum / tden : NAN;
|
||||
s_G = isfinite(G) ? fmax(0.0, G) : 1.0;
|
||||
}
|
||||
__syncthreads();
|
||||
if (!(s_G >= 0.0)) { if (threadIdx.x == 0) { g[f] = 1.0; scaled[f] = 1; } return; }
|
||||
|
||||
for (int iter = 0; iter < 30; ++iter) {
|
||||
const double G = s_G;
|
||||
num = 0.0; den = 0.0;
|
||||
for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) {
|
||||
if (!sco_ok[i]) continue;
|
||||
const double coeff = sco_coeff[i];
|
||||
const double w = SafeInvD(sigma[i], 1.0);
|
||||
const double w2 = w * w;
|
||||
const double res = w * (G * coeff - double(I[i]));
|
||||
const double rw = 1.0 / (1.0 + res * res / k2);
|
||||
num += rw * w2 * coeff * double(I[i]);
|
||||
den += rw * w2 * coeff * coeff;
|
||||
}
|
||||
tnum = BlockReduceSum(num, sh); __syncthreads();
|
||||
tden = BlockReduceSum(den, sh);
|
||||
if (threadIdx.x == 0) {
|
||||
const double G_next = tden > 0.0 ? tnum / tden : NAN;
|
||||
if (isfinite(G_next)) {
|
||||
const double Gn = fmax(0.0, G_next);
|
||||
s_cnt = (fabs(Gn - G) <= 1e-7 * fmax(Gn, 1.0)) ? 1.0 : 0.0; // reuse s_cnt as "converged"
|
||||
s_G = Gn;
|
||||
} else {
|
||||
s_cnt = 1.0;
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
if (s_cnt != 0.0) break;
|
||||
}
|
||||
if (threadIdx.x == 0) { g[f] = s_G; scaled[f] = 1; }
|
||||
}
|
||||
|
||||
// corr = rlp / (partiality * G[frame]) for fitted frames; unchanged otherwise (grid-stride).
|
||||
__global__ void UpdateCorrKernel(int n_obs, const int32_t *__restrict__ frame,
|
||||
const float *__restrict__ rlp, const float *__restrict__ partiality,
|
||||
const double *__restrict__ g, const uint8_t *__restrict__ scaled,
|
||||
float *__restrict__ corr) {
|
||||
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
|
||||
const int f = frame[i];
|
||||
if (!scaled[f]) continue;
|
||||
const double denom = double(partiality[i]) * g[f];
|
||||
corr[i] = (isfinite(double(rlp[i])) && isfinite(denom) && denom > 0.0)
|
||||
? float(rlp[i] / denom) : NAN;
|
||||
}
|
||||
}
|
||||
|
||||
void CudaCheck(cudaError_t e, const char *what) {
|
||||
if (e != cudaSuccess)
|
||||
throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
|
||||
std::string("RotationScaleMergeGPU: ") + what + ": " + cudaGetErrorString(e));
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void Upload(CudaDevicePtr<T> &dst, const T *src, int n) {
|
||||
dst = CudaDevicePtr<T>(std::max(1, n));
|
||||
if (n > 0)
|
||||
CudaCheck(cudaMemcpy(dst.get(), src, size_t(n) * sizeof(T), cudaMemcpyHostToDevice), "upload");
|
||||
}
|
||||
}
|
||||
|
||||
struct RotationScaleMergeGPU::Impl {
|
||||
bool available = false;
|
||||
int n_obs = 0, n_frames = 0, n_groups = 0;
|
||||
|
||||
// immutable per-obs
|
||||
CudaDevicePtr<float> I, sigma, rlp, partiality, zeta;
|
||||
CudaDevicePtr<uint8_t> on_ice;
|
||||
CudaDevicePtr<int32_t> frame;
|
||||
CudaDevicePtr<float> corr; // mutable, resident across iterations
|
||||
CudaDevicePtr<int32_t> frame_start, frame_count;
|
||||
// per space group
|
||||
CudaDevicePtr<int32_t> group, group_perm, group_start, group_count;
|
||||
// scratch
|
||||
CudaDevicePtr<double> group_mean, g;
|
||||
CudaDevicePtr<uint8_t> scaled;
|
||||
CudaDevicePtr<float> sco_coeff;
|
||||
CudaDevicePtr<uint8_t> sco_ok;
|
||||
};
|
||||
|
||||
RotationScaleMergeGPU::RotationScaleMergeGPU() : impl_(std::make_unique<Impl>()) {
|
||||
if (get_gpu_count() > 0) {
|
||||
set_gpu(0);
|
||||
impl_->available = true;
|
||||
}
|
||||
}
|
||||
|
||||
RotationScaleMergeGPU::~RotationScaleMergeGPU() = default;
|
||||
|
||||
bool RotationScaleMergeGPU::Available() const { return impl_->available; }
|
||||
|
||||
void RotationScaleMergeGPU::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) {
|
||||
auto &d = *impl_;
|
||||
d.n_obs = n_obs;
|
||||
d.n_frames = n_frames;
|
||||
Upload(d.I, I, n_obs); Upload(d.sigma, sigma, n_obs); Upload(d.rlp, rlp, n_obs);
|
||||
Upload(d.partiality, partiality, n_obs); Upload(d.zeta, zeta, n_obs); Upload(d.on_ice, on_ice, n_obs);
|
||||
Upload(d.frame, frame, n_obs); Upload(d.corr, corr0, n_obs);
|
||||
Upload(d.frame_start, frame_start, n_frames); Upload(d.frame_count, frame_count, n_frames);
|
||||
d.g = CudaDevicePtr<double>(n_frames);
|
||||
d.scaled = CudaDevicePtr<uint8_t>(n_frames);
|
||||
d.sco_coeff = CudaDevicePtr<float>(n_obs);
|
||||
d.sco_ok = CudaDevicePtr<uint8_t>(n_obs);
|
||||
}
|
||||
|
||||
void RotationScaleMergeGPU::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) {
|
||||
auto &d = *impl_;
|
||||
d.n_groups = n_groups;
|
||||
Upload(d.group, group, d.n_obs);
|
||||
Upload(d.group_perm, group_perm, n_group_perm); // obs with group >= 0, in group order
|
||||
Upload(d.group_start, group_start, n_groups);
|
||||
Upload(d.group_count, group_count, n_groups);
|
||||
d.group_mean = CudaDevicePtr<double>(std::max(1, n_groups));
|
||||
}
|
||||
|
||||
void RotationScaleMergeGPU::SetCorr(const float *corr) {
|
||||
CudaCheck(cudaMemcpy(impl_->corr.get(), corr, size_t(impl_->n_obs) * sizeof(float),
|
||||
cudaMemcpyHostToDevice), "upload corr");
|
||||
}
|
||||
|
||||
void RotationScaleMergeGPU::ScalePartials(int iters, double robust_k, double min_partiality,
|
||||
bool /*has_d_min*/) {
|
||||
auto &d = *impl_;
|
||||
CudaCheck(cudaMemset(d.scaled.get(), 0, size_t(d.n_frames) * sizeof(uint8_t)), "memset scaled");
|
||||
CudaCheck(cudaMemset(d.g.get(), 0, size_t(d.n_frames) * sizeof(double)), "memset g"); // unscaled g unused
|
||||
const int obs_blocks = (d.n_obs + BLK - 1) / BLK;
|
||||
const int upd_blocks = std::min(65535, obs_blocks);
|
||||
const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK);
|
||||
for (int it = 0; it < iters; ++it) {
|
||||
ReduceGroupMeansKernel<<<grp_blocks, BLK>>>(d.n_groups, min_partiality,
|
||||
d.group_perm.get(), d.group_start.get(), d.group_count.get(),
|
||||
d.I.get(), d.sigma.get(), d.partiality.get(), d.corr.get(), d.group_mean.get());
|
||||
PrepScaleObsKernel<<<obs_blocks, BLK>>>(d.n_obs, d.group.get(), d.partiality.get(), d.rlp.get(),
|
||||
d.zeta.get(), d.on_ice.get(), d.group_mean.get(), d.sco_coeff.get(), d.sco_ok.get());
|
||||
FitPerFrameGKernel<<<d.n_frames, BLK>>>(d.n_frames, robust_k, d.frame_start.get(), d.frame_count.get(),
|
||||
d.I.get(), d.sigma.get(), d.sco_coeff.get(), d.sco_ok.get(), d.g.get(), d.scaled.get());
|
||||
UpdateCorrKernel<<<upd_blocks, BLK>>>(d.n_obs, d.frame.get(), d.rlp.get(), d.partiality.get(),
|
||||
d.g.get(), d.scaled.get(), d.corr.get());
|
||||
}
|
||||
CudaCheck(cudaGetLastError(), "kernel launch");
|
||||
CudaCheck(cudaDeviceSynchronize(), "scale sync");
|
||||
}
|
||||
|
||||
void RotationScaleMergeGPU::GetCorr(float *corr_out) const {
|
||||
CudaCheck(cudaMemcpy(corr_out, impl_->corr.get(), size_t(impl_->n_obs) * sizeof(float),
|
||||
cudaMemcpyDeviceToHost), "download corr");
|
||||
}
|
||||
|
||||
void RotationScaleMergeGPU::GetG(double *g_out, uint8_t *scaled_out) const {
|
||||
CudaCheck(cudaMemcpy(g_out, impl_->g.get(), size_t(impl_->n_frames) * sizeof(double),
|
||||
cudaMemcpyDeviceToHost), "download g");
|
||||
CudaCheck(cudaMemcpy(scaled_out, impl_->scaled.get(), size_t(impl_->n_frames) * sizeof(uint8_t),
|
||||
cudaMemcpyDeviceToHost), "download scaled");
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
// 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_;
|
||||
};
|
||||
Reference in New Issue
Block a user