SIMD 在最后一个峰之后搜索谷

SIMD search for trough after the last peak

我需要找到比最后一个滚动最大峰值低 X 或更多 % 的值的索引。

峰值是一个数组 (highs) 中元素的滚动最大值,而值在另一个数组 (lows) 中。数组具有相同的长度,并且值保证 <= peaks 数组中的对应项,没有 0、NAN 或无穷大元素。 since 保证小于 till.

迭代实现很简单:

inline
size_t trail_max_intern(double *highs,
        double *lows,
        double max,
        double trail,
        size_t since,
        size_t till)
{
    for (; since < till; ++since) {
        if (max < highs[since]) {
            max = highs[since];
        }

        if (lows[since] / max <= trail) {
            break;
        }
    }

    return since;
}

size_t trail_max_iter(double *highs, double *lows, double trail, size_t since, size_t till)
{
    trail = 1 - trail;
    return trail_max_intern(highs, lows, highs[since], trail, since, till);
}

这基本上是一个线性搜索,条件有所修改。由于任务上下文(任意 since、till 和 trail 值),无法使用其他结构或算法。

为了获得一些加速,我想使用 AVX2 矢量化扩展,看看会发生什么。我的结果如下所示:

size_t trail_max_vec(double *highs, double *lows, double trail, size_t since, size_t till)
{
    double max = highs[since];
    trail = 1 - trail;

    if (till - since > 4) {
        __m256d maxv = _mm256_set1_pd(max);
        __m256d trailv = _mm256_set1_pd(trail);

        for (size_t last = till & ~3; since < last; since += 4) {
            __m256d chunk = _mm256_loadu_pd(highs + since); // load peak block

            // peak rolling maximum computation
            maxv = _mm256_max_pd(maxv, chunk);
            // propagating the maximum value to places 2, 3, and 4
            maxv = _mm256_max_pd(maxv, _mm256_permute4x64_pd(maxv, 0b10010000));
            maxv = _mm256_max_pd(maxv, _mm256_permute4x64_pd(maxv, 0b01000000));
            maxv = _mm256_max_pd(maxv, _mm256_permute4x64_pd(maxv, 0b00000000));

            // divide lows by rolling maximum
            __m256d res = _mm256_div_pd(_mm256_loadu_pd(lows + since), maxv);
            // and if it is lower than the given fraction return its index
            int found = _mm256_movemask_pd(_mm256_cmp_pd(res, trailv, _CMP_LE_OQ));
            if (found) {
                return since + __builtin_ctz(found);
            }

            maxv = _mm256_set1_pd(maxv[3]);
        }

        max = maxv[3];
    }

    // make sure trailing elements are seen
    return trail_max_intern(highs, lows, max, trail, since, till);
}

它产生了正确的结果,但它比迭代版本慢 ~2 倍。我肯定做错了什么,但我不知道是什么。

所以我的问题是,我的方法有什么问题,我该如何解决?

P. S. 带有基准的完整源代码可在 https://godbolt.org/z/e5YrTo

获得

这会创建一个很长的依赖链:

        // peak rolling maximum computation
        maxv = _mm256_max_pd(maxv, chunk);
        // propagating the maximum value to places 2, 3, and 4
        maxv = _mm256_max_pd(maxv, _mm256_permute4x64_pd(maxv, 0b10010000));
        maxv = _mm256_max_pd(maxv, _mm256_permute4x64_pd(maxv, 0b01000000));
        maxv = _mm256_max_pd(maxv, _mm256_permute4x64_pd(maxv, 0b00000000));

        // ...

        maxv = _mm256_set1_pd(maxv[3]);

即,您有 8 条指令,全部取决于前一条指令的结果。假设每个都有 3 个周期的延迟,仅此一项就需要 4 个元素的 24 个周期(其余操作可能会在该时间内发生,特别是如果您比较 lows[since] <= trail * max——假设 max > 0)。

为了减少依赖链,您应该首先计算 chunk 内的“局部”滚动最大值,然后计算 maxv 的最大值:

        chunk = _mm256_max_pd(chunk, _mm256_movedup_pd(chunk));                  // [c0, c0c1, c2, c2c3]
        chunk = _mm256_max_pd(chunk, _mm256_permute4x64_pd(chunk, 0b01010100));  // [c0, c0c1, c0c1c2, c0c1c2c3]

        __m256d max_local = _mm256_max_pd(maxv, chunk);

