himada-dispatch 0.1.0

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

use himada_dispatch::kernels::*;

const SIZES: &[usize] = &[100, 1_000, 10_000, 100_000, 1_000_000];

fn bench_dot_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("dot_f64");
    for &n in SIZES {
        let a: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let b: Vec<f64> = (0..n).map(|i| (i * 2) as f64).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| black_box(dot_scalar(black_box(&a), black_box(&b))))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| black_box(dot_neon(black_box(&a), black_box(&b))))
        });
    }
    group.finish();
}

fn bench_dot_f32(c: &mut Criterion) {
    let mut group = c.benchmark_group("dot_f32");
    for &n in SIZES {
        let a: Vec<f32> = (0..n).map(|i| i as f32).collect();
        let b: Vec<f32> = (0..n).map(|i| (i * 2) as f32).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| black_box(dot_f32_scalar(black_box(&a), black_box(&b))))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| black_box(dot_f32_neon(black_box(&a), black_box(&b))))
        });
    }
    group.finish();
}

fn bench_reduce_sum_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("reduce_sum_f64");
    for &n in SIZES {
        let a: Vec<f64> = (0..n).map(|i| (i % 100) as f64).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b, _| {
            b.iter(|| black_box(reduce_sum_scalar(black_box(&a))))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b, _| {
            b.iter(|| black_box(reduce_sum_neon(black_box(&a))))
        });
    }
    group.finish();
}

fn bench_reduce_sum_f32(c: &mut Criterion) {
    let mut group = c.benchmark_group("reduce_sum_f32");
    for &n in SIZES {
        let a: Vec<f32> = (0..n).map(|i| (i % 100) as f32).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b, _| {
            b.iter(|| black_box(reduce_sum_f32_scalar(black_box(&a))))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b, _| {
            b.iter(|| black_box(reduce_sum_f32_neon(black_box(&a))))
        });
    }
    group.finish();
}

fn bench_reduce_max_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("reduce_max_f64");
    for &n in SIZES {
        let a: Vec<f64> = (0..n).map(|i| (i % 100) as f64).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b, _| {
            b.iter(|| black_box(reduce_max_scalar(black_box(&a))))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b, _| {
            b.iter(|| black_box(reduce_max_neon(black_box(&a))))
        });
    }
    group.finish();
}

fn bench_reduce_max_f32(c: &mut Criterion) {
    let mut group = c.benchmark_group("reduce_max_f32");
    for &n in SIZES {
        let a: Vec<f32> = (0..n).map(|i| (i % 100) as f32).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b, _| {
            b.iter(|| black_box(reduce_max_f32_scalar(black_box(&a))))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b, _| {
            b.iter(|| black_box(reduce_max_f32_neon(black_box(&a))))
        });
    }
    group.finish();
}

fn bench_abs_max_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("abs_max_f64");
    for &n in SIZES {
        let a: Vec<f64> = (0..n).map(|i| (i % 200) as f64 - 100.0).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b, _| {
            b.iter(|| black_box(abs_max_f64_scalar(black_box(&a))))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b, _| {
            b.iter(|| black_box(abs_max_f64_neon(black_box(&a))))
        });
    }
    group.finish();
}

fn bench_abs_max_f32(c: &mut Criterion) {
    let mut group = c.benchmark_group("abs_max_f32");
    for &n in SIZES {
        let a: Vec<f32> = (0..n).map(|i| (i % 200) as f32 - 100.0).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b, _| {
            b.iter(|| black_box(abs_max_f32_scalar(black_box(&a))))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b, _| {
            b.iter(|| black_box(abs_max_f32_neon(black_box(&a))))
        });
    }
    group.finish();
}

fn bench_argmax_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("argmax_f64");
    for &n in SIZES {
        let a: Vec<f64> = (0..n).map(|i| (i % 1000) as f64).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b, _| {
            b.iter(|| black_box(argmax_f64_scalar(black_box(&a))))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b, _| {
            b.iter(|| black_box(argmax_f64_neon(black_box(&a))))
        });
    }
    group.finish();
}

fn bench_argmax_f32(c: &mut Criterion) {
    let mut group = c.benchmark_group("argmax_f32");
    for &n in SIZES {
        let a: Vec<f32> = (0..n).map(|i| (i % 1000) as f32).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b, _| {
            b.iter(|| black_box(argmax_f32_scalar(black_box(&a))))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b, _| {
            b.iter(|| black_box(argmax_f32_neon(black_box(&a))))
        });
    }
    group.finish();
}

