c-只保留16位字中的10个有用位



我有_m256i向量,这些向量包含16位整数中的10位字(因此16*16位仅包含16*10个有用位)。只提取这10个比特并将其打包以产生10个比特值的输出比特流的最佳/最快方法是什么?

这是我的尝试。

还没有进行基准测试,但我认为它总体上应该工作得很快:没有太多指令,所有指令在现代处理器上都有1个延迟周期。存储也是有效的,2个存储指令用于20字节的数据。

该代码仅使用3个常量。如果在循环中调用此函数,那么好的编译器应该在循环外加载所有三个函数,并将它们保存在寄存器中。

// bitwise blend according to a mask
inline void combineHigh( __m256i& vec, __m256i high, const __m256i lowMask )
{
vec = _mm256_and_si256( vec, lowMask );
high = _mm256_andnot_si256( lowMask, high );
vec = _mm256_or_si256( vec, high );
}
// Store 10-bit pieces from each of the 16-bit lanes of the AVX2 vector.
// The function writes 20 bytes to the pointer.
inline void store_10x16_avx2( __m256i v, uint8_t* rdi )
{
// Pack pairs of 10 bits into 20, into 32-bit lanes
__m256i high = _mm256_srli_epi32( v, 16 - 10 );
const __m256i low10 = _mm256_set1_epi32( ( 1 << 10 ) - 1 ); // Bitmask of 10 lowest bits in 32-bit lanes
combineHigh( v, high, low10 );
// Now the vector contains 32-bit lanes with 20 payload bits / each
// Pack pairs of 20 bits into 40, into 64-bit lanes
high = _mm256_srli_epi64( v, 32 - 20 );
const __m256i low20 = _mm256_set1_epi64x( ( 1 << 20 ) - 1 ); // Bitmask of 20 lowest bits in 64-bit lanes
combineHigh( v, high, low20 );
// Now the vector contains 64-bit lanes with 40 payload bits / each
// 40 bits = 5 bytes, store initial 4 bytes of the result
_mm_storeu_si32( rdi, _mm256_castsi256_si128( v ) );
// Shuffle the remaining 16 bytes of payload into correct positions.
// The indices of the payload bytes are [ 0 .. 4 ] and [ 8 .. 12 ]
// _mm256_shuffle_epi8 can only move data within 16-byte lanes
const __m256i shuffleIndices = _mm256_setr_epi8(
// 6 remaining payload bytes from the lower half of the vector
4, 8, 9, 10, 11, 12,
// 10 bytes gap, will be zeros
-1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
// 6 bytes gap, will be zeros
-1, -1, -1, -1, -1, -1,
// 10 payload bytes from the higher half of the vector
0, 1, 2, 3, 4,
8, 9, 10, 11, 12
);
v = _mm256_shuffle_epi8( v, shuffleIndices );
// Combine and store the final 16 bytes of payload
const __m128i low16 = _mm256_castsi256_si128( v );
const __m128i high16 = _mm256_extracti128_si256( v, 1 );
const __m128i result = _mm_or_si128( low16, high16 );
_mm_storeu_si128( ( __m128i* )( rdi + 4 ), result );
}

此代码截断未使用的值的较高6位。


如果你想饱和,你还需要一条指令_mm256_min_epu16

此外,如果这样做,函数的第一步可以使用pmaddwd。这是一个完整的函数,它使源数字饱和,并进行了一些额外的调整。

