Files
Jungfraujoch/image_analysis/indexing/FFBIDXIndexerCPU.cpp
T
leonarski_f ea383f9e83 FFBIDXIndexerCPU: the fast-feedback indexer on the CPU, used where there is no GPU
A port of the fast-feedback indexer's CUDA kernels to plain C++/Eigen: the same half-sphere and
rotation samples, the same scores and the same 32 output cells for Refine(). The loops are
reorganised so the inner one runs over samples and vectorises (GCC/Clang, no intrinsics), and the
cell search separates the rotation about the candidate vector out of the per-spot terms. The
vector refinement reproduces the library's warp sums - partial sums its shuffle reduction leaves in
the other lanes carry into the next round - and its float arithmetic (FMA contraction, the
Fibonacci points in float), because the candidates searched depend on it: on the library's 27 test
spot lists and 4 synthetic stills the CPU and GPU best cells agree to ~1e-6 relative, the top-32
sets in 30 of 31 cases completely.

FFBIDX now resolves to the GPU library where a GPU is visible and to the CPU port otherwise (no
GPU, --no-gpu, or a build without CUDA); Auto takes FFBIDX for stills whenever a cell is known,
also without a GPU, instead of FFTW. With a GPU nothing changes (a stills set and a rotation set
reproduce the baseline exactly). Without a GPU the stills indexer pool gets one indexer per worker
instead of eight, since the CPU indexer holds no card memory.

Single core, synthetic still, loaded 5950X: ~8 ms per call at 30 spots, ~15 ms at 80, ~35 ms at
250; 16 threads ~1300/700/330 calls/s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_015CeEGgpKdfkCTJb3YNMJcu
(cherry picked from commit 9a8eb8669d1be6c3ec16db6786280affd6830ee5)
2026-10-23 17:56:50 +02:00