要计算下一个 maxv,您可以广播 max_local[3](总延迟约 6 个周期),或者首先广播 chunk[3] 并使用 maxv(这将在依赖链中只留下一个 maxpd,并且您显然会受到吞吐量的限制)。在 Godbolt 上对此进行基准测试会产生很多噪音,以决定哪个更适合您的情况。

此外,您可以考虑处理更大的块(即加载两个连续的 __m256ds,为这些计算区域设置滚动最大值等)

特别是在标量版本中,我会在找到一个新的最大值时计算一次max * trail,然后另一个条件只是lows[since] <= threshold,不涉及任何计算。
这甚至比用 max * (1./trail) 替换 max / trail 更好(在评论中建议)。除法的延迟更高,如果超过除法器单元的有限吞吐量 (which is worst for 256-bit double vectors) 可能会成为瓶颈。

OTOH,GCC 正在将您当前的 if 优化为 maxsd,因此 new-max 是无分支的。另一种选择是 if (lows[since] * (1./trail) <= max) 欺骗编译器在每次迭代时使用 maxsdmulsd,而不是有条件地跳过 if (lows[since] <= max * trail) 中的乘法。但是对于您在 Godbolt 上的测试微基准测试,使用 threshold 似乎是最好的,即使这意味着更多的分支。 (但在常见情况下每次迭代的工作量较少,尽管 GCC 没有展开并且它是一个小测试,其数据可能太“简单”。)


在无分支 SIMD 中这没有帮助1 因为无论如何你每次都必须做所有的工作,但是 如果一个新的 max是罕见的(例如 highs[] 大且均匀分布)可能值得对此进行分支,以便在常见情况下减少 很多 洗牌和最大化工作。 特别是如果你在 CPU 上使用超线程/SMT 调整到 运行;另一个逻辑核心的微指令可以在核心从分支预测错误中恢复时保持忙碌。 (但是,如果新的最大值非常普遍,比如 highs 平均增加,而不是均匀增加,这将很糟糕。)

注 1:假设您已经在关注@chtz 的回答,其中展示了如何使大部分工作远离关键路径(循环携带的依赖链)延迟瓶颈。


常见/快速的情况将通过向量进行,只需检查两个条件highs > maxlows <= thresh)。如果其中任何一个为真,您 运行“昂贵的”代码在向量中执行正确的扫描顺序,然后检查阈值。您仍然需要正确处理 high[1] 是新最大值并且 lows[1..3] 低于新阈值的情况,即使它们低于旧最大值.还有 lows[0] 低于旧阈值但不低于新阈值的情况。当然,一个向量中的多个新最大值也是可能的。

        __m256d maxv = _mm256_set1_pd(max);
        __m256d thresh = _mm256_mul_pd(maxv, _mm256_set1_pd(trail));
        for() {
            __m256d vhi = _mm256_loadu_pd(highs + since);
            __m256d vlo = _mm256_loadu_pd(lows + since);
    
            __m256d vtrough = _mm256_cmp_pd(vlo, thresh, _CMP_LE_OQ);  // possible candidate
            __m256d newmax = _mm256_cmp_pd(vhi, maxv, _CMP_GT_OQ);     // definite new max

            // __m256d special_case = _mm256_or_pd(vtrough, newmax);
            // The special case needs trough separately, and it's only slightly more expensive to movmskpd twice.
            // AVX512 could use kortest between two compare results.

            unsigned trough = _mm256_movemask_pd(vtrough);
            // add/jnz can macro-fuse on Intel, unlike or/jnz, so use that to check for either one being non-zero.  AMD Zen doesn't care.
            if (trough + _mm256_movemask_pd(newmax)) {
                unsigned trough = _mm256_movemask_pd(vtrough);
                if (trough) {
                    // ... full work to verify or reject the candidate, hopefully rare.
                    // Perhaps extract this to a helper function
                    int found = ...;
                    // This becomes trivial (found = trough) in the no-new-max case, but that may be rare enough not to be worth special-casing here.
                    if (found) {
                        return since + __builtin_ctz(found);
                    }
                    maxv = _mm256_permute4x64_pd(maxv, _MM_SHUFFLE(3,3,3,3));
                } else {
                    // just a new max, no possible trough even with the lowest threshold
                    // horizontal max-broadcast of the current chunk, replacing old maxv (no data dependency on it)
                    maxv = _mm256_max_pd(vhi, _mm256_shuffle_pd(vhi,vhi, _MM_SHUFFLE(2,3,0,1)));   // in-lane swap pairs so both hold max of that 128-bit chunk
                    maxv = _mm256_max_pd(maxv, _mm256_permute4x64_pd(maxv, _MM_SHUFFLE(1,0,3,2)));    // swap low and high halves.  vperm2f128 would also work, but slower on Zen1.
                }

                thresh = _mm256_mul_pd(maxv, _mm256_set1_pd(trail));
            }
        }

