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不是我的强项，这是我需要做得更好的事情--请温柔一点。我想可能有许多调整，我可以应用到当前的矢量化代码(建议欢迎！)

发展环境：

v8.3.0

• c++14

• LinuxUbuntu18.04
• G++G++

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将展开循环(两个内部迭代)，这将稍微提高足够大的输入的性能。