使用 AVX-512 模拟 64 字节的移位
Emulating shifts on 64 bytes with AVX-512
我的问题是前一个问题的延伸: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
and _mm_srli_si128
。
这是一个使用临时数组的有效解决方案:
__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
正如 Peter 所指出的,这个存储加载技巧将导致所有带有 AVX512 的 CPU 上的存储转发停止。最有效的转发情况(约 6 个周期延迟)仅在所有加载字节都来自一个存储时才有效。如果负载超出了与其完全重叠的最新存储,则它有额外的延迟(如约 16 个周期)来扫描存储缓冲区,并在需要时从 L1d 缓存中合并字节。有关详细信息,请参阅 Can modern x86 implementations store-forward from more than one prior store? and Agner Fog's microarch guide。这种额外的扫描过程可能会同时发生在多个负载上,并且至少不会拖延其他事情(例如正常的存储转发或管道的其余部分),因此它可能不是吞吐量问题。
如果您想要同一数据的多个移位偏移量,一次存储和多次重新加载在不同的对齐方式应该很好。
但是如果延迟是您的主要问题,您应该尝试基于 valignd
的解决方案(另外,如果您想要移动 4 字节的倍数,这显然是一个更简单的解决方案)。或者对于恒定的班次计数,vpermw
的矢量控制可以工作。
为了完整性,这是一个基于 valignd
和 valignr
的版本,适用于从 0 到 64 的班次,在编译时已知(使用 C++17——但您可以轻松地避免 if constexpr
这只是因为 static_assert
)。您可以传递第二个寄存器,而不是移入零(即,如果它跨车道对齐,它的行为就像 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-link,其中包含一些示例程序集以及每个可能的 shift_right
操作的输出:https://godbolt.org/z/xmKJvA
GCC 忠实地将其转换为 valignd
和 valignr
指令——但可能会执行不必要的 vpxor
指令(例如在 shiftleft_49
示例中),Clang 会执行一些疯狂的替换(虽然不确定它们是否真的有所作为)。
可以扩展代码以移位任意序列的寄存器(始终携带来自前一个寄存器的字节)。
对于那些需要恰好移动64位的用户,可以使用直接在寄存器中工作的permute指令。对于 8 位的倍数的移位,您可以使用字节洗牌(请参阅 VPSHUFB
并查看转换函数,如果您正在处理浮点数,因为洗牌使用整数)。
这是一个移动 64 位的例子 ("SHR zmm1, 64")。掩码用于清除前 64 位。如果你想 ROR
喜欢的功能,你可以使用没有掩码的版本。请注意,也可以向左移动。只需根据需要更改索引即可。
#include <immintrin.h>
#include <iostream>
void show(char const * msg, double *v)
{
std::cout
<< msg
<< ": "
<< v[0]
<< " "
<< v[1]
<< " "
<< v[2]
<< " "
<< v[3]
<< " "
<< v[4]
<< " "
<< v[5]
<< " "
<< v[6]
<< " "
<< v[7]
<< "\n";
}
int main(int argc, char * argv[])
{
double v[8] = { 1., 2., 3., 4., 5., 6., 7., 8. };
double q[8] = {};
alignas(64) std::uint64_t indexes[8] = { 1, 2, 3, 4, 5, 6, 7, 0 };
show("init", v);
show("q", q);
// load
__m512d a(_mm512_loadu_pd(v));
__m512i i(_mm512_load_epi64(indexes));
// shift
//__m512d b(_mm512_permutex_pd(a, 0x39)); // can't cross between 4 low and 4 high with immediate
//__m512d b(_mm512_permutexvar_pd(i, a)); // ROR
__m512d b(_mm512_maskz_permutexvar_pd(0x7F, i, a)); // LSR on a double basis
// store
_mm512_storeu_pd(q, b);
show("shifted", q);
show("original", v);
}
完全优化的输出 (-O3) 将整个移位减少到 3 条指令(在输出中与其他指令混合):
96a: 62 f1 fd 48 6f 85 10 vmovdqa64 -0xf0(%rbp),%zmm0
971: ff ff ff
974: b8 7f 00 00 00 mov [=11=]x7f,%eax # mask
979: 48 8d 3d 10 04 00 00 lea 0x410(%rip),%rdi # d90 <_IO_stdin_used+0x10>
980: c5 f9 92 c8 kmovb %eax,%k1 # special k1 register
984: 4c 89 e6 mov %r12,%rsi
987: 62 f2 fd c9 16 85 d0 vpermpd -0x130(%rbp),%zmm0,%zmm0{%k1}{z} # "shift"
98e: fe ff ff
991: 62 f1 fd 48 11 45 fe vmovupd %zmm0,-0x80(%rbp)
在我的例子中,我想在循环中使用它,加载 (vmovdqa64
) 和存储 (vmovupd
) 将在循环之前和之后,在循环内部,它会非常快。 (在我需要保存结果之前,它需要这样旋转 4,400 次)。
正如Peter所指出的,我们也可以使用valignq
指令:
// this is in place of the permute, without the need for the indexes
__m512i b(_mm512_maskz_alignr_epi64(0xFF, _mm512_castpd_si512(a), _mm512_castpd_si512(a), 1));
结果是这样的一条指令:
979: 62 f1 fd 48 6f 85 d0 vmovdqa64 -0x130(%rbp),%zmm0
980: fe ff ff
983: 48 8d 75 80 lea -0x80(%rbp),%rsi
987: 48 8d 3d 02 04 00 00 lea 0x402(%rip),%rdi # d90 <_IO_stdin_used+0x10>
98e: 62 f3 fd 48 03 c0 01 valignq [=13=]x1,%zmm0,%zmm0,%zmm0
995: 62 f1 fd 48 11 45 fd vmovupd %zmm0,-0xc0(%rbp)
重要的一点是,使用较少的寄存器也更好,因为它增加了我们在寄存器中获得 100% 完全优化的机会,而不必使用内存(512 位传输到内存和从内存传输很多)。
我的问题是前一个问题的延伸: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
and _mm_srli_si128
。
这是一个使用临时数组的有效解决方案:
__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
正如 Peter 所指出的,这个存储加载技巧将导致所有带有 AVX512 的 CPU 上的存储转发停止。最有效的转发情况(约 6 个周期延迟)仅在所有加载字节都来自一个存储时才有效。如果负载超出了与其完全重叠的最新存储,则它有额外的延迟(如约 16 个周期)来扫描存储缓冲区,并在需要时从 L1d 缓存中合并字节。有关详细信息,请参阅 Can modern x86 implementations store-forward from more than one prior store? and Agner Fog's microarch guide。这种额外的扫描过程可能会同时发生在多个负载上,并且至少不会拖延其他事情(例如正常的存储转发或管道的其余部分),因此它可能不是吞吐量问题。
如果您想要同一数据的多个移位偏移量,一次存储和多次重新加载在不同的对齐方式应该很好。
但是如果延迟是您的主要问题,您应该尝试基于 valignd
的解决方案(另外,如果您想要移动 4 字节的倍数,这显然是一个更简单的解决方案)。或者对于恒定的班次计数,vpermw
的矢量控制可以工作。
为了完整性,这是一个基于 valignd
和 valignr
的版本,适用于从 0 到 64 的班次,在编译时已知(使用 C++17——但您可以轻松地避免 if constexpr
这只是因为 static_assert
)。您可以传递第二个寄存器,而不是移入零(即,如果它跨车道对齐,它的行为就像 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-link,其中包含一些示例程序集以及每个可能的 shift_right
操作的输出:https://godbolt.org/z/xmKJvA
GCC 忠实地将其转换为 valignd
和 valignr
指令——但可能会执行不必要的 vpxor
指令(例如在 shiftleft_49
示例中),Clang 会执行一些疯狂的替换(虽然不确定它们是否真的有所作为)。
可以扩展代码以移位任意序列的寄存器(始终携带来自前一个寄存器的字节)。
对于那些需要恰好移动64位的用户,可以使用直接在寄存器中工作的permute指令。对于 8 位的倍数的移位,您可以使用字节洗牌(请参阅 VPSHUFB
并查看转换函数,如果您正在处理浮点数,因为洗牌使用整数)。
这是一个移动 64 位的例子 ("SHR zmm1, 64")。掩码用于清除前 64 位。如果你想 ROR
喜欢的功能,你可以使用没有掩码的版本。请注意,也可以向左移动。只需根据需要更改索引即可。
#include <immintrin.h>
#include <iostream>
void show(char const * msg, double *v)
{
std::cout
<< msg
<< ": "
<< v[0]
<< " "
<< v[1]
<< " "
<< v[2]
<< " "
<< v[3]
<< " "
<< v[4]
<< " "
<< v[5]
<< " "
<< v[6]
<< " "
<< v[7]
<< "\n";
}
int main(int argc, char * argv[])
{
double v[8] = { 1., 2., 3., 4., 5., 6., 7., 8. };
double q[8] = {};
alignas(64) std::uint64_t indexes[8] = { 1, 2, 3, 4, 5, 6, 7, 0 };
show("init", v);
show("q", q);
// load
__m512d a(_mm512_loadu_pd(v));
__m512i i(_mm512_load_epi64(indexes));
// shift
//__m512d b(_mm512_permutex_pd(a, 0x39)); // can't cross between 4 low and 4 high with immediate
//__m512d b(_mm512_permutexvar_pd(i, a)); // ROR
__m512d b(_mm512_maskz_permutexvar_pd(0x7F, i, a)); // LSR on a double basis
// store
_mm512_storeu_pd(q, b);
show("shifted", q);
show("original", v);
}
完全优化的输出 (-O3) 将整个移位减少到 3 条指令(在输出中与其他指令混合):
96a: 62 f1 fd 48 6f 85 10 vmovdqa64 -0xf0(%rbp),%zmm0
971: ff ff ff
974: b8 7f 00 00 00 mov [=11=]x7f,%eax # mask
979: 48 8d 3d 10 04 00 00 lea 0x410(%rip),%rdi # d90 <_IO_stdin_used+0x10>
980: c5 f9 92 c8 kmovb %eax,%k1 # special k1 register
984: 4c 89 e6 mov %r12,%rsi
987: 62 f2 fd c9 16 85 d0 vpermpd -0x130(%rbp),%zmm0,%zmm0{%k1}{z} # "shift"
98e: fe ff ff
991: 62 f1 fd 48 11 45 fe vmovupd %zmm0,-0x80(%rbp)
在我的例子中,我想在循环中使用它,加载 (vmovdqa64
) 和存储 (vmovupd
) 将在循环之前和之后,在循环内部,它会非常快。 (在我需要保存结果之前,它需要这样旋转 4,400 次)。
正如Peter所指出的,我们也可以使用valignq
指令:
// this is in place of the permute, without the need for the indexes
__m512i b(_mm512_maskz_alignr_epi64(0xFF, _mm512_castpd_si512(a), _mm512_castpd_si512(a), 1));
结果是这样的一条指令:
979: 62 f1 fd 48 6f 85 d0 vmovdqa64 -0x130(%rbp),%zmm0
980: fe ff ff
983: 48 8d 75 80 lea -0x80(%rbp),%rsi
987: 48 8d 3d 02 04 00 00 lea 0x402(%rip),%rdi # d90 <_IO_stdin_used+0x10>
98e: 62 f3 fd 48 03 c0 01 valignq [=13=]x1,%zmm0,%zmm0,%zmm0
995: 62 f1 fd 48 11 45 fd vmovupd %zmm0,-0xc0(%rbp)
重要的一点是,使用较少的寄存器也更好,因为它增加了我们在寄存器中获得 100% 完全优化的机会,而不必使用内存(512 位传输到内存和从内存传输很多)。