基于比较将浮点变量设置为 0.0f 或 1.0f 的 SSE 代码

Question

我有两个数组：char* c 和 float* f 我需要做这个操作：

// Compute float mask
float* f;
char* c;
char c_thresh;
int n;

for ( int i = 0; i < n; ++i )
{
    if ( c[i] < c_thresh ) f[i] = 0.0f;
    else                   f[i] = 1.0f;
}

我正在寻找一种快速的方法：没有条件并尽可能使用 SSE（4.2 或 AVX）。

如果使用 float 而不是 char 可以产生更快的代码，我可以将代码更改为仅使用浮点数：

// Compute float mask
float* f;
float* c;
float c_thresh;
int n;

for ( int i = 0; i < n; ++i )
{
    if ( c[i] < c_thresh ) f[i] = 0.0f;
    else                   f[i] = 1.0f;
}

谢谢

Answer 1

非常简单，只需进行比较，将字节转换为 dword，并使用 1.0f：（未经测试，这并不是要复制和粘贴代码，而是要展示你是如何做到的）

movd xmm0, [c]          ; read 4 bytes from c
pcmpgtb xmm0, threshold ; compare (note: comparison is >, not >=, so adjust threshold)
pmovzxbd xmm0, xmm0     ; convert bytes to dwords
pand xmm0, one          ; AND all four elements with 1.0f
movdqa [f], xmm0        ; save result

应该很容易转换为内在函数。

Answer 2

以下代码使用 SSE2（我认为）。

它在一条指令中执行 16 次字节比较 (_mm_cmpgt_epi8)。它假定char 已签名；如果您的 char 未签名，则需要额外的摆弄（翻转每个 char 的最高有效位）。

它唯一不标准的做法是使用幻数3f80 来表示浮点常量1.0。幻数实际上是0x3f800000，但是 16 LSB 为零这一事实使得可以更有效地进行位摆弄（使用 16 位掩码而不是 32 位掩码）。

// load (assuming the pointer is aligned)
__m128i input = *(const __m128i*)c;
// compare
__m128i cmp = _mm_cmpgt_epi8(input, _mm_set1_epi8(c_thresh - 1));
// convert to 16-bit
__m128i c0 = _mm_unpacklo_epi8(cmp, cmp);
__m128i c1 = _mm_unpackhi_epi8(cmp, cmp);
// convert ffff to 3f80
c0 = _mm_and_si128(c0, _mm_set1_epi16(0x3f80));
c1 = _mm_and_si128(c1, _mm_set1_epi16(0x3f80));
// convert to 32-bit and write (assuming the pointer is aligned)
__m128i* result = (__m128i*)f;
result[0] = _mm_unpacklo_epi16(_mm_setzero_si128(), c0);
result[1] = _mm_unpackhi_epi16(_mm_setzero_si128(), c0);
result[2] = _mm_unpacklo_epi16(_mm_setzero_si128(), c1);
result[3] = _mm_unpackhi_epi16(_mm_setzero_si128(), c1);

Answer 3

通过切换到浮点数，您可以在 GCC 中自动矢量化循环，而不必担心内在函数。以下代码将执行您想要的操作并自动矢量化。

void foo(float *f, float*c, float c_thresh, const int n) {
    for (int i = 0; i < n; ++i) {
        f[i] = (float)(c[i] >= c_thresh);
    }
}

编译

g++  -O3 -Wall  -pedantic -march=native main.cpp -ftree-vectorizer-verbose=1 

您可以在coliru 查看结果并自己编辑/编译代码。但是，MSVC2013 没有对循环进行矢量化。

Answer 4

怎么样：

f[i] = (c[i] >= c_thresh);

至少这消除了条件。

Answer 5

AVX 版本：

void floatSelect(float* f, const char* c, size_t n, char c_thresh) {
    for (size_t i = 0; i < n; ++i) {
        if (c[i] < c_thresh) f[i] = 0.0f;
        else f[i] = 1.0f;
    }
}

