c-提取AVX2 16x16位矩阵的边缘



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

提取"顶部"one_answers"底部"很容易——这只是矢量的第一个和最后一个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
}

但这只是将一个令人惊讶的恼人问题简化为另一个令人震惊的恼人问题——我不知道如何廉价地对__m256i进行全比特转置。

同上,可能有_mm256_movemask_epi8风格的东西可以做到这一点,但没有什么能打动我

有更好的方法吗?

使用快速BMI2pext(Haswell或Zen 3及更高版本),如果您从vpmovmskb+shift+vpmovmskb开始获取边的位(与垃圾位交织,因为我们希望每16位一次,但我们每8位一次),这是一个选项。

9个uop用于前端,其中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;
}

这可以编译为英特尔CPU(和Zen 4)的10个uop,包括将所有内容都返回到一个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

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

或者,如果我们vpsllw+vpacksswb,那么使用正确的掩码对比特进行分组,如0x00ff00ff。但这越来越接近我的非文本想法,也许即使我们有快速的pext也会更好

在没有快速BMI2pext饱和的情况下,将矢量压缩为8位元素

即使pext很快,这也可能更快

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

除了AVX2vpacksswb ymm是两个独立的通道内分组操作之外,因此我们最终得到了交错的8元素块。我们可以在用vpermq打包后立即修复,但它在Zen 1到Zen 3上有多个uop,我们可以在将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)
}

这编译得很好(请参阅前面的Godbolt链接),尽管GCC4.9和更高版本(以及clang)将我的vmovd+vpblendd简化为vpinsrd,即使使用-march=haswell或Skylake,其中端口5为2个uops(https://uops.info/)当函数中的大多数其他指令也是仅在端口5上运行的混洗时。(对于英特尔CPU来说,这是一个更重的洗牌。)

使用vpblendd而不是vpalignr将使英特尔(如__m128i bottom_top = _mm_blend_epi32(vhi, vlo, 0b0001);)不那么糟糕,即使在Zen 1上也能达到与下面的vpermq版本相同的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

因此,对于Intel上的端口5,这是6个uops,每6个周期就有1个吞吐量瓶颈。而PEXT版本是需要端口0的3个uop、需要端口5的3个。但这对于前端来说总共只有8个uop,而对于pext版本来说只有9个。vpermq版本在Intel上又保存了一个,假设GCC在内联后不会浪费vmovdqa

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

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

如果你只打算有一个非GFNI AVX2版本,这可能是一个很好的折衷方案

vpermq在Zen 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)
}
# 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

在Intel Haswell/Skylake上,对于端口5,这是5个uops,加上移位(p01)和vpmovmskb(p0)。所以总共有7个uop。(不包括应该通过内联消除的ret或浪费的vmovdqa。)

在Ice Lake和更高版本中,vpinsrd中的一个uop可以在p15上运行,如果你在循环中这样做,可以减轻该端口上的一个压力。CCD_ 40是Alder湖E核上的一个单核。

Ice Lake(及以后的版本)也可以在p1/p5上运行vpshufb,进一步降低端口5的压力,降至7个uop中的3个。端口5可以处理任何混洗,端口1可以处理一些但不是所有的混洗uop。它可以连接到512位混洗单元的上半部分,以便为一些256位和更窄的混洗提供额外的吞吐量,比如p0/p1 FMA单元如何在p0上作为单个512位FMA单元工作。它不处理vpermqvpacksswb;这些仍然是p5只在冰/奥尔德湖。

因此,这个版本在当前一代和未来的英特尔CPU上是相当合理的。Alder Lake E-cores以2个uops的方式运行vpermq ymm,具有7个周期的延迟。但是,如果他们可以通过更有限的无序调度(大的ROB,但每个端口的队列没有那么长)来隐藏延迟,那么将vpinsrd作为单个uop运行有助于弥补前端吞吐量。

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


我曾考虑过将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

Zen 1具有非常快的vextracti128——对于任何端口都是单个uop,128位矢量操作是1个uop,而对于__m256i操作是2个uop。我们已经在进行提取,以将顶部和底部结合在一起。

但它仍然会导致更多的标量运算,尤其是如果您希望将结果组合到向量中。2xvpinsrw或在CCD_。

#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

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

  • 如何有效地解交织比特(逆Morton)
  • 如何去交错位(UnMortonization?)

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


AVX-512和/或GFNI

这里有两个有趣的策略:

  • vpmovw2m和/或vptestmwmb作为更方便的vpmovmskb。仅需要AVX-512BW(Skylake-ax512)
  • 将8位压缩到每个qword的底部,然后进行混洗。可能只适用于GFNI+AVX512VBMI,如Ice Lake/Zen4及更高版本。也许只是GFNI+AVX2,就像在瘫痪的奥尔德湖(没有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上的单个uop,但在Intel上,它的成本为2个uop,与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)

其中xmm0可以由vpermq设置为在低位16字节中具有top:bottom;所有采用AVX-512的CPU都具有高效的CCD_ 69。因此,在我手写的asm的5个uop之上又增加了1个uop(用内部函数编写应该很简单,我只是觉得在找到可用的asm指令后,不想再花额外的步骤查找正确的内部函数。)

将位封装在qword中,然后进行混洗:对于vpermb,GFNI和可能的AVX-512VBMI

(需要AVX512VBMI意味着Ice Lake或Zen 4,所以vpermb将是单个uop。除非未来某个带E核的英特尔CPU支持较慢的AVX-512,但vpermb ymm希望不会太差。)

可能按左:右顺序打包(每个半字节1个),然后进行字节洗牌。如果我们可以在交替的字节中进行left:rightright:left,则字节混洗(类似于vpermbvpermt2b)应该能够设置用于vprolw在每个16位字内旋转到组8";"左";按正确顺序排列的位。

在一个qword中移动位:Harold使用SIMD将位包ascii字符串的答案转换为7位二进制blob,显示_mm256_gf2p8affine_epi64_epi8将每个字节中的1位放在每个qword的顶部。(并打包剩余的7位字段,这就是答案中的目标。)

如果这是可行的,它可能会比戴口罩和返回口罩更少的uop和更好的延迟。

有了Alder Lake(GFNI,但AVX-512被禁用,除非你设法避免英特尔的努力削弱这个惊人的CPU),这可能仍然有用,因为它有用于_mm256_gf2p8affine_epi64_epi8的AVX+GFNI。vpshufb+vpermd可以代替vpermb。但你不会让单词旋转;不过,像ABAB这样的混洗字节可以让你使用纯左移来获得你想要的窗口,然后再次混洗。

最新更新