Struct Lut 

pub struct Lut { /* private fields */ }

The document-side lookup state for one residual codec: a table turning each packed residual byte directly into its 8/nbits int8 bucket weights, plus the dequantisation scale.

Build once per index (it depends only on the bucket weights and the packing), share across threads.

Implementations

impl Lut

pub fn new<P: Packing>( packing: &P, bucket_weights: &[f32], ) -> Result<Self, Error>

Build the table from a packing layout and the codec’s 2^nbits bucket weights (f32, in bucket-index order).

Weights are quantised symmetrically to int8 with scale = max|w| / 127.

Examples found in repository

examples/bench.rs (line 73)

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

pub fn colbert(nbits: usize, bucket_weights: &[f32]) -> Result<Self, Error>

Convenience for the ColBERT / PLAID layout.

Examples found in repository

examples/real_index.rs (line 170)

fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();
    if args.len() < 2 {
        eprintln!("usage: real_index <index_dir> <queries.npy> [n_queries=20] [ref_docs=200]");
        std::process::exit(2);
    }
    let dir = Path::new(&args[0]);
    let n_queries: usize = args.get(2).map(|s| s.parse().unwrap()).unwrap_or(20);
    let ref_docs: usize = args.get(3).map(|s| s.parse().unwrap()).unwrap_or(200);

    let meta = fs::read_to_string(dir.join("metadata.json")).unwrap();
    let nbits = json_usize(&meta, "nbits");
    let num_chunks = json_usize(&meta, "num_chunks");
    let dim = json_usize(&meta, "embedding_dim");
    let (cshape, centroids) = npy_f32(&dir.join("centroids.npy"));
    let ncent = cshape[0];
    assert_eq!(cshape[1], dim);
    let (_, weights) = npy_f32(&dir.join("bucket_weights.npy"));
    assert_eq!(weights.len(), 1 << nbits);

    let mut codes: Vec<u32> = Vec::new();
    let mut inv: Vec<f32> = Vec::new();
    let mut packed: Vec<u8> = Vec::new();
    let mut doclens: Vec<usize> = Vec::new();
    let pdim = dim * nbits / 8;
    for c in 0..num_chunks {
        codes.extend(npy_codes(&dir.join(format!("{c}.codes.npy"))));
        inv.extend(npy_f32(&dir.join(format!("{c}.inv_norms.npy"))).1);
        let (rshape, r) = npy_u8(&dir.join(format!("{c}.residuals.npy")));
        assert_eq!(rshape[1], pdim, "residual row width");
        packed.extend(r);
        let dl = fs::read_to_string(dir.join(format!("doclens.{c}.json"))).unwrap();
        doclens.extend(
            dl.trim()
                .trim_matches(['[', ']'])
                .split(',')
                .map(|s| s.trim().parse::<usize>().unwrap()),
        );
    }
    let ntok_total: usize = doclens.iter().sum();
    assert_eq!(codes.len(), ntok_total);
    assert_eq!(inv.len(), ntok_total);
    assert_eq!(packed.len(), ntok_total * pdim);
    let mut offsets = Vec::with_capacity(doclens.len() + 1);
    offsets.push(0usize);
    for &l in &doclens {
        offsets.push(offsets.last().unwrap() + l);
    }

    let (qshape, queries) = npy_f32(Path::new(&args[1]));
    let (nq_rows, qdim) = (qshape[1], qshape[2]);
    assert_eq!(qdim, dim);
    let n_queries = n_queries.min(qshape[0]);

    let lut = Lut::colbert(nbits, &weights).unwrap();
    println!(
        "index: {} docs, {} tokens, dim {dim}, nbits {nbits}, {ncent} centroids, {} B/token packed\nkernel: {}\nqueries: {n_queries} × {nq_rows} rows",
        doclens.len(),
        ntok_total,
        pdim,
        lut.kernel(dim)
    );
    let max_w = weights.iter().fold(0f32, |m, &w| m.max(w.abs()));

    let mut t_cdot = 0.0f64;
    let mut t_score = 0.0f64;
    let mut max_abs_err = 0f32;
    let mut sum_rel_err = 0f64;
    let mut n_ref = 0usize;
    let mut skippable = 0usize;
    let mut skip_seen = 0usize;
    let mut bound_sum = 0f64;
    let mut resid_abs_sum = 0f64;
    let mut resid_n = 0usize;

    for qi in 0..n_queries {
        let qf = &queries[qi * nq_rows * dim..(qi + 1) * nq_rows * dim];
        let q = PreparedQuery::new(&lut, qf, nq_rows, dim).unwrap();

        // Stage-1 product the host would already have: centroid-major [ncent × nq].
        let t = Instant::now();
        let mut cdot = vec![0f32; ncent * nq_rows];
        for c in 0..ncent {
            let cv = &centroids[c * dim..(c + 1) * dim];
            for r in 0..nq_rows {
                let qr = &qf[r * dim..(r + 1) * dim];
                cdot[c * nq_rows + r] = cv.iter().zip(qr).map(|(a, b)| a * b).sum();
            }
        }
        t_cdot += t.elapsed().as_secs_f64();

        let scorer = Scorer::new(&lut, &q).with_centroid_term(&cdot, ncent).unwrap();
        let mut scores = vec![0f32; doclens.len()];
        let t = Instant::now();
        for (d, s) in scores.iter_mut().enumerate() {
            let (a, b) = (offsets[d], offsets[d + 1]);
            *s = scorer.score(
                DocView::new(&packed[a * pdim..b * pdim], b - a, pdim)
                    .codes(Codes::U32(&codes[a..b]))
                    .inv_norms(&inv[a..b]),
            );
        }
        t_score += t.elapsed().as_secs_f64();

        // Float reference on the first `ref_docs` docs: decompress token =
        // centroid + bucket weight per dim, exact f32 MaxSim with inv norms.
        let mut tok = vec![0f32; dim];
        for d in 0..ref_docs.min(doclens.len()) {
            let (a, b) = (offsets[d], offsets[d + 1]);
            let mut best = vec![f32::NEG_INFINITY; nq_rows];
            for t in a..b {
                let cv = &centroids[codes[t] as usize * dim..(codes[t] as usize + 1) * dim];
                let row = &packed[t * pdim..(t + 1) * pdim];
                for dd in 0..dim {
                    // ColBERT packing: key k of byte i is bits (7 - k*nbits ..), bucket bit-reversed.
                    let kpb = 8 / nbits;
                    let (i, k) = (dd / kpb, dd % kpb);
                    let shift = 8 - nbits * (k + 1);
                    let seg = (row[i] >> shift) as usize & ((1 << nbits) - 1);
                    let mut bucket = 0usize;
                    for bit in 0..nbits {
                        if seg & (1 << bit) != 0 {
                            bucket |= 1 << (nbits - 1 - bit);
                        }
                    }
                    tok[dd] = cv[dd] + weights[bucket];
                }
                for r in 0..nq_rows {
                    let qr = &qf[r * dim..(r + 1) * dim];
                    let s: f32 = tok.iter().zip(qr).map(|(x, y)| x * y).sum::<f32>() * inv[t];
                    if s > best[r] {
                        best[r] = s;
                    }
                }
            }
            let reference: f32 = best.iter().sum();
            let err = (scores[d] - reference).abs();
            max_abs_err = max_abs_err.max(err);
            sum_rel_err += (err / reference.abs().max(1e-6)) as f64;
            n_ref += 1;
        }

        // Exact skip bound: |residual term| <= sqw[r] * 127 * Σ|q̂| = max|w| * ||q_r||_1
        // (up to quantisation). Simulate a sequential pass over each doc's tokens and
        // count tokens where every row could have been skipped.
        let l1: Vec<f32> = (0..nq_rows)
            .map(|r| qf[r * dim..(r + 1) * dim].iter().map(|x| x.abs()).sum())
            .collect();
        let bounds: Vec<f32> = l1.iter().map(|&l| l * max_w).collect();
        bound_sum += bounds.iter().map(|&b| b as f64).sum::<f64>() / nq_rows as f64;
        for d in 0..ref_docs.min(doclens.len()) {
            let (a, b) = (offsets[d], offsets[d + 1]);
            let mut best = vec![f32::NEG_INFINITY; nq_rows];
            for t in a..b {
                let crow = &cdot[codes[t] as usize * nq_rows..(codes[t] as usize + 1) * nq_rows];
                let can_skip = (0..nq_rows).all(|r| (crow[r] + bounds[r]) * inv[t] <= best[r]);
                skip_seen += 1;
                if can_skip {
                    skippable += 1;
                }
                // Actual residual term magnitude for the record.
                let cv = &centroids[codes[t] as usize * dim..(codes[t] as usize + 1) * dim];
                let row = &packed[t * pdim..(t + 1) * pdim];
                let kpb = 8 / nbits;
                for r in 0..nq_rows {
                    let qr = &qf[r * dim..(r + 1) * dim];
                    let mut resid = 0f32;
                    for dd in 0..dim {
                        let (i, k) = (dd / kpb, dd % kpb);
                        let shift = 8 - nbits * (k + 1);
                        let seg = (row[i] >> shift) as usize & ((1 << nbits) - 1);
                        let mut bucket = 0usize;
                        for bit in 0..nbits {
                            if seg & (1 << bit) != 0 {
                                bucket |= 1 << (nbits - 1 - bit);
                            }
                        }
                        resid += qr[dd] * weights[bucket];
                    }
                    resid_abs_sum += resid.abs() as f64;
                    resid_n += 1;
                    let s =
                        (crow[r] + (cv.iter().zip(qr).map(|(x, y)| x * y).sum::<f32>() - crow[r]) + resid)
                            * inv[t];
                    if s > best[r] {
                        best[r] = s;
                    }
                }
            }
        }
    }

    let tokens_scored = ntok_total as f64 * n_queries as f64;
    println!(
        "\nexhaustive stage-2 over all docs: {:.2} ms/query, {:.2} ns/token (kernel only; stage-1 cdot GEMM excluded: {:.1} ms/query naive)",
        t_score / n_queries as f64 * 1e3,
        t_score / tokens_scored * 1e9,
        t_cdot / n_queries as f64 * 1e3
    );
    println!(
        "parity vs float decompression on {n_ref} (query, doc) pairs: max |Δscore| {max_abs_err:.4}, mean rel {:.2e}",
        sum_rel_err / n_ref.max(1) as f64
    );
    println!(
        "exact skip bound: mean bound {:.3} vs mean |residual term| {:.4}; tokens skippable {}/{} ({:.2}%)",
        bound_sum / n_queries as f64,
        resid_abs_sum / resid_n.max(1) as f64,
        skippable,
        skip_seen,
        100.0 * skippable as f64 / skip_seen.max(1) as f64
    );
}

