AVX512 - 如何将所有设置位向右移动?

AVX512 - How to move all set bits to the right?

如何将掩码寄存器的所有设置位右移? (到底部,最不重要的位置)。

例如:

__mmask16 mask = _mm512_cmpeq_epi32_mask(vload, vlimit); // mask = 1101110111011101

如果我们将所有设置位向右移动,我们将得到:1101110111011101 -> 0000111111111111

我怎样才能有效地做到这一点?

下面您可以看到我是如何尝试获得相同结果的,但效率很低:

__mmask16 mask = 56797;
// mask: 1101110111011101
__m512i vbrdcast = _mm512_maskz_broadcastd_epi32(mask, _mm_set1_epi32(~0));
// vbrdcast: -1 0 -1 -1 -1 0 -1 -1 -1 0 -1 -1 -1 0 -1 -1
__m512i vcompress = _mm512_maskz_compress_epi32(mask, vbrdcast);
// vcompress:-1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 -1 0 0 0 0 
__mmask16 right_packed_mask =   _mm512_movepi32_mask(vcompress);   
// right_packed_mask: 0000111111111111                         

最好的方法是什么?

BMI2 pextv[p]compressd/q/ps/pd.
的按位等效标量 在掩码值上使用它,将它们左打包到值的底部。

  mask = _pext_u32(-1U, mask);    // or _pext_u64(-1ULL, mask64)  for __mmask64
// costs 3 asm instructions (kmov + pext + kmov) if you need to use the result as a mask
// not including putting -1 in a register.

__mmask16(在 GCC 中又名 uint16_t)和 uint32_t 之间的隐式转换有效。
如果您愿意,可以使用 _cvtu32_mask16_cvtu32_mask16 使 KMOVW 显式化。

有关以这种方式使用 pext/pdep 的更多信息,请参阅

所有当前使用 AVX-512 的 CPUs 也具有快速 BMI2 pext(包括 Xeon Phi),与 popcnt 相同的性能。 AMD 在 Zen 3 之前一直很慢 pext,但是 if/when AMD 曾经推出过 AVX-512 CPU 它应该很快 pext/pdep.

对于没有 AVX512 的早期 AMD,您可能需要 (1ULL << __builtin_popcount(mask)) - 1,但如果所有位都已设置,请注意溢出。 1ULL << 64 是未定义的行为,在为 x86-64 编译时可能会产生 1 而不是 0


如果您打算使用 vpcompressd,请注意源向量可以是全一 _mm512_set1_epi32(-1); compress 不关心掩码为零的元素,它们不需要已经为零。

(打包 -1 并不重要;一旦您使用布尔值,true 之间就没有区别了来自原始位掩码的常量 true 只是坐在那里,您生成的更便宜,不依赖于输入掩码。同样的推理适用于 pext,为什么您可以使用 -1U 作为源数据而不是 pdep。即 -1 或设置位没有身份;它与任何其他 -1 或设置位相同)。

所以让我们尝试两种方法,看看 good/bad asm 是怎样的。

inline
__mmask16 leftpack_k(__mmask16 mask){
    return _pdep_u32(-1U, mask);
}

inline
__mmask16 leftpack_comp(__mmask16 mask) {
    __m512i v = _mm512_maskz_compress_epi32(mask, _mm512_set1_epi32(-1));
    return _mm512_movepi32_mask(v);
}

查看这些的独立版本没有用,因为 __mmask16unsigned short 的类型定义,因此在整数寄存器中是 passed/returned,而不是 k 注册。这使得 pext 版本看起来非常好,当然,但我们想看看它如何内联到我们生成和使用带有 AVX-512 内在函数的掩码的情况。

// not a useful function, just something that compiles to asm in an obvious way
void use_leftpack_compress(void *dst, __m512i v){
    __mmask16 m = _mm512_test_epi32_mask(v,v);
    m = leftpack_comp(m);
    _mm512_mask_storeu_epi32(dst, m, v);
}

注释掉 m = pack(m),这只是生成并使用掩码的简单 2 条指令。

use_mask_nocompress(void*, long long __vector(8)):
        vptestmd        k1, zmm0, zmm0
        vmovdqu32       ZMMWORD PTR [rdi]{k1}, zmm0
        ret

所以任何额外的指令都是由于左包装(压缩)掩码造成的。 GCC 和 clang 彼此制作相同的 asm,唯一不同的是 clang 避免 kmovw 而总是 kmovdGodbolt

# GCC10.3 -O3 -march=skylake-avx512
use_leftpack_k(void*, long long __vector(8)):
        vptestmd        k0, zmm0, zmm0
        mov     eax, -1                        # could be hoisted out of a loop
        kmovd   edx, k0
        pdep    eax, eax, edx
        kmovw   k1, eax
        vmovdqu32       ZMMWORD PTR [rdi]{k1}, zmm0
        ret
use_leftpack_compress(void*, long long __vector(8)):
        vptestmd        k1, zmm0, zmm0
        vpternlogd      zmm2, zmm2, zmm2, 0xFF   # set1(-1)  could be hoisted out of a loop
        vpcompressd     zmm1{k1}{z}, zmm2
        vpmovd2m        k1, zmm1
        vmovdqu32       ZMMWORD PTR [rdi]{k1}, zmm0
        ret

所以不可吊装的部分是

  • kmov r,k(端口 0)/pext(端口 1)/kmov k,r(端口 5)= 3 微指令,每个执行端口一个。 (包括端口 1,它的矢量 ALU 在 512 位 uops 运行时关闭)。 kmov/kmov 往返有 4 cycle latency on SKXpext 是 3 个周期延迟,总共有 7 个周期延迟

  • vpcompressd zmm{k}{z}, z (2 p5) / vpmovd2m (port 0) = 3 uops,两个用于端口 5。vpmovd2m3 cycle latency on SKX / ICL, and vpcompressd-zeroing-into-zmm has 6 cycle from the k input to the zmm output (SKX 和集成电路)。因此,总共有 9 个周期的延迟,而且 uops 的端口分布更差。

此外,可提升的部分通常更差(vpternlogdmov r32, imm32 更长并且竞争更少的端口),除非你的函数已经需要一个全一向量来处理某些事情而不是全部-ones 注册。

结论:BMI2pext方式一点都不差,好几个。 (除非周围的代码 严重 在端口 1 微指令上出现瓶颈,如果使用 512 位向量,这是不太可能的,因为在那种情况下它只能是 运行 标量整数微指令,例如 3-循环 LEA、IMUL、LZCNT,当然还有简单的 1 循环整数,例如 add/sub/and/or).