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

RESULT-CHANGING: ModelScaleGPU: J^T J and J^T r summed in float, the R and the robust residuals' scale in float

Browse Source 
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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HVeDQqN9qRD57vnShzTC37
This commit is contained in:
leonarski_fandClaude Opus 5.5 committed 2026-10-28 09:46:26 +01:00
1 parent c5b7f18f67
commit fdafc009bc
1 file changed
+39 -26
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/structure_refinement/ModelScaleGPU.cu
+39 -26
View File
@@ -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<float>(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<float>(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<size_t>(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<float>(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<float>(dy), w_f = static_cast<float>(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<size_t>(f.column) * n;
float uf[6];
for (int k = 0; k < 6; k++)
uf[k] = static_cast<float>(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<float>(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<float>(f.k_overall * expf(-0.25f * r_u_r_f));
residual[static_cast<size_t>(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<size_t>(f.column) * n;
const float c = static_cast<float>(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<float>(1.0 / sqrt(1.0 + q * q));
const float q = residual[row + i] / c;
residual[row + i] = 1.0f / sqrtf(1.0f + q * q);
}
}