rustfft/avx/
avx64_utils.rs

1use std::arch::x86_64::*;
2
3// Treat the input like the rows of a 2x2 array, and transpose said rows to the columns
4#[inline(always)]
5pub unsafe fn transpose_2x2_f64(rows: [__m256d; 2]) -> [__m256d; 2] {
6    let col0 = _mm256_permute2f128_pd(rows[0], rows[1], 0x20);
7    let col1 = _mm256_permute2f128_pd(rows[0], rows[1], 0x31);
8
9    [col0, col1]
10}
11
12// Treat the input like the rows of a 4x2 array, and transpose it to a 2x4 array
13#[inline(always)]
14pub unsafe fn transpose_4x2_to_2x4_f64(rows0: [__m256d; 2], rows1: [__m256d; 2]) -> [__m256d; 4] {
15    let output00 = transpose_2x2_f64(rows0);
16    let output01 = transpose_2x2_f64(rows1);
17
18    [output00[0], output00[1], output01[0], output01[1]]
19}
20
21// Treat the input like the rows of a 3x3 array, and transpose it
22// The assumption here is that it's very likely that the caller wants to do some more AVX operations on the columns of the transposed array, so the output is arranged to make that more convenient
23#[inline(always)]
24pub unsafe fn transpose_3x3_f64(
25    rows0: [__m128d; 3],
26    rows1: [__m256d; 3],
27) -> ([__m128d; 3], [__m256d; 3]) {
28    // the first column of output will be made up of the first row of input
29    let output0 = [
30        rows0[0],
31        _mm256_castpd256_pd128(rows1[0]),
32        _mm256_extractf128_pd(rows1[0], 1),
33    ];
34
35    // the second column of output will be made of the second 2 rows of input
36    let output10 = _mm256_permute2f128_pd(
37        _mm256_castpd128_pd256(rows0[1]),
38        _mm256_castpd128_pd256(rows0[2]),
39        0x20,
40    );
41    let lower_chunk = [rows1[1], rows1[2]];
42    let lower_transposed = transpose_2x2_f64(lower_chunk);
43    let output1 = [output10, lower_transposed[0], lower_transposed[1]];
44
45    (output0, output1)
46}
47
48// Treat the input like the rows of a 3x4 array, and transpose it to a 4x3 array
49// The assumption here is that the caller wants to do some more AVX operations on the columns of the transposed array, so the output is arranged to make that more convenient
50#[inline(always)]
51pub unsafe fn transpose_3x4_to_4x3_f64(
52    rows0: [__m128d; 4],
53    rows1: [__m256d; 4],
54) -> ([__m256d; 3], [__m256d; 3]) {
55    // the top row of each output array will come from the first column, and the second 2 rows will come from 2x2 transposing the rows1 array
56    let merged0 = _mm256_permute2f128_pd(
57        _mm256_castpd128_pd256(rows0[0]),
58        _mm256_castpd128_pd256(rows0[1]),
59        0x20,
60    );
61    let merged1 = _mm256_permute2f128_pd(
62        _mm256_castpd128_pd256(rows0[2]),
63        _mm256_castpd128_pd256(rows0[3]),
64        0x20,
65    );
66
67    let chunk0 = [rows1[0], rows1[1]];
68    let chunk1 = [rows1[2], rows1[3]];
69
70    let lower0 = transpose_2x2_f64(chunk0);
71    let lower1 = transpose_2x2_f64(chunk1);
72
73    (
74        [merged0, lower0[0], lower0[1]],
75        [merged1, lower1[0], lower1[1]],
76    )
77}
78
79// Treat the input like the rows of a 3x6 array, and transpose it to a 6x3 array
80// The assumption here is that caller wants to do some more AVX operations on the columns of the transposed array, so the output is arranged to make that more convenient
81#[inline(always)]
82pub unsafe fn transpose_3x6_to_6x3_f64(
83    rows0: [__m128d; 6],
84    rows1: [__m256d; 6],
85) -> ([__m256d; 3], [__m256d; 3], [__m256d; 3]) {
86    let chunk0 = [rows1[0], rows1[1]];
87    let chunk1 = [rows1[2], rows1[3]];
88    let chunk2 = [rows1[4], rows1[5]];
89
90    let transposed0 = transpose_2x2_f64(chunk0);
91    let transposed1 = transpose_2x2_f64(chunk1);
92    let transposed2 = transpose_2x2_f64(chunk2);
93
94    let output0 = [
95        _mm256_insertf128_pd(_mm256_castpd128_pd256(rows0[0]), rows0[1], 1),
96        transposed0[0],
97        transposed0[1],
98    ];
99    let output1 = [
100        _mm256_insertf128_pd(_mm256_castpd128_pd256(rows0[2]), rows0[3], 1),
101        transposed1[0],
102        transposed1[1],
103    ];
104    let output2 = [
105        _mm256_insertf128_pd(_mm256_castpd128_pd256(rows0[4]), rows0[5], 1),
106        transposed2[0],
107        transposed2[1],
108    ];
109
110    (output0, output1, output2)
111}
112
113// Treat the input like the rows of a 9x3 array, and transpose it to a 3x9 array
114// The assumption here is that caller wants to do some more AVX operations on the columns of the transposed array, so the output is arranged to make that more convenient
115#[inline(always)]
116pub unsafe fn transpose_9x3_to_3x9_f64(
117    rows0: [__m128d; 3],
118    rows1: [__m256d; 3],
119    rows2: [__m256d; 3],
120    rows3: [__m256d; 3],
121    rows4: [__m256d; 3],
122) -> ([__m128d; 9], [__m256d; 9]) {
123    let chunk1 = [rows1[1], rows1[2]];
124    let chunk2 = [rows2[1], rows2[2]];
125    let chunk3 = [rows3[1], rows3[2]];
126    let chunk4 = [rows4[1], rows4[2]];
127
128    let transposed1 = transpose_2x2_f64(chunk1);
129    let transposed2 = transpose_2x2_f64(chunk2);
130    let transposed3 = transpose_2x2_f64(chunk3);
131    let transposed4 = transpose_2x2_f64(chunk4);
132
133    let output0 = [
134        rows0[0],
135        _mm256_castpd256_pd128(rows1[0]),
136        _mm256_extractf128_pd(rows1[0], 1),
137        _mm256_castpd256_pd128(rows2[0]),
138        _mm256_extractf128_pd(rows2[0], 1),
139        _mm256_castpd256_pd128(rows3[0]),
140        _mm256_extractf128_pd(rows3[0], 1),
141        _mm256_castpd256_pd128(rows4[0]),
142        _mm256_extractf128_pd(rows4[0], 1),
143    ];
144    let output1 = [
145        _mm256_insertf128_pd(_mm256_castpd128_pd256(rows0[1]), rows0[2], 1),
146        transposed1[0],
147        transposed1[1],
148        transposed2[0],
149        transposed2[1],
150        transposed3[0],
151        transposed3[1],
152        transposed4[0],
153        transposed4[1],
154    ];
155
156    (output0, output1)
157}
158
159// Treat the input like the rows of a 4x4 array, and transpose said rows to the columns
160// The assumption here is that the caller wants to do some more AVX operations on the columns of the transposed array, so the output is arranged to make that more convenient
161#[inline(always)]
162pub unsafe fn transpose_4x4_f64(
163    rows0: [__m256d; 4],
164    rows1: [__m256d; 4],
165) -> ([__m256d; 4], [__m256d; 4]) {
166    let chunk00 = [rows0[0], rows0[1]];
167    let chunk01 = [rows0[2], rows0[3]];
168    let chunk10 = [rows1[0], rows1[1]];
169    let chunk11 = [rows1[2], rows1[3]];
170
171    let output00 = transpose_2x2_f64(chunk00);
172    let output01 = transpose_2x2_f64(chunk10);
173    let output10 = transpose_2x2_f64(chunk01);
174    let output11 = transpose_2x2_f64(chunk11);
175
176    (
177        [output00[0], output00[1], output01[0], output01[1]],
178        [output10[0], output10[1], output11[0], output11[1]],
179    )
180}
181
182// Treat the input like the rows of a 6x6 array, and transpose said rows to the columns
183// The assumption here is that the caller wants to do some more AVX operations on the columns of the transposed array, so the output is arranged to make that more convenient
184#[inline(always)]
185pub unsafe fn transpose_6x6_f64(
186    rows0: [__m256d; 6],
187    rows1: [__m256d; 6],
188    rows2: [__m256d; 6],
189) -> ([__m256d; 6], [__m256d; 6], [__m256d; 6]) {
190    let chunk00 = [rows0[0], rows0[1]];
191    let chunk01 = [rows0[2], rows0[3]];
192    let chunk02 = [rows0[4], rows0[5]];
193    let chunk10 = [rows1[0], rows1[1]];
194    let chunk11 = [rows1[2], rows1[3]];
195    let chunk12 = [rows1[4], rows1[5]];
196    let chunk20 = [rows2[0], rows2[1]];
197    let chunk21 = [rows2[2], rows2[3]];
198    let chunk22 = [rows2[4], rows2[5]];
199
200    let output00 = transpose_2x2_f64(chunk00);
201    let output01 = transpose_2x2_f64(chunk10);
202    let output02 = transpose_2x2_f64(chunk20);
203    let output10 = transpose_2x2_f64(chunk01);
204    let output11 = transpose_2x2_f64(chunk11);
205    let output12 = transpose_2x2_f64(chunk21);
206    let output20 = transpose_2x2_f64(chunk02);
207    let output21 = transpose_2x2_f64(chunk12);
208    let output22 = transpose_2x2_f64(chunk22);
209
210    (
211        [
212            output00[0],
213            output00[1],
214            output01[0],
215            output01[1],
216            output02[0],
217            output02[1],
218        ],
219        [
220            output10[0],
221            output10[1],
222            output11[0],
223            output11[1],
224            output12[0],
225            output12[1],
226        ],
227        [
228            output20[0],
229            output20[1],
230            output21[0],
231            output21[1],
232            output22[0],
233            output22[1],
234        ],
235    )
236}
237
238// Treat the input like the rows of a 6x4 array, and transpose said rows to the columns
239// The assumption here is that the caller wants to do some more AVX operations on the columns of the transposed array, so the output is arranged to make that more convenient
240#[inline(always)]
241pub unsafe fn transpose_6x4_to_4x6_f64(
242    rows0: [__m256d; 4],
243    rows1: [__m256d; 4],
244    rows2: [__m256d; 4],
245) -> ([__m256d; 6], [__m256d; 6]) {
246    let chunk00 = [rows0[0], rows0[1]];
247    let chunk01 = [rows0[2], rows0[3]];
248    let chunk10 = [rows1[0], rows1[1]];
249    let chunk11 = [rows1[2], rows1[3]];
250    let chunk20 = [rows2[0], rows2[1]];
251    let chunk21 = [rows2[2], rows2[3]];
252
253    let output00 = transpose_2x2_f64(chunk00);
254    let output01 = transpose_2x2_f64(chunk10);
255    let output02 = transpose_2x2_f64(chunk20);
256    let output10 = transpose_2x2_f64(chunk01);
257    let output11 = transpose_2x2_f64(chunk11);
258    let output12 = transpose_2x2_f64(chunk21);
259
260    (
261        [
262            output00[0],
263            output00[1],
264            output01[0],
265            output01[1],
266            output02[0],
267            output02[1],
268        ],
269        [
270            output10[0],
271            output10[1],
272            output11[0],
273            output11[1],
274            output12[0],
275            output12[1],
276        ],
277    )
278}
279
280// Treat the input like the rows of a 8x4 array, and transpose it to a 4x8 array
281// The assumption here is that it's very likely that the caller wants to do some more AVX operations on the columns of the transposed array, so the output is arranged to make that more convenient
282#[inline(always)]
283pub unsafe fn transpose_8x4_to_4x8_f64(
284    rows0: [__m256d; 4],
285    rows1: [__m256d; 4],
286    rows2: [__m256d; 4],
287    rows3: [__m256d; 4],
288) -> ([__m256d; 8], [__m256d; 8]) {
289    let chunk00 = [rows0[0], rows0[1]];
290    let chunk01 = [rows0[2], rows0[3]];
291    let chunk10 = [rows1[0], rows1[1]];
292    let chunk11 = [rows1[2], rows1[3]];
293    let chunk20 = [rows2[0], rows2[1]];
294    let chunk21 = [rows2[2], rows2[3]];
295    let chunk30 = [rows3[0], rows3[1]];
296    let chunk31 = [rows3[2], rows3[3]];
297
298    let output00 = transpose_2x2_f64(chunk00);
299    let output01 = transpose_2x2_f64(chunk10);
300    let output02 = transpose_2x2_f64(chunk20);
301    let output03 = transpose_2x2_f64(chunk30);
302    let output10 = transpose_2x2_f64(chunk01);
303    let output11 = transpose_2x2_f64(chunk11);
304    let output12 = transpose_2x2_f64(chunk21);
305    let output13 = transpose_2x2_f64(chunk31);
306
307    (
308        [
309            output00[0],
310            output00[1],
311            output01[0],
312            output01[1],
313            output02[0],
314            output02[1],
315            output03[0],
316            output03[1],
317        ],
318        [
319            output10[0],
320            output10[1],
321            output11[0],
322            output11[1],
323            output12[0],
324            output12[1],
325            output13[0],
326            output13[1],
327        ],
328    )
329}
rustfft/avx/avx64_utils.rs

rustfft/avx/
avx64_utils.rs
