diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index d5da6e73f..38d52965a 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -5,6 +5,7 @@ #include #include +#include #include #include @@ -640,11 +641,28 @@ namespace { std::string("RotationScaleMergeGPU: ") + what + ": " + cudaGetErrorString(e)); } - template - void Upload(CudaDevicePtr &dst, const T *src, int n) { - dst = CudaDevicePtr(std::max(1, n), ALLOC); - if (n > 0) - CudaCheck(cudaMemcpy(dst.get(), src, size_t(n) * sizeof(T), cudaMemcpyHostToDevice), "upload"); + // 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; + + // 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 + // buffer that does not fit ends the run, and says why and what to do instead. + [[noreturn]] void ThrowTooLargeForGPU(int64_t n_obs) { + size_t free_bytes = 0, total_bytes = 0; + cudaMemGetInfo(&free_bytes, &total_bytes); + cudaGetLastError(); + char msg[640]; + snprintf(msg, sizeof(msg), + "Scaling %lld partial observations needs more GPU memory than this card has: %.1f GB for " + "their per-observation arrays alone (%zu bytes each), before the merge's other buffers, " + "on a %.1f GB card. This data set is too large for GPU scaling on this card - run it on " + "the CPU (CUDA_VISIBLE_DEVICES= rugnux ..., or a CPU-only build)", + static_cast(n_obs), double(n_obs) * OBS_BYTES / 1e9, OBS_BYTES, + double(total_bytes) / 1e9); + throw JFJochException(JFJochExceptionCategory::GPUCUDAError, msg); } } @@ -653,6 +671,22 @@ struct RotationScaleMergeGPU::Impl { bool available = false; int n_obs = 0, n_frames = 0, n_groups = 0; + template + CudaDevicePtr Alloc(size_t n) const { + try { + return CudaDevicePtr(n, ALLOC); + } catch (const JFJochException &) { + ThrowTooLargeForGPU(n_obs); + } + } + + template + void Upload(CudaDevicePtr &dst, const T *src, int n) const { + dst = Alloc(std::max(1, n)); + if (n > 0) + CudaCheck(cudaMemcpy(dst.get(), src, size_t(n) * sizeof(T), cudaMemcpyHostToDevice), "upload"); + } + // immutable per-obs CudaDevicePtr I, sigma, prescaling_corr, partiality, zeta; CudaDevicePtr on_ice, clipped; @@ -764,24 +798,30 @@ void RotationScaleMergeGPU::SetPartialsLayout(int n_obs, int n_frames, auto &d = *impl_; d.n_obs = n_obs; d.n_frames = n_frames; - Upload(d.frame_start, frame_start, n_frames); Upload(d.frame_count, frame_count, n_frames); + // Known now, before anything is allocated: arrays that could not fit even an empty card fail here, + // not after gigabytes of them have been. + size_t free_bytes = 0, total_bytes = 0; + CudaCheck(cudaMemGetInfo(&free_bytes, &total_bytes), "memory info"); + if (size_t(std::max(0, n_obs)) * OBS_BYTES > total_bytes) + 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 = CudaDevicePtr(n, ALLOC); d.sigma = CudaDevicePtr(n, ALLOC); - d.prescaling_corr = CudaDevicePtr(n, ALLOC); d.partiality = CudaDevicePtr(n, ALLOC); - d.zeta = CudaDevicePtr(n, ALLOC); d.corr = CudaDevicePtr(n, ALLOC); - d.bkg = CudaDevicePtr(n, ALLOC); d.var_bkg = CudaDevicePtr(n, ALLOC); - d.image_number = CudaDevicePtr(n, ALLOC); d.d_obs = CudaDevicePtr(n, ALLOC); - d.px_obs = CudaDevicePtr(n, ALLOC); d.py_obs = CudaDevicePtr(n, ALLOC); - d.frame = CudaDevicePtr(n, ALLOC); - d.on_ice = CudaDevicePtr(n, ALLOC); - d.clipped = CudaDevicePtr(n, ALLOC); - d.g = CudaDevicePtr(n_frames, ALLOC); - d.scaled = CudaDevicePtr(n_frames, ALLOC); - d.inv_sigma = CudaDevicePtr(n_obs, ALLOC); - d.sco_coeff = CudaDevicePtr(n_obs, ALLOC); - d.sco_ok = CudaDevicePtr(n_obs, ALLOC); - d.cc = CudaDevicePtr(std::max(1, n_frames), ALLOC); - d.cc_n = CudaDevicePtr(std::max(1, n_frames), ALLOC); + d.I = d.Alloc(n); d.sigma = d.Alloc(n); + d.prescaling_corr = d.Alloc(n); d.partiality = d.Alloc(n); + d.zeta = d.Alloc(n); d.corr = d.Alloc(n); + d.bkg = d.Alloc(n); d.var_bkg = d.Alloc(n); + d.image_number = d.Alloc(n); d.d_obs = d.Alloc(n); + d.px_obs = d.Alloc(n); d.py_obs = d.Alloc(n); + d.frame = d.Alloc(n); + d.on_ice = d.Alloc(n); + d.clipped = d.Alloc(n); + d.g = d.Alloc(n_frames); + d.scaled = d.Alloc(n_frames); + d.inv_sigma = d.Alloc(n_obs); + d.sco_coeff = d.Alloc(n_obs); + d.sco_ok = d.Alloc(n_obs); + d.cc = d.Alloc(std::max(1, n_frames)); + d.cc_n = d.Alloc(std::max(1, n_frames)); } namespace { @@ -833,11 +873,11 @@ void RotationScaleMergeGPU::SetGroups(int n_groups, const int32_t *group, const DeviceGuard guard(impl_->device, impl_->available); auto &d = *impl_; d.n_groups = n_groups; - Upload(d.group, group, d.n_obs); - Upload(d.group_perm, group_perm, n_group_perm); // obs with group >= 0, in group order - Upload(d.group_start, group_start, n_groups); - Upload(d.group_count, group_count, n_groups); - d.group_mean = CudaDevicePtr(std::max(1, n_groups), ALLOC); + d.Upload(d.group, group, d.n_obs); + d.Upload(d.group_perm, group_perm, n_group_perm); // obs with group >= 0, in group order + d.Upload(d.group_start, group_start, n_groups); + d.Upload(d.group_count, group_count, n_groups); + d.group_mean = d.Alloc(std::max(1, n_groups)); } void RotationScaleMergeGPU::SetCorr(const float *corr) { @@ -891,7 +931,7 @@ void RotationScaleMergeGPU::GetG(double *g_out, uint8_t *scaled_out) const { void RotationScaleMergeGPU::SetFrameCellOk(const uint8_t *frame_cell_ok) { DeviceGuard guard(impl_->device, impl_->available); - Upload(impl_->frame_cell_ok, frame_cell_ok, impl_->n_frames); + impl_->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 @@ -904,23 +944,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 = CudaDevicePtr(std::max(1, ng), ALLOC); d.m_swI = CudaDevicePtr(std::max(1, ng), ALLOC); - d.m_em_mean = CudaDevicePtr(std::max(1, ng), ALLOC); d.m_cnt = CudaDevicePtr(std::max(1, ng), ALLOC); + d.m_sw = d.Alloc(std::max(1, ng)); d.m_swI = d.Alloc(std::max(1, ng)); + d.m_em_mean = d.Alloc(std::max(1, ng)); d.m_cnt = d.Alloc(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. if (hand && has_hands) { - d.m_swh = CudaDevicePtr(std::max(1, 2 * ng), ALLOC); d.m_swIh = CudaDevicePtr(std::max(1, 2 * ng), ALLOC); - d.m_cnth = CudaDevicePtr(std::max(1, 2 * ng), ALLOC); - if (nf > 0) Upload(d.m_hand, hand, nf); - if (ng > 0) Upload(d.m_has_hands, has_hands, ng); + d.m_swh = d.Alloc(std::max(1, 2 * ng)); d.m_swIh = d.Alloc(std::max(1, 2 * ng)); + d.m_cnth = d.Alloc(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(); d.m_swIh = CudaDevicePtr(); d.m_cnth = CudaDevicePtr(); d.m_hand = CudaDevicePtr(); d.m_has_hands = CudaDevicePtr(); } - d.m_s2 = CudaDevicePtr(std::max(1, nf), ALLOC); d.m_I2 = CudaDevicePtr(std::max(1, nf), ALLOC); - d.m_dev2 = CudaDevicePtr(std::max(1, nf), ALLOC); d.m_valid = CudaDevicePtr(std::max(1, nf), ALLOC); + d.m_s2 = d.Alloc(std::max(1, nf)); d.m_I2 = d.Alloc(std::max(1, nf)); + d.m_dev2 = d.Alloc(std::max(1, nf)); d.m_valid = d.Alloc(std::max(1, nf)); MergeParams p{}; p.n_groups = ng; p.min_partiality = min_partiality; @@ -962,19 +1002,19 @@ void RotationScaleMergeGPU::MergeAccum(double error_model_a, double error_model_ DeviceGuard guard(impl_->device, impl_->available); auto &d = *impl_; const int ng = d.n_groups; - d.a_swI = CudaDevicePtr(std::max(1, ng), ALLOC); d.a_sw = CudaDevicePtr(std::max(1, ng), ALLOC); - d.a_swIh0 = CudaDevicePtr(std::max(1, ng), ALLOC); d.a_swIh1 = CudaDevicePtr(std::max(1, ng), ALLOC); - d.a_swh0 = CudaDevicePtr(std::max(1, ng), ALLOC); d.a_swh1 = CudaDevicePtr(std::max(1, ng), ALLOC); - d.a_swht0 = CudaDevicePtr(std::max(1, ng), ALLOC); d.a_swht1 = CudaDevicePtr(std::max(1, ng), ALLOC); - d.a_nh0 = CudaDevicePtr(std::max(1, ng), ALLOC); d.a_nh1 = CudaDevicePtr(std::max(1, ng), ALLOC); - d.a_d = CudaDevicePtr(std::max(1, ng), ALLOC); d.a_rejected = CudaDevicePtr(std::max(1, ng), ALLOC); - d.a_on_ice = CudaDevicePtr(std::max(1, ng), ALLOC); - Upload(d.m_rejected, rejected_obs, d.n_fulls); - Upload(d.reject_median, reject_median, ng); - if (reject_var_add && ng > 0) Upload(d.reject_var_add, reject_var_add, ng); + d.a_swI = d.Alloc(std::max(1, ng)); d.a_sw = d.Alloc(std::max(1, ng)); + d.a_swIh0 = d.Alloc(std::max(1, ng)); d.a_swIh1 = d.Alloc(std::max(1, ng)); + d.a_swh0 = d.Alloc(std::max(1, ng)); d.a_swh1 = d.Alloc(std::max(1, ng)); + d.a_swht0 = d.Alloc(std::max(1, ng)); d.a_swht1 = d.Alloc(std::max(1, ng)); + d.a_nh0 = d.Alloc(std::max(1, ng)); d.a_nh1 = d.Alloc(std::max(1, ng)); + d.a_d = d.Alloc(std::max(1, ng)); d.a_rejected = d.Alloc(std::max(1, ng)); + d.a_on_ice = d.Alloc(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(); - Upload(d.m_half, half, d.n_fulls); - Upload(d.cc_factor, frame_cc_factor, d.n_frames); + d.Upload(d.m_half, half, d.n_fulls); + d.Upload(d.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{}; @@ -1027,11 +1067,11 @@ void RotationScaleMergeGPU::MergeRmeas(const double *merged_I, double *absdev, d DeviceGuard guard(impl_->device, impl_->available); auto &d = *impl_; const int ng = d.n_groups; - Upload(d.merged_I, merged_I, ng); - d.r_absdev = CudaDevicePtr(std::max(1, ng), ALLOC); d.r_sumI = CudaDevicePtr(std::max(1, ng), ALLOC); - d.r_wabsdev = CudaDevicePtr(std::max(1, ng), ALLOC); d.r_wsumI = CudaDevicePtr(std::max(1, ng), ALLOC); - d.r_sumv = CudaDevicePtr(std::max(1, ng), ALLOC); d.r_sumv2 = CudaDevicePtr(std::max(1, ng), ALLOC); - d.r_n = CudaDevicePtr(std::max(1, ng), ALLOC); d.r_nusable = CudaDevicePtr(std::max(1, ng), ALLOC); + d.Upload(d.merged_I, merged_I, ng); + d.r_absdev = d.Alloc(std::max(1, ng)); d.r_sumI = d.Alloc(std::max(1, ng)); + d.r_wabsdev = d.Alloc(std::max(1, ng)); d.r_wsumI = d.Alloc(std::max(1, ng)); + d.r_sumv = d.Alloc(std::max(1, ng)); d.r_sumv2 = d.Alloc(std::max(1, ng)); + d.r_n = d.Alloc(std::max(1, ng)); d.r_nusable = d.Alloc(std::max(1, ng)); MergeParams p{}; p.n_groups = ng; p.min_partiality = d.merge_min_part; @@ -1063,8 +1103,8 @@ void RotationScaleMergeGPU::MergeRmeas(const double *merged_I, double *absdev, d void RotationScaleMergeGPU::SmoothCorr(const uint8_t *apply, const double *ratio) { DeviceGuard guard(impl_->device, impl_->available); auto &d = *impl_; - Upload(d.smooth_apply, apply, d.n_frames); - Upload(d.smooth_ratio, ratio, d.n_frames); + d.Upload(d.smooth_apply, apply, d.n_frames); + d.Upload(d.smooth_ratio, ratio, d.n_frames); const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK); SmoothCorrKernel<<>>(d.n_obs, d.frame.get(), d.smooth_apply.get(), d.smooth_ratio.get(), d.corr.get()); @@ -1076,8 +1116,8 @@ void RotationScaleMergeGPU::SmoothFullsCorr(const uint8_t *apply, const double * DeviceGuard guard(impl_->device, impl_->available); auto &d = *impl_; if (d.n_fulls == 0) return; - Upload(d.smooth_apply, apply, d.n_frames); - Upload(d.smooth_ratio, ratio, d.n_frames); + d.Upload(d.smooth_apply, apply, d.n_frames); + d.Upload(d.smooth_ratio, ratio, d.n_frames); const int blocks = std::min(65535, (d.n_fulls + BLK - 1) / BLK); SmoothCorrKernel<<>>(d.n_fulls, d.f_frame.get(), d.smooth_apply.get(), d.smooth_ratio.get(), d.f_corr.get()); @@ -1088,7 +1128,7 @@ void RotationScaleMergeGPU::SmoothFullsCorr(const uint8_t *apply, const double * int64_t RotationScaleMergeGPU::FilterCorrByZeta(double min_zeta) { DeviceGuard guard(impl_->device, impl_->available); auto &d = *impl_; - CudaDevicePtr dropped(1, ALLOC); + CudaDevicePtr dropped = d.Alloc(1); CudaCheck(cudaMemset(dropped.get(), 0, sizeof(unsigned long long)), "zero zeta drop count"); const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK); FilterZetaKernel<<>>(d.n_obs, min_zeta, d.zeta.get(), d.corr.get(), dropped.get()); @@ -1103,7 +1143,7 @@ int64_t RotationScaleMergeGPU::FilterCorrByZeta(double min_zeta) { void RotationScaleMergeGPU::FilterCorrByFrame(const uint8_t *reject) { DeviceGuard guard(impl_->device, impl_->available); auto &d = *impl_; - Upload(d.filter_reject, reject, d.n_frames); + d.Upload(d.filter_reject, reject, d.n_frames); const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK); FilterFrameKernel<<>>(d.n_obs, d.frame.get(), d.filter_reject.get(), d.corr.get()); CudaCheck(cudaGetLastError(), "frame filter launch"); @@ -1138,16 +1178,16 @@ void RotationScaleMergeGPU::SetRawRuns(int n_runs, int n_perm, const int32_t *pe auto &d = *impl_; d.n_runs = n_runs; d.n_perm = n_perm; - Upload(d.perm, perm, n_perm); - Upload(d.rr_start, rr_start, n_runs); - Upload(d.rr_count, rr_count, n_runs); - Upload(d.rr_h, rr_h, n_runs); - Upload(d.rr_k, rr_k, n_runs); - Upload(d.rr_l, rr_l, n_runs); + d.Upload(d.perm, perm, n_perm); + d.Upload(d.rr_start, rr_start, n_runs); + d.Upload(d.rr_count, rr_count, n_runs); + 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.rr_group = CudaDevicePtr(std::max(1, n_runs), ALLOC); - d.rr_nevents = CudaDevicePtr(std::max(1, n_runs), ALLOC); - d.rr_offset = CudaDevicePtr(std::max(1, n_runs), ALLOC); + d.rr_group = d.Alloc(std::max(1, n_runs)); + d.rr_nevents = d.Alloc(std::max(1, n_runs)); + d.rr_offset = d.Alloc(std::max(1, n_runs)); } int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_partiality, @@ -1189,17 +1229,17 @@ int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_parti // Allocate the fulls SoA and emit. const int nf = std::max(1, d.n_fulls); - d.f_h = CudaDevicePtr(nf, ALLOC); d.f_k = CudaDevicePtr(nf, ALLOC); d.f_l = CudaDevicePtr(nf, ALLOC); - d.f_frame = CudaDevicePtr(nf, ALLOC); d.f_group = CudaDevicePtr(nf, ALLOC); - d.f_I = CudaDevicePtr(nf, ALLOC); d.f_sigma = CudaDevicePtr(nf, ALLOC); - d.f_d = CudaDevicePtr(nf, ALLOC); d.f_img = CudaDevicePtr(nf, ALLOC); - d.f_px = CudaDevicePtr(nf, ALLOC); d.f_py = CudaDevicePtr(nf, ALLOC); - d.f_var_bkg = CudaDevicePtr(nf, ALLOC); d.f_var_per_I = CudaDevicePtr(nf, ALLOC); - d.f_on_ice = CudaDevicePtr(nf, ALLOC); d.f_clipped = CudaDevicePtr(nf, ALLOC); - d.f_corr = CudaDevicePtr(nf, ALLOC); d.f_partiality = CudaDevicePtr(nf, ALLOC); - d.f_rlp = CudaDevicePtr(nf, ALLOC); d.f_zeta = CudaDevicePtr(nf, ALLOC); - d.f_inv_sigma = CudaDevicePtr(nf, ALLOC); - d.f_sco_coeff = CudaDevicePtr(nf, ALLOC); d.f_sco_ok = CudaDevicePtr(nf, ALLOC); + d.f_h = d.Alloc(nf); d.f_k = d.Alloc(nf); d.f_l = d.Alloc(nf); + d.f_frame = d.Alloc(nf); d.f_group = d.Alloc(nf); + d.f_I = d.Alloc(nf); d.f_sigma = d.Alloc(nf); + d.f_d = d.Alloc(nf); d.f_img = d.Alloc(nf); + d.f_px = d.Alloc(nf); d.f_py = d.Alloc(nf); + d.f_var_bkg = d.Alloc(nf); d.f_var_per_I = d.Alloc(nf); + d.f_on_ice = d.Alloc(nf); d.f_clipped = d.Alloc(nf); + d.f_corr = d.Alloc(nf); d.f_partiality = d.Alloc(nf); + d.f_rlp = d.Alloc(nf); d.f_zeta = d.Alloc(nf); + d.f_inv_sigma = d.Alloc(nf); + d.f_sco_coeff = d.Alloc(nf); d.f_sco_ok = d.Alloc(nf); CudaCheck(cudaMemcpy(d.rr_offset.get(), offset.data(), size_t(d.n_runs) * sizeof(int32_t), cudaMemcpyHostToDevice), "upload offset"); @@ -1250,18 +1290,18 @@ void RotationScaleMergeGPU::SetFullsFrameCSR(const int32_t *frame_perm, int n_pe const int32_t *frame_start, const int32_t *frame_count) { DeviceGuard guard(impl_->device, impl_->available); auto &d = *impl_; - Upload(d.f_frame_perm, frame_perm, n_perm); - Upload(d.f_frame_start, frame_start, d.n_frames); - Upload(d.f_frame_count, frame_count, d.n_frames); + d.Upload(d.f_frame_perm, frame_perm, n_perm); + d.Upload(d.f_frame_start, frame_start, d.n_frames); + d.Upload(d.f_frame_count, frame_count, d.n_frames); } void RotationScaleMergeGPU::SetFullsGroups(const int32_t *gperm, int n_gperm, const int32_t *gstart, const int32_t *gcount) { DeviceGuard guard(impl_->device, impl_->available); auto &d = *impl_; - Upload(d.f_gperm, gperm, n_gperm); - Upload(d.f_gstart, gstart, d.n_groups); - Upload(d.f_gcount, gcount, d.n_groups); + d.Upload(d.f_gperm, gperm, n_gperm); + d.Upload(d.f_gstart, gstart, d.n_groups); + d.Upload(d.f_gcount, gcount, d.n_groups); } void RotationScaleMergeGPU::ResetFullsScale() {