Struct Scorer 

pub struct Scorer<'a> { /* private fields */ }

A Lut and a PreparedQuery bound together, with the host’s optional centroid term, ready to score documents.

Copy, Send + Sync: make one per query and share it across the threads scoring that query’s candidates.

Implementations

impl<'a> Scorer<'a>

pub fn new(lut: &'a Lut, query: &'a PreparedQuery) -> Self

Bind a table and a query. Panics if the query was prepared against a different table (use Scorer::try_new to get an error instead).

Examples found in repository

examples/bench.rs (line 135)

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 207)

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 try_new(lut: &'a Lut, query: &'a PreparedQuery) -> Result<Self, Error>

Bind a table and a query.

pub fn with_centroid_term( self, cdot_centroid_major: &'a [f32], num_centroids: usize, ) -> Result<Self, Error>

Supply the host’s query × centroid scores, centroid-major: cdot[cid · n_query_tokens + q]. One centroid’s scores across all query rows are then contiguous, so the vectorised fold loads them as one vector. Hosts that hold the [n_query_tokens, num_centroids] orientation transpose once per query.

Examples found in repository

examples/bench.rs (line 136)

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 207)

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 lut(&self) -> &'a Lut

The table this scorer uses.

pub fn query(&self) -> &'a PreparedQuery

The query this scorer uses.

pub fn score(&self, doc: DocView<'_>) -> f32

MaxSim of the query against one document. Panics on a shape violation (see Scorer::try_score for the checked form); every SIMD path returns the same bits as the scalar reference.

Examples found in repository

examples/real_index.rs (lines 212-216)

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 try_score(&self, doc: DocView<'_>) -> Result<f32, Error>

MaxSim of the query against one document, with shape validation.

pub fn score_many<'d, I>(&self, docs: I, out: &mut [f32])

where I: IntoIterator<Item = DocView<'d>>,

Score many documents into out (out.len() must equal the number of documents). Sequential; wrap the call in the host’s parallel iterator over chunks to fan out.

Panics if the counts disagree in either direction. Silently dropping the tail of a candidate list is the kind of bug that shows up as a slightly worse recall number months later, not as a failure.

Examples found in repository

examples/bench.rs (line 138)

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<'a> Clone for Scorer<'a>

fn clone(&self) -> Scorer<'a>

Returns a duplicate of the value.

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

Performs copy-assignment from source.

impl<'a> Copy for Scorer<'a>

impl<'a> Debug for Scorer<'a>

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

Formats the value using the given formatter.