ARM Neon：将非零字节的第 n 个位置存储在 8 字节向量通道中

Question

我想转换一个 Neon 64 位向量通道以获得非零（又名 0xFF）8 位值的第 n 个位置，然后用向量的其余部分填充零。以下是一些示例：

    0  1  2  3  4  5  6  7

d0: 00 FF 00 FF 00 00 00 FF
d1: 1  3  7  0  0  0  0  0

d0: 00 FF FF FF 00 00 FF 00
d1: 1  2  3  6  0  0  0  0

d0: FF FF FF FF FF FF FF FF
d1: 0  1  2  3  4  5  6  7

d0: FF 00 00 00 00 00 00 00
d1: 0  0  0  0  0  0  0  0

d0: 00 00 00 00 00 00 00 00
d1: 0  0  0  0  0  0  0  0

我觉得它可能是一个或两个位移 Neon 指令和另一个“好”向量。我该怎么做？

Answer 1

事实证明这并不简单。

简单有效的方法从简单地获取索引开始（只需使用位掩码加载0 1 2 3 4 5 6 7 和vand 的静态向量）。然而，为了在输出向量的一端收集它们——在与它们所代表的输入通道不同的通道中——你需要一些任意的置换操作。只有一条指令能够任意置换向量，vtbl（或vtbx，本质上是一回事）。但是，vtbl 采用目标顺序的源索引向量，结果与您尝试生成的内容完全相同。因此，为了产生最终结果，您需要使用最终结果，因此天真的有效解决方案是不可能的； QED。

根本问题是，您实际上所做的是排序一个向量，这本质上不是并行 SIMD 操作。 NEON 是专为媒体处理而设计的并行 SIMD 指令集，实际上并不适用于更一般的矢量处理的任何数据相关/水平/分散收集操作。

为了证明这一点，我确实设法在纯 NEON 中做到了这一点，根本没有任何标量代码，这可怕；我能想到的最好的“一个或两个位移位 NEON 指令”是一些基于条件选择的旋转位掩码累积技巧。如果不清楚，我建议在调试器或模拟器中逐步完成它的工作（example）：

// d0 contains input vector
vmov.u8 d1, #0
vmov.u8 d2, #0
vmvn.u8 d3, #0
vdup.u8 d4, d0[0]
vext.u8 d5, d2, d3, #7
vbit.u8 d3, d5, d4
vsub.u8 d1, d1, d3
vdup.u8 d4, d0[1]
vext.u8 d5, d2, d3, #7
vbit.u8 d3, d5, d4
vsub.u8 d1, d1, d3
vdup.u8 d4, d0[2]
vext.u8 d5, d2, d3, #7
vbit.u8 d3, d5, d4
vsub.u8 d1, d1, d3
vdup.u8 d4, d0[3]
vext.u8 d5, d2, d3, #7
vbit.u8 d3, d5, d4
vsub.u8 d1, d1, d3
vdup.u8 d4, d0[4]
vext.u8 d5, d2, d3, #7
vbit.u8 d3, d5, d4
vsub.u8 d1, d1, d3
vdup.u8 d4, d0[5]
vext.u8 d5, d2, d3, #7
vbit.u8 d3, d5, d4
vsub.u8 d1, d1, d3
vdup.u8 d4, d0[6]
vext.u8 d5, d2, d3, #7
vbit.u8 d3, d5, d4
vsub.u8 d1, d1, d3
vdup.u8 d4, d0[7]
vext.u8 d5, d2, d3, #7
vbit.u8 d3, d5, d4
vbic.u8 d1, d1, d3
// d1 contains output vector

作弊和使用循环（这需要以相反的方向旋转d0，以便我们可以通过d0[0] 访问每个原始车道）使其更小，但实际上并没有那么糟糕：

vmov.u8 d1, #0
vmov.u8 d2, #0
vmvn.u8 d3, #0
mov r0, #8
1:
vdup.u8 d4, d0[0]
vext.u8 d5, d2, d3, #7
vbit.u8 d3, d5, d4
subs r0, r0, #1
vext.u8 d0, d0, d0, #1
vsub.u8 d1, d1, d3
bne 1b
vbic.u8 d1, d1, d3

理想情况下，如果有可能对算法的其他部分进行返工以避免需要向量的非常量排列，则改为这样做。

Answer 2

我计划通过使用可变班次的分而治之的技术来实现这一点。 在每一步中，输入被视为“高”和“低”部分，其中“高”部分需要先右移 0 或 1 个字节，然后再右移 0-2 个字节，然后0-4 字节。

该解决方案允许在所有指令中使用“q”变体，从而允许并行执行两个独立的压缩。

///  F F 0 F F 0 0 F   mask
///  7 6 - 4 3 - - 0   mask & idx
///  7 6 - 4 - 3 - 0   pairwise shift.16 right by 1-count_left bytes
///    2   1   1   1   count_left + count_right = vpaddl[q]_s8 -> s16
///  - 7 6 4 - - 3 0   pairwise shift.32 right by 2-count_left bytes
///        3       2   count_left + count_right = vpaddl[q]_s16 -> s32
///  - - - 7 6 4 3 0   shift.64 right by 4 - count_left bytes
///                5   number of elements to write = vpaddl[q]_s32 -> s64

