diff --git a/image_analysis/structure_refinement/RefinementGPU.cu b/image_analysis/structure_refinement/RefinementGPU.cu index 33fa3baff..56a5c00e1 100644 --- a/image_analysis/structure_refinement/RefinementGPU.cu +++ b/image_analysis/structure_refinement/RefinementGPU.cu @@ -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(nu / 2 + 1) * (imod(k[1], nv) + static_cast(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 &p, size_t &capacity, size_t n) { p = CudaDevicePtr(capacity, CudaAlloc::Synchronous); } +// h.t of a row and an operator, in units of 1 / den. +int HT(const std::array &h, const std::array &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 nll, nll_sum; CudaDevicePtr reduce_temp; size_t reduce_bytes = 0; - CudaDevicePtr scatter_point, scatter_start, scatter_term; + CudaDevicePtr scatter_point, scatter_start, scatter_term, scatter_phase; + CudaDevicePtr phase; // exp(2 pi i h.t) over the integer h.t from phase_min + int phase_min = 0; CudaDevicePtr grad; std::unique_ptr density; std::unique_ptr mask; @@ -452,6 +481,11 @@ struct RefinementGPUEngineImpl { scatter_point = CudaDevicePtr(std::max(setup.scatter_point.size(), 1), sync); scatter_start = CudaDevicePtr(std::max(setup.scatter_start.size(), 1), sync); scatter_term = CudaDevicePtr(std::max(setup.scatter_term.size(), 1), sync); + scatter_phase = CudaDevicePtr(std::max(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(phase_count, sync); density = std::make_unique(stream, setup.max_atoms, setup.max_pairs, ModelDensityGPU::Bricks(setup.grid.nu, setup.grid.nv, setup.grid.nw)); mask = std::make_unique(stream, GridPoints(setup.mask_grid)); @@ -467,12 +501,24 @@ struct RefinementGPUEngineImpl { mask_spectrum = CudaDevicePtr(ComplexPoints(setup.mask_grid), sync); } + // the rotations in whole units (gemmi's are den times an integer matrix), the translations as given std::vector 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 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(setup.sym_ops.size()); + std::vector 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<<>>(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<<>>( spectrum, Norm(), setup.grid.nu, setup.grid.nv, setup.grid.nw, ops, static_cast(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(setup.scatter_point.size()); if (points > 0) { - scatter_kernel<<>>(coef, row_hkl, ops, - static_cast(setup.sym_ops.size()), setup.den, - scatter_point, scatter_start, scatter_term, - points, spectrum); + scatter_kernel<<>>(coef, phase, scatter_point, scatter_start, + scatter_term, scatter_phase, points, spectrum); cuda_err(cudaGetLastError()); } cufft_err(inverse.C2R(reinterpret_cast(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(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))