Skip to content

Commit 69d7900

Browse files
committed
fix(jina): use crate::simd::F32x16 for Base17 L1 distance hot path
L1 distance between Base17 tokens now uses F32x16: 16 dims via SIMD (load, subtract, abs, reduce_sum) 1 dim scalar remainder (the 17th dimension) Consumer never sees hardware — F32x16 dispatches via LazyLock to AVX-512/AVX2/scalar automatically. 12 tests passing. https://claude.ai/code/session_01Y69Vnw751w75iVSBRws7o7
1 parent bd5f14e commit 69d7900

1 file changed

Lines changed: 19 additions & 4 deletions

File tree

src/hpc/jina/codec.rs

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -48,13 +48,28 @@ impl Base17Token {
4848
}
4949

5050
/// L1 distance between two Base17 tokens.
51+
///
52+
/// 17 dimensions: 16 via `crate::simd::F32x16` + 1 scalar remainder.
53+
/// Consumer never sees hardware — `F32x16` dispatches to AVX-512/AVX2/scalar.
5154
#[inline(always)]
5255
pub fn l1(&self, other: &Base17Token) -> u32 {
53-
let mut d = 0u32;
54-
for i in 0..BASE_DIM {
55-
d += (self.dims[i] as i32 - other.dims[i] as i32).unsigned_abs();
56+
// SIMD path: load 16 dims as f32, compute abs diff, reduce.
57+
// The 17th dim is scalar.
58+
let mut a_f32 = [0.0f32; 16];
59+
let mut b_f32 = [0.0f32; 16];
60+
for i in 0..16 {
61+
a_f32[i] = self.dims[i] as f32;
62+
b_f32[i] = other.dims[i] as f32;
5663
}
57-
d
64+
let va = F32x16::from_slice(&a_f32);
65+
let vb = F32x16::from_slice(&b_f32);
66+
let diff = va - vb;
67+
let abs_diff = diff.abs();
68+
let simd_sum = abs_diff.reduce_sum() as u32;
69+
70+
// 17th dimension: scalar
71+
let d16 = (self.dims[16] as i32 - other.dims[16] as i32).unsigned_abs();
72+
simd_sum + d16
5873
}
5974
}
6075

0 commit comments

Comments
 (0)