How to do an AVX shuffle by a variable

I want to do implement a fixed lookup table search through instruction.

The instruction _mm_shuffle_epi32(table, index) suit my requirement. but it need a immediate number.

if I want to use similar instruction, I should extend the int8_t to an __m128i

Is there a better way to handle this situation

I have a look up table like

int32_t table = {t0, t1, t2, t3};
an index vector int8_t [100] = vindex;
the c code is like

for (uint32_t i = 0; i < 100; i++){
    pdata[i * 4] = table[vindex[i] & 3];
    pdata[i * 4 +1] = table[vindex[i] >> 2];
    pdata[i * 4 +2] = table[vindex[i] >> 4];
    pdata[i * 4 +3] = table[vindex[i] >> 6];}

This is consistent with the instruction _mm_shuffle_epi32() performance.

First, your baseline has some bugs if I'm not mistaken. This is the version I use for benchmarking and validation:

void baseline(
        int* pdata, const int* table, const uint8_t* vindex, ptrdiff_t n)
{
    ptrdiff_t i;
    for(i = 0; i < n; ++i) {
        pdata[i * 4] = table[vindex[i] & 3];
        pdata[i * 4 + 1] = table[(vindex[i] >> 2) & 3];
        pdata[i * 4 + 2] = table[(vindex[i] >> 4) & 3];
        pdata[i * 4 + 3] = table[vindex[i] >> 6];
    }
}

If we only have AVX-1, we need to do some bit-twiddling tricks to get the lookup bits into shape. We can view a 64 bit integer as a packed 4x16 bit vector. What we want is the lowest 2 bits per 16 bit word to be the shuffle for that element. We can do this via a single multiplication because it can act like a broadcast with a variable shift as long as there is no overlap between elements.

