我已经用 0-63 的字节整数数组填充了 zmm 寄存器。 这些数字充当矩阵的索引。 非零元素表示矩阵中包含数据的行。 并非所有行都包含数据。
我正在寻找一个指令(或多个指令)来对 zmm 寄存器中的字节数组执行与 VPCOMPRESSQ 对 zmm 寄存器中的 qword 数组相同的操作。
考虑第 3、6、7、9 和 12 行包含数据而所有其他行(总共 64 行或更少)为空的情况。 掩码寄存器 k1 现在包含 64 位,设置为 001001001001000 ... 0。
使用 VPCOMPRESSQ 指令,我可以使用 k1 来压缩 zmm 寄存器,以便非零元素从寄存器的开头连续排列,但对于字节没有等效的指令。
PSHUFB 作为候选出现,但随机控制掩码必须是整数,而不是位。 我可以获得一个整数数组 0 0 3 0 0 6 等,但我仍然将元素设置为零,并且我需要它们从 zmm 寄存器的开头连续排列。
所以我的问题是,给定一个带有字节数组的 zmm 寄存器,其中 k1 设置如上所示,如何才能从 zmm 寄存器的开头连续排列非零元素,就像 VPCOMPRESSQ 对 qwords 所做的那样?
VPCOMPRESSB
,但仅在 Ice Lake 及更高版本中可用。 Skylake-avx512 / Cascade Lake 可以做什么?
这是我用作 vpcompressb 和 vpexpandb 的 AVX2 后备的内容。它基于上一个问题中发布的解决方案。
请注意,跨页面时 Decompress() 可能会加载越界和错误。如果这是一个问题,可以将输入复制到临时缓冲区中,就像在 Compress() 中所做的那样。
void MaskCompress_AVX2(const uint8_t src[64], uint8_t* dest, uint64_t mask) {
alignas(64) uint8_t temp[64];
int j = 0;
for (int i = 0; i < 64; i += 16) {
// ABC... -> AAAABBBBCCCC...: replicate each bit to fill a nibble
uint64_t extended_mask = _pdep_u64(mask, 0x1111'1111'1111'1111) * 0xF;
uint64_t dest_idx_nib = _pext_u64(0xFEDC'BA98'7654'3210, extended_mask);
__m128i dest_idx = _mm_cvtsi64_si128((int64_t)dest_idx_nib);
dest_idx = _mm_unpacklo_epi8(dest_idx, _mm_srli_epi16(dest_idx, 4));
dest_idx = _mm_and_si128(dest_idx, _mm_set1_epi8(15));
// load will never be out of bounds because `j < i = always true`
auto values = _mm_shuffle_epi8(_mm_loadu_si128((__m128i*)&src[i]), dest_idx);
_mm_storeu_si128((__m128i*)&temp[j], values);
j += _mm_popcnt_u32(mask & 0xFFFF);
mask >>= 16;
}
__builtin_memcpy(dest, temp, (uint32_t)j);
}
void MaskDecompress_AVX2(const uint8_t* src, uint8_t dest[64], uint64_t mask) {
for (int i = 0, j = 0; i < 64; i += 16) {
// ABC... -> AAAABBBBCCCC...: replicate each bit to fill a nibble
uint64_t extended_mask = _pdep_u64(mask, 0x1111'1111'1111'1111) * 0xF;
uint64_t dest_idx_nib = _pdep_u64(0xFEDC'BA98'7654'3210, extended_mask);
__m128i dest_idx = _mm_cvtsi64_si128((int64_t)dest_idx_nib);
dest_idx = _mm_unpacklo_epi8(dest_idx, _mm_srli_epi16(dest_idx, 4));
dest_idx = _mm_and_si128(dest_idx, _mm_set1_epi8(15));
__m128i zero_mask = _mm_cvtsi64_si128((int64_t)~extended_mask);
zero_mask = _mm_unpacklo_epi8(zero_mask, _mm_srli_epi16(zero_mask, 4));
// shuffle_epi8 outputs zeroes when high index bit is set
dest_idx = _mm_or_si128(dest_idx, _mm_slli_epi16(zero_mask, 4));
// load will never be out of bounds because `j < i = always true`
auto values = _mm_shuffle_epi8(_mm_loadu_si128((__m128i*)&src[j]), dest_idx);
_mm_storeu_si128((__m128i*)&dest[i], values);
j += _mm_popcnt_u32(mask & 0xFFFF);
mask >>= 16;
}
}
void MaskCompress_Scalar(const uint8_t src[64], uint8_t* dest, uint64_t mask) {
for (uint32_t i = 0, j = 0; mask != 0; i++) {
dest[j] = src[i];
j += (mask & 1);
mask >>= 1;
}
}
void MaskDecompress_Scalar(const uint8_t* src, uint8_t dest[64], uint64_t mask) {
for (uint32_t i = 0, j = 0; i < 64; i++) {
uint8_t v = src[j];
dest[i] = (mask & 1) ? v : 0;
j += (mask & 1);
mask >>= 1;
}
}
我的 Tigerlake CPU 的一些基准测试(掩码是随机生成的):
ns/op | 操作/s | 错误% | 总计 | 基准 |
---|---|---|---|---|
2.69 | 371,063,085.72 | 1.9% | 0.04 |
|
10.22 | 97,820,644.12 | 1.2% | 0.17 |
|
24.36 | 41,053,932.84 | 0.7% | 0.39 |
|
1.09 | 918,672,592.02 | 1.2% | 0.03 |
|
6.44 | 155,358,119.20 | 0.5% | 0.11 |
|
24.83 | 40,274,581.73 | 0.7% | 0.40 |
|
compressstore
显然比 compress+store
慢两倍,因此尽可能避免使用它可能是个好主意。