高效转置二维nibble矩阵？

Question

给定一个2D的4x8的nibble矩阵，用16字节的uint8_t数组表示。对于每一对nibbles i和j，字节计算如下：(j << 4) | i。

例如，给定以下矩阵：

    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扩展。

这是我目前基于 Nibble shuffling with x64 SIMD的C实现。它将矩阵分成两个64位输入，解压半字节，重新排列并重新打包。

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

    //Interlave high nibbles with full nibbles [0000 hi, hi lo, ...] and clear high
    __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);

    //Transpose
    __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);

    //Pack nibbles
    __m128i pack_mask = _mm_set1_epi16(0x1001);
    data_hi = _mm_maddubs_epi16(data_hi, pack_mask);  //even bytes are multiplied by 0x10, odd bytes by 0x01
    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);
}

- ValentiMS

---

@DavidRanieri：void*指针数学运算是GNU扩展（类似于char*），但问题已经被编辑为可移植的。我们可以删除这些评论。 - Peter Cordes

感觉应该有一种方法可以在不进行解包/重新打包的情况下完成此操作。比如将字节洗牌，使所需的对齐，然后使用AND/ANDN/OR进行垂直位混合（然后再次洗牌以将合并的字节放置在其位置）。但是当我们需要从两个输入中获取相同的半字节时，例如输出的第一个字节为0x40，来自每个输入字节的低半字节（前两行的开头）0x10；0x54。也许将一个输入左移4位可以与2个混合配合使用？ - Peter Cordes

我怀疑AVX-512 VBMI对于vpmultishiftqb（并行位域提取）可能会有所帮助，也许还可以结合16位SIMD旋转（也是AVX-512）。或者使用Galois Field GFNI进行位重排（http://0x80.pl/articles/avx512-galois-field-for-bit-shuffling.html）。因此，如果您有兴趣为这些ISA扩展制作版本，可能会有所收获。 - Peter Cordes

我认为我有一个解决方案，使用1个pshufb、1个pmaddwd和4个位运算（不幸的是，pmaddwd需要有符号输入）。我会在今天晚些时候写出答案。 - chtz

Answer 1

让我们按照以下方式命名字节（所有内容均按小端序）：

X0 Y0 X1 Y1 X2 Y2 X3 Y3
Z0 W0 Z1 W1 Z2 W2 Z3 W3
X4 Y4 X5 Y5 X6 Y6 X7 Y7
Z4 W4 Z5 W5 Z6 W6 Z7 W7

在转置后，我们注意到X个半字节留在低半字节中，W个半字节留在高半字节中，Y个半字节从高半字节移动到低半字节，Z个半字节从低半字节移动到高半字节：

X0 Z0 X4 Z4
Y0 W0 Y4 W4
X1 Z1 X5 Z5
Y1 W1 Y5 W5
X2 Z2 X6 Z6
Y2 W2 Y6 W6
X3 Z3 X7 Z7
Y3 W3 Y7 W7

这意味着可以通过简单的pshufb正确地放置X和W nibble。所有Z nibble都需要向上移位（或乘以0x10），Y nibble需要向下移位（或将uint16块乘以0x1000并获取结果的上半部分）。一个00 Z0 00 Z4 Y0 00 Y4 00的块实际上是一个32位整数，我们几乎可以直接通过单个带有0x10和0x1000的pmaddwd指令从Z0 00 Z4 00和00 Y0 00 Y4中获取它。

00 Z0 00 Z4 Y0 00 Y4 00 = (00 Y0 00 Y4)* 0x1000 + (Z0 00 Z4 00) * 0x10

这些 nibbles 恰好在同一个字节中，即 X0，X4 和 W0，W4，因此只需要一个 pshufb 就可以按照要求排列字节。但不幸的是，如果 Y4>7，我们将得到一个负整数，需要再次屏蔽掉一些位（至少可以重复使用相同的掩码）。

总体而言，以下函数应该能胜任：

void transpose_4x8_nibbles(uint8_t const *src, uint8_t *dst) {
    __m128i const input = _mm_loadu_si128((__m128i const*)src);

    __m128i const shuff = _mm_shuffle_epi8(input, _mm_setr_epi8(0, 8, 4, 12, 1, 9, 5, 13, 2, 10, 6, 14, 3, 11, 7, 15));
    __m128i const mask = _mm_set1_epi32(0x0f0ff0f0);
    __m128i const XW = _mm_andnot_si128(mask, shuff);
    __m128i const YZ = _mm_and_si128(mask, shuff);
    __m128i const YZ_trans = _mm_madd_epi16(YZ, _mm_set1_epi32(0x00101000));
    __m128i const result = _mm_or_si128(XW, _mm_and_si128(mask, YZ_trans));

    _mm_storeu_si128((__m128i*)dst, result);
}

Godbolt演示（仅需要SSSE3，因为有pshufb）：https://godbolt.org/z/c43oTz43r

- chtz