用 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()536us531us
vmovemask_u8_David()189us208us
vmovemask_u8_EasyasPi()92us340us
vmovemask_u8_inspirit()286us389us
vmovemask_u8_solrex()137us166us
表1:内联函数调用,重复 10 万次

分析

重复处理单变量场景下,对一个固定的 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 的思路限制,感兴趣的同学可以自己钻研一下。

发表评论

您的电子邮箱地址不会被公开。 必填项已用*标注