fn bench_euclidean_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("euclidean_f64");
    for &n in SIZES {
        let a: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let b: Vec<f64> = (n..2 * n).map(|i| i as f64).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| black_box(euclidean_distance_f64_scalar(black_box(&a), black_box(&b))))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| black_box(euclidean_distance_f64_neon(black_box(&a), black_box(&b))))
        });
    }
    group.finish();
}

fn bench_euclidean_f32(c: &mut Criterion) {
    let mut group = c.benchmark_group("euclidean_f32");
    for &n in SIZES {
        let a: Vec<f32> = (0..n).map(|i| i as f32).collect();
        let b: Vec<f32> = (n..2 * n).map(|i| i as f32).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| black_box(euclidean_distance_f32_scalar(black_box(&a), black_box(&b))))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| black_box(euclidean_distance_f32_neon(black_box(&a), black_box(&b))))
        });
    }
    group.finish();
}

fn bench_cosine_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("cosine_similarity_f64");
    for &n in SIZES[..SIZES.len() - 1].iter() {
        let a: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let b: Vec<f64> = (n..2 * n).map(|i| i as f64).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| black_box(cosine_similarity_f64_scalar(black_box(&a), black_box(&b))))
        });
    }
    group.finish();
}

fn bench_cosine_f32(c: &mut Criterion) {
    let mut group = c.benchmark_group("cosine_similarity_f32");
    for &n in SIZES[..SIZES.len() - 1].iter() {
        let a: Vec<f32> = (0..n).map(|i| i as f32).collect();
        let b: Vec<f32> = (n..2 * n).map(|i| i as f32).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| black_box(cosine_similarity_f32_scalar(black_box(&a), black_box(&b))))
        });
    }
    group.finish();
}

fn bench_hadamard_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("hadamard_f64");
    for &n in SIZES {
        let a: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let b: Vec<f64> = (0..n).map(|i| (i * 2) as f64).collect();
        let mut c_ = vec![0.0; n];
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| hadamard_product_f64_scalar(black_box(&a), black_box(&b), black_box(&mut c_)))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| hadamard_product_f64_neon(black_box(&a), black_box(&b), black_box(&mut c_)))
        });
    }
    group.finish();
}

fn bench_hadamard_f32(c: &mut Criterion) {
    let mut group = c.benchmark_group("hadamard_f32");
    for &n in SIZES {
        let a: Vec<f32> = (0..n).map(|i| i as f32).collect();
        let b: Vec<f32> = (0..n).map(|i| (i * 2) as f32).collect();
        let mut c_ = vec![0.0; n];
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| hadamard_product_f32_scalar(black_box(&a), black_box(&b), black_box(&mut c_)))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| hadamard_product_f32_neon(black_box(&a), black_box(&b), black_box(&mut c_)))
        });
    }
    group.finish();
}

fn bench_add_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("add_f64");
    for &n in SIZES {
        let a: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let b: Vec<f64> = (0..n).map(|i| (i * 2) as f64).collect();
        let mut c_ = vec![0.0; n];
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| add_f64_scalar(black_box(&a), black_box(&b), black_box(&mut c_)))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| hadamard_product_f64_neon(black_box(&a), black_box(&b), black_box(&mut c_)))
        });
    }
    group.finish();
}

fn bench_add_f32(c: &mut Criterion) {
    let mut group = c.benchmark_group("add_f32");
    for &n in SIZES {
        let a: Vec<f32> = (0..n).map(|i| i as f32).collect();
        let b: Vec<f32> = (0..n).map(|i| (i * 2) as f32).collect();
        let mut c_ = vec![0.0; n];
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| add_f32_scalar(black_box(&a), black_box(&b), black_box(&mut c_)))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| add_f32_neon(black_box(&a), black_box(&b), black_box(&mut c_)))
        });
    }
    group.finish();
}

fn bench_sub_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("sub_f64");
    for &n in SIZES {
        let a: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let b: Vec<f64> = (0..n).map(|i| (i * 2) as f64).collect();
        let mut c_ = vec![0.0; n];
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| sub_f64_scalar(black_box(&a), black_box(&b), black_box(&mut c_)))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| hadamard_product_f64_neon(black_box(&a), black_box(&b), black_box(&mut c_)))
        });
    }
    group.finish();
}

