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

RefinementGPU: the symmetry phases from a table instead of sincospi in the compose and scatter kernels

Browse Source 
exp(2 pi i h.t) depends only on the integer h.t (in units of 1/den), so the engine tabulates it once over the
range the rows and operators make, with the same op_phase() at the same double argument - the same bits. The
compose kernel reads the table at h.t; the scatter terms carry their table entry, and their row directly, so
neither hot kernel divides: the rotations go to the device in whole units (gemmi's are den times an integer
matrix), the term's row no longer has to be split from its operator. The table is a few thousand entries, the
scatter terms one int wider.

Outputs are byte-identical; the X-ray evaluation counter drops 7-10% on the demo-sized sets.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
This commit is contained in:
leonarski_fandClaude Fable 5.1 committed 2026-10-26 20:58:10 +01:00
1 parent dfb1536fe2
commit 128e3bb9ae
1 file changed
+81 -31
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/structure_refinement/RefinementGPU.cu
+81 -31
View File
@@ -46,19 +46,31 @@ __device__ __forceinline__ float expapprox(float x) {
return a * (0.509871020f + b * (0.312146713f + b * (0.166617139f + b * (-2.190619930e-3f + b * 1.3555747234e-2f))));
}
// exp(+2 pi i h.t) of row h and operator o, exact for the fractions a translation is made of.
__device__ __forceinline__ void op_phase(const int *h, const int *tran, int den, double &c, double &s) {
sincospi(2.0 * (h[0] * tran[0] + h[1] * tran[1] + h[2] * tran[2]) / den, &s, &c);
// exp(+2 pi i h.t) for h.t = ht / den, exact for the fractions a translation is made of.
__device__ __forceinline__ void op_phase(int ht, int den, double &c, double &s) {
sincospi(2.0 * ht / den, &s, &c);
}
// The phases of every h.t the rows and operators make, as a table over the integer h.t from `min`: the hot
// kernels read it in place of op_phase() - the same function at the same argument, so the same bits.
__global__ void phase_table_kernel(int min, int count, int den, double2 *phase) {
const int i = blockIdx.x * blockDim.x + threadIdx.x;
if (i >= count)
return;
double s, c;
op_phase(min + i, den, c, s);
phase[i] = make_double2(c, s);
}
// F of the atoms at each row, composed from the half-u transform of one copy (ModelStructureFactorsGPU.cu's
// compose_kernel), and with `evaluate` the target's term of a working row and the coefficient
// conj(dT/dF) n_cen unblur the gradient's map is made of.
// conj(dT/dF) n_cen unblur the gradient's map is made of. `ops` holds each operator's rotation in whole
// units and its translation in units of 1 / den.
__global__ void compose_target_kernel(const float2 *f1, float norm, int nu, int nv, int nw, const int *ops, int n_ops,
int den, const int *row_hkl, const double *row_scale, int rows, bool evaluate,
const float *k_total, const float2 *fmask, const float *ml_d, const float *ml_var,
const float *fo, const char *work, const char *centric, float2 *f, double *nll,
float2 *coef) {
const double2 *phase, int phase_min, const int *row_hkl, const double *row_scale,
int rows, bool evaluate, const float *k_total, const float2 *fmask,
const float *ml_d, const float *ml_var, const float *fo, const char *work,
const char *centric, float2 *f, double *nll, float2 *coef) {
const int m = blockIdx.x * blockDim.x + threadIdx.x;
if (m >= rows)
return;
@@ -68,17 +80,16 @@ __global__ void compose_target_kernel(const float2 *f1, float norm, int nu, int
const int *rot = ops + 12 * o, *tran = rot + 9;
int k[3];
for (int i = 0; i < 3; i++)
k[i] = (rot[i] * h[0] + rot[3 + i] * h[1] + rot[6 + i] * h[2]) / den;
k[i] = rot[i] * h[0] + rot[3 + i] * h[1] + rot[6 + i] * h[2];
const bool conj = k[0] < 0;
if (conj)
for (int i = 0; i < 3; i++)
k[i] = -k[i];
const float2 out = f1[k[0] + static_cast<size_t>(nu / 2 + 1) * (imod(k[1], nv) + static_cast<size_t>(nv) * imod(k[2], nw))];
const double vr = out.x * norm, vi = (conj ? out.y : -out.y) * norm;
double s, c;
op_phase(h, tran, den, c, s);
sr += c * vr - s * vi;
si += c * vi + s * vr;
const double2 p = phase[h[0] * tran[0] + h[1] * tran[1] + h[2] * tran[2] - phase_min];
sr += p.x * vr - p.y * vi;
si += p.x * vi + p.y * vr;
}
const double scale = row_scale[m];
sr *= scale;
@@ -105,9 +116,10 @@ __global__ void compose_target_kernel(const float2 *f1, float norm, int nu, int
}
// The gradient's coefficients on the half-u spectrum: each point the sum, in order, of the terms
// (coef[row] exp(2 pi i h.t_op)) / 2 that land on it, conjugated where they land as a Friedel mate.
__global__ void scatter_kernel(const float2 *coef, const int *row_hkl, const int *ops, int n_ops, int den,
const int *point, const int *start, const int *term, int points, float2 *spectrum) {
// (coef[row] exp(2 pi i h.t_op)) / 2 that land on it, conjugated where they land as a Friedel mate. A term is
// its row * 2 + conj, and its phase the table's entry at term_phase.
__global__ void scatter_kernel(const float2 *coef, const double2 *phase, const int *point, const int *start,
const int *term, const int *term_phase, int points, float2 *spectrum) {
const int s = blockIdx.x * blockDim.x + threadIdx.x;
if (s >= points)
return;
@@ -115,10 +127,9 @@ __global__ void scatter_kernel(const float2 *coef, const int *row_hkl, const int
for (int e = start[s]; e < start[s + 1]; e++) {
const int t = term[e];
const bool conj = t & 1;
const int m = (t >> 1) / n_ops, o = (t >> 1) % n_ops;
double c, sn;
op_phase(row_hkl + 3 * m, ops + 12 * o + 9, den, c, sn);
const double ar = 0.5 * (coef[m].x * c - coef[m].y * sn), ai = 0.5 * (coef[m].x * sn + coef[m].y * c);
const int m = t >> 1;
const double2 p = phase[term_phase[e]];
const double ar = 0.5 * (coef[m].x * p.x - coef[m].y * p.y), ai = 0.5 * (coef[m].x * p.y + coef[m].y * p.x);
sr += ar;
si += conj ? -ai : ai;
}
@@ -359,6 +370,22 @@ void Ensure(CudaDevicePtr<T> &p, size_t &capacity, size_t n) {
p = CudaDevicePtr<T>(capacity, CudaAlloc::Synchronous);
}
// h.t of a row and an operator, in units of 1 / den.
int HT(const std::array<int, 3> &h, const std::array<int, 12> &op) {
return h[0] * op[9] + h[1] * op[10] + h[2] * op[11];
}
// The range of h.t over the rows and operators: the phase table's.
void PhaseRange(const RefinementGPUSetup &setup, int &min, int &max) {
min = max = 0;
for (const auto &h : setup.rows)
for (const auto &op : setup.sym_ops) {
const int ht = HT(h, op);
min = std::min(min, ht);
max = std::max(max, ht);
}
}
// cub's temporary storage to reduce n doubles.
size_t ReduceBytes(size_t n) {
size_t bytes = 0;
@@ -410,7 +437,9 @@ struct RefinementGPUEngineImpl {
CudaDevicePtr<double> nll, nll_sum;
CudaDevicePtr<char> reduce_temp;
size_t reduce_bytes = 0;
CudaDevicePtr<int> scatter_point, scatter_start, scatter_term;
CudaDevicePtr<int> scatter_point, scatter_start, scatter_term, scatter_phase;
CudaDevicePtr<double2> phase; // exp(2 pi i h.t) over the integer h.t from phase_min
int phase_min = 0;
CudaDevicePtr<float4> grad;
std::unique_ptr<ModelDensityGPU> density;
std::unique_ptr<ModelMaskGPU> mask;
@@ -452,6 +481,11 @@ struct RefinementGPUEngineImpl {
scatter_point = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_point.size(), 1), sync);
scatter_start = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_start.size(), 1), sync);
scatter_term = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_term.size(), 1), sync);
scatter_phase = CudaDevicePtr<int>(std::max<size_t>(setup.scatter_term.size(), 1), sync);
int phase_max = 0;
PhaseRange(setup, phase_min, phase_max);
const int phase_count = phase_max - phase_min + 1;
phase = CudaDevicePtr<double2>(phase_count, sync);
density = std::make_unique<ModelDensityGPU>(stream, setup.max_atoms, setup.max_pairs,
ModelDensityGPU::Bricks(setup.grid.nu, setup.grid.nv, setup.grid.nw));
mask = std::make_unique<ModelMaskGPU>(stream, GridPoints(setup.mask_grid));
@@ -467,12 +501,24 @@ struct RefinementGPUEngineImpl {
mask_spectrum = CudaDevicePtr<float2>(ComplexPoints(setup.mask_grid), sync);
}
// the rotations in whole units (gemmi's are den times an integer matrix), the translations as given
std::vector<int> flat_ops;
for (const auto &o : setup.sym_ops)
flat_ops.insert(flat_ops.end(), o.begin(), o.end());
for (const auto &o : setup.sym_ops) {
for (int i = 0; i < 9; i++)
flat_ops.push_back(o[i] / setup.den);
flat_ops.insert(flat_ops.end(), o.begin() + 9, o.end());
}
std::vector<int> hkl;
for (const auto &h : setup.rows)
hkl.insert(hkl.end(), h.begin(), h.end());
// each scatter term as row * 2 + conj, and its phase's place in the table
const int n_ops = static_cast<int>(setup.sym_ops.size());
std::vector<int> term, term_phase;
for (const int t : setup.scatter_term) {
const int m = (t >> 1) / n_ops, o = (t >> 1) % n_ops;
term.push_back(m * 2 + (t & 1));
term_phase.push_back(HT(setup.rows[m], setup.sym_ops[o]) - phase_min);
}
Upload(ops, flat_ops, stream);
Upload(row_hkl, hkl, stream);
Upload(fo, setup.fo, stream);
@@ -480,7 +526,10 @@ struct RefinementGPUEngineImpl {
Upload(centric, setup.centric, stream);
Upload(scatter_point, setup.scatter_point, stream);
Upload(scatter_start, setup.scatter_start, stream);
Upload(scatter_term, setup.scatter_term, stream);
Upload(scatter_term, term, stream);
Upload(scatter_phase, term_phase, stream);
phase_table_kernel<<<blocks(phase_count), THREADS, 0, stream>>>(phase_min, phase_count, setup.den, phase);
cuda_err(cudaGetLastError());
const double inv_n[3] = {1.0 / setup.grid.nu, 1.0 / setup.grid.nv, 1.0 / setup.grid.nw};
geom.nu = setup.grid.nu;
@@ -522,8 +571,8 @@ struct RefinementGPUEngineImpl {
return;
compose_target_kernel<<<blocks(rows), THREADS, 0, stream>>>(
spectrum, Norm(), setup.grid.nu, setup.grid.nv, setup.grid.nw, ops, static_cast<int>(setup.sym_ops.size()),
setup.den, row_hkl, row_scale, rows, evaluate, k_total, fmask_scaled, ml_d, ml_var, fo, work, centric, f,
nll, coef);
phase, phase_min, row_hkl, row_scale, rows, evaluate, k_total, fmask_scaled, ml_d, ml_var, fo, work, centric,
f, nll, coef);
cuda_err(cudaGetLastError());
}
@@ -543,10 +592,8 @@ struct RefinementGPUEngineImpl {
cuda_err(cudaMemsetAsync(spectrum, 0, ComplexPoints(g) * sizeof(float2), stream));
const int points = static_cast<int>(setup.scatter_point.size());
if (points > 0) {
scatter_kernel<<<blocks(points), THREADS, 0, stream>>>(coef, row_hkl, ops,
static_cast<int>(setup.sym_ops.size()), setup.den,
scatter_point, scatter_start, scatter_term,
points, spectrum);
scatter_kernel<<<blocks(points), THREADS, 0, stream>>>(coef, phase, scatter_point, scatter_start,
scatter_term, scatter_phase, points, spectrum);
cuda_err(cudaGetLastError());
}
cufft_err(inverse.C2R(reinterpret_cast<cufftComplex *>(spectrum.get()), this->grid));
@@ -615,7 +662,10 @@ size_t RefinementGPUEngine::DeviceBytes(const RefinementGPUSetup &s, size_t fft_
size_t bytes = na * (2 * sizeof(float4) + 3 * sizeof(int) + sizeof(ModelMaskAtom) + sizeof(ModelDensityAtom));
bytes += GridPoints(s.grid) * sizeof(float) + ComplexPoints(s.grid) * sizeof(float2) + fft_work_bytes;
bytes += nr * (3 * sizeof(int) + 2 * sizeof(double) + 4 * sizeof(float) + 4 * sizeof(float2) + 2);
bytes += (s.scatter_point.size() + s.scatter_start.size() + s.scatter_term.size()) * sizeof(int);
bytes += (s.scatter_point.size() + s.scatter_start.size() + 2 * s.scatter_term.size()) * sizeof(int);
int phase_min = 0, phase_max = 0;
PhaseRange(s, phase_min, phase_max);
bytes += static_cast<size_t>(phase_max - phase_min + 1) * sizeof(double2);
bytes += ModelDensityGPU::DeviceBytes(s.max_atoms, s.max_pairs, ModelDensityGPU::Bricks(s.grid.nu, s.grid.nv, s.grid.nw));
bytes += ModelMaskGPU::DeviceBytes(GridPoints(s.mask_grid));
if (!SameGrid(s.mask_grid, s.grid))