1use alloc::format;
2use alloc::vec::Vec;
3use core::hint::black_box;
4use core::ops::Div;
5
6use criterion::{BatchSize, Criterion};
7use p3_field::{Algebra, Field, PrimeCharacteristicRing, chunked_linear_combination};
8use rand::distr::StandardUniform;
9use rand::prelude::Distribution;
10use rand::rngs::SmallRng;
11use rand::{RngExt, SeedableRng};
12
13pub fn benchmark_mul<F: Field>(c: &mut Criterion, name: &str)
16where
17 StandardUniform: Distribution<F>,
18{
19 let mut rng = SmallRng::seed_from_u64(1);
20 let x = rng.random::<F>();
21 let y = rng.random::<F>();
22 c.bench_function(&format!("{name} mul"), |b| {
23 b.iter(|| black_box(black_box(x) * black_box(y)));
24 });
25}
26
27pub fn benchmark_square<F: Field>(c: &mut Criterion, name: &str)
28where
29 StandardUniform: Distribution<F>,
30{
31 let mut rng = SmallRng::seed_from_u64(1);
32 let x = rng.random::<F>();
33 c.bench_function(&format!("{name} square"), |b| {
34 b.iter(|| black_box(black_box(x).square()));
35 });
36}
37
38pub fn benchmark_inv<F: Field>(c: &mut Criterion, name: &str)
39where
40 StandardUniform: Distribution<F>,
41{
42 let mut rng = SmallRng::seed_from_u64(1);
43 let x = rng.random::<F>();
44 c.bench_function(&format!("{name} inv"), |b| {
45 b.iter(|| black_box(black_box(x)).inverse());
46 });
47}
48
49pub fn benchmark_sqrt<F: Field>(c: &mut Criterion, name: &str)
50where
51 StandardUniform: Distribution<F>,
52{
53 let mut rng = SmallRng::seed_from_u64(1);
54 let x = rng.random::<F>().square();
57 c.bench_function(&format!("{name} sqrt"), |b| {
58 b.iter(|| black_box(black_box(x).try_sqrt()));
59 });
60}
61
62pub fn benchmark_mul_2exp<R: PrimeCharacteristicRing + Copy, const REPS: usize>(
63 c: &mut Criterion,
64 name: &str,
65 val: u64,
66) where
67 StandardUniform: Distribution<R>,
68{
69 let mut rng = SmallRng::seed_from_u64(1);
70 let mut input = Vec::new();
71 for _ in 0..REPS {
72 input.push(rng.random::<R>());
73 }
74 c.bench_function(&format!("{name} mul_2exp_u64 {val}"), |b| {
75 b.iter(|| input.iter_mut().for_each(|i| *i = i.mul_2exp_u64(val)));
76 });
77}
78
79pub fn benchmark_halve<F: Field, const REPS: usize>(c: &mut Criterion, name: &str)
80where
81 StandardUniform: Distribution<F>,
82{
83 let mut rng = SmallRng::seed_from_u64(1);
84 let mut input = Vec::new();
85 for _ in 0..REPS {
86 input.push(rng.random::<F>());
87 }
88 c.bench_function(&format!("{name} halve. Num Reps: {REPS}"), |b| {
89 b.iter(|| input.iter_mut().for_each(|i| *i = i.halve()));
90 });
91}
92
93pub fn benchmark_div_2exp<F: Field, const REPS: usize>(c: &mut Criterion, name: &str, val: u64)
94where
95 StandardUniform: Distribution<F>,
96{
97 let mut rng = SmallRng::seed_from_u64(1);
98 let mut input = Vec::new();
99 for _ in 0..REPS {
100 input.push(rng.random::<F>());
101 }
102 c.bench_function(&format!("{name} div_2exp_u64 {val}"), |b| {
103 b.iter(|| input.iter_mut().for_each(|i| *i = i.div_2exp_u64(val)));
104 });
105}
106
107pub fn benchmark_iter_sum<R: PrimeCharacteristicRing + Copy, const N: usize, const REPS: usize>(
113 c: &mut Criterion,
114 name: &str,
115) where
116 StandardUniform: Distribution<R>,
117{
118 let mut rng = SmallRng::seed_from_u64(1);
119 let mut input = Vec::new();
120 for _ in 0..REPS {
121 input.push(rng.random::<[R; N]>());
122 }
123 c.bench_function(&format!("{name} sum/{REPS}, {N}"), |b| {
124 b.iter(|| {
125 let mut acc = R::ZERO;
126 for row in &mut input {
127 acc += row.iter().copied().sum();
128 }
129 acc
130 });
131 });
132}
133
134pub fn benchmark_sum_array<R: PrimeCharacteristicRing + Copy, const N: usize, const REPS: usize>(
140 c: &mut Criterion,
141 name: &str,
142) where
143 StandardUniform: Distribution<R>,
144{
145 let mut rng = SmallRng::seed_from_u64(1);
146 let mut input = Vec::new();
147 for _ in 0..REPS {
148 input.push(rng.random::<[R; N]>());
149 }
150 c.bench_function(&format!("{name} tree sum/{REPS}, {N}"), |b| {
151 b.iter(|| {
152 let mut acc = R::ZERO;
153 for row in &mut input {
154 acc += R::sum_array::<N>(row);
155 }
156 acc
157 });
158 });
159}
160
161pub fn benchmark_dot_array<R: PrimeCharacteristicRing + Copy, const N: usize>(
166 c: &mut Criterion,
167 name: &str,
168) where
169 StandardUniform: Distribution<R>,
170{
171 let mut rng = SmallRng::seed_from_u64(1);
172 let lhs = rng.random::<[R; N]>();
173 let rhs = rng.random::<[R; N]>();
174
175 c.bench_function(&format!("{name} dot product/{N}"), |b| {
176 b.iter(|| black_box(R::dot_product(black_box(&lhs), black_box(&rhs))));
177 });
178}
179
180pub fn benchmark_mixed_dot_array<A: Algebra<F> + Copy, F: Field, const N: usize>(
182 c: &mut Criterion,
183 name: &str,
184) where
185 StandardUniform: Distribution<A> + Distribution<F>,
186{
187 let mut rng = SmallRng::seed_from_u64(1);
188 let a = rng.random::<[A; N]>();
189 let f = rng.random::<[F; N]>();
190 c.bench_function(&format!("{name} mixed dot product/{N}"), |b| {
191 b.iter(|| black_box(A::mixed_dot_product(black_box(&a), black_box(&f))));
192 });
193}
194
195pub fn benchmark_add_slices<F: Field, const LENGTH: usize>(c: &mut Criterion, name: &str)
197where
198 StandardUniform: Distribution<F>,
199{
200 let mut rng = SmallRng::seed_from_u64(1);
201 let mut slice_1 = Vec::new();
202 let mut slice_2 = Vec::new();
203 for _ in 0..LENGTH {
204 slice_1.push(rng.random());
205 slice_2.push(rng.random());
206 }
207 c.bench_function(&format!("{name} add slices/{LENGTH}"), |b| {
208 let mut in_slice = slice_1.clone();
209 b.iter(|| {
210 F::add_slices(&mut in_slice, &slice_2);
211 });
212 });
213}
214
215pub fn benchmark_add_latency<R: PrimeCharacteristicRing + Copy, const N: usize>(
216 c: &mut Criterion,
217 name: &str,
218) where
219 StandardUniform: Distribution<R>,
220{
221 c.bench_function(&format!("add-latency/{N} {name}"), |b| {
222 b.iter_batched(
223 || {
224 let mut rng = SmallRng::seed_from_u64(1);
225 let mut vec = Vec::new();
226 for _ in 0..N {
227 vec.push(rng.random::<R>());
228 }
229 vec
230 },
231 |x| x.iter().fold(R::ZERO, |x, y| x + *y),
232 BatchSize::SmallInput,
233 );
234 });
235}
236
237pub fn benchmark_add_throughput<R: PrimeCharacteristicRing + Copy, const N: usize>(
238 c: &mut Criterion,
239 name: &str,
240) where
241 StandardUniform: Distribution<R>,
242{
243 c.bench_function(&format!("add-throughput/{N} {name}"), |b| {
244 b.iter_batched(
245 || {
246 let mut rng = SmallRng::seed_from_u64(1);
247 (
248 rng.random::<R>(),
249 rng.random::<R>(),
250 rng.random::<R>(),
251 rng.random::<R>(),
252 rng.random::<R>(),
253 rng.random::<R>(),
254 rng.random::<R>(),
255 rng.random::<R>(),
256 rng.random::<R>(),
257 rng.random::<R>(),
258 )
259 },
260 |(mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h, mut i, mut j)| {
261 for _ in 0..N {
262 (a, b, c, d, e, f, g, h, i, j) = (
263 a + b,
264 b + c,
265 c + d,
266 d + e,
267 e + f,
268 f + g,
269 g + h,
270 h + i,
271 i + j,
272 j + a,
273 );
274 }
275 (a, b, c, d, e, f, g, h, i, j)
276 },
277 BatchSize::SmallInput,
278 );
279 });
280}
281
282pub fn benchmark_sub_latency<R: PrimeCharacteristicRing + Copy, const N: usize>(
283 c: &mut Criterion,
284 name: &str,
285) where
286 StandardUniform: Distribution<R>,
287{
288 c.bench_function(&format!("sub-latency/{N} {name}"), |b| {
289 b.iter_batched(
290 || {
291 let mut rng = SmallRng::seed_from_u64(1);
292 let mut vec = Vec::new();
293 for _ in 0..N {
294 vec.push(rng.random::<R>());
295 }
296 vec
297 },
298 |x| x.iter().fold(R::ZERO, |x, y| x - *y),
299 BatchSize::SmallInput,
300 );
301 });
302}
303
304pub fn benchmark_sub_throughput<R: PrimeCharacteristicRing + Copy, const N: usize>(
305 c: &mut Criterion,
306 name: &str,
307) where
308 StandardUniform: Distribution<R>,
309{
310 c.bench_function(&format!("sub-throughput/{N} {name}"), |b| {
311 b.iter_batched(
312 || {
313 let mut rng = SmallRng::seed_from_u64(1);
314 (
315 rng.random::<R>(),
316 rng.random::<R>(),
317 rng.random::<R>(),
318 rng.random::<R>(),
319 rng.random::<R>(),
320 rng.random::<R>(),
321 rng.random::<R>(),
322 rng.random::<R>(),
323 rng.random::<R>(),
324 rng.random::<R>(),
325 )
326 },
327 |(mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h, mut i, mut j)| {
328 for _ in 0..N {
329 (a, b, c, d, e, f, g, h, i, j) = (
330 a - b,
331 b - c,
332 c - d,
333 d - e,
334 e - f,
335 f - g,
336 g - h,
337 h - i,
338 i - j,
339 j - a,
340 );
341 }
342 (a, b, c, d, e, f, g, h, i, j)
343 },
344 BatchSize::SmallInput,
345 );
346 });
347}
348
349pub fn benchmark_mul_latency<R: PrimeCharacteristicRing + Copy, const N: usize>(
350 c: &mut Criterion,
351 name: &str,
352) where
353 StandardUniform: Distribution<R>,
354{
355 c.bench_function(&format!("mul-latency/{N} {name}"), |b| {
356 b.iter_batched(
357 || {
358 let mut rng = SmallRng::seed_from_u64(1);
359 let mut vec = Vec::new();
360 for _ in 0..N {
361 vec.push(rng.random::<R>());
362 }
363 vec
364 },
365 |x| x.iter().fold(R::ONE, |x, y| x * *y),
366 BatchSize::SmallInput,
367 );
368 });
369}
370
371pub fn benchmark_mul_throughput<R: PrimeCharacteristicRing + Copy, const N: usize>(
372 c: &mut Criterion,
373 name: &str,
374) where
375 StandardUniform: Distribution<R>,
376{
377 c.bench_function(&format!("mul-throughput/{N} {name}"), |b| {
378 b.iter_batched(
379 || {
380 let mut rng = SmallRng::seed_from_u64(1);
381 (
382 rng.random::<R>(),
383 rng.random::<R>(),
384 rng.random::<R>(),
385 rng.random::<R>(),
386 rng.random::<R>(),
387 rng.random::<R>(),
388 rng.random::<R>(),
389 rng.random::<R>(),
390 rng.random::<R>(),
391 rng.random::<R>(),
392 )
393 },
394 |(mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h, mut i, mut j)| {
395 for _ in 0..N {
396 (a, b, c, d, e, f, g, h, i, j) = (
397 a * b,
398 b * c,
399 c * d,
400 d * e,
401 e * f,
402 f * g,
403 g * h,
404 h * i,
405 i * j,
406 j * a,
407 );
408 }
409 (a, b, c, d, e, f, g, h, i, j)
410 },
411 BatchSize::SmallInput,
412 );
413 });
414}
415
416pub fn benchmark_div_latency<R: PrimeCharacteristicRing + Copy + Div<Output = R>, const N: usize>(
417 c: &mut Criterion,
418 name: &str,
419) where
420 StandardUniform: Distribution<R>,
421{
422 c.bench_function(&format!("div-latency/{N} {name}"), |b| {
423 b.iter_batched(
424 || {
425 let mut rng = SmallRng::seed_from_u64(1);
426 let init = rng.random::<R>();
427 let mut vec = Vec::with_capacity(N);
428 for _ in 0..N {
429 vec.push(rng.random::<R>());
430 }
431 (init, vec)
432 },
433 |(init, vec)| vec.iter().fold(init, |x, y| x / *y),
434 BatchSize::SmallInput,
435 );
436 });
437}
438
439pub fn benchmark_div_throughput<
440 R: PrimeCharacteristicRing + Copy + Div<Output = R>,
441 const N: usize,
442>(
443 c: &mut Criterion,
444 name: &str,
445) where
446 StandardUniform: Distribution<R>,
447{
448 c.bench_function(&format!("div-throughput/{N} {name}"), |b| {
449 b.iter_batched(
450 || {
451 let mut rng = SmallRng::seed_from_u64(1);
452 (
453 rng.random::<R>(),
454 rng.random::<R>(),
455 rng.random::<R>(),
456 rng.random::<R>(),
457 rng.random::<R>(),
458 rng.random::<R>(),
459 rng.random::<R>(),
460 rng.random::<R>(),
461 rng.random::<R>(),
462 rng.random::<R>(),
463 )
464 },
465 |(mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h, mut i, mut j)| {
466 for _ in 0..N {
467 (a, b, c, d, e, f, g, h, i, j) = (
468 a / b,
469 b / c,
470 c / d,
471 d / e,
472 e / f,
473 f / g,
474 g / h,
475 h / i,
476 i / j,
477 j / a,
478 );
479 }
480 (a, b, c, d, e, f, g, h, i, j)
481 },
482 BatchSize::SmallInput,
483 );
484 });
485}
486
487pub fn benchmark_base_mul_latency<F: Field, A: Algebra<F> + Copy, const N: usize>(
488 c: &mut Criterion,
489 name: &str,
490) where
491 StandardUniform: Distribution<F> + Distribution<A>,
492{
493 c.bench_function(&format!("base_mul-latency/{N} {name}"), |b| {
494 b.iter_batched(
495 || {
496 let mut rng = SmallRng::seed_from_u64(1);
497 let mut vec = Vec::new();
498 for _ in 0..N {
499 vec.push(rng.random::<F>());
500 }
501 let init_val = rng.random::<A>();
502 (vec, init_val)
503 },
504 |(x, init_val)| x.iter().fold(init_val, |x, y| x * *y),
505 BatchSize::SmallInput,
506 );
507 });
508}
509
510pub fn benchmark_base_mul_throughput<F: Field, A: Algebra<F> + Copy, const N: usize>(
511 c: &mut Criterion,
512 name: &str,
513) where
514 StandardUniform: Distribution<F> + Distribution<A>,
515{
516 c.bench_function(&format!("base_mul-throughput/{N} {name}"), |b| {
517 b.iter_batched(
518 || {
519 let mut rng = SmallRng::seed_from_u64(1);
520 let a_tuple = (
521 rng.random::<A>(),
522 rng.random::<A>(),
523 rng.random::<A>(),
524 rng.random::<A>(),
525 rng.random::<A>(),
526 rng.random::<A>(),
527 rng.random::<A>(),
528 rng.random::<A>(),
529 rng.random::<A>(),
530 rng.random::<A>(),
531 );
532 let f_tuple = (
533 rng.random::<F>(),
534 rng.random::<F>(),
535 rng.random::<F>(),
536 rng.random::<F>(),
537 rng.random::<F>(),
538 rng.random::<F>(),
539 rng.random::<F>(),
540 rng.random::<F>(),
541 rng.random::<F>(),
542 rng.random::<F>(),
543 );
544 (a_tuple, f_tuple)
545 },
546 |(
547 (mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h, mut i, mut j),
548 (a_f, b_f, c_f, d_f, e_f, f_f, g_f, h_f, i_f, j_f),
549 )| {
550 for _ in 0..N {
551 (a, b, c, d, e, f, g, h, i, j) = (
552 a * a_f,
553 b * b_f,
554 c * c_f,
555 d * d_f,
556 e * e_f,
557 f * f_f,
558 g * g_f,
559 h * h_f,
560 i * i_f,
561 j * j_f,
562 );
563 }
564 (a, b, c, d, e, f, g, h, i, j)
565 },
566 BatchSize::SmallInput,
567 );
568 });
569}
570
571pub fn benchmark_exp_const<R: PrimeCharacteristicRing + Copy, const POWER: u64, const REPS: usize>(
576 c: &mut Criterion,
577 name: &str,
578) where
579 StandardUniform: Distribution<R>,
580{
581 let mut rng = SmallRng::seed_from_u64(1);
582 let input: Vec<R> = (0..REPS).map(|_| rng.random()).collect();
583
584 c.bench_function(&format!("{name} exp_const<{POWER}>/{REPS}"), |b| {
585 b.iter_batched(
586 || input.clone(),
587 |mut data| {
588 for x in data.iter_mut() {
589 *x = x.exp_const_u64::<POWER>();
590 }
591 black_box(data);
592 },
593 BatchSize::SmallInput,
594 );
595 });
596}
597
598pub fn benchmark_neg_latency<R: PrimeCharacteristicRing + Copy, const N: usize>(
599 c: &mut Criterion,
600 name: &str,
601) where
602 StandardUniform: Distribution<R>,
603{
604 c.bench_function(&format!("neg-latency/{N} {name}"), |b| {
605 b.iter_batched(
606 || {
607 let mut rng = SmallRng::seed_from_u64(1);
608 let mut vec = Vec::new();
609 for _ in 0..N {
610 vec.push(rng.random::<R>());
611 }
612 vec
613 },
614 |x| x.iter().fold(R::ZERO, |acc, y| -(acc + *y)),
615 BatchSize::SmallInput,
616 );
617 });
618}
619
620pub fn benchmark_neg_throughput<R: PrimeCharacteristicRing + Copy, const N: usize>(
621 c: &mut Criterion,
622 name: &str,
623) where
624 StandardUniform: Distribution<R>,
625{
626 c.bench_function(&format!("neg-throughput/{N} {name}"), |b| {
627 b.iter_batched(
628 || {
629 let mut rng = SmallRng::seed_from_u64(1);
630 (
631 rng.random::<R>(),
632 rng.random::<R>(),
633 rng.random::<R>(),
634 rng.random::<R>(),
635 rng.random::<R>(),
636 rng.random::<R>(),
637 rng.random::<R>(),
638 rng.random::<R>(),
639 rng.random::<R>(),
640 rng.random::<R>(),
641 )
642 },
643 |(mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h, mut i, mut j)| {
644 for _ in 0..N {
645 (a, b, c, d, e, f, g, h, i, j) = (-a, -b, -c, -d, -e, -f, -g, -h, -i, -j);
646 }
647 (a, b, c, d, e, f, g, h, i, j)
648 },
649 BatchSize::SmallInput,
650 );
651 });
652}
653
654pub fn benchmark_double_latency<R: PrimeCharacteristicRing + Copy, const N: usize>(
655 c: &mut Criterion,
656 name: &str,
657) where
658 StandardUniform: Distribution<R>,
659{
660 c.bench_function(&format!("double-latency/{N} {name}"), |b| {
661 b.iter_batched(
662 || {
663 let mut rng = SmallRng::seed_from_u64(1);
664 black_box(rng.random::<R>())
665 },
666 |x| {
667 let mut acc = x;
668 for _ in 0..N {
669 acc = acc.double();
670 }
671 acc
672 },
673 BatchSize::SmallInput,
674 );
675 });
676}
677
678pub fn benchmark_double_throughput<R: PrimeCharacteristicRing + Copy, const N: usize>(
679 c: &mut Criterion,
680 name: &str,
681) where
682 StandardUniform: Distribution<R>,
683{
684 c.bench_function(&format!("double-throughput/{N} {name}"), |b| {
685 b.iter_batched(
686 || {
687 let mut rng = SmallRng::seed_from_u64(1);
688 (
689 rng.random::<R>(),
690 rng.random::<R>(),
691 rng.random::<R>(),
692 rng.random::<R>(),
693 rng.random::<R>(),
694 rng.random::<R>(),
695 rng.random::<R>(),
696 rng.random::<R>(),
697 rng.random::<R>(),
698 rng.random::<R>(),
699 )
700 },
701 |(mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h, mut i, mut j)| {
702 for _ in 0..N {
703 (a, b, c, d, e, f, g, h, i, j) = (
704 a.double(),
705 b.double(),
706 c.double(),
707 d.double(),
708 e.double(),
709 f.double(),
710 g.double(),
711 h.double(),
712 i.double(),
713 j.double(),
714 );
715 }
716 (a, b, c, d, e, f, g, h, i, j)
717 },
718 BatchSize::SmallInput,
719 );
720 });
721}
722
723pub fn benchmark_chunked_linear_combination<F: Field, A: Algebra<F> + Copy, const LEN: usize>(
726 c: &mut Criterion,
727 name: &str,
728) where
729 StandardUniform: Distribution<F> + Distribution<A>,
730{
731 let mut rng = SmallRng::seed_from_u64(1);
732 let values: Vec<A> = (0..LEN).map(|_| rng.random()).collect();
733 let coeffs: Vec<F> = (0..LEN).map(|_| rng.random()).collect();
734
735 macro_rules! bench_chunk {
736 ($($chunk:literal),*) => {$(
737 c.bench_function(
738 &format!("{name} batched_lc/chunk={}, len={LEN}", $chunk),
739 |b| {
740 b.iter(|| {
741 chunked_linear_combination::<$chunk, A, F>(
742 black_box(values.as_slice()),
743 black_box(coeffs.as_slice()),
744 )
745 });
746 },
747 );
748 )*};
749 }
750 bench_chunk!(1, 2, 4, 8, 16, 32, 64);
751}
752
753#[macro_export]
775macro_rules! bench_packed_extension_field {
776 ($base:ty, $($label:ident = $ef:ty),+ $(,)?) => {
777 const REPS: usize = 100;
779 const L_REPS: usize = 10 * REPS;
780
781 $(
782 fn $label(c: &mut criterion::Criterion) {
783 type Packed = <$ef as p3_field::ExtensionField<$base>>::ExtensionPacking;
784 let name = stringify!($ef);
785 $crate::bench_func::benchmark_add_throughput::<Packed, REPS>(c, name);
786 $crate::bench_func::benchmark_add_latency::<Packed, L_REPS>(c, name);
787 $crate::bench_func::benchmark_mul_throughput::<Packed, REPS>(c, name);
788 $crate::bench_func::benchmark_mul_latency::<Packed, L_REPS>(c, name);
789 $crate::bench_func::benchmark_div_throughput::<Packed, REPS>(c, name);
790 $crate::bench_func::benchmark_div_latency::<Packed, L_REPS>(c, name);
791 }
792 )+
793
794 criterion::criterion_group!(packed_ext_benches, $($label),+);
795 criterion::criterion_main!(packed_ext_benches);
796 };
797}