用 ARM NEON 实现 _mm_movemask_epi8 的几种方法

背景

上一篇文章中描述了一种使用 SIMD 指令进行并行查找的 B16 哈希表，我让它支持 ARM 时遇到了一些指令集兼容的问题，对这个问题小小地探索了一下。

SSE2 指令集提供了 _mm_movemask_epi8 (pmovmskb) 指令，作用是取所有 8 bit 操作数最高 bit，然后把它们存储到返回值里。对包含 16 个 8 bit 数的 128 bit 输入，取得高位 16 个 bit，存入 32 位的返回值里，并且将返回值的高位置 0。

但是在 ARM 的指令集中，没有这条指令，只能想其它办法替代。

已有实现

通过搜索，找到这个 StackOverflow 问题的回答，里面提到了四种实现方法，我整理了一下接口，分列如下：

// Yves Daoust 的回答 (7 votes): 与 _mm_movemask_epi8 略有不符，要求输入的每个 8 bits 全 0 或全 1
inline uint32_t vmovemask_u8_YvesDaoust(uint8x16_t a) {
    const uint8_t __attribute__ ((aligned (16))) _Powers[16]=
        { 1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128 };
    // Set the powers of 2 (do it once for all, if applicable)
    uint8x16_t Powers= vld1q_u8(_Powers);
    // Compute the mask from the input
    uint64x2_t Mask= vpaddlq_u32(vpaddlq_u16(vpaddlq_u8(vandq_u8(a, Powers))));
    // Get the resulting bytes
    uint32_t Output;
    vst1q_lane_u8((uint8_t*)&Output + 0, (uint8x16_t)Mask, 0);
    vst1q_lane_u8((uint8_t*)&Output + 1, (uint8x16_t)Mask, 8);
    return Output;
}
​
// David 对 Yves Daoust 回答最后三行进行了一些改进
inline uint32_t vmovemask_u8_David(uint8x16_t a) {
    const uint8_t __attribute__ ((aligned (16))) _Powers[16]=
        { 1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128 };
    // Set the powers of 2 (do it once for all, if applicable)
    uint8x16_t Powers= vld1q_u8(_Powers);
    // Compute the mask from the input
    uint64x2_t Mask= vpaddlq_u32(vpaddlq_u16(vpaddlq_u8(vandq_u8(a, Powers))));
    // Get the resulting bytes
    uint32_t Output = vgetq_lane_u64(Mask, 0) + (vgetq_lane_u64(Mask, 1) << 8);
    return Output;
}
​
// EasyasPi 的回答（4 votes): 标准实现了 _mm_movemask_epi8，被 simde 库采纳，link:
// https://github.com/simd-everywhere/simde/blob/master/simde/x86/sse2.h
inline uint32_t vmovemask_u8_EasyasPi(uint8x16_t input)
{
    // Example input (half scale):
    // 0x89 FF 1D C0 00 10 99 33
    // Shift out everything but the sign bits
    // 0x01 01 00 01 00 00 01 00
    uint16x8_t high_bits = vreinterpretq_u16_u8(vshrq_n_u8(input, 7));
    // Merge the even lanes together with vsra. The '??' bytes are garbage.
    // vsri could also be used, but it is slightly slower on aarch64.
    // 0x??03 ??02 ??00 ??01
    uint32x4_t paired16 = vreinterpretq_u32_u16(vsraq_n_u16(high_bits, high_bits, 7));
    // Repeat with wider lanes.
    // 0x??????0B ??????04
    uint64x2_t paired32 = vreinterpretq_u64_u32(vsraq_n_u32(paired16, paired16, 14));
    // 0x??????????????4B
    uint8x16_t paired64 = vreinterpretq_u8_u64(vsraq_n_u64(paired32, paired32, 28));
    // Extract the low 8 bits from each lane and join.
    // 0x4B
    return vgetq_lane_u8(paired64, 0) | ((uint32_t)vgetq_lane_u8(paired64, 8) << 8);
}
​
// inspirit 的回答 (1 vote): 标准实现了 _mm_movemask_epi8，但分了上下半边，指令很多
inline uint32_t vmovemask_u8_inspirit(uint8x16_t input)
{
    const int8_t __attribute__ ((aligned (16))) xr[8] = {-7,-6,-5,-4,-3,-2,-1,0};
    uint8x8_t mask_and = vdup_n_u8(0x80);
    int8x8_t mask_shift = vld1_s8(xr);
​
    uint8x8_t lo = vget_low_u8(input);
    uint8x8_t hi = vget_high_u8(input);
​
    lo = vand_u8(lo, mask_and);
    lo = vshl_u8(lo, mask_shift);
​
    hi = vand_u8(hi, mask_and);
    hi = vshl_u8(hi, mask_shift);
​
    lo = vpadd_u8(lo,lo);
    lo = vpadd_u8(lo,lo);
    lo = vpadd_u8(lo,lo);
​
    hi = vpadd_u8(hi,hi);
    hi = vpadd_u8(hi,hi);
    hi = vpadd_u8(hi,hi);
​
    return ((hi[0] << 8) | (lo[0] & 0xFF));
}