pub fn force_scalar(self, yes: bool) -> Self

Pin every score to the scalar reference kernel. For tests and for measuring what the SIMD is worth; the results are bit-identical either way. The environment variable MAXSIM_LUT_FORCE_SCALAR=1 has the same effect process-wide.

Examples found in repository

examples/bench.rs (line 123)

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

pub fn pin_kernel(self, kernel: Option<Kernel>) -> Self

Pin dispatch to one kernel instead of the calibrated choice.

Only a kernel this CPU can execute is honoured (check with crate::supported_kernels); anything else is ignored and dispatch proceeds normally, because a kernel the CPU cannot run has no meaningful behaviour to fall back to. Scoring is bit-identical whichever kernel runs, so this only affects speed.

Use it to benchmark one path, or to make dispatch deterministic on a fleet of mixed cores. None restores the default.

Examples found in repository

examples/bench.rs (line 108)

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

pub fn nbits(&self) -> usize

Code width this table was built for.

pub fn keys_per_byte(&self) -> usize

8 / nbits: how many dims one packed byte carries.

pub fn scale(&self) -> f32

Dequantisation scale: fused as f32 · scale ≈ bucket_weight.

pub fn expand(&self, byte: u8) -> &[i8]

The 8/nbits int8 weights a packed byte expands to, in dim order.