void vecFloatSelect(float* f, const char* c, size_t n, char c_thresh) {
    const auto thresh = _mm_set1_epi8(c_thresh);
    const auto zeros = _mm256_setzero_ps();
    const auto ones = _mm256_set1_ps(1.0f);
    const auto shuffle0 = _mm_set_epi8(3, -1, -1, -1, 2, -1, -1, -1, 1, -1, -1, -1, 0, -1, -1, -1);
    const auto shuffle1 = _mm_set_epi8(7, -1, -1, -1, 6, -1, -1, -1, 5, -1, -1, -1, 4, -1, -1, -1);
    const auto shuffle2 = _mm_set_epi8(11, -1, -1, -1, 10, -1, -1, -1, 9, -1, -1, -1, 8, -1, -1, -1);
    const auto shuffle3 = _mm_set_epi8(15, -1, -1, -1, 14, -1, -1, -1, 13, -1, -1, -1, 12, -1, -1, -1);

    const size_t nVec = (n / 16) * 16;
    for (size_t i = 0; i < nVec; i += 16) {
        const auto chars = _mm_loadu_si128(reinterpret_cast<const __m128i*>(c + i));
        const auto mask = _mm_cmplt_epi8(chars, thresh);
        const auto floatMask0 = _mm_shuffle_epi8(mask, shuffle0);
        const auto floatMask1 = _mm_shuffle_epi8(mask, shuffle1);
        const auto floatMask2 = _mm_shuffle_epi8(mask, shuffle2);
        const auto floatMask3 = _mm_shuffle_epi8(mask, shuffle3);
        const auto floatMask01 = _mm256_set_m128i(floatMask1, floatMask0);
        const auto floatMask23 = _mm256_set_m128i(floatMask3, floatMask2);
        const auto floats0 = _mm256_blendv_ps(ones, zeros, _mm256_castsi256_ps(floatMask01));
        const auto floats1 = _mm256_blendv_ps(ones, zeros, _mm256_castsi256_ps(floatMask23));
        _mm256_storeu_ps(f + i, floats0);
        _mm256_storeu_ps(f + i + 8, floats1);
    }
    floatSelect(f + nVec, c + nVec, n % 16, c_thresh);
}

Answer 6

转换为

f[i] = (float)(c[i] >= c_thresh);

- 也可以通过英特尔编译器自动矢量化（其他人提到的 gcc 也是如此）

如果您需要自动矢量化一些一般的分支循环， - 您也可以尝试 #pragma ivdep 或 pragma simd (最后一个是Intel Cilk Plus 和 OpenMP 4.0 标准的一部分）。这些编译指示 auto-vectorize 以可移植的方式为 SSE、AVX 和未来的矢量扩展（如 AVX512）提供代码。 Intel Compiler（所有已知版本）、Cray 和 PGI 编译器（仅限 ivdep）、可能即将发布的 GCC4.9 版本支持这些 pragma，并且从 VS2012 开始部分支持 MSVC（仅限 ivdep）。 p>

对于给定的示例，我没有更改任何内容（保留 if 和 char*），只是添加了 pragma ivdep：

void foo(float *f, char*c, char c_thresh, const int n) {
    #pragma ivdep
    for ( int i = 0; i < n; ++i )
    {
        if ( c[i] < c_thresh ) f[i] = 0.0f;
        else                   f[i] = 1.0f;
    }
}

在不支持 AVX 的 Core i5 上（仅限 SSE3），对于 n = 32K (32000000)，随机生成 c[i] 并使用 c_thresh 等于 0（我们使用有符号字符），给定代码提供大约 ~5x由于 ICL 的矢量化而加速。

完整测试（带有额外的测试用例正确性检查）可用here（它是 coliru，即仅 gcc4.8，没有 ICL/Cray；这就是它不在 coliru env 中矢量化的原因）。

应该可以通过处理更多的预取、对齐和类型转换编译指示/优化来进一步优化性能。对于给定的简单情况，也可以使用添加限制关键字（或 restrict，具体取决于使用的编译器）代替 ivdep/simd，而对于更一般的情况 - 编译指示 simd/ivdep 是最强大的。

注意：实际上#pragma ivdep“指示编译器忽略假定的跨迭代依赖项”（粗略说，如果您并行化相同的循环，那些会导致数据竞争的人）。由于众所周知的原因，编译器在这些假设中非常保守。在给定的情况下，显然没有读后写或写后读依赖。如果需要，可以使用诸如Advisor XE 正确性分析之类的动态工具，至少在给定的工作负载上验证这种依赖关系的存在，就像下面我的 cmets 中显示的那样。