Skip to content

Commit

Permalink
Add benchmarks for SIMD m31
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewmilson committed May 16, 2024
1 parent 9f9d708 commit a4fa5c7
Showing 1 changed file with 54 additions and 15 deletions.
69 changes: 54 additions & 15 deletions crates/prover/benches/field.rs
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
use criterion::Criterion;
use num_traits::One;
use rand::rngs::SmallRng;
use rand::{Rng, SeedableRng};
use stwo_prover::core::backend::simd::m31::N_LANES;
use stwo_prover::core::fields::cm31::CM31;
use stwo_prover::core::fields::m31::M31;
use stwo_prover::core::fields::m31::{BaseField, M31};
use stwo_prover::core::fields::qm31::SecureField;

pub const N_ELEMENTS: usize = 1 << 16;
Expand Down Expand Up @@ -153,23 +155,60 @@ pub fn avx512_m31_operations_bench(c: &mut criterion::Criterion) {
});
}

pub fn simd_m31_operations_bench(c: &mut criterion::Criterion) {
use stwo_prover::core::backend::simd::m31::PackedBaseField;

let mut rng = SmallRng::seed_from_u64(0);
let elements: Vec<PackedBaseField> = (0..N_ELEMENTS / N_LANES).map(|_| rng.gen()).collect();
let mut states = vec![PackedBaseField::broadcast(BaseField::one()); N_STATE_ELEMENTS];

c.bench_function("mul_simd", |b| {
b.iter(|| {
for elem in elements.iter() {
for _ in 0..128 {
for state in states.iter_mut() {
*state *= *elem;
}
}
}
})
});

c.bench_function("add_simd", |b| {
b.iter(|| {
for elem in elements.iter() {
for _ in 0..128 {
for state in states.iter_mut() {
*state += *elem;
}
}
}
})
});

c.bench_function("sub_simd", |b| {
b.iter(|| {
for elem in elements.iter() {
for _ in 0..128 {
for state in states.iter_mut() {
*state -= *elem;
}
}
}
})
});
}

#[cfg(target_arch = "x86_64")]
criterion::criterion_group!(
name=m31_benches;
name = benches;
config = Criterion::default().sample_size(10);
targets=
m31_operations_bench,
avx512_m31_operations_bench
);
targets = m31_operations_bench, cm31_operations_bench, qm31_operations_bench,
avx512_m31_operations_bench, simd_m31_operations_bench);
#[cfg(not(target_arch = "x86_64"))]
criterion::criterion_group!(m31_benches, m31_operations_bench);

criterion::criterion_group!(
name=field_comparison;
name = benches;
config = Criterion::default().sample_size(10);
targets=
m31_operations_bench,
cm31_operations_bench,
qm31_operations_bench
);
criterion::criterion_main!(field_comparison, m31_benches);
targets = m31_operations_bench, cm31_operations_bench, qm31_operations_bench,
simd_m31_operations_bench);
criterion::criterion_main!(benches);

0 comments on commit a4fa5c7

Please sign in to comment.