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

Rotation scaling on the GPU: smaller device footprint, output unchanged

Browse Source 
The scaling engine held far more on the card than any step needed, and the
image passes' pooled memory was still reserved when it started.

- The image passes' allocation pool is trimmed before the scaling's arrays
  are allocated (they are not pooled): 4.3-8.8 GB of the 16 GB card was held
  by others when the scaling started, now 1.5-2.3 GB (context, cuFFT and the
  run's FFT indexers, which live for the whole run).
- A new grouping or combine frees the previous one's fulls, CSRs and merge
  results before allocating its own; the final merge of a large cell used to
  hold the P1 search's per-group arrays beside its own.
- The partials' group CSR goes at the combine; the arrays only the fulls'
  scaling and download read (records, zeta, frame CSR, group means, image
  number, px/py, clipped) go once the fulls are on the host.
- Merge buffers live only in the call that needs them: the error-model sums
  and sample places, the outlier medians, half-sets and CC factors, the R_meas
  sums; MergeAccum's per-group results go once downloaded (MergeAccumRelease).
- The combine's per-run scratch is local to the combine.
- Ingest: both radix sorts run on cub::DoubleBuffer (no cub copies of keys and
  values), the sort outputs are allocated after the raw-hkl key is gathered and
  freed, the uncut runs and old order go before the observations are cut, and
  the delta_phi re-solve takes the hkl from its run, so the per-observation
  64-bit key no longer outlives the sort.

Measured with per-array device accounting on an RTX 5080, engine peak after
the GPU/CPU decision (MB, before -> after; largest engine of each set):
  4nwv  32.7 M obs, 15.4 M hkl:  7536 -> 4265
  8qaw 122.7 M obs,  8.6 M hkl: 11716 -> 9327 (ingest 12533 -> 10756)
  9v16  54.7 M obs,  5.7 M hkl:  5831 -> 4283
  7qij  78.1 M obs,  2.2 M hkl:  6108 -> 5242 (ingest  7978 ->  6768)
  8sa8  50.0 M obs,  5.0 M hkl:  5490 -> 4205
The ingest's own peak went from ~104 to 82-100 bytes per reflection.

GPUScalingPeakBytes refitted on fourteen engines of these five sets to the
peak after the decision it makes (the ingest has passed by then): 81 B/obs +
291 B/hkl -> 63 + 115, within 12%. DeviceBytes refitted as an upper bound:
67/260/200 -> 66/94/52 per obs/event/hkl (up to 18% over), ingest
103/58 -> 89/56 per reflection/run. With it the all-observation arm now runs
beside the first merge on a 16 GB card for 4nwv, 9v16 and 8sa8.

p.hkl byte-identical to the rc175 baseline on 15 battery sets (the five above,
11if, 5reo, 6oel, 9jzo, 9min, metformin, lyso_x06da_ref, thau_bl1a_3p8keV,
insu_I_x06da_5keV_2, myob_x06da), p.mtz identical apart from the history;
battery compare shows no change. [RotationScale] and
ErrorModel_DeviceFitIsTheHostFit pass.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
leonarski_f
2026-10-15 09:49:11 +02:00
co-authored by Claude Opus 5.5
parent 228f0fc00f
commit c130c79236
4 changed files with 214 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/scale_merge/RotationScaleMerge.cpp
+20 -16
View File
@@ -720,30 +720,32 @@ bool RotationScaleMerge::OnGPU() const {
}
size_t RotationScaleMerge::GPUScalingPeakBytes() const {
// The device scaling's peak, fitted to what nsys measured on four large data sets (33-123 M
// observations, 4-16 M distinct hkl): within 5% on each. Unlike DeviceBytes, an estimate and not
// an upper bound - it decides where the GPU clearly cannot take the scaling.
constexpr size_t PEAK_BYTES_PER_OBS = 81, PEAK_BYTES_PER_HKL = 291;
// The device scaling's peak from this decision on, fitted to the engine's device buffers measured at
// their peak on fourteen engines of five large data sets (23-123 M observations, 1.6-15 M distinct
// hkl): within 12% on each, the furthest under a P1 merge of a large cell. The ingest's own peak, 82
// to 100 bytes per reflection, has passed by the time this decides. Unlike DeviceBytes, an estimate
// and not an upper bound - it decides where the GPU clearly cannot take the scaling.
constexpr size_t PEAK_BYTES_PER_OBS = 63, PEAK_BYTES_PER_HKL = 115;
return PEAK_BYTES_PER_OBS * static_cast<size_t>(n_partials_obs) + PEAK_BYTES_PER_HKL * rawrun_start.size();
}
size_t RotationScaleMerge::DeviceBytes() const {
// What each costs at the peak, measured: an observation's arrays and its group CSR; a full's arrays
// with the surface terms, or with the outlier test's and the error model's samples; an hkl's raw run
// and its ASU group's sums. The events are counted on the device path's ingest; elsewhere every
// observation may be an event of its own. Per observation, 13 bytes less than measured: the
// device scaling no longer keeps 1/sigma, the fit coefficient and its flag per observation.
constexpr size_t BYTES_PER_OBS = 67, BYTES_PER_EVENT = 260, BYTES_PER_HKL = 200;
// What each costs at the peak: an observation's arrays, and its group CSR until the combine; a full's
// arrays with its error-model sample, or with the surface terms; an hkl's raw run and its ASU group's
// sums. Fitted as an upper bound to the engine's device buffers measured at their peak on fourteen
// engines of five large data sets (23-123 M observations): up to 18% over. The events are counted
// on the device path's ingest; elsewhere every observation may be an event of its own.
constexpr size_t BYTES_PER_OBS = 66, BYTES_PER_EVENT = 94, BYTES_PER_HKL = 52;
const auto events = static_cast<size_t>(resident_ingest ? rocking_events_at_ingest : n_partials_obs);
const size_t bytes = BYTES_PER_OBS * static_cast<size_t>(n_partials_obs) + BYTES_PER_EVENT * events
+ BYTES_PER_HKL * rawrun_start.size();
// The resident ingest's own peak comes before any of that, over every reflection read and every run
// before the resolution cut, measured with cudaMemGetInfo. Beside the observation arrays, the sort key
// and the integrated delta_phi of each (66 bytes): the second radix sort holds the frame order, the
// permutation, both keys and cub's alternate key and value buffers, 102.6 bytes per reflection in
// all; the cut (IngestSort) 98 per reflection and 56 per run. The upload before them holds the 66
// bytes and the two 64 MB staging buffers.
constexpr size_t INGEST_BYTES_PER_REFLECTION = 103, INGEST_BYTES_PER_RUN = 58;
// before the resolution cut, measured on the same engines. Beside the observation arrays and the
// integrated delta_phi (58 bytes), each radix sort holds two pairs of keys and values, 82 bytes per
// reflection in all; the cut (IngestSort) of the runs 74 per reflection and 56 per run, that of the
// observations after it 89 per reflection and 24 per run kept. The upload before them holds the
// observation arrays, the sort key and delta_phi (66 bytes) and the two 64 MB staging buffers.
constexpr size_t INGEST_BYTES_PER_REFLECTION = 89, INGEST_BYTES_PER_RUN = 56;
constexpr size_t INGEST_ARRAY_BYTES = 66, INGEST_STAGING_BYTES = size_t(160) << 20;
const size_t ingest = std::max(INGEST_BYTES_PER_REFLECTION * ingest_reflections + INGEST_BYTES_PER_RUN * ingest_runs,
INGEST_ARRAY_BYTES * ingest_reflections + INGEST_STAGING_BYTES);
@@ -5335,6 +5337,7 @@ RotationScaleMerge::Result RotationScaleMerge::MergeAndStats(int n_groups, bool
rejected_total += rej;
});
}
gpu_->MergeAccumRelease();
reject_count = rejected_total;
did_gpu_acc = true;
}
@@ -6709,6 +6712,7 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool
gpu_->GetFullsPxPy(st.px.data(), st.py.data());
gpu_->GetFullsVariance(st.var_bkg.data(), st.var_per_I.data(), st.capture.data());
if (scaled_fulls_on_gpu) gpu_->GetFullsCorr(st.corr.data());
gpu_->FullsDownloaded();
ParallelChunks(nf, ThreadsForWork(nf, nthreads), [&](int lo, int hi) {
for (int i = lo; i < hi; ++i) {
Obs &o = fulls[i];
image_analysis/scale_merge/RotationScaleMergeGPU.cu
+182 -135
View File
@@ -1220,7 +1220,10 @@ struct RotationScaleMergeGPU::Impl {
void AllocObservations(int n);
// (frame_start null: the CSR is left to the caller to fill.)
void AllocFrames(const int32_t *frame_start, const int32_t *frame_count);
void AllocRunScratch();
// The previous grouping's fulls, their CSRs and what the merge left of them, before a new grouping
// or a new combine: nothing reads them again, and on a large cell they are gigabytes beside the
// arrays that replace them.
void ReleaseFulls();
CudaDevicePtr<float> zeta;
CudaDevicePtr<uint8_t> on_ice, clipped;
CudaDevicePtr<int32_t> frame;
@@ -1237,25 +1240,17 @@ struct RotationScaleMergeGPU::Impl {
CudaDevicePtr<double> smooth_ratio;
CudaDevicePtr<uint8_t> filter_reject; // per-frame --min-image-cc rejection flag, length n_frames
// merge / error-model reductions over the resident fulls (reuse the fulls group CSR f_gperm/...)
// What one merge keeps between its calls; every other merge buffer lives only in the call that
// needs it.
CudaDevicePtr<uint8_t> frame_cell_ok;
CudaDevicePtr<double> m_sw, m_swI, m_em_mean; // per group (n_groups)
CudaDevicePtr<int32_t> m_cnt;
CudaDevicePtr<double> m_swh, m_swIh; // per (group, hand), 2 * n_groups
CudaDevicePtr<int32_t> m_cnth;
CudaDevicePtr<double> m_em_mean; // per group (n_groups)
CudaDevicePtr<uint8_t> m_hand; // per full: Bijvoet hand (0 = I(+))
CudaDevicePtr<uint8_t> m_has_hands; // per group: acentric, and the merge pools the hands
CudaDevicePtr<int32_t> m_sample_at; // per full (n_fulls)
CudaDevicePtr<ErrorModelSample> m_sample; // the compacted error-model samples
CudaDevicePtr<uint8_t> m_rejected; // per-full outlier-rejected flag (MergeAccum -> MergeRmeas)
CudaDevicePtr<uint8_t> m_half; // per-full CC1/2 half-set, assigned on the host
CudaDevicePtr<double> a_swI, a_sw, a_swIh0, a_swIh1, a_swh0, a_swh1, a_swht0, a_swht1, a_d; // merge accum per group
CudaDevicePtr<double> cc_factor; // per-frame CC1/2 weight factor (length n_frames)
// merge accum per group, from MergeAccum to MergeAccumRelease
CudaDevicePtr<double> a_swI, a_sw, a_swIh0, a_swIh1, a_swh0, a_swh1, a_swht0, a_swht1, a_d;
CudaDevicePtr<int32_t> a_nh0, a_nh1, a_rejected;
CudaDevicePtr<uint8_t> a_on_ice; // per group: any contributing full on an ice ring
CudaDevicePtr<float> reject_median;
CudaDevicePtr<float> reject_var_add; // the pooled test's widening
CudaDevicePtr<double> merged_I, r_absdev, r_sumI, r_wabsdev, r_wsumI, r_sumv, r_sumv2; // R_meas per group (merged_I uploaded)
CudaDevicePtr<int32_t> r_n, r_nusable;
int merge_for_search = 0; // filter context for one MergeAndStats call
double merge_min_part = 0.0;
double merge_em_a = 1.0, merge_em_b = 0.0; // the error model of the last MergeAccum, for MergeRmeas
@@ -1264,8 +1259,7 @@ struct RotationScaleMergeGPU::Impl {
// combine: extra per-obs inputs + the one-time raw-hkl run layout
CudaDevicePtr<float> bkg, var_bkg, image_number, d_obs, px_obs, py_obs;
int n_runs = 0, n_perm = 0;
CudaDevicePtr<int32_t> perm, rr_start, rr_count, rr_h, rr_k, rr_l, rr_group;
CudaDevicePtr<int32_t> rr_nevents, rr_noverloaded, rr_offset;
CudaDevicePtr<int32_t> perm, rr_start, rr_count, rr_h, rr_k, rr_l;
// combine: resident fulls SoA (rebuilt each Combine)
int n_fulls = 0;
CudaDevicePtr<int32_t> f_h, f_k, f_l, f_frame, f_group;
@@ -1292,8 +1286,8 @@ struct RotationScaleMergeGPU::Impl {
CudaDevicePtr<double> s_A, s_sw, s_swI, s_tcross, s_tref2, s_cross, s_ref2;
CudaDevicePtr<double> s_w_Is, s_w_Iref, s_Iref; // per term of a subset, in segment order
// the resident ingest (IngestUpload .. IngestFinish): each frame's first reflection
// (host, length n_frames + 1), and per observation its raw hkl as the sort key and its delta_phi as
// integrated, until the smoothing has used them
// (host, length n_frames + 1), per observation its raw hkl as the sort key until the sort has used
// it, and its delta_phi as integrated until the smoothing has
std::vector<int32_t> in_src_start;
CudaDevicePtr<uint64_t> in_key;
CudaDevicePtr<float> in_dphi;
@@ -1358,6 +1352,11 @@ void RotationScaleMergeGPU::SetPartialsLayout(int n_obs, int n_frames,
}
void RotationScaleMergeGPU::Impl::AllocObservations(int n_in) {
// What the image passes' engines handed back to the allocation pool is still reserved by it when
// the scaling starts - 5 to 7 GB of a 16 GB card on a large data set, measured - and the scaling's
// buffers are not pooled, so it goes back to the driver here rather than only once an allocation
// has failed for want of it.
cuda_trim_allocation_pool();
const int n = std::max(1, n_in);
obs = Alloc<ObsRecord>(n);
zeta = Alloc<float>(n); corr = Alloc<float>(n);
@@ -1385,11 +1384,17 @@ void RotationScaleMergeGPU::Impl::AllocFrames(const int32_t *fstart, const int32
cc_n = Alloc<int64_t>(std::max(1, n_frames));
}
void RotationScaleMergeGPU::Impl::AllocRunScratch() {
rr_group = Alloc<int32_t>(std::max(1, n_runs));
rr_nevents = Alloc<int32_t>(std::max(1, n_runs));
rr_noverloaded = Alloc<int32_t>(std::max(1, n_runs));
rr_offset = Alloc<int32_t>(std::max(1, n_runs));
void RotationScaleMergeGPU::Impl::ReleaseFulls() {
n_fulls = 0;
f_h = {}; f_k = {}; f_l = {}; f_frame = {}; f_group = {};
f_I = {}; f_sigma = {}; f_d = {}; f_img = {}; f_px = {}; f_py = {};
f_var_bkg = {}; f_var_per_I = {}; f_capture = {}; f_on_ice = {}; f_clipped = {};
f_corr = {}; f_partiality = {}; f_zeta = {}; f_obs = {};
f_frame_perm = {}; f_frame_start = {}; f_frame_count = {};
f_gperm = {}; f_gstart = {}; f_gcount = {};
m_em_mean = {}; m_hand = {}; m_has_hands = {}; m_rejected = {};
a_swI = {}; a_sw = {}; a_swIh0 = {}; a_swIh1 = {}; a_swh0 = {}; a_swh1 = {}; a_swht0 = {}; a_swht1 = {};
a_d = {}; a_nh0 = {}; a_nh1 = {}; a_rejected = {}; a_on_ice = {};
}
namespace {
@@ -1462,7 +1467,6 @@ void RotationScaleMergeGPU::CopyObservationsFrom(const RotationScaleMergeGPU &fr
copy(d.rr_start, f.rr_start, d.n_runs); copy(d.rr_count, f.rr_count, d.n_runs);
copy(d.rr_h, f.rr_h, d.n_runs); copy(d.rr_k, f.rr_k, d.n_runs); copy(d.rr_l, f.rr_l, d.n_runs);
CudaCheck(cudaStreamSynchronize(d.s()), "copy observations");
d.AllocRunScratch();
}
void RotationScaleMergeGPU::SetObsFrame(int offset, int count, const int32_t *frame) {
@@ -1485,9 +1489,12 @@ void RotationScaleMergeGPU::BuildGroups(int n_groups, const int32_t *rawrun_grou
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
d.n_groups = n_groups;
CopyAndWait(d.rr_group.get(), rawrun_group, size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyHostToDevice, impl_->s(), "upload rr_group");
CudaDevicePtr<int32_t> sorted, first;
// The previous grouping goes before this one is allocated: the old CSRs (Combine frees the partials'
// one) and the fulls grouped by it.
d.ReleaseFulls();
d.group_perm = {}; d.group_start = {}; d.group_count = {}; d.group_mean = {};
CudaDevicePtr<int32_t> rr_group, sorted, first;
d.Upload(rr_group, rawrun_group, d.n_runs);
d.Upload(sorted, sorted_run, n_sorted);
d.Upload(first, group_first, n_groups + 1);
if (d.group.get() == nullptr) d.group = d.Alloc<int32_t>(std::max(1, d.n_obs));
@@ -1496,7 +1503,7 @@ void RotationScaleMergeGPU::BuildGroups(int n_groups, const int32_t *rawrun_grou
GroupParams p{};
p.n_runs = d.n_runs; p.n_groups = n_groups;
p.perm = d.perm.get(); p.rr_start = d.rr_start.get(); p.rr_count = d.rr_count.get(); p.rr_group = d.rr_group.get();
p.perm = d.perm.get(); p.rr_start = d.rr_start.get(); p.rr_count = d.rr_count.get(); p.rr_group = rr_group.get();
p.sorted_run = sorted.get(); p.group_first = first.get();
p.obs = d.obs.get();
p.group = d.group.get(); p.group_start = d.group_start.get(); p.group_count = d.group_count.get();
@@ -1591,22 +1598,23 @@ void RotationScaleMergeGPU::MergeEmSamples(bool for_search, double min_partialit
auto &d = *impl_;
const int ng = d.n_groups, nf = d.n_fulls;
d.merge_for_search = for_search ? 1 : 0; d.merge_min_part = min_partiality;
d.m_sw = d.Alloc<double>(std::max(1, ng)); d.m_swI = d.Alloc<double>(std::max(1, ng));
d.m_em_mean = d.Alloc<double>(std::max(1, ng)); d.m_cnt = d.Alloc<int32_t>(std::max(1, ng));
// Only em_mean and the hands outlive this call.
d.m_em_mean = d.Alloc<double>(std::max(1, ng));
CudaDevicePtr<double> sw = d.Alloc<double>(std::max(1, ng)), swI = d.Alloc<double>(std::max(1, ng));
CudaDevicePtr<int32_t> cnt = d.Alloc<int32_t>(std::max(1, ng));
// The error model is fitted on the Bijvoet hands (see the host obs_hand block). Absent - a merge
// that already separates them - the hand sums are not allocated and the kernels take the pooled
// group, which is what they did before.
CudaDevicePtr<double> swh, swIh;
CudaDevicePtr<int32_t> cnth;
d.m_hand = {}; d.m_has_hands = {};
if (hand && has_hands) {
d.m_swh = d.Alloc<double>(std::max(1, 2 * ng)); d.m_swIh = d.Alloc<double>(std::max(1, 2 * ng));
d.m_cnth = d.Alloc<int32_t>(std::max(1, 2 * ng));
swh = d.Alloc<double>(std::max(1, 2 * ng)); swIh = d.Alloc<double>(std::max(1, 2 * ng));
cnth = d.Alloc<int32_t>(std::max(1, 2 * ng));
if (nf > 0) d.Upload(d.m_hand, hand, nf);
if (ng > 0) d.Upload(d.m_has_hands, has_hands, ng);
} else {
d.m_swh = CudaDevicePtr<double>(); d.m_swIh = CudaDevicePtr<double>();
d.m_cnth = CudaDevicePtr<int32_t>();
d.m_hand = CudaDevicePtr<uint8_t>(); d.m_has_hands = CudaDevicePtr<uint8_t>();
}
d.m_sample_at = d.Alloc<int32_t>(std::max(1, nf));
CudaDevicePtr<int32_t> sample_at = d.Alloc<int32_t>(std::max(1, nf));
MergeParams p{};
p.n_groups = ng; p.min_partiality = min_partiality;
@@ -1614,14 +1622,13 @@ void RotationScaleMergeGPU::MergeEmSamples(bool for_search, double min_partialit
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.half = d.m_half.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.swh = d.m_swh.get(); p.swIh = d.m_swIh.get(); p.cnth = d.m_cnth.get();
p.em_mean = d.m_em_mean.get(); p.sw = sw.get(); p.swI = swI.get();
p.em_mean_out = d.m_em_mean.get(); p.cnt = cnt.get();
p.swh = swh.get(); p.swIh = swIh.get(); p.cnth = cnth.get();
p.hand = d.m_hand.get(); p.has_hands = d.m_has_hands.get();
p.var_bkg = d.f_var_bkg.get(); p.var_per_I = d.f_var_per_I.get(); p.capture = d.f_capture.get();
p.sample_at = d.m_sample_at.get();
p.sample_at = sample_at.get();
const int grp_blocks = std::min(65535, (ng + BLK - 1) / BLK);
const int obs_blocks = std::min(65535, (nf + BLK - 1) / BLK);
@@ -1632,32 +1639,31 @@ void RotationScaleMergeGPU::MergeEmSamples(bool for_search, double min_partialit
int n_samples = 0;
if (nf > 0) {
int32_t last_flag = 0;
CopyAndWait(&last_flag, d.m_sample_at.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost,
CopyAndWait(&last_flag, sample_at.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost,
impl_->s(), "dl last sample flag");
size_t bytes = 0;
CudaCheck(cub::DeviceScan::ExclusiveSum(nullptr, bytes, d.m_sample_at.get(), d.m_sample_at.get(), nf,
CudaCheck(cub::DeviceScan::ExclusiveSum(nullptr, bytes, sample_at.get(), sample_at.get(), nf,
impl_->s()), "sample scan size");
auto scratch = d.Alloc<uint8_t>(std::max<size_t>(bytes, 1));
CudaCheck(cub::DeviceScan::ExclusiveSum(scratch.get(), bytes, d.m_sample_at.get(), d.m_sample_at.get(), nf,
CudaCheck(cub::DeviceScan::ExclusiveSum(scratch.get(), bytes, sample_at.get(), sample_at.get(), nf,
impl_->s()), "sample scan");
CopyAndWait(&n_samples, d.m_sample_at.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost,
CopyAndWait(&n_samples, sample_at.get() + nf - 1, sizeof(int32_t), cudaMemcpyDeviceToHost,
impl_->s(), "dl sample count");
n_samples += last_flag;
}
d.m_sample = d.Alloc<ErrorModelSample>(std::max(1, n_samples));
p.sample = d.m_sample.get();
CudaDevicePtr<ErrorModelSample> sample = d.Alloc<ErrorModelSample>(std::max(1, n_samples));
p.sample = sample.get();
MergeSampleWriteKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(nf, p);
CudaCheck(cudaGetLastError(), "merge samples write launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "merge em/samples sync");
CopyAndWait(em_mean_out, d.m_em_mean.get(), size_t(ng) * sizeof(double),
cudaMemcpyDeviceToHost, impl_->s(), "dl em_mean");
CopyAndWait(cnt_out, d.m_cnt.get(), size_t(ng) * sizeof(int32_t),
CopyAndWait(cnt_out, cnt.get(), size_t(ng) * sizeof(int32_t),
cudaMemcpyDeviceToHost, impl_->s(), "dl cnt");
samples.resize(n_samples);
if (n_samples > 0)
CopyAndWait(samples.data(), d.m_sample.get(), size_t(n_samples) * sizeof(ErrorModelSample),
CopyAndWait(samples.data(), sample.get(), size_t(n_samples) * sizeof(ErrorModelSample),
cudaMemcpyDeviceToHost, impl_->s(), "dl samples");
d.m_sample = CudaDevicePtr<ErrorModelSample>();
}
void RotationScaleMergeGPU::MergeAccum(double error_model_a, double error_model_b, bool error_model_active,
@@ -1677,11 +1683,13 @@ void RotationScaleMergeGPU::MergeAccum(double error_model_a, double error_model_
d.a_d = d.Alloc<double>(std::max(1, ng)); d.a_rejected = d.Alloc<int32_t>(std::max(1, ng));
d.a_on_ice = d.Alloc<uint8_t>(std::max(1, ng));
d.Upload(d.m_rejected, rejected_obs, d.n_fulls);
d.Upload(d.reject_median, reject_median, ng);
if (reject_var_add && ng > 0) d.Upload(d.reject_var_add, reject_var_add, ng);
else d.reject_var_add = CudaDevicePtr<float>();
d.Upload(d.m_half, half, d.n_fulls);
d.Upload(d.cc_factor, frame_cc_factor, d.n_frames);
CudaDevicePtr<float> median, var_add;
d.Upload(median, reject_median, ng);
if (reject_var_add && ng > 0) d.Upload(var_add, reject_var_add, ng);
CudaDevicePtr<uint8_t> d_half;
d.Upload(d_half, half, d.n_fulls);
CudaDevicePtr<double> cc_factor;
d.Upload(cc_factor, frame_cc_factor, d.n_frames);
d.merge_em_a = error_model_a; d.merge_em_b = error_model_b; d.merge_em_active = error_model_active ? 1 : 0;
MergeParams p{};
@@ -1692,15 +1700,15 @@ void RotationScaleMergeGPU::MergeAccum(double error_model_a, double error_model_
p.reject_outliers = reject_outliers ? 1 : 0; p.reject_nsigma = reject_nsigma;
p.reject_up = std::expm1(reject_nsigma * reject_mult_sd);
p.reject_down = -std::expm1(-reject_nsigma * reject_mult_sd);
p.reject_median = d.reject_median.get();
p.reject_var_add = d.reject_var_add.get();
p.reject_median = median.get();
p.reject_var_add = var_add.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.half = d.m_half.get();
p.half = d_half.get();
p.var_bkg = d.f_var_bkg.get(); p.var_per_I = d.f_var_per_I.get(); p.capture = d.f_capture.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.cc_factor = d.cc_factor.get();
p.em_mean = d.m_em_mean.get(); p.cc_factor = cc_factor.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_swht0 = d.a_swht0.get(); p.a_swht1 = d.a_swht1.get(); p.a_nh0 = d.a_nh0.get(); p.a_nh1 = d.a_nh1.get();
@@ -1730,27 +1738,35 @@ void RotationScaleMergeGPU::MergeAccumRange(int g0, int n, double *swI, double *
dl(nh0, d.a_nh0); dl(nh1, d.a_nh1); dl(rejected, d.a_rejected); dl(on_ice_out, d.a_on_ice);
}
void RotationScaleMergeGPU::MergeAccumRelease() {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
d.a_swI = {}; d.a_sw = {}; d.a_swIh0 = {}; d.a_swIh1 = {}; d.a_swh0 = {}; d.a_swh1 = {};
d.a_swht0 = {}; d.a_swht1 = {}; d.a_d = {}; d.a_nh0 = {}; d.a_nh1 = {}; d.a_rejected = {}; d.a_on_ice = {};
}
void RotationScaleMergeGPU::MergeRmeas(const double *merged_I, double *absdev, double *sumI, double *wabsdev,
double *wsumI, double *sumv, double *sumv2, int32_t *n,
int32_t *nusable) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
const int ng = d.n_groups;
d.Upload(d.merged_I, merged_I, ng);
d.r_absdev = d.Alloc<double>(std::max(1, ng)); d.r_sumI = d.Alloc<double>(std::max(1, ng));
d.r_wabsdev = d.Alloc<double>(std::max(1, ng)); d.r_wsumI = d.Alloc<double>(std::max(1, ng));
d.r_sumv = d.Alloc<double>(std::max(1, ng)); d.r_sumv2 = d.Alloc<double>(std::max(1, ng));
d.r_n = d.Alloc<int32_t>(std::max(1, ng)); d.r_nusable = d.Alloc<int32_t>(std::max(1, ng));
CudaDevicePtr<double> d_merged_I;
d.Upload(d_merged_I, merged_I, ng);
CudaDevicePtr<double> r_absdev = d.Alloc<double>(std::max(1, ng)), r_sumI = d.Alloc<double>(std::max(1, ng)),
r_wabsdev = d.Alloc<double>(std::max(1, ng)), r_wsumI = d.Alloc<double>(std::max(1, ng)),
r_sumv = d.Alloc<double>(std::max(1, ng)), r_sumv2 = d.Alloc<double>(std::max(1, ng));
CudaDevicePtr<int32_t> r_n = d.Alloc<int32_t>(std::max(1, ng)), r_nusable = d.Alloc<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();
p.r_wabsdev = d.r_wabsdev.get(); p.r_wsumI = d.r_wsumI.get();
p.r_sumv = d.r_sumv.get(); p.r_sumv2 = d.r_sumv2.get();
p.merged_I = d_merged_I.get();
p.r_absdev = r_absdev.get(); p.r_sumI = r_sumI.get(); p.r_n = r_n.get(); p.r_nusable = r_nusable.get();
p.r_wabsdev = r_wabsdev.get(); p.r_wsumI = r_wsumI.get();
p.r_sumv = r_sumv.get(); p.r_sumv2 = r_sumv2.get();
p.error_model_a = d.merge_em_a; p.error_model_b = d.merge_em_b; p.error_model_active = d.merge_em_active;
p.em_mean = d.m_em_mean.get(); p.var_bkg = d.f_var_bkg.get(); p.var_per_I = d.f_var_per_I.get();
p.capture = d.f_capture.get();
@@ -1760,14 +1776,14 @@ void RotationScaleMergeGPU::MergeRmeas(const double *merged_I, double *absdev, d
MergeRmeasKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(p);
CudaCheck(cudaGetLastError(), "merge rmeas launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "merge rmeas sync");
CopyAndWait(absdev, d.r_absdev.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl absdev");
CopyAndWait(sumI, d.r_sumI.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sumI");
CopyAndWait(wabsdev, d.r_wabsdev.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl wabsdev");
CopyAndWait(wsumI, d.r_wsumI.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl wsumI");
CopyAndWait(sumv, d.r_sumv.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sumv");
CopyAndWait(sumv2, d.r_sumv2.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sumv2");
CopyAndWait(n, d.r_n.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(), "dl rn");
CopyAndWait(nusable, d.r_nusable.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(), "dl rnusable");
CopyAndWait(absdev, r_absdev.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl absdev");
CopyAndWait(sumI, r_sumI.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sumI");
CopyAndWait(wabsdev, r_wabsdev.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl wabsdev");
CopyAndWait(wsumI, r_wsumI.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl wsumI");
CopyAndWait(sumv, r_sumv.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sumv");
CopyAndWait(sumv2, r_sumv2.get(), size_t(ng) * sizeof(double), cudaMemcpyDeviceToHost, impl_->s(), "dl sumv2");
CopyAndWait(n, r_n.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(), "dl rn");
CopyAndWait(nusable, r_nusable.get(), size_t(ng) * sizeof(int32_t), cudaMemcpyDeviceToHost, impl_->s(), "dl rnusable");
}
void RotationScaleMergeGPU::MergeBijvoetExcess(double error_model_a, double error_model_b, float *group_d,
@@ -1941,7 +1957,6 @@ void RotationScaleMergeGPU::SetRawRuns(int n_runs, int n_perm, const int32_t *pe
d.Upload(d.rr_h, rr_h, n_runs);
d.Upload(d.rr_k, rr_k, n_runs);
d.Upload(d.rr_l, rr_l, n_runs);
d.AllocRunScratch();
}
int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_partiality,
@@ -1949,8 +1964,14 @@ int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_parti
float max_frame_gap, int64_t &n_overloaded) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
CopyAndWait(d.rr_group.get(), rawrun_group, size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyHostToDevice, impl_->s(), "upload rr_group");
// The partials' group CSR is read only by their scaling, which is over (BuildGroups makes the next
// one); the group means stay for the fulls' scaling. The previous fulls go before the new ones are
// allocated.
d.group = {}; d.group_perm = {}; d.group_start = {}; d.group_count = {};
d.ReleaseFulls();
CudaDevicePtr<int32_t> rr_group, rr_nevents = d.Alloc<int32_t>(std::max(1, d.n_runs)),
rr_noverloaded = d.Alloc<int32_t>(std::max(1, d.n_runs)), rr_offset;
d.Upload(rr_group, rawrun_group, d.n_runs);
CombineParams p{};
p.n_runs = d.n_runs;
@@ -1963,9 +1984,9 @@ int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_parti
p.px = d.px_obs.get(); p.py = d.py_obs.get();
p.frame = d.frame.get(); p.on_ice = d.on_ice.get(); p.clipped = d.clipped.get();
p.perm = d.perm.get(); p.rr_start = d.rr_start.get(); p.rr_count = d.rr_count.get();
p.rr_h = d.rr_h.get(); p.rr_k = d.rr_k.get(); p.rr_l = d.rr_l.get(); p.rr_group = d.rr_group.get();
p.rr_nevents = d.rr_nevents.get();
p.rr_noverloaded = d.rr_noverloaded.get();
p.rr_h = d.rr_h.get(); p.rr_k = d.rr_k.get(); p.rr_l = d.rr_l.get(); p.rr_group = rr_group.get();
p.rr_nevents = rr_nevents.get();
p.rr_noverloaded = rr_noverloaded.get();
const int blocks = std::min(65535, (d.n_runs + BLK - 1) / BLK);
@@ -1975,17 +1996,18 @@ int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_parti
// Exclusive prefix sum on the host (deterministic) -> per-run output offset + total fulls.
std::vector<int32_t> nevents(d.n_runs);
CopyAndWait(nevents.data(), d.rr_nevents.get(), size_t(d.n_runs) * sizeof(int32_t),
CopyAndWait(nevents.data(), rr_nevents.get(), size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyDeviceToHost, impl_->s(), "download nevents");
std::vector<int32_t> offset(d.n_runs);
int64_t acc = 0;
for (int r = 0; r < d.n_runs; ++r) { offset[r] = static_cast<int32_t>(acc); acc += nevents[r]; }
d.n_fulls = static_cast<int>(acc);
std::vector<int32_t> noverloaded(d.n_runs);
CopyAndWait(noverloaded.data(), d.rr_noverloaded.get(), size_t(d.n_runs) * sizeof(int32_t),
CopyAndWait(noverloaded.data(), rr_noverloaded.get(), size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyDeviceToHost, impl_->s(), "download noverloaded");
n_overloaded = 0;
for (int r = 0; r < d.n_runs; ++r) n_overloaded += noverloaded[r];
rr_nevents = {}; rr_noverloaded = {};
// Allocate the fulls SoA and emit.
const int nf = std::max(1, d.n_fulls);
@@ -1999,10 +2021,9 @@ int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_parti
d.f_corr = d.Alloc<float>(nf); d.f_partiality = d.Alloc<float>(nf);
d.f_zeta = d.Alloc<float>(nf);
d.f_obs = d.Alloc<ObsRecord>(nf);
CopyAndWait(d.rr_offset.get(), offset.data(), size_t(d.n_runs) * sizeof(int32_t),
cudaMemcpyHostToDevice, impl_->s(), "upload offset");
d.Upload(rr_offset, offset.data(), d.n_runs);
p.rr_offset = d.rr_offset.get();
p.rr_offset = rr_offset.get();
p.f_h = d.f_h.get(); p.f_k = d.f_k.get(); p.f_l = d.f_l.get();
p.f_frame = d.f_frame.get(); p.f_group = d.f_group.get();
p.f_I = d.f_I.get(); p.f_sigma = d.f_sigma.get(); p.f_d = d.f_d.get(); p.f_img = d.f_img.get();
@@ -2104,6 +2125,14 @@ void RotationScaleMergeGPU::GetFullsVariance(float *var_bkg, float *var_per_I, f
CopyAndWait(capture, d.f_capture.get(), bytes, cudaMemcpyDeviceToHost, impl_->s(), "download f_capture");
}
void RotationScaleMergeGPU::FullsDownloaded() {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
d.f_img = {}; d.f_px = {}; d.f_py = {}; d.f_clipped = {};
d.f_obs = {}; d.f_zeta = {}; d.f_frame_perm = {}; d.f_frame_start = {}; d.f_frame_count = {};
d.group_mean = {};
}
void RotationScaleMergeGPU::SetFullsCorr(const float *corr) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
@@ -2616,39 +2645,45 @@ namespace {
float m2_S0, m3_S0;
};
// RotationScaleMerge::SmoothedDeltaPhi / SolvePhi.
__global__ void IngestDeltaPhiKernel(int n, const int32_t *__restrict__ frame, const uint64_t *__restrict__ key,
// RotationScaleMerge::SmoothedDeltaPhi / SolvePhi. One thread per raw-hkl run, which gives the hkl of
// its observations; each observation is solved on its own.
__global__ void IngestDeltaPhiKernel(int n_runs, const int32_t *__restrict__ start,
const int32_t *__restrict__ count, const int32_t *__restrict__ perm,
const int32_t *__restrict__ rr_h, const int32_t *__restrict__ rr_k,
const int32_t *__restrict__ rr_l, const int32_t *__restrict__ frame,
DeltaPhiGeometry g, float *__restrict__ dphi,
unsigned long long *__restrict__ n_resolved) {
constexpr float PI_F = 3.14159265358979323846f;
unsigned long long resolved = 0;
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) {
const int o = frame[i];
if (!g.ok[o]) continue;
const float *rc = g.recip + 9 * o;
const Vec3 a{rc[0], rc[1], rc[2]}, b{rc[3], rc[4], rc[5]}, c{rc[6], rc[7], rc[8]};
const float h = float(RawHklIndex(key[i], 32)), k = float(RawHklIndex(key[i], 16)),
l = float(RawHklIndex(key[i], 0));
const Vec3 p0 = VAdd(VAdd(VScale(a, h), VScale(b, k)), VScale(c, l));
const float phi_near = __fdiv_rn(__fmul_rn(dphi[i], PI_F), 180.0f);
const float p0_sq = VDot(p0, p0);
if (!(p0_sq > 0.0f)) continue;
const float p0_m1 = VDot(p0, g.m1), p0_m2 = VDot(p0, g.m2), p0_m3 = VDot(p0, g.m3);
const float rho_sq = __fsub_rn(p0_sq, __fmul_rn(p0_m2, p0_m2));
const float p_m3 = __fdiv_rn(__fsub_rn(__fdiv_rn(-p0_sq, 2.0f), __fmul_rn(p0_m2, g.m2_S0)), g.m3_S0);
if (!(rho_sq > __fmul_rn(p_m3, p_m3))) continue;
const float root = __fsqrt_rn(__fsub_rn(rho_sq, __fmul_rn(p_m3, p_m3)));
float best = 0.0f;
bool have = false;
for (const float p_m1 : {root, -root}) {
const float cosphi = __fdiv_rn(__fadd_rn(__fmul_rn(p_m1, p0_m1), __fmul_rn(p_m3, p0_m3)), rho_sq);
const float sinphi = __fdiv_rn(__fsub_rn(__fmul_rn(p_m1, p0_m3), __fmul_rn(p_m3, p0_m1)), rho_sq);
const float phi = -atan2f(sinphi, cosphi);
if (!have || fabsf(__fsub_rn(phi, phi_near)) < fabsf(__fsub_rn(best, phi_near))) { best = phi; have = true; }
for (int r = blockIdx.x * blockDim.x + threadIdx.x; r < n_runs; r += gridDim.x * blockDim.x) {
const float h = float(rr_h[r]), k = float(rr_k[r]), l = float(rr_l[r]);
for (int q = start[r]; q < start[r] + count[r]; ++q) {
const int i = perm[q];
const int o = frame[i];
if (!g.ok[o]) continue;
const float *rc = g.recip + 9 * o;
const Vec3 a{rc[0], rc[1], rc[2]}, b{rc[3], rc[4], rc[5]}, c{rc[6], rc[7], rc[8]};
const Vec3 p0 = VAdd(VAdd(VScale(a, h), VScale(b, k)), VScale(c, l));
const float phi_near = __fdiv_rn(__fmul_rn(dphi[i], PI_F), 180.0f);
const float p0_sq = VDot(p0, p0);
if (!(p0_sq > 0.0f)) continue;
const float p0_m1 = VDot(p0, g.m1), p0_m2 = VDot(p0, g.m2), p0_m3 = VDot(p0, g.m3);
const float rho_sq = __fsub_rn(p0_sq, __fmul_rn(p0_m2, p0_m2));
const float p_m3 = __fdiv_rn(__fsub_rn(__fdiv_rn(-p0_sq, 2.0f), __fmul_rn(p0_m2, g.m2_S0)), g.m3_S0);
if (!(rho_sq > __fmul_rn(p_m3, p_m3))) continue;
const float root = __fsqrt_rn(__fsub_rn(rho_sq, __fmul_rn(p_m3, p_m3)));
float best = 0.0f;
bool have = false;
for (const float p_m1 : {root, -root}) {
const float cosphi = __fdiv_rn(__fadd_rn(__fmul_rn(p_m1, p0_m1), __fmul_rn(p_m3, p0_m3)), rho_sq);
const float sinphi = __fdiv_rn(__fsub_rn(__fmul_rn(p_m1, p0_m3), __fmul_rn(p_m3, p0_m1)), rho_sq);
const float phi = -atan2f(sinphi, cosphi);
if (!have || fabsf(__fsub_rn(phi, phi_near)) < fabsf(__fsub_rn(best, phi_near))) { best = phi; have = true; }
}
if (!isfinite(best)) continue;
dphi[i] = __fdiv_rn(__fmul_rn(best, 180.0f), PI_F);
++resolved;
}
if (!isfinite(best)) continue;
dphi[i] = __fdiv_rn(__fmul_rn(best, 180.0f), PI_F);
++resolved;
}
atomicAdd(n_resolved, resolved);
}
@@ -2869,15 +2904,23 @@ int RotationScaleMergeGPU::IngestSort(const double *frame_flux, std::optional<do
// (h, k, l, image_number, index) is a total order, so two stable radix passes - by the frame position
// with the index as the value, then by the raw hkl - give the one permutation std::sort gives on the
// host, whatever either algorithm does on the way.
const auto sort_pairs = [&](auto *key_in, auto *key_out, const int32_t *val_in, int32_t *val_out, int end_bit) {
// host, whatever either algorithm does on the way. Each pass goes back and forth between the two
// buffers of its keys and of its values, as neither input is needed after it, so that cub holds no
// copy of them of its own: the sorts are where the ingest's footprint peaks. The same stable sort, so
// the same order; it is left in key_out and val_out.
const auto sort_pairs = [&](auto &key_in, auto &key_out, CudaDevicePtr<int32_t> &val_in,
CudaDevicePtr<int32_t> &val_out, int end_bit) {
using K = std::remove_reference_t<decltype(*key_in.get())>;
cub::DoubleBuffer<K> keys(key_in.get(), key_out.get());
cub::DoubleBuffer<int32_t> values(val_in.get(), val_out.get());
size_t bytes = 0;
CudaCheck(cub::DeviceRadixSort::SortPairs(nullptr, bytes, key_in, key_out, val_in, val_out, n, 0, end_bit,
d.s()), "sort temp size");
CudaCheck(cub::DeviceRadixSort::SortPairs(nullptr, bytes, keys, values, n, 0, end_bit, d.s()),
"sort temp size");
CudaDevicePtr<uint8_t> temp = d.Alloc<uint8_t>(std::max<size_t>(1, bytes));
CudaCheck(cub::DeviceRadixSort::SortPairs(temp.get(), bytes, key_in, key_out, val_in, val_out, n, 0, end_bit,
d.s()), "sort");
CudaCheck(cub::DeviceRadixSort::SortPairs(temp.get(), bytes, keys, values, n, 0, end_bit, d.s()), "sort");
CudaCheck(cudaStreamSynchronize(d.s()), "sort");
if (keys.Current() != key_out.get()) std::swap(key_in, key_out);
if (values.Current() != val_out.get()) std::swap(val_in, val_out);
};
CudaDevicePtr<int32_t> by_image = d.Alloc<int32_t>(na);
{
@@ -2885,17 +2928,18 @@ int RotationScaleMergeGPU::IngestSort(const double *frame_flux, std::optional<do
CudaDevicePtr<int32_t> idx = d.Alloc<int32_t>(na);
IngestImageKeyKernel<<<Blocks(n), BLK, 0, d.s()>>>(n, d.image_number.get(), ikey.get(), idx.get());
CudaCheck(cudaGetLastError(), "IngestImageKeyKernel launch");
if (n > 0) sort_pairs(ikey.get(), ikey_sorted.get(), idx.get(), by_image.get(), 32);
if (n > 0) sort_pairs(ikey, ikey_sorted, idx, by_image, 32);
}
CudaDevicePtr<uint64_t> key = d.Alloc<uint64_t>(na);
IngestGatherKeyKernel<<<Blocks(n), BLK, 0, d.s()>>>(n, by_image.get(), d.in_key.get(), key.get());
CudaCheck(cudaGetLastError(), "IngestGatherKeyKernel launch");
// The runs the sort cuts carry the raw hkl from here on.
CudaCheck(cudaStreamSynchronize(d.s()), "gather key");
d.in_key = {};
CudaDevicePtr<int32_t> perm = d.Alloc<int32_t>(na);
CudaDevicePtr<uint64_t> key_sorted = d.Alloc<uint64_t>(na);
{
CudaDevicePtr<uint64_t> key = d.Alloc<uint64_t>(na);
IngestGatherKeyKernel<<<Blocks(n), BLK, 0, d.s()>>>(n, by_image.get(), d.in_key.get(), key.get());
CudaCheck(cudaGetLastError(), "IngestGatherKeyKernel launch");
if (n > 0) sort_pairs(key.get(), key_sorted.get(), by_image.get(), perm.get(), 48);
}
by_image = CudaDevicePtr<int32_t>();
if (n > 0) sort_pairs(key, key_sorted, by_image, perm, 48);
key = {}; by_image = {};
// The raw-hkl runs of the sorted order.
int n_runs = 0;
@@ -2995,6 +3039,10 @@ int RotationScaleMergeGPU::IngestSort(const double *frame_flux, std::optional<do
d.rr_start.get(), d.rr_count.get(), d.rr_h.get(), d.rr_k.get(), d.rr_l.get(), d.perm.get(), rr_d.get()};
IngestCutRunsKernel<<<Blocks(n_runs), BLK, 0, d.s()>>>(c);
CudaCheck(cudaGetLastError(), "IngestCutRunsKernel launch");
CudaCheck(cudaStreamSynchronize(d.s()), "cut runs");
// The runs are cut: what was kept of them is in rr_*, so the uncut ones and the old order go before
// the observations are cut below, one field held twice at a time.
perm = {}; run_key = {}; run_count = {}; run_start = {}; run_d = {}; keep_run = {}; run_idx = {}; run_at = {};
// The frame CSR of what survived.
auto &L = d.in_layout;
@@ -3026,11 +3074,10 @@ int RotationScaleMergeGPU::IngestSort(const double *frame_flux, std::optional<do
field = std::move(out);
};
cut(d.obs); cut(d.zeta); cut(d.corr); cut(d.bkg); cut(d.var_bkg); cut(d.image_number); cut(d.d_obs);
cut(d.px_obs); cut(d.py_obs); cut(d.frame); cut(d.on_ice); cut(d.clipped); cut(d.in_key); cut(d.in_dphi);
cut(d.px_obs); cut(d.py_obs); cut(d.frame); cut(d.on_ice); cut(d.clipped); cut(d.in_dphi);
}
d.n_obs = n_keep;
d.AllocFrames(L.frame_start.data(), L.frame_count.data());
d.AllocRunScratch();
L.run_start.resize(n_keep_run); L.run_count.resize(n_keep_run);
L.run_h.resize(n_keep_run); L.run_k.resize(n_keep_run); L.run_l.resize(n_keep_run); L.run_d.resize(n_keep_run);
@@ -3058,8 +3105,9 @@ int64_t RotationScaleMergeGPU::IngestDeltaPhi(const DeltaPhiFrame &g) {
CudaCheck(cudaMemsetAsync(n_resolved.get(), 0, sizeof(unsigned long long), d.s()), "memset");
const auto v = [](const float *a) { return Vec3{a[0], a[1], a[2]}; };
const DeltaPhiGeometry dg{recip.get(), ok.get(), v(g.S0), v(g.m1), v(g.m2), v(g.m3), g.m2_S0, g.m3_S0};
IngestDeltaPhiKernel<<<Blocks(d.n_obs), BLK, 0, d.s()>>>(d.n_obs, d.frame.get(), d.in_key.get(), dg,
d.in_dphi.get(), n_resolved.get());
IngestDeltaPhiKernel<<<Blocks(d.n_runs), BLK, 0, d.s()>>>(d.n_runs, d.rr_start.get(), d.rr_count.get(),
d.perm.get(), d.rr_h.get(), d.rr_k.get(), d.rr_l.get(),
d.frame.get(), dg, d.in_dphi.get(), n_resolved.get());
CudaCheck(cudaGetLastError(), "IngestDeltaPhiKernel launch");
unsigned long long out = 0;
CopyAndWait(&out, n_resolved.get(), sizeof(out), cudaMemcpyDeviceToHost, d.s(), "resolved count");
@@ -3114,7 +3162,6 @@ void RotationScaleMergeGPU::IngestFinish(float *corr) {
if (d.n_obs > 0)
CopyAndWait(corr, d.corr.get(), size_t(d.n_obs) * sizeof(float), cudaMemcpyDeviceToHost, d.s(),
"download corr");
d.in_key = CudaDevicePtr<uint64_t>();
d.in_dphi = CudaDevicePtr<float>();
d.in_src_start.clear();
d.in_layout = RotationScaleMergeGPU::IngestLayout{};
image_analysis/scale_merge/RotationScaleMergeGPU.h
+11 -1
View File
@@ -191,6 +191,9 @@ public:
double *swh0, double *swh1, double *swh_typ0, double *swh_typ1,
int32_t *nh0, int32_t *nh1, double *d_out,
int32_t *rejected, uint8_t *on_ice_out);
// Hand back what MergeAccum left - thirteen numbers per group - once every range is downloaded, so
// that the sums the merge goes on to take do not stand beside them.
void MergeAccumRelease();
// Per-group R_meas accumulators: sum|I_corr-merged_I| and sum I, the same with each observation
// weighted by its merge weight v = 1/sigma^2 under the error model of the last MergeAccum,
@@ -220,7 +223,9 @@ public:
// id per raw hkl (it becomes the full's group). Deterministic: fulls are emitted in raw-run-major,
// event order (a count pass -> host prefix sum -> emit-at-offset), matching the CPU path. Returns the
// number of fulls (call GetFulls with buffers of that length); n_overloaded receives how many rocking
// events were dropped for a saturated pixel (see the host combine).
// events were dropped for a saturated pixel (see the host combine). The previous fulls and the
// partials' group CSR go here, so the partials' scaling (ScalePartials, ComputePartialCC) needs
// BuildGroups again after it.
int Combine(const int32_t *rawrun_group, double min_partiality, double capture_uncertainty_coeff,
double min_captured_fraction, float max_frame_gap, int64_t &n_overloaded);
@@ -237,6 +242,11 @@ public:
// Download the fulls' variance model, var(I) = var_bkg + var_per_I * I + (capture * I)^2. Length = n_fulls.
void GetFullsVariance(float *var_bkg, float *var_per_I, float *capture) const;
// Hand back the arrays only the fulls' scaling and the downloads above read (the records, zeta, the
// frame CSR, the group means, image number, px/py and the clipped flags), once the fulls are on the
// host: the merge and the surfaces after it read none of them, and the next arm makes its own.
void FullsDownloaded();
// Re-upload the fulls' working corr (length n_fulls) after the host correction surfaces (decay /
// absorption) mutate it, so the resident merge reads the corrected scale.
void SetFullsCorr(const float *corr);
rugnux/RugnuxScaleMerge.cpp
+1 -1
View File
@@ -1162,7 +1162,7 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) {
// free: the answer must not depend on the moment. What else stays on the card meanwhile - the
// context, the run's indexers - was measured at up to 2.3 GB; a speculative pass running beside
// holds more, but a merge that does not fit beside it waits for it to end. On a 16 GB card that
// is the data sets up to about 50 M observations at high multiplicity; the large ones merge
// is the data sets up to about 70 M observations at high multiplicity; the large ones merge
// their arms one after the other.
if (!geometry_prepass && !experiment_.GetGemmiSpaceGroup().has_value() && rsm->GetSearchMinZeta() > 0.0
&& cards > 0 && rsm->OnGPU()) {