在同一个 SIMD 寄存器中串行添加值

Posted 2023-02-16

【中文标题】在同一个 SIMD 寄存器中串行添加值

【英文标题】：Add values serially in the same SIMD register

【发布时间】：2021-09-24 15:04:50

【问题描述】：

我正在尝试将其转换为 AVX2：

// parallel arrays
int16_t* Nums = ...
int16_t* Capacities = ...
int** Data = ...

int* freePointer = ...

for (int i = 0; i < n; i++)

    if (Nums[i] == 0)
        Capacities[i] = 0;
    else
    
        Data[i] = freePointer;
        freePointer += Capacities[i];
    

但没有走得太远：

for (int i = 0; i < n; i += 4) // 4 as Data is 64 bits

    const __m256i nums = _mm256_loadu_si256((__m256i*)&Nums[i]);
    const __m256i bZeroes = _mm256_cmpeq_epi16(nums, ZEROES256);
    const __m256i capacities = _mm256_loadu_si256((__m256i*)&Capacities[i]);
    const __m256i zeroedCapacities = _mm256_andnot_si256(bZeroes, capacities);
    _mm256_storeu_si256((__m256i*)&Capacities[i], zeroedCapacities);

卡在else 分支，不确定如何将（前缀总和？...）Capacities 添加到freePointer 并将“串行”结果分配给同一 256 位 SIMD 寄存器中的Data . 我的术语可能是错误的，我希望代码能够理解我想要完成的任务。

lane0: freePointer
lane1: freePointer + Capacities[i + 0]
lane2: freePointer + Capacities[i + 0] + Capacities[i + 1]
lane3: freePointer + Capacities[i + 0] + Capacities[i + 1] + Capacities[i + 2]

基本上，如果可能的话，这就是我想用尽可能少的指令来做的事情。目标是 AVX2。

【问题讨论】：

这就是所谓的“前缀和”。 parallel prefix (cumulative) sum with SSE 显示 8x 32 位 FP 元素，或 4x 32 位 FP 元素与 SSE。您也许可以针对 64 位整数元素调整 shuffle，但 shuffle 延迟可能是个问题。 （有关 hsum 的链接，请参见 Fastest way to do horizontal SSE vector sum (or other reduction) 中的链接，这需要一些相同的改组。）

哦，您需要即时屏蔽传入的数据，但我认为您已经通过 cmpeq_epi16 / andnot 解决了这个问题。但是，是的，然后用vpmovzx 扩大到 64 位指针宽度我猜。 （如果您想要更好的数据密度，也可以使用 32 位数组偏移量。）

谢谢！所以你是说即使我知道怎么做，它也很有可能太慢了，对吗？

值得尝试矢量化，尤其是如果您可以将Data[] 更改为相对于数组的uint32_t 偏移量数组。 （或者如果指针是 32 位的，例如 gcc -mx32 或实际的 32 位模式）。但即使使用 64 位指针，您也可以使用快速 128 位随机播放以 4 个元素为一组进行一些前缀求和，并且仅在最后扩大到 64 位。

【参考方案1】：

你可以在这里找到很多细节：https://***.com/a/69452433/5021064

在这里你可以插入任何类型而不是 T 并查看生成的 x86 和 arm 的 asm