Packing and de-interleaving two __m256 registers

Question

I have a row-wise array of floats (~20 cols x ~1M rows) from which I need to extract two columns at a time into two __m256 registers.

...a0.........b0......
...a1.........b1......
// ...
...a7.........b7......
// end first __m256

A naive way to do this is

__m256i vindex = _mm256_setr_epi32(
    0,
    1 * stride,
    2 * stride,
    // ...
    7 * stride);
__m256 colA = _mm256_i32gather_ps(baseAddrColA, vindex, sizeof(float));
__m256 colB = _mm256_i32gather_ps(baseAddrColB, vindex, sizeof(float));

However, I was wondering if I would get better performance by retrieving a0, b0, a1, b1, a2, b2, a3, b3 in one gather, and a4, b4, ... a7, b7 in another because they're closer in memory, and then de-interleave them. That is:

// __m256   lo = a0 b0 a1 b1 a2 b2 a3 b3 // load proximal elements
// __m256   hi = a4 b4 a5 b5 a6 b6 a7 b7
// __m256 colA = a0 a1 a2 a3 a4 a5 a6 a7 // goal
// __m256 colB = b0 b1 b2 b3 b4 b5 b6 b7

I can't figure out how to nicely interleave lo and hi. I basically need the opposite of _mm256_unpacklo_ps. The best I've come up with is something like:

__m256i idxA = _mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7);
__m256i idxB = _mm256_setr_epi32(1, 3, 5, 7, 0, 2, 4, 6);

__m256 permLA = _mm256_permutevar8x32_ps(lo, idxA);        // a0 a1 a2 a3 b0 b1 b2 b3
__m256 permHB = _mm256_permutevar8x32_ps(hi, idxB);        // b4 b5 b6 b7 a4 a5 a6 a7
__m256 colA = _mm256_blend_ps(permLA, permHB, 0b11110000); // a0 a1 a2 a3 a4 a5 a6 a7
__m256 colB = _mm256_setr_m128(
                          _mm256_extractf128_ps(permLA, 1), 
                          _mm256_castps256_ps128(permHB)); // b0 b1 b2 b3 b4 b5 b6 b7

That's 13 cycles. Is there a better way?

(For all I know, prefetch is already optimizing the naive approach as best as possible, but lacking that knowledge, I was hoping to benchmark the second approach. If anyone already knows what the result of this would be, please do share. With the above de-interlacing method, it's about 8% slower than the naive approach.)

Edit Even without the de-interlacing, the "proximal" gather method is about 6% slower than the naive, constant-stride gather method. I take that to mean that this access pattern confuses hardware prefetch too much to be a worthwhile optimization.

Answer 1

// __m256   lo = a0 b0 a1 b1 a2 b2 a3 b3 // load proximal elements
// __m256   hi = a4 b4 a5 b5 a6 b6 a7 b7
// __m256 colA = a0 a1 a2 a3 a4 a5 a6 a7 // goal
// __m256 colB = b0 b1 b2 b3 b4 b5 b6 b7

It seems we can do this shuffle even faster than my orginal answer:

void unpack_cols(__m256i lo, __m256i hi, __m256i& colA, __m256i& colB) {
    const __m256i mask = _mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7);
    // group cols crossing lanes: 
    // a0 a1 a2 a3 b0 b1 b2 b3
    // a4 a5 a6 a7 b4 b5 b6 b7
    auto lo_grouped = _mm256_permutevar8x32_epi32(lo, mask);
    auto hi_grouped = _mm256_permutevar8x32_epi32(hi, mask);

    // swap lanes: 
    // a0 a1 a2 a3 a4 a5 a6 a7
    // b0 b1 b2 b3 b4 b5 b6 b7
    colA = _mm256_permute2x128_si256(lo_grouped, hi_grouped, 0 | (2 << 4));
    colB = _mm256_permute2x128_si256(lo_grouped, hi_grouped, 1 | (3 << 4));
}

While both instructions have a 3 cycles latency on Haswell (see Agner Fog) they have a single cycle throughput. This means it has a throughput of 4 cycles and 8 cycles latency. If you have a spare register which can keep the mask this should be better. Doing only two of these in parallel allows you to completly hide its latency. See godbolt and rextester.

---

Old answer, kept for reference:

The fastest way to do this shuffle is the following:

void unpack_cols(__m256i lo, __m256i hi, __m256i& colA, __m256i& colB) {
    // group cols within lanes: 
    // a0 a1 b0 b1 a2 a3 b2 b3
    // a4 a5 b4 b5 a6 a7 b6 b7
    auto lo_shuffled = _mm256_shuffle_epi32(lo, _MM_SHUFFLE(3, 1, 2, 0));
    auto hi_shuffled = _mm256_shuffle_epi32(hi, _MM_SHUFFLE(3, 1, 2, 0));

    // unpack lo + hi a 64 bit
    // a0 a1 a4 a5 a2 a3 a6 a7
    // b0 b1 b4 b5 b2 b3 b6 b7
    auto colA_shuffled = _mm256_unpacklo_epi64(lo_shuffled, hi_shuffled);
    auto colB_shuffled = _mm256_unpackhi_epi64(lo_shuffled, hi_shuffled);

    // swap crossing lanes: 
    // a0 a1 a2 a3 a4 a5 a6 a7
    // b0 b1 b2 b3 b4 b5 b6 b7
    colA = _mm256_permute4x64_epi64(colA_shuffled, _MM_SHUFFLE(3, 1, 2, 0));
    colB = _mm256_permute4x64_epi64(colB_shuffled, _MM_SHUFFLE(3, 1, 2, 0));
}

Starting with Haswell this has a throughput of 6 cycles (sadly six instructions on port 5). According to Agner Fog _mm256_permute4x64_epi64 has a latency of 3 cycles. This means unpack_cols has a latency of 11 8 cycles.

You can check the code on godbolt.org or test it at rextester which has AVX2 support but sadly no permalinks like godbolt.

---

Note that this is also very close to the problem I had where I gathered 64 bit ints and needed the high and low 32 bits separated.

---

Note that gather performance is really bad in Haswell but according to Agner Fog Skylake got a lot better at it (~12 cycles throughput down to ~5). Still shuffling around such simple patterns should still be a lot faster than gathering.

Answer 2

In order to load columns of 32-bit float type you could use intrinsics _mm256_setr_pd and _mm256_shuffle_ps (it takes 10 cycles):

#include <iostream>
#include <immintrin.h>

inline void Print(const __m256 & v)
{
    float b[8];
    _mm256_storeu_ps(b, v);
    for (int i = 0; i < 8; i++)
        std::cout << b[i] << " ";
    std::cout << std::endl;
}

int main()
{
    const size_t stride = 100;
    float m[stride * 8];
    for (size_t i = 0; i < stride*8; ++i)
        m[i] = (float)i;

    const size_t stride2 = stride / 2;
    double * p = (double*)m;

    __m256 ab0145 = _mm256_castpd_ps(_mm256_setr_pd(p[0 * stride2], p[1 * stride2], p[4 * stride2], p[5 * stride2]));
    __m256 ab2367 = _mm256_castpd_ps(_mm256_setr_pd(p[2 * stride2], p[3 * stride2], p[6 * stride2], p[7 * stride2]));

    __m256 a = _mm256_shuffle_ps(ab0145, ab2367, 0x88);
    __m256 b = _mm256_shuffle_ps(ab0145, ab2367, 0xDD);

    Print(a);
    Print(b);

    return 0;
}

Output:

0 100 200 300 400 500 600 700
1 101 201 301 401 501 601 701

Concerning to performance of intrinsic _mm256_i32gather_ps I would recommend to see here.