1pub const TQ_CODEC_VERSION: u32 = 2;
20pub const TQ_BITS: u32 = 4;
22pub const TQ_BLOCK_LANES: usize = 16;
24pub const TQ_MIN_PADDED_DIM: usize = 8;
27
28const TQ_STAGE1_LEVELS: usize = 8;
29const TQ_STAGE1_SEED: u64 = 0x7154_5354_4147_4531; const TQ_QJL_SEED: u64 = 0x7154_514a_4c53_4b31; const TQ_LLOYD_GRID: usize = 8192;
32const TQ_LLOYD_MAX_ITERATIONS: usize = 64;
33const TQ_LLOYD_TOLERANCE: f64 = 1e-9;
34#[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
37const TQ_ACCUMULATE_CHUNK_DIMS: usize = 128;
38
39#[inline]
45pub fn tq_padded_dim(dim: usize) -> usize {
46 dim.next_multiple_of(2).max(TQ_MIN_PADDED_DIM)
47}
48
49#[inline]
51pub fn tq_expected_fingerprint(dim: usize) -> u64 {
52 tq_fingerprint(dim, tq_padded_dim(dim))
53}
54
55pub fn tq_shared_codec(dim: usize) -> std::sync::Arc<TqCodec> {
59 static CODECS: std::sync::OnceLock<
60 std::sync::Mutex<rustc_hash::FxHashMap<usize, std::sync::Arc<TqCodec>>>,
61 > = std::sync::OnceLock::new();
62 let cache = CODECS.get_or_init(Default::default);
63 let mut guard = cache
64 .lock()
65 .unwrap_or_else(|poisoned| poisoned.into_inner());
66 std::sync::Arc::clone(
67 guard
68 .entry(dim)
69 .or_insert_with(|| std::sync::Arc::new(TqCodec::new(dim))),
70 )
71}
72
73#[inline]
75pub const fn tq_block_bytes(code_size: usize) -> usize {
76 TQ_BLOCK_LANES * (size_of::<f32>() + code_size)
77}
78
79#[inline]
81pub const fn tq_codes_column_len(count: usize, code_size: usize) -> usize {
82 count.div_ceil(TQ_BLOCK_LANES) * tq_block_bytes(code_size)
83}
84
85#[inline]
87pub fn tq_codes_column_len_checked(count: usize, code_size: usize) -> Option<usize> {
88 count
89 .div_ceil(TQ_BLOCK_LANES)
90 .checked_mul(tq_block_bytes(code_size))
91}
92
93#[inline]
96pub const fn tq_ivf_block_bytes(code_size: usize) -> usize {
97 TQ_BLOCK_LANES * (2 * size_of::<f32>() + code_size)
98}
99
100#[inline]
102pub fn tq_ivf_codes_column_len_checked(count: usize, code_size: usize) -> Option<usize> {
103 count
104 .div_ceil(TQ_BLOCK_LANES)
105 .checked_mul(tq_ivf_block_bytes(code_size))
106}
107
108#[inline]
109fn splitmix64(state: &mut u64) -> u64 {
110 *state = state.wrapping_add(0x9e37_79b9_7f4a_7c15);
111 let mut z = *state;
112 z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
113 z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
114 z ^ (z >> 31)
115}
116
117const TQ_ROTATION_ROUNDS: usize = 3;
122
123#[derive(Debug, Clone)]
130pub struct TqRotation {
131 input_dim: usize,
132 padded_dim: usize,
133 fwht_len: usize,
135 signs: Vec<f32>,
137 perms: Vec<u32>,
139}
140
141impl TqRotation {
142 pub fn new(input_dim: usize, seed: u64) -> Self {
143 let padded_dim = tq_padded_dim(input_dim);
144 let fwht_len = if padded_dim.is_power_of_two() {
145 padded_dim
146 } else {
147 padded_dim.next_power_of_two() / 2
148 };
149 let mut state = seed;
150 let mut signs = Vec::with_capacity(TQ_ROTATION_ROUNDS * padded_dim);
151 let mut perms = Vec::with_capacity(TQ_ROTATION_ROUNDS * padded_dim);
152 for _ in 0..TQ_ROTATION_ROUNDS {
153 signs.extend((0..padded_dim).map(|_| {
154 if splitmix64(&mut state) & 1 == 1 {
155 1.0f32
156 } else {
157 -1.0f32
158 }
159 }));
160 let round_base = perms.len();
161 perms.extend(0..padded_dim as u32);
162 for i in (1..padded_dim).rev() {
163 let j = (splitmix64(&mut state) % (i as u64 + 1)) as usize;
164 perms.swap(round_base + i, round_base + j);
165 }
166 }
167 Self {
168 input_dim,
169 padded_dim,
170 fwht_len,
171 signs,
172 perms,
173 }
174 }
175
176 #[inline]
177 pub fn padded_dim(&self) -> usize {
178 self.padded_dim
179 }
180
181 pub fn apply(&self, input: &[f32], scratch: &mut Vec<f32>, output: &mut [f32]) {
185 debug_assert!(input.len() == self.input_dim || input.len() == self.padded_dim);
186 debug_assert_eq!(output.len(), self.padded_dim);
187 let padded_dim = self.padded_dim;
188 scratch.clear();
189 scratch.resize(padded_dim, 0.0);
190 let signs = &self.signs[..padded_dim];
193 for (slot, (sign, index)) in scratch.iter_mut().zip(signs.iter().zip(0..)) {
194 *slot = input.get(index).copied().unwrap_or(0.0) * sign;
195 }
196 fwht_normalized(&mut scratch[..self.fwht_len]);
197 let perm = &self.perms[..padded_dim];
198 for (slot, &source) in output.iter_mut().zip(perm) {
199 *slot = scratch[source as usize];
200 }
201 for round in 1..TQ_ROTATION_ROUNDS {
202 let signs = &self.signs[round * padded_dim..(round + 1) * padded_dim];
203 for (slot, (&value, sign)) in scratch.iter_mut().zip(output.iter().zip(signs)) {
204 *slot = value * sign;
205 }
206 fwht_normalized(&mut scratch[..self.fwht_len]);
207 let perm = &self.perms[round * padded_dim..(round + 1) * padded_dim];
208 for (slot, &source) in output.iter_mut().zip(perm) {
209 *slot = scratch[source as usize];
210 }
211 }
212 }
213}
214
215fn fwht_normalized(values: &mut [f32]) {
217 let len = values.len();
218 debug_assert!(len.is_power_of_two());
219 let mut step = 1;
220 while step < len {
221 let mut base = 0;
222 while base < len {
223 let (left_half, right_half) = values[base..base + step * 2].split_at_mut(step);
224 for (left, right) in left_half.iter_mut().zip(right_half.iter_mut()) {
225 let sum = *left + *right;
226 let difference = *left - *right;
227 *left = sum;
228 *right = difference;
229 }
230 base += step * 2;
231 }
232 step *= 2;
233 }
234 let scale = 1.0 / (len as f32).sqrt();
235 for value in values.iter_mut() {
236 *value *= scale;
237 }
238}
239
240#[derive(Debug, Clone)]
244pub struct TqCodebook {
245 levels: [f32; TQ_STAGE1_LEVELS],
246 boundaries: [f32; TQ_STAGE1_LEVELS - 1],
248}
249
250impl TqCodebook {
251 pub fn analytic(padded_dim: usize) -> Self {
252 assert!(
253 padded_dim >= TQ_MIN_PADDED_DIM,
254 "TQ codebook requires padded_dim >= {TQ_MIN_PADDED_DIM}, got {padded_dim}"
255 );
256 let exponent = (padded_dim as f64 - 3.0) / 2.0;
257 let cell = 2.0 / TQ_LLOYD_GRID as f64;
258 let mut weights = vec![0.0f64; TQ_LLOYD_GRID];
262 let mut positions = vec![0.0f64; TQ_LLOYD_GRID];
263 for index in 0..TQ_LLOYD_GRID {
264 let t = -1.0 + (index as f64 + 0.5) * cell;
265 positions[index] = t;
266 let log_density = exponent * (1.0 - t * t).max(f64::MIN_POSITIVE).ln();
267 weights[index] = log_density.exp();
268 }
269
270 let total_mass: f64 = weights.iter().sum();
272 let mut levels = [0.0f64; TQ_STAGE1_LEVELS];
273 let mut boundaries = [0.0f64; TQ_STAGE1_LEVELS - 1];
274 let mut accumulated = 0.0f64;
275 let mut next_boundary = 0usize;
276 for index in 0..TQ_LLOYD_GRID {
277 accumulated += weights[index];
278 while next_boundary < TQ_STAGE1_LEVELS - 1
279 && accumulated
280 >= total_mass * (next_boundary as f64 + 1.0) / TQ_STAGE1_LEVELS as f64
281 {
282 boundaries[next_boundary] = positions[index];
283 next_boundary += 1;
284 }
285 }
286
287 for _ in 0..TQ_LLOYD_MAX_ITERATIONS {
290 let mut mass = [0.0f64; TQ_STAGE1_LEVELS];
291 let mut moment = [0.0f64; TQ_STAGE1_LEVELS];
292 let mut bucket = 0usize;
293 for index in 0..TQ_LLOYD_GRID {
294 let t = positions[index];
295 while bucket < TQ_STAGE1_LEVELS - 1 && t > boundaries[bucket] {
296 bucket += 1;
297 }
298 mass[bucket] += weights[index];
299 moment[bucket] += weights[index] * t;
300 }
301 let mut shift = 0.0f64;
302 for level in 0..TQ_STAGE1_LEVELS {
303 if mass[level] > 0.0 {
304 let updated = moment[level] / mass[level];
305 shift = shift.max((updated - levels[level]).abs());
306 levels[level] = updated;
307 }
308 }
309 for boundary in 0..TQ_STAGE1_LEVELS - 1 {
310 boundaries[boundary] = 0.5 * (levels[boundary] + levels[boundary + 1]);
311 }
312 if shift < TQ_LLOYD_TOLERANCE {
313 break;
314 }
315 }
316
317 for index in 0..TQ_STAGE1_LEVELS / 2 {
321 let magnitude = 0.5 * (levels[TQ_STAGE1_LEVELS - 1 - index] - levels[index]);
322 levels[index] = -magnitude;
323 levels[TQ_STAGE1_LEVELS - 1 - index] = magnitude;
324 }
325 for boundary in 0..TQ_STAGE1_LEVELS - 1 {
326 boundaries[boundary] = 0.5 * (levels[boundary] + levels[boundary + 1]);
327 }
328
329 Self {
330 levels: levels.map(|level| level as f32),
331 boundaries: boundaries.map(|boundary| boundary as f32),
332 }
333 }
334
335 #[inline]
337 pub fn encode_coordinate(&self, value: f32) -> u8 {
338 let mut code = 0u8;
339 for &boundary in &self.boundaries {
340 code += u8::from(value > boundary);
341 }
342 code
343 }
344}
345
346#[derive(Debug, Clone)]
349pub struct TqCodec {
350 dim: usize,
351 padded_dim: usize,
352 stage1_rotation: TqRotation,
353 qjl_rotation: TqRotation,
354 codebook: TqCodebook,
355 qjl_scale: f32,
357 fingerprint: u64,
358}
359
360impl TqCodec {
361 pub fn new(dim: usize) -> Self {
362 assert!(dim > 0, "TQ codec requires a non-zero dimension");
363 let stage1_rotation = TqRotation::new(dim, TQ_STAGE1_SEED);
364 let padded_dim = stage1_rotation.padded_dim();
365 let qjl_rotation = TqRotation::new(padded_dim, TQ_QJL_SEED);
366 debug_assert_eq!(qjl_rotation.padded_dim(), padded_dim);
367 let codebook = TqCodebook::analytic(padded_dim);
368 let qjl_scale = (std::f64::consts::PI / 2.0).sqrt() as f32 / (padded_dim as f32).sqrt();
369 let fingerprint = tq_fingerprint(dim, padded_dim);
370 Self {
371 dim,
372 padded_dim,
373 stage1_rotation,
374 qjl_rotation,
375 codebook,
376 qjl_scale,
377 fingerprint,
378 }
379 }
380
381 #[inline]
382 pub fn dim(&self) -> usize {
383 self.dim
384 }
385
386 #[inline]
387 pub fn padded_dim(&self) -> usize {
388 self.padded_dim
389 }
390
391 #[inline]
393 pub fn code_size(&self) -> usize {
394 self.padded_dim / 2
395 }
396
397 #[inline]
399 pub fn fingerprint(&self) -> u64 {
400 self.fingerprint
401 }
402
403 pub fn estimated_memory_bytes(&self) -> usize {
406 2 * self.padded_dim * (size_of::<f32>() + size_of::<u32>()) + size_of::<TqCodebook>()
407 }
408
409 pub fn encode_into(
413 &self,
414 vector: &[f32],
415 nibbles: &mut [u8],
416 scratch: &mut TqEncodeScratch,
417 ) -> f32 {
418 self.encode_residual_into(vector, nibbles, scratch).1
419 }
420
421 pub fn encode_residual_into(
427 &self,
428 vector: &[f32],
429 nibbles: &mut [u8],
430 scratch: &mut TqEncodeScratch,
431 ) -> (f32, f32) {
432 assert_eq!(vector.len(), self.dim, "TQ encode dimension mismatch");
433 assert_eq!(nibbles.len(), self.padded_dim, "TQ nibble buffer mismatch");
434 let norm = crate::structures::simd::dot_product_f32(vector, vector, vector.len()).sqrt();
435 if !norm.is_finite() || norm <= 0.0 {
436 nibbles.fill(0);
437 return (0.0, 0.0);
438 }
439 scratch.normalized.clear();
440 scratch
441 .normalized
442 .extend(vector.iter().map(|value| value / norm));
443
444 scratch.rotated.resize(self.padded_dim, 0.0);
445 let (normalized, rotated, fwht) =
446 (&scratch.normalized, &mut scratch.rotated, &mut scratch.fwht);
447 self.stage1_rotation.apply(normalized, fwht, rotated);
448
449 scratch.residual.resize(self.padded_dim, 0.0);
451 let mut residual_norm_sq = 0.0f32;
452 for ((&value, nibble), residual_slot) in scratch
453 .rotated
454 .iter()
455 .zip(nibbles.iter_mut())
456 .zip(scratch.residual.iter_mut())
457 {
458 let code = self.codebook.encode_coordinate(value);
459 *nibble = code << 1;
460 let residual = value - self.codebook.levels[code as usize];
461 *residual_slot = residual;
462 residual_norm_sq += residual * residual;
463 }
464
465 scratch.rotated_residual.resize(self.padded_dim, 0.0);
467 let (residual, rotated_residual, fwht) = (
468 &scratch.residual,
469 &mut scratch.rotated_residual,
470 &mut scratch.fwht,
471 );
472 self.qjl_rotation.apply(residual, fwht, rotated_residual);
473 for (nibble, &rotated) in nibbles.iter_mut().zip(scratch.rotated_residual.iter()) {
474 *nibble |= u8::from(rotated >= 0.0);
475 }
476 (norm, residual_norm_sq.sqrt())
477 }
478}
479
480#[derive(Debug, Default)]
482pub struct TqEncodeScratch {
483 normalized: Vec<f32>,
484 rotated: Vec<f32>,
485 residual: Vec<f32>,
486 rotated_residual: Vec<f32>,
487 fwht: Vec<f32>,
488}
489
490fn tq_fingerprint(dim: usize, padded_dim: usize) -> u64 {
491 let mut hash = 0xcbf2_9ce4_8422_2325u64; let mut mix = |bytes: &[u8]| {
493 for &byte in bytes {
494 hash ^= u64::from(byte);
495 hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
496 }
497 };
498 mix(b"hermes-tq");
499 mix(&TQ_CODEC_VERSION.to_le_bytes());
500 mix(&TQ_BITS.to_le_bytes());
501 mix(&(dim as u64).to_le_bytes());
502 mix(&(padded_dim as u64).to_le_bytes());
503 mix(&TQ_STAGE1_SEED.to_le_bytes());
504 mix(&TQ_QJL_SEED.to_le_bytes());
505 if hash == 0 { 1 } else { hash }
506}
507
508pub struct TqQueryPlan {
516 padded_dim: usize,
517 fingerprint: u64,
518 query_bits: Box<[u32]>,
525 base_lut_i8: Vec<i8>,
526 qjl_lut_i8: Vec<i8>,
527 base_dequant: f32,
528 qjl_dequant: f32,
529 kernel: lut16::Lut16Kernel,
532 #[cfg(test)]
534 reference_luts: (Vec<f32>, Vec<f32>),
535}
536
537impl std::fmt::Debug for TqQueryPlan {
538 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
539 formatter
540 .debug_struct("TqQueryPlan")
541 .field("padded_dim", &self.padded_dim)
542 .field("fingerprint", &self.fingerprint)
543 .field("query_dim", &self.query_bits.len())
544 .finish()
545 }
546}
547
548impl TqQueryPlan {
549 pub fn build(codec: &TqCodec, query: &[f32]) -> Self {
550 assert_eq!(query.len(), codec.dim, "TQ query dimension mismatch");
551 let padded_dim = codec.padded_dim;
552 let norm = crate::structures::simd::dot_product_f32(query, query, query.len()).sqrt();
553 let inverse_norm = if norm.is_finite() && norm > 0.0 {
554 1.0 / norm
555 } else {
556 0.0
557 };
558 let normalized: Vec<f32> = query.iter().map(|value| value * inverse_norm).collect();
559 let mut fwht = Vec::with_capacity(padded_dim);
560 let mut rotated = vec![0.0f32; padded_dim];
561 codec
562 .stage1_rotation
563 .apply(&normalized, &mut fwht, &mut rotated);
564 let mut qjl_rotated = vec![0.0f32; padded_dim];
565 codec
566 .qjl_rotation
567 .apply(&rotated, &mut fwht, &mut qjl_rotated);
568
569 let mut base_lut = vec![0.0f32; padded_dim * 16];
570 let mut qjl_lut = vec![0.0f32; padded_dim * 16];
571 for dim in 0..padded_dim {
572 for nibble in 0..16 {
573 let level = codec.codebook.levels[nibble >> 1];
574 let sign = if nibble & 1 == 1 { 1.0 } else { -1.0 };
575 base_lut[dim * 16 + nibble] = rotated[dim] * level;
576 qjl_lut[dim * 16 + nibble] = sign * codec.qjl_scale * qjl_rotated[dim];
577 }
578 }
579 let (base_lut_i8, base_dequant) = quantize_lut(&base_lut);
580 let (qjl_lut_i8, qjl_dequant) = quantize_lut(&qjl_lut);
581 Self {
582 padded_dim,
583 fingerprint: codec.fingerprint,
584 query_bits: query
585 .iter()
586 .map(|value| value.to_bits())
587 .collect::<Vec<_>>()
588 .into_boxed_slice(),
589 base_lut_i8,
590 qjl_lut_i8,
591 base_dequant,
592 qjl_dequant,
593 kernel: lut16::Lut16Kernel::resolve(),
594 #[cfg(test)]
595 reference_luts: (base_lut, qjl_lut),
596 }
597 }
598
599 #[inline]
600 pub fn padded_dim(&self) -> usize {
601 self.padded_dim
602 }
603
604 #[inline]
605 pub fn fingerprint(&self) -> u64 {
606 self.fingerprint
607 }
608
609 #[inline]
611 pub(crate) fn matches_query(&self, query: &[f32]) -> bool {
612 self.query_bits.len() == query.len()
613 && self
614 .query_bits
615 .iter()
616 .zip(query)
617 .all(|(&bits, value)| bits == value.to_bits())
618 }
619
620 #[cfg(test)]
623 pub(crate) fn estimate_row(&self, nibbles: &[u8], gamma: f32) -> f32 {
624 debug_assert_eq!(nibbles.len(), self.padded_dim);
625 let (base_lut, qjl_lut) = &self.reference_luts;
626 let mut base = 0.0f32;
627 let mut qjl = 0.0f32;
628 for (dim, &nibble) in nibbles.iter().enumerate() {
629 base += base_lut[dim * 16 + nibble as usize];
630 qjl += qjl_lut[dim * 16 + nibble as usize];
631 }
632 base + gamma * qjl
633 }
634}
635
636fn quantize_lut(values: &[f32]) -> (Vec<i8>, f32) {
637 let max_abs = values
638 .iter()
639 .fold(0.0f32, |acc, &value| acc.max(value.abs()));
640 if !max_abs.is_finite() || max_abs <= 0.0 {
641 return (vec![0i8; values.len()], 0.0);
642 }
643 let quantize_scale = 127.0 / max_abs;
644 let quantized = values
645 .iter()
646 .map(|&value| (value * quantize_scale).round().clamp(-127.0, 127.0) as i8)
647 .collect();
648 (quantized, max_abs / 127.0)
649}
650
651pub fn tq_score_block(plan: &TqQueryPlan, block: &[u8], scores: &mut [f32; TQ_BLOCK_LANES]) {
655 debug_assert_eq!(block.len(), tq_block_bytes(plan.padded_dim / 2));
656 let (gamma_bytes, nibble_bytes) = block.split_at(TQ_BLOCK_LANES * size_of::<f32>());
657 let mut base = [0i32; TQ_BLOCK_LANES];
658 let mut qjl = [0i32; TQ_BLOCK_LANES];
659 plan.kernel.accumulate(
660 &plan.base_lut_i8,
661 &plan.qjl_lut_i8,
662 nibble_bytes,
663 plan.padded_dim,
664 &mut base,
665 &mut qjl,
666 );
667 lut16::finish_block(
668 gamma_bytes,
669 &base,
670 &qjl,
671 plan.base_dequant,
672 plan.qjl_dequant,
673 scores,
674 );
675}
676
677pub fn tq_pack_block(rows: &[&[u8]], gammas: &[f32], padded_dim: usize, output: &mut Vec<u8>) {
680 assert!(rows.len() <= TQ_BLOCK_LANES && rows.len() == gammas.len());
681 for lane in 0..TQ_BLOCK_LANES {
682 let gamma = gammas.get(lane).copied().unwrap_or(0.0);
683 output.extend_from_slice(&gamma.to_le_bytes());
684 }
685 pack_nibble_rows(rows, padded_dim, output);
686}
687
688pub fn tq_pack_ivf_block(
690 rows: &[&[u8]],
691 scales: &[f32],
692 gammas: &[f32],
693 padded_dim: usize,
694 output: &mut Vec<u8>,
695) {
696 assert!(rows.len() <= TQ_BLOCK_LANES && rows.len() == gammas.len());
697 assert_eq!(scales.len(), gammas.len());
698 for lane in 0..TQ_BLOCK_LANES {
699 let scale = scales.get(lane).copied().unwrap_or(0.0);
700 output.extend_from_slice(&scale.to_le_bytes());
701 }
702 for lane in 0..TQ_BLOCK_LANES {
703 let gamma = gammas.get(lane).copied().unwrap_or(0.0);
704 output.extend_from_slice(&gamma.to_le_bytes());
705 }
706 pack_nibble_rows(rows, padded_dim, output);
707}
708
709fn pack_nibble_rows(rows: &[&[u8]], padded_dim: usize, output: &mut Vec<u8>) {
710 for dim in 0..padded_dim {
711 for byte_index in 0..TQ_BLOCK_LANES / 2 {
712 let low = rows.get(byte_index).map_or(0, |row| row[dim] & 0x0F);
713 let high = rows
714 .get(byte_index + TQ_BLOCK_LANES / 2)
715 .map_or(0, |row| row[dim] & 0x0F);
716 output.push(low | (high << 4));
717 }
718 }
719}
720
721pub fn tq_score_ivf_block(
725 plan: &TqQueryPlan,
726 block: &[u8],
727 cluster_dot: f32,
728 scores: &mut [f32; TQ_BLOCK_LANES],
729) {
730 debug_assert_eq!(block.len(), tq_ivf_block_bytes(plan.padded_dim() / 2));
731 let lane_f32 = TQ_BLOCK_LANES * size_of::<f32>();
732 let (scale_bytes, rest) = block.split_at(lane_f32);
733 let (gamma_bytes, nibble_bytes) = rest.split_at(lane_f32);
734 let mut base = [0i32; TQ_BLOCK_LANES];
735 let mut qjl = [0i32; TQ_BLOCK_LANES];
736 plan.kernel.accumulate(
737 &plan.base_lut_i8,
738 &plan.qjl_lut_i8,
739 nibble_bytes,
740 plan.padded_dim,
741 &mut base,
742 &mut qjl,
743 );
744 lut16::finish_ivf_block(
745 scale_bytes,
746 gamma_bytes,
747 &base,
748 &qjl,
749 plan.base_dequant,
750 plan.qjl_dequant,
751 cluster_dot,
752 scores,
753 );
754}
755
756mod lut16 {
757 #[cfg(any(target_arch = "aarch64", target_arch = "x86_64"))]
758 use super::TQ_ACCUMULATE_CHUNK_DIMS;
759 use super::TQ_BLOCK_LANES;
760
761 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
765 pub(super) enum Lut16Kernel {
766 #[cfg(target_arch = "aarch64")]
767 Neon,
768 #[cfg(target_arch = "x86_64")]
769 Avx2,
770 #[cfg(target_arch = "x86_64")]
771 Ssse3,
772 #[cfg_attr(target_arch = "aarch64", allow(dead_code))]
775 Scalar,
776 }
777
778 impl Lut16Kernel {
779 pub(super) fn resolve() -> Self {
780 #[cfg(target_arch = "aarch64")]
781 {
782 Self::Neon
783 }
784 #[cfg(target_arch = "x86_64")]
785 {
786 if std::arch::is_x86_feature_detected!("avx2") {
787 return Self::Avx2;
788 }
789 if std::arch::is_x86_feature_detected!("ssse3") {
790 return Self::Ssse3;
791 }
792 Self::Scalar
793 }
794 #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
795 {
796 Self::Scalar
797 }
798 }
799
800 #[inline]
804 pub(super) fn accumulate(
805 self,
806 base_lut: &[i8],
807 qjl_lut: &[i8],
808 nibble_bytes: &[u8],
809 padded_dim: usize,
810 base: &mut [i32; TQ_BLOCK_LANES],
811 qjl: &mut [i32; TQ_BLOCK_LANES],
812 ) {
813 match self {
814 #[cfg(target_arch = "aarch64")]
815 Self::Neon => unsafe {
816 accumulate_block_neon(base_lut, qjl_lut, nibble_bytes, padded_dim, base, qjl)
817 },
818 #[cfg(target_arch = "x86_64")]
819 Self::Avx2 => unsafe {
820 accumulate_block_avx2(base_lut, qjl_lut, nibble_bytes, padded_dim, base, qjl)
821 },
822 #[cfg(target_arch = "x86_64")]
823 Self::Ssse3 => unsafe {
824 accumulate_block_ssse3(base_lut, qjl_lut, nibble_bytes, padded_dim, base, qjl)
825 },
826 Self::Scalar => {
827 accumulate_block_scalar(base_lut, qjl_lut, nibble_bytes, padded_dim, base, qjl)
828 }
829 }
830 }
831 }
832
833 #[cfg(test)]
836 pub(super) fn accumulate_block(
837 base_lut: &[i8],
838 qjl_lut: &[i8],
839 nibble_bytes: &[u8],
840 padded_dim: usize,
841 base: &mut [i32; TQ_BLOCK_LANES],
842 qjl: &mut [i32; TQ_BLOCK_LANES],
843 ) {
844 Lut16Kernel::resolve().accumulate(base_lut, qjl_lut, nibble_bytes, padded_dim, base, qjl)
845 }
846
847 #[inline]
853 pub(super) fn finish_block(
854 gamma_bytes: &[u8],
855 base: &[i32; TQ_BLOCK_LANES],
856 qjl: &[i32; TQ_BLOCK_LANES],
857 base_dequant: f32,
858 qjl_dequant: f32,
859 scores: &mut [f32; TQ_BLOCK_LANES],
860 ) {
861 assert!(gamma_bytes.len() >= TQ_BLOCK_LANES * 4);
862 #[cfg(target_arch = "aarch64")]
863 {
864 unsafe { finish_block_neon(gamma_bytes, base, qjl, base_dequant, qjl_dequant, scores) }
865 return;
866 }
867 #[cfg(target_arch = "x86_64")]
868 {
869 finish_block_sse2(gamma_bytes, base, qjl, base_dequant, qjl_dequant, scores);
871 return;
872 }
873 #[allow(unreachable_code)]
874 finish_block_scalar(gamma_bytes, base, qjl, base_dequant, qjl_dequant, scores);
875 }
876
877 #[inline]
879 #[allow(clippy::too_many_arguments)]
880 pub(super) fn finish_ivf_block(
881 scale_bytes: &[u8],
882 gamma_bytes: &[u8],
883 base: &[i32; TQ_BLOCK_LANES],
884 qjl: &[i32; TQ_BLOCK_LANES],
885 base_dequant: f32,
886 qjl_dequant: f32,
887 cluster_dot: f32,
888 scores: &mut [f32; TQ_BLOCK_LANES],
889 ) {
890 assert!(scale_bytes.len() >= TQ_BLOCK_LANES * 4);
891 assert!(gamma_bytes.len() >= TQ_BLOCK_LANES * 4);
892 #[cfg(target_arch = "aarch64")]
893 {
894 unsafe {
895 finish_ivf_block_neon(
896 scale_bytes,
897 gamma_bytes,
898 base,
899 qjl,
900 base_dequant,
901 qjl_dequant,
902 cluster_dot,
903 scores,
904 )
905 }
906 return;
907 }
908 #[cfg(target_arch = "x86_64")]
909 {
910 finish_ivf_block_sse2(
911 scale_bytes,
912 gamma_bytes,
913 base,
914 qjl,
915 base_dequant,
916 qjl_dequant,
917 cluster_dot,
918 scores,
919 );
920 return;
921 }
922 #[allow(unreachable_code)]
923 finish_ivf_block_scalar(
924 scale_bytes,
925 gamma_bytes,
926 base,
927 qjl,
928 base_dequant,
929 qjl_dequant,
930 cluster_dot,
931 scores,
932 );
933 }
934
935 #[cfg_attr(any(target_arch = "aarch64", target_arch = "x86_64"), allow(dead_code))]
938 #[inline]
939 fn lane_f32(bytes: &[u8], lane: usize) -> f32 {
940 f32::from_le_bytes(
941 bytes[lane * 4..lane * 4 + 4]
942 .try_into()
943 .expect("lane slice is 4 bytes"),
944 )
945 }
946
947 #[cfg_attr(any(target_arch = "aarch64", target_arch = "x86_64"), allow(dead_code))]
949 pub(super) fn finish_block_scalar(
950 gamma_bytes: &[u8],
951 base: &[i32; TQ_BLOCK_LANES],
952 qjl: &[i32; TQ_BLOCK_LANES],
953 base_dequant: f32,
954 qjl_dequant: f32,
955 scores: &mut [f32; TQ_BLOCK_LANES],
956 ) {
957 for lane in 0..TQ_BLOCK_LANES {
958 let gamma = lane_f32(gamma_bytes, lane);
959 scores[lane] =
960 base[lane] as f32 * base_dequant + gamma * qjl[lane] as f32 * qjl_dequant;
961 }
962 }
963
964 #[cfg_attr(any(target_arch = "aarch64", target_arch = "x86_64"), allow(dead_code))]
966 #[allow(clippy::too_many_arguments)]
967 pub(super) fn finish_ivf_block_scalar(
968 scale_bytes: &[u8],
969 gamma_bytes: &[u8],
970 base: &[i32; TQ_BLOCK_LANES],
971 qjl: &[i32; TQ_BLOCK_LANES],
972 base_dequant: f32,
973 qjl_dequant: f32,
974 cluster_dot: f32,
975 scores: &mut [f32; TQ_BLOCK_LANES],
976 ) {
977 for lane in 0..TQ_BLOCK_LANES {
978 let scale = lane_f32(scale_bytes, lane);
979 let gamma = lane_f32(gamma_bytes, lane);
980 scores[lane] = cluster_dot
981 + scale
982 * (base[lane] as f32 * base_dequant + gamma * qjl[lane] as f32 * qjl_dequant);
983 }
984 }
985
986 #[cfg(target_arch = "aarch64")]
987 #[target_feature(enable = "neon")]
988 unsafe fn finish_block_neon(
989 gamma_bytes: &[u8],
990 base: &[i32; TQ_BLOCK_LANES],
991 qjl: &[i32; TQ_BLOCK_LANES],
992 base_dequant: f32,
993 qjl_dequant: f32,
994 scores: &mut [f32; TQ_BLOCK_LANES],
995 ) {
996 use std::arch::aarch64::*;
997 unsafe {
1002 let bd = vdupq_n_f32(base_dequant);
1003 let qd = vdupq_n_f32(qjl_dequant);
1004 for quarter in 0..4 {
1005 let gamma = vreinterpretq_f32_u8(vld1q_u8(gamma_bytes.as_ptr().add(quarter * 16)));
1006 let b = vcvtq_f32_s32(vld1q_s32(base.as_ptr().add(quarter * 4)));
1007 let j = vcvtq_f32_s32(vld1q_s32(qjl.as_ptr().add(quarter * 4)));
1008 let score = vaddq_f32(vmulq_f32(b, bd), vmulq_f32(vmulq_f32(gamma, j), qd));
1009 vst1q_f32(scores.as_mut_ptr().add(quarter * 4), score);
1010 }
1011 }
1012 }
1013
1014 #[cfg(target_arch = "aarch64")]
1015 #[target_feature(enable = "neon")]
1016 #[allow(clippy::too_many_arguments)]
1017 unsafe fn finish_ivf_block_neon(
1018 scale_bytes: &[u8],
1019 gamma_bytes: &[u8],
1020 base: &[i32; TQ_BLOCK_LANES],
1021 qjl: &[i32; TQ_BLOCK_LANES],
1022 base_dequant: f32,
1023 qjl_dequant: f32,
1024 cluster_dot: f32,
1025 scores: &mut [f32; TQ_BLOCK_LANES],
1026 ) {
1027 use std::arch::aarch64::*;
1028 unsafe {
1031 let bd = vdupq_n_f32(base_dequant);
1032 let qd = vdupq_n_f32(qjl_dequant);
1033 let cd = vdupq_n_f32(cluster_dot);
1034 for quarter in 0..4 {
1035 let scale = vreinterpretq_f32_u8(vld1q_u8(scale_bytes.as_ptr().add(quarter * 16)));
1036 let gamma = vreinterpretq_f32_u8(vld1q_u8(gamma_bytes.as_ptr().add(quarter * 16)));
1037 let b = vcvtq_f32_s32(vld1q_s32(base.as_ptr().add(quarter * 4)));
1038 let j = vcvtq_f32_s32(vld1q_s32(qjl.as_ptr().add(quarter * 4)));
1039 let inner = vaddq_f32(vmulq_f32(b, bd), vmulq_f32(vmulq_f32(gamma, j), qd));
1040 vst1q_f32(
1041 scores.as_mut_ptr().add(quarter * 4),
1042 vaddq_f32(cd, vmulq_f32(scale, inner)),
1043 );
1044 }
1045 }
1046 }
1047
1048 #[cfg(target_arch = "x86_64")]
1049 fn finish_block_sse2(
1050 gamma_bytes: &[u8],
1051 base: &[i32; TQ_BLOCK_LANES],
1052 qjl: &[i32; TQ_BLOCK_LANES],
1053 base_dequant: f32,
1054 qjl_dequant: f32,
1055 scores: &mut [f32; TQ_BLOCK_LANES],
1056 ) {
1057 use std::arch::x86_64::*;
1058 unsafe {
1062 let bd = _mm_set1_ps(base_dequant);
1063 let qd = _mm_set1_ps(qjl_dequant);
1064 for quarter in 0..4 {
1065 let gamma = _mm_loadu_ps(gamma_bytes.as_ptr().add(quarter * 16).cast::<f32>());
1066 let b = _mm_cvtepi32_ps(_mm_loadu_si128(base.as_ptr().add(quarter * 4).cast()));
1067 let j = _mm_cvtepi32_ps(_mm_loadu_si128(qjl.as_ptr().add(quarter * 4).cast()));
1068 let score = _mm_add_ps(_mm_mul_ps(b, bd), _mm_mul_ps(_mm_mul_ps(gamma, j), qd));
1069 _mm_storeu_ps(scores.as_mut_ptr().add(quarter * 4), score);
1070 }
1071 }
1072 }
1073
1074 #[cfg(target_arch = "x86_64")]
1075 #[allow(clippy::too_many_arguments)]
1076 fn finish_ivf_block_sse2(
1077 scale_bytes: &[u8],
1078 gamma_bytes: &[u8],
1079 base: &[i32; TQ_BLOCK_LANES],
1080 qjl: &[i32; TQ_BLOCK_LANES],
1081 base_dequant: f32,
1082 qjl_dequant: f32,
1083 cluster_dot: f32,
1084 scores: &mut [f32; TQ_BLOCK_LANES],
1085 ) {
1086 use std::arch::x86_64::*;
1087 unsafe {
1089 let bd = _mm_set1_ps(base_dequant);
1090 let qd = _mm_set1_ps(qjl_dequant);
1091 let cd = _mm_set1_ps(cluster_dot);
1092 for quarter in 0..4 {
1093 let scale = _mm_loadu_ps(scale_bytes.as_ptr().add(quarter * 16).cast::<f32>());
1094 let gamma = _mm_loadu_ps(gamma_bytes.as_ptr().add(quarter * 16).cast::<f32>());
1095 let b = _mm_cvtepi32_ps(_mm_loadu_si128(base.as_ptr().add(quarter * 4).cast()));
1096 let j = _mm_cvtepi32_ps(_mm_loadu_si128(qjl.as_ptr().add(quarter * 4).cast()));
1097 let inner = _mm_add_ps(_mm_mul_ps(b, bd), _mm_mul_ps(_mm_mul_ps(gamma, j), qd));
1098 _mm_storeu_ps(
1099 scores.as_mut_ptr().add(quarter * 4),
1100 _mm_add_ps(cd, _mm_mul_ps(scale, inner)),
1101 );
1102 }
1103 }
1104 }
1105
1106 pub(super) fn accumulate_block_scalar(
1109 base_lut: &[i8],
1110 qjl_lut: &[i8],
1111 nibble_bytes: &[u8],
1112 padded_dim: usize,
1113 base: &mut [i32; TQ_BLOCK_LANES],
1114 qjl: &mut [i32; TQ_BLOCK_LANES],
1115 ) {
1116 for dim in 0..padded_dim {
1117 let row = &nibble_bytes[dim * 8..dim * 8 + 8];
1118 let base_table = &base_lut[dim * 16..dim * 16 + 16];
1119 let qjl_table = &qjl_lut[dim * 16..dim * 16 + 16];
1120 for (lane, &byte) in row.iter().enumerate() {
1121 let low = (byte & 0x0F) as usize;
1122 let high = (byte >> 4) as usize;
1123 base[lane] += i32::from(base_table[low]);
1124 base[lane + 8] += i32::from(base_table[high]);
1125 qjl[lane] += i32::from(qjl_table[low]);
1126 qjl[lane + 8] += i32::from(qjl_table[high]);
1127 }
1128 }
1129 }
1130
1131 #[cfg(target_arch = "aarch64")]
1132 #[target_feature(enable = "neon")]
1133 unsafe fn accumulate_block_neon(
1134 base_lut: &[i8],
1135 qjl_lut: &[i8],
1136 nibble_bytes: &[u8],
1137 padded_dim: usize,
1138 base: &mut [i32; TQ_BLOCK_LANES],
1139 qjl: &mut [i32; TQ_BLOCK_LANES],
1140 ) {
1141 use std::arch::aarch64::*;
1142 unsafe {
1143 let mask = vdup_n_u8(0x0F);
1144 let mut base_lo_i32 = [vdupq_n_s32(0); 4];
1145 let mut qjl_lo_i32 = [vdupq_n_s32(0); 4];
1146 let mut dim = 0;
1147 while dim < padded_dim {
1148 let chunk_end = (dim + TQ_ACCUMULATE_CHUNK_DIMS).min(padded_dim);
1149 let mut base_acc = [vdupq_n_s16(0); 2];
1150 let mut qjl_acc = [vdupq_n_s16(0); 2];
1151 while dim < chunk_end {
1152 let row = vld1_u8(nibble_bytes.as_ptr().add(dim * 8));
1153 let low = vand_u8(row, mask);
1154 let high = vshr_n_u8::<4>(row);
1155 let lanes = vcombine_u8(low, high);
1156 let base_table = vld1q_s8(base_lut.as_ptr().add(dim * 16));
1157 let qjl_table = vld1q_s8(qjl_lut.as_ptr().add(dim * 16));
1158 let base_values =
1159 vreinterpretq_s8_u8(vqtbl1q_u8(vreinterpretq_u8_s8(base_table), lanes));
1160 let qjl_values =
1161 vreinterpretq_s8_u8(vqtbl1q_u8(vreinterpretq_u8_s8(qjl_table), lanes));
1162 base_acc[0] = vaddw_s8(base_acc[0], vget_low_s8(base_values));
1163 base_acc[1] = vaddw_s8(base_acc[1], vget_high_s8(base_values));
1164 qjl_acc[0] = vaddw_s8(qjl_acc[0], vget_low_s8(qjl_values));
1165 qjl_acc[1] = vaddw_s8(qjl_acc[1], vget_high_s8(qjl_values));
1166 dim += 1;
1167 }
1168 for half in 0..2 {
1169 base_lo_i32[half * 2] =
1170 vaddw_s16(base_lo_i32[half * 2], vget_low_s16(base_acc[half]));
1171 base_lo_i32[half * 2 + 1] =
1172 vaddw_s16(base_lo_i32[half * 2 + 1], vget_high_s16(base_acc[half]));
1173 qjl_lo_i32[half * 2] =
1174 vaddw_s16(qjl_lo_i32[half * 2], vget_low_s16(qjl_acc[half]));
1175 qjl_lo_i32[half * 2 + 1] =
1176 vaddw_s16(qjl_lo_i32[half * 2 + 1], vget_high_s16(qjl_acc[half]));
1177 }
1178 }
1179 for quarter in 0..4 {
1180 vst1q_s32(base.as_mut_ptr().add(quarter * 4), base_lo_i32[quarter]);
1181 vst1q_s32(qjl.as_mut_ptr().add(quarter * 4), qjl_lo_i32[quarter]);
1182 }
1183 }
1184 }
1185
1186 #[cfg(target_arch = "x86_64")]
1190 #[target_feature(enable = "avx2")]
1191 unsafe fn accumulate_block_avx2(
1192 base_lut: &[i8],
1193 qjl_lut: &[i8],
1194 nibble_bytes: &[u8],
1195 padded_dim: usize,
1196 base: &mut [i32; TQ_BLOCK_LANES],
1197 qjl: &mut [i32; TQ_BLOCK_LANES],
1198 ) {
1199 use std::arch::x86_64::*;
1200 unsafe {
1201 let mask = _mm_set1_epi8(0x0F);
1202 let zero256 = _mm256_setzero_si256();
1203 let mut base_i32 = [zero256; 2];
1204 let mut qjl_i32 = [zero256; 2];
1205 let mut dim = 0;
1206 while dim < padded_dim {
1207 let chunk_end = (dim + TQ_ACCUMULATE_CHUNK_DIMS).min(padded_dim);
1208 let mut base_acc = zero256;
1209 let mut qjl_acc = zero256;
1210 while dim + 2 <= chunk_end {
1211 let rows = _mm_loadu_si128(nibble_bytes.as_ptr().add(dim * 8).cast());
1213 let low = _mm_and_si128(rows, mask);
1214 let high = _mm_and_si128(_mm_srli_epi16(rows, 4), mask);
1215 let lanes_first = _mm_unpacklo_epi64(low, high);
1217 let lanes_second = _mm_unpackhi_epi64(low, high);
1218 let lanes = _mm256_inserti128_si256(
1219 _mm256_castsi128_si256(lanes_first),
1220 lanes_second,
1221 1,
1222 );
1223 let base_tables = _mm256_loadu_si256(base_lut.as_ptr().add(dim * 16).cast());
1224 let qjl_tables = _mm256_loadu_si256(qjl_lut.as_ptr().add(dim * 16).cast());
1225 let base_values = _mm256_shuffle_epi8(base_tables, lanes);
1226 let qjl_values = _mm256_shuffle_epi8(qjl_tables, lanes);
1227 base_acc = _mm256_add_epi16(
1228 base_acc,
1229 _mm256_cvtepi8_epi16(_mm256_castsi256_si128(base_values)),
1230 );
1231 base_acc = _mm256_add_epi16(
1232 base_acc,
1233 _mm256_cvtepi8_epi16(_mm256_extracti128_si256(base_values, 1)),
1234 );
1235 qjl_acc = _mm256_add_epi16(
1236 qjl_acc,
1237 _mm256_cvtepi8_epi16(_mm256_castsi256_si128(qjl_values)),
1238 );
1239 qjl_acc = _mm256_add_epi16(
1240 qjl_acc,
1241 _mm256_cvtepi8_epi16(_mm256_extracti128_si256(qjl_values, 1)),
1242 );
1243 dim += 2;
1244 }
1245 while dim < chunk_end {
1247 let row = _mm_loadl_epi64(nibble_bytes.as_ptr().add(dim * 8).cast());
1248 let low = _mm_and_si128(row, mask);
1249 let high = _mm_and_si128(_mm_srli_epi16(row, 4), mask);
1250 let lanes = _mm_unpacklo_epi64(low, high);
1251 let base_table = _mm_loadu_si128(base_lut.as_ptr().add(dim * 16).cast());
1252 let qjl_table = _mm_loadu_si128(qjl_lut.as_ptr().add(dim * 16).cast());
1253 base_acc = _mm256_add_epi16(
1254 base_acc,
1255 _mm256_cvtepi8_epi16(_mm_shuffle_epi8(base_table, lanes)),
1256 );
1257 qjl_acc = _mm256_add_epi16(
1258 qjl_acc,
1259 _mm256_cvtepi8_epi16(_mm_shuffle_epi8(qjl_table, lanes)),
1260 );
1261 dim += 1;
1262 }
1263 for (accumulators, chunk) in [(&mut base_i32, base_acc), (&mut qjl_i32, qjl_acc)] {
1264 accumulators[0] = _mm256_add_epi32(
1265 accumulators[0],
1266 _mm256_cvtepi16_epi32(_mm256_castsi256_si128(chunk)),
1267 );
1268 accumulators[1] = _mm256_add_epi32(
1269 accumulators[1],
1270 _mm256_cvtepi16_epi32(_mm256_extracti128_si256(chunk, 1)),
1271 );
1272 }
1273 }
1274 _mm256_storeu_si256(base.as_mut_ptr().cast(), base_i32[0]);
1275 _mm256_storeu_si256(base.as_mut_ptr().add(8).cast(), base_i32[1]);
1276 _mm256_storeu_si256(qjl.as_mut_ptr().cast(), qjl_i32[0]);
1277 _mm256_storeu_si256(qjl.as_mut_ptr().add(8).cast(), qjl_i32[1]);
1278 }
1279 }
1280
1281 #[cfg(target_arch = "x86_64")]
1282 #[target_feature(enable = "ssse3")]
1283 unsafe fn accumulate_block_ssse3(
1284 base_lut: &[i8],
1285 qjl_lut: &[i8],
1286 nibble_bytes: &[u8],
1287 padded_dim: usize,
1288 base: &mut [i32; TQ_BLOCK_LANES],
1289 qjl: &mut [i32; TQ_BLOCK_LANES],
1290 ) {
1291 use std::arch::x86_64::*;
1292 unsafe {
1293 let mask = _mm_set1_epi8(0x0F);
1294 let zero = _mm_setzero_si128();
1295 let mut base_i32 = [zero; 4];
1296 let mut qjl_i32 = [zero; 4];
1297 let mut dim = 0;
1298 while dim < padded_dim {
1299 let chunk_end = (dim + TQ_ACCUMULATE_CHUNK_DIMS).min(padded_dim);
1300 let mut base_acc = [zero; 2];
1301 let mut qjl_acc = [zero; 2];
1302 while dim < chunk_end {
1303 let row = _mm_loadl_epi64(nibble_bytes.as_ptr().add(dim * 8).cast());
1304 let low = _mm_and_si128(row, mask);
1305 let high = _mm_and_si128(_mm_srli_epi16(row, 4), mask);
1306 let lanes = _mm_unpacklo_epi64(low, high);
1307 let base_table = _mm_loadu_si128(base_lut.as_ptr().add(dim * 16).cast());
1308 let qjl_table = _mm_loadu_si128(qjl_lut.as_ptr().add(dim * 16).cast());
1309 let base_values = _mm_shuffle_epi8(base_table, lanes);
1310 let qjl_values = _mm_shuffle_epi8(qjl_table, lanes);
1311 let base_sign = _mm_cmpgt_epi8(zero, base_values);
1313 let qjl_sign = _mm_cmpgt_epi8(zero, qjl_values);
1314 base_acc[0] =
1315 _mm_add_epi16(base_acc[0], _mm_unpacklo_epi8(base_values, base_sign));
1316 base_acc[1] =
1317 _mm_add_epi16(base_acc[1], _mm_unpackhi_epi8(base_values, base_sign));
1318 qjl_acc[0] = _mm_add_epi16(qjl_acc[0], _mm_unpacklo_epi8(qjl_values, qjl_sign));
1319 qjl_acc[1] = _mm_add_epi16(qjl_acc[1], _mm_unpackhi_epi8(qjl_values, qjl_sign));
1320 dim += 1;
1321 }
1322 for half in 0..2 {
1323 let base_sign = _mm_cmpgt_epi16(zero, base_acc[half]);
1324 let qjl_sign = _mm_cmpgt_epi16(zero, qjl_acc[half]);
1325 base_i32[half * 2] = _mm_add_epi32(
1326 base_i32[half * 2],
1327 _mm_unpacklo_epi16(base_acc[half], base_sign),
1328 );
1329 base_i32[half * 2 + 1] = _mm_add_epi32(
1330 base_i32[half * 2 + 1],
1331 _mm_unpackhi_epi16(base_acc[half], base_sign),
1332 );
1333 qjl_i32[half * 2] = _mm_add_epi32(
1334 qjl_i32[half * 2],
1335 _mm_unpacklo_epi16(qjl_acc[half], qjl_sign),
1336 );
1337 qjl_i32[half * 2 + 1] = _mm_add_epi32(
1338 qjl_i32[half * 2 + 1],
1339 _mm_unpackhi_epi16(qjl_acc[half], qjl_sign),
1340 );
1341 }
1342 }
1343 for quarter in 0..4 {
1344 _mm_storeu_si128(base.as_mut_ptr().add(quarter * 4).cast(), base_i32[quarter]);
1345 _mm_storeu_si128(qjl.as_mut_ptr().add(quarter * 4).cast(), qjl_i32[quarter]);
1346 }
1347 }
1348 }
1349}
1350
1351#[cfg(feature = "native")]
1358pub struct TqFlatBuilder {
1359 codec: std::sync::Arc<TqCodec>,
1360 pub doc_ids: Vec<u32>,
1361 pub ordinals: Vec<u16>,
1362 pub codes: Vec<u8>,
1363 pending_rows: Vec<Vec<u8>>,
1364 pending_gammas: Vec<f32>,
1365}
1366
1367#[cfg(feature = "native")]
1368impl TqFlatBuilder {
1369 pub fn new(codec: std::sync::Arc<TqCodec>) -> Self {
1370 Self {
1371 codec,
1372 doc_ids: Vec::new(),
1373 ordinals: Vec::new(),
1374 codes: Vec::new(),
1375 pending_rows: Vec::with_capacity(TQ_BLOCK_LANES),
1376 pending_gammas: Vec::with_capacity(TQ_BLOCK_LANES),
1377 }
1378 }
1379
1380 #[inline]
1381 pub fn codec(&self) -> &TqCodec {
1382 &self.codec
1383 }
1384
1385 #[inline]
1386 pub fn len(&self) -> usize {
1387 self.doc_ids.len()
1388 }
1389
1390 #[inline]
1391 pub fn is_empty(&self) -> bool {
1392 self.doc_ids.is_empty()
1393 }
1394
1395 pub fn add_batch(
1398 &mut self,
1399 labels: &[(u32, u16)],
1400 vectors: &[f32],
1401 ) -> Result<(), &'static str> {
1402 use rayon::prelude::*;
1403
1404 let dim = self.codec.dim();
1405 let vector_count = labels.len();
1406 let expected = vector_count
1407 .checked_mul(dim)
1408 .ok_or("TQ input size overflow")?;
1409 if vectors.len() != expected {
1410 return Err("TQ vector and label matrices are inconsistent");
1411 }
1412 let padded_dim = self.codec.padded_dim();
1413 let codec = std::sync::Arc::clone(&self.codec);
1414 let mut rows = vec![0u8; vector_count * padded_dim];
1418 let mut gammas = vec![0.0f32; vector_count];
1419 vectors
1420 .par_chunks_exact(dim)
1421 .zip(rows.par_chunks_exact_mut(padded_dim))
1422 .zip(gammas.par_iter_mut())
1423 .for_each_init(
1424 TqEncodeScratch::default,
1425 |scratch, ((vector, row), gamma)| {
1426 *gamma = codec.encode_into(vector, row, scratch);
1427 },
1428 );
1429
1430 let mut index = 0;
1433 while index < vector_count && !self.pending_rows.is_empty() {
1434 let (doc_id, ordinal) = labels[index];
1435 self.doc_ids.push(doc_id);
1436 self.ordinals.push(ordinal);
1437 self.pending_rows
1438 .push(rows[index * padded_dim..(index + 1) * padded_dim].to_vec());
1439 self.pending_gammas.push(gammas[index]);
1440 if self.pending_rows.len() == TQ_BLOCK_LANES {
1441 self.flush_block();
1442 }
1443 index += 1;
1444 }
1445 while vector_count - index >= TQ_BLOCK_LANES {
1446 let row_refs: Vec<&[u8]> = (0..TQ_BLOCK_LANES)
1447 .map(|lane| {
1448 let row = index + lane;
1449 &rows[row * padded_dim..(row + 1) * padded_dim]
1450 })
1451 .collect();
1452 tq_pack_block(
1453 &row_refs,
1454 &gammas[index..index + TQ_BLOCK_LANES],
1455 padded_dim,
1456 &mut self.codes,
1457 );
1458 for &(doc_id, ordinal) in &labels[index..index + TQ_BLOCK_LANES] {
1459 self.doc_ids.push(doc_id);
1460 self.ordinals.push(ordinal);
1461 }
1462 index += TQ_BLOCK_LANES;
1463 }
1464 for row in index..vector_count {
1465 let (doc_id, ordinal) = labels[row];
1466 self.doc_ids.push(doc_id);
1467 self.ordinals.push(ordinal);
1468 self.pending_rows
1469 .push(rows[row * padded_dim..(row + 1) * padded_dim].to_vec());
1470 self.pending_gammas.push(gammas[row]);
1471 }
1472 Ok(())
1473 }
1474
1475 fn flush_block(&mut self) {
1476 let rows: Vec<&[u8]> = self.pending_rows.iter().map(Vec::as_slice).collect();
1477 tq_pack_block(
1478 &rows,
1479 &self.pending_gammas,
1480 self.codec.padded_dim(),
1481 &mut self.codes,
1482 );
1483 self.pending_rows.clear();
1484 self.pending_gammas.clear();
1485 }
1486
1487 pub fn finish(&mut self) {
1489 if !self.pending_rows.is_empty() {
1490 self.flush_block();
1491 }
1492 debug_assert_eq!(
1493 self.codes.len(),
1494 tq_codes_column_len(self.doc_ids.len(), self.codec.code_size())
1495 );
1496 }
1497}
1498
1499#[cfg(test)]
1500mod tests {
1501 use super::*;
1502
1503 fn seeded_unit_vector(dim: usize, seed: u64) -> Vec<f32> {
1504 let mut state = seed;
1506 let mut values: Vec<f32> = (0..dim)
1507 .map(|_| {
1508 let a = (splitmix64(&mut state) >> 11) as f64 / (1u64 << 53) as f64;
1509 let b = (splitmix64(&mut state) >> 11) as f64 / (1u64 << 53) as f64;
1510 let gaussian = (-2.0 * (1.0 - a).max(f64::MIN_POSITIVE).ln()).sqrt()
1511 * (2.0 * std::f64::consts::PI * b).cos();
1512 gaussian as f32
1513 })
1514 .collect();
1515 let norm = values.iter().map(|v| v * v).sum::<f32>().sqrt();
1516 values.iter_mut().for_each(|v| *v /= norm);
1517 values
1518 }
1519
1520 #[test]
1521 fn rotation_is_orthonormal_and_deterministic() {
1522 let rotation = TqRotation::new(100, 42);
1523 assert_eq!(rotation.padded_dim(), 100);
1524 let mut fwht_scratch = Vec::new();
1525 let mut probe = vec![0.0f32; 100];
1526 assert_eq!(TqRotation::new(99, 42).padded_dim(), 100);
1528 TqRotation::new(99, 42).apply(&vec![1.0; 99], &mut fwht_scratch, &mut probe);
1529 let input = seeded_unit_vector(100, 7);
1530 let mut fwht = Vec::new();
1531 let mut output = vec![0.0f32; 100];
1532 rotation.apply(&input, &mut fwht, &mut output);
1533 let norm: f32 = output.iter().map(|v| v * v).sum();
1534 assert!(
1535 (norm - 1.0).abs() < 1e-5,
1536 "rotation must preserve norm, got {norm}"
1537 );
1538
1539 let mut second = vec![0.0f32; 100];
1540 TqRotation::new(100, 42).apply(&input, &mut fwht, &mut second);
1541 assert_eq!(output, second, "rotation must be deterministic");
1542
1543 let other = seeded_unit_vector(100, 8);
1545 let mut other_rotated = vec![0.0f32; 100];
1546 rotation.apply(&other, &mut fwht, &mut other_rotated);
1547 let dot_before: f32 = input.iter().zip(&other).map(|(a, b)| a * b).sum();
1548 let dot_after: f32 = output.iter().zip(&other_rotated).map(|(a, b)| a * b).sum();
1549 assert!(
1550 (dot_before - dot_after).abs() < 1e-4,
1551 "rotation must preserve inner products: {dot_before} vs {dot_after}"
1552 );
1553 }
1554
1555 #[test]
1556 fn analytic_codebook_is_symmetric_and_monotonic() {
1557 for padded_dim in [8, 128, 1024] {
1558 let codebook = TqCodebook::analytic(padded_dim);
1559 let levels = &codebook.levels;
1560 for pair in levels.windows(2) {
1561 assert!(pair[0] < pair[1], "levels must be strictly increasing");
1562 }
1563 for index in 0..TQ_STAGE1_LEVELS / 2 {
1564 assert_eq!(
1565 levels[index],
1566 -levels[TQ_STAGE1_LEVELS - 1 - index],
1567 "levels must be exactly symmetric for P={padded_dim}: {levels:?}"
1568 );
1569 }
1570 assert!(levels[TQ_STAGE1_LEVELS - 1] < 1.0);
1571 let scale = 1.0 / (padded_dim as f32).sqrt();
1574 assert!(
1575 levels[TQ_STAGE1_LEVELS - 1] < 6.0 * scale,
1576 "top level {} is implausibly large for P={padded_dim}",
1577 levels[TQ_STAGE1_LEVELS - 1]
1578 );
1579 }
1580 }
1581
1582 #[test]
1583 fn encode_coordinate_matches_nearest_level() {
1584 let codebook = TqCodebook::analytic(256);
1585 for step in -1000i32..=1000 {
1586 let value = step as f32 / 1000.0;
1587 let code = codebook.encode_coordinate(value) as usize;
1588 let nearest = codebook
1589 .levels
1590 .iter()
1591 .enumerate()
1592 .min_by(|a, b| (a.1 - value).abs().total_cmp(&(b.1 - value).abs()))
1593 .unwrap()
1594 .0;
1595 assert_eq!(
1596 code, nearest,
1597 "value {value} coded {code}, nearest {nearest}"
1598 );
1599 }
1600 }
1601
1602 #[test]
1603 fn estimator_is_unbiased_and_tight() {
1604 let dim = 96;
1605 let codec = TqCodec::new(dim);
1606 let mut scratch = TqEncodeScratch::default();
1607 let mut nibbles = vec![0u8; codec.padded_dim()];
1608
1609 let pairs = 512;
1610 let mut signed_error_sum = 0.0f64;
1611 let mut squared_error_sum = 0.0f64;
1612 let mut stage1_signed_error_sum = 0.0f64;
1613 for pair in 0..pairs {
1614 let vector = seeded_unit_vector(dim, 1000 + pair);
1615 let query = seeded_unit_vector(dim, 900_000 + pair);
1616 let gamma = codec.encode_into(&vector, &mut nibbles, &mut scratch);
1617 let plan = TqQueryPlan::build(&codec, &query);
1618 let estimate = plan.estimate_row(&nibbles, gamma);
1619 let stage1_only = plan.estimate_row(&nibbles, 0.0);
1620 let truth: f32 = vector.iter().zip(&query).map(|(a, b)| a * b).sum();
1621 signed_error_sum += f64::from(estimate - truth);
1622 squared_error_sum += f64::from(estimate - truth).powi(2);
1623 stage1_signed_error_sum += f64::from(stage1_only - truth);
1624 }
1625 let mean_error = signed_error_sum / pairs as f64;
1626 let rmse = (squared_error_sum / pairs as f64).sqrt();
1627 let stage1_mean_error = stage1_signed_error_sum / pairs as f64;
1628 assert!(
1629 mean_error.abs() < 3e-3,
1630 "QJL-corrected estimator must be unbiased: mean error {mean_error}"
1631 );
1632 assert!(rmse < 0.05, "estimator RMSE too large: {rmse}");
1633 assert!(
1634 mean_error.abs() <= stage1_mean_error.abs() + 1e-4,
1635 "QJL correction must not increase bias: {mean_error} vs stage-1 {stage1_mean_error}"
1636 );
1637 }
1638
1639 #[test]
1640 fn block_scoring_matches_row_estimates() {
1641 let dim = 100;
1642 let codec = TqCodec::new(dim);
1643 let mut scratch = TqEncodeScratch::default();
1644 let query = seeded_unit_vector(dim, 3);
1645 let plan = TqQueryPlan::build(&codec, &query);
1646
1647 let lanes = 13; let mut rows = Vec::new();
1649 let mut gammas = Vec::new();
1650 for lane in 0..lanes {
1651 let vector = seeded_unit_vector(dim, 50 + lane as u64);
1652 let mut nibbles = vec![0u8; codec.padded_dim()];
1653 let gamma = codec.encode_into(&vector, &mut nibbles, &mut scratch);
1654 rows.push(nibbles);
1655 gammas.push(gamma);
1656 }
1657 let row_refs: Vec<&[u8]> = rows.iter().map(Vec::as_slice).collect();
1658 let mut block = Vec::new();
1659 tq_pack_block(&row_refs, &gammas, codec.padded_dim(), &mut block);
1660 assert_eq!(block.len(), tq_block_bytes(codec.code_size()));
1661
1662 let mut scores = [0.0f32; TQ_BLOCK_LANES];
1663 tq_score_block(&plan, &block, &mut scores);
1664 for lane in 0..lanes {
1665 let expected = plan.estimate_row(&rows[lane], gammas[lane]);
1666 assert!(
1668 (scores[lane] - expected).abs() < 0.02,
1669 "lane {lane}: block score {} vs row estimate {expected}",
1670 scores[lane]
1671 );
1672 }
1673 }
1674
1675 #[test]
1676 fn simd_accumulation_matches_scalar_reference_exactly() {
1677 let padded_dim = 192; let mut state = 99u64;
1679 let mut base_lut = vec![0i8; padded_dim * 16];
1680 let mut qjl_lut = vec![0i8; padded_dim * 16];
1681 for value in base_lut.iter_mut().chain(qjl_lut.iter_mut()) {
1682 *value = (splitmix64(&mut state) as i32 % 255 - 127) as i8;
1683 }
1684 let mut nibble_bytes = vec![0u8; padded_dim * 8];
1685 for byte in nibble_bytes.iter_mut() {
1686 *byte = splitmix64(&mut state) as u8;
1687 }
1688
1689 let mut base_reference = [0i32; TQ_BLOCK_LANES];
1690 let mut qjl_reference = [0i32; TQ_BLOCK_LANES];
1691 lut16::accumulate_block_scalar(
1692 &base_lut,
1693 &qjl_lut,
1694 &nibble_bytes,
1695 padded_dim,
1696 &mut base_reference,
1697 &mut qjl_reference,
1698 );
1699 let mut base_dispatch = [0i32; TQ_BLOCK_LANES];
1700 let mut qjl_dispatch = [0i32; TQ_BLOCK_LANES];
1701 lut16::accumulate_block(
1702 &base_lut,
1703 &qjl_lut,
1704 &nibble_bytes,
1705 padded_dim,
1706 &mut base_dispatch,
1707 &mut qjl_dispatch,
1708 );
1709 assert_eq!(
1710 base_reference, base_dispatch,
1711 "base sums must match exactly"
1712 );
1713 assert_eq!(qjl_reference, qjl_dispatch, "qjl sums must match exactly");
1714 }
1715
1716 #[test]
1717 fn epilogue_matches_scalar_reference_bit_exactly() {
1718 let mut state = 4242u64;
1719 for trial in 0..256 {
1720 let mut base = [0i32; TQ_BLOCK_LANES];
1721 let mut qjl = [0i32; TQ_BLOCK_LANES];
1722 let mut scale_bytes = Vec::with_capacity(TQ_BLOCK_LANES * 4);
1723 let mut gamma_bytes = Vec::with_capacity(TQ_BLOCK_LANES * 4);
1724 for lane in 0..TQ_BLOCK_LANES {
1725 base[lane] = (splitmix64(&mut state) % 40_000) as i32 - 20_000;
1726 qjl[lane] = (splitmix64(&mut state) % 40_000) as i32 - 20_000;
1727 let scale = (splitmix64(&mut state) % 10_000) as f32 / 10_000.0;
1728 let gamma = (splitmix64(&mut state) % 10_000) as f32 / 10_000.0;
1729 let scale = match lane {
1731 15 => 0.0,
1732 7 => -scale,
1733 _ => scale,
1734 };
1735 scale_bytes.extend_from_slice(&scale.to_le_bytes());
1736 gamma_bytes.extend_from_slice(&gamma.to_le_bytes());
1737 }
1738 let base_dequant = 1e-4 + (trial as f32) * 1e-6;
1739 let qjl_dequant = 3e-5 + (trial as f32) * 1e-7;
1740 let cluster_dot = (trial as f32) / 256.0 - 0.5;
1741
1742 let mut expected = [0.0f32; TQ_BLOCK_LANES];
1743 let mut actual = [0.0f32; TQ_BLOCK_LANES];
1744 lut16::finish_ivf_block_scalar(
1745 &scale_bytes,
1746 &gamma_bytes,
1747 &base,
1748 &qjl,
1749 base_dequant,
1750 qjl_dequant,
1751 cluster_dot,
1752 &mut expected,
1753 );
1754 lut16::finish_ivf_block(
1755 &scale_bytes,
1756 &gamma_bytes,
1757 &base,
1758 &qjl,
1759 base_dequant,
1760 qjl_dequant,
1761 cluster_dot,
1762 &mut actual,
1763 );
1764 assert_eq!(
1765 expected.map(f32::to_bits),
1766 actual.map(f32::to_bits),
1767 "IVF epilogue must be bit-identical to the scalar reference (trial {trial})"
1768 );
1769
1770 lut16::finish_block_scalar(
1771 &gamma_bytes,
1772 &base,
1773 &qjl,
1774 base_dequant,
1775 qjl_dequant,
1776 &mut expected,
1777 );
1778 lut16::finish_block(
1779 &gamma_bytes,
1780 &base,
1781 &qjl,
1782 base_dequant,
1783 qjl_dequant,
1784 &mut actual,
1785 );
1786 assert_eq!(
1787 expected.map(f32::to_bits),
1788 actual.map(f32::to_bits),
1789 "flat epilogue must be bit-identical to the scalar reference (trial {trial})"
1790 );
1791 }
1792 }
1793
1794 #[test]
1795 fn fingerprint_pins_codec_constants() {
1796 let codec = TqCodec::new(768);
1797 assert_eq!(codec.padded_dim(), 768, "codec v2 is padding-free");
1798 assert_eq!(codec.code_size(), 384);
1799 assert_ne!(codec.fingerprint(), 0);
1800 assert_eq!(codec.fingerprint(), TqCodec::new(768).fingerprint());
1801 assert_ne!(codec.fingerprint(), TqCodec::new(769).fingerprint());
1802 assert_eq!(
1806 codec.fingerprint(),
1807 GOLDEN_FINGERPRINT_768,
1808 "TQ fingerprint for dim 768 changed; existing segments would be \
1809 rejected. Bump TQ_CODEC_VERSION deliberately instead."
1810 );
1811 }
1812
1813 #[test]
1814 fn query_plan_identity_uses_exact_float_bits() {
1815 let codec = TqCodec::new(4);
1816 let query = [0.0, -0.0, 1.0, f32::from_bits(0x7fc0_0001)];
1817 let plan = TqQueryPlan::build(&codec, &query);
1818
1819 assert!(plan.matches_query(&query));
1820 assert!(!plan.matches_query(&query[..3]));
1821
1822 let mut different_zero = query;
1823 different_zero[1] = 0.0;
1824 assert!(!plan.matches_query(&different_zero));
1825
1826 let mut different_nan = query;
1827 different_nan[3] = f32::from_bits(0x7fc0_0002);
1828 assert!(!plan.matches_query(&different_nan));
1829 }
1830
1831 const GOLDEN_FINGERPRINT_768: u64 = 7026088428300072418;
1832
1833 #[cfg(feature = "native")]
1834 #[test]
1835 fn builder_packs_blocks_in_input_order() {
1836 let dim = 32;
1837 let codec = std::sync::Arc::new(TqCodec::new(dim));
1838 let mut builder = TqFlatBuilder::new(std::sync::Arc::clone(&codec));
1839 let count = 37; let labels: Vec<(u32, u16)> = (0..count).map(|i| (i as u32, (i % 3) as u16)).collect();
1841 let mut vectors = Vec::new();
1842 for i in 0..count {
1843 vectors.extend(seeded_unit_vector(dim, 7_000 + i as u64));
1844 }
1845 builder.add_batch(&labels, &vectors).unwrap();
1846 builder.finish();
1847 assert_eq!(builder.doc_ids.len(), count);
1848 assert_eq!(
1849 builder.codes.len(),
1850 tq_codes_column_len(count, codec.code_size())
1851 );
1852
1853 let query = seeded_unit_vector(dim, 1);
1855 let plan = TqQueryPlan::build(&codec, &query);
1856 let block_bytes = tq_block_bytes(codec.code_size());
1857 let mut scratch = TqEncodeScratch::default();
1858 let mut scores = [0.0f32; TQ_BLOCK_LANES];
1859 for index in 0..count {
1860 let block = &builder.codes[(index / TQ_BLOCK_LANES) * block_bytes..][..block_bytes];
1861 tq_score_block(&plan, block, &mut scores);
1862 let mut nibbles = vec![0u8; codec.padded_dim()];
1863 let gamma = codec.encode_into(
1864 &vectors[index * dim..(index + 1) * dim],
1865 &mut nibbles,
1866 &mut scratch,
1867 );
1868 let expected = plan.estimate_row(&nibbles, gamma);
1869 assert!(
1870 (scores[index % TQ_BLOCK_LANES] - expected).abs() < 0.02,
1871 "vector {index} scored {} expected {expected}",
1872 scores[index % TQ_BLOCK_LANES]
1873 );
1874 }
1875 }
1876}