c++ - AVX2 SIMD 内联 16 位到 8 位，反之亦然

Question

我有一个 C++(或类 C)函数，我正在尝试对其进行向量化。该函数是图像合成的众多变体之一，它采用色度 444 子采样的 Y、U 或 V 图像平面，并将 src 图像合成/覆盖到 dst 图像上(其中 src 图像还包含 alpha 透明度)。

#include <cstdint>


void composite(uint8_t *__restrict__ pSrc,  // Source plane
               uint8_t *__restrict__ pSrcA, // Source alpha plane 
               uint8_t *__restrict__ pDst,  // Destination plane
               const std::size_t nCount)    // Number of component pixels to process.
{
    for (std::size_t k = 0; k < nCount; ++k)
    {
        uint16_t w = (pSrc[k] * pSrcA[k]);
        uint16_t x = (255 - pSrcA[k]) * pDst[k];
        uint16_t y = w+x;
        uint16_t z = y / uint16_t{255};
        pDst[k] = static_cast<uint8_t>(z);
    }
}

在 AVX2 矢量化等效项中，我很难理解如何有效地读取 8 位、转换为 16 位，以及(在处理/合成之后)最终将 16 位样本转换回 8 位以存储回来内存。在读取方面，我使用中间 xmm 寄存器 - 这似乎不是最好的方法；我猜测混合寄存器系列时会出现性能损失。

我想出了(不完整):

#include <cstdint>

#include <immintrin.h>
#include <emmintrin.h>


///////////////////////////////////////////////////////////////////////////
// Credit: https://stackoverflow.com/questions/35285324/how-to-divide-16-bit-integer-by-255-with-using-sse
#define AVX2_DIV255_U16(x) _mm256_srli_epi16(_mm256_mulhi_epu16(x, _mm256_set1_epi16((short)0x8081)), 7)

///////////////////////////////////////////////////////////////////////////
/// Blends/composites/overlays two planes of Y, U, or V plane with 4:4:4 chroma subsampling over the other.
/// \param d The destination Y, U , or V component
/// \param s The source Y, U, or V component
/// \param sa The source alpha component
/// \param pixels The number of pixels that require processing.
/// \return The number of pixels processed.
int blend_plane_pixels_444_vectorized(uint8_t *__restrict__ d,
                                      uint8_t *__restrict__ s,
                                      uint8_t *__restrict__ sa,
                                      const int pixels)
{
    int n = 0; // Return number of component pixels processed.
    for (int k = 0; k + 32 <= pixels; k += 32)
    {
        // Load first 16 (unaligned) of d, s, sa
        // TODO: This efficient mixing xmm registers with ymm??
        auto vecD0 = _mm256_cvtepu8_epi16(_mm_loadu_si128((__m128i_u *)d));
        auto vecS0 = _mm256_cvtepu8_epi16(_mm_loadu_si128((__m128i_u *)s));
        auto vecSa0 = _mm256_cvtepu8_epi16(_mm_loadu_si128((__m128i_u *)sa));

        // Load second 16 (unaligned) of d, s, sa
        auto vd1 = _mm256_cvtepu8_epi16(_mm_loadu_si128((__m128i_u *)d+16));
        auto vs1 = _mm256_cvtepu8_epi16(_mm_loadu_si128((__m128i_u *)s+16));
        auto vsa1 = _mm256_cvtepu8_epi16(_mm_loadu_si128((__m128i_u *)sa+16));

        // Load 255 into register
        auto vec255 = _mm256_set1_epi16(255);

        // uint16_t w = (pSrc[k] * pSrcA[k]);
        auto vecW0 = _mm256_mullo_epi16(vecS0, vecSa0);
        auto vecW1 = _mm256_mullo_epi16(vs1, vsa1);

        // uint16_t x = (255 - pSrcA[k]) * pDst[k];
        auto vecX0 = _mm256_mullo_epi16(_mm256_subs_epu16(vec255, vecSa0), vecD0);
        auto vecX1 = _mm256_mullo_epi16(_mm256_subs_epu16(vec255, vsa1), vd1);

        // Load 127 into register
        auto vec127 = _mm256_set1_epi16(127);

        // uint16_t y = w+x;
        auto vecY0 = _mm256_adds_epu16(_mm256_adds_epu16(vecW0, vecX0), vec127);
        auto vecY1 = _mm256_adds_epu16(_mm256_adds_epu16(vecW1, vecX1), vec127);

        // uint16_t z = y / uint16_t{255};
        auto vecZ0 = AVX2_DIV255_U16(vecY0);
        auto vecZ1 = AVX2_DIV255_U16(vecY1);

        // TODO: How to get this back into 8-bit samples so that it can be stored
        //       back into array.
        auto vecResult = _mm256_blendv_epi8(vecZ0, vecZ1, _mm256_set1_epi16(127));

        // Write data back to memory (unaligned)
        _mm256_storeu_si256((__m256i*)d, vecResult);

        d += 32;
        s += 32;
        sa += 32;
        n += 32;
    }

    return n;
}

