SIMD位重新排序打包的12位整数数组

Question

我有一个大而紧密排列的12位整数数组，其具有以下重复的位压缩模式：（其中An/Bn中的n表示位数，A和B是数组中的前两个12位整数）

|           byte0           |            byte1          |           byte2         | etc..
| A11 A10 A9 A8 A7 A6 A5 A4 | B11 B10 B9 B8 B7 B6 B5 B4 | B3 B2 B1 B0 A3 A2 A1 A0 | etc..

我稍微重新排列成以下模式：

|           byte0           |            byte1          |           byte2         | etc..
| A11 A10 A9 A8 A7 A6 A5 A4 | A3 A2 A1 A0 B11 B10 B9 B8 | B7 B6 B5 B4 B3 B2 B1 B0 | etc..

我已经用以下代码在一个每3字节循环中使其工作。

void CSI2toBE12(uint8_t* pCSI2, uint8_t* pBE, uint8_t* pCSI2LineEnd)
{
    while (pCSI2 < pCSI2LineEnd) {
        pBE[0] = pCSI2[0];
        pBE[1] = ((pCSI2[2] & 0xf) << 4) | (pCSI2[1] >> 4);
        pBE[2] = ((pCSI2[1] & 0xf) << 4) | (pCSI2[2] >> 4);
        
        // Go to next 12-bit pixel pair (3 bytes)
        pCSI2 += 3;
        pBE += 3;
    }
}

但是以字节粒度进行操作对性能来说并不理想。目标 CPU 是 64 位 ARM Cortex-A72（树莓派计算模块 4）。为了背景，这段代码将 MIPI CSI-2 位打包的原始图像数据转换为 Adobe DNG 的位打包形式。

我希望能够通过使用 SIMD 指令集获得显著的性能改进，但我不太确定从何处开始。我已经有了用于翻译指令集的 SIMDe 头文件，因此欢迎使用 AVX/AVX2 解决方案。

- Russell Newman

---

SO不是一个免费的代码编写服务。由于这个问题本身的性质，你可能会收到一些负面评价。要手动使用内部函数进行矢量化，你应该使用目标ISA，即NEON。由于你正在处理3字节的倍数，你可能会在有效的方式下遇到一些困难来进行矢量化。 - Simon Goater

ARM64（NEON）在这里有一个理想的ld3指令，可以加载48个字节并将它们解压到三个寄存器中。因此，您将在一个寄存器中拥有所有的byte0，以此类推。我认为相应的内部函数是vld3q_u8，返回uint8x16x3_t。然后，您可以像当前代码一样进行移位和掩码操作，并使用st3（vst3q_u8）重新压缩和存储。 - Nate Eldredge

我对AVX不太了解，但我认为它没有ld3/st3的等效指令，所以SIMDe可能无法发挥作用。 - Nate Eldredge

在https://gcc.gnu.org/bugzilla/show_bug.cgi?id=110780和https://github.com/llvm/llvm-project/issues/64050分别报告了奇怪的多余位移负载。 - Nate Eldredge

你的天真函数在使用 gcc -O3 -msse4 编译时，在 Intel 上实现了向量化。在我的旧笔记本上，从没有 -msse4 的每字节 4.5 个周期减少到在一个 96MB 缓冲区上测试时的 1.3 个周期。 - Simon Goater

Answer 1

NEON的ld3指令非常适合这个任务；它可以加载48字节并将其解压缩到三个NEON寄存器中。然后只需要进行一些移位和或操作。

我想出了以下解决方案：

void vectorized(const uint8_t* pCSI2, uint8_t* pBE, const uint8_t* pCSI2LineEnd)
{
    while (pCSI2 < pCSI2LineEnd) {
        uint8x16x3_t in = vld3q_u8(pCSI2);
        uint8x16x3_t out;
        out.val[0] = in.val[0];
        out.val[1] = vorrq_u8(vshlq_n_u8(in.val[2], 4), vshrq_n_u8(in.val[1], 4));
        out.val[2] = vorrq_u8(vshlq_n_u8(in.val[1], 4), vshrq_n_u8(in.val[2], 4));
        vst3q_u8(pBE, out);
        pCSI2 += 48;
        pBE += 48;
    }
}

试试godbolt。

使用gcc编译器，生成的汇编代码看起来符合预期。（有一个mov指令可以通过更好的寄存器分配来消除，但这只是次要问题。）

不幸的是，clang编译器似乎存在一个奇怪的优化错误，将4位右移操作分解为3位和1位的右移操作。我已经提交了一个错误报告。

原则上，我们可以通过使用sli（Shift Left and Insert）来稍微改进，将OR操作与其中一个右移操作合并：

out.val[1] = vsliq_n_u8(vshrq_n_u8(in.val[1], 4), in.val[2], 4);
out.val[2] = vsliq_n_u8(vshrq_n_u8(in.val[2], 4), in.val[1], 4);

但由于它覆盖了源操作数，我们需要付出额外的几个mov指令作为代价。https://godbolt.org/z/TbzEEd1Pn。clang更巧妙地分配寄存器，只需要一个mov指令。

另一个选项，可能会稍微快一些，是使用sra指令，即右移并累加，它执行的是加法而不是插入操作。由于相关位已经为零，这具有相同的效果。奇怪的是没有sla指令。

out.val[1] = vsraq_n_u8(vshlq_n_u8(in.val[2], 4), in.val[1], 4);
out.val[2] = vsraq_n_u8(vshlq_n_u8(in.val[1], 4), in.val[2], 4);

- Nate Eldredge

Answer 2

我建议你从一个图表开始。

关于NEON我无法说，因此我会描述如何编写符合你要求的AVX2代码（然而，你应该使用目标指令集来实现它；最好别用转换器，如果你的目标是编写新代码）。x64内置函数有很好的文档；这里是我使用的一个示例。

AVX2寄存器有256位，即32字节。也就是说，你的24位数据可以容纳10个单位。做一个图表（对我来说最好是在纸上）：画出从内存中读取时256位寄存器包含哪些位。然后画出你希望在变换后在其中得到哪些位。用线连接它们。识别具有相同相对位置的位块。

然后编写代码来分离相关的位块（_mm256_and_si256），将它们移动（_mm256_slli_si256，可能还有_mm256_bslli_epi128或其他方式），然后组合它们（_mm256_or_si256）。AVX2对于移位操作特别古怪，所以我确信NEON代码会更容易编写。

您的主循环应该包含读取、处理和写入3个寄存器或768位。如果您为第一个寄存器绘制一个图表，可能可以类似地实现其他两个。当然，对于循环剩余部分（最后几个数据元素），您需要对它们使用常规的C代码进行特殊处理。

- anatolyg