From a67402961d7ed08592253ed319f4ac17d50544a6 Mon Sep 17 00:00:00 2001 From: Duncan Tourolle Date: Sat, 29 Aug 2026 11:58:07 +0200 Subject: [PATCH] Let the merge engine's dot product use the machine's kernel too MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Engine::cross is the one place a dot product is computed during agglomeration — when two groups become adjacent through a third and their sub-threshold pairs, never summed because they were never interesting, have to be accounted for. It was calling the portable loop while the scan beside it had AVX2 or NEON, which on the reference library was 1,753,514 dot products taking 0.54s. Co-Authored-By: Claude Opus 5 (1M context) --- core/dr-face/src/cluster.rs | 7 ++++++- core/dr-face/src/neighbours.rs | 4 ++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/core/dr-face/src/cluster.rs b/core/dr-face/src/cluster.rs index 8f578e3..065f81d 100644 --- a/core/dr-face/src/cluster.rs +++ b/core/dr-face/src/cluster.rs @@ -370,6 +370,10 @@ impl Ord for Pending { struct Engine<'a> { faces: &'a [Candidate], cal: &'a Calibration, + /// The same kernel [`neighbours`] scans with. [`Engine::cross`] is the one + /// place a dot product is computed during agglomeration, and it was using + /// the portable loop while the scan beside it had the machine's SIMD. + dot: neighbours::DotFn, min_probability: f32, groups: Vec, links: HashMap<(usize, usize), Link>, @@ -394,6 +398,7 @@ impl<'a> Engine<'a> { Self { faces, cal, + dot: neighbours::fastest_dot(), min_probability, groups, links: HashMap::new(), @@ -581,7 +586,7 @@ impl<'a> Engine<'a> { let mut count = 0.0_f64; for &i in members { for &j in &self.groups[group].members { - let cos = neighbours::dot(&self.faces[i].embedding, &self.faces[j].embedding); + let cos = (self.dot)(&self.faces[i].embedding, &self.faces[j].embedding); let min_crop = self.faces[i].crop_px.min(self.faces[j].crop_px); sum += self.cal.probability(cos, min_crop, 0.0) as f64; count += 1.0; diff --git a/core/dr-face/src/neighbours.rs b/core/dr-face/src/neighbours.rs index 03d0dd8..b6b4213 100644 --- a/core/dr-face/src/neighbours.rs +++ b/core/dr-face/src/neighbours.rs @@ -340,7 +340,7 @@ fn loosest_cosine(crop_px: &[f32], cal: &Calibration, min_probability: f32) -> f /// A dot product over two equal-length, L2-normalised rows. /// /// Chosen once per scan rather than per pair — see [`fastest_dot`]. -type DotFn = fn(&[f32], &[f32]) -> f32; +pub(crate) type DotFn = fn(&[f32], &[f32]) -> f32; /// The widest dot product this machine can actually run. /// @@ -369,7 +369,7 @@ type DotFn = fn(&[f32], &[f32]) -> f32; /// last bit between them. That only matters for a pair sitting exactly on the /// threshold, and the portable kernel already sums in eight accumulators rather /// than one, so the module was never bit-comparable with a naive sum. -fn fastest_dot() -> DotFn { +pub(crate) fn fastest_dot() -> DotFn { #[cfg(target_arch = "x86_64")] { if std::arch::is_x86_feature_detected!("avx2") && std::arch::is_x86_feature_detected!("fma")