Struct ColbertPacking 

pub struct ColbertPacking { /* private fields */ }

The ColBERT / PLAID residual layout, shared by ColBERTv2’s ResidualCodec, PLAID, fast-plaid, next-plaid and WARP.

quantize_residuals writes each dimension’s bucket bits MSB-first into the row, bit 0 of the bucket first. So within a byte, key k occupies bits 7 - k·nbits down to 8 - (k+1)·nbits, and the bucket index is that nbits-wide segment with its bits reversed. The original decoders express this as a byte_reversed_bits_map followed by a group split; this is the same function written directly.

Implementations

impl ColbertPacking

pub fn new(nbits: usize) -> Result<Self, Error>

nbits must be 1, 2, 4 or 8.

Examples found in repository

examples/bench.rs (line 69)

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"
        }
    );
}

Trait Implementations

impl Clone for ColbertPacking

fn clone(&self) -> ColbertPacking

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Copy for ColbertPacking

impl Debug for ColbertPacking

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Eq for ColbertPacking

impl Packing for ColbertPacking

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.

fn keys_per_byte(&self) -> usize

8 / nbits.

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.

impl PartialEq for ColbertPacking

fn eq(&self, other: &ColbertPacking) -> bool

Equality operator ==.

fn ne(&self, other: &Rhs) -> bool

Inequality operator !=.

impl StructuralPartialEq for ColbertPacking