如何在两个 AVX2 向量之间交换 128 位部分

Posted 2023-02-16

【中文标题】如何在两个 AVX2 向量之间交换 128 位部分

【英文标题】：How to swap 128-bit parts between two AVX2 vectors

【发布时间】：2020-05-18 09:44:20

【问题描述】：

问题：我有 4 个 256 位 AVX2 向量（A、B、C、D），我需要在它们各自的 128 位部分以及两个不同的向量之间执行交换操作.这是我需要做的转变。

             Original                      Transformed
    || Low Lane || High Lane||     || Low Lane || High Lane||
A = ||    L1    ||    H1    || = > ||    L1    ||    L2    ||
B = ||    L2    ||    H2    || = > ||    H1    ||    H2    ||
C = ||    L3    ||    H3    || = > ||    L3    ||    L4    ||
D = ||    L4    ||    H4    || = > ||    H3    ||    H4    ||

Visualization

基本上我需要将输出按以下顺序 L1、L2、L3、L4、H1、H2、H3、H4 存储到数组中。

我目前的解决方案是： 4x _mm256_blend_epi32（最坏情况：延迟 1，吞吐量 0.35） 4x _mm256_permute2x128_si256（最坏情况：延迟 3，吞吐量 1）

// (a, c) = block0, (b, d) = block1
a = Avx2.Permute2x128(a, a, 1);
var template = Avx2.Blend(a, b, 0b1111_0000); // H1 H2
a = Avx2.Blend(a, b, 0b0000_1111); // L2 l1
a = Avx2.Permute2x128(a, a, 1); // L1 l2
b = template;

c = Avx2.Permute2x128(c, c, 1);
template = Avx2.Blend(c, d, 0b1111_0000); // H3 H4
c = Avx2.Blend(c, d, 0b0000_1111);  // L4 L3
c = Avx2.Permute2x128(c, c, 1); // L3 l4
d = template;

// Store keystream into buffer (in corrected order = [block0, block1])
Avx2.Store(outputPtr, a);
Avx2.Store(outputPtr + Vector256<uint>.Count, c);
Avx2.Store(outputPtr + Vector256<uint>.Count * 2, b);
Avx2.Store(outputPtr + Vector256<uint>.Count * 3, d);

注意：如果您想知道的话，我正在使用 C#/NetCore 来执行 AVX2！随意使用 C/C++ 中的示例。

有没有更好或更有效的方法来做到这一点？

编辑

接受的答案为 C#

var tmp = Avx2.Permute2x128(a, b, 0x20);
b = Avx2.Permute2x128(a, b, 0x31);
a = tmp;
tmp = Avx2.Permute2x128(c, d, 0x20);
d = Avx2.Permute2x128(c, d, 0x31);
c = tmp;

【参考方案1】：

如果我对您的理解正确，我认为您可以在没有此 2x4 转置的混合说明的情况下逃脱，创建新变量来选择您想要的通道。比如：

__m256i a;    // L1 H1
__m256i b;    // L2 H2
__m256i c;    // L3 H3
__m256i d;    // L4 H4

__m256i A = _mm256_permute2x128_si256(a, b, 0x20);  // L1 L2
__m256i B = _mm256_permute2x128_si256(a, b, 0x31);  // H1 H2
__m256i C = _mm256_permute2x128_si256(c, d, 0x20);  // L3 L4
__m256i D = _mm256_permute2x128_si256(c, d, 0x31);  // H3 H4