// Store 10-bit pieces from 16-bit lanes of the AVX2 vector, with saturation.
// The function writes 20 bytes to the pointer.
inline void store_10x16_avx2( __m256i v, uint8_t* rdi )
{
const __m256i low10 = _mm256_set1_epi16( ( 1 << 10 ) - 1 );
#if 0
// Truncate higher 6 bits; pmaddwd won't truncate, it needs zeroes in the unused higher bits.
v = _mm256_and_si256( v, low10 );
#else
// Saturate numbers into the range instead of truncating
v = _mm256_min_epu16( v, low10 );
#endif
// Pack pairs of 10 bits into 20, into 32-bit lanes
// pmaddwd computes a[ 0 ] * b[ 0 ] + a[ 1 ] * b[ 1 ] for pairs of 16-bit lanes, making a single 32-bit number out of two pairs.
// Initializing multiplier with pairs of [ 1, 2^10 ] to implement bit shifts + packing
const __m256i multiplier = _mm256_set1_epi32( 1 | ( 1 << ( 10 + 16 ) ) );
v = _mm256_madd_epi16( v, multiplier );
// Now the vector contains 32-bit lanes with 20 payload bits / each
// Pack pairs of 20 bits into 40 in 64-bit lanes
__m256i low = _mm256_slli_epi32( v, 12 );
v = _mm256_blend_epi32( v, low, 0b01010101 );
v = _mm256_srli_epi64( v, 12 );
// Now the vector contains 64-bit lanes with 40 payload bits / each
// 40 bits = 5 bytes, store initial 4 bytes of the result
_mm_storeu_si32( rdi, _mm256_castsi256_si128( v ) );
// Shuffle the remaining 16 bytes of payload into correct positions.
const __m256i shuffleIndices = _mm256_setr_epi8(
// Lower half
4, 8, 9, 10, 11, 12,
-1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
// Higher half
-1, -1, -1, -1, -1, -1,
0, 1, 2, 3, 4,
8, 9, 10, 11, 12
);
v = _mm256_shuffle_epi8( v, shuffleIndices );
// Combine and store the final 16 bytes of payload
const __m128i low16 = _mm256_castsi256_si128( v );
const __m128i high16 = _mm256_extracti128_si256( v, 1 );
const __m128i result = _mm_or_si128( low16, high16 );
_mm_storeu_si128( ( __m128i* )( rdi + 4 ), result );
}

这可能会稍快或稍慢,具体取决于处理器、编译器和调用函数的代码,但肯定有助于代码大小。没有人再关心二进制大小了,但CPU的L1I和µop缓存有限。


为了完整性,这里有另一个使用SSE2和可选的SSSE3而不是AVX2的,在实践中只稍微慢一点。

// Compute v = ( v & lowMask ) | ( high & ( ~lowMask ) ), for 256 bits of data in two registers
inline void combineHigh( __m128i& v1, __m128i& v2, __m128i h1, __m128i h2, const __m128i lowMask )
{
v1 = _mm_and_si128( v1, lowMask );
v2 = _mm_and_si128( v2, lowMask );
h1 = _mm_andnot_si128( lowMask, h1 );
h2 = _mm_andnot_si128( lowMask, h2 );
v1 = _mm_or_si128( v1, h1 );
v2 = _mm_or_si128( v2, h2 );
}
inline void store_10x16_sse( __m128i v1, __m128i v2, uint8_t* rdi )
{
// Pack pairs of 10 bits into 20, in 32-bit lanes
__m128i h1 = _mm_srli_epi32( v1, 16 - 10 );
__m128i h2 = _mm_srli_epi32( v2, 16 - 10 );
const __m128i low10 = _mm_set1_epi32( ( 1 << 10 ) - 1 );
combineHigh( v1, v2, h1, h2, low10 );
// Pack pairs of 20 bits into 40, in 64-bit lanes
h1 = _mm_srli_epi64( v1, 32 - 20 );
h2 = _mm_srli_epi64( v2, 32 - 20 );
const __m128i low20 = _mm_set1_epi64x( ( 1 << 20 ) - 1 );
combineHigh( v1, v2, h1, h2, low20 );
#if 1
// 40 bits is 5 bytes, for the final shuffle we use pshufb instruction from SSSE3 set
// If you don't have SSSE3, below under `#else` there's SSE2-only workaround.
const __m128i shuffleIndices = _mm_setr_epi8(
0, 1, 2, 3, 4,
8, 9, 10, 11, 12,
-1, -1, -1, -1, -1, -1 );
v1 = _mm_shuffle_epi8( v1, shuffleIndices );
v2 = _mm_shuffle_epi8( v2, shuffleIndices );
#else
// SSE2-only version of the above, uses 8 instructions + 2 constants to emulate 2 instructions + 1 constant
// Need two constants because after this step we want zeros in the unused higher 6 bytes.
h1 = _mm_srli_si128( v1, 3 );
h2 = _mm_srli_si128( v2, 3 );
const __m128i low40 = _mm_setr_epi8( -1, -1, -1, -1, -1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 );
const __m128i high40 = _mm_setr_epi8( 0, 0, 0, 0, 0, -1, -1, -1, -1, -1, 0, 0, 0, 0, 0, 0 );
const __m128i l1 = _mm_and_si128( v1, low40 );
const __m128i l2 = _mm_and_si128( v2, low40 );
h1 = _mm_and_si128( h1, high40 );
h2 = _mm_and_si128( h2, high40 );
v1 = _mm_or_si128( h1, l1 );
v2 = _mm_or_si128( h2, l2 );
#endif
// Now v1 and v2 vectors contain densely packed 10 bytes / each.
// Produce final result: 16 bytes in the low part, 4 bytes in the high part
__m128i low16 = _mm_or_si128( v1, _mm_slli_si128( v2, 10 ) );
__m128i high16 = _mm_srli_si128( v2, 6 );
// Store these 20 bytes with 2 instructions
_mm_storeu_si128( ( __m128i* )rdi, low16 );
_mm_storeu_si32( rdi + 16, high16 );
}