fn bench_mul_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("mul_f64");
    for &n in SIZES {
        let a: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let b: Vec<f64> = (0..n).map(|i| (i * 2) as f64).collect();
        let mut c_ = vec![0.0; n];
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| mul_f64_scalar(black_box(&a), black_box(&b), black_box(&mut c_)))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| hadamard_product_f64_neon(black_box(&a), black_box(&b), black_box(&mut c_)))
        });
    }
    group.finish();
}

fn bench_negate_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("negate_f64");
    for &n in SIZES {
        let a: Vec<f64> = (0..n).map(|i| i as f64).collect();
        let mut c_ = vec![0.0; n];
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| negate_f64_scalar(black_box(&a), black_box(&mut c_)))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| negate_f64_neon(black_box(&a), black_box(&mut c_)))
        });
    }
    group.finish();
}

fn bench_negate_f32(c: &mut Criterion) {
    let mut group = c.benchmark_group("negate_f32");
    for &n in SIZES {
        let a: Vec<f32> = (0..n).map(|i| i as f32).collect();
        let mut c_ = vec![0.0; n];
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| negate_f32_scalar(black_box(&a), black_box(&mut c_)))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| negate_f32_neon(black_box(&a), black_box(&mut c_)))
        });
    }
    group.finish();
}

fn bench_clamp_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("clamp_f64");
    for &n in SIZES {
        let a: Vec<f64> = (0..n).map(|i| (i % 200) as f64 - 100.0).collect();
        let mut c_ = vec![0.0; n];
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| clamp_f64_scalar(black_box(&a), -20.0, 20.0, black_box(&mut c_)))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| clamp_f64_neon(black_box(&a), -20.0, 20.0, black_box(&mut c_)))
        });
    }
    group.finish();
}

fn bench_clamp_f32(c: &mut Criterion) {
    let mut group = c.benchmark_group("clamp_f32");
    for &n in SIZES {
        let a: Vec<f32> = (0..n).map(|i| (i % 200) as f32 - 100.0).collect();
        let mut c_ = vec![0.0; n];
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| clamp_f32_scalar(black_box(&a), -20.0, 20.0, black_box(&mut c_)))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| clamp_f32_neon(black_box(&a), -20.0, 20.0, black_box(&mut c_)))
        });
    }
    group.finish();
}

fn bench_softmax_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("softmax_f64");
    for &n in SIZES[..SIZES.len() - 1].iter() {
        let input: Vec<f64> = (0..n).map(|i| (i % 100) as f64).collect();
        let mut output = vec![0.0; n];
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| softmax_scalar(black_box(&input), black_box(&mut output)))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| softmax_neon(black_box(&input), black_box(&mut output)))
        });
    }
    group.finish();
}

fn bench_softmax_f32(c: &mut Criterion) {
    let mut group = c.benchmark_group("softmax_f32");
    for &n in SIZES[..SIZES.len() - 1].iter() {
        let input: Vec<f32> = (0..n).map(|i| (i % 100) as f32).collect();
        let mut output = vec![0.0; n];
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| softmax_f32_scalar(black_box(&input), black_box(&mut output)))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| softmax_f32_neon(black_box(&input), black_box(&mut output)))
        });
    }
    group.finish();
}

fn bench_memchr(c: &mut Criterion) {
    let mut group = c.benchmark_group("memchr");
    for &n in SIZES {
        let data: Vec<u8> = (0..n).map(|i| (i % 251) as u8).collect();
        let target = 42u8;
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| black_box(memchr_scalar(target, black_box(&data))))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| black_box(memchr_neon(target, black_box(&data))))
        });
    }
    group.finish();
}

const MATMUL_SIZES: &[usize] = &[32, 64, 128, 256];

fn bench_matmul_f64(c: &mut Criterion) {
    let mut group = c.benchmark_group("matmul_f64");
    for &n in MATMUL_SIZES {
        let a: Vec<f64> = (0..n * n).map(|i| (i % 5) as f64).collect();
        let b: Vec<f64> = (0..n * n).map(|i| (i % 7) as f64).collect();
        let mut c_ = vec![0.0; n * n];
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            c_.fill(0.0);
            b_.iter(|| matmul_scalar(black_box(&a), black_box(&b), black_box(&mut c_), n))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            c_.fill(0.0);
            b_.iter(|| matmul_neon(black_box(&a), black_box(&b), black_box(&mut c_), n))
        });
        group.bench_with_input(BenchmarkId::new("cache_tiled", n), &n, |b_, _| {
            c_.fill(0.0);
            b_.iter(|| matmul_cache_tiled(black_box(&a), black_box(&b), black_box(&mut c_), n))
        });
    }
    group.finish();
}

