AVX2 收集加载两个整数的结构

AVX2 gather load a struct of two ints

我目前正在尝试实现我现有的一些标量代码的 AVX2 版本 (Haswell CPU)。它实现了这样的步骤:

struct entry {
  uint32_t low, high;
};

// both filled with "random" data in previous loops
std::vector<entry> table;
std::vector<int>   queue;  // this is strictly increasing but
                           // without a constant delta

for (auto index : queue) {
  auto v = table[index];
  uint32_t rank = v.high + __builtin_popcount(_bzhi_u32(v.low, index % 32));
  use_rank(rank); // contains a lot of integer operations which nicely map to avx2
}

我用 2 个收集指令实现了这个,每个指令都像这样加载一个 int32:

__m256iv_low  = _mm256_i32gather_epi32 (reinterpret_cast<int *>(table.data()) + 0, index, 8);
__m256i v_high = _mm256_i32gather_epi32 (reinterpret_cast<int *>(table.data()) + 1, index, 8);

有没有更快的方法来加载这些值?我考虑过使用 2 个 64 位负载(它只发出一半的读取量 => 执行端口的流量更少)然后洗牌结果向量以获得 v_low 和 v_high 例如但遗憾的是据我所知,大多数洗牌功能只允许单独洗牌 128 位。

为 Paul R 编辑: 此代码是使用我在压缩算法中使用的 Burrows Wheeler 变换的子串枚举例程的一部分。 table 包含位向量上的排名数据。高位部分包含先前条目中的 1 的数量,低位部分被屏蔽掉并进行弹出计数,然后添加以获得给定索引前面的设置位的最终数量。之后发生了更多的计算,幸运的是可以很好地并行化。

队列中的增量在开始和结束时非常高(由于算法的性质)。这导致了很多缓存未命中,这也是我使用轮班从 SoA 切换到 AoS 以减少标量代码中加载端口压力的原因。

使用 SoA 也会产生相同的独立收集指令,但会使访问的缓存行数量加倍。

编辑(部分回答): 我尝试使用两个 _mm_i32gather_epi64 来减少一半的内存访问次数(因此循环数,请参见 here)。

__m256i index; // contains the indices
__m128i low = _mm256_extractf128_si256(index, 0);
__m128i high = _mm256_extractf128_si256(index, 1);
__m256i v_part1 = _mm256_i32gather_epi64(reinterpret_cast<long long int*>(table.data()), low , 8);
__m256i v_part2 = _mm256_i32gather_epi64(reinterpret_cast<long long int*>(table.data()), high, 8);

将我的数据加载到两个 ymm 寄存器中,这种格式(无 c++):

register v_part1:
[v[0].low][v[0].high][v[1].low][v[1].high][v[2].low][v[2].high][v[3].low][v[3].high]
register v_part2:
[v[4].low][v[4].high][v[5].low][v[5].high][v[6].low][v[6].high][v[7].low][v[7].high]

有没有一种有效的方法来交错它们以获得原始格式:

register v_low:
[v[0].low][v[1].low][v[2].low][v[3].low][v[4].low][v[5].low][v[6].low][v[7].low]
register v_high:
[v[0].high][v[1].high][v[2].high][v[3].high][v[4].high][v[5].high][v[6].high][v[7].high]

我自己找到了使用 5 条指令对值重新排序的方法:

// this results in [01][45][23][67] when gathering
index = _mm256_permute4x64_epi64(index, _MM_SHUFFLE(3,1,2,0));

// gather the values
__m256i v_part1 = _mm256_i32gather_epi64(i, _mm256_extractf128_si256(index, 0), 8);
__m256i v_part2 = _mm256_i32gather_epi64(i, _mm256_extractf128_si256(index, 1), 8);

// seperates low and high values
v_part1 = _mm256_shuffle_epi32(v_part1, _MM_SHUFFLE(3,1,2,0));
v_part2 = _mm256_shuffle_epi32(v_part2, _MM_SHUFFLE(3,1,2,0));

// unpack merges lows and highs: [01][23][45][56]
o1 = _mm256_unpackhi_epi64(v_part1, v_part2);
o2 = _mm256_unpacklo_epi64(v_part1, v_part2);