可能不值得进一步对 no-newmax 情况进行特殊处理,以仅使用当前的 maxv,从而使用您已经计算的当前 trough,除非您希望找到一个低谷,如果不超过一个新的最大值(如果这是一个逐渐变化的单个数组(不是分开的高和低),那么 很常见:很难找到一个包含新最大值 的单个向量和一个低谷。)

maxv = _mm256_set1_pd(maxv[3]); 可能 编译效率高,但有些编译器在发明随机播放以实现 set1 时比其他编译器更笨。此外,vector[] 特定于 GNU C 实现 (gcc/clang) 定义 __m256d 的方式,并且不可移植。
使用maxv = _mm256_permute4x64_pd(maxv, _MM_SHUFFLE(3,3,3,3));vpermpd ymm, ymm, imm8指令广播高位元素

只是为了看看它是否编译,以及速度的粗略想法,我把它放在 Godbolt 上(int found = trough; 错误处理了有新的最大值和可能的谷值候选的情况).
在 Godbolt 的 Skylake 服务器 CPUs 上,显然基准测试结果因 AWS 上的负载而产生噪声,但矢量化版本在 100 次重复时从 33 毫秒到 55 毫秒不等,而标量版本约为 95 毫秒。在优化标量版本以使用 lows[since] <= thresh 并在 new-max 情况下更新 thresh = max*trail; 之前,标量的时间从 ~120 毫秒反弹到 127 毫秒。 (可能 CPU 频率和其他预热效应使得第一个测试在如此短的时间间隔内变化更大。)

我使用 -march=haswell 而不是 -mavx2 因为 GCC 的默认调整设置对于现代 CPUs 具有 256 位可能未对齐的负载很差(适用于 Sandybridge,不适用于 Haswell 或禅 2): .

循环的实际 asm 像我预期的那样高效:

# gcc -O3 -march=haswell, inner loop
        vmovupd ymm0, YMMWORD PTR [rdi+rdx*8]
        vmovupd ymm7, YMMWORD PTR [rsi+rdx*8]
        vcmppd  ymm3, ymm0, ymm2, 30
        vcmppd  ymm1, ymm7, ymm4, 18
        vmovmskpd       ecx, ymm3
        vmovmskpd       r8d, ymm1
        add     ecx, r8d                # destroys movemask(newmax) and sets FLAGS
        je      .L16

        test    r8d, r8d
        jne     .L36                    # jump if trough, else fall through to new max.  In this hacky version, finding non-zero trough always exits the loop because it ignores the possibility of also being a new max.  The asm will *probably* look similar once handled properly, though.

        vshufpd ymm1, ymm0, ymm0, 1
        vmaxpd  ymm0, ymm0, ymm1
        vpermpd ymm2, ymm0, 78
        vmaxpd  ymm2, ymm0, ymm2       # new max broadcast into YMM2
        vmulpd  ymm4, ymm2, ymm6       # update thresh from max

.L16:
        add     rdx, 4
        cmp     r9, rdx
        ja      .L19

请注意,常见情况(只是比较)没有循环携带的数据依赖性(指针增量除外)。

new-max 案例也没有对旧 max 的数据依赖性,只有控制依赖性(如果我们看到可预测的 运行 向量一个新的最大值,但不是一个低谷候选者。)

所以处理槽候选的复杂案例代码不必太担心数据依赖的长度,尽管新的 max 准备得越早,乱序执行程序就可以越早开始检查以后迭代的分支条件。


另请注意,如果需要,您可以使用 SIMD 安全地执行最后的 0..2 迭代,因为重做您已经看过的元素是安全的。 max 是单调递增的(因此阈值也是如此)所以你永远不会接受你不应该接受的 lows[]。所以你可以做一个最终的向量,它以你的范围的最后一个元素结束,只要总大小 >= 4 这样你就不会读取数组之外的内容。

在这种情况下,它可能不值得(因为每个 SIMD 迭代比标量昂贵得多),除了可能只是快速检查一下是否有可能的新最大值 and/or 低候选在 运行ning 标量清理之前。