Question

如何从__m256向量中的索引中提取单个float，并将其广播到结果向量？

伪代码：

__m256 input = { 2, 3, 4, 5, 6, 7, 8, 9 };
__m256 output = __mm256_selectidx(input, 2);

// output [0 .. 7] now consists of input[2], that is, {4, 4, 4, 4, 4, 4, 4, 4}

相关功能似乎是提取/插入或置换，但文档很少/我不太了解它。广播系列看起来也不错，但只适用于内存操作数？

Answer 1

对于AVX-only（即没有AVX2），您可以这样做：

#include <stdio.h>
#include <immintrin.h>

#define _mm256_selectidx(v, i) ({\
     __m256 vt = _mm256_permute2f128_ps(v, v, (i >> 2) | ((i >> 2) << 4)); \
     vt = _mm256_permute_ps(vt, _MM_SHUFFLE(i & 3, i & 3, i & 3, i & 3)); \
})

int main(void)
{
    __m256 v0 = _mm256_setr_ps(2, 3, 4, 5, 6, 7, 8, 9);
    __m256 v1 = _mm256_selectidx(v0, 2);
    float f0[8], f1[8];

    _mm256_storeu_ps(f0, v0);
    _mm256_storeu_ps(f1, v1);

    printf("v0: %g %g %g %g %g %g %g %g\n", f0[0], f0[1], f0[2], f0[3], f0[4], f0[5], f0[6], f0[7]);
    printf("v1: %g %g %g %g %g %g %g %g\n", f1[0], f1[1], f1[2], f1[3], f1[4], f1[5], f1[6], f1[7]);

    return 0;
}

测试：

$ gcc -Wall -mavx test_avx_select.c && ./a.out
v0: 2 3 4 5 6 7 8 9
v1: 4 4 4 4 4 4 4 4

请注意，此代码对宏使用gcc扩展，其行为类似于函数 - 如果您使用的编译器不支持此扩展，那么您可能需要使用内联函数并希望编译器可以处理使用AVX内在函数所需的编译时常量。

Answer 2

如果您有AVX2，则可以使用_mm256_permutevar8x32_ps：

#define _mm256_selectidx(v, i) _mm256_permutevar8x32_ps(v, _mm256_set1_epi32(i))

显然这会生成几条指令，具体取决于编译器处理_mm256_set1_epi32的方式以及元素索引是否为编译时常量。

演示：

#include <stdio.h>
#include <immintrin.h>

#define _mm256_selectidx(v, i) _mm256_permutevar8x32_ps(v, _mm256_set1_epi32(i))

int main(void)
{
    __m256 v0 = _mm256_setr_ps(2, 3, 4, 5, 6, 7, 8, 9);
    __m256 v1 = _mm256_selectidx(v0, 2);
    float f0[8], f1[8];

    _mm256_storeu_ps(f0, v0);
    _mm256_storeu_ps(f1, v1);

    printf("v0: %g %g %g %g %g %g %g %g\n", f0[0], f0[1], f0[2], f0[3], f0[4], f0[5], f0[6], f0[7]);
    printf("v1: %g %g %g %g %g %g %g %g\n", f1[0], f1[1], f1[2], f1[3], f1[4], f1[5], f1[6], f1[7]);

    return 0;
}

测试：

$ gcc -Wall -mavx2 test_avx2_select.c && ./a.out
v0: 2 3 4 5 6 7 8 9
v1: 4 4 4 4 4 4 4 4

Answer 3

对于SSE来说，它更简单 - 你可以使用_mm_shuffle_ps：

#include <stdio.h>
#include <xmmintrin.h>

#define _mm_selectidx(v, i) _mm_shuffle_ps(v, v, _MM_SHUFFLE(i, i, i, i))

int main(void)
{
    __m128 v0 = _mm_setr_ps(2, 3, 4, 5);
    __m128 v1 = _mm_selectidx(v0, 2);
    float f0[4], f1[4];

    _mm_storeu_ps(f0, v0);
    _mm_storeu_ps(f1, v1);

    printf("v0: %g %g %g %g\n", f0[0], f0[1], f0[2], f0[3]);
    printf("v1: %g %g %g %g\n", f1[0], f1[1], f1[2], f1[3]);

    return 0;
}

测试：

$ gcc -Wall -msse test_sse_select.c && ./a.out
v0: 2 3 4 5
v1: 4 4 4 4

向量广播特定元素到另一个向量

3 个答案: