Skip to content

Commit 941d617

Browse files
ShnatselShnatsel
committed
targeted conversion to a for loop for just enough inlining
1 parent dd00bef commit 941d617

1 file changed

Lines changed: 72 additions & 72 deletions

File tree

‎src/kernels/dit.rs‎

Lines changed: 72 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -996,44 +996,44 @@ fn fft_dit_chunk_n_simd_f64<S: Simd>(
996996
let chunk_size = dist * 2;
997997
assert!(chunk_size >= LANES * 2);
998998

999-
reals
999+
for (reals_chunk, imags_chunk) in reals
10001000
.chunks_exact_mut(chunk_size)
10011001
.zip(imags.chunks_exact_mut(chunk_size))
1002-
.for_each(|(reals_chunk, imags_chunk)| {
1003-
let (reals_s0, reals_s1) = reals_chunk.split_at_mut(dist);
1004-
let (imags_s0, imags_s1) = imags_chunk.split_at_mut(dist);
1005-
1006-
(reals_s0.as_chunks_mut::<LANES>().0.iter_mut())
1007-
.zip(reals_s1.as_chunks_mut::<LANES>().0.iter_mut())
1008-
.zip(imags_s0.as_chunks_mut::<LANES>().0.iter_mut())
1009-
.zip(imags_s1.as_chunks_mut::<LANES>().0.iter_mut())
1010-
.zip(twiddles_re.as_chunks::<LANES>().0.iter())
1011-
.zip(twiddles_im.as_chunks::<LANES>().0.iter())
1012-
.for_each(|(((((re_s0, re_s1), im_s0), im_s1), tw_re), tw_im)| {
1013-
let two = f64x8::splat(simd, 2.0);
1014-
let in0_re = f64x8::simd_from(simd, *re_s0);
1015-
let in1_re = f64x8::simd_from(simd, *re_s1);
1016-
let in0_im = f64x8::simd_from(simd, *im_s0);
1017-
let in1_im = f64x8::simd_from(simd, *im_s1);
1018-
1019-
let tw_re = f64x8::simd_from(simd, *tw_re);
1020-
let tw_im = f64x8::simd_from(simd, *tw_im);
1021-
1022-
// out0.re = (in0.re + tw_re * in1.re) - tw_im * in1.im
1023-
let out0_re = tw_im.mul_add(-in1_im, tw_re.mul_add(in1_re, in0_re));
1024-
// out0.im = (in0.im + tw_re * in1.im) + tw_im * in1.re
1025-
let out0_im = tw_im.mul_add(in1_re, tw_re.mul_add(in1_im, in0_im));
1026-
1027-
// Use FMA for out1 = 2*in0 - out0
1028-
let out1_re = two.mul_sub(in0_re, out0_re);
1029-
let out1_im = two.mul_sub(in0_im, out0_im);
1030-
1031-
out0_re.store_slice(re_s0);
1032-
out0_im.store_slice(im_s0);
1033-
out1_re.store_slice(re_s1);
1034-
out1_im.store_slice(im_s1);
1035-
});
1036-
});
1002+
{
1003+
let (reals_s0, reals_s1) = reals_chunk.split_at_mut(dist);
1004+
let (imags_s0, imags_s1) = imags_chunk.split_at_mut(dist);
1005+
1006+
(reals_s0.as_chunks_mut::<LANES>().0.iter_mut())
1007+
.zip(reals_s1.as_chunks_mut::<LANES>().0.iter_mut())
1008+
.zip(imags_s0.as_chunks_mut::<LANES>().0.iter_mut())
1009+
.zip(imags_s1.as_chunks_mut::<LANES>().0.iter_mut())
1010+
.zip(twiddles_re.as_chunks::<LANES>().0.iter())
1011+
.zip(twiddles_im.as_chunks::<LANES>().0.iter())
1012+
.for_each(|(((((re_s0, re_s1), im_s0), im_s1), tw_re), tw_im)| {
1013+
let two = f64x8::splat(simd, 2.0);
1014+
let in0_re = f64x8::simd_from(simd, *re_s0);
1015+
let in1_re = f64x8::simd_from(simd, *re_s1);
1016+
let in0_im = f64x8::simd_from(simd, *im_s0);
1017+
let in1_im = f64x8::simd_from(simd, *im_s1);
1018+
1019+
let tw_re = f64x8::simd_from(simd, *tw_re);
1020+
let tw_im = f64x8::simd_from(simd, *tw_im);
1021+
1022+
// out0.re = (in0.re + tw_re * in1.re) - tw_im * in1.im
1023+
let out0_re = tw_im.mul_add(-in1_im, tw_re.mul_add(in1_re, in0_re));
1024+
// out0.im = (in0.im + tw_re * in1.im) + tw_im * in1.re
1025+
let out0_im = tw_im.mul_add(in1_re, tw_re.mul_add(in1_im, in0_im));
1026+
1027+
// Use FMA for out1 = 2*in0 - out0
1028+
let out1_re = two.mul_sub(in0_re, out0_re);
1029+
let out1_im = two.mul_sub(in0_im, out0_im);
1030+
1031+
out0_re.store_slice(re_s0);
1032+
out0_im.store_slice(im_s0);
1033+
out1_re.store_slice(re_s1);
1034+
out1_im.store_slice(im_s1);
1035+
});
1036+
}
10371037
}
10381038

