提取 AVX2 16x16 位矩阵的边缘

TLW*_*TLW 5 c bit-manipulation intrinsics avx2

有没有一种相对便宜的方法将存储在 a 中的 16x16 位矩阵的四个边(第 0 行和第 15 行,以及第 0 行和第 15 列)提取到 a__m256i的四个 16b 通道中__m256i?我不关心输出到哪个通道,或者寄存器的其余部分是否有垃圾。轻度偏好所有这些都处于下半部分,但只是轻度。

提取“顶部”和“底部”很容易 - 只需向量的第一个和最后 16b 个元素即可完成 - 但侧面是另一回事。您需要每个 16b 元素的第一位和最后一位,这会变得很复杂。

您可以使用完整的位转置来完成此操作,如下所示:

// Full bit-transpose of input viewed as a 16x16 bitmatrix.
extern __m256i transpose(__m256i m);

__m256i get_edges(__m256i m) {
    __m256i t = transpose(m);
    // We only care about first and last u16 of each
    // m = [abcdefghijklmnop]
    // t = [ABCDEFGHIJKLMNOP]
    m = _mm256_permutevar8x32_epi32(m, _mm256_set_epi32(0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x7, 0x0));
    // m = [............a..p]
    t = _mm256_permutevar8x32_epi32(t, _mm256_set_epi32(0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x7, 0x0));
    // m = [............A..P]

    __m256i r = _mm256_unpacklo_epi16(t, m);
    // r = [........aA....pP]
    return r; // output in low and high dwords of low half
}
Run Code Online (Sandbox Code Playgroud)

...但这只是将一个令人惊讶的恼人问题减少为另一个令人惊讶的恼人问题 - 我不知道如何便宜地对一个令人惊讶的完整位转置__m256i.

同上,可能有什么_mm256_movemask_epi8类似的东西可以做到这一点,但我什么也没注意到。

有更好的方法吗?

Pet*_*des 7

对于快速 BMI2pextvpmovmskb (Haswell 或 Zen 3 及更高版本),如果您从+ shift +开始vpmovmskb获取边缘位(与垃圾位交错,因为我们想要每 16 个,但我们得到每 8 个),这是一种选择。

前端有 9 个 uops,其中 6 个需要 Intel Skylake 系列上的端口 5。(不计算整数常量设置,假设您会在循环中执行此操作。如果没有,这也会对此产生影响。)

__m128i edges_zen3_intel(__m256i v)
{
    __m128i vtop_bottom = _mm256_castsi256_si128( 
          _mm256_permute4x64_epi64(v, _MM_SHUFFLE(0,0, 3, 0)) );
    // vpermq: 3 uops on Zen1, 2 on Zen2&3, 1 on Zen4 and Intel.

   // side bits interleaved with garbage
   // without AVX-512 we can only extract a bit per byte, dword, or qword
   unsigned left = _mm256_movemask_epi8(v);   // high bit of each element
   unsigned right = _mm256_movemask_epi8( _mm256_slli_epi16(v, 15) );  // low<<15

//   left = _pext_u32(left, 0xAAAAAAAAul);  // take every other bit starting with #1
//   right = _pext_u32(right, 0xAAAAAAAAul);
    // then combine or do whatever

   uint64_t lr = ((uint64_t)left << 32) | right;
   lr = _pext_u64(lr, 0xAAAAAAAAAAAAAAAAull);

   //__m128i vsides = _mm_cvtsi32_si128(lr);
   __m128i vtblr = _mm_insert_epi32(vtop_bottom, lr, 1);  // into an unused space
   // u16 elems: [ top | x | x | x | left | right | x | bottom ]

   return vtblr;
}
Run Code Online (Sandbox Code Playgroud)

对于 Intel CPU(和 Zen 4),这会编译为 10 uops,包括将所有内容重新放入一个 SIMD 向量中。可以将其movabs吊出环路。SHL/OR 不竞争 SIMD 执行端口吞吐量(能够在 Intel 上的端口 6 上运行),但会竞争前端。 神箭

