模拟 AVX512 VPCOMPESSB 字节打包，无需 AVX512_VBMI2

Question

我已经用 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 所做的那样？

AVX512VBMI2 有

VPCOMPRESSB

，但仅在 Ice Lake 及更高版本中可用。 Skylake-avx512 / Cascade Lake 可以做什么？

Answer 1

这是我用作 vpcompressb 和 vpexpandb 的 AVX2 后备的内容。它基于上一个问题中发布的解决方案。

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	Compress AVX512
10.22	97,820,644.12	1.2%	0.17	Compress AVX2
24.36	41,053,932.84	0.7%	0.39	Compress Scalar
1.09	918,672,592.02	1.2%	0.03	Decompress AVX512
6.44	155,358,119.20	0.5%	0.11	Decompress AVX2
24.83	40,274,581.73	0.7%	0.40	Decompress Scalar

compressstore

显然比

compress+store

慢两倍，因此尽可能避免使用它可能是个好主意。