From 46ad2c10ed1e1bd6aef88bd2941a28162a7404ea Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sun, 25 Oct 2026 20:35:03 +0100 Subject: [PATCH] Refinement GPU: the robust scale fit on the card, the pool back in every trial slot The score's robust fine pass (pseudo-Huber IRLS of ModelScaleBox::robust) on the device, so that refinement on a GPU no longer fits its scale on the host every cycle - with the trials side by side that was a robust grid search on one thread per slot (13-133 s per refinement on the seven timing sets, against 1.7-8 s before the score commit). - ModelScaleGPU: the weights multiply the WSSR and J^T J / J^T r in scale_sums (a weight row per fit, reserved at the first robust fit); robust_residuals / robust_weights kernels; the median of each fit's first residuals taken on the host as FitModelScale takes it (nth_element at n/2, on the fits' threads); LevMarBatch split out of FitBatch so that the rounds restart from where a fit stands. FitSolvent(..., robust) refits the coarse winner in the fine pass, as FitModelScale's robust fit does; the plain fit is unchanged. - ModelScaleGPUPool::Fit(..., robust); its streams at the highest priority: a fit is a chain of small launches the host waits on, and on a card shared with the trials' long kernels each wait queued behind them (refit.scale 58-87 ms -> 14-45 ms per cycle, the result the same to the bit). - Refinement: the main context and every trial slot fit on their pool (robust); the CPU engine keeps FitModelScale's robust fit. The card's fit is FitModelScale's to rounding (float residuals and weights, the steps' anisotropic factor in float), deterministic run to run and whichever engine a fit takes: R-work/R-free of the 14 parity sets on the GPU equal the CPU engine's to 4 digits (largest difference 1.4e-6 on the 3.6 A set). Test: ModelScaleGPUPool_RobustMatchesFitModelScale (one reflection in fifty a 4x outlier, five cells). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu (cherry picked from commit d15f048cbd7ed2d6b430f9f9cb5018ca284da8f0) --- .../structure_refinement/ModelScaleGPU.cu | 124 +++++++++++++++--- .../structure_refinement/ModelScaleGPU.h | 15 ++- .../structure_refinement/ModelScaling.cpp | 23 +++- .../structure_refinement/ModelScaling.h | 12 +- .../structure_refinement/Refinement.cpp | 19 ++- tests/ModelScaleGPUTest.cpp | 39 ++++++ 6 files changed, 202 insertions(+), 30 deletions(-) diff --git a/image_analysis/structure_refinement/ModelScaleGPU.cu b/image_analysis/structure_refinement/ModelScaleGPU.cu index f26a72744..0665fddce 100644 --- a/image_analysis/structure_refinement/ModelScaleGPU.cu +++ b/image_analysis/structure_refinement/ModelScaleGPU.cu @@ -10,6 +10,7 @@ #include #include "../../common/JFJochException.h" +#include "../../common/ParallelFor.h" #include "RigidBodyGPUEngine.h" // ModelScaleGPUTooFewReflections namespace { @@ -25,6 +26,8 @@ constexpr int THREADS = 128; constexpr int BLOCKS = 48; constexpr int SLOTS = 36; // 1 + 28 + 7: the WSSR, the lower half of J^T J and J^T r for 7 parameters constexpr int MAX_FITS = 88; // FitModelScale's coarse pass, the larger of its two +constexpr int MAX_ROBUST_FITS = 32; // FitModelScale's fine pass, at most 5 x 5 points +constexpr int ROBUST_ROUNDS = 4; // as ModelScaling.cpp // What a launch sums over the points for one fit. constexpr int MODE_ISOTROPIC = 0; // fit_isotropic_b_approximately(): sx, sy, sxx, sxy, n @@ -54,7 +57,7 @@ __global__ void solvent_amplitudes(const double *__restrict__ stol2, const float template __global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict__ stol2, const float *__restrict__ fobs, const unsigned char *__restrict__ strong, - const float *__restrict__ f_abs_all, int n, + const float *__restrict__ f_abs_all, const float *__restrict__ weight_all, int n, const ModelScaleFitState *__restrict__ fits, Constraints c, double *__restrict__ partial) { constexpr int NS = 1 + NA * (NA + 1) / 2 + NA; // WSSR, J^T J lower half, J^T r const ModelScaleFitState f = fits[blockIdx.y]; @@ -105,10 +108,12 @@ __global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict const float r_u_r_f = fx * fx * uf[0] + fy * fy * uf[1] + fz * fz * uf[2] + 2 * (fx * fy * uf[3] + fx * fz * uf[4] + fy * fz * uf[5]); const float k_aniso = expf(-0.25f * r_u_r_f); + // the robust fit's weight (ModelScaling.cpp, ComputeLMMatrices and ComputeWSSR multiply by it) + const double w = f.weighted ? weight_all[static_cast(f.column) * n + i] : 1.0; if (f.mode == MODE_WSSR) { const float value = f_abs * static_cast(f.k_overall * k_aniso); const double dy = fo - value; - acc[0] += dy * dy; + acc[0] += dy * dy * w; continue; } @@ -124,13 +129,13 @@ __global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict dy_da[j] = r[0] * du[0] + r[1] * du[1] + r[2] * du[2] + r[3] * du[3] + r[4] * du[4] + r[5] * du[5]; } const double dy = fo - y; - acc[0] += dy * dy; + acc[0] += dy * dy * w; int s = 1; for (int j = 0; j < NA; j++) for (int k = 0; k <= j; k++) - acc[s++] += dy_da[j] * dy_da[k]; + acc[s++] += dy_da[j] * dy_da[k] * w; for (int j = 0; j < NA; j++) - acc[s++] += dy * dy_da[j]; + acc[s++] += dy * dy_da[j] * w; } // The block's sum of each slot the mode fills: a fixed shuffle tree within each warp, then the warps @@ -155,6 +160,36 @@ __global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict } } +// The robust fit's residuals |Fo - value| of each fit (blockIdx.y) at its parameters, compute_value() as in MODE_R. +__global__ void robust_residuals(const int *__restrict__ hkl, const float *__restrict__ fobs, + const float *__restrict__ f_abs_all, int n, + const ModelScaleFitState *__restrict__ fits, float *__restrict__ residual) { + const ModelScaleFitState f = fits[blockIdx.y]; + const float *f_abs_fit = f_abs_all + static_cast(f.column) * n; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n; i += BLOCKS * THREADS) { + const double hx = hkl[3 * i], hy = hkl[3 * i + 1], hz = hkl[3 * i + 2]; + const double *u = f.b_star; + const double r_u_r = hx * hx * u[0] + hy * hy * u[1] + hz * hz * u[2] + + 2 * (hx * hy * u[3] + hx * hz * u[4] + hy * hz * u[5]); + const float value = f_abs_fit[i] * static_cast(f.k_overall * exp(-0.25 * r_u_r)); + residual[static_cast(f.column) * n + i] = fabsf(fobs[i] - value); + } +} + +// The robust fit's weights 1 / sqrt(1 + (r / c)^2) from the residuals (ModelScaling.cpp, RobustReweight); a fit +// whose c is not positive keeps its weights. +__global__ void robust_weights(const float *__restrict__ residual, int n, const ModelScaleFitState *__restrict__ fits, + float *__restrict__ weight) { + const ModelScaleFitState f = fits[blockIdx.y]; + if (!(f.c > 0)) + return; + const size_t row = static_cast(f.column) * n; + for (int i = blockIdx.x * THREADS + threadIdx.x; i < n; i += BLOCKS * THREADS) { + const double q = residual[row + i] / f.c; + weight[row + i] = static_cast(1.0 / sqrt(1.0 + q * q)); + } +} + // Each fit's block partials summed in block order. __global__ void sum_blocks(const double *__restrict__ partial, double *__restrict__ sums) { const int s = threadIdx.x; @@ -303,6 +338,10 @@ struct ModelScaleGPU::LevMarRun { } }; +size_t ModelScaleGPU::RobustBytes(size_t max_points) { + return 2 * MAX_ROBUST_FITS * max_points * sizeof(float); +} + size_t ModelScaleGPU::DeviceBytes(size_t max_points) { return max_points * (3 * sizeof(int) + sizeof(double) + sizeof(float) + sizeof(unsigned char)) + MAX_FITS * max_points * sizeof(float) @@ -368,12 +407,12 @@ const double *ModelScaleGPU::Reduce(int nfits) { c.row[j][k] = static_cast(constraints_[j][k]); const dim3 grid(BLOCKS, nfits); switch (n_params_) { - case 2: scale_sums<2><<>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break; - case 3: scale_sums<3><<>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break; - case 4: scale_sums<4><<>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break; - case 5: scale_sums<5><<>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break; - case 6: scale_sums<6><<>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break; - case 7: scale_sums<7><<>>(hkl_, stol2_, fobs_, strong_, f_abs_, n_, fits_, c, partial_); break; + case 2: scale_sums<2><<>>(hkl_, stol2_, fobs_, strong_, f_abs_, weight_, n_, fits_, c, partial_); break; + case 3: scale_sums<3><<>>(hkl_, stol2_, fobs_, strong_, f_abs_, weight_, n_, fits_, c, partial_); break; + case 4: scale_sums<4><<>>(hkl_, stol2_, fobs_, strong_, f_abs_, weight_, n_, fits_, c, partial_); break; + case 5: scale_sums<5><<>>(hkl_, stol2_, fobs_, strong_, f_abs_, weight_, n_, fits_, c, partial_); break; + case 6: scale_sums<6><<>>(hkl_, stol2_, fobs_, strong_, f_abs_, weight_, n_, fits_, c, partial_); break; + case 7: scale_sums<7><<>>(hkl_, stol2_, fobs_, strong_, f_abs_, weight_, n_, fits_, c, partial_); break; default: throw std::logic_error("ModelScaleGPU: unexpected number of parameters"); } cuda_err(cudaGetLastError()); @@ -414,7 +453,11 @@ void ModelScaleGPU::FitBatch(const float2 *d_fcmol, const float2 *d_fmask, std:: const double b[6] = {elem(0, 0), elem(1, 1), elem(2, 2), elem(0, 1), elem(0, 2), elem(1, 2)}; std::copy(b, b + 6, fits[i].b_star); } + LevMarBatch(fits); +} +void ModelScaleGPU::LevMarBatch(std::vector &fits) { + const size_t nfits = fits.size(); // get_parameters() and set_parameters() with k_sol and b_sol fixed const int na = n_params_; auto get_parameters = [&](const ModelScaleFitState &f) { @@ -459,6 +502,51 @@ void ModelScaleGPU::FitBatch(const float2 *d_fcmol, const float2 *d_fmask, std:: set_parameters(fits[i], runs[i].Result()); } +// ModelScaling.cpp's robust rounds of each fit: the residuals at its parameters, their median taken once as c, +// the weights, and gemmi's Levenberg-Marquardt again from where the fit stands. The median is the element +// std::nth_element puts at n / 2, as there. +void ModelScaleGPU::RobustBatch(std::vector &fits, size_t nthreads) { + const size_t nfits = fits.size(); + if (nfits > static_cast(MAX_ROBUST_FITS)) + throw std::logic_error("ModelScaleGPU: too many robust fits in one launch"); + if (!weight_) { + weight_ = CudaDevicePtr(MAX_ROBUST_FITS * max_points_); + residual_ = CudaDevicePtr(MAX_ROBUST_FITS * max_points_); + host_residual_ = CudaHostPtr(MAX_ROBUST_FITS * max_points_); + } + for (size_t i = 0; i < nfits; i++) { + fits[i].column = static_cast(i); + fits[i].weighted = 0; + fits[i].c = 0; + } + const dim3 grid(BLOCKS, static_cast(nfits)); + for (int round = 0; round < ROBUST_ROUNDS; round++) { + UploadFits(fits); + robust_residuals<<>>(hkl_, fobs_, f_abs_, n_, fits_, residual_); + cuda_err(cudaGetLastError()); + if (round == 0) { + cuda_err(cudaMemcpyAsync(host_residual_, residual_, nfits * n_ * sizeof(float), cudaMemcpyDeviceToHost, + stream_)); + cuda_err(cudaStreamSynchronize(stream_)); + ParallelFor(static_cast(nfits), nthreads, [&](int i) { + std::vector r(host_residual_.get() + static_cast(i) * n_, + host_residual_.get() + static_cast(i + 1) * n_); + std::nth_element(r.begin(), r.begin() + r.size() / 2, r.end()); + fits[i].c = r[r.size() / 2]; + }); + UploadFits(fits); + } + robust_weights<<>>(residual_, n_, fits_, weight_); + cuda_err(cudaGetLastError()); + cuda_err(cudaStreamSynchronize(stream_)); // host_fits_ is written again below + for (ModelScaleFitState &f : fits) + f.weighted = f.c > 0; + LevMarBatch(fits); + } + for (ModelScaleFitState &f : fits) + f.weighted = 0; +} + ModelScaleParams ModelScaleGPU::Fit(const float2 *d_fcmol, const float2 *d_fmask, double k_sol, double b_sol) { ModelScaleParams out; if (n_ == 0) @@ -473,7 +561,8 @@ ModelScaleParams ModelScaleGPU::Fit(const float2 *d_fcmol, const float2 *d_fmask // FitModelScale's grid (ModelScaling.cpp, default ModelScaleBox with b_sol up to b_hi): a coarse pass // over the box and a finer one around its winner, every grid point its own fit, the winner the lowest // finite R in grid order. The fits of a pass run together, one launch per Levenberg-Marquardt step. -ModelSolventFit ModelScaleGPU::FitSolvent(const float2 *d_fcmol, const float2 *d_fmask, double b_hi) { +ModelSolventFit ModelScaleGPU::FitSolvent(const float2 *d_fcmol, const float2 *d_fmask, double b_hi, bool robust, + size_t nthreads) { ModelSolventFit out; if (n_ < 20) return out; @@ -483,8 +572,10 @@ ModelSolventFit ModelScaleGPU::FitSolvent(const float2 *d_fcmol, const float2 *d const double k_lo = 0.10, k_hi = 0.60, b_lo = 10.0; // ModelScaleBox{} const double b_step = std::max(10.0, (b_hi - b_lo) / 7.0); double best_r = -1; - auto try_points = [&](std::vector &fits) { + auto try_points = [&](std::vector &fits, bool robust_pass) { FitBatch(d_fcmol, d_fmask, fits); + if (robust_pass) + RobustBatch(fits, nthreads); for (auto &f : fits) f.mode = MODE_R; const double *r = Sums(fits); @@ -509,7 +600,7 @@ ModelSolventFit ModelScaleGPU::FitSolvent(const float2 *d_fcmol, const float2 *d for (double ks = k_lo; ks <= k_hi + 1e-9; ks += 0.05) for (double bs = b_lo; bs <= b_hi + 1e-9; bs += b_step) coarse.push_back(point(ks, bs)); - try_points(coarse); + try_points(coarse, false); const double k0 = out.k_sol, b0 = out.b_sol; const double k_hi2 = std::min(k_hi, k0 + 0.05); const double b_hi2 = std::min(b_hi, b0 + b_step); @@ -517,7 +608,10 @@ ModelSolventFit ModelScaleGPU::FitSolvent(const float2 *d_fcmol, const float2 *d for (double ks = std::max(k_lo, k0 - 0.05); ks <= k_hi2 + 1e-9; ks += 0.025) for (double bs = std::max(b_lo, b0 - b_step); bs <= b_hi2 + 1e-9; bs += b_step / 2.0) fine.push_back(point(ks, bs)); - try_points(fine); + // a robust fit takes the fine pass robustly, the coarse winner included (as FitModelScale) + if (robust) + best_r = -1; + try_points(fine, robust); out.r = best_r; return out; } diff --git a/image_analysis/structure_refinement/ModelScaleGPU.h b/image_analysis/structure_refinement/ModelScaleGPU.h index 3906b5c53..e211a0aa6 100644 --- a/image_analysis/structure_refinement/ModelScaleGPU.h +++ b/image_analysis/structure_refinement/ModelScaleGPU.h @@ -36,6 +36,8 @@ struct ModelScaleFitState { double k_sol, b_sol; int mode; int column; // the fit's row of |Fcalc + solvent| on the device + int weighted; // the robust fit's weights (row `column`) multiply the WSSR and J^T J, J^T r; 0: weight 1 + double c; // the robust fit's scale of the residuals (ModelScaling.cpp, RobustReweight) }; class ModelScaleGPU { @@ -59,7 +61,12 @@ public: // FitModelScale (ModelScaling.cpp) with the default box, b_sol up to b_hi: the same grid, the same // winner. Throws ModelScaleGPUTooFewReflections where FitModelScale would chain the grid points // (five or fewer reflections for the isotropic fit). - ModelSolventFit FitSolvent(const float2 *d_fcmol, const float2 *d_fmask, double b_hi = 80.0); + // robust: the fine pass by robust least squares, as FitModelScale with ModelScaleBox::robust - the + // reweighting on the device, the median of each fit's first residuals on the host on `nthreads` threads. + ModelSolventFit FitSolvent(const float2 *d_fcmol, const float2 *d_fmask, double b_hi = 80.0, + bool robust = false, size_t nthreads = 1); + // What a robust FitSolvent takes on the card beyond DeviceBytes(), reserved at its first call. + static size_t RobustBytes(size_t max_points); private: struct LevMarRun; @@ -72,6 +79,10 @@ private: const double *Reduce(int nfits); const double *Sums(const std::vector &fits); void FitBatch(const float2 *d_fcmol, const float2 *d_fmask, std::vector &fits); + // gemmi's Levenberg-Marquardt for every fit in `fits` from the parameters it holds. + void LevMarBatch(std::vector &fits); + // The robust rounds of a batch FitBatch() has fitted: reweight, fit again, ROBUST_ROUNDS times. + void RobustBatch(std::vector &fits, size_t nthreads); cudaStream_t stream_; size_t max_points_; @@ -91,4 +102,6 @@ private: CudaDevicePtr sums_; // per fit, per slot CudaHostPtr host_fits_; CudaHostPtr host_sums_; + CudaDevicePtr weight_, residual_; // the robust fit's, per fit of a batch, per point + CudaHostPtr host_residual_; }; diff --git a/image_analysis/structure_refinement/ModelScaling.cpp b/image_analysis/structure_refinement/ModelScaling.cpp index 022a9fa73..6fb1dfb84 100644 --- a/image_analysis/structure_refinement/ModelScaling.cpp +++ b/image_analysis/structure_refinement/ModelScaling.cpp @@ -471,15 +471,26 @@ double SolventBMax(const gemmi::Model &model) { } #ifdef JFJOCH_USE_CUDA +// A fit is a chain of small launches, each waited for by the host's Levenberg-Marquardt: on a card shared with +// long kernels (refinement's trials side by side) every wait would queue behind them, so the stream goes first. +int HighestStreamPriority() { + int least = 0, greatest = 0; + cudaDeviceGetStreamPriorityRange(&least, &greatest); + return greatest; +} + struct ModelScaleGPUPool::Engine { CudaStream stream; ModelScaleGPU scale; CudaDevicePtr fcalc, fmask; - explicit Engine(size_t max_points) : scale(stream, max_points), fcalc(max_points), fmask(max_points) {} + explicit Engine(size_t max_points) + : stream(cudaStreamNonBlocking, HighestStreamPriority()), scale(stream, max_points), fcalc(max_points), + fmask(max_points) {} }; -size_t ModelScaleGPUPool::DeviceBytes(size_t max_points, size_t engines) { - return engines * (ModelScaleGPU::DeviceBytes(max_points) + 2 * max_points * sizeof(float2)); +size_t ModelScaleGPUPool::DeviceBytes(size_t max_points, size_t engines, bool robust) { + return engines * (ModelScaleGPU::DeviceBytes(max_points) + 2 * max_points * sizeof(float2) + + (robust ? ModelScaleGPU::RobustBytes(max_points) : 0)); } ModelScaleGPUPool::ModelScaleGPUPool(int device, size_t max_points, size_t engines) @@ -500,7 +511,7 @@ ModelScaleGPUPool::~ModelScaleGPUPool() { set_gpu(current); } -ModelScaleReport ModelScaleGPUPool::Fit(gemmi::Scaling &scaling, double b_hi, size_t nthreads) { +ModelScaleReport ModelScaleGPUPool::Fit(gemmi::Scaling &scaling, double b_hi, size_t nthreads, bool robust) { const size_t n = scaling.points.size(); std::vector> hkl(n); std::vector stol2(n); @@ -543,11 +554,11 @@ ModelScaleReport ModelScaleGPUPool::Fit(gemmi::Scaling &scaling, double b if (cudaMemcpyAsync(e->fcalc, fc.data(), n * sizeof(float2), cudaMemcpyHostToDevice, e->stream) != cudaSuccess || cudaMemcpyAsync(e->fmask, fm.data(), n * sizeof(float2), cudaMemcpyHostToDevice, e->stream) != cudaSuccess) throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "ModelScaleGPUPool: the upload failed"); - fit = e->scale.FitSolvent(e->fcalc, e->fmask, b_hi); + fit = e->scale.FitSolvent(e->fcalc, e->fmask, b_hi, robust, nthreads); } catch (const ModelScaleGPUTooFewReflections &) { set_gpu(current); release(); - return FitModelScale(scaling, {.b_hi = b_hi}, nthreads); + return FitModelScale(scaling, {.b_hi = b_hi, .robust = robust}, nthreads); } catch (...) { set_gpu(current); release(); diff --git a/image_analysis/structure_refinement/ModelScaling.h b/image_analysis/structure_refinement/ModelScaling.h index bd1434c04..9ec565a24 100644 --- a/image_analysis/structure_refinement/ModelScaling.h +++ b/image_analysis/structure_refinement/ModelScaling.h @@ -63,15 +63,17 @@ ModelScaleReport FitModelScale(gemmi::Scaling &scaling, ModelScaleBox box class ModelScaleGPUPool { public: // What `engines` engines for fits of up to `max_points` reflections take on the card. - static size_t DeviceBytes(size_t max_points, size_t engines); + static size_t DeviceBytes(size_t max_points, size_t engines, bool robust = false); // On `device`. Throws JFJochException on a CUDA failure. ModelScaleGPUPool(int device, size_t max_points, size_t engines); ~ModelScaleGPUPool(); - // FitModelScale(scaling, {.b_hi = b_hi}, nthreads). Where FitModelScale would walk the grid point to - // point (five or fewer reflections for the isotropic start), it is FitModelScale on the host. Safe to - // call from several threads. Throws JFJochException on a CUDA failure. - ModelScaleReport Fit(gemmi::Scaling &scaling, double b_hi, size_t nthreads); + // FitModelScale(scaling, {.b_hi = b_hi, .robust = robust}, nthreads). Where FitModelScale would walk the + // grid point to point (five or fewer reflections for the isotropic start), it is FitModelScale on the host. + // Safe to call from several threads. Throws JFJochException on a CUDA failure. The robust fit's reweighting + // runs on the device; it is FitModelScale's robust fit to rounding (float residuals and weights, the + // anisotropic factor of the steps in float), deterministic run to run. + ModelScaleReport Fit(gemmi::Scaling &scaling, double b_hi, size_t nthreads, bool robust = false); // What the pool holds on the card. size_t Bytes() const { return DeviceBytes(max_points_, engines_.size()); } diff --git a/image_analysis/structure_refinement/Refinement.cpp b/image_analysis/structure_refinement/Refinement.cpp index 244fe68fe..e5547eda1 100644 --- a/image_analysis/structure_refinement/Refinement.cpp +++ b/image_analysis/structure_refinement/Refinement.cpp @@ -100,6 +100,9 @@ struct Context { RefinementReflections refl; std::unique_ptr engine; std::unique_ptr mask; +#ifdef JFJOCH_USE_CUDA + std::unique_ptr scale_gpu; +#endif size_t nthreads = 1; int device = -1; // the engine's GPU size_t engine_bytes = 0; // and what it holds there @@ -193,8 +196,13 @@ struct Context { gemmi::Scaling scaling(cell, sg); scaling.use_solvent = true; scaling.prepare_points(fc, fo, &fm); - // the scale of the score: by robust least squares (ModelScaleBox::robust), on the CPU - FitModelScale(scaling, {.b_hi = SolventBMax(Model()), .robust = true}, nthreads); + // the scale of the score: by robust least squares (ModelScaleBox::robust) +#ifdef JFJOCH_USE_CUDA + if (scale_gpu) + scale_gpu->Fit(scaling, SolventBMax(Model()), nthreads, true); + else +#endif + FitModelScale(scaling, {.b_hi = SolventBMax(Model()), .robust = true}, nthreads); s3.reset(); k_sol = scaling.k_sol; b_sol = scaling.b_sol; @@ -668,6 +676,7 @@ void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logg const RefinementStage stage("engine"); ctx.engine.reset(); #ifdef JFJOCH_USE_CUDA + ctx.scale_gpu.reset(); if (gpu && get_gpu_count() > 0) { try { const int device = get_gpu(); @@ -679,12 +688,15 @@ void MakeEngine(Context &ctx, bool gpu, const RefinementSettings &settings, Logg ctx.b_ceiling, budget, &bytes, settings.nthreads); ctx.device = device; ctx.engine_bytes = bytes; - if (!ctx.engine) + if (ctx.engine) + ctx.scale_gpu = std::make_unique(device, ctx.refl.Size(), 1); + else logger.Info("Refinement: on the CPU - the GPU cannot take this case or {:.1f} GB is over its budget", bytes / 1e9); } catch (const JFJochException &e) { logger.Warning("Refinement: the GPU failed ({}), on the CPU", e.what()); ctx.engine.reset(); + ctx.scale_gpu.reset(); } } #else @@ -752,6 +764,7 @@ std::unique_ptr MakeOwnSlot(const gemmi::Structure &before_topology, co ctx.b_ceiling, SIZE_MAX, nullptr, nthreads); if (!ctx.engine) return nullptr; + ctx.scale_gpu = std::make_unique(device, ctx.refl.Size(), 1); } catch (const JFJochException &) { return nullptr; } diff --git a/tests/ModelScaleGPUTest.cpp b/tests/ModelScaleGPUTest.cpp index 3ceb54200..d9e608ce1 100644 --- a/tests/ModelScaleGPUTest.cpp +++ b/tests/ModelScaleGPUTest.cpp @@ -392,6 +392,45 @@ TEST_CASE("ModelScaleGPUPool_MatchesFitModelScale", "[ModelValidation][gpu]") { } } +// Refinement's robust fit (ModelScaleBox::robust) on the card: the reweighting on the device, the answer +// FitModelScale's to rounding where one reflection in fifty is a gross outlier - and the same, to the bit, run to +// run and whichever engine a fit got. +TEST_CASE("ModelScaleGPUPool_RobustMatchesFitModelScale", "[ModelValidation][gpu]") { + if (get_gpu_count() == 0) + SKIP("No GPU"); + for (const Case &c : CASES) { + INFO(c.name); + gemmi::UnitCell cell(c.a, c.b, c.c, c.alpha, c.beta, c.gamma); + const gemmi::SpaceGroup &sg = *gemmi::find_spacegroup_by_name(c.hm); + gemmi::Scaling points = MakeScaling(cell, sg, 2.8); + for (size_t i = 0; i < points.points.size(); i += 50) + points.points[i].fobs *= 4.0f; + gemmi::Scaling cpu = points, cpu_plain = points; + FitModelScale(cpu, {.robust = true}, 8); + FitModelScale(cpu_plain, {}, 8); + + ModelScaleGPUPool pool(0, points.points.size(), 2); + std::vector> gpu(4, points); + std::vector> fits; + for (size_t i = 0; i < gpu.size(); i++) + fits.push_back(std::async(std::launch::async, [&, i] { pool.Fit(gpu[i], 80.0, 2, true); })); + for (auto &f : fits) + f.get(); + + // the robust answer is not the plain one, and the card's is the host's to rounding + CHECK(std::fabs(cpu.k_overall - cpu_plain.k_overall) > 1e-3 * cpu.k_overall); + CHECK(gpu[0].k_overall == Catch::Approx(cpu.k_overall).epsilon(1e-3)); + CHECK(gpu[0].k_sol == Catch::Approx(cpu.k_sol).margin(0.026)); + CHECK(std::fabs(gpu[0].calculate_r_factor() - cpu.calculate_r_factor()) < 1e-4); + for (size_t i = 1; i < gpu.size(); i++) { + CHECK(gpu[i].k_sol == gpu[0].k_sol); + CHECK(gpu[i].b_sol == gpu[0].b_sol); + CHECK(gpu[i].k_overall == gpu[0].k_overall); + CHECK(std::memcmp(&gpu[i].b_star, &gpu[0].b_star, sizeof(gpu[0].b_star)) == 0); + } + } +} + // ModelScaling fits past a reflection the solvent cancels exactly (ModelScalingTest.cpp), on the card. TEST_CASE("ModelScaleGPUPool_FitsPastAReflectionTheSolventCancelsExactly", "[ModelValidation][gpu]") { if (get_gpu_count() == 0)