leonarski_fandClaude Opus 5.5
04655bd8bf
On macOS, where a Metal device may be used (not under --no-gpu), each merge's rotation stage - the sphere and the peaks' refinement - runs on the GPU: the same trilinear interpolation of the same transform in float, one threadgroup per orientation, every bin's sums of |F|^2, |F|^4 and |F|^2 (E^2 - 1) in a fixed order, and the correlation made of them in double on the host. The GPU reads its own copy of the transform in 2 x 2 x 2 bricks (1.3x faster than the grid's rows, as the eight corners then lie on fewer cache lines), made for the rotation stage and given back before the translation function; the seeds' orientations later on stay on the CPU. Where the GPU fails, the stage goes on on the CPU with a warning. The kernel is MolecularReplacementMetal.metal, carried in the binary as the Bragg integration's are. Elsewhere nothing changes. Measured on an Apple M6 (16 GB), alternating with rc176 + the CPU prefetch, two runs each: - 4nwv (the timing benchmark's moved model): rotation 78.9 -> 30.7 s, wall 344.5 -> 298.8 s; peak footprint 16.1-16.8 -> 18.5 GB. - 9s02: rotation 2.99 -> 0.28 s, wall 79.2 -> 76.4 s. p.hkl, the placed model and the maps are identical to the CPU's on both, and so is every line of the molecular-replacement log (RF Z, TF Z, R values). The RF itself agrees with the CPU's to 1e-4, with the same best orientation (MolecularReplacementMetal_MatchesCPU, [portable]). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_018My7tqt8a4LyKKjP3FhScT
93 lines
4.1 KiB
Metal
93 lines
4.1 KiB
Metal
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
// The rotation function's kernel (MolecularReplacementMetal.h), compiled when it is first used from this source,
|
|
// which CMakeLists.txt copies into the binary.
|
|
|
|
#include <metal_stdlib>
|
|
using namespace metal;
|
|
|
|
struct Params {
|
|
int nu, nv, bu, bv; // the grid, and its bricks along u and v
|
|
uint nops, bins;
|
|
float box, blur;
|
|
};
|
|
|
|
// MolecularTransform::At()'s corner, from the bricked grid.
|
|
static float2 transform_at(device const float2 *g, constant Params &p, int i, int j, int l) {
|
|
bool conj = false;
|
|
if (l < 0) { // the half-l grid holds l >= 0; F(-h) = conj F(h)
|
|
i = -i; j = -j; l = -l;
|
|
conj = true;
|
|
}
|
|
if (i < 0 || i >= p.nu) { i %= p.nu; if (i < 0) i += p.nu; }
|
|
if (j < 0 || j >= p.nv) { j %= p.nv; if (j < 0) j += p.nv; }
|
|
const uint brick = uint(i >> 1) + uint(p.bu) * (uint(j >> 1) + uint(p.bv) * uint(l >> 1));
|
|
float2 v = g[brick * 8 + uint(i & 1) + 2 * uint(j & 1) + 4 * uint(l & 1)];
|
|
if (conj)
|
|
v.y = -v.y;
|
|
return v;
|
|
}
|
|
|
|
// One threadgroup per orientation, as MolecularReplacementGPU.cu's rotation_function; the reflections in the order
|
|
// of their bins, every bin's sums taken in a fixed order: each thread's own, then the SIMD group's, then the groups'.
|
|
// For orientation r, bin b: out[(r bins + b) 3 + k] = the sum over the bin's reflections of ic, ic^2 and
|
|
// ic (E^2 - 1 - mean), ic being |F|^2 summed over the copies.
|
|
kernel void rotation_function(device const float2 *g [[buffer(0)]], constant Params &p [[buffer(1)]],
|
|
device const float3x3 *rot [[buffer(2)]], device const packed_float3 *s [[buffer(3)]],
|
|
device const float *w [[buffer(4)]], device const uint *bin_first [[buffer(5)]],
|
|
device float *out [[buffer(6)]],
|
|
uint r [[threadgroup_position_in_grid]], uint t [[thread_index_in_threadgroup]],
|
|
uint nt [[threads_per_threadgroup]], uint lane [[thread_index_in_simdgroup]],
|
|
uint sg [[simdgroup_index_in_threadgroup]], uint nsg [[simdgroups_per_threadgroup]]) {
|
|
threadgroup float sh[3][32];
|
|
const float3x3 m = rot[r]; // R^T: the model turned by R has F_R(s) = F_model(R^T s)
|
|
for (uint b = 0; b < p.bins; b++) {
|
|
float s1 = 0, s2 = 0, sw = 0;
|
|
for (uint i = bin_first[b] + t; i < bin_first[b + 1]; i += nt) {
|
|
float ic = 0;
|
|
for (uint j = 0; j < p.nops; j++) {
|
|
const float3 v = m * float3(s[i * p.nops + j]);
|
|
const float3 u = v * p.box;
|
|
const float3 fl = floor(u);
|
|
const float3 f = u - fl;
|
|
const int3 i0 = int3(fl);
|
|
float2 acc = 0;
|
|
for (int dx = 0; dx < 2; dx++)
|
|
for (int dy = 0; dy < 2; dy++)
|
|
for (int dz = 0; dz < 2; dz++) {
|
|
const float wt = (dx ? f.x : 1 - f.x) * (dy ? f.y : 1 - f.y) * (dz ? f.z : 1 - f.z);
|
|
acc += wt * transform_at(g, p, i0.x + dx, i0.y + dy, i0.z + dz);
|
|
}
|
|
// |F exp(blur s^2 / 4)|^2
|
|
ic += dot(acc, acc) * exp(p.blur * dot(v, v) / 2.0f);
|
|
}
|
|
s1 += ic;
|
|
s2 += ic * ic;
|
|
sw += ic * w[i];
|
|
}
|
|
s1 = simd_sum(s1);
|
|
s2 = simd_sum(s2);
|
|
sw = simd_sum(sw);
|
|
if (lane == 0) {
|
|
sh[0][sg] = s1;
|
|
sh[1][sg] = s2;
|
|
sh[2][sg] = sw;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
if (t == 0) {
|
|
float a = 0, c = 0, d = 0;
|
|
for (uint k = 0; k < nsg; k++) {
|
|
a += sh[0][k];
|
|
c += sh[1][k];
|
|
d += sh[2][k];
|
|
}
|
|
device float *o = out + (r * p.bins + b) * 3;
|
|
o[0] = a;
|
|
o[1] = c;
|
|
o[2] = d;
|
|
}
|
|
threadgroup_barrier(mem_flags::mem_threadgroup);
|
|
}
|
|
}
|