AVX2 SIMD Instrinsics 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 位半（或内存）获取输入，但结果在高位和低位一半），并且没有（简单的）解决方案将来自不同通道的 16 位值连接回一个（直到 AVX512 通道交叉单寄存器包指令具有饱和或截断）。

在您的情况下，您实际上可以将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;
}

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

Clang 将展开循环（到两个内部迭代），这将略微提高足够大输入的性能。