vperm2i128 指令的 3 周期延迟仍然存在，但当数据跨越 128 位通道时，您总是会遇到这种情况。这 4 个 shuffle 是独立的，因此它们可以流水线化（ILP）； Intel 和 Zen 2 对 vperm2i128 (https://agner.org/optimize/, https://uops.info/) 有 1/clock 吞吐量。

如果幸运的话，编译器会将 L1、L2 和 L3、L4 洗牌优化为vinserti128，AMD Zen 1 运行效率更高（1 uop 而不是 8；车道交叉洗牌被分成多个 128 -bit uops。）

---

这 4 次 shuffle 需要 4 uop 用于 shuffle 端口（Intel 上的端口 5）； Intel 和 Zen2 对这些 shuffle 仅有 1/clock shuffle 吞吐量。如果这将成为您循环中的瓶颈，请考虑@chtz 的答案，该答案通过执行 2 次洗牌来排列需要移动的 4 个通道以准备廉价混合 (vpblendd)，从而花费更多的前端吞吐量。相关：What considerations go into predicting latency for operations on modern superscalar processors and how can I calculate them by hand?

【讨论】：

这是消除 4x 混合操作的绝佳解决方案！还在学习AVX2集..

我可以使用 _mm256_permute2x128_si256 代替 _mm256_permute2f128_ps 吗？还是有特定的原因？好的，它也适用于第一个变体......

不，您可以使用整数版本。我误读了您的原始示例，并没有发现您拥有整数数据的事实。我会编辑。

【参考方案2】：

您可以使用两个置换和 4 个混合进行操作，从而提供 2 个周期的绝对吞吐量：

void foo(
    __m256i a,    // L1 H1
    __m256i b,    // L2 H2
    __m256i c,    // L3 H3
    __m256i d,    // L4 H4
    __m256i* outputPtr
)

    // permute. Port usage: 1*p5, Latency 3 on both inputs
    __m256i BA = _mm256_permute2x128_si256(a, b, 0x21);  // H1 L2 
    __m256i DC = _mm256_permute2x128_si256(c, d, 0x21);  // H3 L4

    // blend. Port usage: 1*p015, Latency 1 on both inputs
    __m256i A = _mm256_blend_epi32(a, BA, 0xf0);  // L1 L2
    __m256i B = _mm256_blend_epi32(BA, b, 0xf0);  // H1 H2
    __m256i C = _mm256_blend_epi32(c, DC, 0xf0);  // L3 L4
    __m256i D = _mm256_blend_epi32(DC, d, 0xf0);  // H3 H4

    _mm256_store_si256(outputPtr+0, A);
    _mm256_store_si256(outputPtr+1, B);
    _mm256_store_si256(outputPtr+2, C);
    _mm256_store_si256(outputPtr+3, D);

但是，根据上下文（特别是如果 a、...、d 最初是从内存中读取的），使用 vmovdqu 和 vinserti128 指令和 @987654329 的序列可能会更好@内存操作数。您将有两倍的负载，但没有通道间延迟，也没有端口 5 上的瓶颈——关于延迟和端口使用，基于内存的 vinsert128 表现得像一个混合体。

【讨论】：

有趣，因为我对代码向量化很陌生，我不熟悉术语“端口”，你能详细说明一下吗？我正在使用 broutcast 从内存中读取。 uint* state = new uint[32] 0, 1, 2, 3, 20, 21, 22, 23, 4, 5, 6, 7, 24, 25, 26, 27, 8, 9, 10, 11, 28, 29, 30, 31, 12, 13, 14, 15, 32, 33, 34, 35 ; a = Avx2.BroadcastVector128ToVector256(state); b = Avx2.BroadcastVector128ToVector256(state + Vector128<uint>.Count); c = Avx2.BroadcastVector128ToVector256(state + Vector128<uint>.Count * 2); d = Avx2.BroadcastVector128ToVector256(state + Vector128<uint>.Count * 3);

我有理由使用“广播”，因为我正在同时计算 2 个 chacha20 密码块，只是为了澄清。另外，您如何计算绝对吞吐量？ giving an absolute throughput of 2 cycles

@xtremertx：我添加了指向 Jason 答案的链接以及 3 周期延迟数的来源。 “吞吐量”只对包括周边代码在内的整个区块很重要；如果您仍然不会在 shuffle-port 吞吐量上遇到瓶颈，而是在前端 uop 吞吐量上（或更糟糕的是延迟），那么在 Jason 的答案中使用 4 指令方式，而不是在此答案中使用 6 指令方式。例如如果您在循环中的这些 shuffle 步骤之间有很多 AND / OR / shift 工作，可能会针对更少的指令进行优化。

@xtremertx 广播负载是否直接在“转置”之前发生？或者中间是否有指示发生？另外，只是我正确理解了 C#-AVX 语法：广播后a=0,1,2,3, 0,1,2,3、b=20,21,22,23, 20,21,22,23 等等？

@chtz 很可能有多种解决方案如何使用 AVX2 对 chacha20 进行矢量化，但是这样的讨论需要一个新的话题。我正在使用关注paper-pdf，它更详细地解释了一些事情，并且在发布时它们的性能比铬项目更好。我基本上是在使用 AVX2 同时计算 2 个密钥流块（他们在论文中称之为双四轮）。