【发布时间】：2021-12-10 03:54:29
【问题描述】：
我正在尝试将一些标量代码（下面的calc_offsets）转换为 AVX2 等效项。它需要一系列“计数”并生成一个偏移位置表，从一些提供的基值开始。
我尝试将其转换为 AVX2 (avx2_calc_offsets)，我认为这是正确的，它的速度似乎是简单数组方法的一半左右。这是将更大的热部分（瓶颈）代码转换为 AVX2 指令的努力的一部分，我希望将偏移量作为向量进一步处理。对于这样的操作，我想避免在 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.
【问题讨论】：
-
看起来像 Prefix Sum，又名包容性扫描，又名累积和。加上
counts[8]中的一些固定偏移量，它们并没有真正改变依赖模式。有关 FP 版本（包括 AVX），请参阅 parallel prefix (cumulative) sum with SSE，但正如我在那里评论的那样，整数添加延迟较低意味着预期的加速可能较小。 -
从您的问题标题中，我预计它将是Fastest way to do horizontal SSE vector sum (or other reduction)，您只想要总数。可能想在标题中的某处放置“总和”或“前缀总和”。
-
@Peter 谢谢，如果我能在标量数组版本附近的某个地方得到它，即使稍微慢一点，我也会很高兴。主要目的是避免基于向量的基于数组的跳转多次。
标签： c vectorization x86-64 intrinsics avx2