# Haswell/Sklake uop counts
edges_zen3_intel(long long __vector(4)):
        vpsllw  ymm2, ymm0, 15             # p0 (or p01 on Skylake)
        vpmovmskb       eax, ymm0          # p0
        vpermq  ymm1, ymm0, 12             # p5
        vpmovmskb       edx, ymm2          # p0
        sal     rax, 32                    # p06
        or      rax, rdx                   # p0156
        movabs  rdx, -6148914691236517206  # p0156 (and can be hoisted out of loops)
        pext    rax, rax, rdx              # p1
        vpinsrd xmm0, xmm1, eax, 1         # 2 p5.  On Intel, both uops compete with shuffles
        ret
Run Code Online (Sandbox Code Playgroud)

vpmovmskb作为一种变体,如果我们可以左移奇数字节而不是偶数字节,我们也许可以将左边缘和右边缘放在一起?可能不是,_mm256_maddubs_epi16因为_mm256_set1_epi16(0x0180)不能这样做,它添加了水平对,并且左移 7 (0x80 = 1<<7) 是不够的,我们需要 8 才能将顶部位返回到顶部。

或者,如果我们vpsllw+ vpacksswb,则使用正确的掩码对位进行分组,例如0x00ff00ff。但这越来越接近我的非 pext 想法,即使我们确实有快速的,也许更好pext

没有快速 BMI2 pext- 饱和打包向量以减少到 8 位元素

即使pext速度很快,这也可能更快。

带符号饱和打包始终保留符号位,因此您可以将 16 位缩小到 8 位,而不会丢失想要保留的信息。我们想要对每个字的高位和低位(16 位元素)执行此操作,因此与原始 和 的 2:1 包v<<15是完美的。

除了AVX2vpacksswb ymm是两个独立的通道内打包操作这一事实之外,因此我们最终会得到交错的 8 元素块。我们可以在打包后立即解决这个问题,但它在 Zen 1 到 Zen 3 上有多个微指令,我们可以在将结果放回向量寄存器vpermq后对字节进行洗牌。movemask(同样vpshufb可以在高和低元素周围移动。)

// avoiding PEXT because it's slow on Zen 2 and Zen 1 (and Excavator)
// This might be good on Intel and Zen 3, maybe comparable to using PEXT
__m128i edges_no_pext(__m256i v)
{
    __m128i vhi = _mm256_extract_si128(v, 1);  // contains top, as vhi.u16[7]
    __m128i vlo = _mm256_castsi256_si128(v);   // contains bottom, as vlo.u16[0], contiguous if concatenated the right way
    __m128i bottom_top = _mm_alignr_epi8(vhi, vlo, 12);  // rotate bottom :top down to the 2nd dword [ x | x | bottom:top | x]

   // vpermq ymm, ymm, imm would also work to get them into the low 128
   // but that's 3 uops on Zen1, 2 on Zen2&3, 1 on Zen4 and Intel.
   // and would need a slightly more expensive vpinsrd instead of vmovd+vpblendd

   // On Intel CPUs (and Zen4) vpermq is better; we pshufb later so we can get the bytes where we want them.
   // A compromise is to use vextracti128+vpblendd here, vpinsrd later
   //   __m128i bottom_top = _mm_blend_epi32(vhi, vlo, 0b0001);
                    // [ hi | x | x | x   |   x | x | x | lo ]

    __m256i vright = _mm256_slli_epi16(v, 15);
    __m256i vpacked = _mm256_packs_epi16(v, vright);   // pack now, shuffle bytes later.
    unsigned bits = _mm256_extract_epi8(vpacked);    // [ left_hi | right_hi | left_lo | right_lo ]

    __m128i vsides = _mm_cvtsi32_si128(bits);
    __m128i vtblr = _mm_blend_epi32(top_bottom, vsides, 0b0001);  // vpinsrd xmm0, eax, 0 but the merge can run on more ports

    __m128i shuffle = _mm_set_epi8(-1,-1,-1,-1, -1,-1,-1,-1,
                                   7,6,5,4, 3,1, 2,0);
     // swap middle 2 bytes of the low dword, fixing up the in-lane pack
     vtblr = _mm_shuffle_epi8(vtblr, shuffle);
     return vtblr;   // low 4 u16 elements are (MSB) top | bottom | left | right  (LSB)
}
Run Code Online (Sandbox Code Playgroud)

