turbovec 0.4.1

Fast vector quantization with 2-4 bit compression and SIMD search
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57

//! Dump the rotation matrix and Lloyd-Max codebook for a given (dim, bits) to
//! a raw little-endian f32 file the Python POC can load with np.fromfile.
//!
//! Used by benchmarks/rabitq_poc/poc_apples_to_apples.py to consume Rust's
//! exact rotation and codebook so the Python pipeline matches the Rust
//! pipeline byte-for-byte at the parameter level.
//!
//! Layout per (dim, bits) tuple, written to `<out>/state_d{dim}_b{bits}.bin`:
//!   - rotation matrix:   dim * dim   f32, row-major
//!   - boundaries:        2^bits - 1  f32
//!   - centroids:         2^bits      f32

use std::env;
use std::fs::{create_dir_all, File};
use std::io::{BufWriter, Write};

extern crate blas_src;

use turbovec::codebook::codebook;
use turbovec::rotation::make_rotation_matrix;

fn write_f32s<W: Write>(w: &mut W, vals: &[f32]) -> std::io::Result<()> {
    for &v in vals {
        w.write_all(&v.to_le_bytes())?;
    }
    Ok(())
}

fn main() {
    let args: Vec<String> = env::args().collect();
    if args.len() < 2 {
        eprintln!("usage: dump_state <out_dir>");
        std::process::exit(1);
    }
    let out_dir = &args[1];
    create_dir_all(out_dir).unwrap();

    let configs = [(200usize, 2usize), (200, 4), (1536, 2), (1536, 4), (3072, 2), (3072, 4)];

    for &(dim, bits) in &configs {
        eprintln!("Dumping d={} bits={}", dim, bits);
        let rotation = make_rotation_matrix(dim);
        let (boundaries, centroids) = codebook(bits, dim);

        assert_eq!(rotation.len(), dim * dim);
        assert_eq!(boundaries.len(), (1 << bits) - 1);
        assert_eq!(centroids.len(), 1 << bits);

        let path = format!("{}/state_d{}_b{}.bin", out_dir, dim, bits);
        let mut w = BufWriter::new(File::create(&path).unwrap());
        write_f32s(&mut w, &rotation).unwrap();
        write_f32s(&mut w, &boundaries).unwrap();
        write_f32s(&mut w, &centroids).unwrap();
        w.flush().unwrap();
        eprintln!("  -> {}", path);
    }
}