使用英特尔 AVX 进行掩码洗牌

Posted 2023-02-16

【中文标题】使用英特尔 AVX 进行掩码洗牌

【英文标题】：Shuffling by mask with Intel AVX

【发布时间】：2018-04-30 10:45:15

【问题描述】：

我是 AVX 编程新手。我有一个需要洗牌的寄存器。我想将 256 位寄存器 R1 中的几个字节洗牌到空寄存器 R2。我想定义一个掩码，它告诉洗牌操作应该将旧寄存器（R1）中的哪个字节复制到新寄存器的哪个位置。

掩码应如下所示（Src:Byte Pos in R1, Target:Byte Pos in R2）：

(0,0),(1,1),(1,4),(2,5),...

这意味着几个字节被复制了两次。

我不能 100% 确定我应该为此使用哪个函数。我对这两个 AVX 功能进行了一些尝试，第二个只使用了 2 个通道。

__m256 _mm256_permute_ps (__m256 a, int imm8)
__m256 _mm256_shuffle_ps (__m256 a, __m256 b, const int imm8)

我对 imm8 中的 Shuffle Mask 以及如何设计它以使其按上述方式工作感到非常困惑。

我查看了这个 slides（第 26 页），其中描述了 _MM_SHUFFLE，但我找不到解决问题的方法。

有没有关于如何设计这种面具的教程？或者这两种方法的示例函数来深入理解它们？

提前感谢您的提示

【参考方案1】：

TL:DR：您可能需要多次洗牌来处理车道交叉，或者如果您的模式继续完全一样，您可以使用_mm256_cvtepu16_epi32 (vpmovzxwd) 然后_mm256_blend_epi16。

---

对于 x86 shuffle（我认为与大多数 SIMD 指令集一样），目标位置是隐含的。 shuffle-control 常量仅具有目标顺序的源索引，无论是被编译+组装成 asm 指令的imm8，还是每个元素中都有索引的向量。

每个目标位置仅读取一个源位置，但可以多次读取同一源位置。每个目标元素都从随机播放源中获取一个值。

请参阅 Convert _mm_shuffle_epi32 to C expression for the permutation? 以获取 dst = _mm_shuffle_epi32(src, _MM_SHUFFLE(d,c,b,a)) 的纯 C 版本，显示如何使用控制字节。

（对于pshufb / _mm_shuffle_epi8，具有高位的元素将目标位置设置为零而不是读取任何源元素，但其他 x86 shuffle 会忽略 shuffle-control 向量中的所有高位。）

如果没有 AVX512 合并掩码，则不会有混入目标的随机播放。有一些像_mm256_shuffle_ps (vshufps) 这样的双源混洗可以将来自两个源的元素混洗在一起以产生单个结果向量。 如果您想保留一些未写入的目标元素，您可能必须先洗牌然后混合，例如使用_mm256_blendv_epi8，或者如果你可以使用16位粒度的混合，你可以使用更有效的即时混合_mm256_blend_epi16，甚至更好的_mm256_blend_epi32（AVX2vpblendd与英特尔CPU上的_mm256_and_si256一样便宜，如果您确实需要混合，如果它可以完成工作，这是最好的选择；请参阅http://agner.org/optimize/)

---

对于您的问题（在 Cannonlake 中没有 AVX512VBMI vpermb），您不能通过单个操作将单个字节从低 16“通道”转移到 __m256i 向量的高 16“通道” .

AVX shuffle 不像完整的 256 位 SIMD，它们更像是两个 128 位并行操作。唯一的例外是一些具有 32 位或更大粒度的 AVX2 车道交叉洗牌，例如 vpermd (_mm256_permutevar8x32_epi32)。还有pmovzx / pmovsx 的 AVX2 版本，例如pmovzxbq 将 XMM 寄存器的低 4 个字节零扩展为 YMM 寄存器的 4 个 qword，而不是 YMM 寄存器的每一半的低 2 个字节。这使得它对内存源操作数更有用。

但无论如何，pshufb (_mm256_shuffle_epi8) 的 AVX2 版本在 256 位向量的两个通道中执行两个单独的 16x16 字节混洗。

---

你可能会想要这样的东西：

// Intrinsics have different types for integer, float and double vectors
// the asm uses the same registers either way
__m256i  shuffle_and_blend(__m256i dst, __m256i src)

    // setr takes element in low to high order, like a C array init
    // unlike the standard Intel notation where high element is first
    const __m256i  shuffle_control = _mm256_setr_epi8(
          0,      1,  -1, -1,   1,      2, ...);
    // (0,0),  (1,1), (zero)  (1,4), (2,5),...  in your src,dst notation
    // Use -1 or 0x80 or anything with the high bit set
    //  for positions you want to leave unmodified in dst
   // blendv uses the high bit as a blend control, so the same vector can do double duty

    // maybe need some lane-crossing stuff depending on the pattern of your shuffle.
    __m256i  shuffled = _mm256_shuffle_epi8(src, shuffle_control);

    // or if the pattern continues, and you're just leaving 2 bytes between every 2-byte group:
    shuffled = _mm256_cvtepu16_epi32(src);  // if src is a __m128i

    __m256i  blended = _mm256_blendv_epi8(shuffled, dst, shuffle_control);
    // blend dst elements we want to keep into the shuffled src result.
    return blended;
    

请注意，pshufb 编号从 0 重新开始，用于第 2 个 16 字节。 __m256i 的两半可以不同，但​​它们不能从另一半读取元素。如果您需要高通道中的位置以从低通道获取字节，则需要更多的洗牌 + 混合（例如，包括 vinserti128 或 vperm2i128，或者可能是 vpermd 通道交叉双字洗牌）来获取所有您需要的字节按一些顺序放入一个 16 字节组中。

(实际上_mm256_shuffle_epi8 (PSHUFB) 忽略了 shuffle 索引中的第 4..6 位，所以写17 与1 相同，但非常具有误导性。它实际上是在执行%16，只要高位未设置。如果在随机控制向量中设置了高位，它将将该元素归零。我们在这里不需要该功能；_mm256_blendv_epi8 不关心它所在元素的旧值替换）

无论如何，这个简单的 2 指令示例仅在模式不继续时才有效。如果您在设计真正的 shuffle 时需要帮助，则必须提出更具体的问题。

---

顺便说一句，我注意到您的混合模式使用了 2 个新字节，然后 2 个跳过了 2 个。如果这种情况继续下去，您可以使用vpblendw_mm256_blend_epi16 而不是blendv，因为该指令在 Intel CPU 上仅在 1 uop 而不是 2 中运行。它还允许您使用 AVX512BW vpermw，这是当前 Skylake-AVX512 CPU 中可用的 16 位 shuffle，而不是可能更慢的 AVX512VBMI vpermb。

或者实际上，它可能会让您使用 vpmovzxwd (_mm256_cvtepu16_epi32) 将 16 位元素零扩展为 32 位，作为车道交叉洗牌。然后与dst 混合。

【讨论】：

谢谢你这么详细的回答，对我帮助很大。

@Thorgas：感谢您的反馈，让我知道它对初学者实际上很有用。你迷路的任何部分可以更清楚地表达吗？我添加了指向AVX tag wiki 和sse 的链接。 （我特意写了这个答案的前半部分作为 AVX shuffle 的一般指南，因此希望它对未来的其他 shuffle 的读者有用。）