Trait Packing 

pub trait Packing {
    // Required methods
    fn nbits(&self) -> usize;
    fn bucket_index(&self, byte: u8, key: usize) -> usize;

    // Provided methods
    fn keys_per_byte(&self) -> usize { ... }
    fn pack_row(&self, buckets: &[usize], out: &mut [u8]) { ... }
}

Describes a bit-packing layout of nbits-wide bucket indices.

Key k of a byte is the k-th embedding dimension that byte carries, in dimension order: a token’s packed row bytes[0..dim·nbits/8] holds dims i·keys_per_byte + k at (bytes[i], key k).

Required Methods

fn nbits(&self) -> usize

Code width: 1, 2, 4 or 8.

fn bucket_index(&self, byte: u8, key: usize) -> usize

Bucket index (0 .. 2^nbits) stored at key position key (0 .. 8/nbits) of byte.

Provided Methods

fn keys_per_byte(&self) -> usize

8 / nbits.

Examples found in repository

examples/bench.rs (line 79)

fn main() {
    let args: Vec<usize> = std::env::args()
        .skip(1)
        .map(|a| a.parse().expect("integer arg"))
        .collect();
    let dim = args.first().copied().unwrap_or(128);
    let nbits = args.get(1).copied().unwrap_or(4);
    let nq = args.get(2).copied().unwrap_or(32);
    let ntok = args.get(3).copied().unwrap_or(240);
    let ndocs = args.get(4).copied().unwrap_or(1024);
    let ncent = 16_384;
    let reps = 9;
    let mut rng = Rng(0x9E3779B97F4A7C15);

    let p = ColbertPacking::new(nbits).unwrap();
    let nb = 1usize << nbits;
    let mut w: Vec<f32> = (0..nb).map(|_| rng.f32(-0.4, 0.4)).collect();
    w.sort_by(|a, b| a.total_cmp(b));
    let lut = Lut::new(&p, &w).unwrap();

    let query: Vec<f32> = (0..nq * dim).map(|_| rng.f32(-1.0, 1.0)).collect();
    let q = PreparedQuery::new(&lut, &query, nq, dim).unwrap();
    let cdot: Vec<f32> = (0..ncent * nq).map(|_| rng.f32(-1.0, 1.0)).collect();

    let pdim = dim / p.keys_per_byte();
    let mut packed = vec![0u8; ndocs * ntok * pdim];
    for b in packed.iter_mut() {
        *b = (rng.next() >> 56) as u8;
    }
    let codes: Vec<u32> = (0..ndocs * ntok)
        .map(|_| (rng.next() % ncent as u64) as u32)
        .collect();
    let inv: Vec<f32> = (0..ndocs * ntok).map(|_| rng.f32(0.8, 1.2)).collect();
    let docs: Vec<DocView> = (0..ndocs)
        .map(|d| {
            DocView::new(&packed[d * ntok * pdim..(d + 1) * ntok * pdim], ntok, pdim)
                .codes(Codes::U32(&codes[d * ntok..(d + 1) * ntok]))
                .inv_norms(&inv[d * ntok..(d + 1) * ntok])
        })
        .collect();

    // One arm per executable kernel, plus the scalar reference. `dispatch`
    // is what an unpinned host would get, named so the calibrated choice is
    // visible next to the kernels it chose between.
    let dispatched = lut.kernel(dim);
    let mut arms: Vec<Arm> = Vec::new();
    for &k in supported_kernels() {
        arms.push(Arm {
            label: if k == dispatched {
                format!("{k} (dispatched)")
            } else {
                format!("{k}")
            },
            lut: lut.clone().pin_kernel(Some(k)),
            ns: Vec::new(),
            checksum: 0.0,
        });
    }
    if !dispatched.is_simd() {
        arms.push(Arm {
            label: format!("{dispatched} (dispatched)"),
            lut: lut.clone(),
            ns: Vec::new(),
            checksum: 0.0,
        });
    }
    arms.push(Arm {
        label: "scalar reference".to_string(),
        lut: lut.clone().force_scalar(true),
        ns: Vec::new(),
        checksum: 0.0,
    });

    println!(
        "dim {dim}, nbits {nbits}, {nq} query tokens, {ndocs} docs × {ntok} tokens, {ncent} centroids\narch {}, dispatched kernel: {dispatched}, {reps} interleaved rounds",
        std::env::consts::ARCH,
    );

    let mut out = vec![0.0f32; ndocs];
    for arm in arms.iter_mut() {
        let s = Scorer::new(&arm.lut, &q)
            .with_centroid_term(&cdot, ncent)
            .unwrap();
        s.score_many(docs.iter().copied(), &mut out); // warm caches and branch predictors
    }
    for _ in 0..reps {
        for arm in arms.iter_mut() {
            let s = Scorer::new(&arm.lut, &q)
                .with_centroid_term(&cdot, ncent)
                .unwrap();
            let t = Instant::now();
            s.score_many(docs.iter().copied(), &mut out);
            arm.ns.push(t.elapsed().as_nanos() as f64 / (ndocs * ntok) as f64);
            arm.checksum = out.iter().map(|&v| v as f64).sum();
        }
    }

    let reference = arms.last().expect("at least the scalar arm");
    let (slow, want) = {
        let mut v = reference.ns.clone();
        v.sort_by(f64::total_cmp);
        (v[0], reference.checksum)
    };
    let mut worst_spread = 0.0f64;
    for arm in &arms {
        let mut v = arm.ns.clone();
        v.sort_by(f64::total_cmp);
        let (best, med) = (v[0], median(&v));
        let spread = (med - best) / best;
        worst_spread = worst_spread.max(spread);
        assert_eq!(
            arm.checksum.to_bits(),
            want.to_bits(),
            "{}: checksum {} differs from the scalar reference {want}",
            arm.label,
            arm.checksum
        );
        println!(
            "{:>26}: {best:7.2} ns/token  (median {med:7.2}, {:5.1} µs/doc, {:5.2}x scalar)",
            arm.label,
            best * ntok as f64 / 1e3,
            slow / best,
        );
    }
    println!(
        "{:>26}: {:.1}% median-vs-best spread — {}",
        "noise",
        worst_spread * 100.0,
        if worst_spread < 0.10 {
            "quiet enough to compare kernels"
        } else {
            "TOO NOISY, differences under ~2x are not real; free the machine and rerun"
        }
    );
}

fn pack_row(&self, buckets: &[usize], out: &mut [u8])

Reference packer, the inverse of Packing::bucket_index, for hosts that want to produce rows the same way the tests do. Not optimised.

Dyn Compatibility

This trait is dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors

impl Packing for ColbertPacking