simd - 使用 AVX-512 模拟 64 字节的移位

Question

我的问题是上一个问题的延伸:Emulating shifts on 32 bytes with AVX .

如何使用 AVX-512 在 64 字节上实现类似的移位？具体来说，我应该如何实现:

__m512i _mm512_slli_si512(__m512i a, int imm8)

__m512i _mm512_srli_si512(__m512i a, int imm8)

对应于 SSE2 方法 _mm_slli_si128 和 _mm_srli_si128 .

Answer 1

这是使用临时数组的工作解决方案:

__m512i _mm512_slri_si512(__m512i a, size_t imm8)
{
    // set up temporary array and set upper half to zero 
    // (this needs to happen outside any critical loop)
    alignas(64) char temp[128];
    _mm512_store_si512(temp+64, _mm512_setzero_si512());

    // store input into lower half
    _mm512_store_si512(temp, a);

    // load shifted register
    return _mm512_loadu_si512(temp+imm8);
}

__m512i _mm512_slli_si512(__m512i a, size_t imm8)
{
    // set up temporary array and set lower half to zero 
    // (this needs to happen outside any critical loop)
    alignas(64) char temp[128];
    _mm512_store_si512(temp, _mm512_setzero_si512());

    // store input into upper half
    _mm512_store_si512(temp+64, a);

    // load shifted register
    return _mm512_loadu_si512(temp+(64-imm8));
}

如果 imm8，这也应该有效在编译时不知道，但它不做任何越界检查。
您实际上可以使用 3*64临时并在左移和右移方法之间共享它(两者也适用于负输入)。

当然，如果你在函数体之外共享一个临时的，你必须确保它不会被多个线程同时访问。

Godbolt-Link 与使用演示:https://godbolt.org/z/LSgeWZ

正如彼得指出的那样，这种存储加载技巧将导致所有带有 AVX512 的 CPU 上的存储转发停顿 .最有效的转发情况(约 6 个周期延迟)仅在所有加载字节都来自一个存储时才有效。如果负载超出了与它重叠的最新存储，则扫描存储缓冲区并在需要时从 L1d 缓存中合并字节会有额外的延迟(如约 16 个周期)。见 Can modern x86 implementations store-forward from more than one prior store?和 Agner Fog's microarch guide更多细节。这种额外的扫描过程可能会并行发生在多个负载上，并且至少不会拖延其他事情(例如正常的存储转发或管道的其余部分)，因此它可能不是吞吐量问题。

如果您想要相同数据的多个移位偏移，则以不同的对齐方式进行一次存储和多次重新加载应该是好的。

但是，如果延迟是您的主要问题，您应该尝试基于 valignd 的解决方案。 (此外，如果您想移动 4 个字节的倍数，这显然是一个更简单的解决方案)。或者对于恒定移位计数，vpermw 的矢量控制可以工作。

为了完整起见，这里是一个基于 valignd 的版本和 valignr用于从 0 到 64 的转换，在编译时已知(使用 C++17 - 但您可以轻松避免 if constexpr 这只是因为 static_assert 在这里)。您可以传递第二个寄存器(即，如果它跨 channel 对齐，它的行为就像 valignr 将行为一样)，而不是移入零。

template<int N>
__m512i shift_right(__m512i a, __m512i carry = _mm512_setzero_si512())
{
  static_assert(0 <= N && N <= 64);
  if constexpr(N   == 0) return a;
  if constexpr(N   ==64) return carry;
  if constexpr(N%4 == 0) return _mm512_alignr_epi32(carry, a, N / 4);
  else
  {
    __m512i a0 = shift_right< (N/16 + 1)*16>(a, carry);  // 16, 32, 48, 64
    __m512i a1 = shift_right< (N/16    )*16>(a, carry);  //  0, 16, 32, 48
    return _mm512_alignr_epi8(a0, a1, N % 16);
  }
}

template<int N>
__m512i shift_left(__m512i a, __m512i carry = _mm512_setzero_si512())
{
  return shift_right<64-N>(carry, a);
}

这是一个带有一些示例程序集的 Godbolt 链接以及每个可能的输出 shift_right操作:https://godbolt.org/z/xmKJvA

GCC 忠实地将其翻译成 valignd和 valignr说明——但可能会做不必要的 vpxor指令(例如在 shiftleft_49 示例中)，Clang 做了一些疯狂的替换(不过不确定它们是否真的有所作为)。

该代码可以扩展为移位任意寄存器序列(始终携带来自前一个寄存器的字节)。