如何高效实现2D半字节矩阵转置？（支持AVX2扩展）

问题说明 #

我们有一个以16字节uint8_t数组存储的4×8半字节矩阵，数组中每个字节遵循公式(j << 4) | i生成——即高4位对应矩阵中的半字节值j，低4位对应另一个半字节值i。

举个具体例子：
如下4×8矩阵：

0  1  2  3  3  7  1  9 
    4  5  6  7  4  1  6  15
    8  9  10 11 3  14 6  11
    12 13 14 15 8  10 7  4

对应的存储数组为：

const uint8_t matrix[] = {
    0x10, 0x32, 0x73, 0x91,
    0x54, 0x76, 0x14, 0xf6,
    0x98, 0xba, 0xe3, 0xb6,
    0xdc, 0xfe, 0xa8, 0x47,
};

我们需要将其转换为目标数组：

const uint8_t result[] = {
    0x40, 0xc8, 0x51, 0xd9,
    0x62, 0xea, 0x73, 0xfb,
    0x43, 0x83, 0x17, 0xae,
    0x61, 0x76, 0xf9, 0x4b,
};

要求实现最高效的转换函数，且支持AVX2指令集。

现有实现代码（基于x64 SIMD半字节洗牌技术） #

以下是基于《Nibble shuffling with x64 SIMD》实现的C语言代码，核心逻辑是将矩阵拆分为两个64位输入，拆分半字节、打乱顺序后重新打包：

__m128i unpack_nibbles(__m128i src) {
    __m128i nibbles_hi = _mm_srli_epi64(src, 4);

    // 交错高半字节与完整字节，再清除每个字节的高4位
    __m128i unpacked = _mm_unpacklo_epi8(src, nibbles_hi);
    return _mm_and_si128(unpacked, _mm_set1_epi8(0xf));
}

void transpose_4x8_nibbles(uint8_t *src, uint8_t *dst) {
    uint8_t *src_lo = src + 0x8;

    __m128i data_hi = _mm_loadl_epi64((__m128i*)src);
    __m128i data_lo = _mm_loadl_epi64((__m128i*)src_lo);

    data_hi = unpack_nibbles(data_hi);
    data_lo = unpack_nibbles(data_lo);

    // 执行转置的洗牌操作
    __m128i transpose_mask = _mm_setr_epi8(0, 0x8, 0x1, 0x9, 0x2, 0xa, 0x3, 0xb, 0x4, 0xc, 0x5, 0xd, 0x6, 0xe, 0x7, 0xf);
    data_hi = _mm_shuffle_epi8(data_hi, transpose_mask);
    data_lo = _mm_shuffle_epi8(data_lo, transpose_mask);

    // 重新将半字节打包为完整字节
    __m128i pack_mask = _mm_set1_epi16(0x1001);
    data_hi = _mm_maddubs_epi16(data_hi, pack_mask);  // 偶数位置字节乘0x10（移到高4位），奇数位置乘0x01（保留低4位）
    data_lo = _mm_maddubs_epi16(data_lo, pack_mask);
    
    __m128i data = _mm_packus_epi16(data_hi, data_lo);
    data = _mm_shuffle_epi8(data, transpose_mask);
    
    _mm_store_si128((__m128i*) dst, data);
}

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