1use core::marker::PhantomData;
48use core::ops::{Add, AddAssign, Neg, Sub, SubAssign};
49
50use p3_field::{Algebra, Field};
51
52pub trait ConvolutionElt:
56 Add<Output = Self> + AddAssign + Copy + Neg<Output = Self> + Sub<Output = Self> + SubAssign
57{
58}
59
60impl<T> ConvolutionElt for T where
61 T: Add<Output = T> + AddAssign + Copy + Neg<Output = T> + Sub<Output = T> + SubAssign
62{
63}
64
65pub trait ConvolutionRhs:
69 Add<Output = Self> + Copy + Neg<Output = Self> + Sub<Output = Self>
70{
71}
72
73impl<T> ConvolutionRhs for T where T: Add<Output = T> + Copy + Neg<Output = T> + Sub<Output = T> {}
74
75pub trait Convolve<F, T: ConvolutionElt, U: ConvolutionRhs> {
97 const T_ZERO: T;
102
103 const U_ZERO: U;
108
109 fn halve(val: T) -> T;
114
115 fn read(input: F) -> T;
118
119 fn parity_dot<const N: usize>(lhs: [T; N], rhs: [U; N]) -> T;
126
127 fn reduce(z: T) -> F;
130
131 #[inline(always)]
136 fn apply<const N: usize, C: Fn([T; N], [U; N], &mut [T])>(
137 lhs: [F; N],
138 rhs: [U; N],
139 conv: C,
140 ) -> [F; N] {
141 let lhs = lhs.map(Self::read);
142 let mut output = [Self::T_ZERO; N];
143 conv(lhs, rhs, &mut output);
144 output.map(Self::reduce)
145 }
146
147 #[inline(always)]
148 fn conv3(lhs: [T; 3], rhs: [U; 3], output: &mut [T]) {
149 output[0] = Self::parity_dot(lhs, [rhs[0], rhs[2], rhs[1]]);
150 output[1] = Self::parity_dot(lhs, [rhs[1], rhs[0], rhs[2]]);
151 output[2] = Self::parity_dot(lhs, [rhs[2], rhs[1], rhs[0]]);
152 }
153
154 #[inline(always)]
155 fn negacyclic_conv3(lhs: [T; 3], rhs: [U; 3], output: &mut [T]) {
156 output[0] = Self::parity_dot(lhs, [rhs[0], -rhs[2], -rhs[1]]);
157 output[1] = Self::parity_dot(lhs, [rhs[1], rhs[0], -rhs[2]]);
158 output[2] = Self::parity_dot(lhs, [rhs[2], rhs[1], rhs[0]]);
159 }
160
161 #[inline(always)]
162 fn conv4(lhs: [T; 4], rhs: [U; 4], output: &mut [T]) {
163 let u_p = [lhs[0] + lhs[2], lhs[1] + lhs[3]];
166 let u_m = [lhs[0] - lhs[2], lhs[1] - lhs[3]];
167 let v_p = [rhs[0] + rhs[2], rhs[1] + rhs[3]];
168 let v_m = [rhs[0] - rhs[2], rhs[1] - rhs[3]];
169
170 output[0] = Self::parity_dot(u_m, [v_m[0], -v_m[1]]);
171 output[1] = Self::parity_dot(u_m, [v_m[1], v_m[0]]);
172 output[2] = Self::parity_dot(u_p, v_p);
173 output[3] = Self::parity_dot(u_p, [v_p[1], v_p[0]]);
174
175 output[0] += output[2];
176 output[1] += output[3];
177
178 output[0] = Self::halve(output[0]);
179 output[1] = Self::halve(output[1]);
180
181 output[2] -= output[0];
182 output[3] -= output[1];
183 }
184
185 #[inline(always)]
186 fn negacyclic_conv4(lhs: [T; 4], rhs: [U; 4], output: &mut [T]) {
187 output[0] = Self::parity_dot(lhs, [rhs[0], -rhs[3], -rhs[2], -rhs[1]]);
188 output[1] = Self::parity_dot(lhs, [rhs[1], rhs[0], -rhs[3], -rhs[2]]);
189 output[2] = Self::parity_dot(lhs, [rhs[2], rhs[1], rhs[0], -rhs[3]]);
190 output[3] = Self::parity_dot(lhs, [rhs[3], rhs[2], rhs[1], rhs[0]]);
191 }
192
193 #[inline(always)]
196 fn conv_n_recursive<const N: usize, const HALF_N: usize, C, NC>(
197 lhs: [T; N],
198 rhs: [U; N],
199 output: &mut [T],
200 inner_conv: C,
201 inner_negacyclic_conv: NC,
202 ) where
203 C: Fn([T; HALF_N], [U; HALF_N], &mut [T]),
204 NC: Fn([T; HALF_N], [U; HALF_N], &mut [T]),
205 {
206 debug_assert_eq!(2 * HALF_N, N);
207 let mut lhs_pos = [Self::T_ZERO; HALF_N]; let mut lhs_neg = [Self::T_ZERO; HALF_N]; let mut rhs_pos = [Self::U_ZERO; HALF_N]; let mut rhs_neg = [Self::U_ZERO; HALF_N]; for i in 0..HALF_N {
213 let s = lhs[i];
214 let t = lhs[i + HALF_N];
215
216 lhs_pos[i] = s + t;
217 lhs_neg[i] = s - t;
218
219 let s = rhs[i];
220 let t = rhs[i + HALF_N];
221
222 rhs_pos[i] = s + t;
223 rhs_neg[i] = s - t;
224 }
225
226 let (left, right) = output.split_at_mut(HALF_N);
227
228 inner_negacyclic_conv(lhs_neg, rhs_neg, left);
230
231 inner_conv(lhs_pos, rhs_pos, right);
233
234 for i in 0..HALF_N {
235 left[i] += right[i]; left[i] = Self::halve(left[i]); right[i] -= left[i]; }
239 }
240
241 #[inline(always)]
244 fn negacyclic_conv_n_recursive<const N: usize, const HALF_N: usize, NC>(
245 lhs: [T; N],
246 rhs: [U; N],
247 output: &mut [T],
248 inner_negacyclic_conv: NC,
249 ) where
250 NC: Fn([T; HALF_N], [U; HALF_N], &mut [T]),
251 {
252 debug_assert_eq!(2 * HALF_N, N);
253 let mut lhs_even = [Self::T_ZERO; HALF_N];
254 let mut lhs_odd = [Self::T_ZERO; HALF_N];
255 let mut lhs_sum = [Self::T_ZERO; HALF_N];
256 let mut rhs_even = [Self::U_ZERO; HALF_N];
257 let mut rhs_odd = [Self::U_ZERO; HALF_N];
258 let mut rhs_sum = [Self::U_ZERO; HALF_N];
259
260 for i in 0..HALF_N {
261 let s = lhs[2 * i];
262 let t = lhs[2 * i + 1];
263 lhs_even[i] = s;
264 lhs_odd[i] = t;
265 lhs_sum[i] = s + t;
266
267 let s = rhs[2 * i];
268 let t = rhs[2 * i + 1];
269 rhs_even[i] = s;
270 rhs_odd[i] = t;
271 rhs_sum[i] = s + t;
272 }
273
274 let mut even_s_conv = [Self::T_ZERO; HALF_N];
275 let (left, right) = output.split_at_mut(HALF_N);
276
277 inner_negacyclic_conv(lhs_even, rhs_even, &mut even_s_conv);
280 inner_negacyclic_conv(lhs_odd, rhs_odd, left);
281 inner_negacyclic_conv(lhs_sum, rhs_sum, right);
282
283 right[0] -= even_s_conv[0] + left[0];
286 even_s_conv[0] -= left[HALF_N - 1];
287
288 for i in 1..HALF_N {
289 right[i] -= even_s_conv[i] + left[i];
290 even_s_conv[i] += left[i - 1];
291 }
292
293 for i in 0..HALF_N {
295 output[2 * i] = even_s_conv[i];
296 output[2 * i + 1] = output[i + HALF_N];
297 }
298 }
299
300 #[inline(always)]
301 fn conv6(lhs: [T; 6], rhs: [U; 6], output: &mut [T]) {
302 Self::conv_n_recursive(lhs, rhs, output, Self::conv3, Self::negacyclic_conv3);
303 }
304
305 #[inline(always)]
306 fn negacyclic_conv6(lhs: [T; 6], rhs: [U; 6], output: &mut [T]) {
307 Self::negacyclic_conv_n_recursive(lhs, rhs, output, Self::negacyclic_conv3);
308 }
309
310 #[inline(always)]
311 fn conv8(lhs: [T; 8], rhs: [U; 8], output: &mut [T]) {
312 Self::conv_n_recursive(lhs, rhs, output, Self::conv4, Self::negacyclic_conv4);
313 }
314
315 #[inline(always)]
316 fn negacyclic_conv8(lhs: [T; 8], rhs: [U; 8], output: &mut [T]) {
317 Self::negacyclic_conv_n_recursive(lhs, rhs, output, Self::negacyclic_conv4);
318 }
319
320 #[inline(always)]
321 fn conv12(lhs: [T; 12], rhs: [U; 12], output: &mut [T]) {
322 Self::conv_n_recursive(lhs, rhs, output, Self::conv6, Self::negacyclic_conv6);
323 }
324
325 #[inline(always)]
326 fn negacyclic_conv12(lhs: [T; 12], rhs: [U; 12], output: &mut [T]) {
327 Self::negacyclic_conv_n_recursive(lhs, rhs, output, Self::negacyclic_conv6);
328 }
329
330 #[inline(always)]
331 fn conv16(lhs: [T; 16], rhs: [U; 16], output: &mut [T]) {
332 Self::conv_n_recursive(lhs, rhs, output, Self::conv8, Self::negacyclic_conv8);
333 }
334
335 #[inline(always)]
336 fn negacyclic_conv16(lhs: [T; 16], rhs: [U; 16], output: &mut [T]) {
337 Self::negacyclic_conv_n_recursive(lhs, rhs, output, Self::negacyclic_conv8);
338 }
339
340 #[inline(always)]
341 fn conv24(lhs: [T; 24], rhs: [U; 24], output: &mut [T]) {
342 Self::conv_n_recursive(lhs, rhs, output, Self::conv12, Self::negacyclic_conv12);
343 }
344
345 #[inline(always)]
346 fn conv32(lhs: [T; 32], rhs: [U; 32], output: &mut [T]) {
347 Self::conv_n_recursive(lhs, rhs, output, Self::conv16, Self::negacyclic_conv16);
348 }
349
350 #[inline(always)]
351 fn negacyclic_conv32(lhs: [T; 32], rhs: [U; 32], output: &mut [T]) {
352 Self::negacyclic_conv_n_recursive(lhs, rhs, output, Self::negacyclic_conv16);
353 }
354
355 #[inline(always)]
356 fn conv64(lhs: [T; 64], rhs: [U; 64], output: &mut [T]) {
357 Self::conv_n_recursive(lhs, rhs, output, Self::conv32, Self::negacyclic_conv32);
358 }
359}
360
361struct FieldConvolve<F, A>(PhantomData<(F, A)>);
366
367impl<F: Field, A: Algebra<F> + Copy> Convolve<A, A, F> for FieldConvolve<F, A> {
368 const T_ZERO: A = A::ZERO;
369 const U_ZERO: F = F::ZERO;
370
371 #[inline(always)]
372 fn halve(val: A) -> A {
373 val.halve()
374 }
375
376 #[inline(always)]
377 fn read(input: A) -> A {
378 input
379 }
380
381 #[inline(always)]
382 fn parity_dot<const N: usize>(lhs: [A; N], rhs: [F; N]) -> A {
383 A::mixed_dot_product(&lhs, &rhs)
384 }
385
386 #[inline(always)]
387 fn reduce(z: A) -> A {
388 z
389 }
390}
391
392#[inline]
394pub fn mds_circulant_karatsuba_8<F: Field, A: Algebra<F> + Copy>(state: &mut [A; 8], col: &[F; 8]) {
395 let input = *state;
396 FieldConvolve::<F, A>::conv8(input, *col, state.as_mut_slice());
397}
398
399#[inline]
401pub fn mds_circulant_karatsuba_12<F: Field, A: Algebra<F> + Copy>(
402 state: &mut [A; 12],
403 col: &[F; 12],
404) {
405 let input = *state;
406 FieldConvolve::<F, A>::conv12(input, *col, state.as_mut_slice());
407}
408
409#[inline]
411pub fn mds_circulant_karatsuba_16<F: Field, A: Algebra<F> + Copy>(
412 state: &mut [A; 16],
413 col: &[F; 16],
414) {
415 let input = *state;
416 FieldConvolve::<F, A>::conv16(input, *col, state.as_mut_slice());
417}
418
419#[inline]
421pub fn mds_circulant_karatsuba_24<F: Field, A: Algebra<F> + Copy>(
422 state: &mut [A; 24],
423 col: &[F; 24],
424) {
425 let input = *state;
426 FieldConvolve::<F, A>::conv24(input, *col, state.as_mut_slice());
427}
428
429#[cfg(test)]
430mod tests {
431 use p3_baby_bear::BabyBear;
432 use p3_field::PrimeCharacteristicRing;
433 use proptest::prelude::*;
434
435 use super::*;
436
437 type F = BabyBear;
438
439 fn arb_f() -> impl Strategy<Value = F> {
440 prop::num::u32::ANY.prop_map(F::from_u32)
441 }
442
443 fn naive_cyclic_conv<const N: usize>(lhs: [F; N], rhs: [F; N]) -> [F; N] {
444 core::array::from_fn(|i| {
446 let mut acc = F::ZERO;
447 for j in 0..N {
448 acc += lhs[j] * rhs[(N + i - j) % N];
449 }
450 acc
451 })
452 }
453
454 fn naive_negacyclic_conv<const N: usize>(lhs: [F; N], rhs: [F; N]) -> [F; N] {
455 let mut out = [F::ZERO; N];
458 for (i, &l) in lhs.iter().enumerate() {
459 for (j, &r) in rhs.iter().enumerate() {
460 let k = i + j;
461 if k < N {
462 out[k] += l * r;
463 } else {
464 out[k - N] -= l * r;
465 }
466 }
467 }
468 out
469 }
470
471 fn check_conv<const N: usize>(
472 lhs: [F; N],
473 rhs: [F; N],
474 conv_fn: fn([F; N], [F; N], &mut [F]),
475 naive_fn: fn([F; N], [F; N]) -> [F; N],
476 ) {
477 let expected = naive_fn(lhs, rhs);
478 let mut output = [F::ZERO; N];
479 conv_fn(lhs, rhs, &mut output);
480 assert_eq!(output, expected, "convolution mismatch");
481 }
482
483 macro_rules! conv_test {
484 ($name:ident, $n:expr, $conv:expr, $naive:expr, $arr:ident) => {
485 proptest! {
486 #[test]
487 fn $name(
488 lhs in prop::array::$arr(arb_f()),
489 rhs in prop::array::$arr(arb_f()),
490 ) {
491 check_conv::<$n>(lhs, rhs, $conv, $naive);
492 }
493 }
494 };
495 }
496
497 conv_test!(
499 conv3_matches_naive,
500 3,
501 FieldConvolve::<F, F>::conv3,
502 naive_cyclic_conv,
503 uniform3
504 );
505 conv_test!(
506 negacyclic_conv3_matches_naive,
507 3,
508 FieldConvolve::<F, F>::negacyclic_conv3,
509 naive_negacyclic_conv,
510 uniform3
511 );
512
513 conv_test!(
515 conv4_matches_naive,
516 4,
517 FieldConvolve::<F, F>::conv4,
518 naive_cyclic_conv,
519 uniform4
520 );
521 conv_test!(
522 negacyclic_conv4_matches_naive,
523 4,
524 FieldConvolve::<F, F>::negacyclic_conv4,
525 naive_negacyclic_conv,
526 uniform4
527 );
528
529 conv_test!(
531 conv6_matches_naive,
532 6,
533 FieldConvolve::<F, F>::conv6,
534 naive_cyclic_conv,
535 uniform6
536 );
537 conv_test!(
538 negacyclic_conv6_matches_naive,
539 6,
540 FieldConvolve::<F, F>::negacyclic_conv6,
541 naive_negacyclic_conv,
542 uniform6
543 );
544
545 conv_test!(
547 conv8_matches_naive,
548 8,
549 FieldConvolve::<F, F>::conv8,
550 naive_cyclic_conv,
551 uniform8
552 );
553 conv_test!(
554 negacyclic_conv8_matches_naive,
555 8,
556 FieldConvolve::<F, F>::negacyclic_conv8,
557 naive_negacyclic_conv,
558 uniform8
559 );
560
561 conv_test!(
563 conv12_matches_naive,
564 12,
565 FieldConvolve::<F, F>::conv12,
566 naive_cyclic_conv,
567 uniform12
568 );
569 conv_test!(
570 negacyclic_conv12_matches_naive,
571 12,
572 FieldConvolve::<F, F>::negacyclic_conv12,
573 naive_negacyclic_conv,
574 uniform12
575 );
576
577 conv_test!(
579 conv16_matches_naive,
580 16,
581 FieldConvolve::<F, F>::conv16,
582 naive_cyclic_conv,
583 uniform16
584 );
585 conv_test!(
586 negacyclic_conv16_matches_naive,
587 16,
588 FieldConvolve::<F, F>::negacyclic_conv16,
589 naive_negacyclic_conv,
590 uniform16
591 );
592
593 conv_test!(
595 conv24_matches_naive,
596 24,
597 FieldConvolve::<F, F>::conv24,
598 naive_cyclic_conv,
599 uniform24
600 );
601
602 conv_test!(
604 conv32_matches_naive,
605 32,
606 FieldConvolve::<F, F>::conv32,
607 naive_cyclic_conv,
608 uniform32
609 );
610 conv_test!(
611 negacyclic_conv32_matches_naive,
612 32,
613 FieldConvolve::<F, F>::negacyclic_conv32,
614 naive_negacyclic_conv,
615 uniform32
616 );
617
618 #[test]
619 fn conv64_matches_naive_fixed() {
620 let lhs: [F; 64] = core::array::from_fn(|i| F::from_u32(i as u32 + 1));
621 let rhs: [F; 64] = core::array::from_fn(|i| F::from_u32(64 - i as u32));
622 check_conv::<64>(lhs, rhs, FieldConvolve::<F, F>::conv64, naive_cyclic_conv);
623 }
624
625 #[test]
626 fn conv64_all_ones() {
627 let ones = [F::ONE; 64];
628 let expected = naive_cyclic_conv(ones, ones);
629 let mut output = [F::ZERO; 64];
630 FieldConvolve::<F, F>::conv64(ones, ones, &mut output);
631 assert_eq!(output, expected);
632 }
633
634 proptest! {
635 #[test]
636 fn karatsuba_16_matches_naive(
637 col in prop::array::uniform16(arb_f()),
638 state in prop::array::uniform16(arb_f()),
639 ) {
640 let expected = naive_cyclic_conv(state, col);
641 let mut actual = state;
642 mds_circulant_karatsuba_16(&mut actual, &col);
643 prop_assert_eq!(actual, expected);
644 }
645
646 #[test]
647 fn karatsuba_24_matches_naive(
648 col in prop::array::uniform24(arb_f()),
649 state in prop::array::uniform24(arb_f()),
650 ) {
651 let expected = naive_cyclic_conv(state, col);
652 let mut actual = state;
653 mds_circulant_karatsuba_24(&mut actual, &col);
654 prop_assert_eq!(actual, expected);
655 }
656 }
657
658 proptest! {
659 #[test]
660 fn conv8_commutative(
661 a in prop::array::uniform8(arb_f()),
662 b in prop::array::uniform8(arb_f()),
663 ) {
664 let mut ab = [F::ZERO; 8];
666 let mut ba = [F::ZERO; 8];
667 FieldConvolve::<F, F>::conv8(a, b, &mut ab);
668 FieldConvolve::<F, F>::conv8(b, a, &mut ba);
669 prop_assert_eq!(ab, ba);
670 }
671
672 #[test]
673 fn conv8_identity(a in prop::array::uniform8(arb_f())) {
674 let mut id = [F::ZERO; 8];
676 id[0] = F::ONE;
677 let mut out = [F::ZERO; 8];
678 FieldConvolve::<F, F>::conv8(a, id, &mut out);
679 prop_assert_eq!(out, a);
680 }
681
682 #[test]
683 fn conv8_zero(a in prop::array::uniform8(arb_f())) {
684 let zeros = [F::ZERO; 8];
686 let mut out = [F::ZERO; 8];
687 FieldConvolve::<F, F>::conv8(a, zeros, &mut out);
688 prop_assert_eq!(out, zeros);
689 }
690 }
691}