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

Posted

技术标签:

【中文标题】将向量的特定元素广播到另一个向量【英文标题】:Broadcast specific element of vector to another vector 【发布时间】:2014-10-01 11:45:23 【问题描述】:

如何从 __m256 向量中的索引中提取单个浮点数,并将其广播到结果向量?

伪代码:

__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

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

【问题讨论】:

元素索引(上面示例中的 2)是编译时文字常量,还是需要是运行时值(即变量)? 这是一个立即数/常数 【参考方案1】:

对于仅 AVX(即没有 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 内部函数所需的编译时常量。

【讨论】:

聪明的解决方案(+1)。我想知道 AVX2 版本与仅 AVX 版本相比的性能如何。 谢谢 - AVX 版本可能更高效,但它仅限于元素索引的编译时常量,而 AVX2 版本将使用运行时元素索引。 正是我想要的,谢谢。这个解决方案(经过一些修改)也可以用于 SSE2 吗? 当然 - 不过 SSE 更容易 - 你只需要 _mm_shuffle_ps - 我会用 SSE 版本发布另一个答案。【参考方案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

【讨论】:

不幸的是,我只使用 AVX...也欢迎仅适用于 SSE2 的解决方案 好的 - 我现在添加了一个单独的 AVX-only 答案,它适用于元素索引的编译时间常数。【参考方案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

【讨论】:

以上是关于将向量的特定元素广播到另一个向量的主要内容,如果未能解决你的问题,请参考以下文章

C++ 将元素从一个向量移动到另一个向量

如何将唯一指针从一个向量移动到另一个唯一指针向量? [关闭]

将向量映射到特定范围

将指向向量向量的指针移动到特定位置

C++ 将对象向量中的元素复制到具有此元素的向量中

将项目从一个向量复制到另一个向量