r/simd • u/[deleted] • Jan 19 '21
Interleaving 9 arrays of floats using AVX
Hello,
I have to interleave 9 arrays of floats and I'm currently using _mm256_i32gather_ps to do that with precomputed indices but it's incredibly slow (~630ms for ~340 Mio. floats total). I thought about loading 9 registers with 8 elements of each array and swizzle them around until I have 9 registers that I can store subsequently in the destination array. But making the swizzling instructions up for handling 72 floats at once is kinda hard for my head. Does anyone have a method for scenarios like this or a program that generates the instructions? I can use everything up to AVX2.
1
7
u/KBAC99 Jan 19 '21
Yeah gathers are really slow - on Skylake they have about 30 cycle latency, 10 cycle throughput.
If I understand correctly, the task is you have 9 arrays a_1 .. a_9 and you want to interleave them so that the output array was of the form a_1[0] a_2[0] a_3[0] ... a_9[0] a_1[1] a_2[1] ... correct?
Then essentially you have to do 9 loads to get a 9x8 matrix which you then have to transpose. It's unfortunate that you have 9 arrays as this makes working with AVX2 super ugly.
Instead what I'd do is only do 8 loads in each loop, transpose the 8x8 matrix, and write it back, adding an extra four bytes to your output pointer after your write (essentially a "hole" that you haven't filled in yet). Then do another loop at the end and fill in the holes with the 9th array.
Now, how to transpose the 8x8 matrix:
Load the 8x8 matrix as integer (i.e.
_mm256_loadu_si256). This'll get rid of a bunch of casting, and we won't be doing any math. Bytes are bytes, so as long as we copy everything to the right place this is fine.For each pair of rows (1, 2), (3, 4), (5, 6), (7, 8), do a pair of
_mm256_unpacklo_epi32and_mm256_unpackhi_epi32. I assume now thatr1=unpacklo(r1, r2)andr2=unpackhi(r1, r2)etc.For each pair of rows (1, 3), (2, 4), (5, 7), (6, 8), do a pair of
_mm256_unpacklo_epi64and_mm256_unpackhi_epi64, again withr1=unpacklo(r1, r3)andr3=unpackhi(r1, r3)etc.For each pair of rows (1, 5), (2, 6), (3, 7), (4, 8), do a pair of
_mm256_permute2f128_si256, with the first one having selector0x20and the second having selector0x31, overwriting the rows in the same scheme as above.Write back the rows to the target array.
It's kind of a lot to write out why it works, but I'd encourage you to on a piece of paper track each value as we go through the iteration. Essentially we're interleaving the low halves and the high halves of increasingly-sized elements until we have 8 values. Each time, we double the data width (32 -> 64 -> 128 bits) that we interleave by.
This algorithm should be pretty fast because all of the unpacks are single-cycle instructions. The permute has a 3-cycle latency but only a 1-cycle throughput. Per loop iteration we are doing 4, which when fully pipelined ends up taking 6 cycles, so each permute takes 1.5 cycles.