432 lines
18 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
// Ported from the CUDA kernels of the fast-feedback indexer,
// https://github.com/paulscherrerinstitute/fast-feedback-indexer (indexer/src/indexer_gpu.cu)
// (c) 2022 Paul Scherrer Institute, BSD-3-Clause
//
// The samples and the scores are the library's; the loops are reorganised for a CPU. Each block of
// samples is scored against one spot at a time, so the inner loop runs over samples and vectorises,
// and in the cell search the rotation about the candidate vector is separated out of the spot terms.
#include "FFBIDXIndexerCPU.h"
#include <algorithm>
#include <array>
#include <bit>
#include <cmath>
#include <cstdint>
#include <limits>
#include <numbers>
#include <numeric>
#include "PostIndexingRefinement.h"
namespace {
// The GPU library's runtime defaults (config_runtime in ffbidx/indexer.h) and rugnux's persistent
// configuration of it (FFBIDXIndexer.h). Keep them in step: the CPU vs GPU test compares the two.
constexpr int NUM_HALFSPHERE_POINTS = 32 * 1024;
constexpr int NUM_ANGLE_POINTS = 1024; // the library's own choice for 32k half-sphere points
constexpr int NUM_CANDIDATE_VECTORS = 32;
constexpr int NUM_OUTPUT_CELLS = 32;
constexpr int MIN_SPOTS = 6; // vector refinement
constexpr float TRIML = 0.001f;
constexpr float TRIMH = 0.3f;
constexpr float DELTA = 0.1f;
constexpr float DIST1 = 0.2f;
constexpr float DIST3 = 0.15f;
constexpr float LENGTH_THRESHOLD = 1e-9f;
// Half-sphere directions are scored in blocks of this many, so a block's accumulators stay in L1
constexpr int BLOCK = 1024;
static_assert(NUM_HALFSPHERE_POINTS % BLOCK == 0);
// Every factor of a score's product is in [TRIML + DELTA, TRIMH + DELTA]; after 8 of them the
// product is still a normal float
constexpr int RENORMALISE_EVERY = 8;
struct Tables {
std::array<std::vector<float>, 3> dir; // Fibonacci points on the half sphere z >= 0
std::vector<float> cos_alpha, sin_alpha; // rotation angles about a cell vector
};
// (a, rest) = a + x * y + rest, a Kahan sum, with the product fused into the first addition as
// nvcc compiles the library's
inline void KahanAdd(float &a, float &rest, float x, float y = 1.0f) {
const float s = std::fma(x, y, rest);
const float t = a;
a = t + s;
rest = s - (a - t);
}
Tables MakeTables() {
constexpr double pi2 = 2.0 * std::numbers::pi;
const double dl = 3.0 - std::sqrt(5.0);
Tables t;
for (auto &v: t.dir)
v.resize(NUM_HALFSPHERE_POINTS);
std::array<float, 32> dl2n;
for (int n = 0; n < 32; n++)
dl2n[n] = static_cast<float>(std::fmod(dl * std::ldexp(1.0, n), pi2));
for (int i = 0; i < NUM_HALFSPHERE_POINTS; i++) {
const float z = 1.0f - static_cast<float>(i) / static_cast<float>(NUM_HALFSPHERE_POINTS - 1);
const float r = std::sqrt(std::fma(-z, z, 1.0f));
float l = 0, rest = 0;
for (int n = 0; n < 32; n++)
if (i & (1 << n))
KahanAdd(l, rest, dl2n[n]);
t.dir[0][i] = std::cos(l) * r;
t.dir[1][i] = std::sin(l) * r;
t.dir[2][i] = z;
}
for (int q = 0; q < NUM_ANGLE_POINTS; q++) {
const double alpha = q * pi2 / NUM_ANGLE_POINTS;
t.cos_alpha.push_back(static_cast<float>(std::cos(alpha)));
t.sin_alpha.push_back(static_cast<float>(std::sin(alpha)));
}
return t;
}
const Tables &GetTables() {
static const Tables t = MakeTables();
return t;
}
inline float Dist2Int(float x) {
return std::abs(x - std::rint(x));
}
// One spot's share of a sample's score: a factor of the product, and whether it counts as good
inline void AddSpot(float d, float good_limit, float &product, int &good) {
good += static_cast<int>(d < good_limit);
product *= std::clamp(d, TRIML, TRIMH) + DELTA;
}
// Moves the power of two of the product into the exponent, leaving the product in [1, 2). The
// library takes a log2 every 8 spots instead; this is the same sum, and it vectorises.
inline void Renormalise(float &product, int &exponent) {
const auto bits = std::bit_cast<uint32_t>(product);
exponent += static_cast<int>(bits >> 23) - 127;
product = std::bit_cast<float>((bits & 0x007fffffu) | 0x3f800000u);
}
// exp2(mean log2(trim(d) + delta)) - delta - (good + 1): lower is better. A NaN, from a degenerate
// candidate, goes last instead of breaking the sort.
inline float Score(float log2_sum, int good, int nspots) {
const float s = std::exp2(log2_sum / static_cast<float>(nspots)) - DELTA - static_cast<float>(good + 1);
return std::isnan(s) ? std::numeric_limits<float>::infinity() : s;
}
// Indices of the k lowest scores, best first
std::vector<int> Best(const std::vector<float> &score, int k) {
std::vector<int> idx(score.size());
std::iota(idx.begin(), idx.end(), 0);
auto lower = [&score](int a, int b) { return score[a] < score[b] || (score[a] == score[b] && a < b); };
std::nth_element(idx.begin(), idx.begin() + k, idx.end(), lower);
idx.resize(k);
std::sort(idx.begin(), idx.end(), lower);
return idx;
}
// Stage 1: every half-sphere direction as the direction of a cell vector of this length, scored on
// how close the spots' projections on it come to integers
std::vector<int> FindCandidates(const Tables &t, const std::vector<Eigen::Vector3f> &spots, float length) {
const int nspots = static_cast<int>(spots.size());
std::vector<float> score(NUM_HALFSPHERE_POINTS);
std::array<float, BLOCK> product;
std::array<int, BLOCK> exponent, good;
for (int b = 0; b < NUM_HALFSPHERE_POINTS; b += BLOCK) {
const float *vx = t.dir[0].data() + b;
const float *vy = t.dir[1].data() + b;
const float *vz = t.dir[2].data() + b;
product.fill(1.0f);
exponent.fill(0);
good.fill(0);
for (int i = 0; i < nspots; i++) {
const Eigen::Vector3f s = length * spots[i];
for (int k = 0; k < BLOCK; k++)
AddSpot(Dist2Int(s.x() * vx[k] + s.y() * vy[k] + s.z() * vz[k]), DIST1, product[k], good[k]);
if ((i + 1) % RENORMALISE_EVERY == 0)
for (int k = 0; k < BLOCK; k++)
Renormalise(product[k], exponent[k]);
}
for (int k = 0; k < BLOCK; k++)
score[b + k] = Score(static_cast<float>(exponent[k]) + std::log2(product[k]), good[k], nspots);
}
return Best(score, NUM_CANDIDATE_VECTORS);
}
// The library's dot product: rounding matters here, as a spot on the edge of a tolerance decides
// which spots the candidate is refined against
inline float Dot(const Eigen::Vector3f &a, const Eigen::Vector3f &b) {
return std::fma(a[0], b[0], std::fma(a[1], b[1], a[2] * b[2]));
}
// A sum as the library forms it on a warp of 32 lanes: spot i is added on lane i % 32, and the
// lanes are reduced by shuffles that leave partial sums behind in every lane but the first. Those
// stay in the lanes and enter the next round's sum. Reproduced as it is, because the refined
// candidates - and so the cells searched - depend on it.
struct WarpSum {
std::array<float, 32> value{}, rest{};
void Add(size_t spot, float x, float y) {
KahanAdd(value[spot % 32], rest[spot % 32], x, y);
}
float Reduce() {
std::array<float, 32> r{};
for (int d = 1; d < 32; d *= 2) {
const auto old_value = value;
const auto old_r = r;
for (int l = 0; l < 32; l++) {
const int from = (l % (2 * d) < d) ? l + d : l;
r[l] = old_r[from];
KahanAdd(value[l], r[l], old_value[from]);
}
}
return value[0];
}
};
// x = A^-1 v for symmetric A given by its upper triangle (00 01 02 11 12 22), by cofactors, as
// the library solves it (products fused as nvcc compiles them)
Eigen::Vector3f SymSolve(const std::array<float, 6> &A, const std::array<float, 3> &v) {
auto cofactor = [](float a, float b, float c, float d) { return std::fma(a, b, -(c * d)); };
const float co00 = cofactor(A[3], A[5], A[4], A[4]);
const float co01 = cofactor(A[1], A[5], A[2], A[4]);
const float co02 = cofactor(A[1], A[4], A[2], A[3]);
const float co11 = cofactor(A[0], A[5], A[2], A[2]);
const float co12 = cofactor(A[0], A[4], A[1], A[2]);
const float co22 = cofactor(A[0], A[3], A[1], A[1]);
// a0 * b0 + a1 * b1 + a2 * b2
auto sum3 = [](float a0, float b0, float a1, float b1, float a2, float b2) {
float t = a0 * b0;
float rest = 0;
KahanAdd(t, rest, a1, b1);
KahanAdd(t, rest, a2, b2);
return t;
};
const float r_det = 1.0f / sum3(A[0], co00, -A[1], co01, A[2], co02);
return {sum3(co00, v[0], -co01, v[1], co02, v[2]) * r_det,
sum3(-co01, v[0], co11, v[1], -co12, v[2]) * r_det,
sum3(co02, v[0], -co12, v[1], co22, v[2]) * r_det};
}
// Stage 2: a candidate refined by least squares against the spots it nearly indexes, in four rounds
// of tightening tolerance; the Miller indices stay those of the unrefined vector
Eigen::Vector3f RefineCandidate(const std::vector<Eigen::Vector3f> &spots, const Eigen::Vector3f &z0) {
std::array<WarpSum, 6> STS;
std::array<WarpSum, 3> STm;
Eigen::Vector3f z = z0;
for (const float beta: {0.1f, 0.05f, 0.025f, 0.01f}) {
int good = 0;
for (size_t i = 0; i < spots.size(); i++) {
const Eigen::Vector3f &s = spots[i];
const float m = std::rint(Dot(s, z0));
if (std::abs(Dot(s, z) - m) >= beta)
continue;
good++;
int e = 0;
for (int j = 0; j < 3; j++) {
STm[j].Add(i, s[j], m);
for (int k = j; k < 3; k++)
STS[e++].Add(i, s[j], s[k]);
}
}
if (good < MIN_SPOTS)
break;
std::array<float, 6> A;
std::array<float, 3> v;
for (int e = 0; e < 3; e++)
v[e] = STm[e].Reduce();
for (int e = 0; e < 6; e++)
A[e] = STS[e].Reduce();
z = SymSolve(A, v);
}
return z;
}
// The input cell turned so its vector cell_vec lies along c, in the frame (x, y, z = c/|c|) the
// rotation about c is sampled in: a(alpha) = a_z z + a_xy (cos alpha x + sin alpha y), and b the
// same at alpha + delta.
struct CellFrame {
int cell_vec;
Eigen::Vector3f x, y, z;
float c_length;
float a_z, a_xy, b_z, b_xy;
float cos_delta, sin_delta;
Eigen::Vector3f A(float cos_a, float sin_a) const {
return a_z * z + a_xy * (cos_a * x + sin_a * y);
}
Eigen::Vector3f B(float cos_a, float sin_a) const {
const float c = cos_a * cos_delta - sin_a * sin_delta;
const float s = sin_a * cos_delta + cos_a * sin_delta;
return b_z * z + b_xy * (c * x + s * y);
}
};
CellFrame MakeFrame(const Eigen::Matrix3f &cell, int cell_vec, const Eigen::Vector3f &c) {
CellFrame f;
f.cell_vec = cell_vec;
f.c_length = c.norm();
f.z = c / f.c_length;
// Mirroring on the bisector of the input vector and z takes one onto the other
const Eigen::Vector3f v = cell.row(cell_vec);
Eigen::Vector3f a = cell.row((cell_vec + 1) % 3);
Eigen::Vector3f b = cell.row((cell_vec + 2) % 3);
Eigen::Vector3f t = v + v.norm() * f.z;
if (t.norm() >= std::numeric_limits<float>::epsilon())
t.normalize();
else
t = f.z.cross(a).normalized();
a = 2.0f * a.dot(t) * t - a;
b = 2.0f * b.dot(t) * t - b;
f.a_z = a.dot(f.z);
f.x = (a - f.a_z * f.z).normalized();
f.a_xy = a.dot(f.x);
f.y = f.z.cross(f.x);
f.b_z = b.dot(f.z);
const Eigen::Vector3f b_dir = (b - f.b_z * f.z).normalized();
f.b_xy = b.dot(b_dir);
f.cos_delta = b_dir.dot(f.x);
f.sin_delta = b_dir.dot(f.y);
return f;
}
// Stage 3: the cell of one frame at every rotation angle, scored on how close each spot comes to
// a lattice point. The spot's projections on z, x and y are taken once; c does not turn at all.
void ScoreAngles(const Tables &t, const CellFrame &f, const std::vector<Eigen::Vector3f> &spots, float *score) {
const int nspots = static_cast<int>(spots.size());
const float *cos_a = t.cos_alpha.data();
const float *sin_a = t.sin_alpha.data();
std::array<float, NUM_ANGLE_POINTS> product;
std::array<int, NUM_ANGLE_POINTS> exponent, good;
product.fill(1.0f);
exponent.fill(0);
good.fill(0);
for (int i = 0; i < nspots; i++) {
const float zs = f.z.dot(spots[i]);
const float xs = f.x.dot(spots[i]);
const float ys = f.y.dot(spots[i]);
const float dc = Dist2Int(f.c_length * zs);
const float a0 = f.a_z * zs;
const float a1 = f.a_xy * xs;
const float a2 = f.a_xy * ys;
const float b0 = f.b_z * zs;
const float b1 = f.b_xy * (f.cos_delta * xs + f.sin_delta * ys);
const float b2 = f.b_xy * (f.cos_delta * ys - f.sin_delta * xs);
for (int q = 0; q < NUM_ANGLE_POINTS; q++) {
const float da = Dist2Int(a0 + a1 * cos_a[q] + a2 * sin_a[q]);
const float db = Dist2Int(b0 + b1 * cos_a[q] + b2 * sin_a[q]);
AddSpot(std::sqrt(dc * dc + da * da + db * db), DIST3, product[q], good[q]);
}
if ((i + 1) % RENORMALISE_EVERY == 0)
for (int q = 0; q < NUM_ANGLE_POINTS; q++)
Renormalise(product[q], exponent[q]);
}
// The library starts this sum at 1, not 0; kept, so the scores agree with the GPU
for (int q = 0; q < NUM_ANGLE_POINTS; q++)
score[q] = Score(1.0f + static_cast<float>(exponent[q]) + std::log2(product[q]), good[q], nspots);
}
}
// Following the fast-feedback indexer (H.-C. Stadler, PSI) and TORO, Gasparotto et al. (2024)
// J. Appl. Cryst. 57, 931-944
FFBIDXCells FFBIDXFindCells(const Eigen::Matrix3f &cell, const std::vector<Coord> &coord, size_t nspots) {
const Tables &t = GetTables();
std::vector<Eigen::Vector3f> spots(nspots);
for (size_t i = 0; i < nspots; i++)
spots[i] = Eigen::Vector3f(coord[i].x, coord[i].y, coord[i].z);
// Cell vectors of the same length share one set of candidates
std::array<int, 3> group{};
std::vector<float> group_length;
for (int v = 0; v < 3; v++) {
const float length = cell.row(v).norm();
auto it = std::find_if(group_length.begin(), group_length.end(),
[length](float l) { return std::abs(l - length) < LENGTH_THRESHOLD; });
group[v] = static_cast<int>(it - group_length.begin());
if (it == group_length.end())
group_length.push_back(length);
}
std::vector<std::vector<Eigen::Vector3f>> candidates(group_length.size());
for (size_t g = 0; g < group_length.size(); g++) {
for (const int i: FindCandidates(t, spots, group_length[g])) {
const Eigen::Vector3f dir(t.dir[0][i], t.dir[1][i], t.dir[2][i]);
candidates[g].push_back(RefineCandidate(spots, group_length[g] * dir));
}
}
// Every cell vector along every candidate of its length, both ways round
std::vector<CellFrame> frames;
for (int v = 0; v < 3; v++) {
for (const auto &c: candidates[group[v]]) {
frames.push_back(MakeFrame(cell, v, c));
frames.push_back(MakeFrame(cell, v, -c));
}
}
std::vector<float> score(frames.size() * NUM_ANGLE_POINTS);
for (size_t i = 0; i < frames.size(); i++)
ScoreAngles(t, frames[i], spots, score.data() + i * NUM_ANGLE_POINTS);
FFBIDXCells ret;
ret.cells.resize(3 * NUM_OUTPUT_CELLS, 3);
ret.scores.resize(NUM_OUTPUT_CELLS);
const auto best = Best(score, NUM_OUTPUT_CELLS);
for (int k = 0; k < NUM_OUTPUT_CELLS; k++) {
const CellFrame &f = frames[best[k] / NUM_ANGLE_POINTS];
const int q = best[k] % NUM_ANGLE_POINTS;
ret.scores(k) = score[best[k]];
ret.cells.row(3 * k + f.cell_vec) = f.c_length * f.z;
ret.cells.row(3 * k + (f.cell_vec + 1) % 3) = f.A(t.cos_alpha[q], t.sin_alpha[q]);
ret.cells.row(3 * k + (f.cell_vec + 2) % 3) = f.B(t.cos_alpha[q], t.sin_alpha[q]);
}
return ret;
}
void FFBIDXIndexerCPU::SetupUnitCell(const std::optional<UnitCell> &cell) {
if (!cell.has_value())
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
"FFBIDX requires unit cell");
reference_unit_cell = cell;
CrystalLattice l(cell.value());
input_cell << l.Vec0().x, l.Vec0().y, l.Vec0().z,
l.Vec1().x, l.Vec1().y, l.Vec1().z,
l.Vec2().x, l.Vec2().y, l.Vec2().z;
}
std::vector<CrystalLattice> FFBIDXIndexerCPU::RunInternal(const std::vector<Coord> &coord, size_t nspots) {
if (nspots > coord.size())
nspots = coord.size();
if (nspots < viable_cell_min_spots)
return {};
auto found = FFBIDXFindCells(input_cell, coord, nspots);
RefineParameters parameters{
.viable_cell_min_spots = viable_cell_min_spots,
.dist_tolerance_vs_reference = dist_tolerance_vs_reference,
.reference_unit_cell = reference_unit_cell,
.min_length_A = 1, // doesn't matter
.max_length_A = 1000, // doesn't matter
.min_angle_deg = 30,
.max_angle_deg = 150,
.indexing_tolerance = indexing_tolerance
};
return Refine(coord, nspots, found.cells, found.scores, parameters);
}