fn bench_matmul_f32(c: &mut Criterion) {
    let mut group = c.benchmark_group("matmul_f32");
    for &n in MATMUL_SIZES {
        let a: Vec<f32> = (0..n * n).map(|i| (i % 5) as f32).collect();
        let b: Vec<f32> = (0..n * n).map(|i| (i % 7) as f32).collect();
        let mut c_ = vec![0.0; n * n];
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            c_.fill(0.0);
            b_.iter(|| matmul_f32_scalar(black_box(&a), black_box(&b), black_box(&mut c_), n))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            c_.fill(0.0);
            b_.iter(|| matmul_f32_neon(black_box(&a), black_box(&b), black_box(&mut c_), n))
        });
    }
    group.finish();
}

fn bench_fused_dot_clamp_reduce(c: &mut Criterion) {
    let mut group = c.benchmark_group("fused_dot_clamp_reduce");
    for &n in SIZES {
        let a: Vec<f64> = (0..n).map(|i| (i % 10) as f64 * 0.5).collect();
        let b: Vec<f64> = (0..n).map(|i| (i % 10) as f64 * 0.3).collect();
        group.bench_with_input(BenchmarkId::new("scalar", n), &n, |b_, _| {
            b_.iter(|| black_box(dot_clamp_reduce_scalar(black_box(&a), black_box(&b))))
        });
        group.bench_with_input(BenchmarkId::new("neon", n), &n, |b_, _| {
            b_.iter(|| black_box(dot_clamp_reduce_neon(black_box(&a), black_box(&b))))
        });
    }
    group.finish();
}

fn bench_pipeline_dot_clamp_softmax_reduce(c: &mut Criterion) {
    let mut group = c.benchmark_group("pipeline_dot_clamp_softmax_reduce");
    for &n in &[1_000, 10_000, 100_000] {
        let a: Vec<f64> = (0..n).map(|i| (i % 100) as f64).collect();
        let b: Vec<f64> = (0..n).map(|i| (i % 100) as f64 * 0.5).collect();
        group.bench_with_input(BenchmarkId::new("separate", n), &n, |b_, _| {
            b_.iter(|| {
                let mut temp = vec![0.0; n];
                hadamard_product_f64_scalar(&a, &b, &mut temp);
                let max = temp.iter().copied().fold(f64::NEG_INFINITY, f64::max);
                let mut sum = 0.0;
                for v in temp.iter_mut() {
                    let clamped = v.clamp(0.0, 1.0);
                    let e = (clamped - max).exp();
                    *v = e;
                    sum += e;
                }
                let inv = 1.0 / sum;
                for v in temp.iter_mut() {
                    *v *= inv;
                }
                black_box(temp.iter().sum::<f64>())
            })
        });
        group.bench_with_input(BenchmarkId::new("fused", n), &n, |b_, _| {
            b_.iter(|| {
                black_box(dot_clamp_reduce_scalar(&a, &b))
            })
        });
    }
    group.finish();
}

criterion_group!(
    benches,
    bench_dot_f64,
    bench_dot_f32,
    bench_reduce_sum_f64,
    bench_reduce_sum_f32,
    bench_reduce_max_f64,
    bench_reduce_max_f32,
    bench_abs_max_f64,
    bench_abs_max_f32,
    bench_argmax_f64,
    bench_argmax_f32,
    bench_euclidean_f64,
    bench_euclidean_f32,
    bench_cosine_f64,
    bench_cosine_f32,
    bench_hadamard_f64,
    bench_hadamard_f32,
    bench_add_f64,
    bench_add_f32,
    bench_sub_f64,
    bench_mul_f64,
    bench_negate_f64,
    bench_negate_f32,
    bench_clamp_f64,
    bench_clamp_f32,
    bench_softmax_f64,
    bench_softmax_f32,
    bench_memchr,
    bench_matmul_f64,
    bench_matmul_f32,
    bench_fused_dot_clamp_reduce,
    bench_pipeline_dot_clamp_softmax_reduce,
);
criterion_main!(benches);