From fdafc009bce937f31d8f4496f2fc78616d2df8bb Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Wed, 28 Oct 2026 09:46:26 +0100 Subject: [PATCH] RESULT-CHANGING: ModelScaleGPU: J^T J and J^T r summed in float, the R and the robust residuals' scale in float scale_sums accumulated every Levenberg-Marquardt slot in double, which on a card with little double throughput (T4, consumer cards) is most of a step. J^T J and J^T r now go in float, per point, per thread and through the warp's shuffle tree, into the double sums across warps and blocks; they only set the direction of the next step, which the WSSR - still in double - accepts or not. MODE_R and robust_residuals take the anisotropic scale in float as the WSSR already does, and robust_weights its 1 / sqrt(1 + q^2) in float. Fixed launch shape and order as before: deterministic. cytochrome C refine run (50k reflections) on an RTX 5080 under nsys: scale_sums<3> 59 -> 38 us per launch. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01HVeDQqN9qRD57vnShzTC37 --- .../structure_refinement/ModelScaleGPU.cu | 65 +++++++++++-------- 1 file changed, 39 insertions(+), 26 deletions(-) diff --git a/image_analysis/structure_refinement/ModelScaleGPU.cu b/image_analysis/structure_refinement/ModelScaleGPU.cu index 00685adb1..de11ff941 100644 --- a/image_analysis/structure_refinement/ModelScaleGPU.cu +++ b/image_analysis/structure_refinement/ModelScaleGPU.cu @@ -69,6 +69,12 @@ __global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict double acc[NS]; for (int s = 0; s < NS; s++) acc[s] = 0; + // MODE_MATRICES' J^T J and J^T r, summed in float: a card with little double throughput spends most of a step + // on these otherwise, and they only set the direction of the next step, which the WSSR (in double) then accepts + // or not + float acc_f[NS]; + for (int s = 0; s < NS; s++) + acc_f[s] = 0; for (int i = blockIdx.x * THREADS + threadIdx.x; i < n; i += BLOCKS * THREADS) { const float f_abs = f_abs_fit[i]; @@ -88,18 +94,6 @@ __global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict continue; } - if (f.mode == MODE_R) { - 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]); - // compute_value(): |F| times the scale cast to float, in float - const float value = f_abs * static_cast(f.k_overall * exp(-0.25 * r_u_r)); - acc[0] += fabs(fo - value); - acc[1] += fo; - continue; - } - // Below, the anisotropic factor is taken in float rather than in double as gemmi takes it: double // arithmetic is most of the cost of a step on a card with little double throughput. In // compute_value() gemmi casts the scale to float anyway, so the WSSR moves by the last bit of some @@ -109,6 +103,13 @@ __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); + if (f.mode == MODE_R) { + // compute_value(): |F| times the scale cast to float, in float + const float value = f_abs * static_cast(f.k_overall * k_aniso); + acc[0] += fabs(fo - value); + acc[1] += fo; + continue; + } // 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) { @@ -118,12 +119,12 @@ __global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict continue; } - // compute_value_and_derivatives() with k_sol and b_sol fixed, per point in float, summed in double + // compute_value_and_derivatives() with k_sol and b_sol fixed, per point in float const float fe = f_abs * k_aniso; const float y = static_cast(f.k_overall) * fe; const float du[6] = {-0.25f * y * (fx * fx), -0.25f * y * (fy * fy), -0.25f * y * (fz * fz), -0.5f * y * (fx * fy), -0.5f * y * (fx * fz), -0.5f * y * (fy * fz)}; - double dy_da[NA]; + float dy_da[NA]; dy_da[0] = fe; for (int j = 1; j < NA; j++) { const float *r = c.row[j - 1]; @@ -131,12 +132,13 @@ __global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict } const double dy = fo - y; acc[0] += dy * dy * w; + const float dy_f = static_cast(dy), w_f = static_cast(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] * w; + acc_f[s++] += dy_da[j] * dy_da[k] * w_f; for (int j = 0; j < NA; j++) - acc[s++] += dy * dy_da[j] * w; + acc_f[s++] += dy_f * dy_da[j] * w_f; } // The block's sum of each slot the mode fills: a fixed shuffle tree within each warp, then the warps @@ -145,9 +147,17 @@ __global__ void scale_sums(const int *__restrict__ hkl, const double *__restrict __shared__ double warp_sum[THREADS / 32][NS]; const int lane = threadIdx.x % 32, warp = threadIdx.x / 32; for (int s = 0; s < used; s++) { - double v = acc[s]; - for (int offset = 16; offset > 0; offset /= 2) - v += __shfl_down_sync(0xffffffffu, v, offset); + double v; + if (f.mode == MODE_MATRICES && s > 0) { + float vf = acc_f[s]; + for (int offset = 16; offset > 0; offset /= 2) + vf += __shfl_down_sync(0xffffffffu, vf, offset); + v = vf; + } else { + v = acc[s]; + for (int offset = 16; offset > 0; offset /= 2) + v += __shfl_down_sync(0xffffffffu, v, offset); + } if (lane == 0) warp_sum[warp][s] = v; } @@ -167,12 +177,14 @@ __global__ void robust_residuals(const int *__restrict__ hkl, const float *__res 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; + float uf[6]; + for (int k = 0; k < 6; k++) + uf[k] = static_cast(f.b_star[k]); 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)); + const float fx = hkl[3 * i], fy = hkl[3 * i + 1], fz = hkl[3 * i + 2]; + 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 value = f_abs_fit[i] * static_cast(f.k_overall * expf(-0.25f * r_u_r_f)); residual[static_cast(f.column) * n + i] = fabsf(fobs[i] - value); } } @@ -184,9 +196,10 @@ __global__ void robust_weights(float *residual, int n, const ModelScaleFitState if (!(f.c > 0)) return; const size_t row = static_cast(f.column) * n; + const float c = static_cast(f.c); for (int i = blockIdx.x * THREADS + threadIdx.x; i < n; i += BLOCKS * THREADS) { - const double q = residual[row + i] / f.c; - residual[row + i] = static_cast(1.0 / sqrt(1.0 + q * q)); + const float q = residual[row + i] / c; + residual[row + i] = 1.0f / sqrtf(1.0f + q * q); } }