c++ - 使用 SIMD 解决循环数据依赖性 - 在 sgn 值的 int8_t 数组中查找 -1 和 +1 之间的转换

Question

我尝试实现性能提升，并在 SIMD 方面取得了一些良好的经验。到目前为止，我正在使用 OMP，并希望使用内在函数进一步提高我的技能。

在以下场景中，由于元素 n+1 测试所需的 last_value 的数据依赖性，我未能改进(甚至矢量化)。

环境是具有 AVX2 的 x64，因此想要找到一种对这样的函数进行矢量化和 SIMDfy 的方法。

inline static size_t get_indices_branched(size_t* _vResultIndices, size_t _size, const int8_t* _data) {
    size_t index = 0;
    int8_t last_value = 0;
    for (size_t i = 0; i < _size; ++i) {
        if ((_data[i] != 0) && (_data[i] != last_value)) {
            // add to _vResultIndices
            _vResultIndices[index] = i;
            last_value = _data[i];
            ++index;
        }
    }
    return index;
}

输入是一个有符号 1 字节值的数组。每个元素都是<0,1,-1>之一。 输出是输入值(或指针)的索引数组，表示更改为 1 或 -1。

输入/输出示例

in: { 0,0,1,0,1,1,-1,1, 0,-1,-1,1,0,0,1,1, 1,0,-1,0,0,0,0,0, 0,1,1,1,-1,-1,0,0, ... }
out { 2,6,7,9,11,18,25,28, ... }

我的第一次尝试是尝试各种无分支版本，并通过比较汇编输出来查看自动矢量化或 OMP 是否能够将其转换为 SIMDish 代码。

示例尝试

int8_t* rgLast = (int8_t*)alloca((_size + 1) * sizeof(int8_t));
rgLast[0] = 0;

#pragma omp simd safelen(1)
for (size_t i = 0; i < _size; ++i) {
    bool b = (_data[i] != 0) & (_data[i] != rgLast[i]);
    _vResultIndices[index] = i;
    rgLast[i + 1] = (b * _data[i]) + (!b * rgLast[i]);
    index += b;
}

由于没有实验产生 SIMD 输出，因此我开始尝试内在函数，目标是将条件部分转换为掩码。

对于 != 0 部分来说非常简单:

__m256i* vData = (__m256i*)(_data);
__m256i vHasSignal = _mm256_cmpeq_epi8(vData[i], _mm256_set1_epi8(0)); // elmiminate 0's

测试“最后一次翻转”的条件方面我还没有找到方法。

为了解决以下输出打包问题，我假设 AVX2 what is the most efficient way to pack left based on a mask?可以工作。

更新1

深入研究这个主题就会发现，分离 1/-1 并去掉 0 是有益的。 幸运的是，就我而言，我可以直接从预处理中获取它，并使用 _mm256_xor_si256 跳过处理到 <1,0,-1>，例如，将 2 个输入 vector 分隔为 gt0(全部为 1) ) 和 lt0(全部为 -1)。这还允许将数据打包得更紧 4 倍。

我可能想要结束这样的过程 现在的挑战是如何根据 gt0 和 lt0 掩码创建过渡掩码。

更新2

显然，一种将 1 和 -1 分成 2 个流的方法(参见答案如何)，在访问元素以进行交替扫描时引入了依赖关系: How to efficiently scan 2 bit masks alternating each iteration

按照 @aqrit 使用方法创建过渡蒙版
转换掩码 = ((~lt + gt) & lt) | ((~gt + lt) & gt) 是可能的。尽管这增加了相当多的指令，但它似乎是消除数据依赖性的有益权衡。我假设寄存器越大增益就会增加(可能取决于芯片)。

更新3

通过矢量化转换掩码 = ((~lt + gt) & lt) | ((~gt + lt) & gt) 我可以编译这个输出

vmovdqu     ymm5,ymmword ptr transition_mask[rax]  
vmovdqu     ymm4,ymm5  
vpandn      ymm0,ymm5,ymm6  
vpaddb      ymm1,ymm0,ymm5  
vpand       ymm3,ymm1,ymm5  
vpandn      ymm2,ymm5,ymm6  
vpaddb      ymm0,ymm2,ymm5  
vpand       ymm1,ymm0,ymm5  
vpor        ymm3,ymm1,ymm3  
vmovdqu     ymmword ptr transition_mask[rax],ymm3

乍一看，与后处理的潜在条件相关陷阱(垂直扫描+附加到输出)相比，它似乎更高效，尽管处理 2 个流而不是 1 个流似乎是正确且符合逻辑的。

