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

Merge branch 'merge-layout' into rc175: merge engine observations as a packed 16-byte record, the scale preparation folded into the per-frame fit

Browse Source 
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
leonarski_f
2026-10-12 13:37:40 +02:00
co-authored by Claude Opus 5.5
parent d712eba00f 7d656c13de
commit b28dff47e4
2 changed files with 146 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/scale_merge/RotationScaleMerge.cpp
+3 -2
View File
@@ -677,8 +677,9 @@ 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.
constexpr size_t BYTES_PER_OBS = 80, BYTES_PER_EVENT = 260, BYTES_PER_HKL = 200;
// 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;
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();
image_analysis/scale_merge/RotationScaleMergeGPU.cu
+143 -132
View File
@@ -51,6 +51,11 @@ namespace {
return s[0];
}
// The four immutable fields of an observation that the scaling, the combine and the grouping read
// together, in one 16-byte record: the kernels that reach an observation through a permutation
// (the group reduction above all) then fetch one sector for it instead of one per field.
struct alignas(16) ObsRecord { float I, sigma, partiality, prescaling_corr; }; // floats 0, 1, 2, 3
// The observation's inverse-variance weight in the scaling reference (0 = not in it), and its
// scaled intensity. The same filter as the host's ReferenceWeight.
__device__ __forceinline__ double ReferenceWeight(float I, float sigma, float partiality, float corr,
@@ -72,8 +77,7 @@ namespace {
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 ObsRecord *__restrict__ obs,
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) {
@@ -81,8 +85,9 @@ namespace {
double sw = 0.0, swI = 0.0;
for (int p = lo; p < hi; ++p) {
const int i = group_perm[p];
const ObsRecord o = obs[i];
float I_corr;
const double w = ReferenceWeight(I[i], sigma[i], partiality[i], corr[i], min_partiality, I_corr);
const double w = ReferenceWeight(o.I, o.sigma, o.partiality, corr[i], min_partiality, I_corr);
if (w <= 0.0) continue;
sw += w;
swI += w * I_corr;
@@ -91,29 +96,20 @@ namespace {
}
}
// Per-observation scale-fit coefficient (rotation model) and accept flag, recomputed each scaling
// iteration once the group means are known. coeff = partiality * (1/prescaling_corr) * mean[group].
__global__ void PrepScaleObsKernel(int n_obs, double min_partiality, const int32_t *__restrict__ group,
const float *__restrict__ partiality, const float *__restrict__ prescaling_corr,
const float *__restrict__ zeta, const uint8_t *__restrict__ on_ice,
const double *__restrict__ group_mean,
const float *__restrict__ sigma, double *__restrict__ inv_sigma,
float *__restrict__ sco_coeff, uint8_t *__restrict__ sco_ok) {
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= n_obs) return;
// The scale-fit coefficient of an observation (rotation model), recomputed each scaling iteration
// once the group means are known: coeff = partiality * (1/prescaling_corr) * mean[group], or false
// where the observation is not one the reference is built from (the host FitPerFrameG says why).
__device__ __forceinline__ bool ScaleCoeff(int i, double min_partiality, const int32_t *__restrict__ group,
const ObsRecord &o, const float *__restrict__ zeta,
const uint8_t *__restrict__ on_ice,
const double *__restrict__ group_mean, float &coeff) {
const int g = group[i];
// Only the observations the reference is built from (the host FitPerFrameG says why).
bool ok = (g >= 0) && !on_ice[i] && isfinite(zeta[i]) && zeta[i] > 0.0f && partiality[i] >= min_partiality;
double mean = 0.0;
if (ok) {
mean = group_mean[g];
ok = isfinite(mean);
}
// sigma never changes once it is uploaded, so its reciprocal is the same in every scaling
// pass; hoisted as the CPU does (ScaleObs::weight).
inv_sigma[i] = SafeInvD(sigma[i], 1.0);
sco_ok[i] = ok ? 1 : 0;
sco_coeff[i] = ok ? float(double(partiality[i]) * SafeInvD(prescaling_corr[i], 1.0) * mean) : 0.0f;
if (!((g >= 0) && !on_ice[i] && isfinite(zeta[i]) && zeta[i] > 0.0f && o.partiality >= min_partiality))
return false;
const double mean = group_mean[g];
if (!isfinite(mean)) return false;
coeff = float(double(o.partiality) * SafeInvD(o.prescaling_corr, 1.0) * mean);
return true;
}
// One block per frame: the per-frame scale G as the weighted least-squares slope over the frame's
@@ -122,13 +118,15 @@ namespace {
// for the partials, whose arrays are already frame-contiguous) maps a position in the frame's
// [lo,hi) range to the obs index, so the same kernel scales the fulls (emit-ordered) through a
// frame-grouping permutation without physically reordering the fulls arrays. info[f] receives the
// fit's information on G, sum w^2 c^2 (0 for a frame left alone).
// fit's information on G, sum w^2 c^2 (0 for a frame left alone). The weight is 1/sigma, as the
// CPU's ScaleObs::weight.
__global__ void FitPerFrameGKernel(int n_frames,
const int32_t *__restrict__ frame_start,
const int32_t *__restrict__ frame_count,
const float *__restrict__ I, const double *__restrict__ inv_sigma,
const float *__restrict__ sco_coeff, const uint8_t *__restrict__ sco_ok,
const int32_t *__restrict__ perm, long min_obs,
const int32_t *__restrict__ perm, double min_partiality,
const int32_t *__restrict__ group, const ObsRecord *__restrict__ obs,
const float *__restrict__ zeta, const uint8_t *__restrict__ on_ice,
const double *__restrict__ group_mean, long min_obs,
double *__restrict__ g, uint8_t *__restrict__ scaled,
double *__restrict__ info) {
const int f = blockIdx.x;
@@ -137,25 +135,28 @@ namespace {
__shared__ double sh[BLK];
if (threadIdx.x == 0) info[f] = 0.0;
// The sums are formed in the same pass as the count; a frame with too few observations
// discards them.
long cnt_local = 0;
for (int i = lo + threadIdx.x; i < hi; i += blockDim.x)
if (sco_ok[perm ? perm[i] : i]) ++cnt_local;
double num = 0.0, den = 0.0;
for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) {
const int a = perm ? perm[i] : i;
const ObsRecord o = obs[a];
float coeff_f;
if (!ScaleCoeff(a, min_partiality, group, o, zeta, on_ice, group_mean, coeff_f)) continue;
++cnt_local;
const double coeff = coeff_f;
const double w = SafeInvD(o.sigma, 1.0);
const double w2 = w * w;
num += w2 * coeff * double(o.I);
den += w2 * coeff * coeff;
}
const double cnt = BlockReduceSum(double(cnt_local), sh);
__shared__ double s_cnt;
if (threadIdx.x == 0) s_cnt = cnt;
__syncthreads();
if (s_cnt < min_obs) return; // leave g[f]/scaled[f] as-is
double num = 0.0, den = 0.0;
for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) {
const int a = perm ? perm[i] : i;
if (!sco_ok[a]) continue;
const double coeff = sco_coeff[a];
const double w = inv_sigma[a];
const double w2 = w * w;
num += w2 * coeff * double(I[a]);
den += w2 * coeff * coeff;
}
const double tnum = BlockReduceSum(num, sh); __syncthreads();
const double tden = BlockReduceSum(den, sh);
if (threadIdx.x == 0) {
@@ -172,8 +173,7 @@ namespace {
__global__ void PerFrameCCKernel(int n_frames, double min_partiality,
const int32_t *__restrict__ frame_start,
const int32_t *__restrict__ frame_count,
const float *__restrict__ I, const float *__restrict__ sigma,
const float *__restrict__ partiality, const float *__restrict__ corr,
const ObsRecord *__restrict__ obs, const float *__restrict__ corr,
const uint8_t *__restrict__ on_ice, const int32_t *__restrict__ group,
const double *__restrict__ group_mean,
double *__restrict__ cc_out, int64_t *__restrict__ cc_n_out) {
@@ -187,13 +187,14 @@ namespace {
if (on_ice[i]) continue;
const int g = group[i];
if (g < 0) continue;
if (partiality[i] < min_partiality) continue;
const ObsRecord o = obs[i];
if (o.partiality < min_partiality) continue;
const float c = corr[i];
if (!isfinite(I[i]) || !isfinite(c) || !(c > 0.0f)) continue;
if (!isfinite(sigma[i]) || !(sigma[i] > 0.0f)) continue;
if (!isfinite(o.I) || !isfinite(c) || !(c > 0.0f)) continue;
if (!isfinite(o.sigma) || !(o.sigma > 0.0f)) continue;
const double mean = group_mean[g];
if (!isfinite(mean)) continue;
const double img = double(I[i]) * c;
const double img = double(o.I) * c;
sx += img; sy += mean; sx2 += img * img; sy2 += mean * mean; sxy += img * mean; ++nl;
}
const double tsx = BlockReduceSum(sx, sh); __syncthreads();
@@ -250,15 +251,16 @@ namespace {
// corr = prescaling_corr / (partiality * G[frame]) for fitted frames; unchanged otherwise (grid-stride).
__global__ void UpdateCorrKernel(int n_obs, const int32_t *__restrict__ frame,
const float *__restrict__ prescaling_corr, const float *__restrict__ partiality,
const ObsRecord *__restrict__ obs,
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(prescaling_corr[i])) && isfinite(denom) && denom > 0.0)
? float(prescaling_corr[i] / denom) : NAN;
const ObsRecord o = obs[i];
const double denom = double(o.partiality) * g[f];
corr[i] = (isfinite(double(o.prescaling_corr)) && isfinite(denom) && denom > 0.0)
? float(o.prescaling_corr / denom) : NAN;
}
}
@@ -269,11 +271,10 @@ namespace {
__device__ __forceinline__ float Fmax(float a, float b) { return (a < b) ? b : a; }
// A partial is usable for the combine iff its corr and (I, sigma) are finite with corr>0, sigma>0.
__device__ __forceinline__ bool CombineUsable(int i, const float *I, const float *sigma,
const float *corr) {
__device__ __forceinline__ bool CombineUsable(int i, const ObsRecord *obs, const float *corr) {
const float c = corr[i];
if (!(c > 0.0f) || !isfinite(c)) return false;
return isfinite(I[i]) && isfinite(sigma[i]) && sigma[i] > 0.0f;
return isfinite(obs[i].I) && isfinite(obs[i].sigma) && obs[i].sigma > 0.0f;
}
// RotationScaleMerge::OBS_CLIPPED / OBS_OVERLOADED, the flag bits a partial's `clipped` byte carries.
@@ -287,8 +288,8 @@ namespace {
// RotationScaleMerge::max_frame_gap, passed in rather than recomputed here so that the host
// and the device compare the same float and the combine stays bit-identical on both paths.
float max_frame_gap;
const float *__restrict__ I, *__restrict__ sigma, *__restrict__ corr,
*__restrict__ partiality, *__restrict__ bkg, *__restrict__ var_bkg,
const ObsRecord *__restrict__ obs;
const float *__restrict__ corr, *__restrict__ bkg, *__restrict__ var_bkg,
*__restrict__ image_number, *__restrict__ d, *__restrict__ px, *__restrict__ py;
const int32_t *__restrict__ frame;
const uint8_t *__restrict__ on_ice, *__restrict__ clipped;
@@ -317,14 +318,14 @@ namespace {
int n_emit = 0, n_overloaded = 0;
int cursor = lo;
while (cursor < hi) {
while (cursor < hi && !CombineUsable(p.perm[cursor], p.I, p.sigma, p.corr)) ++cursor;
while (cursor < hi && !CombineUsable(p.perm[cursor], p.obs, p.corr)) ++cursor;
if (cursor >= hi) break;
const int ev_start = cursor; // first usable position of the event
int ev_end = cursor; // last usable position (inclusive), extended below
float last_img = p.image_number[p.perm[cursor]];
int probe = cursor + 1;
while (probe < hi) {
while (probe < hi && !CombineUsable(p.perm[probe], p.I, p.sigma, p.corr)) ++probe;
while (probe < hi && !CombineUsable(p.perm[probe], p.obs, p.corr)) ++probe;
if (probe >= hi) break;
const float img = p.image_number[p.perm[probe]];
if (img - last_img > p.max_frame_gap) break;
@@ -340,16 +341,16 @@ namespace {
int n_pool = 0;
for (int m = ev_start; m <= ev_end; ++m) {
const int i = p.perm[m];
if (!CombineUsable(i, p.I, p.sigma, p.corr)) continue;
if (!CombineUsable(i, p.obs, p.corr)) continue;
const float b = p.bkg[i];
if (isfinite(b)) { pooled_bkg += b; ++n_pool; }
}
pooled_bkg = n_pool > 0 ? pooled_bkg / n_pool : 0.0;
auto pooled_I = [&](int i) -> double {
const double n_bkg = Dmax(0.0, double(p.sigma[i]) * p.sigma[i] - p.I[i])
const double n_bkg = Dmax(0.0, double(p.obs[i].sigma) * p.obs[i].sigma - p.obs[i].I)
/ Fmax(p.bkg[i], 1.0f);
return double(p.I[i]) + n_bkg * (double(p.bkg[i]) - pooled_bkg);
return double(p.obs[i].I) + n_bkg * (double(p.bkg[i]) - pooled_bkg);
};
// Pass B: seed F (inverse-variance mean of pooled_I*corr), plus peak / d / on_ice / partiality.
@@ -368,14 +369,14 @@ namespace {
const bool clipped = flags & OBS_CLIPPED_FLAG;
for (int m = ev_start; m <= ev_end; ++m) {
const int i = p.perm[m];
if (!CombineUsable(i, p.I, p.sigma, p.corr)) continue;
const double sigma_corr = double(p.sigma[i]) * p.corr[i];
if (!CombineUsable(i, p.obs, p.corr)) continue;
const double sigma_corr = double(p.obs[i].sigma) * p.corr[i];
const double w = 1.0 / (sigma_corr * sigma_corr);
sum_w += w;
sum_wI += w * pooled_I(i) * p.corr[i];
sum_partiality += p.partiality[i];
if (p.partiality[i] > peak_partiality) {
peak_partiality = p.partiality[i];
sum_partiality += p.obs[i].partiality;
if (p.obs[i].partiality > peak_partiality) {
peak_partiality = p.obs[i].partiality;
peak_outcome = p.frame[i];
peak_frame = p.image_number[i];
peak_px = p.px[i]; peak_py = p.py[i];
@@ -391,10 +392,10 @@ namespace {
sum_w = 0.0; sum_wI = 0.0; sum_wb = 0.0; sum_cwb = 0.0;
for (int m = ev_start; m <= ev_end; ++m) {
const int i = p.perm[m];
if (!CombineUsable(i, p.I, p.sigma, p.corr)) continue;
if (!CombineUsable(i, p.obs, p.corr)) continue;
const double corr = p.corr[i];
const double I_corr = pooled_I(i) * corr;
const double sigma_corr = double(p.sigma[i]) * corr;
const double sigma_corr = double(p.obs[i].sigma) * corr;
// The integrator's own non-signal variance (see the host combine).
const double bkg_var = corr * corr * (double) p.var_bkg[i];
const double a_var = bkg_var > 0.0 ? bkg_var : sigma_corr * sigma_corr;
@@ -459,15 +460,17 @@ namespace {
// ===== the per-space-group grouping, built on the device from the raw-hkl runs =====
// The host's finiteness test (finite_ok in RotationScaleMerge::Ingest), on the same uploaded values.
__device__ __forceinline__ bool ObsFinite(int i, const float *I, const float *sigma, const float *pcorr) {
return isfinite(I[i]) && isfinite(pcorr[i]) && pcorr[i] != 0.0f && isfinite(sigma[i]) && sigma[i] > 0.0f;
__device__ __forceinline__ bool ObsFinite(int i, const ObsRecord *obs) {
const ObsRecord o = obs[i];
return isfinite(o.I) && isfinite(o.prescaling_corr) && o.prescaling_corr != 0.0f && isfinite(o.sigma)
&& o.sigma > 0.0f;
}
struct GroupParams {
int n_runs, n_groups;
const int32_t *perm, *rr_start, *rr_count, *rr_group;
const int32_t *sorted_run, *group_first; // group g = runs sorted_run[group_first[g] .. [g+1])
const float *I, *sigma, *pcorr;
const ObsRecord *obs;
int32_t *group, *group_start, *group_count, *group_perm;
};
@@ -478,7 +481,7 @@ namespace {
const int g = p.rr_group[r];
for (int q = p.rr_start[r]; q < p.rr_start[r] + p.rr_count[r]; ++q) {
const int i = p.perm[q];
p.group[i] = (g >= 0 && ObsFinite(i, p.I, p.sigma, p.pcorr)) ? g : -1;
p.group[i] = (g >= 0 && ObsFinite(i, p.obs)) ? g : -1;
}
}
}
@@ -489,7 +492,7 @@ namespace {
for (int j = p.group_first[g]; j < p.group_first[g + 1]; ++j) {
const int r = p.sorted_run[j];
for (int q = p.rr_start[r]; q < p.rr_start[r] + p.rr_count[r]; ++q)
c += ObsFinite(p.perm[q], p.I, p.sigma, p.pcorr) ? 1 : 0;
c += ObsFinite(p.perm[q], p.obs) ? 1 : 0;
}
p.group_count[g] = c;
}
@@ -506,7 +509,7 @@ namespace {
const int r = p.sorted_run[j];
for (int q = p.rr_start[r]; q < p.rr_start[r] + p.rr_count[r]; ++q) {
const int i = p.perm[q];
if (ObsFinite(i, p.I, p.sigma, p.pcorr)) a[n++] = i;
if (ObsFinite(i, p.obs)) a[n++] = i;
}
}
auto sift = [a](int root, int end) {
@@ -526,6 +529,21 @@ namespace {
}
}
// One field of a slice of the records, from a contiguous array: rec[i].field = v[i], the field being
// the record's float number `field`.
__global__ void SetRecordFieldKernel(int n, const float *__restrict__ v, ObsRecord *__restrict__ rec, int field) {
float *f = reinterpret_cast<float *>(rec);
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
f[size_t(i) * 4 + field] = v[i];
}
// The fulls' records for their Unity scaling: partiality and prescaling_corr 1.
__global__ void FullsRecordKernel(int n, const float *__restrict__ I, const float *__restrict__ sigma,
ObsRecord *__restrict__ obs) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x)
obs[i] = {I[i], sigma[i], 1.0f, 1.0f};
}
template <typename T>
__global__ void FillKernel(T *p, int n, T v) {
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) p[i] = v;
@@ -1088,11 +1106,10 @@ namespace {
std::string("RotationScaleMergeGPU: ") + what + ": " + cudaGetErrorString(e));
}
// Device bytes per observation in the arrays SetPartialsLayout allocates: twelve floats (I, sigma,
// prescaling_corr, partiality, zeta, corr, bkg, var_bkg, image_number, d, px, py), frame and
// sco_coeff (4 bytes each), inv_sigma (8), on_ice, clipped and sco_ok (1 each). The merge's other
// buffers come on top of these.
constexpr size_t OBS_BYTES = 12 * 4 + 4 + 4 + 8 + 3 * 1;
// Device bytes per observation in the arrays SetPartialsLayout allocates: twelve floats (the record's
// I, sigma, partiality and prescaling_corr; zeta, corr, bkg, var_bkg, image_number, d, px, py), frame
// (4 bytes), on_ice and clipped (1 each). The merge's other buffers come on top of these.
constexpr size_t OBS_BYTES = 12 * 4 + 4 + 2 * 1;
// Every buffer here is sized by the data set, and the scaling does not fall back to the CPU
// mid-run: the two paths differ in the last bits, so the result would depend on the card. A
@@ -1190,7 +1207,9 @@ struct RotationScaleMergeGPU::Impl {
}
// immutable per-obs
CudaDevicePtr<float> I, sigma, prescaling_corr, partiality, zeta;
CudaDevicePtr<ObsRecord> obs;
void UploadRecordChunk(int field, int offset, int count, const float *v);
CudaDevicePtr<float> zeta;
CudaDevicePtr<uint8_t> on_ice, clipped;
CudaDevicePtr<int32_t> frame;
CudaDevicePtr<float> corr; // mutable, resident across iterations
@@ -1200,9 +1219,6 @@ struct RotationScaleMergeGPU::Impl {
// scratch
CudaDevicePtr<double> group_mean, g, info;
CudaDevicePtr<uint8_t> scaled;
CudaDevicePtr<double> inv_sigma; // 1/sigma, hoisted out of the IRLS loop (sigma never changes)
CudaDevicePtr<float> sco_coeff;
CudaDevicePtr<uint8_t> sco_ok;
CudaDevicePtr<double> cc; // per-frame CC (diagnostic), length n_frames
CudaDevicePtr<int64_t> cc_n;
CudaDevicePtr<uint8_t> smooth_apply; // per-frame smooth-G apply flag + ratio, length n_frames
@@ -1243,12 +1259,12 @@ struct RotationScaleMergeGPU::Impl {
CudaDevicePtr<int32_t> f_h, f_k, f_l, f_frame, f_group;
CudaDevicePtr<float> f_I, f_sigma, f_d, f_img, f_px, f_py, f_var_bkg, f_var_per_I, f_capture;
CudaDevicePtr<uint8_t> f_on_ice, f_clipped;
// scale-fulls (Unity model, kept resident): all-ones partiality/prescaling_corr/zeta so the shared scaling kernels
// yield coeff=mean, plus the working corr, the per-obs scale scratch, and the fulls frame/group CSRs
// (built on the host from the small f_frame/f_group key arrays, over the emit-ordered fulls).
CudaDevicePtr<float> f_corr, f_partiality, f_rlp, f_zeta, f_sco_coeff;
CudaDevicePtr<double> f_inv_sigma;
CudaDevicePtr<uint8_t> f_sco_ok;
// scale-fulls (Unity model, kept resident): all-ones partiality/zeta and records with partiality and
// prescaling_corr 1, so the shared scaling kernels yield coeff=mean, plus the working corr and the
// fulls frame/group CSRs (built on the host from the small f_frame/f_group key arrays, over the
// emit-ordered fulls).
CudaDevicePtr<float> f_corr, f_partiality, f_zeta;
CudaDevicePtr<ObsRecord> f_obs;
CudaDevicePtr<int32_t> f_frame_perm, f_frame_start, f_frame_count;
CudaDevicePtr<int32_t> f_gperm, f_gstart, f_gcount;
// correction-surface fit (one ApplyCellSurface call at a time): its terms and their group CSR, the
@@ -1320,8 +1336,7 @@ void RotationScaleMergeGPU::SetPartialsLayout(int n_obs, int n_frames,
ThrowTooLargeForGPU(n_obs);
d.Upload(d.frame_start, frame_start, n_frames); d.Upload(d.frame_count, frame_count, n_frames);
const int n = std::max(1, n_obs);
d.I = d.Alloc<float>(n); d.sigma = d.Alloc<float>(n);
d.prescaling_corr = d.Alloc<float>(n); d.partiality = d.Alloc<float>(n);
d.obs = d.Alloc<ObsRecord>(n);
d.zeta = d.Alloc<float>(n); d.corr = d.Alloc<float>(n);
d.bkg = d.Alloc<float>(n); d.var_bkg = d.Alloc<float>(n);
d.image_number = d.Alloc<float>(n); d.d_obs = d.Alloc<float>(n);
@@ -1332,9 +1347,6 @@ void RotationScaleMergeGPU::SetPartialsLayout(int n_obs, int n_frames,
d.g = d.Alloc<double>(n_frames);
d.scaled = d.Alloc<uint8_t>(n_frames);
d.info = d.Alloc<double>(n_frames);
d.inv_sigma = d.Alloc<double>(n_obs);
d.sco_coeff = d.Alloc<float>(n_obs);
d.sco_ok = d.Alloc<uint8_t>(n_obs);
d.cc = d.Alloc<double>(std::max(1, n_frames));
d.cc_n = d.Alloc<int64_t>(std::max(1, n_frames));
}
@@ -1346,16 +1358,30 @@ namespace {
CopyAndWait(dst.get() + offset, v, size_t(count) * sizeof(T), cudaMemcpyHostToDevice, s,
"upload chunk");
}
}
// A slice of one field of the records (`field` its float number), staged contiguously and then put
// in place on the device: a strided copy from pageable memory goes element by element.
void RotationScaleMergeGPU::Impl::UploadRecordChunk(int field, int offset, int count, const float *v) {
if (count == 0)
return;
auto staged = Alloc<float>(count);
CopyAndWait(staged.get(), v, size_t(count) * sizeof(float), cudaMemcpyHostToDevice, s(), "upload record chunk");
SetRecordFieldKernel<<<std::min(65535, (count + BLK - 1) / BLK), BLK, 0, s()>>>(count, staged.get(),
obs.get() + offset, field);
CudaCheck(cudaGetLastError(), "record field launch");
CudaCheck(cudaStreamSynchronize(s()), "upload record chunk");
}
void RotationScaleMergeGPU::SetObsField(ObsField f, int offset, int count, const float *v) {
DeviceGuard guard(impl_->device, impl_->available);
auto &d = *impl_;
switch (f) {
case ObsField::I: UploadChunk(d.I, offset, count, v, impl_->s()); break;
case ObsField::Sigma: UploadChunk(d.sigma, offset, count, v, impl_->s()); break;
case ObsField::PrescalingCorr: UploadChunk(d.prescaling_corr, offset, count, v, impl_->s()); break;
case ObsField::Partiality: UploadChunk(d.partiality, offset, count, v, impl_->s()); break;
case ObsField::I: d.UploadRecordChunk(0, offset, count, v); break;
case ObsField::Sigma: d.UploadRecordChunk(1, offset, count, v); break;
case ObsField::Partiality: d.UploadRecordChunk(2, offset, count, v); break;
case ObsField::PrescalingCorr: d.UploadRecordChunk(3, offset, count, v); break;
case ObsField::Zeta: UploadChunk(d.zeta, offset, count, v, impl_->s()); break;
case ObsField::Corr0: UploadChunk(d.corr, offset, count, v, impl_->s()); break;
case ObsField::Bkg: UploadChunk(d.bkg, offset, count, v, impl_->s()); break;
@@ -1375,8 +1401,7 @@ void RotationScaleMergeGPU::CopyObservationsFrom(const RotationScaleMergeGPU &fr
CudaCheck(cudaMemcpyPeerAsync(dst.get(), d.device, src.get(), f.device,
size_t(d.n_obs) * sizeof(*src.get()), d.s()), "copy observations");
};
copy(d.I, f.I); copy(d.sigma, f.sigma); copy(d.prescaling_corr, f.prescaling_corr);
copy(d.partiality, f.partiality); copy(d.zeta, f.zeta); copy(d.corr, f.corr);
copy(d.obs, f.obs); copy(d.zeta, f.zeta); copy(d.corr, f.corr);
copy(d.bkg, f.bkg); copy(d.var_bkg, f.var_bkg); copy(d.image_number, f.image_number);
copy(d.d_obs, f.d_obs); copy(d.px_obs, f.px_obs); copy(d.py_obs, f.py_obs);
copy(d.frame, f.frame); copy(d.on_ice, f.on_ice); copy(d.clipped, f.clipped);
@@ -1416,7 +1441,7 @@ void RotationScaleMergeGPU::BuildGroups(int n_groups, const int32_t *rawrun_grou
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.sorted_run = sorted.get(); p.group_first = first.get();
p.I = d.I.get(); p.sigma = d.sigma.get(); p.pcorr = d.prescaling_corr.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();
const int run_blocks = std::min(65535, (d.n_runs + BLK - 1) / BLK);
const int grp_blocks = std::min(65535, (n_groups + BLK - 1) / BLK);
@@ -1456,23 +1481,18 @@ void RotationScaleMergeGPU::ScalePartials(int iters, double min_partiality, bool
// call did not fit must read as unfitted, not as fitted with the value of the call before.
CudaCheck(cudaMemsetAsync(d.scaled.get(), 0, size_t(d.n_frames) * sizeof(uint8_t), impl_->s()), "memset scaled");
CudaCheck(cudaMemsetAsync(d.g.get(), 0, size_t(d.n_frames) * sizeof(double), impl_->s()), "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 upd_blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK);
const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK);
for (int it = 0; it < iters; ++it) {
ReduceGroupMeansKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(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());
d.obs.get(), d.corr.get(), d.group_mean.get());
CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch");
PrepScaleObsKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(d.n_obs, min_partiality, d.group.get(), d.partiality.get(), d.prescaling_corr.get(),
d.zeta.get(), d.on_ice.get(), d.group_mean.get(), d.sigma.get(), d.inv_sigma.get(),
d.sco_coeff.get(), d.sco_ok.get());
CudaCheck(cudaGetLastError(), "PrepScaleObsKernel launch");
FitPerFrameGKernel<<<d.n_frames, BLK, 0, impl_->s()>>>(d.n_frames, d.frame_start.get(), d.frame_count.get(),
d.I.get(), d.inv_sigma.get(), d.sco_coeff.get(), d.sco_ok.get(), nullptr, long(MIN_REFLECTIONS), d.g.get(), d.scaled.get(),
d.info.get());
nullptr, min_partiality, d.group.get(), d.obs.get(), d.zeta.get(), d.on_ice.get(), d.group_mean.get(),
long(MIN_REFLECTIONS), d.g.get(), d.scaled.get(), d.info.get());
CudaCheck(cudaGetLastError(), "FitPerFrameGKernel launch");
UpdateCorrKernel<<<upd_blocks, BLK, 0, impl_->s()>>>(d.n_obs, d.frame.get(), d.prescaling_corr.get(), d.partiality.get(),
UpdateCorrKernel<<<upd_blocks, BLK, 0, impl_->s()>>>(d.n_obs, d.frame.get(), d.obs.get(),
d.g.get(), d.scaled.get(), d.corr.get());
}
CudaCheck(cudaGetLastError(), "kernel launch");
@@ -1838,11 +1858,11 @@ void RotationScaleMergeGPU::ComputePartialCC(double min_partiality, double *cc_o
// per-frame CC over the resident partials. Only the tiny per-frame cc/cc_n come back to the host.
ReduceGroupMeansKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(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());
d.obs.get(), d.corr.get(), d.group_mean.get());
CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch");
PerFrameCCKernel<<<d.n_frames, BLK, 0, impl_->s()>>>(d.n_frames, min_partiality,
d.frame_start.get(), d.frame_count.get(), d.I.get(), d.sigma.get(), d.partiality.get(),
d.corr.get(), d.on_ice.get(), d.group.get(), d.group_mean.get(), d.cc.get(), d.cc_n.get());
d.frame_start.get(), d.frame_count.get(), d.obs.get(), d.corr.get(), d.on_ice.get(), d.group.get(),
d.group_mean.get(), d.cc.get(), d.cc_n.get());
CudaCheck(cudaGetLastError(), "partial CC launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "partial CC sync");
CopyAndWait(cc_out, d.cc.get(), size_t(d.n_frames) * sizeof(double),
@@ -1885,7 +1905,7 @@ int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_parti
p.capture_uncertainty_coeff = capture_uncertainty_coeff;
p.min_captured_fraction = min_captured_fraction;
p.max_frame_gap = max_frame_gap;
p.I = d.I.get(); p.sigma = d.sigma.get(); p.corr = d.corr.get(); p.partiality = d.partiality.get();
p.obs = d.obs.get(); p.corr = d.corr.get();
p.bkg = d.bkg.get(); p.var_bkg = d.var_bkg.get(); p.image_number = d.image_number.get(); p.d = d.d_obs.get();
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();
@@ -1924,9 +1944,8 @@ int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_parti
d.f_var_bkg = d.Alloc<float>(nf); d.f_var_per_I = d.Alloc<float>(nf); d.f_capture = d.Alloc<float>(nf);
d.f_on_ice = d.Alloc<uint8_t>(nf); d.f_clipped = d.Alloc<uint8_t>(nf);
d.f_corr = d.Alloc<float>(nf); d.f_partiality = d.Alloc<float>(nf);
d.f_rlp = d.Alloc<float>(nf); d.f_zeta = d.Alloc<float>(nf);
d.f_inv_sigma = d.Alloc<double>(nf);
d.f_sco_coeff = d.Alloc<float>(nf); d.f_sco_ok = d.Alloc<uint8_t>(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");
@@ -2002,8 +2021,8 @@ void RotationScaleMergeGPU::ResetFullsScale() {
CudaCheck(cudaGetLastError(), "FillKernel launch");
FillKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(d.f_partiality.get(), nf, 1.0f);
CudaCheck(cudaGetLastError(), "FillKernel launch");
FillKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(d.f_rlp.get(), nf, 1.0f);
CudaCheck(cudaGetLastError(), "FillKernel launch");
FullsRecordKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(nf, d.f_I.get(), d.f_sigma.get(), d.f_obs.get());
CudaCheck(cudaGetLastError(), "FullsRecordKernel launch");
FillKernel<<<obs_blocks, BLK, 0, impl_->s()>>>(d.f_zeta.get(), nf, 1.0f);
CudaCheck(cudaGetLastError(), "FillKernel launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "reset fulls scale sync");
@@ -2022,19 +2041,11 @@ void RotationScaleMergeGPU::FitFullsScale(double min_partiality) {
ReduceGroupMeansKernel<<<grp_blocks, BLK, 0, impl_->s()>>>(d.n_groups, min_partiality,
d.f_gperm.get(), d.f_gstart.get(), d.f_gcount.get(),
d.f_I.get(), d.f_sigma.get(), d.f_partiality.get(), d.f_corr.get(), d.group_mean.get());
d.f_obs.get(), d.f_corr.get(), d.group_mean.get());
CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch");
// Not grid-stride, so its grid has to cover every full - unlike the grid-stride kernels
// below, which the 65535 cap is there for. Capped, it would silently leave the tail of
// sco_coeff/sco_ok stale above 16.8M fulls.
PrepScaleObsKernel<<<(nf + BLK - 1) / BLK, BLK, 0, impl_->s()>>>(nf, min_partiality, d.f_group.get(), d.f_partiality.get(),
d.f_rlp.get(), d.f_zeta.get(), d.f_on_ice.get(), d.group_mean.get(),
d.f_sigma.get(), d.f_inv_sigma.get(), d.f_sco_coeff.get(), d.f_sco_ok.get());
CudaCheck(cudaGetLastError(), "PrepScaleObsKernel launch");
FitPerFrameGKernel<<<d.n_frames, BLK, 0, impl_->s()>>>(d.n_frames,
d.f_frame_start.get(), d.f_frame_count.get(), d.f_I.get(), d.f_inv_sigma.get(),
d.f_sco_coeff.get(), d.f_sco_ok.get(), d.f_frame_perm.get(), 1L, d.g.get(), d.scaled.get(),
d.info.get());
FitPerFrameGKernel<<<d.n_frames, BLK, 0, impl_->s()>>>(d.n_frames, d.f_frame_start.get(), d.f_frame_count.get(),
d.f_frame_perm.get(), min_partiality, d.f_group.get(), d.f_obs.get(), d.f_zeta.get(), d.f_on_ice.get(),
d.group_mean.get(), 1L, d.g.get(), d.scaled.get(), d.info.get());
CudaCheck(cudaGetLastError(), "scale fulls launch");
CudaCheck(cudaStreamSynchronize(impl_->s()), "scale fulls sync");
}