qdrant/lib/quantization/benches/p_square.rs

use std::hint::black_box;

use criterion::{Criterion, criterion_group, criterion_main};
use quantization::p_square::P2Quantile;
use quantization::quantile::find_quantile_interval_per_coordinate_with_preprocess;
use rand::RngExt;
use rand::rngs::SmallRng;

fn p_square(c: &mut Criterion) {
    let mut group = c.benchmark_group("p_square");

    let count = 10_000;
    let mut rng = rand::make_rng::<SmallRng>();
    let data = (0..count)
        .map(|_| rng.random::<f64>())
        .collect::<Vec<f64>>();
    let quantile = 0.99;

    group.bench_function("p_square_5", |b| {
        b.iter(|| {
            let mut p2 = P2Quantile::<5>::new(quantile).unwrap();
            for &x in &data {
                p2.push(x);
            }
            black_box(p2.estimate());
        });
    });

    group.bench_function("p_square_7", |b| {
        b.iter(|| {
            let mut p2 = P2Quantile::<7>::new(quantile).unwrap();
            for &x in &data {
                p2.push(x);
            }
            black_box(p2.estimate());
        });
    });

    group.bench_function("p_square_9", |b| {
        b.iter(|| {
            let mut p2 = P2Quantile::<9>::new(quantile).unwrap();
            for &x in &data {
                p2.push(x);
            }
            black_box(p2.estimate());
        });
    });
}

fn p_square_vectors(c: &mut Criterion) {
    let mut group = c.benchmark_group("p_square_vectors");

    let count = 10_000;
    let dim = 1536;
    let mut rng = rand::make_rng::<SmallRng>();
    let data = (0..count)
        .map(|_| (0..dim).map(|_| rng.random::<f32>()).collect::<Vec<f32>>())
        .collect::<Vec<Vec<f32>>>();
    let quantile = 0.99;

    group.bench_function("p_square_vectors", |b| {
        b.iter(|| {
            black_box(
                find_quantile_interval_per_coordinate_with_preprocess(
                    data.iter(),
                    dim,
                    dim,
                    count,
                    quantile,
                    4,
                    8_192,
                    |raw, scratch| {
                        for (s, &v) in scratch.iter_mut().zip(raw.iter()) {
                            *s = f64::from(v);
                        }
                    },
                    &std::sync::atomic::AtomicBool::new(false),
                )
                .unwrap(),
            );
        });
    });
}

criterion_group! {
    name = benches;
    config = Criterion::default().sample_size(10);
    targets = p_square, p_square_vectors,
}

criterion_main!(benches);