1use std::convert::TryInto;
2use std::sync::Arc;
3use std::{any::TypeId, arch::x86_64::*};
4
5use num_complex::Complex;
6use num_integer::{div_ceil, Integer};
7use num_traits::Zero;
8use primal_check::miller_rabin;
9use strength_reduce::StrengthReducedUsize;
10
11use crate::{array_utils, FftDirection};
12use crate::{math_utils, twiddles};
13use crate::{Direction, Fft, FftNum, Length};
14
15use super::avx_vector;
16use super::{
17 avx_vector::{AvxArray, AvxArrayMut, AvxVector, AvxVector128, AvxVector256},
18 AvxNum,
19};
20
21#[derive(Clone)]
24struct VectorizedMultiplyMod {
25 b: __m256i,
26 divisor: __m256i,
27 intermediate: __m256i,
28}
29
30impl VectorizedMultiplyMod {
31 #[target_feature(enable = "avx")]
32 unsafe fn new(b: u32, divisor: u32) -> Self {
33 assert!(
34 divisor.leading_zeros() > 0,
35 "divisor must be less than {}, got {}",
36 1 << 31,
37 divisor
38 );
39
40 let b = b % divisor;
41 let intermediate = ((b as i64) << 32) / divisor as i64;
42
43 Self {
44 b: _mm256_set1_epi64x(b as i64),
45 divisor: _mm256_set1_epi64x(divisor as i64),
46 intermediate: _mm256_set1_epi64x(intermediate),
47 }
48 }
49
50 #[allow(unused)]
53 #[inline(always)]
54 unsafe fn mul_rem(&self, a: __m256i) -> __m256i {
55 let masked_divisor = _mm256_blend_epi32(self.divisor, _mm256_setzero_si256(), 0xAA);
60
61 let quotient = _mm256_srli_epi64(_mm256_mul_epu32(a, self.intermediate), 32);
63
64 let numerator = _mm256_mul_epu32(a, self.b);
66 let quotient_product = _mm256_mul_epu32(quotient, masked_divisor);
67
68 let remainder = _mm256_sub_epi64(numerator, quotient_product);
70
71 let casted_remainder = _mm256_castsi256_pd(remainder);
76 let subtracted_remainder = _mm256_castsi256_pd(_mm256_sub_epi64(remainder, masked_divisor));
77 let wrapped_remainder = _mm256_castpd_si256(_mm256_blendv_pd(
78 subtracted_remainder,
79 casted_remainder,
80 subtracted_remainder,
81 ));
82 wrapped_remainder
83 }
84}
85
86pub struct RadersAvx2<A: AvxNum, T> {
97 input_index_multiplier: VectorizedMultiplyMod,
98 input_index_init: __m256i,
99
100 output_index_mapping: Box<[__m128i]>,
101 twiddles: Box<[A::VectorType]>,
102
103 inner_fft: Arc<dyn Fft<T>>,
104
105 len: usize,
106
107 inplace_scratch_len: usize,
108 outofplace_scratch_len: usize,
109 immut_scratch_len: usize,
110 direction: FftDirection,
111
112 _phantom: std::marker::PhantomData<T>,
113}
114
115impl<A: AvxNum, T: FftNum> RadersAvx2<A, T> {
116 #[inline]
122 pub fn new(inner_fft: Arc<dyn Fft<T>>) -> Result<Self, ()> {
123 let id_a = TypeId::of::<A>();
127 let id_t = TypeId::of::<T>();
128 assert_eq!(id_a, id_t);
129
130 let has_avx = is_x86_feature_detected!("avx");
131 let has_avx2 = is_x86_feature_detected!("avx2");
132 let has_fma = is_x86_feature_detected!("fma");
133 if has_avx && has_avx2 && has_fma {
134 Ok(unsafe { Self::new_with_avx(inner_fft) })
136 } else {
137 Err(())
138 }
139 }
140
141 #[target_feature(enable = "avx")]
142 unsafe fn new_with_avx(inner_fft: Arc<dyn Fft<T>>) -> Self {
143 let inner_fft_len = inner_fft.len();
144 let len = inner_fft_len + 1;
145 assert!(miller_rabin(len as u64), "For raders algorithm, inner_fft.len() + 1 must be prime. Expected prime number, got {} + 1 = {}", inner_fft_len, len);
146
147 let direction = inner_fft.fft_direction();
148 let reduced_len = StrengthReducedUsize::new(len);
149
150 let primitive_root = math_utils::primitive_root(len as u64).unwrap() as usize;
152
153 let gcd_data = i64::extended_gcd(&(primitive_root as i64), &(len as i64));
157 let primitive_root_inverse = if gcd_data.x >= 0 {
158 gcd_data.x
159 } else {
160 gcd_data.x + len as i64
161 } as usize;
162
163 let inner_fft_scale = T::one() / T::from_usize(inner_fft_len).unwrap();
165 let mut inner_fft_input = vec![Complex::zero(); inner_fft_len];
166 let mut twiddle_input = 1;
167 for input_cell in &mut inner_fft_input {
168 let twiddle = twiddles::compute_twiddle(twiddle_input, len, direction);
169 *input_cell = twiddle * inner_fft_scale;
170
171 twiddle_input = (twiddle_input * primitive_root_inverse) % reduced_len;
172 }
173
174 let required_inner_scratch = inner_fft.get_inplace_scratch_len();
175 let extra_inner_scratch = if required_inner_scratch <= inner_fft_len {
176 0
177 } else {
178 required_inner_scratch
179 };
180
181 let mut inner_fft_scratch = vec![Zero::zero(); required_inner_scratch];
183 inner_fft.process_with_scratch(&mut inner_fft_input, &mut inner_fft_scratch);
184
185 let conjugation_mask =
187 AvxVector256::broadcast_complex_elements(Complex::new(A::zero(), -A::zero()));
188
189 let inner_fft_multiplier: Box<[_]> = {
190 let transmuted_inner_input: &mut [Complex<A>] =
192 array_utils::workaround_transmute_mut(&mut inner_fft_input);
193
194 transmuted_inner_input
195 .chunks(A::VectorType::COMPLEX_PER_VECTOR)
196 .map(|chunk| {
197 let chunk_vector = match chunk.len() {
198 1 => chunk.load_partial1_complex(0).zero_extend(),
199 2 => {
200 if chunk.len() == A::VectorType::COMPLEX_PER_VECTOR {
201 chunk.load_complex(0)
202 } else {
203 chunk.load_partial2_complex(0).zero_extend()
204 }
205 }
206 3 => chunk.load_partial3_complex(0),
207 4 => chunk.load_complex(0),
208 _ => unreachable!(),
209 };
210 AvxVector::xor(chunk_vector, conjugation_mask) })
212 .collect()
213 };
214
215 const NUM_POWERS: usize = 5;
217 let mut root_powers = [0; NUM_POWERS];
218 let mut current_power = 1;
219 for i in 0..NUM_POWERS {
220 root_powers[i] = current_power;
221 current_power = (current_power * primitive_root) % reduced_len;
222 }
223
224 let (input_index_multiplier, input_index_init) = if A::VectorType::COMPLEX_PER_VECTOR == 4 {
225 (
226 VectorizedMultiplyMod::new(root_powers[4] as u32, len as u32),
227 _mm256_loadu_si256(root_powers.as_ptr().add(1) as *const __m256i),
228 )
229 } else {
230 let duplicated_powers = [
231 root_powers[1],
232 root_powers[1],
233 root_powers[2],
234 root_powers[2],
235 ];
236 (
237 VectorizedMultiplyMod::new(root_powers[2] as u32, len as u32),
238 _mm256_loadu_si256(duplicated_powers.as_ptr() as *const __m256i),
239 )
240 };
241
242 let mapping_size = 1 + div_ceil(len, A::VectorType::COMPLEX_PER_VECTOR)
247 * A::VectorType::COMPLEX_PER_VECTOR;
248 let mut output_mapping_inverse: Vec<i32> = vec![0; mapping_size];
249 let mut output_index = 1;
250 for i in 1..len {
251 output_index = (output_index * primitive_root_inverse) % reduced_len;
252 output_mapping_inverse[output_index] = i.try_into().unwrap();
253 }
254
255 let output_index_mapping = if A::VectorType::COMPLEX_PER_VECTOR == 4 {
257 (&output_mapping_inverse[1..])
258 .chunks_exact(A::VectorType::COMPLEX_PER_VECTOR)
259 .map(|chunk| _mm_loadu_si128(chunk.as_ptr() as *const __m128i))
260 .collect::<Box<[__m128i]>>()
261 } else {
262 (&output_mapping_inverse[1..])
263 .chunks_exact(A::VectorType::COMPLEX_PER_VECTOR)
264 .map(|chunk| {
265 let duplicated_indexes = [chunk[0], chunk[0], chunk[1], chunk[1]];
266 _mm_loadu_si128(duplicated_indexes.as_ptr() as *const __m128i)
267 })
268 .collect::<Box<[__m128i]>>()
269 };
270 let inplace_scratch_len = len + extra_inner_scratch;
271 Self {
272 input_index_multiplier,
273 input_index_init,
274
275 output_index_mapping,
276
277 inner_fft: inner_fft,
278 twiddles: inner_fft_multiplier,
279
280 len,
281
282 inplace_scratch_len,
283 outofplace_scratch_len: extra_inner_scratch,
284 immut_scratch_len: inner_fft_len + required_inner_scratch + 1,
285 direction,
286
287 _phantom: std::marker::PhantomData,
288 }
289 }
290
291 #[target_feature(enable = "avx2", enable = "avx", enable = "fma")]
293 unsafe fn prepare_raders(&self, input: &[Complex<A>], output: &mut [Complex<A>]) {
294 let mut indexes = self.input_index_init;
295
296 let index_multiplier = self.input_index_multiplier.clone();
297
298 let mut chunks_iter =
300 (&mut output[1..]).chunks_exact_mut(A::VectorType::COMPLEX_PER_VECTOR);
301 for mut chunk in chunks_iter.by_ref() {
302 let gathered_elements =
303 A::VectorType::gather_complex_avx2_index64(input.as_ptr(), indexes);
304
305 indexes = index_multiplier.mul_rem(indexes);
307
308 chunk.store_complex(gathered_elements, 0);
310 }
311
312 let mut output_remainder = chunks_iter.into_remainder();
314 if output_remainder.len() == 2 {
315 let half_data = AvxVector128::gather64_complex_avx2(
316 input.as_ptr(),
317 _mm256_castsi256_si128(indexes),
318 );
319
320 output_remainder.store_partial2_complex(half_data, 0);
322 }
323 }
324
325 #[target_feature(enable = "avx2", enable = "avx", enable = "fma")]
327 unsafe fn finalize_raders(&self, input: &[Complex<A>], output: &mut [Complex<A>]) {
328 let conjugation_mask =
330 AvxVector256::broadcast_complex_elements(Complex::new(A::zero(), -A::zero()));
331
332 let mut chunks_iter =
333 (&mut output[1..]).chunks_exact_mut(A::VectorType::COMPLEX_PER_VECTOR);
334 for (i, mut chunk) in chunks_iter.by_ref().enumerate() {
335 let index_chunk = *self.output_index_mapping.get_unchecked(i);
336 let gathered_elements =
337 A::VectorType::gather_complex_avx2_index32(input.as_ptr(), index_chunk);
338
339 let conjugated_elements = AvxVector::xor(gathered_elements, conjugation_mask);
340 chunk.store_complex(conjugated_elements, 0);
341 }
342
343 let mut output_remainder = chunks_iter.into_remainder();
345 if output_remainder.len() == 2 {
346 let index_chunk = *self
347 .output_index_mapping
348 .get_unchecked(self.output_index_mapping.len() - 1);
349 let half_data = AvxVector128::gather32_complex_avx2(input.as_ptr(), index_chunk);
350
351 let conjugated_elements = AvxVector::xor(half_data, conjugation_mask.lo());
352 output_remainder.store_partial2_complex(conjugated_elements, 0);
353 }
354 }
355
356 fn perform_fft_immut(
357 &self,
358 input: &[Complex<T>],
359 output: &mut [Complex<T>],
360 scratch: &mut [Complex<T>],
361 ) {
362 unsafe {
363 let transmuted_input: &[Complex<A>] = array_utils::workaround_transmute(input);
365 let transmuted_output: &mut [Complex<A>] =
366 array_utils::workaround_transmute_mut(output);
367 self.prepare_raders(transmuted_input, transmuted_output)
368 }
369
370 let (first_input, _) = input.split_first().unwrap();
371 let (first_output, inner_output) = output.split_first_mut().unwrap();
372 let (scratch2, extra_scratch) = scratch.split_at_mut(self.len());
373 let (_, scratch) = scratch2.split_first_mut().unwrap();
374
375 self.inner_fft.process_with_scratch(inner_output, scratch);
376
377 *first_output = inner_output[0] + *first_input;
379
380 unsafe {
384 let transmuted_inner_input: &mut [Complex<A>] =
386 array_utils::workaround_transmute_mut(scratch);
387 let transmuted_inner_output: &mut [Complex<A>] =
388 array_utils::workaround_transmute_mut(inner_output);
389 avx_vector::pairwise_complex_mul_conjugated(
390 transmuted_inner_output,
391 transmuted_inner_input,
392 &self.twiddles,
393 )
394 };
395
396 scratch[0] = scratch[0] + first_input.conj();
399
400 self.inner_fft.process_with_scratch(scratch, extra_scratch);
401 scratch2[0] = *first_input;
402
403 unsafe {
405 let transmuted_input: &mut [Complex<A>] =
407 array_utils::workaround_transmute_mut(scratch2);
408 let transmuted_output: &mut [Complex<A>] =
409 array_utils::workaround_transmute_mut(output);
410 self.finalize_raders(transmuted_input, transmuted_output);
411 }
412 }
413
414 fn perform_fft_out_of_place(
415 &self,
416 input: &mut [Complex<T>],
417 output: &mut [Complex<T>],
418 scratch: &mut [Complex<T>],
419 ) {
420 unsafe {
421 let transmuted_input: &mut [Complex<A>] = array_utils::workaround_transmute_mut(input);
423 let transmuted_output: &mut [Complex<A>] =
424 array_utils::workaround_transmute_mut(output);
425 self.prepare_raders(transmuted_input, transmuted_output)
426 }
427
428 let (first_input, inner_input) = input.split_first_mut().unwrap();
429 let (first_output, inner_output) = output.split_first_mut().unwrap();
430
431 let inner_scratch = if scratch.len() > 0 {
433 &mut scratch[..]
434 } else {
435 &mut inner_input[..]
436 };
437 self.inner_fft
438 .process_with_scratch(inner_output, inner_scratch);
439
440 *first_output = inner_output[0] + *first_input;
442
443 unsafe {
447 let transmuted_inner_output: &mut [Complex<A>] =
449 array_utils::workaround_transmute_mut(inner_input);
450 let transmuted_inner_input: &mut [Complex<A>] =
451 array_utils::workaround_transmute_mut(inner_output);
452 avx_vector::pairwise_complex_mul_conjugated(
453 transmuted_inner_input,
454 transmuted_inner_output,
455 &self.twiddles,
456 )
457 };
458
459 inner_input[0] = inner_input[0] + first_input.conj();
462
463 let inner_scratch = if scratch.len() > 0 {
465 scratch
466 } else {
467 &mut inner_output[..]
468 };
469 self.inner_fft
470 .process_with_scratch(inner_input, inner_scratch);
471
472 unsafe {
474 let transmuted_input: &mut [Complex<A>] = array_utils::workaround_transmute_mut(input);
476 let transmuted_output: &mut [Complex<A>] =
477 array_utils::workaround_transmute_mut(output);
478 self.finalize_raders(transmuted_input, transmuted_output);
479 }
480 }
481 fn perform_fft_inplace(&self, buffer: &mut [Complex<T>], scratch: &mut [Complex<T>]) {
482 let (scratch, extra_scratch) = scratch.split_at_mut(self.len());
483 unsafe {
484 let transmuted_scratch: &mut [Complex<A>] =
486 array_utils::workaround_transmute_mut(scratch);
487 let transmuted_buffer: &mut [Complex<A>] =
488 array_utils::workaround_transmute_mut(buffer);
489 self.prepare_raders(transmuted_buffer, transmuted_scratch)
490 }
491
492 let first_input = buffer[0];
493
494 let truncated_scratch = &mut scratch[1..];
495
496 let inner_scratch = if extra_scratch.len() > 0 {
498 extra_scratch
499 } else {
500 &mut buffer[..]
501 };
502 self.inner_fft
503 .process_with_scratch(truncated_scratch, inner_scratch);
504
505 let first_output = first_input + truncated_scratch[0];
507
508 unsafe {
512 let transmuted_scratch: &mut [Complex<A>] =
514 array_utils::workaround_transmute_mut(truncated_scratch);
515 avx_vector::pairwise_complex_mul_assign_conjugated(transmuted_scratch, &self.twiddles)
516 };
517
518 truncated_scratch[0] = truncated_scratch[0] + first_input.conj();
521
522 self.inner_fft
524 .process_with_scratch(truncated_scratch, inner_scratch);
525
526 buffer[0] = first_output;
528 unsafe {
529 let transmuted_scratch: &mut [Complex<A>] =
531 array_utils::workaround_transmute_mut(scratch);
532 let transmuted_buffer: &mut [Complex<A>] =
533 array_utils::workaround_transmute_mut(buffer);
534 self.finalize_raders(transmuted_scratch, transmuted_buffer);
535 }
536 }
537}
538boilerplate_avx_fft!(
539 RadersAvx2,
540 |this: &RadersAvx2<_, _>| this.len,
541 |this: &RadersAvx2<_, _>| this.inplace_scratch_len,
542 |this: &RadersAvx2<_, _>| this.outofplace_scratch_len,
543 |this: &RadersAvx2<_, _>| this.immut_scratch_len
544);
545
546#[cfg(test)]
547mod unit_tests {
548 use num_traits::Float;
549 use rand::distributions::uniform::SampleUniform;
550
551 use super::*;
552 use crate::algorithm::Dft;
553 use crate::test_utils::check_fft_algorithm;
554 use std::sync::Arc;
555
556 #[test]
557 fn test_raders_avx_f32() {
558 for len in 3..100 {
559 if miller_rabin(len as u64) {
560 test_raders_with_length::<f32>(len, FftDirection::Forward);
561 test_raders_with_length::<f32>(len, FftDirection::Inverse);
562 }
563 }
564 }
565
566 #[test]
567 fn test_raders_avx_f64() {
568 for len in 3..100 {
569 if miller_rabin(len as u64) {
570 test_raders_with_length::<f64>(len, FftDirection::Forward);
571 test_raders_with_length::<f64>(len, FftDirection::Inverse);
572 }
573 }
574 }
575
576 fn test_raders_with_length<T: AvxNum + Float + SampleUniform>(
577 len: usize,
578 direction: FftDirection,
579 ) {
580 let inner_fft = Arc::new(Dft::new(len - 1, direction));
581 let fft = RadersAvx2::<T, T>::new(inner_fft).unwrap();
582
583 check_fft_algorithm::<T>(&fft, len, direction);
584 }
585}