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位值重新连接回1（直到AVX512通道交叉具有饱和或截断的单寄存器组指令）。

在您的情况下，您实际上可以将dand s值包装在一个寄存器中，并将aand 255-a值包装在另一个寄存器中，并vpmaddubsw用于乘法和加法。将一个和参数打包在一起之前，您需要从d和s值中减去128 ，因为一个参数需要是一个符号int8。结果将以关闭128*255，但无法补偿，特别是如果您127仍然为舍入而加。（如果不这样做，则可以在除（四舍五入的有符号除法）并重新打包后，向每个字节添加128。

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

// /sf/ask/2469972711/
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-Link：https ://godbolt.org/z/EYzLw2 注意，-march=haswell无论您要支持哪种架构，都至关重要，因为否则gcc不会将未对齐的数据用作内存源操作数。当然，一般的矢量化规则都适用，即，如果您可以控制对齐方式，则最好分配对齐的数据。如果没有，则可以剥离第一个未对齐的字节（例如从d），以使至少一个负载和存储对齐。

Clang将展开循环（进行两次内部迭代），这对于输入足够大的内容会稍微提高性能。