MolecularReplacement: the rotation function on the Apple GPU (Metal), equal to rounding
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
This commit is contained in:
Claude Opus 5.5
@@ -21,3 +21,23 @@ ADD_LIBRARY(JFJochStructureRefinement STATIC
|
||||
# JFJochIndexing for CUDAMemHelpers.h and, under CUDA, to put cuFFT (RigidBodyGPU.cu) after these
|
||||
# objects on the link line.
|
||||
TARGET_LINK_LIBRARIES(JFJochStructureRefinement JFJochScaleMerge JFJochIndexing JFJochCommon gemmi fftw3f)
|
||||
|
||||
IF (APPLE)
|
||||
# The rotation function on the Apple GPU (MolecularReplacementMetal.h). Its kernel is compiled when it is first
|
||||
# used, from source carried in the binary, as image_analysis/bragg_integration does: the .metal file copied into
|
||||
# a generated .cpp here; editing it re-runs the configure.
|
||||
FILE(READ ${CMAKE_CURRENT_SOURCE_DIR}/MolecularReplacementMetal.metal text)
|
||||
SET_PROPERTY(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/MolecularReplacementMetal.metal)
|
||||
# Written only when it changed, so a configure does not rebuild the library.
|
||||
FILE(WRITE ${CMAKE_CURRENT_BINARY_DIR}/MolecularReplacementMetalSource.cpp.tmp
|
||||
"// Generated from image_analysis/structure_refinement by its CMakeLists.txt - do not edit.\n"
|
||||
"extern const char *const mr_metal_source;\n"
|
||||
"const char *const mr_metal_source = R\"MR_METAL(\n${text})MR_METAL\";\n")
|
||||
CONFIGURE_FILE(${CMAKE_CURRENT_BINARY_DIR}/MolecularReplacementMetalSource.cpp.tmp
|
||||
${CMAKE_CURRENT_BINARY_DIR}/MolecularReplacementMetalSource.cpp COPYONLY)
|
||||
TARGET_SOURCES(JFJochStructureRefinement PRIVATE MolecularReplacementMetal.mm MolecularReplacementMetal.h
|
||||
MolecularReplacementMetal.metal ${CMAKE_CURRENT_BINARY_DIR}/MolecularReplacementMetalSource.cpp)
|
||||
SET_SOURCE_FILES_PROPERTIES(MolecularReplacementMetal.mm PROPERTIES COMPILE_OPTIONS -fobjc-arc)
|
||||
SET_SOURCE_FILES_PROPERTIES(MolecularReplacementMetal.metal PROPERTIES HEADER_FILE_ONLY ON)
|
||||
TARGET_LINK_LIBRARIES(JFJochStructureRefinement "-framework Metal" "-framework Foundation")
|
||||
ENDIF()
|
||||
|
||||
@@ -38,6 +38,9 @@
|
||||
#include "ModelValidation.h"
|
||||
#include "RigidBodyRefine.h"
|
||||
#include "MolecularReplacementGPU.h"
|
||||
#ifdef __APPLE__
|
||||
#include "MolecularReplacementMetal.h"
|
||||
#endif
|
||||
#ifdef JFJOCH_USE_CUDA
|
||||
#include "RigidBodyGPU.h"
|
||||
#include "RigidBodyGPUEngine.h"
|
||||
@@ -2616,7 +2619,36 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
|
||||
const SearchReflections lrefl = MakeSearchReflections(*merge.merged, cell, merge.space_group, d_llg,
|
||||
result.d_max_search);
|
||||
plan_gpu(m);
|
||||
const std::vector<double> rf = rotation_function(refl, cell, merge.space_group, rotations);
|
||||
#ifdef __APPLE__
|
||||
// On the Apple GPU where there is one, from its copy of the transform, made for this stage and given back
|
||||
// before the translation function: the seeds' few orientations later on are scored on the CPU. Where the
|
||||
// GPU fails, on the CPU from there on - a Mac has no other GPU to try.
|
||||
std::unique_ptr<MolecularReplacementMetal> metal;
|
||||
const auto metal_failed = [&](const JFJochException &e) {
|
||||
logger.Warning("Molecular replacement: the rotation function on the Apple GPU failed, on the CPU instead: {}",
|
||||
e.what());
|
||||
metal.reset();
|
||||
};
|
||||
try {
|
||||
if (MolecularReplacementMetal::Available(ft))
|
||||
metal = std::make_unique<MolecularReplacementMetal>(ft, nthreads);
|
||||
} catch (const JFJochException &e) {
|
||||
metal_failed(e);
|
||||
}
|
||||
#endif
|
||||
const auto rotation_stage = [&](const std::vector<gemmi::Mat33> &r) {
|
||||
#ifdef __APPLE__
|
||||
if (metal) {
|
||||
try {
|
||||
return metal->RotationFunction(refl, cell, merge.space_group, r);
|
||||
} catch (const JFJochException &e) {
|
||||
metal_failed(e);
|
||||
}
|
||||
}
|
||||
#endif
|
||||
return rotation_function(refl, cell, merge.space_group, r);
|
||||
};
|
||||
const std::vector<double> rf = rotation_stage(rotations);
|
||||
double rf_mean = 0, rf_sd = 0;
|
||||
for (double x : rf)
|
||||
rf_mean += x;
|
||||
@@ -2658,7 +2690,7 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
|
||||
for (int axis = 0; axis < 3; axis++)
|
||||
for (double sign : {1.0, -1.0})
|
||||
trial.push_back(axis_rotation(axis, sign * step).multiply(peak_rotations[k]));
|
||||
const std::vector<double> v = rotation_function(refl, cell, merge.space_group, trial);
|
||||
const std::vector<double> v = rotation_stage(trial);
|
||||
size_t next = 0;
|
||||
for (size_t k = 0; k < peaks.size(); k++) {
|
||||
if (!active[k])
|
||||
@@ -2677,10 +2709,17 @@ MolecularReplacementResult search_model(const SearchModel &model, const std::vec
|
||||
}
|
||||
}
|
||||
}
|
||||
#ifdef __APPLE__
|
||||
const bool on_metal = metal != nullptr;
|
||||
metal.reset();
|
||||
#else
|
||||
const bool on_metal = false;
|
||||
#endif
|
||||
result.t_rotation += seconds_since(t0);
|
||||
logger.Info("Molecular replacement: {} in {}: {} reflections {:.1f}-{:.1f} A, {} orientations, best RF Z {:.1f}",
|
||||
logger.Info("Molecular replacement: {} in {}: {} reflections {:.1f}-{:.1f} A, {} orientations{}, best RF Z {:.1f}",
|
||||
merge.name, merge.space_group.xhm(), refl.hkl.size(), result.d_max_search, result.d_min_search,
|
||||
rotations.size(), rf_sd > 0 ? (peak_rf.front() - rf_mean) / rf_sd : 0.0);
|
||||
rotations.size(), on_metal ? " (rotation function on the Apple GPU)" : "",
|
||||
rf_sd > 0 ? (peak_rf.front() - rf_mean) / rf_sd : 0.0);
|
||||
|
||||
// --- translation, for every group and every peak orientation ---
|
||||
t0 = std::chrono::steady_clock::now();
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#pragma once
|
||||
|
||||
// The rotation function of MolecularReplacement.h on an Apple GPU (macOS only): the same trilinear interpolation
|
||||
// of the same molecular transform, the same bins and the same correlation, in float on the device. It agrees with
|
||||
// the CPU to rounding, not bit for bit. The device sums every bin's |F|^2, |F|^4 and |F|^2 (E^2 - 1) for each
|
||||
// orientation in a fixed order, and the host makes the correlation of them in double.
|
||||
//
|
||||
// The device reads its own copy of the transform, in 2 x 2 x 2 bricks: the eight corners of an interpolation then
|
||||
// lie on fewer cache lines than in the grid's rows, which the GPU, waiting on memory as the CPU does, reads 1.3x
|
||||
// faster. The copy takes as much memory as the transform and is held as long as this object.
|
||||
|
||||
#include <memory>
|
||||
#include <vector>
|
||||
|
||||
#include "MolecularReplacement.h"
|
||||
|
||||
class MolecularReplacementMetal {
|
||||
public:
|
||||
// Whether the rotation function of `ft` can run on the GPU: a Metal device the process may use (rugnux
|
||||
// --no-gpu hides it, see hide_gpus) whose buffers hold the transform.
|
||||
static bool Available(const MolecularTransform &ft);
|
||||
|
||||
// Copies the transform into the device's layout. Throws JFJochException on a Metal failure, as every call
|
||||
// below does.
|
||||
MolecularReplacementMetal(const MolecularTransform &ft, size_t nthreads);
|
||||
~MolecularReplacementMetal();
|
||||
|
||||
// RotationFunction() of MolecularReplacement.h; may be called from several threads at once.
|
||||
std::vector<double> RotationFunction(const SearchReflections &refl, const gemmi::UnitCell &cell,
|
||||
const gemmi::SpaceGroup &sg, const std::vector<gemmi::Mat33> &rotations) const;
|
||||
|
||||
private:
|
||||
struct State;
|
||||
std::unique_ptr<State> state_;
|
||||
};
|
||||
@@ -0,0 +1,92 @@
|
||||
// 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);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,215 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#import <Metal/Metal.h>
|
||||
|
||||
#include "MolecularReplacementMetal.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <cstdint>
|
||||
#include <limits>
|
||||
#include <numeric>
|
||||
#include <string>
|
||||
|
||||
#include "../../common/CUDAWrapper.h" // gpus_are_hidden
|
||||
#include "../../common/JFJochException.h"
|
||||
#include "../../common/ParallelFor.h"
|
||||
|
||||
// The kernel's source, MolecularReplacementMetal.metal, copied in at configure time (CMakeLists.txt) so the
|
||||
// binary needs no file beside it.
|
||||
extern const char *const mr_metal_source;
|
||||
|
||||
namespace {
|
||||
|
||||
// The grid in bricks of 2 x 2 x 2 points, brick (i/2, j/2, l/2) holding point (i%2, j%2, l%2) at
|
||||
// (i%2) + 2 (j%2) + 4 (l%2); the bricks in the grid's order. The kernel's index is 32-bit.
|
||||
size_t BrickedPoints(const gemmi::FPhiGrid<float> &g) {
|
||||
return static_cast<size_t>((g.nu + 1) / 2) * ((g.nv + 1) / 2) * ((g.nw + 1) / 2) * 8;
|
||||
}
|
||||
|
||||
constexpr size_t THREADS = 256;
|
||||
// Orientations per command buffer: half a second of the GPU at most on the largest transforms.
|
||||
constexpr size_t BATCH = 256;
|
||||
|
||||
struct Params {
|
||||
int32_t nu, nv, bu, bv;
|
||||
uint32_t nops, bins;
|
||||
float box, blur;
|
||||
};
|
||||
|
||||
struct Kernel {
|
||||
id<MTLDevice> device = nil; // nil: no Metal device
|
||||
id<MTLComputePipelineState> pipeline = nil;
|
||||
};
|
||||
|
||||
// Compiled once per process, the first time it is asked for.
|
||||
const Kernel &GetKernel() {
|
||||
static const Kernel kernel = [] {
|
||||
Kernel k;
|
||||
k.device = MTLCreateSystemDefaultDevice();
|
||||
if (k.device == nil)
|
||||
return k;
|
||||
MTLCompileOptions *options = [MTLCompileOptions new];
|
||||
options.languageVersion = MTLLanguageVersion3_0;
|
||||
NSError *error = nil;
|
||||
id<MTLLibrary> library = [k.device newLibraryWithSource:@(mr_metal_source) options:options error:&error];
|
||||
if (library == nil)
|
||||
throw JFJochException(JFJochExceptionCategory::GPUMetalError,
|
||||
std::string("rotation function kernel: ") + error.localizedDescription.UTF8String);
|
||||
id<MTLFunction> function = [library newFunctionWithName:@"rotation_function"];
|
||||
k.pipeline = [k.device newComputePipelineStateWithFunction:function error:&error];
|
||||
if (k.pipeline == nil)
|
||||
throw JFJochException(JFJochExceptionCategory::GPUMetalError,
|
||||
std::string("rotation_function: ") + error.localizedDescription.UTF8String);
|
||||
return k;
|
||||
}();
|
||||
return kernel;
|
||||
}
|
||||
|
||||
id<MTLBuffer> Upload(id<MTLDevice> device, const void *data, size_t bytes) {
|
||||
id<MTLBuffer> b = [device newBufferWithBytes:data length:std::max<size_t>(bytes, 4)
|
||||
options:MTLResourceStorageModeShared];
|
||||
if (b == nil)
|
||||
throw JFJochException(JFJochExceptionCategory::GPUMetalError, "rotation function: out of GPU memory");
|
||||
return b;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
struct MolecularReplacementMetal::State {
|
||||
id<MTLCommandQueue> queue = nil;
|
||||
id<MTLBuffer> grid = nil;
|
||||
Params params{};
|
||||
};
|
||||
|
||||
bool MolecularReplacementMetal::Available(const MolecularTransform &ft) {
|
||||
if (gpus_are_hidden())
|
||||
return false;
|
||||
const Kernel &k = GetKernel();
|
||||
const size_t points = BrickedPoints(ft.Grid());
|
||||
return k.device != nil && points <= std::numeric_limits<uint32_t>::max()
|
||||
&& points * sizeof(std::complex<float>) <= k.device.maxBufferLength;
|
||||
}
|
||||
|
||||
MolecularReplacementMetal::MolecularReplacementMetal(const MolecularTransform &ft, size_t nthreads)
|
||||
: state_(std::make_unique<State>()) {
|
||||
const Kernel &k = GetKernel();
|
||||
const gemmi::FPhiGrid<float> &g = ft.Grid();
|
||||
const int bu = (g.nu + 1) / 2, bv = (g.nv + 1) / 2, bw = (g.nw + 1) / 2;
|
||||
state_->params = {g.nu, g.nv, bu, bv, 0, 0, static_cast<float>(ft.Box()), static_cast<float>(ft.Blur())};
|
||||
state_->queue = [k.device newCommandQueue];
|
||||
// Points past the grid's edge, in its last bricks, are never read.
|
||||
state_->grid = [k.device newBufferWithLength:BrickedPoints(g) * sizeof(std::complex<float>)
|
||||
options:MTLResourceStorageModeShared];
|
||||
if (state_->queue == nil || state_->grid == nil)
|
||||
throw JFJochException(JFJochExceptionCategory::GPUMetalError, "rotation function: out of GPU memory");
|
||||
auto *out = static_cast<std::complex<float> *>(state_->grid.contents);
|
||||
ParallelFor(bw, nthreads, [&](int bl) {
|
||||
for (int l = 2 * bl; l < std::min(2 * bl + 2, g.nw); l++)
|
||||
for (int j = 0; j < g.nv; j++)
|
||||
for (int i = 0; i < g.nu; i++) {
|
||||
const size_t brick = (i >> 1) + static_cast<size_t>(bu) * ((j >> 1) + static_cast<size_t>(bv) * (l >> 1));
|
||||
out[brick * 8 + (i & 1) + 2 * (j & 1) + 4 * (l & 1)] =
|
||||
g.data[i + static_cast<size_t>(g.nu) * (j + static_cast<size_t>(g.nv) * l)];
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
MolecularReplacementMetal::~MolecularReplacementMetal() = default;
|
||||
|
||||
std::vector<double> MolecularReplacementMetal::RotationFunction(const SearchReflections &refl,
|
||||
const gemmi::UnitCell &cell,
|
||||
const gemmi::SpaceGroup &sg,
|
||||
const std::vector<gemmi::Mat33> &rotations) const {
|
||||
const Kernel &k = GetKernel();
|
||||
const size_t n = refl.hkl.size();
|
||||
const std::vector<gemmi::Vec3> s = CopyVectors(refl, cell, sg);
|
||||
const size_t nops = s.size() / n;
|
||||
double w_mean = 0;
|
||||
for (double e2 : refl.e2)
|
||||
w_mean += e2;
|
||||
w_mean /= static_cast<double>(n);
|
||||
double w_ss = 0;
|
||||
for (double e2 : refl.e2)
|
||||
w_ss += (e2 - w_mean) * (e2 - w_mean);
|
||||
|
||||
// The reflections in the order of their bins, each bin's from bin_first[b] to bin_first[b + 1].
|
||||
std::vector<size_t> order(n);
|
||||
std::iota(order.begin(), order.end(), 0);
|
||||
std::stable_sort(order.begin(), order.end(), [&](size_t a, size_t b) { return refl.bin[a] < refl.bin[b]; });
|
||||
std::vector<uint32_t> bin_first(refl.bins + 1, 0);
|
||||
for (size_t i = 0; i < n; i++)
|
||||
bin_first[refl.bin[i] + 1]++;
|
||||
for (int b = 0; b < refl.bins; b++)
|
||||
bin_first[b + 1] += bin_first[b];
|
||||
std::vector<float> sf(3 * n * nops), w(n);
|
||||
for (size_t i = 0; i < n; i++) {
|
||||
for (size_t j = 0; j < nops; j++) {
|
||||
const gemmi::Vec3 &v = s[order[i] * nops + j];
|
||||
sf[3 * (i * nops + j)] = static_cast<float>(v.x);
|
||||
sf[3 * (i * nops + j) + 1] = static_cast<float>(v.y);
|
||||
sf[3 * (i * nops + j) + 2] = static_cast<float>(v.z);
|
||||
}
|
||||
w[i] = static_cast<float>(refl.e2[order[i]] - w_mean);
|
||||
}
|
||||
// R^T as Metal's float3x3: three columns of four floats.
|
||||
std::vector<float> rot(12 * rotations.size(), 0.0f);
|
||||
for (size_t r = 0; r < rotations.size(); r++)
|
||||
for (int col = 0; col < 3; col++)
|
||||
for (int row = 0; row < 3; row++)
|
||||
rot[12 * r + 4 * col + row] = static_cast<float>(rotations[r].a[col][row]);
|
||||
|
||||
Params params = state_->params;
|
||||
params.nops = static_cast<uint32_t>(nops);
|
||||
params.bins = static_cast<uint32_t>(refl.bins);
|
||||
id<MTLBuffer> d_s = Upload(k.device, sf.data(), sf.size() * sizeof(float));
|
||||
id<MTLBuffer> d_w = Upload(k.device, w.data(), w.size() * sizeof(float));
|
||||
id<MTLBuffer> d_bin_first = Upload(k.device, bin_first.data(), bin_first.size() * sizeof(uint32_t));
|
||||
id<MTLBuffer> d_rot = Upload(k.device, rot.data(), rot.size() * sizeof(float));
|
||||
id<MTLBuffer> d_out = [k.device newBufferWithLength:std::max<size_t>(rotations.size() * refl.bins * 3 * sizeof(float), 4)
|
||||
options:MTLResourceStorageModeShared];
|
||||
if (d_out == nil)
|
||||
throw JFJochException(JFJochExceptionCategory::GPUMetalError, "rotation function: out of GPU memory");
|
||||
std::vector<id<MTLCommandBuffer>> running;
|
||||
for (size_t first = 0; first < rotations.size(); first += BATCH) {
|
||||
const size_t count = std::min(BATCH, rotations.size() - first);
|
||||
id<MTLCommandBuffer> cb = [state_->queue commandBuffer];
|
||||
id<MTLComputeCommandEncoder> e = [cb computeCommandEncoder];
|
||||
[e setComputePipelineState:k.pipeline];
|
||||
[e setBuffer:state_->grid offset:0 atIndex:0];
|
||||
[e setBytes:¶ms length:sizeof(params) atIndex:1];
|
||||
[e setBuffer:d_rot offset:first * 12 * sizeof(float) atIndex:2];
|
||||
[e setBuffer:d_s offset:0 atIndex:3];
|
||||
[e setBuffer:d_w offset:0 atIndex:4];
|
||||
[e setBuffer:d_bin_first offset:0 atIndex:5];
|
||||
[e setBuffer:d_out offset:first * refl.bins * 3 * sizeof(float) atIndex:6];
|
||||
[e dispatchThreadgroups:MTLSizeMake(count, 1, 1) threadsPerThreadgroup:MTLSizeMake(THREADS, 1, 1)];
|
||||
[e endEncoding];
|
||||
[cb commit];
|
||||
running.push_back(cb);
|
||||
}
|
||||
for (id<MTLCommandBuffer> cb : running) {
|
||||
[cb waitUntilCompleted];
|
||||
if (cb.status != MTLCommandBufferStatusCompleted)
|
||||
throw JFJochException(JFJochExceptionCategory::GPUMetalError,
|
||||
cb.error ? cb.error.localizedDescription.UTF8String : "command buffer failed");
|
||||
}
|
||||
|
||||
// The correlation of E^2 - 1 with ic over its bin's mean, from the bins' sums: x = ic / mean_b.
|
||||
const auto *sums = static_cast<const float *>(d_out.contents);
|
||||
std::vector<double> out(rotations.size());
|
||||
for (size_t r = 0; r < rotations.size(); r++) {
|
||||
double x_sum = 0, x_ss = 0, xy = 0;
|
||||
for (int b = 0; b < refl.bins; b++) {
|
||||
const float *o = sums + (r * refl.bins + b) * 3;
|
||||
const double mean = o[0] / static_cast<double>(bin_first[b + 1] - bin_first[b]);
|
||||
x_sum += o[0] / mean;
|
||||
x_ss += o[1] / (mean * mean);
|
||||
xy += o[2] / mean;
|
||||
}
|
||||
const double sxx = x_ss - x_sum * x_sum / static_cast<double>(n);
|
||||
out[r] = sxx > 0 && w_ss > 0 ? xy / std::sqrt(sxx * w_ss) : 0.0;
|
||||
}
|
||||
return out;
|
||||
}
|
||||
@@ -21,6 +21,7 @@ ADD_EXECUTABLE(jfjoch_portable_test EXCLUDE_FROM_ALL
|
||||
BraggIntegrationEngineCompressedImageTest.cpp
|
||||
BraggIntegrationEngineCPUTest.cpp
|
||||
BraggIntegrationEngineMetalTest.cpp
|
||||
MolecularReplacementMetalTest.cpp
|
||||
WriteReflectionsTest.cpp
|
||||
LoadReferenceMtzFreeFlagsTest.cpp
|
||||
LoadReferenceSfCifTest.cpp
|
||||
|
||||
@@ -0,0 +1,121 @@
|
||||
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
||||
// SPDX-License-Identifier: GPL-3.0-only
|
||||
|
||||
#include <catch2/catch_all.hpp>
|
||||
|
||||
#ifdef __APPLE__
|
||||
|
||||
#include <algorithm>
|
||||
#include <cmath>
|
||||
#include <random>
|
||||
#include <vector>
|
||||
|
||||
#include <gemmi/dencalc.hpp>
|
||||
#include <gemmi/it92.hpp>
|
||||
|
||||
#include "../image_analysis/structure_refinement/ModelFFT.h"
|
||||
#include "../image_analysis/structure_refinement/MolecularReplacement.h"
|
||||
#include "../image_analysis/structure_refinement/MolecularReplacementMetal.h"
|
||||
|
||||
// The rotation function on the Apple GPU is the CPU's in float: equal to rounding, with the same best orientation.
|
||||
// A Mac without a Metal device fails rather than skips - the macOS CI job is the only place this can be checked,
|
||||
// as in BraggIntegrationEngineMetalTest.cpp. The search as in MolecularReplacementTest.cpp.
|
||||
|
||||
namespace {
|
||||
// A synthetic body: carbon atoms scattered in a 12 A sphere, centroid at the origin.
|
||||
SearchModel synthetic_model() {
|
||||
SearchModel m;
|
||||
std::mt19937 rng(7);
|
||||
std::uniform_real_distribution<double> u(-12.0, 12.0);
|
||||
gemmi::Model model(1);
|
||||
gemmi::Chain chain("A");
|
||||
gemmi::Residue res;
|
||||
res.name = "UNK";
|
||||
res.seqid = gemmi::SeqId(1, ' ');
|
||||
std::vector<gemmi::Vec3> pos;
|
||||
while (pos.size() < 300) {
|
||||
const gemmi::Vec3 p(u(rng), u(rng), u(rng));
|
||||
if (p.length() < 12.0)
|
||||
pos.push_back(p);
|
||||
}
|
||||
gemmi::Vec3 c;
|
||||
for (const auto &p : pos)
|
||||
c += p;
|
||||
c *= 1.0 / pos.size();
|
||||
for (const auto &p : pos) {
|
||||
gemmi::Atom a;
|
||||
a.name = "CA"; // every atom counts in the clash check
|
||||
a.element = gemmi::El::C;
|
||||
a.pos = gemmi::Position(p - c);
|
||||
a.b_iso = 15.0f;
|
||||
a.occ = 1.0f;
|
||||
res.atoms.push_back(a);
|
||||
m.radius_A = std::max(m.radius_A, (p - c).length());
|
||||
}
|
||||
chain.residues.push_back(res);
|
||||
model.chains.push_back(chain);
|
||||
m.st.models.push_back(model);
|
||||
m.residues_read = m.residues_kept = 1;
|
||||
m.mass_da = 3600;
|
||||
return m;
|
||||
}
|
||||
|
||||
gemmi::Mat33 axis_angle(gemmi::Vec3 axis, double deg) {
|
||||
axis = axis.normalized();
|
||||
const double k = deg * M_PI / 180.0, s = std::sin(k), c1 = 1.0 - std::cos(k);
|
||||
const double x = axis.x, y = axis.y, z = axis.z;
|
||||
return {1 + c1 * (x * x - 1), -s * z + c1 * x * y, s * y + c1 * x * z,
|
||||
s * z + c1 * x * y, 1 + c1 * (y * y - 1), -s * x + c1 * y * z,
|
||||
-s * y + c1 * x * z, s * x + c1 * y * z, 1 + c1 * (z * z - 1)};
|
||||
}
|
||||
|
||||
// |F|^2 of the placed crystal to 4 A, as merged reflections.
|
||||
std::vector<MergedReflection> observed(const gemmi::Structure &crystal) {
|
||||
gemmi::DensityCalculator<gemmi::IT92<float>, float> dc;
|
||||
dc.d_min = 3.5;
|
||||
dc.rate = 1.5;
|
||||
dc.set_grid_cell_and_spacegroup(crystal);
|
||||
dc.put_model_density_on_grid(crystal.models[0]);
|
||||
const auto f = MapToFPhi(dc.grid).prepare_asu_data(3.5, 0, false, false, false);
|
||||
std::vector<MergedReflection> out;
|
||||
for (const auto &hv : f.v) {
|
||||
MergedReflection r;
|
||||
r.h = hv.hkl[0];
|
||||
r.k = hv.hkl[1];
|
||||
r.l = hv.hkl[2];
|
||||
r.I = std::norm(hv.value);
|
||||
r.sigma = 1;
|
||||
r.F = std::abs(hv.value);
|
||||
r.d = 1.0 / std::sqrt(crystal.cell.calculate_1_d2(hv.hkl));
|
||||
out.push_back(r);
|
||||
}
|
||||
return out;
|
||||
}
|
||||
}
|
||||
|
||||
TEST_CASE("MolecularReplacementMetal_MatchesCPU", "[portable][MolecularReplacement][Metal]") {
|
||||
const SearchModel model = synthetic_model();
|
||||
const gemmi::UnitCell cell(42, 50, 58, 90, 90, 90);
|
||||
const gemmi::SpaceGroup &sg = gemmi::get_spacegroup_by_name("P 21 21 21");
|
||||
const gemmi::Mat33 truth = axis_angle({1, 2, 3}, 50);
|
||||
const auto merged = observed(PlaceSearchModel(model, truth, gemmi::Fractional(0.17, 0.31, 0.09), cell, sg));
|
||||
const SearchReflections refl = MakeSearchReflections(merged, cell, sg, 4.0, 12.0);
|
||||
const MolecularTransform ft(model, 4.0, 4);
|
||||
REQUIRE(MolecularReplacementMetal::Available(ft));
|
||||
const MolecularReplacementMetal metal(ft, 4);
|
||||
|
||||
const auto pg = CartesianPointGroup(sg, cell);
|
||||
const auto rotations = RotationAsu(UniformRotations(6.0), pg);
|
||||
const auto rf_cpu = RotationFunction(ft, refl, cell, sg, rotations, 4);
|
||||
const auto rf_metal = metal.RotationFunction(refl, cell, sg, rotations);
|
||||
REQUIRE(rf_metal.size() == rf_cpu.size());
|
||||
double worst = 0;
|
||||
for (size_t i = 0; i < rf_cpu.size(); i++)
|
||||
worst = std::max(worst, std::fabs(rf_metal[i] - rf_cpu[i]));
|
||||
CHECK(worst < 1e-4);
|
||||
const size_t best = std::max_element(rf_metal.begin(), rf_metal.end()) - rf_metal.begin();
|
||||
CHECK(best == static_cast<size_t>(std::max_element(rf_cpu.begin(), rf_cpu.end()) - rf_cpu.begin()));
|
||||
CHECK(RotationDistanceDeg(rotations[best], truth, pg) < 6.0);
|
||||
}
|
||||
|
||||
#endif
|
||||
Reference in New Issue
Block a user