We move this to the vector register and spread it further into 128 bit. This will leave us with a couple of set bits in unexpected locations but the shuffle instructions ignore those (and it would be easy to mask them if they didn't).

void shuffle_avx1(
        int* pdata, const int* table, const uint8_t* vindex, ptrdiff_t n)
{
    const uint64_t magic_mul =
        1 | (1 << 14) | (1 << 28) | (1ull << 42);
    const __m128 vtable = _mm_castsi128_ps(
            _mm_loadu_si128((const __m128i*) table));
    ptrdiff_t i;
    for(i = 0; i < n; ++i) {
        __m128i permut = _mm_cvtsi64_si128(vindex[i] * magic_mul);
        __m128 result;
        permut = _mm_cvtepu16_epi32(permut);
        result = _mm_permutevar_ps(vtable, permut);
        _mm_storeu_ps((float*) pdata + i * 4, result);
    }
}

This is about the factor 2.1 faster. If we have AVX2, we can instead use variable shifts. In my tests it runs faster by a factor of 2.9 compared to the baseline.

void shuffle_avx2(
        int* pdata, const int* table, const uint8_t* vindex, ptrdiff_t n)
{
    const __m128i shifts = _mm_set_epi32(6, 4, 2, 0);
    const __m128 vtable = _mm_castsi128_ps(
            _mm_loadu_si128((const __m128i*) table));
    ptrdiff_t i;
    for(i = 0; i < n; ++i) {
        __m128i result;
        __m128i permut = _mm_set1_epi32(vindex[i]);
        permut = _mm_srlv_epi32(permut, shifts);
        result = _mm_permutevar_ps(vtable, permut);
        _mm_storeu_ps((float*) pdata + i * 4, result);
    }
}

And while we are at it, we can do the same for a full 256 bit vector. My tests show that for this we need to be careful to align the output to a 32 byte boundary if we can (meaning, the output was at least 16 byte aligned to begin with). This increases the speedup to 4.3. I'm not sure why it didn't scale up better. I don't think my Tiger Lake i7-11800H should have any bottlenecks in 256 bit compared to 128 bit. Maybe I'm already memory-bound.

void shuffle_avx2_full(
        int* pdata, const int* table, const uint8_t* vindex, ptrdiff_t n)
{
    const __m256i shifts8 = _mm256_set_epi32(14, 12, 10, 8, 6, 4, 2, 0);
    const __m256 vtable8 = _mm256_castsi256_ps(
            _mm256_broadcastsi128_si256(
            _mm_loadu_si128((const __m128i*) table)));
    const int small_iter = n < 2 ? n : (((ptrdiff_t) pdata & 31) == 16);
    ptrdiff_t i = 0;
    if(small_iter) {
        /*
         * For n == 1, do full output.
         * For larger n, align output to 32 byte boundary if possible
         */
        const __m128i shifts4 = _mm256_castsi256_si128(shifts8);
        const __m128 vtable4 = _mm256_castps256_ps128(vtable8);
        __m128i permut = _mm_set1_epi32(*vindex);
        __m128 result;
        permut = _mm_srlv_epi32(permut, shifts4);
        result = _mm_permutevar_ps(vtable4, permut);
        _mm_storeu_ps((float*) pdata, result);
        i = 1;
    }
    for(; i + 2 <= n; i += 2) {
        __m256i permut = _mm256_set1_epi32(
            *((const uint16_t*) (vindex + i)));;
        __m256 result;
        permut = _mm256_srlv_epi32(permut, shifts8);
        result = _mm256_permutevar_ps(vtable8, permut);
        /*
         * This still needs to be an unaligned store even though
         * we tried to align above, we cannot align if the output
         * address is not at least 16 byte-aligned
         */
        _mm256_storeu_ps((float*) pdata + i * 4, result);
    }
    if(i < n) {
        /* One misaligned iteration for tail */
        i = n - 2;
        __m256i permut = _mm256_set1_epi32(
            *((const uint16_t*) (vindex + i)));;
        __m256 result;
        permut = _mm256_srlv_epi32(permut, shifts8);
        result = _mm256_permutevar_ps(vtable8, permut);
        _mm256_storeu_ps((float*) pdata + i * 4, result);
    }
}

Scaling this up to AVX512 is simple but at least on my system this is is slower. YMMV.

void shuffle_avx512(
        int* pdata, const int* table, const uint8_t* vindex, ptrdiff_t n)
{
    const __m512i shifts16 = _mm512_set_epi32(
            30, 28, 26, 24, 22, 20, 18, 16,
            14, 12, 10,  8,  6,  4,  2,  0);
    const __m128i shifts4 = _mm512_castsi512_si128(shifts16);
    const __m512 vtable16 = _mm512_castsi512_ps(
            _mm512_broadcast_i32x4(
            _mm_loadu_si128((const __m128i*) table)));
    const __m128 vtable4 = _mm512_castps512_ps128(vtable16);
    const int small_count = n < 4 ? n : (ptrdiff_t) pdata >> 4 & 3;
    ptrdiff_t i;
    /* Large n: Align to 64 byte output. Small n: All iterations */
    for(i = 0; i < small_count; ++i) {
        __m128i permut = _mm_set1_epi32(vindex[i]);
        __m128 result;
        permut = _mm_srlv_epi32(permut, shifts4);
        result = _mm_permutevar_ps(vtable4, permut);
        _mm_storeu_ps((float*) pdata + i * 4, result);
    }
    for(; i + 4 < n; i += 4) {
        __m512i permut = _mm512_set1_epi32(*((const int*) (vindex + i)));
        __m512 result;
        permut = _mm512_srlv_epi32(permut, shifts16);
        result = _mm512_permutevar_ps(vtable16, permut);
        /* Output will still be misaligned if pdata isn't 16 byte aligned */
        _mm512_storeu_ps((float*) pdata + i * 4, result);
    }
    if(i < n) {
        /* One overlapping misaligned iteration for tail */
        i = n - 4;
        __m512i permut = _mm512_set1_epi32(*((const int*) (vindex + i)));
        __m512 result;
        permut = _mm512_srlv_epi32(permut, shifts16);
        result = _mm512_permutevar_ps(vtable16, permut);
        _mm512_storeu_ps((float*) pdata + i * 4, result);
    }
}

Side note: Funny enough, without a -march or -mtune setting, GCC decides that the multiplication trick is stupid and uses multiple shifts and ORs instead. That doesn't sound sensible to me though. Which machine has AVX but no fast multiplier?