pub unsafe fn convolve_f64_avx2(
a: impl ExactSizeIterator<Item = f64>,
b: impl ExactSizeIterator<Item = f64>,
range: Range<usize>,
) -> Vec<f64>Examples found in repository?
crates/competitive/src/math/fast_fourier_transform.rs (lines 485-489)
457 fn convolve(a: Self::T, b: Self::T) -> Self::T {
458 let len = (a.len() + b.len()).saturating_sub(1);
459 // Keep accumulation overflow-free and exact in the FFT's f64 representation.
460 if (a.len().min(b.len()) <= 32 || {
461 let size = len.next_power_of_two();
462 let log = size.ilog2();
463 let limit = crate::avx_helper!(@dispatch simd_backend, SimdBackend;
464 3 * log + 16, 2 * log + 16, 4 * log + 16
465 );
466 2 * a.len() as u128 * b.len() as u128 <= limit as u128 * size as u128
467 }) && a.iter().map(|x| x.unsigned_abs()).max().unwrap_or(0) as u128
468 * b.iter().map(|x| x.unsigned_abs()).max().unwrap_or(0) as u128
469 <= (1u128 << 53) / a.len().min(b.len()).max(1) as u128
470 {
471 return crate::avx_helper!(@dispatch simd_backend, SimdBackend;
472 unsafe {
473 if a.len().max(b.len()) < 32 {
474 simd::convolve_i64_avx2(a, b, len)
475 } else {
476 simd::convolve_i64_avx512(a, b, len)
477 }
478 },
479 unsafe { simd::convolve_i64_avx2(a, b, len) },
480 convolve_i64_naive(a, b, len)
481 );
482 }
483 if !a.is_empty() && !b.is_empty() {
484 crate::avx_helper!(@dispatch_avx2_fma return unsafe {
485 simd::convolve_f64_avx2(
486 a.into_iter().map(|x| x as f64),
487 b.into_iter().map(|x| x as f64),
488 0..len,
489 )
490 .into_iter()
491 .map(|x| x.round() as i64)
492 .collect()
493 }, ());
494 }
495 let mut a = Self::transform(a, len);
496 let b = Self::transform(b, len);
497 Self::multiply(&mut a, &b);
498 Self::inverse_transform(a, len)
499 }
500 fn multiply(f: &mut Self::F, g: &Self::F) {
501 assert_eq!(f.len(), g.len());
502 f[0].re *= g[0].re;
503 f[0].im *= g[0].im;
504 for (f, g) in f.iter_mut().zip(g.iter()).skip(1) {
505 *f *= *g;
506 }
507 }
508}
509
510fn middle_product_f64_scalar(
511 a: impl ExactSizeIterator<Item = f64>,
512 b: impl ExactSizeIterator<Item = f64>,
513) -> Vec<f64> {
514 let a_len = a.len();
515 let b_len = b.len();
516 let len = a_len + b_len - 1;
517 let mut a = transform_real(a, len);
518 let b = transform_real(b, len);
519 ConvolveRealFft::multiply(&mut a, &b);
520 inverse_transform_real(a, len)[b_len - 1..a_len].to_vec()
521}
522
523impl ConvolveRealFft {
524 /// Returns coefficients `b.len() - 1..a.len()` of the convolution of `a` and `b`.
525 /// Panics unless `0 < b.len() <= a.len()`.
526 pub fn middle_product_f64(
527 a: impl ExactSizeIterator<Item = f64>,
528 b: impl ExactSizeIterator<Item = f64>,
529 ) -> Vec<f64> {
530 assert!(0 < b.len() && b.len() <= a.len());
531 crate::avx_helper!(@dispatch_avx2_fma return unsafe {
532 let range = b.len() - 1..a.len();
533 simd::convolve_f64_avx2(a, b, range)
534 }, ());
535 middle_product_f64_scalar(a, b)
536 }