在循环中,您可能希望使用部分重叠的存储,这些存储在源数据的每个向量的20字节目标末尾之后写入。这样就省去了在16字节边界上打乱数据以设置16+4字节存储的工作。

(@Soont的更新答案是一个vmovd和一个vmovdqu存储非常好,总共只有两个shuffle uop,包括vpshufbvextracti128。当我最初写这篇文章时,我们还没有想到一个好的方法来避免存储在20字节之外,而不花更多的shuffle UOP,这会造成比前端更糟糕的瓶颈。但vmovdqu+vextracti128 mem, ymm, 1(两个uop未微融合)仍然是略便宜:vpshufb之后的3个uops而不是4个。)

或者展开可能对大型阵列有利,LCM(20,16)=80,因此使用大型展开(以及其中每个位置的不同混洗控制向量),您可以只进行对齐的16字节存储。但这可能需要大量的洗牌,包括可能使用palignr的源块之间的洗牌。


两个重叠的16字节存储的示例

将其用作循环体,其中覆盖过去20个字节是可以的

#include <immintrin.h>
#include <stdint.h>
// Store 10-bit pieces from each of the 16-bit lanes of the AVX2 vector.
// The function writes 20 useful bytes to the pointer
// but actually steps on data out to 26 bytes from dst
void pack10bit_avx2_store26( __m256i v, uint8_t* dst)
{
// clear high garbage if elements aren't already zero-extended   
//v = _mm256_and_si256(v, _mm256_set1_epi16( (1<<10)-1) );
... prep data somehow; pmaddwd + a couple shifts is good for throughput
// Now the vector contains 64-bit lanes with 40 payload bits / each; 40 bits = 5 bytes.
// Shuffle these bytes into a very special order.
// Note _mm256_shuffle_epi8 can only move data within 16-byte lanes.
const __m256i shuffleIndices = _mm256_setr_epi8(
// 6 bytes gap with zeros
// Pack the two 5-byte chunks into the bottom of each 16-byte lane
0, 1, 2, 3, 4,
8, 9, 10, 11, 12,
-1, -1, -1, -1, -1, -1,
0, 1, 2, 3, 4,
8, 9, 10, 11, 12,
-1, -1, -1, -1, -1, -1);
v = _mm256_shuffle_epi8(v, shuffleIndices );
// Split the vector into halves
__m128i low16 = _mm256_castsi256_si128( v );
_mm_storeu_si128( ( __m128i* )dst, low16 );        // vmovdqu      mem, xmm
__m128i high16 = _mm256_extracti128_si256( v, 1 );
_mm_storeu_si128( ( __m128i* )(dst+10), high16 );   // vextracti128 mem, ymm, 1
// An AVX-512 masked store could avoid writing past the end
}

