Benchmark of quad search vs binary search

benchmark.rs

#![feature(portable_simd)]

use std::hint::black_box;

use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};

use std::simd::cmp::SimdPartialEq;
use std::simd::cmp::SimdPartialOrd;
use std::simd::u16x8;

pub fn quad_search(slice: &[u16], val: u16) -> Result<usize, usize> {
    const GAP: usize = u16x8::LEN * 2;

    let (chunks, remaining) = slice.as_chunks::<GAP>();

    if chunks.is_empty() {
        return match remaining
            .iter()
            .copied()
            .enumerate()
            .find(|(_, v)| *v >= val)
        {
            Some((i, v)) if v == val => Ok(i),
            Some((i, _)) => Err(i),
            None => Err(slice.len()),
        };
    }

    let num_blocks = chunks.len();
    let mut base = 0;
    let mut n = num_blocks;
    while n > 3 {
        let quarter = n >> 2; // equivalent to n / 4

        let k1 = chunks[base + quarter][GAP - 1];
        let k2 = chunks[base + 2 * quarter][GAP - 1];
        let k3 = chunks[base + 3 * quarter][GAP - 1];

        let c1 = (k1 < val) as usize;
        let c2 = (k2 < val) as usize;
        let c3 = (k3 < val) as usize;

        base += (c1 + c2 + c3) * quarter;
        n -= 3 * quarter;
    }

    while n > 1 {
        let half = n >> 1; // equivalent to n / 2
        base = if chunks[base + half][GAP - 1] < val {
            base + half
        } else {
            base
        };
        n -= half;
    }

    let lo = if chunks[base][GAP - 1] < val {
        base + 1
    } else {
        base
    };

    if lo < num_blocks {
        let ndl = u16x8::splat(val);
        // I would love to work with arrays here...
        let v0 = u16x8::from_slice(&chunks[lo][..GAP / 2]);
        let v1 = u16x8::from_slice(&chunks[lo][GAP / 2..]);
        let base_index = lo * GAP;
        return match (v0.simd_ge(ndl).first_set(), v1.simd_ge(ndl).first_set()) {
            (Some(i), _) if v0[i] == val => Ok(base_index + i),
            (Some(i), _) => Err(base_index + i),
            (_, Some(i)) if v1[i] == val => Ok(base_index + GAP / 2 + i),
            (_, Some(i)) => Err(base_index + GAP / 2 + i),
            (None, None) => Err(slice.len()),
        };
    }

    match slice
        .iter()
        .copied()
        .enumerate()
        .skip(num_blocks * GAP)
        .find(|(_, v)| *v >= val)
    {
        Some((i, v)) if v == val => Ok(i),
        Some((i, _)) => Err(i),
        None => Err(slice.len()),
    }
}

pub fn quad_contains(slice: &[u16], val: u16) -> bool {
    const GAP: usize = u16x8::LEN * 2;

    let (chunks, remaining) = slice.as_chunks::<GAP>();

    if chunks.is_empty() {
        return match remaining.iter().copied().find(|v| *v >= val) {
            Some(v) => v == val,
            None => false,
        };
    }

    let num_blocks = chunks.len();
    let mut base = 0;
    let mut n = num_blocks;
    while n > 3 {
        let quarter = n >> 2; // equivalent to n / 4

        let k1 = chunks[base + quarter][GAP - 1];
        let k2 = chunks[base + 2 * quarter][GAP - 1];
        let k3 = chunks[base + 3 * quarter][GAP - 1];

        let c1 = (k1 < val) as usize;
        let c2 = (k2 < val) as usize;
        let c3 = (k3 < val) as usize;

        base += (c1 + c2 + c3) * quarter;
        n -= 3 * quarter;
    }

    while n > 1 {
        let half = n >> 1; // equivalent to n / 2
        base = if chunks[base + half][GAP - 1] < val {
            base + half
        } else {
            base
        };
        n -= half;
    }

    let lo = if chunks[base][GAP - 1] < val {
        base + 1
    } else {
        base
    };

    if lo < num_blocks {
        let ndl = u16x8::splat(val);
        // I would love to work with arrays here...
        let v0 = u16x8::from_slice(&chunks[lo][..GAP / 2]);
        let v1 = u16x8::from_slice(&chunks[lo][GAP / 2..]);
        return (v0.simd_eq(ndl) | v1.simd_eq(ndl)).any();
    }

    match slice
        .iter()
        .copied()
        .skip(num_blocks * GAP)
        .find(|v| *v >= val)
    {
        Some(v) => v == val,
        None => false,
    }
}