10391039
/// General DIT butterfly for f32
@@ -1066,42 +1066,42 @@ fn fft_dit_chunk_n_simd_f32<S: Simd>(
10661066
let chunk_size = dist * 2;
10671067
assert!(chunk_size >= LANES * 2);
10681068

1069-
reals
1069+
for (reals_chunk, imags_chunk) in reals
10701070
.chunks_exact_mut(chunk_size)
10711071
.zip(imags.chunks_exact_mut(chunk_size))
1072-
.for_each(|(reals_chunk, imags_chunk)| {
1073-
let (reals_s0, reals_s1) = reals_chunk.split_at_mut(dist);
1074-
let (imags_s0, imags_s1) = imags_chunk.split_at_mut(dist);
1075-
1076-
(reals_s0.as_chunks_mut::<LANES>().0.iter_mut())
1077-
.zip(reals_s1.as_chunks_mut::<LANES>().0.iter_mut())
1078-
.zip(imags_s0.as_chunks_mut::<LANES>().0.iter_mut())
1079-
.zip(imags_s1.as_chunks_mut::<LANES>().0.iter_mut())
1080-
.zip(twiddles_re.as_chunks::<LANES>().0.iter())
1081-
.zip(twiddles_im.as_chunks::<LANES>().0.iter())
1082-
.for_each(|(((((re_s0, re_s1), im_s0), im_s1), tw_re), tw_im)| {
1083-
let two = f32x16::splat(simd, 2.0);
1084-
let in0_re = f32x16::simd_from(simd, *re_s0);
1085-
let in1_re = f32x16::simd_from(simd, *re_s1);
1086-
let in0_im = f32x16::simd_from(simd, *im_s0);
1087-
let in1_im = f32x16::simd_from(simd, *im_s1);
1088-
1089-
let tw_re = f32x16::simd_from(simd, *tw_re);
1090-
let tw_im = f32x16::simd_from(simd, *tw_im);
1091-
1092-
// out0.re = (in0.re + tw_re * in1.re) - tw_im * in1.im
1093-
let out0_re = tw_im.mul_add(-in1_im, tw_re.mul_add(in1_re, in0_re));
1094-
// out0.im = (in0.im + tw_re * in1.im) + tw_im * in1.re
1095-
let out0_im = tw_im.mul_add(in1_re, tw_re.mul_add(in1_im, in0_im));
1096-
1097-
// Use FMA for out1 = 2*in0 - out0
1098-
let out1_re = two.mul_sub(in0_re, out0_re);
1099-
let out1_im = two.mul_sub(in0_im, out0_im);
1100-
1101-
out0_re.store_slice(re_s0);
1102-
out0_im.store_slice(im_s0);
1103-
out1_re.store_slice(re_s1);
1104-
out1_im.store_slice(im_s1);
1105-
});
1106-
});
1072+
{
1073+
let (reals_s0, reals_s1) = reals_chunk.split_at_mut(dist);
1074+
let (imags_s0, imags_s1) = imags_chunk.split_at_mut(dist);
1075+
1076+
(reals_s0.as_chunks_mut::<LANES>().0.iter_mut())
1077+
.zip(reals_s1.as_chunks_mut::<LANES>().0.iter_mut())
1078+
.zip(imags_s0.as_chunks_mut::<LANES>().0.iter_mut())
1079+
.zip(imags_s1.as_chunks_mut::<LANES>().0.iter_mut())
1080+
.zip(twiddles_re.as_chunks::<LANES>().0.iter())
1081+
.zip(twiddles_im.as_chunks::<LANES>().0.iter())
1082+
.for_each(|(((((re_s0, re_s1), im_s0), im_s1), tw_re), tw_im)| {
1083+
let two = f32x16::splat(simd, 2.0);
1084+
let in0_re = f32x16::simd_from(simd, *re_s0);
1085+
let in1_re = f32x16::simd_from(simd, *re_s1);
1086+
let in0_im = f32x16::simd_from(simd, *im_s0);
1087+
let in1_im = f32x16::simd_from(simd, *im_s1);
1088+
1089+
let tw_re = f32x16::simd_from(simd, *tw_re);
1090+
let tw_im = f32x16::simd_from(simd, *tw_im);
1091+
1092+
// out0.re = (in0.re + tw_re * in1.re) - tw_im * in1.im
1093+
let out0_re = tw_im.mul_add(-in1_im, tw_re.mul_add(in1_re, in0_re));
1094+
// out0.im = (in0.im + tw_re * in1.im) + tw_im * in1.re
1095+
let out0_im = tw_im.mul_add(in1_re, tw_re.mul_add(in1_im, in0_im));
1096+
1097+
// Use FMA for out1 = 2*in0 - out0
1098+
let out1_re = two.mul_sub(in0_re, out0_re);
1099+
let out1_im = two.mul_sub(in0_im, out0_im);
1100+
1101+
out0_re.store_slice(re_s0);
1102+
out0_im.store_slice(im_s0);
1103+
out1_re.store_slice(re_s1);
1104+
out1_im.store_slice(im_s1);
1105+
});
1106+
}
11071107
}

0 commit comments

Comments
 (0)