我们可以通过将它编译成一个独立的函数来了解它是如何内联到循环中的(https://godbolt.org/z/8T7KhT)。

# clang -O3 -march=skylake
pack10bit_avx2(long long __vector(4), unsigned char*):
# vpand  commented out
vpmaddwd        ymm0, ymm0, ymmword ptr [rip + .LCPI0_0]
... # work in progress, original PMADDWD idea ignored some limitations!  See Soonts' answer
vpshufb ymm0, ymm0, ymmword ptr [rip + .LCPI0_1] # ymm0 = ymm0[0,1,2,3,4,8,9,10,11,12],zero,zero,zero,zero,zero,zero,ymm0[16,17,18,19,20,24,25,26,27,28],zero,zero,zero,zero,zero,zero
vmovdqu xmmword ptr [rdi], xmm0
vextracti128    xmmword ptr [rdi + 10], ymm0, 1
vzeroupper               # overhead that goes away when inlining into a loop
ret

在循环中,编译器会将这两个向量常量加载到寄存器中,希望使用广播加载。

与一些更宽的整数乘法或水平加法不同,vpmaddwd是有效的,作为一个具有5个周期延迟的单个uop。https://uops.info/

vextracti128存储无法在英特尔上进行微融合,但与vpextrd不同的是,它没有涉及shuffle uop。只需存储地址和数据。Zen2还将其作为2个uops运行,不幸的是,每2个周期的吞吐量为1。(比Zen1更糟糕)。

在冰湖之前,英特尔和AMD每个时钟都可以运行1个存储。


如果你真的想把压缩后的数据放回寄存器,你可能想要使用palignr的@Soont的原始混洗,或者你可以做一个块,然后重新加载。延迟会更高(尤其是因为在重新加载时存储转发暂停),但如果你的块是几个寄存器的数据,那么它应该重叠甚至隐藏延迟,可能会给存储时间提交到L1d,并且在重新加载后不会导致暂停。


BMI2pext

uint64_t packed = _pext_u64(x, 0x03FF03FF03FF03FF);

可能适合标量清理或4个像素的短块或其他什么。这就给您留下了一个5字节存储(或后面有0的8字节存储)的问题。如果使用这种方法,请注意严格的混叠和对齐,例如,使用memcpy进行未对齐可能会将数据混叠到uint64_t中,或制作__attribute__((aligned(1),may_alias))typedef。

pext在Intel上非常有效(1 uop,3c延迟),但在AMD上非常糟糕,比仅使用一个SIMD步骤的低部分差得多。


AVX-512

AVX512VBMI(冰湖)会给你vpermb(车道交叉)而不是vpshufb。(Skylake-X/Cascade Lake上vpermw的AVX512BW需要您已经组合成偶数个字节,即使在vpermb为1的Ice Lake中也是2个uops,所以这非常糟糕。)vpermb可以设置为单个未对齐的32字节存储(有20个有用的字节),您可以在循环中重叠。

AVX-512存储可以被有效地屏蔽,以不实际上覆盖结束,例如使用双字屏蔽。CCD_ 24在Skylake-X上为1μp。但AVX2vmaskmovd即使在英特尔上也只有几个uop,在AMD上也非常昂贵,所以你不想这么做。只有当您为一个存储区准备好所有20个字节时,双字掩码才有效,否则您至少需要16位粒度。

其他AVX-512指令:VBMIvpmultishiftqb,一个并行位字段提取,看起来可能很有用,但它只能从未对齐但连续的源块中写入对齐的8位目标块。我不认为这比我们可以通过可变的移位和旋转来做的更好vpmultishiftqb将允许我们在可能的2条指令中解压缩这种格式(此函数的逆函数):1个shuffle(如vpexpandbvpermb)将所需数据放入向量中的每个qword,1个multi-shift为每个字的底部获取右侧10位字段。

AVX-512具有可变计数移位和旋转,包括字(16位)粒度,因此这将是第一步的vpmaddwd的替代方案使用shift可以免费忽略高垃圾它具有较低的延迟,即时版本的合并掩码可以取代对控制向量的需要。(但你需要一个掩码常量)。

有了掩蔽,等待时间是3个周期,而没有,AVX-512使得从立即广播控制向量到mov reg,imm/kmov kreg, reg的效率大约相同。例如CCD_ 33/CCD_。合并屏蔽还限制优化器覆盖目标寄存器,而不是复制和移位,尽管如果优化器是智能的,这在这里并不重要。这两种方式都不允许将数据加载合并到移位的内存源操作数中:sllvw只能从内存中获取计数,而sllw需要合并到寄存器中的原始操作数中。

移位可以在英特尔的端口0或1上运行(AMD不支持AVX-512)。或者仅512位uop的端口0,在任何512位uops运行时关闭任何矢量ALU uop的1。因此,对于__m512i版本,端口0上存在潜在的吞吐量瓶颈,但对于256位,有足够多的其他uop(shuffle和store,如果对数据阵列这样做,可能会产生循环开销),因此应该相当均匀地分布。

这个移位部分(_mm256_permutexvar_epi8之前)只需要AVX-512BW(+VL),并且将在Skylake-X上工作它将数据保留在与其他方法相同的位置,因此是一种可以与各种策略混合和匹配的替代方法。

// Ice Lake.  Could work on __m512i but then shifts could only run on p0, not p0/p1,
//  and almost every store would be a cache line split.
inline void store_10x16_avx512vbmi( __m256i v, uint8_t* dst )
{
// no _mm256_and_si256 needed, we safely ignore high bits
// v = [ ?(6) ... B[9:0] | ?(6) ... A[9:0] ] repeated
v = _mm256_sllv_epi16(v, _mm256_set1_epi32((0<<16) | 6));  // alternative: simple repeated-pattern control vector
// v =  _mm256_mask_slli_epi16(v, 0x5555, v, 6);   // merge-masking, updating only elements 0,2, etc.
// v = [ ?(6) ... B[9:0] | A[9:0] ... 0(6) ] repeated
v = _mm256_rolv_epi32(v, _mm256_set1_epi64x(((32ULL-6)<<32) | 6));  // top half right, bottom half left
// v = [ 0(6) .. ?(6) .. D[9:0] | C[9:0] | B[9:0] | A[9:0] ... 0(12) ] repeated
v = _mm256_srli_epi64(v, 12);    // 40 bit chunks at the bottom of each qword
const __m256i permb = _mm256_setr_epi8( 0, 1, 2, 3, 4,   8, 9,10,11,12,
16,17,18,19,20,  24,25,26,27,28,
28,28,28,28,28,28,28,28,28,28,28,28 );
// repeat last byte as filler.  vpermb can't zero (except by maskz) but we can do a masked store
v = _mm256_permutexvar_epi8(v, permb);  // AVX512_VBMI
_mm256_mask_storeu_epi32( dst, 0x1F, v);  // 32-bit masking granularity in case that's cheaper for HW.  20 bytes = 5 dwords.
}

编译如下(Godbolt):

# clang -O3 -march=icelake-client.  GCC is essentially the same.
store_10x16_avx512vbmi(long long __vector(4), unsigned char*):
vpsllvw ymm0, ymm0, ymmword ptr [rip + .LCPI0_0]
vprolvd ymm0, ymm0, ymmword ptr [rip + .LCPI0_1]
vpsrlq  ymm0, ymm0, 12
vpermb  ymm0, ymm0, ymmword ptr [rip + .LCPI0_2]
mov     al, 31           # what the heck, clang? partial register false dependency for no reason!
kmovd   k1, eax
vmovdqu32       ymmword ptr [rdi] {k1}, ymm0
# vzeroupper not needed because the caller was using __m256i args.  GCC omits it.
ret

即使您两次使用相同的移位常量向量,使编译器将其保留在寄存器中(而不是直接从内存源操作数中使用),它仍然选择从内存加载,而不是从mov eax,6/vpbroadcast ymm1, eax或其他地方加载。这以需要.rodata中的常量为代价节省了1个uop。公平地说,我们确实需要可能在同一缓存行中的其他常量,但GCC浪费空间的方式是,它们并不都适合一个缓存行!clang注意到该模式并使用vpbroadcastdq加载,gcc则浪费地加载了整整32个字节。(kmov k1, [mem]是3个前端uop,因此它不会保存一个uop来从内存加载掩码常量。)

使用_mm256_mask_slli_epi16(v, 0x5555, v, 6),clang将其优化为具有相同6,0重复常数的vpsllvw ymm0, ymm0, ymmword ptr [rip + .LCPI0_0]。所以我想这是一个好兆头,我做对了。但是GCC编译如下:

store_10x16_avx512vbmi(long long __vector(4), unsigned char*):
mov     eax, 21845
kmovw   k1, eax
vpsllw  ymm0{k1}, ymm0, 6
vprolvd ymm0, ymm0, YMMWORD PTR .LC0[rip]
mov     eax, 31
kmovb   k2, eax
vpsrlq  ymm0, ymm0, 12
vpermb  ymm0, ymm0, YMMWORD PTR .LC1[rip]
vmovdqu32       YMMWORD PTR [rdi]{k2}, ymm0
ret

_mm256_sllv_epi16需要AVX-512BW和AVX-512VL。rolv_epi32只需要AVX-512VL。(或者对于512位版本,只有AVX-512F。)旋转只有32和64个元素大小,而不是16,但AVX-512确实将可变移位粒度扩展到了16(从AVX2中的32或64)。

vpcompressb [rdi]{k1}, ymm0(AVX512VBMI=Ice Lake及更高版本)将是vpermb+存储的替代方案,以将字节打包在寄存器底部(类似于BMI2pext,但用于矢量元素而不是标量寄存器中的位)。但它实际上更昂贵:冰湖上有6个单位,每6厘米就有一个单位。(vpcompressd没有那么糟糕)。

即使vpcompressb进入矢量寄存器也是2 uops,因此对于恒定混洗控制,最好为vpermb加载矢量常数,除非控制矢量的缓存未命中是个问题,例如,如果您每隔一段时间只这样做一次,那么让HW处理k掩码而不是加载。


不带VBMI的AVX-512:2x 16字节存储,不超过20字节范围

...  // same setup as usual, leaving 40-bit chunks at the bottom of each qword

const __m256i shuffleIndices = _mm256_setr_epi8(
// 6 bytes gap with zeros
// Pack the two 5-byte chunks into the bottom of each 16-byte lane
0, 1, 2, 3, 4,
8, 9, 10, 11, 12,
-1, -1, -1, -1, -1, -1,
0, 1, 2, 3, 4,
8, 9, 10, 11, 12,
-1, -1, -1, -1, -1, -1);
v = _mm256_shuffle_epi8(v, shuffleIndices );
// Split the vector into halves
__m128i low16 = _mm256_castsi256_si128( v );
_mm_storeu_si128( ( __m128i* )dst, low16 );        // vmovdqu      mem, xmm  no masking
// An AVX-512BW masked store avoiding writing past the end costs more instructions (and back-end uops), same front-end uops
__m128i high16 = _mm256_extracti128_si256( v, 1 );  // vextracti128 xmm, ymm, 1
_mm_mask_storeu_epi8( dst+10, 0x3FF, high16 );      // vmovdqu8 [mem]{k}, xmm

这需要vextracti128 xmm, ymm, 1vmovdqu8进行设置。与写入26个字节不同,我们不能直接提取到内存中。没有vextracti8x16,只有vextracti32x464x2(以及32x8/64x4 256位提取)。我们需要字节粒度屏蔽,但不能用直接提取到内存的指令来实现,只能通过混洗(vextract进入寄存器)然后vmovdqu8

所以我们得到的asm是

# clang
...        vpshufb result in YMM0
vmovdqu      [rdi], xmm0             # same as before
vextracti128    xmm0, ymm0, 1        # 1 shuffle uop
mov     ax, 1023
kmovd   k1, eax                         # will be hoisted
vmovdqu8     [rdi + 10] {k1}, xmm0   # 1 micro-fused uop

由于vextracti128 [mem], ymm, 1是2个前端uop,所以这不会影响前端吞吐量。(由于shuffle uop,它确实给后端执行端口带来了更大的压力)。

最新更新