pub fn quad_contains_no_bounds_checks(slice: &[u16], val: u16) -> bool {
    const GAP: usize = u16x8::LEN * 2;

    let (chunks, remaining) = slice.as_chunks::<GAP>();

    if chunks.is_empty() {
        return match remaining.iter().copied().find(|v| *v >= val) {
            Some(v) => v == val,
            None => false,
        };
    }

    unsafe {
        let num_blocks = chunks.len();
        let mut base = 0;
        let mut n = num_blocks;
        while n > 3 {
            let quarter = n >> 2; // equivalent to n / 4

            let k1 = chunks.get_unchecked(base + quarter)[GAP - 1];
            let k2 = chunks.get_unchecked(base + 2 * quarter)[GAP - 1];
            let k3 = chunks.get_unchecked(base + 3 * quarter)[GAP - 1];

            let c1 = (k1 < val) as usize;
            let c2 = (k2 < val) as usize;
            let c3 = (k3 < val) as usize;

            base += (c1 + c2 + c3) * quarter;
            n -= 3 * quarter;
        }

        while n > 1 {
            let half = n >> 1; // equivalent to n / 2
            base = if chunks.get_unchecked(base + half)[GAP - 1] < val {
                base + half
            } else {
                base
            };
            n -= half;
        }

        let lo = if chunks.get_unchecked(base)[GAP - 1] < val {
            base + 1
        } else {
            base
        };

        if lo < num_blocks {
            let ndl = u16x8::splat(val);
            let chunk = chunks.get_unchecked(lo);
            // The compiler knows this can't panic
            let v0 = u16x8::from_slice(&chunk[..GAP / 2]);
            let v1 = u16x8::from_slice(&chunk[GAP / 2..]);
            return (v0.simd_eq(ndl) | v1.simd_eq(ndl)).any();
        }

        match slice
            .iter()
            .copied()
            .skip(num_blocks * GAP)
            .find(|v| *v >= val)
        {
            Some(v) => v == val,
            None => false,
        }
    }
}

fn sorted_slice(len: usize) -> Vec<u16> {
    let step = (u16::MAX as usize / len.max(1)).max(1);
    let mut values = Vec::with_capacity(len);
    for i in 0..len {
        values.push((i * step) as u16);
    }
    values
}

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

    for len in [8usize, 16, 64, 256, 1024, 4096] {
        let slice = sorted_slice(len);
        let target = slice[len / 2];
        group.throughput(Throughput::Elements(len as u64));

        group.bench_with_input(
            BenchmarkId::new("quad_search", len),
            &slice[..],
            |b, slice| {
                b.iter(|| quad_search(black_box(slice), black_box(target)));
            },
        );
        group.bench_with_input(
            BenchmarkId::new("binary_search", len),
            &slice[..],
            |b, slice| {
                let target = black_box(target);
                b.iter(|| black_box(slice).partition_point(|&v| v <= target));
            },
        );
    }

    group.finish();
}

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

    for len in [8usize, 16, 64, 256, 1024, 4096] {
        let slice = sorted_slice(len);
        let target = slice[len / 2];
        group.throughput(Throughput::Elements(len as u64));
        group.bench_with_input(
            BenchmarkId::new("quad_contains", len),
            &slice[..],
            |b, slice| {
                b.iter(|| quad_contains(black_box(slice), black_box(target)));
            },
        );
        group.bench_with_input(
            BenchmarkId::new("quad_contains_no_bounds_checks", len),
            &slice[..],
            |b, slice| {
                b.iter(|| quad_contains_no_bounds_checks(black_box(slice), black_box(target)));
            },
        );
        group.bench_with_input(
            BenchmarkId::new("binary_search_contains", len),
            &slice[..],
            |b, slice| {
                let target = black_box(target);
                b.iter(|| black_box(slice).binary_search(&target).is_ok());
            },
        );
    }

    group.finish();
}

criterion_group!(benches, bench_quad_search, bench_quad_contains);
criterion_main!(benches);