Skip to content

Commit a131968

Browse files
committed
perf: disable SAQ for rabitq4
Signed-off-by: usamoi <usamoi@outlook.com>
1 parent b3e306c commit a131968

1 file changed

Lines changed: 12 additions & 5 deletions

File tree

crates/rabitq/src/extended.rs

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -119,8 +119,10 @@ pub fn ugly_code<const BITS: usize>(vector: &[f32]) -> Code {
119119
};
120120
let mut code = Vec::with_capacity(n as _);
121121
for i in 0..n {
122-
let v = scale * normalized_vector[i];
123-
let v = v + v.signum() * delta[i] as f32;
122+
let mut v = scale * normalized_vector[i];
123+
if let Some(delta) = delta.as_ref() {
124+
v = v + v.signum() * delta[i] as f32;
125+
}
124126
let c = v.floor().clamp(min as f32, max as f32) as i32;
125127
code.push((c + (1 << (BITS - 1))) as _);
126128
}
@@ -251,16 +253,21 @@ fn find_scale<const B: usize>(o: &[f32]) -> f32 {
251253
x_m as f32 + f32::EPSILON
252254
}
253255

254-
fn ugly_find_scale<const B: usize>(o: &[f32]) -> (f32, Vec<i32>) {
256+
fn ugly_find_scale<const B: usize>(o: &[f32]) -> (f32, Option<Vec<i32>>) {
255257
assert!((1..=8).contains(&B));
256258

259+
let scale = (1 << (B - 1)) as f32 / f32::reduce_min_max_of_x(o).1;
260+
261+
if B >= 8 {
262+
return (scale, None);
263+
}
264+
257265
let dim = o.len();
258266

259267
let mut code = Vec::<u8>::with_capacity(dim);
260268
let mut numerator_m = 0.0f64;
261269
let mut sqr_denominator_m = 0.0f64;
262270

263-
let scale = (1 << (B - 1)) as f32 / f32::reduce_min_max_of_x(o).1;
264271
for i in 0..dim {
265272
code.push((o[i] as f64 * scale as f64) as u8);
266273
let value = code[i] as f64 + 0.5;
@@ -299,7 +306,7 @@ fn ugly_find_scale<const B: usize>(o: &[f32]) -> (f32, Vec<i32>) {
299306
}
300307
}
301308

302-
(scale, delta)
309+
(scale, Some(delta))
303310
}
304311

305312
pub fn pack_code<const BITS: usize>(input: &[u8]) -> [Vec<u64>; BITS] {

0 commit comments

Comments
 (0)