第一步可以在没有实际班次的情况下完成

int8x8_t step1(int8x8_t mask) {
    auto data = mask & vcreate_u8(0x0706050403020100ull);
    auto shifted = vrev16_u8(data);
    return vbsl_u8(vtst_s16(mask, vdup_n_s16(1)), data, shifted);
}

下一步需要隔离每个uint32_t lane的top 16和bottom 16 bits，将top部分移位-16、-8或0位，然后与隔离的bottom bits组合。

int8x8_t step2(int8x8_t mask, int16x4_t counts) {
    auto top = vshr_n_u32(top, 16);
    auto cnt = vbic_s32(counts, vcreate_u32(0xffff0000ffff0000ull));
    auto bot = vbic_u32(mask, vcreate_u32(0xffff0000ffff0000ull));
    top = vshl_u32(top, cnt);
    return vorr_u8(top, bot); 
}

第三步需要移位 64 位元素。

int8x8_t step3(int8x8_t mask, int32x4_t counts) {
    auto top = vshr_n_u64(top, 32);
    auto cnt = vbic_s64(counts, vcreate_s32(0xffffffff00000000ull));
    auto bot = vbic_u64(mask, vcreate_u32(0xffffffff00000000ull));
    top = vshl_u64(top, cnt);
    return vorr_u8(top, bot); 
}

完整解决方案：

auto cnt8 = vcnt_s8(mask);
mask = step1(mask);
auto counts16 = vpaddl_s8(cnt8, cnt8);
mask = step2(mask, counts16);
auto counts32 = vpaddl_s16(counts16, counts16);
mask = step3(mask, counts32);
auto counts64 = vpaddl_s32(counts32, counts32);

最终的“counts64”实际上应该提前计算，因为计数需要传输到通用寄存器，因为它用于以字节递增流式写入指针：

vst1_u8(ptr, mask); ptr += count64 >> 3;

渐近更好的版本实际上会尝试获得 64（+64= 字节的掩码，将这些字节压缩为位模式（如英特尔的 movmaskb 中），然后使用 8 次迭代到 find leading zeros + clear leading one 或 find + toggle least significant set bit。

这可以通过每次迭代 5 条指令来实现：

// finds + toggles least significant bit
auto prev = mask;
mask &= mask - 1u;
auto index0 = vclz[q]_u8(prev ^ mask);
// assuming reversed bit ordering

这 8 个索引[0..7] 需要转置为 8x8 矩阵，然后按顺序写入；每 64 + 64 字节的指令总数将接近 64 条指令，或每个输出字节 0.5 条指令。

Answer 3

您可以通过对向量进行排序来做到这一点。对于这样的操作，这是一个比您预期的更复杂的操作，但我还没有想出更好的方法。

给定d0 中的00/ff 字节列表和d1 中的常量0, 1, 2, ..., 7，您可以使用vorn 创建一个活动列的可排序列表。

vorn.u8 d0, d1, d0

现在d0 的所有不需要的车道都已替换为0xff，其余的已替换为其车道索引。从那里您可以对该列表进行排序，以将所有不需要的车道聚集在最后。

为此，您必须将列表扩展到 16 个字节：

vmov.u8 d1, #255

然后将它们拆分为奇/偶向量：

vuzp.u8 d0, d1

排序操作由这些向量之间的vmin/vmax 组成，然后是交错操作，然后是另一个vmin/vmax 对在不同对之间进行排序，因此值可以冒泡到合适的位置。像这样：

vmin.u8 d2, d0, d1
vmax.u8 d3, d0, d1
vsri.u64 d2, d2, #8   ; stagger the even lanes (discards d0[0])
vmax.u8 d4, d2, d3    ; dst reg would be d0, but we want d0[0] back...
vmin.u8 d1, d2, d3
vsli.u64 d0, d4, #8   ; revert the stagger, and restore d0[0]

这实现了整个网络的两个阶段，整个区块必须重复四次（八个阶段）才能使d0[7] 中的某些东西在极端情况下一直冒泡到d0[0]最后一个字节是唯一的非零输入，或者如果第一个字节是唯一的零输入，d0[0] 将到达 d0[7]。

完成排序后，将结果重新拼接在一起：

vzip.u8 d0, d1

而且因为您希望剩余车道中的零：

vmov.u8 d1, #0
vmax.s8 d0, d1

现在d0 应该包含结果。

如果您查看维基百科的sorting network 页面，您会发现八车道的理论最小深度只有六个阶段（六对vmin/vmax），因此可能会找到一组置换（替换我的vsli 和vsri 操作）实现六阶段排序网络，而不是我实现的八阶段插入/选择/vbubble 排序。如果确实存在，并且与 NEON 的置换操作兼容，那么它当然值得寻找，但我没有时间看。

还请注意，排序总共需要 16 个字节，这超出了您的需要，如果您使用 q 寄存器，您可以让它在 32 个字节上工作......所以这距离最大吞吐量还有很长的路要走.

哦，即使在这种配置中，我认为您也不需要排序的最后阶段。留给读者的练习。