我的实现

看到上面这几个方法，我就在想，有没有可能找到一种更高效的实现，用更少的 ARM 指令实现这个功能？经过一段时间的思考，我想到了下面这种方法，我感觉这（可能）是指令数最少的一种实现了。

但这个方法和 YvesDaoust 的方法一样，假设每个 8 bits 都是全 0 或者全 1，这在处理向量比较指令（vceq*, vcgt* 等）结果时是可用的，但在其它场景下未必可用。

 // （可能是）指令数最少的实现，要求输入的每个 8 bits 全 0 或全 1
inline uint32_t vmovemask_u8_solrex(uint8x16_t a) {
    // 先取出相邻两个 uint8 的中间 2 bits，1 bit 属于高 uint8，1 bit 属于低 uint8
    uint16x8_t MASK =  vdupq_n_u16(0x180);
    uint16x8_t a_masked = vandq_u16(vreinterpretq_u16_u8(a), MASK);
    // 再将这 8 个 2 bits 按照不同的偏移进行 SHIFT，使得它们加一起能表示最终的 mask
    const int16_t __attribute__ ((aligned (16))) SHIFT_ARR[8]= {-7, -5, -3, -1, 1, 3, 5, 7};
    int16x8_t SHIFT = vld1q_s16(SHIFT_ARR);
    uint16x8_t a_shifted = vshlq_u16(a_masked, SHIFT);
    // 最后把这 8 个数字加起来
    return vaddvq_u16(a_shifted);
}

性能测试

我非常好奇新方法性能如何，所以我对以上几种方法都进行了 benchmark，然后发现结果跟我想的有点不一样：

方法	重复处理单变量	按序处理数组
vmovemask_u8_YvesDaoust()	536us	531us
vmovemask_u8_David()	189us	208us
vmovemask_u8_EasyasPi()	92us	340us
vmovemask_u8_inspirit()	286us	389us
vmovemask_u8_solrex()	137us	166us

分析

重复处理单变量场景下，对一个固定的 uint8x16_t 变量重复计算 movemask，然后把结果累加起来（避免被优化）。这时候，vmovemask_u8_EasyasPi()胜出。这可能是因为 EasyasPi 的方法只有数值计算，没有寄存器 load，而往往 load/store 指令的耗时是比较长的。

按序处理数组场景下，对一个 10 万个元素数组的每个元素计算 movemask，然后把结果累加起来（避免被优化）。这时候，vmovemask_u8_solrex() 胜出。这可能是因为新方法里的 load 操作与数组元素的 load 操作形成了一定的流水线效果，load 的开销被抵消后，指令数少的性能优势就体现出来了。

从与 _mm_movemask_epi8 接口的一致性来说，还是 EasyasPi 给的实现更合适，所以 simde 库在替换 x86 intrinsics 时也用了这个实现。但探索一下不同的实现，还是能让人对向量指令设计和选择更多一些理解。

最后说回哈希表里 SIMD 并行比较的实现，其实 Facebook F14 里的实现更高效，并没有受 movemask 的思路限制，感兴趣的同学可以自己钻研一下。