这缺乏在每个周期生成初始状态的能力(从 0 到 1 或 -1 的转换)。 不确定是否有办法增强transition_mask生成“位旋转”，或者使用Soons在此处使用的自动初始_tzcnt_u32(mask0) > _tzcnt_u32(mask1):https://stackoverflow.com/a/70890642/18030502其中似乎包含一个分支。

结论

@aqrit 分享的方法使用改进的每个 block 加载的位旋转解决方案来查找转换，结果证明是运行时性能最高的。使用 tzcnt 和 blsr 的热内循环只有 9 个 asm 指令长(每 2 个找到的项目与其他方法进行比较)，如下所示

tzcnt       rax,rcx  
mov         qword ptr [rbx+rdx*8],rax  
blsr        rcx,rcx  
tzcnt       rax,rcx  
mov         qword ptr [rbx+rdx*8+8],rax  
blsr        rcx,rcx  
add         rdx,2  
cmp         rdx,r8  
jl          main+2580h (...)  

Answer 1

在 64 位 SIMD channel 之间串行传送状态比在 64 位通用寄存器 (gpr) 之间串行传送状态更昂贵。

实际上，查找表(或 SIMD 左填充)仅限于一次处理 8 个元素。如果数据平均每 64 个元素大约有 6 个保留元素，则左打包会浪费大量处理 (特别是如果我们正在收集偏移量而不执行收集操作)。如果位集很密集，则考虑转向查找表。

正如 @Snoots 所建议的，使用 SIMD 创建 64 位位集并使用 bitscan 内在函数查找所需集位的索引。

分支错误预测:

使用 gt 将大于 ( lt ) 和小于 ( transition_mask = ((~lt + gt) & lt) | ((~gt + lt) & gt) ) 位集压缩为单个位集或者来自 @FalkHüffner transition_mask = (lt ^ (lt - gt)) & (gt ^ (gt – lt)) 的简化.

状态是算术操作之一的进位/进位。我会小心使用 _subborrow_u64因为它是相当不常见的内在函数(并且在旧编译器上有错误)。

这使得唯一剩余的分支在位扫描操作上循环。必须提取所有设置的位..但我们可以展开操作并进行超调以使分支更可预测。超调量需要调整到预期的数据集。

未经测试。未检查汇编。

#include <immintrin.h>
#include <stdint.h>

static inline
uint64_t get_mask (int8_t* src, unsigned char* state) {
    __m256i src0 = _mm256_loadu_si256((__m256i*)(void*)src);
    __m256i src1 = _mm256_loadu_si256((__m256i*)(void*)&src[32]);

    uint64_t lt = (uint32_t)_mm256_movemask_epi8(src0) |
                    (((uint64_t)(uint32_t)_mm256_movemask_epi8(src1)) << 32);

    src0 = _mm256_cmpgt_epi8(src0, _mm256_setzero_si256());
    src1 = _mm256_cmpgt_epi8(src1, _mm256_setzero_si256());

    uint64_t gt = (uint32_t)_mm256_movemask_epi8(src0) |
                    (((uint64_t)(uint32_t)_mm256_movemask_epi8(src1)) << 32);

    // if borrow then greater-than span extends past the msb
    uint64_t m;
    unsigned char s = *state;
    *state = _subborrow_u64(s, lt, gt, (unsigned long long*)&m); // sbb
    return (m ^ lt) & ((gt - (lt + !s)) ^ gt);
}

static inline
size_t bitset_to_index (uint64_t* dst, uint64_t base, uint64_t mask) {
    int64_t cnt = _mm_popcnt_u64(mask);
    int64_t i = 0;
    do { // unroll to taste...
        dst[i + 0] = base + _tzcnt_u64(mask); mask = _blsr_u64(mask);
        dst[i + 1] = base + _tzcnt_u64(mask); mask = _blsr_u64(mask);
        dst[i + 2] = base + _tzcnt_u64(mask); mask = _blsr_u64(mask);
        dst[i + 3] = base + _tzcnt_u64(mask); mask = _blsr_u64(mask);
        i += 4;
    } while (i < cnt);
    return (size_t)cnt;
}

static
uint64_t* get_transition_indices (uint64_t* dst, int8_t* src, size_t len) {
    unsigned char state = 0; // in less-than span
    uint64_t base = 0; // offset into src array
    size_t end = len / 64;
    for (size_t i = 0; i < end; i++) {
        uint64_t mask = get_mask(src, &state);
        src += 64;
        dst += bitset_to_index(dst, base, mask);
        base += 64;
    }
    if (len % 64) {
        ; // todo: tail loop
    }
    return dst;
}