AVX2 SIMD下8/16位互转及图像合成函数向量化优化问询

Hey there! Let's break down your AVX2 vectorization challenges and fix up the code step by step. I'll keep things clear and avoid overly technical jargon where possible.

1.1 Efficient 8-bit ↔ 16-bit Conversion #

• 8-bit to 16-bit: Your current use of _mm256_cvtepu8_epi16 is correct—it zero-extends 16x 8-bit values from an XMM register into 16x 16-bit values in a YMM register. For 32 pixels, splitting into two 16-pixel loads is the right approach since AVX2 doesn't have a 256-bit wide 8→16 extension instruction.
• 16-bit to 8-bit: The missing piece is _mm256_packus_epi16, which safely packs 32x 16-bit unsigned values (each in 0-255 range, since our calculation results are clamped to 8-bit) into 32x 8-bit values in a single YMM register. This instruction handles saturation automatically, so no overflow issues here.

1.2 XMM/YMM Register Mixing Performance #

Don't worry—mixing XMM loads with YMM operations doesn't cause performance loss on AVX2-capable CPUs. The XMM registers are just the lower 128 bits of YMM registers, so extending an XMM value to YMM with _mm256_cvtepu8_epi16 is a zero-cost operation (no extra cycles). The only time you'd see a penalty is if you switch between AVX and legacy SSE instructions without proper state management, which isn't happening here.

1.3 16-bit Result to 8-bit Storage #

Your original _mm256_blendv_epi8 approach was incorrect because it tried to blend 16-bit values as 8-bit. Instead, use _mm256_packus_epi16 to combine vecZ0 and vecZ1 into a single 256-bit register of 8-bit values, then store it directly to memory.

Here are some tweaks to make your vectorized code faster and more robust:

• Move constant registers outside the loop: Values like vec255 and vecDiv255 don't change per iteration, so initialize them once before the loop to avoid redundant instructions.
• Handle remaining pixels: Your current loop only processes multiples of 32 pixels. Add a scalar fallback for any remaining pixels to ensure all data is processed.
• Align memory (if possible): If you can guarantee your input/output pointers are 32-byte aligned, switch from _mm_loadu_si128/_mm256_storeu_si256 to _mm_load_si128/_mm256_store_si256—aligned loads/stores are faster on most CPUs.
• Match scalar behavior: Your original scalar code uses integer division (truncating towards zero), but your vectorized code added vec127 for rounding. I've adjusted the code to match the scalar truncation behavior; if you want rounding instead, just re-add the vec127 addition.

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

///////////////////////////////////////////////////////////////////////////
// Truncating division by 255 for 16-bit unsigned integers (matches scalar behavior)
// Formula: (x * 0x8081) >> 23
#define AVX2_DIV255_U16_TRUNC(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;

    // Initialize constant registers once outside the loop
    const __m256i vec255 = _mm256_set1_epi16(255);

    // Process 32 pixels per iteration
    for (; n + 32 <= pixels; n += 32)
    {
        // Load 16 pixels (unaligned) and extend to 16-bit
        const __m128i xmmD0 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(d));
        const __m128i xmmS0 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(s));
        const __m128i xmmSa0 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(sa));
        const __m256i vecD0 = _mm256_cvtepu8_epi16(xmmD0);
        const __m256i vecS0 = _mm256_cvtepu8_epi16(xmmS0);
        const __m256i vecSa0 = _mm256_cvtepu8_epi16(xmmSa0);

        // Load next 16 pixels (unaligned) and extend to 16-bit
        const __m128i xmmD1 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(d + 16));
        const __m128i xmmS1 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(s + 16));
        const __m128i xmmSa1 = _mm_loadu_si128(reinterpret_cast<const __m128i*>(sa + 16));
        const __m256i vecD1 = _mm256_cvtepu8_epi16(xmmD1);
        const __m256i vecS1 = _mm256_cvtepu8_epi16(xmmS1);
        const __m256i vecSa1 = _mm256_cvtepu8_epi16(xmmSa1);

        // Calculate w = s * sa
        const __m256i vecW0 = _mm256_mullo_epi16(vecS0, vecSa0);
        const __m256i vecW1 = _mm256_mullo_epi16(vecS1, vecSa1);

        // Calculate x = (255 - sa) * d
        const __m256i vecInvSa0 = _mm256_sub_epi16(vec255, vecSa0);
        const __m256i vecInvSa1 = _mm256_sub_epi16(vec255, vecSa1);
        const __m256i vecX0 = _mm256_mullo_epi16(vecInvSa0, vecD0);
        const __m256i vecX1 = _mm256_mullo_epi16(vecInvSa1, vecD1);

        // Calculate y = w + x (no rounding, matches scalar truncation)
        const __m256i vecY0 = _mm256_add_epi16(vecW0, vecX0);
        const __m256i vecY1 = _mm256_add_epi16(vecW1, vecX1);

        // Calculate z = y / 255 (truncating)
        const __m256i vecZ0 = AVX2_DIV255_U16_TRUNC(vecY0);
        const __m256i vecZ1 = AVX2_DIV255_U16_TRUNC(vecY1);

        // Pack 16-bit results into 8-bit values
        const __m256i vecResult = _mm256_packus_epi16(vecZ0, vecZ1);

        // Store back to memory (unaligned)
        _mm256_storeu_si256(reinterpret_cast<__m256i*>(d), vecResult);

        // Advance pointers
        d += 32;
        s += 32;
        sa += 32;
    }

    // Handle remaining pixels with scalar code
    for (; n < pixels; ++n)
    {
        const uint16_t w = static_cast<uint16_t>(s[n]) * sa[n];
        const uint16_t x = static_cast<uint16_t>(255 - sa[n]) * d[n];
        const uint16_t y = w + x;
        const uint16_t z = y / 255;
        d[n] = static_cast<uint8_t>(z);
    }

    return pixels;
}

• Rounding vs Truncation: If you want rounding instead of truncation (to match the behavior of your original vectorized code's vec127 addition), modify the vecY0/vecY1 lines to:

  const __m256i vec127 = _mm256_set1_epi16(127);
  const __m256i vecY0 = _mm256_add_epi16(_mm256_add_epi16(vecW0, vecX0), vec127);
  const __m256i vecY1 = _mm256_add_epi16(_mm256_add_epi16(vecW1, vecX1), vec127);
  

  And keep using the same AVX2_DIV255_U16_TRUNC macro—it will now perform rounded division.
• Compiler Flags: When compiling, use -mavx2 -O3 -std=c++14 to enable AVX2 support and optimizations. G++ 8.3.0 fully supports these flags on Ubuntu 18.04.

内容的提问来源于stack exchange，提问作者ZeroDefect