pub fn advise_huge_pages<T>(_values: &mut Vec<T>)Examples found in repository?
More examples
crates/competitive/src/math/fast_fourier_transform.rs (line 292)
290 fn pack_f64(values: impl Iterator<Item = f64>, n: usize) -> Vec<Complex4> {
291 let mut result = Vec::with_capacity(n / 4);
292 advise_huge_pages(&mut result);
293 result.resize(n / 4, Complex4::default());
294 for (i, value) in values.enumerate() {
295 if i < n {
296 result[i >> 2].re[i & 3] = value;
297 } else {
298 result[(i - n) >> 2].im[i & 3] = value;
299 }
300 }
301 result
302 }crates/competitive/src/math/number_theoretic_transform.rs (line 591)
581fn convolve_large_ntt<M>(a: Vec<MInt<M>>, b: Vec<MInt<M>>) -> Vec<MInt<M>>
582where
583 M: Montgomery32NttModulus,
584{
585 let len = a.len() + b.len() - 1;
586 let ntt_len = 1usize << M::RANK;
587 let block_len = ntt_len / 2;
588 let same = a == b;
589 let transform = |a: &[MInt<M>]| {
590 let mut f = Vec::with_capacity(ntt_len);
591 advise_huge_pages(&mut f);
592 f.extend_from_slice(a);
593 Convolve::<M>::transform_ntt(f, ntt_len)
594 };
595 let fa: Vec<_> = a.chunks(block_len).map(transform).collect();
596 let fb: Option<Vec<_>> = if same {
597 None
598 } else {
599 Some(b.chunks(block_len).map(transform).collect())
600 };
601 let b_blocks = fb.as_ref().map_or(fa.len(), Vec::len);
602 let mut result = vec![MInt::<M>::zero(); len];
603 for diagonal in 0..fa.len() + b_blocks - 1 {
604 let mut spectrum = vec![MInt::<M>::zero(); ntt_len];
605 let start = diagonal.saturating_sub(b_blocks - 1);
606 for i in start..=diagonal.min(fa.len() - 1) {
607 let j = diagonal - i;
608 let g = if let Some(fb) = &fb { &fb[j] } else { &fa[j] };
609 pointwise_multiply_add(&mut spectrum, &fa[i], g);
610 }
611 spectrum = Convolve::<M>::inverse_transform_ntt(spectrum, ntt_len);
612 let offset = diagonal * block_len;
613 for (result, value) in result[offset..].iter_mut().zip(spectrum) {
614 *result += value;
615 }
616 }
617 result
618}
619
620impl<M> ConvolveSteps for Convolve<M>
621where
622 M: Montgomery32NttModulus,
623{
624 const CYCLIC: bool = true;
625
626 type T = Vec<MInt<M>>;
627 type F = Vec<MInt<M>>;
628 fn length(t: &Self::T) -> usize {
629 t.len()
630 }
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}crates/competitive/src/math/mint_fft_convolve.rs (line 43)
32unsafe fn split_coefficients<M>(
33 values: Vec<MInt<M>>,
34 n: usize,
35 modulus: i64,
36 split: i64,
37) -> (Vec<Complex4>, Vec<Complex4>)
38where
39 M: MIntConvert + MIntConvert<u32>,
40{
41 let mut low = Vec::<Complex4>::with_capacity(n / 4);
42 let mut high = Vec::<Complex4>::with_capacity(n / 4);
43 advise_huge_pages(&mut low);
44 advise_huge_pages(&mut high);
45 let divisor = _mm256_set1_pd(split as f64);
46 let split = _mm256_set1_pd(split as f64);
47 for i in (0..values.len()).step_by(4) {
48 let mut centered = [0i32; 4];
49 for lane in 0..4.min(values.len() - i) {
50 let mut value = u32::from(values[i + lane]) as i64;
51 if value * 2 > modulus {
52 value -= modulus;
53 }
54 centered[lane] = value as i32;
55 }
56 let value = _mm256_cvtepi32_pd(_mm_loadu_si128(centered.as_ptr().cast()));
57 let upper = _mm256_round_pd::<{ _MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC }>(
58 _mm256_div_pd(value, divisor),
59 );
60 let lower = _mm256_fnmadd_pd(upper, split, value);
61 if i < n {
62 let mut lo = Complex4::default();
63 let mut hi = Complex4::default();
64 _mm256_store_pd(lo.re.as_mut_ptr(), lower);
65 _mm256_store_pd(hi.re.as_mut_ptr(), upper);
66 low.push(lo);
67 high.push(hi);
68 } else {
69 _mm256_store_pd(low[(i - n) >> 2].im.as_mut_ptr(), lower);
70 _mm256_store_pd(high[(i - n) >> 2].im.as_mut_ptr(), upper);
71 }
72 }
73 low.resize(n / 4, Complex4::default());
74 high.resize(n / 4, Complex4::default());
75 (low, high)
76}
77
78#[target_feature(enable = "avx2,fma")]
79unsafe fn dot_soa(a0: &mut [Complex4], a1: &mut [Complex4], b0: &mut [Complex4], b1: &[Complex4]) {
80 let n = a0.len() * 4;
81 RotateCache::ensure(n / 2);
82 RotateCache::with(|cache| {
83 for i in 0..a0.len() {
84 let (mut cr, mut ci) = load4(&b0[i]);
85 let (mut dr, mut di) = load4(&b1[i]);
86 let mut c0r = _mm256_setzero_pd();
87 let mut c0i = _mm256_setzero_pd();
88 let mut c1r = _mm256_setzero_pd();
89 let mut c1i = _mm256_setzero_pd();
90 let mut c2r = _mm256_setzero_pd();
91 let mut c2i = _mm256_setzero_pd();
92 let w = eval_twiddle(cache, 1, a0.len(), i);
93 let wr = _mm256_setr_pd(w.re, 1.0, 1.0, 1.0);
94 let wi = _mm256_setr_pd(w.im, 0.0, 0.0, 0.0);
95 for lane in 0..4 {
96 let ar = _mm256_set1_pd(a0[i].re[lane]);
97 let ai = _mm256_set1_pd(a0[i].im[lane]);
98 let br = _mm256_set1_pd(a1[i].re[lane]);
99 let bi = _mm256_set1_pd(a1[i].im[lane]);
100 multiply_accumulate4(&mut c0r, &mut c0i, ar, ai, cr, ci);
101 multiply_accumulate4(&mut c1r, &mut c1i, ar, ai, dr, di);
102 multiply_accumulate4(&mut c1r, &mut c1i, br, bi, cr, ci);
103 multiply_accumulate4(&mut c2r, &mut c2i, br, bi, dr, di);
104 if lane != 3 {
105 cr = _mm256_permute4x64_pd::<0x93>(cr);
106 ci = _mm256_permute4x64_pd::<0x93>(ci);
107 dr = _mm256_permute4x64_pd::<0x93>(dr);
108 di = _mm256_permute4x64_pd::<0x93>(di);
109 (cr, ci) = mul4(cr, ci, wr, wi);
110 (dr, di) = mul4(dr, di, wr, wi);
111 }
112 }
113 store4(&mut a0[i], c0r, c0i);
114 store4(&mut a1[i], c1r, c1i);
115 store4(&mut b0[i], c2r, c2i);
116 }
117 });
118}
119
120#[target_feature(enable = "avx2,fma")]
121unsafe fn split_u64_coefficients(values: &[u64], n: usize) -> [Vec<Complex4>; 5] {
122 let mut result: [Vec<Complex4>; 5] = std::array::from_fn(|_| {
123 let mut part = Vec::with_capacity(n / 4);
124 advise_huge_pages(&mut part);
125 part
126 });
127 for (i, chunk) in values.chunks(4).enumerate() {
128 let mut parts = [Complex4::default(); 5];
129 for (lane, mut value) in chunk.iter().copied().enumerate() {
130 for part in &mut parts {
131 let digit = ((value << 51) as i64) >> 51;
132 value = (value >> 13).wrapping_add(u64::from(digit < 0));
133 part.re[lane] = digit as f64;
134 }
135 }
136 for (result, part) in result.iter_mut().zip(parts) {
137 if i < n / 4 {
138 result.push(part);
139 } else {
140 result[i - n / 4].im = part.re;
141 }
142 }
143 }
144 for part in &mut result {
145 part.resize(n / 4, Complex4::default());
146 }
147 result
148}crates/competitive/src/num/mint/simd_matrix.rs (line 238)
211 pub unsafe fn matrix_product_avx2(
212 a: &[Vec<Self>],
213 b: &[Vec<Self>],
214 scale: u32,
215 ) -> Vec<Vec<Self>> {
216 let (n, m, p) = (a.len(), b.len(), b.first().map_or(0, Vec::len));
217 assert!(a.iter().all(|row| row.len() == m));
218 assert!(b.iter().all(|row| row.len() == p));
219 let modulus = M::get_mod();
220 let alignment = if n.min(m).min(p) <= 64 { 8 } else { 32 };
221 let (nn, mm, pp) = (
222 n.div_ceil(alignment) * alignment,
223 m.div_ceil(alignment) * alignment,
224 p.div_ceil(alignment) * alignment,
225 );
226 let mut depth = 0;
227 let (mut x, mut y, mut z) = (nn, mm, pp);
228 while x.min(y).min(z) > 64 && x % 16 == 0 && y % 16 == 0 && z % 16 == 0 {
229 depth += 1;
230 x /= 2;
231 y /= 2;
232 z /= 2;
233 }
234 let entries = nn * mm + mm * pp + nn * pp;
235 // A recursive level uses one quarter of its parent's storage; siblings reuse it.
236 let mut data = if entries + entries / 3 >= 1 << 20 {
237 let mut data = Vec::with_capacity(entries + entries / 3);
238 advise_huge_pages(&mut data);
239 data.resize(entries + entries / 3, 0u32);
240 data
241 } else {
242 vec![0u32; entries + entries / 3]
243 };
244 let quotient = (((scale as u64) << 32) / modulus as u64) as u32;
245 let mut inverse = 1u32;
246 for _ in 0..5 {
247 inverse = inverse.wrapping_mul(2u32.wrapping_sub(modulus.wrapping_mul(inverse)));
248 }
249 let inverse = inverse.wrapping_neg();
250 for (offset, row, col) in blocks(nn, mm, depth) {
251 let (nr, nc) = (nn >> depth, mm >> depth);
252 for i in row..(row + nr).min(n) {
253 // SAFETY: MInt is transparent over u32. Reading raw words preserves Montgomery encoding.
254 let values: &[u32] = unsafe { std::slice::from_raw_parts(a[i].as_ptr().cast(), m) };
255 for j in col..(col + nc).min(m) {
256 let x = values[j];
257 let q = ((x as u64 * quotient as u64) >> 32) as u32;
258 let x = x.wrapping_mul(scale).wrapping_sub(q.wrapping_mul(modulus));
259 data[offset + (i - row) * nc + j - col] = x.min(x.wrapping_sub(modulus));
260 }
261 }
262 }
263 for (offset, row, col) in blocks(mm, pp, depth) {
264 let (nr, nc) = (mm >> depth, pp >> depth);
265 for i in row..(row + nr).min(m) {
266 // SAFETY: MInt is transparent over u32 and row lengths were checked above.
267 let values: &[u32] = unsafe { std::slice::from_raw_parts(b[i].as_ptr().cast(), p) };
268 for j in col..(col + nc).min(p) {
269 data[nn * mm + offset + (i - row) * nc + j - col] = values[j];
270 }
271 }
272 }
273 let kernel = Kernel {
274 modulus,
275 inverse,
276 avx512: avx512_enabled() && is_x86_feature_detected!("avx512f"),
277 };
278 // SAFETY: padding keeps every leaf dimension divisible by eight. The three matrices
279 // and the geometric scratch space are disjoint parts of the allocated buffer.
280 unsafe {
281 let ptr = data.as_mut_ptr();
282 multiply(
283 ptr,
284 ptr.add(nn * mm),
285 ptr.add(nn * mm + mm * pp),
286 (nn, mm, pp),
287 ptr.add(entries),
288 &kernel,
289 );
290 }
291 let mut result = vec![vec![MInt::new_unchecked(M::mod_zero()); p]; n];
292 for (offset, row, col) in blocks(nn, pp, depth) {
293 let (nr, nc) = (nn >> depth, pp >> depth);
294 for i in row..(row + nr).min(n) {
295 for j in col..(col + nc).min(p) {
296 result[i][j] = MInt::new_unchecked(
297 data[nn * mm + mm * pp + offset + (i - row) * nc + j - col],
298 );
299 }
300 }
301 }
302 result
303 }