fn use_block_ntt<M>(len: usize) -> boolwhere
M: Montgomery32NttModulus,Examples found in repository?
crates/competitive/src/math/number_theoretic_transform.rs (line 634)
631 fn transform(mut t: Self::T, len: usize) -> Self::F {
632 t.resize_with(len.max(1).next_power_of_two(), Zero::zero);
633 #[cfg(target_arch = "x86_64")]
634 if use_block_ntt::<M>(t.len()) {
635 unsafe { ntt_simd::transform_blocks_avx2(&mut t) };
636 return t;
637 }
638 ntt(&mut t);
639 t
640 }
641 fn inverse_transform(mut f: Self::F, len: usize) -> Self::T {
642 #[cfg(target_arch = "x86_64")]
643 if use_block_ntt::<M>(f.len()) {
644 unsafe { ntt_simd::inverse_transform_blocks_avx2(&mut f) };
645 f.truncate(len);
646 return f;
647 }
648 intt(&mut f);
649 f.truncate(len);
650 f
651 }
652 fn multiply(f: &mut Self::F, g: &Self::F) {
653 assert_eq!(f.len(), g.len());
654 #[cfg(target_arch = "x86_64")]
655 if use_block_ntt::<M>(f.len()) {
656 unsafe { ntt_simd::multiply_blocks_avx2(f, g) };
657 return;
658 }
659 pointwise_multiply(f, g);
660 }
661 fn square(t: Self::T, len: usize) -> Self::T {
662 let mut f = Self::transform(t, len);
663 let g = f.clone();
664 Self::multiply(&mut f, &g);
665 Self::inverse_transform(f, len)
666 }
667 fn convolve(mut a: Self::T, mut b: Self::T) -> Self::T {
668 let (threshold, naive_threshold) = (100, 60);
669 #[cfg(target_arch = "x86_64")]
670 let (threshold, naive_threshold) = if use_block_ntt::<M>(64) {
671 (
672 60,
673 if M::RANK >= 13 && a.len().max(b.len()) <= 4096 {
674 18
675 } else if M::RANK >= 19 && a.len().max(b.len()) <= 262144 {
676 32
677 } else {
678 34
679 },
680 )
681 } else {
682 (threshold, naive_threshold)
683 };
684 if Self::length(&a).max(Self::length(&b)) <= threshold {
685 return convolve_karatsuba(&a, &b);
686 }
687 if Self::length(&a).min(Self::length(&b)) <= naive_threshold {
688 return convolve_naive(&a, &b);
689 }
690 let len = (Self::length(&a) + Self::length(&b)).saturating_sub(1);
691 let size = len.max(1).next_power_of_two();
692 let max_size = 1usize << M::RANK;
693 #[cfg(target_arch = "x86_64")]
694 let max_size = if use_block_ntt::<M>(size) {
695 max_size << 3
696 } else {
697 max_size
698 };
699 if size > max_size {
700 return convolve_large_ntt(a, b);
701 }
702 if len <= size / 2 + 2 {
703 let xa = a.pop().unwrap();
704 let xb = b.pop().unwrap();
705 let mut c = vec![MInt::<M>::zero(); len];
706 *c.last_mut().unwrap() = xa * xb;
707 for (a, c) in a.iter().zip(&mut c[b.len()..]) {
708 *c += *a * xb;
709 }
710 for (b, c) in b.iter().zip(&mut c[a.len()..]) {
711 *c += *b * xa;
712 }
713 let d = Self::convolve(a, b);
714 for (d, c) in d.into_iter().zip(&mut c) {
715 *c += d;
716 }
717 return c;
718 }
719 let same = a == b;
720 #[cfg(target_arch = "x86_64")]
721 if use_block_ntt::<M>(size) {
722 a.reserve(size - a.len());
723 b.reserve(size - b.len());
724 advise_huge_pages(&mut a);
725 advise_huge_pages(&mut b);
726 a.resize_with(size, Zero::zero);
727 b.resize_with(size, Zero::zero);
728 unsafe { ntt_simd::convolve_blocks_avx2(&mut a, &mut b, same) };
729 a.truncate(len);
730 return a;
731 }
732 let mut a = Self::transform(a, len);
733 if same {
734 for a in a.iter_mut() {
735 *a *= *a;
736 }
737 } else {
738 let b = Self::transform(b, len);
739 Self::multiply(&mut a, &b);
740 }
741 Self::inverse_transform(a, len)
742 }
743}
744
745type MVec<M> = Vec<MInt<M>>;
746
747fn convert_crt_input<M, N1, N2, N3>(t: MVec<M>, capacity: usize) -> (MVec<N1>, MVec<N2>, MVec<N3>)
748where
749 M: MIntConvert<u32>,
750 N1: Montgomery32NttModulus,
751 N2: Montgomery32NttModulus,
752 N3: Montgomery32NttModulus,
753{
754 let mut f = (
755 MVec::<N1>::with_capacity(capacity),
756 MVec::<N2>::with_capacity(capacity),
757 MVec::<N3>::with_capacity(capacity),
758 );
759 advise_huge_pages(&mut f.0);
760 advise_huge_pages(&mut f.1);
761 advise_huge_pages(&mut f.2);
762 for t in t {
763 let t: u32 = t.into();
764 f.0.push(t.into());
765 f.1.push(t.into());
766 f.2.push(t.into());
767 }
768 f
769}
770
771fn reconstruct_mint_crt<M, N1, N2, N3>(f: (MVec<N1>, MVec<N2>, MVec<N3>)) -> MVec<M>
772where
773 M: MIntConvert + MIntConvert<u32>,
774 N1: Montgomery32NttModulus,
775 N2: Montgomery32NttModulus,
776 N3: Montgomery32NttModulus,
777{
778 let t1 = MInt::<N2>::new(N1::get_mod()).inv();
779 let m1_3 = MInt::<N3>::new(N1::get_mod());
780 let t2 = (m1_3 * MInt::<N3>::new(N2::get_mod())).inv();
781 let modulus = <M as MIntConvert<u32>>::mod_into() as u64;
782 let m1 = N1::get_mod() as u64;
783 let m2 = m1 * N2::get_mod() as u64 % modulus;
784 let fits_u64 = (N1::get_mod() - 1) as u128
785 + (N2::get_mod() - 1) as u128 * m1 as u128
786 + (N3::get_mod() - 1) as u128 * m2 as u128
787 <= u64::MAX as u128;
788 f.0.into_iter()
789 .zip(f.1)
790 .zip(f.2)
791 .map(|((c1, c2), c3)| {
792 let d1 = c1.inner();
793 let d2 = ((c2 - MInt::<N2>::from(d1)) * t1).inner();
794 let x = MInt::<N3>::new(d1) + MInt::<N3>::new(d2) * m1_3;
795 let d3 = ((c3 - x) * t2).inner();
796 let value = if fits_u64 {
797 (d1 as u64 + d2 as u64 * m1 + d3 as u64 * m2) % modulus
798 } else {
799 ((d1 as u128 + d2 as u128 * m1 as u128 + d3 as u128 * m2 as u128) % modulus as u128)
800 as u64
801 };
802 MInt::<M>::from(value as u32)
803 })
804 .collect()
805}
806
807impl<M, N1, N2, N3> ConvolveSteps for Convolve<(M, (N1, N2, N3))>
808where
809 M: MIntConvert + MIntConvert<u32>,
810 N1: Montgomery32NttModulus,
811 N2: Montgomery32NttModulus,
812 N3: Montgomery32NttModulus,
813{
814 type T = MVec<M>;
815 type F = (MVec<N1>, MVec<N2>, MVec<N3>);
816 fn length(t: &Self::T) -> usize {
817 t.len()
818 }
819 fn transform(t: Self::T, len: usize) -> Self::F {
820 let npot = len.max(1).next_power_of_two();
821 let f = convert_crt_input(t, npot);
822 (
823 Convolve::<N1>::transform(f.0, npot),
824 Convolve::<N2>::transform(f.1, npot),
825 Convolve::<N3>::transform(f.2, npot),
826 )
827 }
828 fn inverse_transform(f: Self::F, len: usize) -> Self::T {
829 reconstruct_mint_crt((
830 Convolve::<N1>::inverse_transform(f.0, len),
831 Convolve::<N2>::inverse_transform(f.1, len),
832 Convolve::<N3>::inverse_transform(f.2, len),
833 ))
834 }
835 fn multiply(f: &mut Self::F, g: &Self::F) {
836 Convolve::<N1>::multiply(&mut f.0, &g.0);
837 Convolve::<N2>::multiply(&mut f.1, &g.1);
838 Convolve::<N3>::multiply(&mut f.2, &g.2);
839 }
840 fn convolve(a: Self::T, b: Self::T) -> Self::T {
841 let max_len = Self::length(&a).max(Self::length(&b));
842 let min_len = Self::length(&a).min(Self::length(&b));
843 let (balanced, short) = crate::avx_helper!(@dispatch_avx2_fma (30, 10), (384, 128));
844 if max_len <= balanced || min_len <= short {
845 return convolve_karatsuba(&a, &b);
846 }
847 // Limit coefficient growth to leave headroom for FFT roundoff.
848 let fft_limit = crate::avx_helper!(@dispatch_avx2_fma
849 1usize << ((1u64 << 50) / <M as MIntConvert<u32>>::mod_into() as u64).ilog2().min(20), 0);
850 let convolve = |a: Self::T, b: Self::T| {
851 let fft_len = (a.len() + b.len() - 1).next_power_of_two();
852 if fft_len <= 256 && a.len() * b.len() <= fft_len * 8 {
853 return convolve_karatsuba(&a, &b);
854 }
855 if fft_len <= fft_limit {
856 crate::avx_helper!(@dispatch_avx2_fma return unsafe {
857 convolve_mint_avx2(a, b)
858 }, ());
859 }
860 convolve_mint_crt::<M, N1, N2, N3>(a, b)
861 };
862 let block_len = min_len.next_power_of_two() * 8 - min_len + 1;
863 let block_len = if min_len <= fft_limit / 2 {
864 block_len.min(fft_limit - min_len + 1)
865 } else {
866 block_len
867 };
868 if max_len <= block_len {
869 return convolve(a, b);
870 }
871 let (a, b) = if a.len() >= b.len() { (a, b) } else { (b, a) };
872 let mut result = vec![MInt::<M>::zero(); a.len() + b.len() - 1];
873 for (i, a) in a.chunks(block_len).enumerate() {
874 let product = convolve(a.to_vec(), b.clone());
875 for (value, product) in result[i * block_len..].iter_mut().zip(product) {
876 *value += product;
877 }
878 }
879 result
880 }
881}
882
883fn convolve_mint_crt<M, N1, N2, N3>(a: MVec<M>, b: MVec<M>) -> MVec<M>
884where
885 M: MIntConvert + MIntConvert<u32>,
886 N1: Montgomery32NttModulus,
887 N2: Montgomery32NttModulus,
888 N3: Montgomery32NttModulus,
889{
890 let convolve = |a: MVec<M>, b: MVec<M>| {
891 let a_len = a.len();
892 let b_len = b.len();
893 let a = convert_crt_input(a, a_len);
894 let b = convert_crt_input(b, b_len);
895 reconstruct_mint_crt((
896 Convolve::<N1>::convolve(a.0, b.0),
897 Convolve::<N2>::convolve(a.1, b.1),
898 Convolve::<N3>::convolve(a.2, b.2),
899 ))
900 };
901 let modulus = <M as MIntConvert<u32>>::mod_into() as u128;
902 let capacity = N1::MOD as u128 * N2::MOD as u128 * N3::MOD as u128;
903 if a.len().min(b.len()) as u128 * (modulus - 1).pow(2) < capacity {
904 return convolve(a, b);
905 }
906 let block_len = ((capacity - 1) / (modulus - 1).pow(2)) as usize;
907 if block_len == 0 {
908 return convolve_naive(&a, &b);
909 }
910 let mut result = vec![MInt::<M>::zero(); a.len() + b.len() - 1];
911 for (i, a) in a.chunks(block_len).enumerate() {
912 for (j, b) in b.chunks(block_len).enumerate() {
913 let product = convolve(a.to_vec(), b.to_vec());
914 for (value, product) in result[(i + j) * block_len..].iter_mut().zip(product) {
915 *value += product;
916 }
917 }
918 }
919 result
920}
921
922impl<N1, N2, N3> ConvolveSteps for Convolve<(u64, (N1, N2, N3))>
923where
924 N1: Montgomery32NttModulus,
925 N2: Montgomery32NttModulus,
926 N3: Montgomery32NttModulus,
927{
928 type T = Vec<u64>;
929 type F = ([MVec<N1>; 3], [MVec<N2>; 3], [MVec<N3>; 3]);
930
931 fn length(t: &Self::T) -> usize {
932 t.len()
933 }
934
935 fn transform(t: Self::T, len: usize) -> Self::F {
936 let npot = len.max(1).next_power_of_two();
937 assert!(npot <= 1usize << N1::RANK.min(N2::RANK).min(N3::RANK));
938 // The 22-bit fallback needs room for three limb products per coefficient.
939 assert!(
940 3 * npot as u128 * ((1u128 << 22) - 1).pow(2)
941 < N1::MOD as u128 * N2::MOD as u128 * N3::MOD as u128
942 );
943 let bits = if 2 * npot as u128 * (u32::MAX as u128).pow(2)
944 < N1::MOD as u128 * N2::MOD as u128 * N3::MOD as u128
945 {
946 32
947 } else {
948 22
949 };
950 let parts = if bits == 32 && t.iter().all(|&value| value <= u32::MAX as u64) {
951 1
952 } else {
953 64usize.div_ceil(bits)
954 };
955 fn split<M: Montgomery32NttModulus>(
956 t: &[u64],
957 len: usize,
958 bits: usize,
959 parts: usize,
960 ) -> [MVec<M>; 3] {
961 std::array::from_fn(|part| {
962 if part >= parts {
963 return Vec::new();
964 }
965 Convolve::<M>::transform(
966 t.iter()
967 .map(|&t| MInt::from((t >> (part * bits)) & ((1u64 << bits) - 1)))
968 .collect(),
969 len,
970 )
971 })
972 }
973 (
974 split(&t, npot, bits, parts),
975 split(&t, npot, bits, parts),
976 split(&t, npot, bits, parts),
977 )
978 }
979
980 fn inverse_transform(f: Self::F, len: usize) -> Self::T {
981 let bits = if f.0[2].is_empty() { 32 } else { 22 };
982 let t1 = MInt::<N2>::new(N1::get_mod()).inv();
983 let m1 = N1::get_mod() as u64;
984 let m1_3 = MInt::<N3>::new(N1::get_mod());
985 let t2 = (m1_3 * MInt::<N3>::new(N2::get_mod())).inv();
986 let m2 = m1 * N2::get_mod() as u64;
987 let mut result = vec![0u64; len.min(f.0[0].len())];
988 for (part, ((f1, f2), f3)) in f.0.into_iter().zip(f.1).zip(f.2).enumerate() {
989 if f1.is_empty() {
990 continue;
991 }
992 for (value, ((c1, c2), c3)) in result.iter_mut().zip(
993 Convolve::<N1>::inverse_transform(f1, len)
994 .into_iter()
995 .zip(Convolve::<N2>::inverse_transform(f2, len))
996 .zip(Convolve::<N3>::inverse_transform(f3, len)),
997 ) {
998 let d1 = c1.inner();
999 let d2 = ((c2 - MInt::<N2>::from(d1)) * t1).inner();
1000 let x = MInt::<N3>::new(d1) + MInt::<N3>::new(d2) * m1_3;
1001 let d3 = ((c3 - x) * t2).inner();
1002 let limb = (d1 as u64)
1003 .wrapping_add((d2 as u64).wrapping_mul(m1))
1004 .wrapping_add((d3 as u64).wrapping_mul(m2));
1005 *value = value.wrapping_add(limb << (part * bits));
1006 }
1007 }
1008 result
1009 }
1010
1011 fn multiply(f: &mut Self::F, g: &Self::F) {
1012 fn multiply<M: Montgomery32NttModulus>(f: &mut [MVec<M>; 3], g: &[MVec<M>; 3]) {
1013 assert_eq!(f[0].len(), g[0].len());
1014 if f[1].is_empty() || g[1].is_empty() {
1015 if f[1].is_empty() && !g[1].is_empty() {
1016 f[1] = f[0].clone();
1017 Convolve::<M>::multiply(&mut f[1], &g[1]);
1018 } else if !f[1].is_empty() {
1019 Convolve::<M>::multiply(&mut f[1], &g[0]);
1020 }
1021 Convolve::<M>::multiply(&mut f[0], &g[0]);
1022 return;
1023 }
1024 #[cfg(target_arch = "x86_64")]
1025 if use_block_ntt::<M>(f[0].len()) {
1026 for part in (1..if f[2].is_empty() { 2 } else { 3 }).rev() {
1027 let mut sum = f[0].clone();
1028 Convolve::<M>::multiply(&mut sum, &g[part]);
1029 for left in 1..=part {
1030 let mut product = f[left].clone();
1031 Convolve::<M>::multiply(&mut product, &g[part - left]);
1032 for (value, product) in sum.iter_mut().zip(product) {
1033 // Block products contain lazy Montgomery residues.
1034 *value = MInt::new(value.inner() + product.inner());
1035 }
1036 }
1037 f[part] = sum;
1038 }
1039 Convolve::<M>::multiply(&mut f[0], &g[0]);
1040 return;
1041 }
1042 if f[2].is_empty() {
1043 for i in 0..f[0].len() {
1044 f[1][i] = f[0][i] * g[1][i] + f[1][i] * g[0][i];
1045 f[0][i] *= g[0][i];
1046 }
1047 return;
1048 }
1049 for i in 0..f[0].len() {
1050 f[2][i] = f[0][i] * g[2][i] + f[1][i] * g[1][i] + f[2][i] * g[0][i];
1051 f[1][i] = f[0][i] * g[1][i] + f[1][i] * g[0][i];
1052 f[0][i] *= g[0][i];
1053 }
1054 }