这编译得相当好(参见前面的 Godbolt 链接),尽管 GCC4.9 和更高版本(和 clang)将我的vmovd+悲观vpblendd化为vpinsrd,即使使用-march=haswell或 Skylake,端口 5 为 2 uops(https://uops.info/),而大多数情况下该函数中的其他指令也是仅在端口 5 上运行的 shuffle。(这对于 Intel CPU 来说更加需要 shuffle。)

使用vpblendd而不是vpalignr会使英特尔的情况变得不那么糟糕,比如__m128i bottom_top = _mm_blend_epi32(vhi, vlo, 0b0001);,即使在 Zen 1 上也能达到与vpermq下面版本中相同的情况,即使在 Zen 1 上也有 2 uop。但这只是在 Zen 1 上节省了 1 uop,在其他地方都相同或更差。

# GCC12 -O3 -march=haswell
# uop counts for Skylake
edges_no_pext:
        vextracti128    xmm1, ymm0, 0x1        # p5
        vpsllw  ymm2, ymm0, 15                 # p01
        vpalignr        xmm1, xmm1, xmm0, 12   # p5
        vpacksswb       ymm0, ymm0, ymm2       # p5
        vpmovmskb       eax, ymm0              # p0
        vpinsrd xmm0, xmm1, eax, 0             # 2 p5
        vpshufb xmm0, xmm0, XMMWORD PTR .LC0[rip]  # p5
        ret
Run Code Online (Sandbox Code Playgroud)

因此,Intel 上的端口 5 为 6 uops,每 6 个周期 1 个吞吐量瓶颈。与 PEXT 版本相比,3 个 uops 需要端口 0,3 个 uops 需要端口 5。但这对于前端来说总共只有 8 个 uops,而该版本需要 9 个 uops pextvpermq假设 GCC 不会浪费后内联,该版本又比 Intel 节省了 1 个vmovdqa

如果您不关心将输出向量的高 8 字节清零,则可以加载 shuffle 常量,vmovq并且只是 8 字节而不是 16(如果您将高 0 字节全部为零)。但编译器可能不会发现这种优化。

由于编译器坚持vpinsrd在快速的 CPU vpermq(Intel 和 Zen4)上对 , 进行悲观化,我们不妨使用它:

如果您只想拥有一个非 GFNI AVX2 版本,这可能是一个很好的权衡

vpermqZen 1 上的 3 uop 并不比使用 2 条指令模拟我们需要的东西差多少,而且在 Intel CPU 上更差可能是关于 Zen 2 和 Zen 3 的收支平衡,后端端口使用的模数差异。

// for fast vpermq, especially if compilers are going to pessimize vmovd(p5)+vpblendd (p015) into vpinsrd (2p5).
// good on Intel and Zen 4, maybe also Zen 3 and not bad on Zen 2.
__m128i edges_no_pext_fast_vpermq(__m256i v)
{
   __m128i vtop_bottom = _mm256_castsi256_si128( 
            _mm256_permute4x64_epi64(v, _MM_SHUFFLE(0,0, 3, 0)) );
    // 3 uops on Zen1, 2 on Zen2&3, 1 on Zen4 and Intel.

    __m256i vright = _mm256_slli_epi16(v, 15);
    __m256i vpacked = _mm256_packs_epi16(v, vright);   // pack now, shuffle bytes later.
    unsigned bits = _mm256_movemask_epi8(vpacked);    // [ left_hi | right_hi | left_lo | right_lo ]

    __m128i vtblr = _mm_insert_epi32(vtop_bottom, bits, 1);  // into an unused space
    // u16 elems: [ top | x | x | x | lh:rh | ll:rl | x | bottom ]
    __m128i shuffle = _mm_set_epi8(-1,-1,-1,-1, -1,-1,-1,-1,
                                   15,14, 1,0, 7,5, 6,4);
     vtblr = _mm_shuffle_epi8(vtblr, shuffle);
     return vtblr;   // low 4 u16 elements are (MSB) top | bottom | left | right  (LSB)
}
Run Code Online (Sandbox Code Playgroud)
# GCC12.2 -O3 -march=haswell     clang is similar but has vzeroupper despite the caller passing a YMM, but no wasted vmovdqa
edges_no_pext_fast_vpermq(long long __vector(4)):
        vmovdqa ymm1, ymm0
        vpermq  ymm0, ymm0, 12
        vpsllw  ymm2, ymm1, 15
        vpacksswb       ymm1, ymm1, ymm2
        vpmovmskb       eax, ymm1
        vpinsrd xmm0, xmm0, eax, 1
        vpshufb xmm0, xmm0, XMMWORD PTR .LC1[rip]
        ret
Run Code Online (Sandbox Code Playgroud)

在 Intel Haswell/Skylake 上,端口 5 为 5 uops,加上移位 (p01) 和 vpmovmskb (p0)。总共 7 个微指令。vmovdqa(不计算应该通过内联消除的ret 或浪费。)

在 Ice Lake 及更高版本上,其中一个 uopvpinsrd可以在 p15 上运行,如果您在循环中执行此操作,则可以减轻该端口上的一个 uop 压力。 vpinsrd是 Alder Lake E 核上的单微操作。

Ice Lake(及更高版本)还可以vpshufb在 p1/p5 上运行,进一步降低端口 5 的压力,降至 7 个 uops 中的 3 个。端口 5 可以处理任何 shuffle,端口 1 可以处理一些但不是所有 shuffle 微指令。它可以连接到 512 位混洗单元的上半部分,为某些 256 位和更窄的混洗提供额外的吞吐量,例如 p0/p1 FMA 单元如何作为 p0 上的单个 512 位 FMA 单元工作。它不处理vpermqor vpacksswb; 这些仍然是仅在 Ice/Alder Lake 上的 p5。

所以这个版本对于当前一代和未来的 Intel CPU 来说是相当合理的。Alder Lake E 核vpermq ymm以 2 微指令运行,具有 7 个周期延迟。但是,如果他们可以通过更有限的无序调度(大 ROB,但每个端口的队列不那么长)来隐藏延迟,那么vpinsrd作为单个 uop 运行有助于弥补前端吞吐量。

256 位指令(例如vpsllw ymm和 )vpacksswb ymm在 Alder Lake E 核上也各为 2 uop,但vpmovmskb eax,ymm为 1 uop(但延迟可能较高)。因此,即使我们想制作一个针对 Zen1 / Alder E 优化的版本,我们也可能无法通过在vextracti128;之后使用更多 128 位指令来节省它们的总微指令数。我们仍然需要对输入向量的两半进行处理。


我曾研究过以正确的顺序进行打包,以使vpmovmskb xmm每个 16 位组都按正确的顺序排列,但要分开。我曾考虑过用 来实现这一点vperm2i128,但是在 Zen 1 上速度相当慢。

//    __m256i vcombined = _mm256_permute2x128_si256(v, vright, 0x10);  // or something?  Takes two shuffles to get them ordered the right way for pack
Run Code Online (Sandbox Code Playgroud)

Zen 1 的速度非常快vextracti128- 对于任何端口来说都是单微指令,而 128 位向量运算则为 1 微指令与 2 微指令__m256i。我们已经在进行提取以将顶部和底部放在一起。

但它仍然会导致更多的标量工作,特别是如果您希望将结果合并到向量中。2xvpinsrw或之前的额外 SHL/ORvmovd更糟。

#if 0
// Zen 1 has slow vperm2i128, but I didn't end up using it even if it's fast
    __m128i hi = _mm256_extract_si128(v, 1); // vextracti128  - very cheap on Zen1
    __m128i lo = _mm256_castsi256_si128(v);  // no cost
    __m128i vleft = _mm_packs_epi16(lo, hi);  // vpacksswb signed saturation, high bit of each word becomes high bit of byte

    // then shift 2 halves separately and pack again?
#endif
Run Code Online (Sandbox Code Playgroud)

设置向量打包vpmovmskb可能是最好的选择;在考虑这一点之前,我正在考虑vpmovmskb直接在输入上使用并使用标量位黑客来获取奇数位或偶数位:

但这些需要更多的操作,因此速度会更慢,除非您特别遇到 SIMD ALU 的瓶颈,而不是整体前端吞吐量(或 SIMD 和标量 ALU 共享端口的 Intel 上的执行端口吞吐量)。


AVX-512 和/或 GFNI

这里有两个有趣的策略:

  • vpmovw2m和/或vptestmw作为mb更方便的vpmovmskb. 仅需要 AVX-512BW (Skylake-avx512)
  • 将 8 位打包到每个 qword 的底部,然后随机播放。可能仅适用于 GFNI + AVX512VBMI,例如 Ice Lake / Zen4 及更高版本。也许只是 GFNI + AVX2,就像在瘫痪的 Alder Lake 中一样(没有 AVX-512)。

将位提取到掩码:

使用vptestmbset1_epi8(0x8001)我们可以将我们想要的所有位放入一个掩码中,但随后我们需要解交错,可能使用标量pext(这在所有 AVX-512 CPU 上都很快,除了 Knight's Landing,但它没有 AVX-512BW )。

因此,最好提取两个掩码并连接起来。除了等一下,我没有看到将 32 位掩码放入向量寄存器的好方法(无需将其扩展为 0 / -1 元素的向量)。对于 8 位和 16 位掩码,有掩码到向量的广播,例如vpbroadcastmw2d x/y/zmm, k. 它们不支持屏蔽,因此您无法将屏蔽合并到另一个寄存器中。这在 Zen 4 上是单微指令,但在 Intel 上它需要 2 微指令,与kmov eax, k/vpbroadcastd x/y/zmm, eax相同,您应该这样做,这样您就可以将掩码合并到具有顶部和底部边缘的向量中。

  vpmovw2m k1, ymm0                        # left = 16 mask bits from high bits of 16 elements
  vptestmw k2, ymm0, set1_epi16(0x0001)    # right.   pseudocode constant
  kunpckwd k1, k1, k2                      # left:right
     # there's no  vpbroadcastmd2d  only byte/word mask to dword or qword element!
  
    mov    ecx, 0b0010
    kmovb  k7, ecx            # hoist this constant setup out of loops.  If not looping, maybe do something else, like bcast to another register and vpblendd.

  kmovd    eax, k1
  vpbroadcastd xmm0{k7}, eax  # put left:right into the 2nd element of XMM0
                              # leaving other unchanged (merge-masking)
Run Code Online (Sandbox Code Playgroud)

其中 xmm0 可以设置为vpermq在低 16 字节中具有 top:bottom ;所有配备 AVX-512 的 CPU 都具有高效的vpermq. 因此,在我手写的汇编中的 5 个 uop 之上,又多了 1 个 uop(用内在函数编写应该很简单,我只是不想在找到可用的汇编指令后采取额外的步骤来查找正确的内在函数。 )

将位打包到 qwords 中然后进行洗牌:GFNI 和可能的 AVX-512VBMIvpermb

(需要 AVX512VBMI 意味着 Ice Lake 或 Zen 4,因此vpermb将是单微指令。除非未来带有 E 核的英特尔 CPU 支持较慢的 AVX-512,但仍然vpermb ymm希望不会太糟糕。)

可能按左:右顺序打包(每个半字节),然后进行字节洗牌。如果我们可以在交替字节中执行left:rightright:left,则字节洗牌(如vpermbvpermt2b)应该能够设置为vprolw在每个 16 位字内旋转,以按正确的顺序对 8 个“左”位进行分组。

Moveing bits inside a qword : Harold's answer on bitpack ascii string into 7-bit bin blob using SIMD显示 _mm256_gf2p8affine_epi64_epi8将每个字节中的 1 位放在每个 qword 的顶部。(并打包剩余的 7 位字段,这是该答案的目标。)

如果这是可行的,那么与掩模和返回相比,它可能会更少的微指令和明显更好的延迟。

对于 Alder Lake(GFNI 但 AVX-512 已禁用,除非您设法避免英特尔削弱这个令人惊叹的 CPU 的努力),这可能仍然有用,因为它具有AVX+GFNI for_mm256_gf2p8affine_epi64_epi8 . vpshufb+vpermd可以代替vpermb. 但你不会有单词轮换;尽管如此,像ABAB这样的混洗字节可以让你使用简单的左移来获得你想要的窗口,然后再次混洗。