我正在尝试将一些标量代码(下面的calc_offsets)转换为 AVX2 等效代码。它需要一系列“计数”并生成一个偏移位置表，从某个提供的基值开始。
我尝试将其转换为 AVX2 (avx2_calc_offsets)，我认为这是正确的，似乎速度大约是简单数组方法的一半。这是将较大的热点(瓶颈)代码部分转换为 AVX2 指令的努力的一部分，我希望将偏移量进一步处理为 vector 。对于此类操作，我希望避免在 AVX2 和标量代码之间跳转。
提供了一些示例和简单的基准测试代码。阵列版本的运行时间约为 2.15 秒，AVX2 版本的运行时间约为 4.41 秒(在 Ryzen Zen v1 上)。
是否有更好的方法使用 AVX2 来使此操作更快？我需要考虑较旧的 AVX2 CPU，例如 Haswell 和原始 Ryzen 系列。
#include <immintrin.h>
#include <inttypes.h>
#include <stdio.h>
typedef uint32_t u32;
typedef uint64_t u64;
void calc_offsets (const u32 base, const u32 *counts, u32 *offsets)
{
offsets[0] = base;
offsets[1] = offsets[0] + counts[0];
offsets[2] = offsets[1] + counts[1];
offsets[3] = offsets[2] + counts[2];
offsets[4] = offsets[3] + counts[3];
offsets[5] = offsets[4] + counts[4];
offsets[6] = offsets[5] + counts[5];
offsets[7] = offsets[6] + counts[6];
}
__m256i avx2_calc_offsets (const u32 base, const __m256i counts)
{
const __m256i shuff = _mm256_set_epi32 (6, 5, 4, 3, 2, 1, 0, 7);
__m256i v, t;
// shift whole vector `v` 4 bytes left and insert `base`
v = _mm256_permutevar8x32_epi32 (counts, shuff);
v = _mm256_insert_epi32 (v, base, 0);
// accumulate running total within 128-bit sub-lanes
v = _mm256_add_epi32 (v, _mm256_slli_si256 (v, 4));
v = _mm256_add_epi32 (v, _mm256_slli_si256 (v, 8));
// add highest value in right-hand lane to each value in left
t = _mm256_set1_epi32 (_mm256_extract_epi32 (v, 3));
v = _mm256_blend_epi32 (_mm256_add_epi32 (v, t), v, 0x0F);
return v;
}
void main()
{
u32 base = 900000000;
u32 counts[8] = { 5, 50, 500, 5000, 50000, 500000, 5000000, 50000000 };
u32 offsets[8];
calc_offsets (base, &counts[0], &offsets[0]);
printf ("calc_offsets: ");
for (int i = 0; i < 8; i++) printf (" %u", offsets[i]);
printf ("\n-----\n");
__m256i v, t;
v = _mm256_loadu_si256 ((__m256i *) &counts[0]);
t = avx2_calc_offsets (base, v);
_mm256_storeu_si256 ((__m256i *) &offsets[0], t);
printf ("avx2_calc_offsets: ");
for (int i = 0; i < 8; i++) printf (" %u", offsets[i]);
printf ("\n-----\n");
// --- benchmarking ---
#define ITERS 1000000000
// uncomment to benchmark AVX2 version
// #define AVX2_BENCH
#ifdef AVX2_BENCH
// benchmark AVX2 version
for (u64 i = 0; i < ITERS; i++) {
v = avx2_calc_offsets (base, v);
}
_mm256_storeu_si256 ((__m256i *) &offsets[0], v);
#else
// benchmark array version
u32 *c = &counts[0];
u32 *o = &offsets[0];
for (u64 i = 0; i < ITERS; i++) {
calc_offsets (base, c, o);
// feedback results to prevent optimizer 'cleverness'
u32 *tmp = c;
c = o;
o = tmp;
}
#endif
printf ("offsets after benchmark: ");
for (int i = 0; i < 8; i++) printf (" %u", offsets[i]);
printf ("\n-----\n");
}
我正在使用 gcc -O2 -mavx2 ... 进行构建。 Godbolt link .
最佳答案
消除前期_mm256_permutevar8x32_epi32 (vpermd)似乎在这里产生了巨大的差异。这可能是因为它的延迟很大(Ryzen 上有 8 个周期？)以及所有后续指令对其的直接依赖。
我不是预先输入基值，而是在加法过程中将其与在 128 位 channel 之间携带前缀和的相结合。
__m256i avx2_calc_offsets_2 (const u32 base, const __m256i counts)
{
__m256i b, t, v;
v = counts;
// accumulate running totals within 128-bit sub-lanes
v = _mm256_add_epi32 (v, _mm256_slli_si256 (v, 4));
v = _mm256_add_epi32 (v, _mm256_slli_si256 (v, 8));
// extract highest value in right-hand lane and combine with base offset
t = _mm256_set1_epi32 (_mm256_extract_epi32 (v, 3));
b = _mm256_set1_epi32 (base);
t = _mm256_blend_epi32 (_mm256_add_epi32 (b, t), b, 0x0F);
// combine with shifted running totals
v = _mm256_add_epi32 (_mm256_slli_si256 (v, 4), t);
return v;
}
两个版本之间的组装比较:
avx2_calc_offsets:
vmovdqa ymm1, YMMWORD PTR .LC0[rip]
vpermd ymm0, ymm1, ymm0
vpinsrd xmm1, xmm0, edi, 0
vinserti128 ymm0, ymm0, xmm1, 0x0
vpslldq ymm1, ymm0, 4
vpaddd ymm0, ymm0, ymm1
vpslldq ymm1, ymm0, 8
vpaddd ymm0, ymm0, ymm1
vpsrldq xmm1, xmm0, 12
vpbroadcastd ymm1, xmm1
vpaddd ymm1, ymm1, ymm0
vpblendd ymm0, ymm1, ymm0, 15
ret
avx2_calc_offsets_2:
vpslldq ymm1, ymm0, 4
vmovd xmm2, edi
vpaddd ymm1, ymm1, ymm0
vpbroadcastd ymm2, xmm2
vpslldq ymm0, ymm1, 8
vpaddd ymm1, ymm1, ymm0
vpsrldq xmm0, xmm1, 12
vpslldq ymm1, ymm1, 4
vpbroadcastd ymm0, xmm0
vpaddd ymm0, ymm2, ymm0
vpblendd ymm0, ymm0, ymm2, 15
vpaddd ymm0, ymm0, ymm1
ret
总的来说，指令数量相同，只是微指令/延迟成本较低。
使用 avx2_calc_offsets_2 的基准测试现在只需 2.7 秒即可运行，比之前的版本快了约 63%。
更新 1:GCC 将 avx2_calc_offsets_2 内联到基准循环中，进一步解释了性能的提高。作为彼得predicts ，与 _mm256_set1_epi32 (base) 对应的 vmovd/vpbroadcastd 指令确实被提升到循环外部的单个负载中。
循环组装:
...
// loop setup
vmovdqa ymm2, YMMWORD PTR .LC5[rip] // hoisted load of broadcasted base
vmovdqa ymm0, YMMWORD PTR [rbp-176]
vmovdqa ymm1, YMMWORD PTR [rbp-144]
mov eax, 1000000000
jmp .L10
.L17: // loop body
vpslldq ymm1, ymm0, 4
vpaddd ymm0, ymm0, ymm1
vpslldq ymm1, ymm0, 8
vpaddd ymm0, ymm0, ymm1
vpsrldq xmm1, xmm0, 12
vpslldq ymm0, ymm0, 4
vpbroadcastd ymm1, xmm1
vpaddd ymm1, ymm1, ymm2
vpblendd ymm1, ymm1, ymm2, 15
.L10: // loop entry
vpaddd ymm0, ymm1, ymm0
sub rax, 1
jne .L17
...
.LC5: // broadcasted `base`
.long 900000000
.long 900000000
.long 900000000
.long 900000000
.long 900000000
.long 900000000
.long 900000000
.long 900000000
更新 2:
重点关注内联情况，并将 _mm256_blend_epi32/vpblendd 替换为 __m128i 插入到归零的 __m256i 的高 channel 中>，然后添加到最终 vector 可进一步提高性能和代码大小 ( thanks Peter )。
__m256i avx2_calc_offsets_3 (const u32 base, const __m256i counts)
{
const __m256i z = _mm256_setzero_si256 ();
const __m256i b = _mm256_set1_epi32 (base);
__m256i v, t;
__m128i lo;
v = counts;
// accumulate running totals within 128-bit sub-lanes
v = _mm256_add_epi32 (v, _mm256_slli_si256 (v, 4));
v = _mm256_add_epi32 (v, _mm256_slli_si256 (v, 8));
// capture the max total in low-lane and broadcast into high-lane
lo = _mm_shuffle_epi32 (_mm256_castsi256_si128 (v), _MM_SHUFFLE (3, 3, 3, 3));
t = _mm256_inserti128_si256 (z, lo, 1);
// shift totals, add base and low-lane max
v = _mm256_slli_si256 (v, 4);
v = _mm256_add_epi32 (v, b);
v = _mm256_add_epi32 (v, t);
return v;
}
循环中内联版本的程序集现在如下所示:
// compiled with GCC version 10.3: gcc -O2 -mavx2 ...
// loop setup
vmovdqa ymm2, YMMWORD PTR .LC5[rip] // load broadcasted base
vmovdqa ymm0, YMMWORD PTR [rbp-176]
vmovdqa ymm1, YMMWORD PTR [rbp-144]
mov eax, 1000000000
vpxor xmm3, xmm3, xmm3
jmp .L12
.L20: // loop body
vpslldq ymm1, ymm0, 4
vpaddd ymm0, ymm0, ymm1
vpslldq ymm1, ymm0, 8
vpaddd ymm0, ymm0, ymm1
vpshufd xmm1, xmm0, 255
vpslldq ymm0, ymm0, 4
vinserti128 ymm1, ymm3, xmm1, 0x1
.L12: // loop entry
vpaddd ymm0, ymm0, ymm1
vpaddd ymm0, ymm0, ymm2
sub rax, 1
jne .L20
循环体减少到只有 9 个 vector 指令:)。
使用 -O3 时，GCC 中存在优化错误，其中在循环体末尾插入了无关的 vmovdqa ymm0, ymm1，从而使基准测试性能降低了几个百分点。 (至少对于 GCC 版本 11.x、10.x 和 9.x)。
更新 3:另一个轻微的性能提升。如果我们在 128 位插入之前使用 SSE/128 位指令添加低 channel 的最大总数，则 shorten the critical path for v允许更好地使用随机播放端口。
__m256i avx2_calc_offsets_4 (const u32 base, const __m256i counts)
{
const __m256i b = _mm256_set1_epi32 (base);
__m256i v, t;
__m128i lo;
v = counts;
// accumulate running totals within 128-bit sub-lanes
v = _mm256_add_epi32 (v, _mm256_slli_si256 (v, 4));
v = _mm256_add_epi32 (v, _mm256_slli_si256 (v, 8));
// capture the max total in low-lane, broadcast into high-lane and add to base
lo = _mm_shuffle_epi32 (_mm256_castsi256_si128 (v), _MM_SHUFFLE (3, 3, 3, 3));
lo = _mm_add_epi32 (_mm256_castsi256_si128 (b), lo);
t = _mm256_inserti128_si256 (b, lo, 1);
// shift totals, add base and low-lane max
v = _mm256_slli_si256 (v, 4);
v = _mm256_add_epi32 (v, t);
return v;
}
.L23: // loop body
vpslldq ymm1, ymm0, 4
vpaddd ymm0, ymm0, ymm1
vpslldq ymm1, ymm0, 8
vpaddd ymm0, ymm0, ymm1
vpshufd xmm2, xmm0, 255
vpslldq ymm1, ymm0, 4
.L14: // loop entry
vpaddd xmm0, xmm2, xmm3
vinserti128 ymm0, ymm4, xmm0, 0x1
vpaddd ymm0, ymm1, ymm0
sub rax, 1
jne .L23
在我(非专家)看来，这看起来相当理想，至少对于早期的 AVX2 芯片来说是这样。基准时间缩短至约 2.17 秒。
奇怪的是，如果我通过删除以前的函数定义之一来减少源代码的大小，GCC 10 和 11 就会变得有点困惑，并将 3 (!) 条额外的 vmovdqa 指令插入循环(Godbolt)。结果是我的基准测试速度下降了约 18%。 GCC 9.x 似乎不受影响。我不确定这里发生了什么，但这似乎是 GCC 优化器中一个非常讨厌的错误。我会尝试减少它并提交错误。
使用 avx2_calc_offsets_3 的基准测试现在的运行速度实际上与标量版本相同，这对我来说是一个胜利，因为它消除了出于性能原因跳转到标量代码的需要。
关于c - 在 __m256i vector 中水平累积运行总计(前缀总和)，我们在Stack Overflow上找到一个类似的问题： https://stackoverflow.com/questions/69694914/