pub fn fused_table(&self) -> &[i8]

The whole fused table, [256 · keys_per_byte], row b = byte b.

pub fn nibble_tables(&self) -> Option<&NibbleTables>

The nibble-factored tables, if the layout admits them (always, for crate::ColbertPacking at nbits 1, 2 or 4; never at nbits 8).

pub fn kernel(&self, dim: usize) -> Kernel

Which kernel crate::Scorer::score will run for this table and dim on this CPU. Print it next to any benchmark number: a speedup attributed to a path that never executed is the easiest measurement error to make and the hardest to notice.

Examples found in repository

examples/bench.rs (line 99)

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

examples/real_index.rs (line 176)

fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();
    if args.len() < 2 {
        eprintln!("usage: real_index <index_dir> <queries.npy> [n_queries=20] [ref_docs=200]");
        std::process::exit(2);
    }
    let dir = Path::new(&args[0]);
    let n_queries: usize = args.get(2).map(|s| s.parse().unwrap()).unwrap_or(20);
    let ref_docs: usize = args.get(3).map(|s| s.parse().unwrap()).unwrap_or(200);

    let meta = fs::read_to_string(dir.join("metadata.json")).unwrap();
    let nbits = json_usize(&meta, "nbits");
    let num_chunks = json_usize(&meta, "num_chunks");
    let dim = json_usize(&meta, "embedding_dim");
    let (cshape, centroids) = npy_f32(&dir.join("centroids.npy"));
    let ncent = cshape[0];
    assert_eq!(cshape[1], dim);
    let (_, weights) = npy_f32(&dir.join("bucket_weights.npy"));
    assert_eq!(weights.len(), 1 << nbits);

    let mut codes: Vec<u32> = Vec::new();
    let mut inv: Vec<f32> = Vec::new();
    let mut packed: Vec<u8> = Vec::new();
    let mut doclens: Vec<usize> = Vec::new();
    let pdim = dim * nbits / 8;
    for c in 0..num_chunks {
        codes.extend(npy_codes(&dir.join(format!("{c}.codes.npy"))));
        inv.extend(npy_f32(&dir.join(format!("{c}.inv_norms.npy"))).1);
        let (rshape, r) = npy_u8(&dir.join(format!("{c}.residuals.npy")));
        assert_eq!(rshape[1], pdim, "residual row width");
        packed.extend(r);
        let dl = fs::read_to_string(dir.join(format!("doclens.{c}.json"))).unwrap();
        doclens.extend(
            dl.trim()
                .trim_matches(['[', ']'])
                .split(',')
                .map(|s| s.trim().parse::<usize>().unwrap()),
        );
    }
    let ntok_total: usize = doclens.iter().sum();
    assert_eq!(codes.len(), ntok_total);
    assert_eq!(inv.len(), ntok_total);
    assert_eq!(packed.len(), ntok_total * pdim);
    let mut offsets = Vec::with_capacity(doclens.len() + 1);
    offsets.push(0usize);
    for &l in &doclens {
        offsets.push(offsets.last().unwrap() + l);
    }

    let (qshape, queries) = npy_f32(Path::new(&args[1]));
    let (nq_rows, qdim) = (qshape[1], qshape[2]);
    assert_eq!(qdim, dim);
    let n_queries = n_queries.min(qshape[0]);

    let lut = Lut::colbert(nbits, &weights).unwrap();
    println!(
        "index: {} docs, {} tokens, dim {dim}, nbits {nbits}, {ncent} centroids, {} B/token packed\nkernel: {}\nqueries: {n_queries} × {nq_rows} rows",
        doclens.len(),
        ntok_total,
        pdim,
        lut.kernel(dim)
    );
    let max_w = weights.iter().fold(0f32, |m, &w| m.max(w.abs()));

    let mut t_cdot = 0.0f64;
    let mut t_score = 0.0f64;
    let mut max_abs_err = 0f32;
    let mut sum_rel_err = 0f64;
    let mut n_ref = 0usize;
    let mut skippable = 0usize;
    let mut skip_seen = 0usize;
    let mut bound_sum = 0f64;
    let mut resid_abs_sum = 0f64;
    let mut resid_n = 0usize;

    for qi in 0..n_queries {
        let qf = &queries[qi * nq_rows * dim..(qi + 1) * nq_rows * dim];
        let q = PreparedQuery::new(&lut, qf, nq_rows, dim).unwrap();

        // Stage-1 product the host would already have: centroid-major [ncent × nq].
        let t = Instant::now();
        let mut cdot = vec![0f32; ncent * nq_rows];
        for c in 0..ncent {
            let cv = &centroids[c * dim..(c + 1) * dim];
            for r in 0..nq_rows {
                let qr = &qf[r * dim..(r + 1) * dim];
                cdot[c * nq_rows + r] = cv.iter().zip(qr).map(|(a, b)| a * b).sum();
            }
        }
        t_cdot += t.elapsed().as_secs_f64();

        let scorer = Scorer::new(&lut, &q).with_centroid_term(&cdot, ncent).unwrap();
        let mut scores = vec![0f32; doclens.len()];
        let t = Instant::now();
        for (d, s) in scores.iter_mut().enumerate() {
            let (a, b) = (offsets[d], offsets[d + 1]);
            *s = scorer.score(
                DocView::new(&packed[a * pdim..b * pdim], b - a, pdim)
                    .codes(Codes::U32(&codes[a..b]))
                    .inv_norms(&inv[a..b]),
            );
        }
        t_score += t.elapsed().as_secs_f64();

        // Float reference on the first `ref_docs` docs: decompress token =
        // centroid + bucket weight per dim, exact f32 MaxSim with inv norms.
        let mut tok = vec![0f32; dim];
        for d in 0..ref_docs.min(doclens.len()) {
            let (a, b) = (offsets[d], offsets[d + 1]);
            let mut best = vec![f32::NEG_INFINITY; nq_rows];
            for t in a..b {
                let cv = &centroids[codes[t] as usize * dim..(codes[t] as usize + 1) * dim];
                let row = &packed[t * pdim..(t + 1) * pdim];
                for dd in 0..dim {
                    // ColBERT packing: key k of byte i is bits (7 - k*nbits ..), bucket bit-reversed.
                    let kpb = 8 / nbits;
                    let (i, k) = (dd / kpb, dd % kpb);
                    let shift = 8 - nbits * (k + 1);
                    let seg = (row[i] >> shift) as usize & ((1 << nbits) - 1);
                    let mut bucket = 0usize;
                    for bit in 0..nbits {
                        if seg & (1 << bit) != 0 {
                            bucket |= 1 << (nbits - 1 - bit);
                        }
                    }
                    tok[dd] = cv[dd] + weights[bucket];
                }
                for r in 0..nq_rows {
                    let qr = &qf[r * dim..(r + 1) * dim];
                    let s: f32 = tok.iter().zip(qr).map(|(x, y)| x * y).sum::<f32>() * inv[t];
                    if s > best[r] {
                        best[r] = s;
                    }
                }
            }
            let reference: f32 = best.iter().sum();
            let err = (scores[d] - reference).abs();
            max_abs_err = max_abs_err.max(err);
            sum_rel_err += (err / reference.abs().max(1e-6)) as f64;
            n_ref += 1;
        }

        // Exact skip bound: |residual term| <= sqw[r] * 127 * Σ|q̂| = max|w| * ||q_r||_1
        // (up to quantisation). Simulate a sequential pass over each doc's tokens and
        // count tokens where every row could have been skipped.
        let l1: Vec<f32> = (0..nq_rows)
            .map(|r| qf[r * dim..(r + 1) * dim].iter().map(|x| x.abs()).sum())
            .collect();
        let bounds: Vec<f32> = l1.iter().map(|&l| l * max_w).collect();
        bound_sum += bounds.iter().map(|&b| b as f64).sum::<f64>() / nq_rows as f64;
        for d in 0..ref_docs.min(doclens.len()) {
            let (a, b) = (offsets[d], offsets[d + 1]);
            let mut best = vec![f32::NEG_INFINITY; nq_rows];
            for t in a..b {
                let crow = &cdot[codes[t] as usize * nq_rows..(codes[t] as usize + 1) * nq_rows];
                let can_skip = (0..nq_rows).all(|r| (crow[r] + bounds[r]) * inv[t] <= best[r]);
                skip_seen += 1;
                if can_skip {
                    skippable += 1;
                }
                // Actual residual term magnitude for the record.
                let cv = &centroids[codes[t] as usize * dim..(codes[t] as usize + 1) * dim];
                let row = &packed[t * pdim..(t + 1) * pdim];
                let kpb = 8 / nbits;
                for r in 0..nq_rows {
                    let qr = &qf[r * dim..(r + 1) * dim];
                    let mut resid = 0f32;
                    for dd in 0..dim {
                        let (i, k) = (dd / kpb, dd % kpb);
                        let shift = 8 - nbits * (k + 1);
                        let seg = (row[i] >> shift) as usize & ((1 << nbits) - 1);
                        let mut bucket = 0usize;
                        for bit in 0..nbits {
                            if seg & (1 << bit) != 0 {
                                bucket |= 1 << (nbits - 1 - bit);
                            }
                        }
                        resid += qr[dd] * weights[bucket];
                    }
                    resid_abs_sum += resid.abs() as f64;
                    resid_n += 1;
                    let s =
                        (crow[r] + (cv.iter().zip(qr).map(|(x, y)| x * y).sum::<f32>() - crow[r]) + resid)
                            * inv[t];
                    if s > best[r] {
                        best[r] = s;
                    }
                }
            }
        }
    }

    let tokens_scored = ntok_total as f64 * n_queries as f64;
    println!(
        "\nexhaustive stage-2 over all docs: {:.2} ms/query, {:.2} ns/token (kernel only; stage-1 cdot GEMM excluded: {:.1} ms/query naive)",
        t_score / n_queries as f64 * 1e3,
        t_score / tokens_scored * 1e9,
        t_cdot / n_queries as f64 * 1e3
    );
    println!(
        "parity vs float decompression on {n_ref} (query, doc) pairs: max |Δscore| {max_abs_err:.4}, mean rel {:.2e}",
        sum_rel_err / n_ref.max(1) as f64
    );
    println!(
        "exact skip bound: mean bound {:.3} vs mean |residual term| {:.4}; tokens skippable {}/{} ({:.2}%)",
        bound_sum / n_queries as f64,
        resid_abs_sum / resid_n.max(1) as f64,
        skippable,
        skip_seen,
        100.0 * skippable as f64 / skip_seen.max(1) as f64
    );
}

Trait Implementations

impl Clone for Lut

fn clone(&self) -> Lut

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl Debug for Lut

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

Formats the value using the given formatter.