问题描述
我有 _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,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,// 6 bytes gap,// 10 payload bytes from the higher half of the vector
0,1,2,3,4,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,_mm256_castsi256_si128( v ) );
// Shuffle the remaining 16 bytes of payload into correct positions.
const __m256i shuffleIndices = _mm256_setr_epi8(
// Lower half
4,// Higher half
-1,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,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,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 );
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,0 );
const __m128i high40 = _mm_setr_epi8( 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 商店的答案非常好，只有 2 个 shuffle uops，包括 vpshufb 和 vextracti128。当我最初写对此，我们还没有想到一个好方法来避免存储在 20 个字节之外而不花费更多 shuffle uops，这会造成比前端更糟糕的瓶颈。但是 vmovdqu + vextracti128 mem,ymm,1 (2 uops 未微融合）仍然稍微便宜一些：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);
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,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,ymmword ptr [rip + .LCPI0_1] # ymm0 = ymm0[0,12],zero,ymm0[16,17,18,19,20,24,25,26,27,28],zero
vmovdqu xmmword ptr [rdi],xmm0
vextracti128 xmmword ptr [rdi + 10],1
vzeroupper # overhead that goes away when inlining into a loop
ret
在循环中，编译器会将这 2 个向量常量加载到寄存器中，希望使用广播加载。
与一些更宽的整数乘法或水平加法不同，vpmaddwd 的处理效率很高，作为具有 5 个周期延迟的单个 uop。 https://uops.info/
vextracti128 商店无法在 Intel 上进行微融合，但与 vpextrd 不同的是，它不涉及 shuffle uop。只是存储地址和存储数据。 Zen2 也将它作为 2 uop 运行，不幸的是，每 2 个周期的吞吐量为 1。 （比 Zen1 差）。
在 Ice Lake 之前，Intel 和 AMD 都可以每个时钟运行 1 个存储。
如果您确实希望将打包的数据放回寄存器中，您可能需要使用 palignr 进行 @Soont 的原始洗牌，或者您可以先执行此操作，然后重新加载一些数据。延迟会更高（特别是因为重新加载时存储转发停顿），但是如果您的块是几个寄存器的数据，那么它应该重叠甚至隐藏延迟，可能会给存储时间甚至提交到 L1d 而不会导致重新加载时出现停顿。
BMI2 pext
uint64_t packed = _pext_u64(x,0x03FF03FF03FF03FF);
也许适用于标量清理或一小块 4 像素或其他任何东西。这给您留下了进行 5 字节存储（或带有尾随零的 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 uop，所以非常糟糕.) vpermb 可以设置单个未对齐的 32 字节存储（具有 20 个有用字节），您可以在循环中重叠。
AVX-512 存储可以被有效地屏蔽以实际上覆盖末尾，例如使用双字屏蔽。 vmovdqu32 [rdi]{k},ymm0 在 Skylake-X 上为 1 uop。但是 AVX2 vmaskmovd 即使在 Intel 上也只有几个 uops，而且在 AMD 上非常昂贵，所以你不想这样做。并且双字掩码仅在您为一个存储准备好所有 20 个字节时才有效，否则您需要至少 16 位的粒度。
其他 AVX-512 指令：VBMI vpmultishiftqb，一个并行位域提取，似乎很有用，但它只能从未对齐但连续的源块中写入对齐的 8 位目标块。我不认为这比我们可以用可变移位和旋转做的更好。 vpmultishiftqb 会让我们解压这种格式（这个函数的反函数），大概有 2 条指令：1 shuffle（例如 vpexpandb 或 {{1 }}) 将所需数据放入向量中的每个 qword，并进行一次多移位以获取每个单词底部的正确 10 位字段。
AVX-512 具有可变计数的移位和旋转，包括字（16 位）粒度，因此这将是第一步而不是 vpermb 的选项。 使用 shift 可以免费忽略高垃圾。它具有较低的延迟，并且立即版本的合并屏蔽可以取代对控制向量的需求。 （但是你需要一个掩码常量）。
屏蔽延迟为 3 个周期，而没有屏蔽为 1 个周期，而 AVX-512 使得从立即到 vpmaddwd / mov reg,imm 广播控制向量的效率几乎相同。例如kmov kreg,reg / mov reg,imm (1 uop)。合并屏蔽还限制优化器覆盖目标寄存器而不是复制和移位，尽管如果优化器是智能的，这在这里无关紧要。。两种方法都不允许将数据的加载折叠到内存源操作数中进行移位：vpbroadcastd ymm,reg 只能从内存中获取计数，而 sllvw 需要合并到寄存器中的原始操作数中。
Shifts 可以在 Intel 的端口 0 或 1 上运行（并且 AMD 不支持 AVX-512）。或者只有端口 0 用于 512 位 uops，在任何 512 位 uops 正在运行时关闭任何 vector-ALU uop 的端口 1。因此，对于 sllw 版本，端口 0 上存在潜在的吞吐量瓶颈，但对于 256 位，还有足够多的其他 uops（洗牌和存储，如果对数据数组执行此操作，则可能会产生循环开销），因此应该相当均匀地分布。
这个shift部分（在__m512i之前）只需要AVX-512BW（+VL），并且可以在Skylake-X上工作。它和其他方法一样把数据放在同一个地方，所以是一个替代品，您可以混合搭配各种策略。
_mm256_permutexvar_epi8像这样编译 (Godbolt)：
// 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,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,16,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.
}
即使你两次使用同一个移位常量向量让编译器将它保存在一个寄存器中（而不是直接从内存源操作数使用），它仍然选择从内存加载它而不是# clang -O3 -march=icelake-client. GCC is essentially the same.
store_10x16_avx512vbmi(long long __vector(4),unsigned char*):
vpsllvw ymm0,ymmword ptr [rip + .LCPI0_0]
vprolvd ymm0,ymmword ptr [rip + .LCPI0_1]
vpsrlq ymm0,12
vpermb 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 什么的。这以需要 .rodata 中的常量为代价节省了 1 uop。公平地说，我们确实需要在同一个缓存行中的其他常量，但是 GCC 浪费空间的方式它们并不都适合一个缓存行！ clang 注意到该模式并使用 vpbroadcast ymm1,eax 或 vpbroadcastd 加载，gcc 浪费地加载了完整的 32 个字节。 （q 是 3 个前端 uops，因此从内存加载掩码常量不会节省 uop。）
使用 kmov k1,[mem]，clang 将其优化回具有相同 6,0 重复常量的 _mm256_mask_slli_epi16(v,6)。所以我想这是一个好兆头，我做对了。但是 GCC 是按照编写的方式编译的：
vpsllvw ymm0,ymmword ptr [rip + .LCPI0_0]store_10x16_avx512vbmi(long long __vector(4),unsigned char*):
mov eax,21845
kmovw k1,eax
vpsllw ymm0{k1},6
vprolvd ymm0,YMMWORD PTR .LC0[rip]
mov eax,31
kmovb k2,eax
vpsrlq ymm0,YMMWORD PTR .LC1[rip]
vmovdqu32 YMMWORD PTR [rdi]{k2},ymm0
ret
需要 AVX-512BW 和 AVX-512VL。 rolv_epi32 只需要 AVX-512VL。 （或者只是 512 位版本的 AVX-512F。）旋转只有 32 和 64 元素大小，而不是 16，但 AVX-512 确实将可变移位粒度扩展到 16（从 AVX2 中的 32 或 64）。
_mm256_sllv_epi16（AVX512VBMI = Ice Lake 及更高版本）将是 vpermb + store 的替代方案，用于在寄存器底部打包字节（如 BMI2 vpcompressb [rdi]{k1},ymm0 但对于向量元素而不是位标量寄存器）。但它实际上更昂贵：Ice Lake 上的 6 uop，每 6c 吞吐量一个。 （pext 还不错）。
即使 vpcompressd 进入向量寄存器也是 2 uop，因此对于常量 shuffle 控制，最好为 vpcompressb 加载向量常量，除非控制向量的缓存未命中是一个问题，例如如果你只是经常这样做一次，那么让硬件处理一个 k 掩码而不是一个负载。
不带 VBMI 的 AVX-512：2x 16 字节存储，但不超过 20 字节范围
vpermb这需要 ... // 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,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,1
_mm_mask_storeu_epi8( dst+10,0x3FF,high16 ); // vmovdqu8 [mem]{k},xmm
来设置 vextracti128 xmm,1。与写入 26 个字节不同，我们不能直接提取到内存中。没有 vmovdqu8，只有 vextracti8x16 和 vextracti32x4（以及 32x8 / 64x4 256 位提取）。我们需要字节粒度掩码，但不能通过直接提取到内存的指令获得它，只能通过混洗（64x2 到寄存器）然后vextract。
所以我们得到的 asm 是
vmovdqu8因为 # clang
... vpshufb result in YMM0
vmovdqu [rdi],xmm0 # same as before
vextracti128 xmm0,1 # 1 shuffle uop
mov ax,1023
kmovd k1,eax # will be hoisted
vmovdqu8 [rdi + 10] {k1},xmm0 # 1 micro-fused uop
无论如何都是 2 个前端 uops，所以这不会影响前端吞吐量。 （由于 shuffle uop，它确实对后端执行端口造成了更大压力）。