AVX中的AVX2 VPSHUFB仿真

7

在AVX中,只有128位的PSHUFB

VPSHUFB xmm1, xmm2, xmm3/m128

只有 AVX2 具有完整的 PSHUFB,适用于整个 256 位 AVX 寄存器

VPSHUFB ymm1, ymm2, ymm3/m256

如何使用AVX指令高效地模拟此指令?

在这个特殊情况下,源只有8个元素(字节),但是可以在目标的整个32字节内移动。所以只需要运行2次PSHUFB就可以了。

我在使用VPSHUFB时遇到的问题是它将16(0x10)视为0,只有128及以上才填充为零!(最高位设置)。是否可能不添加比较和掩码操作来完成它?


4
也许你没有注意到,但 AVX2 的 VPSHUFB ymm, ymm, ymm/m256 并不是真正的 256 位洗牌,而是两个 128 位的洗牌。 - Marat Dukhan
有趣。谢谢! - alecco
1个回答

10

正如@MaratDukhan所指出的那样,_mm256_shuffle_epi8(即ymm-s的VPSHUFB)并不执行完整的32字节洗牌。对我来说,这真是太遗憾了...

这就是为什么为了在没有AVX2的情况下模拟它,您可以简单地将每个寄存器分成两半,对每个半部分进行排列,然后将它们组合在一起的原因:

//AVX only
__m256i _emu_mm256_shuffle_epi8(__m256i reg, __m256i shuf) {
    __m128i reg0 = _mm256_castsi256_si128(reg);
    __m128i reg1 = _mm256_extractf128_si256(reg, 1);
    __m128i shuf0 = _mm256_castsi256_si128(shuf);
    __m128i shuf1 = _mm256_extractf128_si256(shuf, 1);
    __m128i res0 = _mm_shuffle_epi8(reg0, shuf0);
    __m128i res1 = _mm_shuffle_epi8(reg1, shuf1);
    __m256i res = _mm256_setr_m128i(res0, res1);
    return res;
}

如果您真的想要完全洗牌32字节寄存器,您可以遵循这篇论文中的方法。将每个半部分与另一个半部分混合,然后将结果混合在一起。没有AVX2的话,大致是这样的:

//AVX only
__m256i _emu_mm256_shuffle32_epi8(__m256i reg, __m256i shuf) {
    __m128i reg0 = _mm256_castsi256_si128(reg);
    __m128i reg1 = _mm256_extractf128_si256(reg, 1);
    __m128i shuf0 = _mm256_castsi256_si128(shuf);
    __m128i shuf1 = _mm256_extractf128_si256(shuf, 1);
    __m128i res00 = _mm_shuffle_epi8(reg0, shuf0);
    __m128i res01 = _mm_shuffle_epi8(reg0, shuf1);
    __m128i res10 = _mm_shuffle_epi8(reg1, shuf0);
    __m128i res11 = _mm_shuffle_epi8(reg1, shuf1);
    __m128i res0 = _mm_blendv_epi8(res10, res00, _mm_cmplt_epi8(shuf0, _mm_set1_epi8(16)));
    __m128i res1 = _mm_blendv_epi8(res11, res01, _mm_cmplt_epi8(shuf1, _mm_set1_epi8(16)));
    __m256i res = _mm256_setr_m128i(res0, res1);
    return res;
}

如果你确定只使用了reg的下半部分,那么可以删除reg1res10res11行,并且删除比较和混合。实际上,如果没有AVX2,坚持使用SSE并使用128位寄存器可能更有效率。

使用AVX2可以显着优化通用的32字节重排:

//Uses AVX2
__m256i _ext_mm256_shuffle32_epi8(__m256i reg, __m256i shuf) {
    __m256i regAll0 = _mm256_permute2x128_si256(reg, reg, 0x00);
    __m256i regAll1 = _mm256_permute2x128_si256(reg, reg, 0x11);
    __m256i resR0 = _mm256_shuffle_epi8(regAll0, shuf);
    __m256i resR1 = _mm256_shuffle_epi8(regAll1, shuf);
    __m256i res = _mm256_blendv_epi8(resR1, resR0, _mm256_cmpgt_epi8(_mm256_set1_epi8(16), shuf));
    return res;
}

注意:代码未经测试!


非常详细的回答,比我预期的要多得多。非常感谢! - alecco
你是否意识到你的 AVX 实现中包含 _mm256_extracti128_si256,而该函数仅在 AVX2 中可用? - plasmacel
@stgatilov 请再检查一下,_mm256_extracti128_si256 编译成 vextracti128,这是 AVX2。你说的是 _mm256_extractf128_si256,它编译成 vextractf128。它们很相似,但是 vextractf128 操作浮点数域,而 vextracti128 操作整数域。 - plasmacel
@plasmacel:抱歉,我没有注意到我使用了两种不同的提取方式。希望现在我已经解决了这个问题。 - stgatilov
我在我的代码中的许多地方测试了 _ext_mm256_shuffle32_epi8(),它表现良好。对于经常使用 __m256i 的任何人来说,这是必备的。 - Michel Rouzic

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接