SIMD 不是我的强项，这是我需要提高的地方 - 请保持温柔。我想我可能可以对当前的矢量化代码进行许多调整(欢迎提出建议!)

开发环境:

• Linux Ubuntu 18.04
• G++ v8.3.0
• c++14

Answer 1

通常，如果您需要将结果重新打包为 8 位整数，最好使用 punpcklbw/punpckhbw 解包为零，然后重新打包使用 packuswb 得到的结果。或者有时您可以将奇数和偶数字节屏蔽到单独的寄存器中，一起进行计算和位或结果。

_mm256_cvtepu8_epi16/vpmovzxbw 的“问题”是它是车道交叉的(即，它仅从较低的 128 位一半(或内存)获取输入，但结果在上半部分和下半部分)，并且没有(简单)解决方案将来自不同 channel 的 16 位值连接回 1(直到 AVX512 channel 交叉单寄存器包指令具有饱和或截断)。

在您的情况下，您实际上可以将 d 和 s 值打包在一个寄存器中，并将 a 和 255-a 打包在一起 值，并使用 vpmaddubsw 进行乘法和加法。您需要从 d 和 s 值中减去 128，然后再将它们打包在一起，因为一个参数需要是有符号的 int8。结果将偏离 128*255，但可以进行补偿，特别是如果您添加 127 进行舍入的话。 (如果不这样做，可以在除法(有符号除法向下舍入)和重新打包后为每个字节添加 128。

未经测试的代码，使用与您的尝试相同的签名:

// https://stackoverflow.com/questions/35285324/how-to-divide-16-bit-integer-by-255-with-using-sse
inline __m256i div255_epu16(__m256i x) {
    __m256i mulhi = _mm256_mulhi_epu16(x, _mm256_set1_epi16(0x8081));
    return _mm256_srli_epi16(mulhi, 7);
}

int blend_plane_pixels_444_vectorized(uint8_t *__restrict__ d,
                                      uint8_t *__restrict__ s,
                                      uint8_t *__restrict__ sa,
                                      const int pixels)
{
    int n = 0; // Return number of component pixels processed.
    for (int k = 0; k + 32 <= pixels; k += 32)
    {
        // Load 32 (unaligned) of d, s, sa
        __m256i vecD = _mm256_loadu_si256((__m256i_u *)d);
        __m256i vecS = _mm256_loadu_si256((__m256i_u *)s );
        __m256i vecA = _mm256_loadu_si256((__m256i_u *)sa);

        // subtract 128 from D and S to have them in the signed domain
        // subtracting 128 is equivalent ot xor with 128
        vecD = _mm256_xor_si256(vecD, _mm256_set1_epi8(0x80));
        vecS = _mm256_xor_si256(vecS, _mm256_set1_epi8(0x80));

        // calculate 255-a (equivalent to 255 ^ a):
        __m256i vecA_ = _mm256_xor_si256(vecA, _mm256_set1_epi8(0xFF));

        __m256i vecAA_lo = _mm256_unpacklo_epi8(vecA, vecA_);
        __m256i vecSD_lo = _mm256_unpacklo_epi8(vecS, vecD);
        __m256i vecAA_hi = _mm256_unpackhi_epi8(vecA, vecA_);
        __m256i vecSD_hi = _mm256_unpackhi_epi8(vecS, vecD);

        // R = a * (s-128) + (255-a)*(d-128) = a*s + (255-a)*d - 128*255
        __m256i vecR_lo = _mm256_maddubs_epi16(vecAA_lo,vecSD_lo);
        __m256i vecR_hi = _mm256_maddubs_epi16(vecAA_hi,vecSD_hi);

        // shift back to unsigned domain and add 127 for rounding
        vecR_lo = _mm256_add_epi16(vecR_lo, _mm256_set1_epi16(127+128*255));
        vecR_hi = _mm256_add_epi16(vecR_hi, _mm256_set1_epi16(127+128*255));

        // divide (rounding down)
        vecR_lo = div255_epu16(vecR_lo);
        vecR_hi = div255_epu16(vecR_hi);

        // re-join lower and upper half:
        __m256i vecResult = _mm256_packus_epi16(vecR_lo, vecR_hi);
        // Write data back to memory (unaligned)
        _mm256_storeu_si256((__m256i*)d, vecResult);

        d += 32;
        s += 32;
        sa += 32;
        n += 32;
    }

    return n;
}

Godbolt 链接:https://godbolt.org/z/EYzLw2 请注意，-march=haswell 或您想要支持的任何体系结构都至关重要，因为否则 gcc 将不会使用未对齐的数据作为内存源操作数。当然，适用一般矢量化规则，即，如果您可以控制对齐方式，则更喜欢分配对齐的数据。如果没有，您可以剥离第一个未对齐的字节(例如，来自 d)，以至少进行一次加载和存储对齐。

Clang 将展开循环(到两次内部迭代)，这将稍微提高足够大输入的性能。