mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 66 Wiki Activity

RotationScaleMerge: GPU merge + error-model reductions over the resident fulls

Browse Source 
Port the four fulls-walking reductions of MergeAndStats to the GPU, over the fulls
group-CSR already resident from scale-fulls: the per-group inv-var mean + leverage-
corrected error-model samples, the merge accumulate (inv-var sums + deterministic
half-sets, error-model-corrected sigma, with outlier rejection), and R_meas + the
per-shell usable count. The host keeps the parts that don't parallelise cleanly or are
tiny: the I2-sort + 16-bin (a,b) median fit, the per-group reject median (a per-group
median is awkward on the GPU - cheap on the host from the GPU cnt), the merged export,
the shells and the gemmi completeness. Only per-group arrays (~55k) + the samples
(~n_fulls, for the fit) come back - the fulls are not re-walked on the host.

Device HalfForImage (splitmix64) + IceRingIndex mirror the host; the corrected-sigma
uses (b*I_for_b)^2 (not b^2*I^2) to match the host rounding; the R_meas usable count
requires finite d (the host counts only fulls with a valid shell, and a group's fulls
share d, so the shell is assigned per group). Gated on fulls_resident (GPU
combine+scale-fulls active); reject is fully supported so it runs for the default
rot3d command.

merge+stats ~0.49 -> ~0.37s, taking RSM on lyso to ~0.78s (was ~0.91). Validated across
the battery: 15/15 deterministic crystals bit-identical to the CPU path (SG / ISa /
CC1.2 / completeness / total-obs, and the exact outlier-reject count), only EP_cs_01-24
noise wobbles. The em-sort + a,b fit are the remaining host floor. Non-CUDA build
unaffected (use_gpu_merge is always false there).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
leonarski_f
2026-07-03 10:38:36 +02:00
co-authored by Claude Opus 4.8
parent 786edf9929
commit 34b3c3c4e7
4 changed files with 459 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/scale_merge/RotationScaleMerge.cpp
+131 -63
View File
@@ -271,6 +271,7 @@ void RotationScaleMerge::Ingest() {
gpu_->SetRawRuns(static_cast<int>(rawrun_start.size()), static_cast<int>(perm.size()), perm.data(),
rawrun_start.data(), rawrun_count.data(),
rawrun_h.data(), rawrun_k.data(), rawrun_l.data());
gpu_->SetFrameCellOk(frame_cell_ok.data());
gpu_combine_ = std::getenv("JFJOCH_RSM_CPU_COMBINE") == nullptr;
logger.Info("RotationScaleMerge: GPU partial-scaling{} active",
gpu_combine_ ? " + combine + scale-fulls" : "");
@@ -810,7 +811,8 @@ namespace {
}
RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool for_search,
const std::vector<char> &masked) {
const std::vector<char> &masked,
bool fulls_resident) {
// A full is usable for the merge / error model if it passes AddImage's filters (with the current
// ice/masked-ring context). group >= 0 already encodes "not absent and passes AcceptReflection".
auto masked_ring = [&](const Obs &o) {
@@ -829,24 +831,68 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool
return std::isfinite(I_corr) && std::isfinite(sigma_corr) && sigma_corr > 0.0f;
};
// The em-stats / samples / merge-accumulate / R_meas reductions run on the resident, scaled fulls
// (their group CSR is still on the device from scale-fulls) when fulls_resident; the host keeps the
// I2-sort, the (a,b) fit, the export and the statistics. reject_outliers is excluded upstream.
bool use_gpu_merge = false;
#ifdef JFJOCH_USE_CUDA
use_gpu_merge = fulls_resident && !fulls.empty();
std::vector<uint8_t> gpu_masked(masked.begin(), masked.end());
#endif
// ---- Error model: fit dev2 = a*sigma^2 + b^2*<I>^2 from symmetry-equivalent scatter. ----
std::vector<double> em_mean(n_groups, NAN);
std::vector<float> reject_median(n_groups, NAN);
double error_model_a = 1.0, error_model_b = 0.0, error_model_chi2 = 0.0;
bool error_model_active = false;
{
// Per-group inverse-variance mean over usable fulls (>=2 obs), and the leverage-corrected samples.
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
std::vector<int> cnt(n_groups, 0);
for (const auto &o : fulls) {
if (!usable_merge(o)) continue;
const double sigma_corr = static_cast<double>(o.sigma) * o.corr;
const double w = 1.0 / (sigma_corr * sigma_corr);
sw[o.group] += w; swI[o.group] += w * (static_cast<double>(o.I) * o.corr); cnt[o.group]++;
struct Sample { double s2, I2, dev2; };
std::vector<Sample> samples;
std::vector<int32_t> cnt(n_groups, 0); // per-group usable count (both paths; feeds reject-median)
bool did_gpu = false;
#ifdef JFJOCH_USE_CUDA
if (use_gpu_merge) {
const int nf = static_cast<int>(fulls.size());
std::vector<double> gs2(nf), gI2(nf), gdev2(nf);
std::vector<uint8_t> gvalid(nf);
gpu_->MergeEmSamples(for_search, gpu_masked.data(), static_cast<int>(gpu_masked.size()),
ice_half_width_q, min_partiality, em_mean.data(), cnt.data(),
gs2.data(), gI2.data(), gdev2.data(), gvalid.data());
samples.reserve(nf);
for (int i = 0; i < nf; ++i)
if (gvalid[i]) samples.push_back({gs2[i], gI2[i], gdev2[i]});
did_gpu = true;
}
for (int g = 0; g < n_groups; ++g)
if (cnt[g] >= 2 && sw[g] > 0.0) em_mean[g] = swI[g] / sw[g];
#endif
if (!did_gpu) {
// Per-group inverse-variance mean over usable fulls (>=2 obs), and the leverage-corrected samples.
std::vector<double> sw(n_groups, 0.0), swI(n_groups, 0.0);
for (const auto &o : fulls) {
if (!usable_merge(o)) continue;
const double sigma_corr = static_cast<double>(o.sigma) * o.corr;
const double w = 1.0 / (sigma_corr * sigma_corr);
sw[o.group] += w; swI[o.group] += w * (static_cast<double>(o.I) * o.corr); cnt[o.group]++;
}
for (int g = 0; g < n_groups; ++g)
if (cnt[g] >= 2 && sw[g] > 0.0) em_mean[g] = swI[g] / sw[g];
samples.reserve(fulls.size());
for (const auto &o : fulls) {
if (!usable_merge(o) || cnt[o.group] < 2) continue;
const double mean = em_mean[o.group];
if (!std::isfinite(mean)) continue;
const double sigma_corr = static_cast<double>(o.sigma) * o.corr;
const double s2 = sigma_corr * sigma_corr;
const double w = 1.0 / s2;
const double factor = 1.0 - w / sw[o.group];
if (factor < 0.05) continue;
const double resid = static_cast<double>(o.I) * o.corr - mean;
samples.push_back({s2, mean * mean, resid * resid / factor});
}
}
// Per-group outlier-rejection median of I*corr (host both paths - a per-group median is awkward on
// the GPU; cheap here, cnt >= 2 filter from the em pass). Fed to the merge accumulate.
if (reject_outliers) {
std::vector<std::vector<float>> iv(n_groups);
for (const auto &o : fulls)
@@ -858,23 +904,6 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool
reject_median[g] = iv[g][iv[g].size() / 2];
}
}
struct Sample { double s2, I2, dev2; };
std::vector<Sample> samples;
samples.reserve(fulls.size());
for (const auto &o : fulls) {
if (!usable_merge(o) || cnt[o.group] < 2) continue;
const double mean = em_mean[o.group];
if (!std::isfinite(mean)) continue;
const double sigma_corr = static_cast<double>(o.sigma) * o.corr;
const double s2 = sigma_corr * sigma_corr;
const double w = 1.0 / s2;
const double factor = 1.0 - w / sw[o.group];
if (factor < 0.05) continue;
const double resid = static_cast<double>(o.I) * o.corr - mean;
samples.push_back({s2, mean * mean, resid * resid / factor});
}
constexpr int n_bins = 16;
if (samples.size() >= static_cast<size_t>(8 * n_bins)) {
std::sort(samples.begin(), samples.end(),
@@ -929,30 +958,50 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool
+ (error_model_b * I_for_b) * (error_model_b * I_for_b);
return v > 0.0 ? static_cast<float>(std::sqrt(v)) : sigma_corr;
};
// ---- Merge: per-group inverse-variance sums with corrected sigma + deterministic half sets. ----
struct Accum { double swI = 0, sw = 0, swIh[2] = {0, 0}, swh[2] = {0, 0}; size_t nh[2] = {0, 0}; float d = NAN; };
std::vector<Accum> acc(n_groups);
size_t reject_count = 0;
for (const auto &o : fulls) {
if (!usable_merge(o)) continue;
const float I_corr = o.I * o.corr;
float sigma_corr = o.sigma * o.corr;
sigma_corr = corrected_sigma(I_corr, sigma_corr, o.group);
if (reject_outliers && error_model_active && std::isfinite(reject_median[o.group])
&& std::fabs(I_corr - reject_median[o.group]) > reject_nsigma * sigma_corr) {
++reject_count;
continue;
bool did_gpu_acc = false;
#ifdef JFJOCH_USE_CUDA
if (use_gpu_merge) {
std::vector<double> aswI(n_groups), asw(n_groups), aswIh0(n_groups), aswIh1(n_groups),
aswh0(n_groups), aswh1(n_groups), ad(n_groups);
std::vector<int32_t> anh0(n_groups), anh1(n_groups), arej(n_groups);
gpu_->MergeAccum(error_model_a, error_model_b, error_model_active,
reject_outliers, reject_nsigma, reject_median.data(),
aswI.data(), asw.data(), aswIh0.data(), aswIh1.data(),
aswh0.data(), aswh1.data(), anh0.data(), anh1.data(), ad.data(), arej.data());
for (int g = 0; g < n_groups; ++g) {
Accum &a = acc[g];
a.swI = aswI[g]; a.sw = asw[g]; a.swIh[0] = aswIh0[g]; a.swIh[1] = aswIh1[g];
a.swh[0] = aswh0[g]; a.swh[1] = aswh1[g];
a.nh[0] = static_cast<size_t>(anh0[g]); a.nh[1] = static_cast<size_t>(anh1[g]);
a.d = static_cast<float>(ad[g]);
reject_count += static_cast<size_t>(arej[g]);
}
const double w = 1.0 / (static_cast<double>(sigma_corr) * sigma_corr);
const double wI = w * I_corr;
const int half = HalfForImage(o.frame);
auto &a = acc[o.group];
a.swI += wI; a.sw += w;
a.swIh[half] += wI; a.swh[half] += w; a.nh[half]++;
if (!std::isfinite(a.d) && std::isfinite(o.d) && o.d > 0.0f) a.d = o.d;
did_gpu_acc = true;
}
#endif
if (!did_gpu_acc)
for (const auto &o : fulls) {
if (!usable_merge(o)) continue;
const float I_corr = o.I * o.corr;
float sigma_corr = o.sigma * o.corr;
sigma_corr = corrected_sigma(I_corr, sigma_corr, o.group);
if (reject_outliers && error_model_active && std::isfinite(reject_median[o.group])
&& std::fabs(I_corr - reject_median[o.group]) > reject_nsigma * sigma_corr) {
++reject_count;
continue;
}
const double w = 1.0 / (static_cast<double>(sigma_corr) * sigma_corr);
const double wI = w * I_corr;
const int half = HalfForImage(o.frame);
auto &a = acc[o.group];
a.swI += wI; a.sw += w;
a.swIh[half] += wI; a.swh[half] += w; a.nh[half]++;
if (!std::isfinite(a.d) && std::isfinite(o.d) && o.d > 0.0f) a.d = o.d;
}
// ---- Export merged reflections (+ resolution-shell R-free flags). ----
Result result;
result.isa = error_model_b > 0 ? 1.0 / error_model_b : 0.0;
@@ -997,7 +1046,6 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool
if (reject_count > 0)
logger.Info("Merge outlier rejection: dropped {} observations", reject_count);
// ---- Statistics (10 shells): completeness, multiplicity, <I/sigma>, R_meas, CC1/2. ----
constexpr int n_shells = 10;
float sd_min = std::numeric_limits<float>::max(), sd_max = 0.0f;
@@ -1045,22 +1093,43 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool
// |I_i - <I>| per reflection.
struct RmeasObs { double sum_abs_dev = 0, sum_I = 0; int n = 0, shell = -1; };
std::vector<RmeasObs> rmeas(n_groups);
for (const auto &o : fulls) {
if (o.group < 0) continue;
if (!frame_cell_ok[o.frame]) continue;
if (!(o.corr > 0.0f) || !std::isfinite(o.corr)) continue;
if (o.partiality < min_partiality) continue;
const float I_corr = o.I * o.corr, sigma_corr = o.sigma * o.corr;
if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) continue;
const auto shell = shells.GetShell(o.d);
if (!shell || *shell < 0 || *shell >= n_shells) continue;
sa[*shell].total_obs++;
if (std::isfinite(merged_I[o.group])) {
auto &r = rmeas[o.group];
r.sum_abs_dev += std::fabs(static_cast<double>(I_corr) - merged_I[o.group]);
r.sum_I += I_corr; r.n++; r.shell = *shell;
bool did_gpu_rmeas = false;
#ifdef JFJOCH_USE_CUDA
if (use_gpu_merge) {
// Per-group R_meas + usable count on the GPU; the shell is assigned per group (its fulls share d).
std::vector<double> rabsdev(n_groups), rsumI(n_groups);
std::vector<int32_t> rn(n_groups), rnusable(n_groups);
gpu_->MergeRmeas(merged_I.data(), rabsdev.data(), rsumI.data(), rn.data(), rnusable.data());
for (int g = 0; g < n_groups; ++g) {
if (rnusable[g] == 0) continue;
const auto shell = shells.GetShell(acc[g].d);
if (!shell || *shell < 0 || *shell >= n_shells) continue;
sa[*shell].total_obs += rnusable[g];
if (std::isfinite(merged_I[g]) && rn[g] > 0) {
auto &r = rmeas[g];
r.sum_abs_dev = rabsdev[g]; r.sum_I = rsumI[g]; r.n = rn[g]; r.shell = *shell;
}
}
did_gpu_rmeas = true;
}
#endif
if (!did_gpu_rmeas)
for (const auto &o : fulls) {
if (o.group < 0) continue;
if (!frame_cell_ok[o.frame]) continue;
if (!(o.corr > 0.0f) || !std::isfinite(o.corr)) continue;
if (o.partiality < min_partiality) continue;
const float I_corr = o.I * o.corr, sigma_corr = o.sigma * o.corr;
if (!std::isfinite(I_corr) || !std::isfinite(sigma_corr) || sigma_corr <= 0.0f) continue;
const auto shell = shells.GetShell(o.d);
if (!shell || *shell < 0 || *shell >= n_shells) continue;
sa[*shell].total_obs++;
if (std::isfinite(merged_I[o.group])) {
auto &r = rmeas[o.group];
r.sum_abs_dev += std::fabs(static_cast<double>(I_corr) - merged_I[o.group]);
r.sum_I += I_corr; r.n++; r.shell = *shell;
}
}
std::vector<double> rmeas_num(n_shells, 0.0), rmeas_den(n_shells, 0.0);
double rmeas_num_all = 0.0, rmeas_den_all = 0.0;
for (const auto &r : rmeas) {
@@ -1098,7 +1167,6 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool
overall.cc_half = cc_half_overall.GetCC();
overall.cc_ref = NAN;
overall.r_meas = rmeas_den_all > 0.0 ? rmeas_num_all / rmeas_den_all : NAN;
logger.Info("Merge complete ({} unique reflections)", result.merged.size());
return result;
}
@@ -1274,7 +1342,7 @@ RotationScaleMerge::Result RotationScaleMerge::Run(bool for_search,
lap("scale fulls");
// --- 5. Error model + merge + statistics. ---
auto r = MergeAndStats(n_groups, for_search, masked_ice_rings);
auto r = MergeAndStats(n_groups, for_search, masked_ice_rings, combined_on_gpu && scaled_fulls_on_gpu);
lap("merge+stats");
return r;
}
image_analysis/scale_merge/RotationScaleMerge.h
+4 -1
View File
@@ -192,5 +192,8 @@ private:
const std::vector<uint8_t> &frame_scaled);
// Error model + merge + statistics over the fulls (the last stage). n_groups is the fulls group count.
Result MergeAndStats(int n_groups, bool for_search, const std::vector<char> &masked_ice_rings);
// fulls_resident: the (scaled) fulls + their group CSR are still on the GPU, so the em-stats / samples
// / merge-accumulate / R_meas reductions run there (only per-group + samples come back).
Result MergeAndStats(int n_groups, bool for_search, const std::vector<char> &masked_ice_rings,
bool fulls_resident);
};
image_analysis/scale_merge/RotationScaleMergeGPU.cu
+298
View File
@@ -406,6 +406,172 @@ namespace {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) p[i] = v;
}
// ===== error-model + merge reductions over the resident, scaled fulls (mirror MergeAndStats) =====
// Device copy of common/Definitions.h IceRingIndex (only consulted when the merge has masked rings).
__device__ __forceinline__ int IceRingIndexDev(float d, float hw) {
if (!(d > 0.0f)) return -1;
const float two_pi = 6.283185307f;
const float q = two_pi / d;
const float rings[11] = {3.895f, 3.661f, 3.438f, 2.667f, 2.249f, 2.068f,
1.947f, 1.916f, 1.882f, 1.719f, 1.522f};
for (int i = 0; i < 11; ++i)
if (fabsf(q - two_pi / rings[i]) < hw) return i;
return -1;
}
// Deterministic CC1/2 half from the frame index (splitmix64), matching Merge.cpp / RSM HalfForImage.
__device__ __forceinline__ int HalfForImageDev(long long image_id) {
unsigned long long z = (unsigned long long)image_id + 0x9e3779b97f4a7c15ULL;
z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
z = z ^ (z >> 31);
return (int)(z & 1ULL);
}
struct MergeParams {
int n_groups;
double min_partiality, error_model_a, error_model_b, reject_nsigma;
float ice_hw;
int for_search, n_masked, error_model_active, reject_outliers;
const float *I, *sigma, *corr, *partiality, *d, *reject_median;
const int32_t *group, *frame;
const uint8_t *on_ice, *frame_cell_ok, *masked;
const int32_t *gperm, *gstart, *gcount;
const double *em_mean, *merged_I;
// outputs
double *sw, *swI, *em_mean_out;
int32_t *cnt;
double *s2, *I2, *dev2;
uint8_t *valid;
double *a_swI, *a_sw, *a_swIh0, *a_swIh1, *a_swh0, *a_swh1, *a_d;
int32_t *a_nh0, *a_nh1, *a_rejected;
double *r_absdev, *r_sumI;
int32_t *r_n, *r_nusable;
};
// A full passes the merge / error-model filter (mirrors MergeAndStats::usable_merge).
__device__ __forceinline__ bool MergeUsable(int i, const MergeParams &p) {
const int g = p.group[i];
if (g < 0) return false;
if (!p.frame_cell_ok[p.frame[i]]) return false;
const float c = p.corr[i];
if (!(c > 0.0f) || !isfinite(c)) return false;
if (p.for_search && p.on_ice[i]) return false;
if (p.n_masked > 0) {
const int ring = IceRingIndexDev(p.d[i], p.ice_hw);
if (ring >= 0 && ring < p.n_masked && p.masked[ring]) return false;
}
if (p.partiality[i] < p.min_partiality) return false;
const float I_corr = p.I[i] * c, sigma_corr = p.sigma[i] * c;
return isfinite(I_corr) && isfinite(sigma_corr) && sigma_corr > 0.0f;
}
// The looser R_meas filter (no ice / masked-ring / for_search - Mask = cell only).
__device__ __forceinline__ bool RmeasUsable(int i, const MergeParams &p) {
const int g = p.group[i];
if (g < 0) return false;
if (!p.frame_cell_ok[p.frame[i]]) return false;
const float c = p.corr[i];
if (!(c > 0.0f) || !isfinite(c)) return false;
if (p.partiality[i] < p.min_partiality) return false;
const float I_corr = p.I[i] * c, sigma_corr = p.sigma[i] * c;
return isfinite(I_corr) && isfinite(sigma_corr) && sigma_corr > 0.0f;
}
// One thread per group: inverse-variance sums over the group's usable fulls + the group mean (>=2).
__global__ void MergeEmStatsKernel(MergeParams p) {
for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
const int lo = p.gstart[g], hi = lo + p.gcount[g];
double sw = 0.0, swI = 0.0;
int cnt = 0;
for (int q = lo; q < hi; ++q) {
const int i = p.gperm[q];
if (!MergeUsable(i, p)) continue;
const double sigma_corr = double(p.sigma[i]) * p.corr[i];
const double w = 1.0 / (sigma_corr * sigma_corr);
sw += w; swI += w * (double(p.I[i]) * p.corr[i]); ++cnt;
}
p.sw[g] = sw; p.swI[g] = swI; p.cnt[g] = cnt;
p.em_mean_out[g] = (cnt >= 2 && sw > 0.0) ? swI / sw : NAN;
}
}
// One thread per full: the leverage-corrected error-model sample, or valid=0 if dropped.
__global__ void MergeSamplesKernel(int n_obs, MergeParams p) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) {
p.valid[i] = 0;
if (!MergeUsable(i, p)) continue;
const int g = p.group[i];
if (p.cnt[g] < 2) continue;
const double mean = p.em_mean[g];
if (!isfinite(mean)) continue;
const double sigma_corr = double(p.sigma[i]) * p.corr[i];
const double s2 = sigma_corr * sigma_corr;
const double factor = 1.0 - (1.0 / s2) / p.sw[g];
if (factor < 0.05) continue;
const double resid = double(p.I[i]) * p.corr[i] - mean;
p.s2[i] = s2; p.I2[i] = mean * mean; p.dev2[i] = resid * resid / factor; p.valid[i] = 1;
}
}
// One thread per group: the merge accumulators (inv-var sums + deterministic half-sets), with the
// error-model-corrected sigma. Mirrors MergeAndStats' merge loop (reject path stays on the host).
__global__ void MergeAccumKernel(MergeParams p) {
for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
const int lo = p.gstart[g], hi = lo + p.gcount[g];
double swI = 0, sw = 0, swIh0 = 0, swIh1 = 0, swh0 = 0, swh1 = 0, dd = NAN;
int nh0 = 0, nh1 = 0, rejected = 0;
const float rmed = p.reject_median[g];
for (int q = lo; q < hi; ++q) {
const int i = p.gperm[q];
if (!MergeUsable(i, p)) continue;
const float I_corr = p.I[i] * p.corr[i];
float sigma_corr = p.sigma[i] * p.corr[i];
if (p.error_model_active) {
const double I_for_b = isfinite(p.em_mean[g]) ? p.em_mean[g] : double(I_corr);
const double bi = p.error_model_b * I_for_b;
const double v = p.error_model_a * double(sigma_corr) * sigma_corr + bi * bi;
if (v > 0.0) sigma_corr = float(sqrt(v));
}
if (p.reject_outliers && p.error_model_active && isfinite(rmed)
&& fabsf(I_corr - rmed) > p.reject_nsigma * sigma_corr) { ++rejected; continue; }
const double w = 1.0 / (double(sigma_corr) * sigma_corr);
const double wI = w * I_corr;
const int half = HalfForImageDev(p.frame[i]);
swI += wI; sw += w;
if (half) { swIh1 += wI; swh1 += w; ++nh1; } else { swIh0 += wI; swh0 += w; ++nh0; }
if (!isfinite(dd) && isfinite(p.d[i]) && p.d[i] > 0.0f) dd = p.d[i];
}
p.a_swI[g] = swI; p.a_sw[g] = sw; p.a_swIh0[g] = swIh0; p.a_swIh1[g] = swIh1;
p.a_swh0[g] = swh0; p.a_swh1[g] = swh1; p.a_nh0[g] = nh0; p.a_nh1[g] = nh1; p.a_d[g] = dd;
p.a_rejected[g] = rejected;
}
}
// One thread per group: R_meas accumulators (sum|I_corr - merged_I|, sum_I, n) + usable count for the
// per-shell total_observations. Mirrors MergeAndStats' R_meas re-walk (looser, cell-only filter).
__global__ void MergeRmeasKernel(MergeParams p) {
for (int g = blockIdx.x * blockDim.x + threadIdx.x; g < p.n_groups; g += gridDim.x * blockDim.x) {
const int lo = p.gstart[g], hi = lo + p.gcount[g];
const double mI = p.merged_I[g];
const bool have = isfinite(mI);
double absdev = 0, sumI = 0;
int n = 0, nusable = 0;
for (int q = lo; q < hi; ++q) {
const int i = p.gperm[q];
if (!RmeasUsable(i, p)) continue;
if (!isfinite(p.d[i]) || !(p.d[i] > 0.0f)) continue; // host counts only fulls with a shell
++nusable;
if (have) {
const double I_corr = double(p.I[i]) * p.corr[i];
absdev += fabs(I_corr - mI); sumI += I_corr; ++n;
}
}
p.r_absdev[g] = absdev; p.r_sumI[g] = sumI; p.r_n[g] = n; p.r_nusable[g] = nusable;
}
}
void CudaCheck(cudaError_t e, const char *what) {
if (e != cudaSuccess)
throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
@@ -441,6 +607,20 @@ struct RotationScaleMergeGPU::Impl {
CudaDevicePtr<int64_t> cc_n;
CudaDevicePtr<uint8_t> smooth_apply; // per-frame smooth-G apply flag + ratio, length n_frames
CudaDevicePtr<double> smooth_ratio;
// merge / error-model reductions over the resident fulls (reuse the fulls group CSR f_gperm/...)
CudaDevicePtr<uint8_t> frame_cell_ok, merge_masked;
CudaDevicePtr<double> m_sw, m_swI, m_em_mean; // per group (n_groups)
CudaDevicePtr<int32_t> m_cnt;
CudaDevicePtr<double> m_s2, m_I2, m_dev2; // per full (n_fulls)
CudaDevicePtr<uint8_t> m_valid;
CudaDevicePtr<double> a_swI, a_sw, a_swIh0, a_swIh1, a_swh0, a_swh1, a_d; // merge accum per group
CudaDevicePtr<int32_t> a_nh0, a_nh1, a_rejected;
CudaDevicePtr<float> reject_median;
CudaDevicePtr<double> merged_I, r_absdev, r_sumI; // R_meas per group (merged_I uploaded)
CudaDevicePtr<int32_t> r_n, r_nusable;
int merge_for_search = 0, merge_n_masked = 0; // filter context for one MergeAndStats call
float merge_ice_hw = 0.0f;
double merge_min_part = 0.0;
// combine: extra per-obs inputs + the one-time raw-hkl run layout
CudaDevicePtr<float> bkg, image_number, d_obs;
@@ -544,6 +724,124 @@ void RotationScaleMergeGPU::GetG(double *g_out, uint8_t *scaled_out) const {
cudaMemcpyDeviceToHost), "download scaled");
}
void RotationScaleMergeGPU::SetFrameCellOk(const uint8_t *frame_cell_ok) {
Upload(impl_->frame_cell_ok, frame_cell_ok, impl_->n_frames);
}
// The per-group inv-var mean (em_mean) + the per-full leverage-corrected error-model samples over the
// resident+scaled fulls. Stashes the filter context for the later MergeAccum/MergeRmeas calls.
void RotationScaleMergeGPU::MergeEmSamples(bool for_search, const uint8_t *masked, int n_masked,
double ice_half_width_q, double min_partiality,
double *em_mean_out, int32_t *cnt_out, double *s2_out,
double *I2_out, double *dev2_out, uint8_t *valid_out) {
auto &d = *impl_;
const int ng = d.n_groups, nf = d.n_fulls;
d.merge_for_search = for_search ? 1 : 0; d.merge_n_masked = n_masked;
d.merge_ice_hw = float(ice_half_width_q); d.merge_min_part = min_partiality;
d.m_sw = CudaDevicePtr<double>(std::max(1, ng)); d.m_swI = CudaDevicePtr<double>(std::max(1, ng));
d.m_em_mean = CudaDevicePtr<double>(std::max(1, ng)); d.m_cnt = CudaDevicePtr<int32_t>(std::max(1, ng));
d.m_s2 = CudaDevicePtr<double>(std::max(1, nf)); d.m_I2 = CudaDevicePtr<double>(std::max(1, nf));
d.m_dev2 = CudaDevicePtr<double>(std::max(1, nf)); d.m_valid = CudaDevicePtr<uint8_t>(std::max(1, nf));
Upload(d.merge_masked, masked, n_masked);
MergeParams p{};
p.n_groups = ng; p.min_partiality = min_partiality; p.ice_hw = d.merge_ice_hw;
p.for_search = d.merge_for_search; p.n_masked = n_masked;
p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get(); p.partiality = d.f_partiality.get();
p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get();
p.on_ice = d.f_on_ice.get(); p.frame_cell_ok = d.frame_cell_ok.get(); p.masked = d.merge_masked.get();
p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();
p.em_mean = d.m_em_mean.get(); p.sw = d.m_sw.get(); p.swI = d.m_swI.get();
p.em_mean_out = d.m_em_mean.get(); p.cnt = d.m_cnt.get();
p.s2 = d.m_s2.get(); p.I2 = d.m_I2.get(); p.dev2 = d.m_dev2.get(); p.valid = d.m_valid.get();
const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK);
MergeEmStatsKernel<<<grp_blocks, BLK>>>(p);
MergeSamplesKernel<<<obs_blocks, BLK>>>(nf, p);
CudaCheck(cudaGetLastError(), "merge em/samples launch");
CudaCheck(cudaDeviceSynchronize(), "merge em/samples sync");
CudaCheck(cudaMemcpy(em_mean_out, d.m_em_mean.get(), size_t(ng) * sizeof(double),
cudaMemcpyDeviceToHost), "dl em_mean");
CudaCheck(cudaMemcpy(cnt_out, d.m_cnt.get(), size_t(ng) * sizeof(int32_t),
cudaMemcpyDeviceToHost), "dl cnt");
if (nf > 0) {
CudaCheck(cudaMemcpy(s2_out, d.m_s2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost), "dl s2");
CudaCheck(cudaMemcpy(I2_out, d.m_I2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost), "dl I2");
CudaCheck(cudaMemcpy(dev2_out, d.m_dev2.get(), size_t(nf) * sizeof(double), cudaMemcpyDeviceToHost), "dl dev2");
CudaCheck(cudaMemcpy(valid_out, d.m_valid.get(), size_t(nf) * sizeof(uint8_t), cudaMemcpyDeviceToHost), "dl valid");
}
}
void RotationScaleMergeGPU::MergeAccum(double error_model_a, double error_model_b, bool error_model_active,
bool reject_outliers, double reject_nsigma, const float *reject_median,
double *swI, double *sw, double *swIh0, double *swIh1,
double *swh0, double *swh1, int32_t *nh0, int32_t *nh1, double *d_out,
int32_t *rejected) {
auto &d = *impl_;
const int ng = d.n_groups;
d.a_swI = CudaDevicePtr<double>(std::max(1, ng)); d.a_sw = CudaDevicePtr<double>(std::max(1, ng));
d.a_swIh0 = CudaDevicePtr<double>(std::max(1, ng)); d.a_swIh1 = CudaDevicePtr<double>(std::max(1, ng));
d.a_swh0 = CudaDevicePtr<double>(std::max(1, ng)); d.a_swh1 = CudaDevicePtr<double>(std::max(1, ng));
d.a_nh0 = CudaDevicePtr<int32_t>(std::max(1, ng)); d.a_nh1 = CudaDevicePtr<int32_t>(std::max(1, ng));
d.a_d = CudaDevicePtr<double>(std::max(1, ng)); d.a_rejected = CudaDevicePtr<int32_t>(std::max(1, ng));
Upload(d.reject_median, reject_median, ng);
MergeParams p{};
p.n_groups = ng; p.min_partiality = d.merge_min_part; p.ice_hw = d.merge_ice_hw;
p.for_search = d.merge_for_search; p.n_masked = d.merge_n_masked;
p.error_model_a = error_model_a; p.error_model_b = error_model_b;
p.error_model_active = error_model_active ? 1 : 0;
p.reject_outliers = reject_outliers ? 1 : 0; p.reject_nsigma = reject_nsigma;
p.reject_median = d.reject_median.get();
p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get(); p.partiality = d.f_partiality.get();
p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get();
p.on_ice = d.f_on_ice.get(); p.frame_cell_ok = d.frame_cell_ok.get(); p.masked = d.merge_masked.get();
p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();
p.em_mean = d.m_em_mean.get();
p.a_swI = d.a_swI.get(); p.a_sw = d.a_sw.get(); p.a_swIh0 = d.a_swIh0.get(); p.a_swIh1 = d.a_swIh1.get();
p.a_swh0 = d.a_swh0.get(); p.a_swh1 = d.a_swh1.get(); p.a_nh0 = d.a_nh0.get(); p.a_nh1 = d.a_nh1.get();
p.a_d = d.a_d.get(); p.a_rejected = d.a_rejected.get();
const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
MergeAccumKernel<<<grp_blocks, BLK>>>(p);
CudaCheck(cudaGetLastError(), "merge accum launch");
CudaCheck(cudaDeviceSynchronize(), "merge accum sync");
auto dl = [&](void *h, const CudaDevicePtr<double> &s) {
CudaCheck(cudaMemcpy(h, s.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl accum"); };
dl(swI, d.a_swI); dl(sw, d.a_sw); dl(swIh0, d.a_swIh0); dl(swIh1, d.a_swIh1);
dl(swh0, d.a_swh0); dl(swh1, d.a_swh1); dl(d_out, d.a_d);
CudaCheck(cudaMemcpy(nh0, d.a_nh0.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl nh0");
CudaCheck(cudaMemcpy(nh1, d.a_nh1.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl nh1");
CudaCheck(cudaMemcpy(rejected, d.a_rejected.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl rej");
}
void RotationScaleMergeGPU::MergeRmeas(const double *merged_I, double *absdev, double *sumI,
int32_t *n, int32_t *nusable) {
auto &d = *impl_;
const int ng = d.n_groups;
Upload(d.merged_I, merged_I, ng);
d.r_absdev = CudaDevicePtr<double>(std::max(1, ng)); d.r_sumI = CudaDevicePtr<double>(std::max(1, ng));
d.r_n = CudaDevicePtr<int32_t>(std::max(1, ng)); d.r_nusable = CudaDevicePtr<int32_t>(std::max(1, ng));
MergeParams p{};
p.n_groups = ng; p.min_partiality = d.merge_min_part;
p.I = d.f_I.get(); p.sigma = d.f_sigma.get(); p.corr = d.f_corr.get(); p.partiality = d.f_partiality.get();
p.d = d.f_d.get(); p.group = d.f_group.get(); p.frame = d.f_frame.get(); p.frame_cell_ok = d.frame_cell_ok.get();
p.gperm = d.f_gperm.get(); p.gstart = d.f_gstart.get(); p.gcount = d.f_gcount.get();
p.merged_I = d.merged_I.get();
p.r_absdev = d.r_absdev.get(); p.r_sumI = d.r_sumI.get(); p.r_n = d.r_n.get(); p.r_nusable = d.r_nusable.get();
const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
MergeRmeasKernel<<<grp_blocks, BLK>>>(p);
CudaCheck(cudaGetLastError(), "merge rmeas launch");
CudaCheck(cudaDeviceSynchronize(), "merge rmeas sync");
CudaCheck(cudaMemcpy(absdev, d.r_absdev.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl absdev");
CudaCheck(cudaMemcpy(sumI, d.r_sumI.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost), "dl sumI");
CudaCheck(cudaMemcpy(n, d.r_n.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl rn");
CudaCheck(cudaMemcpy(nusable, d.r_nusable.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost), "dl rnusable");
}
void RotationScaleMergeGPU::SmoothCorr(const uint8_t *apply, const double *ratio) {
auto &d = *impl_;
Upload(d.smooth_apply, apply, d.n_frames);
image_analysis/scale_merge/RotationScaleMergeGPU.h
+26
View File
@@ -58,6 +58,32 @@ public:
// apply[f] (both length n_frames), matching CPU SmoothG. Keeps corr resident (no round-trip).
void SmoothCorr(const uint8_t *apply, const double *ratio);
// --- merge + error-model reductions over the resident, scaled fulls (reuse the fulls group CSR) ---
// The per-frame cell-consistency mask (length n_frames) used by the merge filter. Uploaded once.
void SetFrameCellOk(const uint8_t *frame_cell_ok);
// Per-group inv-var mean (em_mean, length n_groups) + per-full leverage-corrected error-model samples
// (s2/I2/dev2 + valid flag, length n_fulls), mirroring MergeAndStats' first two error-model loops.
// Stashes (for_search, masked, ice, min_partiality) for the MergeAccum/MergeRmeas calls that follow.
void MergeEmSamples(bool for_search, const uint8_t *masked, int n_masked,
double ice_half_width_q, double min_partiality,
double *em_mean_out, int32_t *cnt_out, double *s2_out, double *I2_out,
double *dev2_out, uint8_t *valid_out);
// Per-group merge accumulators (inv-var sums + deterministic half-sets, error-model-corrected sigma
// from a/b). Outputs length n_groups; rejected[g] counts outliers dropped (reject_median uploaded, NAN
// where none). Requires MergeEmSamples first (em_mean resident).
void MergeAccum(double error_model_a, double error_model_b, bool error_model_active,
bool reject_outliers, double reject_nsigma, const float *reject_median,
double *swI, double *sw, double *swIh0, double *swIh1,
double *swh0, double *swh1, int32_t *nh0, int32_t *nh1, double *d_out,
int32_t *rejected);
// Per-group R_meas accumulators (sum|I_corr-merged_I|, sum_I, n, and the usable count for per-shell
// total_observations); merged_I is uploaded. All arrays length n_groups.
void MergeRmeas(const double *merged_I, double *absdev, double *sumI, int32_t *n, int32_t *nusable);
// Post-smooth per-frame diagnostic CC: recompute the group means from the resident (smoothed) corr
// and the Pearson CC of each frame's I*corr vs its group mean, downloading only the per-frame cc /
// cc_n (length n_frames). Mirrors ReduceGroupMeans(partials) + FinalizePerFrameScale's CC loop.