如何用AVX2实现8位灰度图转RGBA_8888预乘alpha图像？

核心思路 #

利用AVX2的256位向量寄存器，一次处理32个灰度字节（对应32个RGBA像素）：

• 将32个灰度字节拆分为4组，每组8个字节
• 对每组8个字节，将每个字节无符号扩展为32位uint
• 分别计算每个通道的预乘值（灰度值×颜色值后右移8位），并移位到RGBA对应的字节位置
• 合并四个通道的值为完整的RGBA uint，批量存储到输出数组

关键指令说明 #

• Avx.LoadVector256：不对齐加载32个灰度字节到向量寄存器（兼容任意内存地址，避免对齐异常）
• Sse2.ZeroExtend：将8字节向量无符号扩展为16位短整型向量
• Avx2.ZeroExtend：将16位短整型向量无符号扩展为32位整型向量
• Avx2.Multiply：向量元素逐位相乘
• Avx2.ShiftRightLogical/Avx2.ShiftLeftLogical：向量元素逻辑移位
• Avx2.Or：向量元素逐位或，合并通道值
• Avx2.Store：将向量结果批量存储到输出数组

完整实现代码 #

using System;
using System.Numerics;
using System.Runtime.CompilerServices;
using System.Runtime.Intrinsics;
using System.Runtime.Intrinsics.X86;

public class Test
{
    const int ImageSize = 2048;
    const int ImageLength = ImageSize * ImageSize;
    private byte[] _bytesGray = new byte[ImageLength];
    private uint[] _pixelsRGBA = new uint[ImageLength];

    private const byte _colorR = 0xFF;
    private const byte _colorG = 0x01;
    private const byte _colorB = 0x02;
    private const byte _colorA = 0xFF;

    [GlobalSetup]
    public void Setup()
    {
        for (int i = 0; i < ImageLength; i++)
        {
            _bytesGray[i] = (byte)(i + 1);
            _pixelsRGBA[i] = 0;
        }
    }

    [Benchmark]
    public unsafe void GrayscaleToColor_Scalar()
    {
        fixed (byte* bytePtr = _bytesGray)
        fixed (uint* pixelPtr = _pixelsRGBA)
        {
            for (int i = 0; i < ImageLength; ++i)
            {
                byte value = bytePtr[i];
                byte r = (byte)((value * _colorR) >> 8);
                byte g = (byte)((value * _colorG) >> 8);
                byte b = (byte)((value * _colorB) >> 8);
                byte a = (byte)((value * _colorA) >> 8);

                pixelPtr[i] = (uint)(r << 24 | g << 16 | b << 8 | a);
            }
        }
    }

    [Benchmark]
    public unsafe void GrayscaleToColor_AVX2()
    {
        if (!Avx2.IsSupported)
            throw new NotSupportedException("AVX2 is not supported on this system.");

        var rConst = Vector256.Create((uint)_colorR);
        var gConst = Vector256.Create((uint)_colorG);
        var bConst = Vector256.Create((uint)_colorB);
        var aConst = Vector256.Create((uint)_colorA);

        fixed (byte* grayPtr = _bytesGray)
        fixed (uint* rgbaPtr = _pixelsRGBA)
        {
            int i = 0;
            // 处理完整的32字节块
            for (; i <= ImageLength - 32; i += 32)
            {
                // 加载32个灰度字节
                var gray32 = Avx.LoadVector256(grayPtr + i);
                
                // 拆分为4组8字节块处理
                var gray8_0 = Sse2.LoadVector64(grayPtr + i);
                var gray8_1 = Sse2.LoadVector64(grayPtr + i + 8);
                var gray8_2 = Sse2.LoadVector64(grayPtr + i + 16);
                var gray8_3 = Sse2.LoadVector64(grayPtr + i + 24);
                
                // 生成每组对应的RGBA向量
                var rgba0 = Process8Pixels(gray8_0, rConst, gConst, bConst, aConst);
                var rgba1 = Process8Pixels(gray8_1, rConst, gConst, bConst, aConst);
                var rgba2 = Process8Pixels(gray8_2, rConst, gConst, bConst, aConst);
                var rgba3 = Process8Pixels(gray8_3, rConst, gConst, bConst, aConst);
                
                // 批量存储结果
                Avx2.Store(rgbaPtr + i, rgba0);
                Avx2.Store(rgbaPtr + i + 8, rgba1);
                Avx2.Store(rgbaPtr + i + 16, rgba2);
                Avx2.Store(rgbaPtr + i + 24, rgba3);
            }
            
            // 处理剩余不足32个的像素（标量兜底）
            for (; i < ImageLength; i++)
            {
                byte value = grayPtr[i];
                byte r = (byte)((value * _colorR) >> 8);
                byte g = (byte)((value * _colorG) >> 8);
                byte b = (byte)((value * _colorB) >> 8);
                byte a = (byte)((value * _colorA) >> 8);
                rgbaPtr[i] = (uint)(r << 24 | g << 16 | b << 8 | a);
            }
        }
    }

    [MethodImpl(MethodImplOptions.AggressiveInlining)]
    private static Vector256<uint> Process8Pixels(Vector64<byte> gray8, Vector256<uint> rConst, Vector256<uint> gConst, Vector256<uint> bConst, Vector256<uint> aConst)
    {
        // 将8个字节无符号扩展为16位短整型（每个元素为0x00xx）
        var grayU16 = Sse2.ZeroExtend(gray8);
        // 进一步扩展为32位整型（每个元素为0x000000xx）
        var grayU32 = Avx2.ZeroExtend(grayU16);

        // 计算R通道：(灰度×R)>>8 后移到最高字节位
        var rVal = Avx2.Multiply(grayU32, rConst);
        rVal = Avx2.ShiftRightLogical(rVal, 8);
        rVal = Avx2.ShiftLeftLogical(rVal, 24);

        // 计算G通道：(灰度×G)>>8 后移到次高字节位
        var gVal = Avx2.Multiply(grayU32, gConst);
        gVal = Avx2.ShiftRightLogical(gVal, 8);
        gVal = Avx2.ShiftLeftLogical(gVal, 16);

        // 计算B通道：(灰度×B)>>8 后移到次低字节位
        var bVal = Avx2.Multiply(grayU32, bConst);
        bVal = Avx2.ShiftRightLogical(bVal, 8);
        bVal = Avx2.ShiftLeftLogical(bVal, 8);

        // 计算A通道：(灰度×A)>>8 保留在最低字节位
        var aVal = Avx2.Multiply(grayU32, aConst);
        aVal = Avx2.ShiftRightLogical(aVal, 8);

        // 合并四个通道为完整RGBA值
        var result = Avx2.Or(rVal, gVal);
        result = Avx2.Or(result, bVal);
        result = Avx2.Or(result, aVal);

        return result;
    }
}

优化要点 #

1. 内存对齐：如果输入数组是32字节对齐的，可将Avx.LoadVector256替换为Avx2.LoadAlignedVector256进一步提升性能
2. 循环效率：按32字节块批量处理，最大化向量寄存器利用率
3. 指令简化：利用ZeroExtend指令替代手动拆包操作，代码更简洁且执行效率更高
4. 边界兼容：剩余不足32个的像素用标量逻辑处理，避免复杂的向量边界判断

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