用于通用 SIMD（SSE、AVX、NEON）测试零匹配的高效 C 向量。（求FP最大绝对值和索引）

Question

我想看看是否可以编写一些可以高效编译的通用 SIMD 代码。主要用于 SSE、AVX 和 NEON。该问题的简化版本是：找到浮点数数组的最大绝对值并返回该值和索引。导致问题的是最后一部分，即最大值的索引。似乎没有一个很好的方法来编写具有分支的代码。

请参阅最后的更新，以获取使用一些建议答案的完成代码。

这是一个示例实现（godbolt上更完整的版本）：

#define VLEN 8
typedef float vNs __attribute__((vector_size(VLEN*sizeof(float))));
typedef int vNb __attribute__((vector_size(VLEN*sizeof(int))));
#define SWAP128 4,5,6,7, 0,1,2,3
#define SWAP64 2,3, 0,1,  6,7, 4,5
#define SWAP32 1, 0,  3, 2,  5, 4,  7, 6

static bool any(vNb x) {
    x = x | __builtin_shufflevector(x,x, SWAP128);
    x = x | __builtin_shufflevector(x,x, SWAP64);
    x = x | __builtin_shufflevector(x,x, SWAP32);
    return x[0];
}

float maxabs(float* __attribute__((aligned(32))) data, unsigned n, unsigned *index) {
    vNs max = {0,0,0,0,0,0,0,0};
    vNs tmax;
    unsigned imax = 0;
    for (unsigned i = 0 ; i < n; i += VLEN) {
        vNs t = *(vNs*)(data + i);
        t = -t < t ? t : -t;  // Absolute value
        vNb cmp = t > max;
        if (any(cmp)) {
            tmax = t; imax = i;
            // broadcast horizontal max of t into every element of max
            vNs tswap128 = __builtin_shufflevector(t,t, SWAP128);
            t = t < tswap128 ? tswap128 : t;
            vNs tswap64 = __builtin_shufflevector(t,t, SWAP64);
            t = t < tswap64 ? tswap64 : t;
            vNs tswap32 = __builtin_shufflevector(t,t, SWAP32);
            max = t < tswap32 ? tswap32 : t;
        }
    }
    // To simplify example, ignore finding index of true value in tmax==max
    *index = imax; // + which(tmax == max);
    return max[0];
}

godbolt 上的代码允许将 VLEN 更改为 8 或 4。

这大多效果很好。对于 AVX/SSE，绝对值变为t & 0x7fffffff使用(v)andps，即清除符号位。对于 NEON，它是用vneg+完成的fmaxnm。查找和广播水平最大值的块变成了排列和最大值指令的有效序列。gcc 能够使用 NEONfabs来获取绝对值。

4 元素 SSE/NEON 目标上的 8 元素向量在 clang 上运行良好。它在两组寄存器上使用一对指令，并且对于 SWAP128 水平操作max或or两个寄存器，没有任何不必要的置换。另一方面，gcc 确实无法处理这个问题，并且生成的大部分是非 SIMD 代码。如果我们将向量长度减少到 4，gcc 对于 SSE 和 NEON 就可以正常工作。

但有一个问题if (any(cmp))。对于 clang + SSE/AVX，它工作得很好，vcmpltps+ vptest，orps在 SSE 上从 8->4。

但是 NEON 上的 gcc 和 clang 会执行所有排列和 OR 操作，然后将结果移至 gp 寄存器进行测试。

除了架构特定的内在函数之外，是否还有一些代码可以ptest与 gcc vmaxvq、clang/gcc 和 NEON 一起使用？

我尝试了一些其他方法，if (x[0] || x[1] || ... x[7])但它们更糟糕。

更新

我创建了一个更新的示例，它显示了两种不同的实现，即原始方法和 chtz 建议的“向量中的索引”方法，并在 Aki Suihkonen 的答案中显示。我们可以看到生成的 SSE 和 NEON 输出。

尽管有些人可能会持怀疑态度，但编译器确实从通用 SIMD（不是自动矢量化！）C++ 代码中生成了非常好的代码。在 SSE/AVX 上，我发现循环中的代码改进空间很小。NEON 版本仍然受到“any()”的次优实现的困扰。

除非数据通常按升序排列，或几乎如此，否则我的原始版本在 SSE/AVX 上仍然是最快的。我还没有在 NEON 上进行过测试。这是因为大多数循环迭代都找不到新的最大值，最好针对这种情况进行优化。“向量中的索引”方法产生更紧密的循环，编译器也做得更好，但常见情况只是在 SSE/AVX 上慢一点。常见情况在 NEON 上可能相同或更快。

有关编写通用 SIMD 代码的一些注意事项。

浮点向量的绝对值可以通过以下方式找到。它在 SSE/AVX（并带有清除符号位的掩码）和 NEON（fabs 指令）上生成最佳代码。

static vNs vabs(vNs x) {
    return -x < x ? x : -x;
}

这将在 SSE/AVX/NEON 上有效地实现垂直最大。它不进行比较；它产生架构的“max”指令。在 NEON 上，将其更改为 use>而不是会<导致编译器生成非常糟糕的标量代码。我猜是带有非规范或异常的东西。

template <typename v>  // Deduce vector type (float, unsigned, etc.)
static v vmax(v a, v b) {
    return a < b ? b : a; // compiles best with "<" as compare op
}

此代码将在寄存器中广播水平最大值。它在 SSE/AVX 上编译得很好。在 NEON 上，如果编译器可以使用水平最大指令然后广播结果，可能会更好。令我印象深刻的是，如果在 SSE/NEON 上使用 8 个元素向量（只有 4 个元素寄存器），编译器就会足够聪明，只使用一个寄存器来广播结果，因为顶部 4 个元素和底部 4 个元素是相同的。

template <typename v>
static v hmax(v x) {
    if (VLEN >= 8)
        x = vmax(x, __builtin_shufflevector(x,x, SWAP128));
    x = vmax(x, __builtin_shufflevector(x,x, SWAP64));
    return vmax(x, __builtin_shufflevector(x,x, SWAP32));
}

这是我发现的最好的“any()”。它在 SSE/AVX 上是最佳的，使用单个 ptest 指令。在 NEON 上，它执行排列和 OR，而不是水平最大指令，但我还没有找到一种方法可以在 NEON 上获得更好的效果。

static bool any(vNb x) {
    if (VLEN >= 8)
        x |= __builtin_shufflevector(x,x, SWAP128);
    x |= __builtin_shufflevector(x,x, SWAP64);
    x |= __builtin_shufflevector(x,x, SWAP32);
    return x[0];
}

同样有趣的是，在 AVX 上，代码i = i + 1将被编译为vpsubd ymmI, ymmI, ymmNegativeOne，即减去 -1。为什么？因为 -1 向量是用 和 生成的，vpcmpeqd ymm0, ymm0, ymm0这比广播 1 向量更快。

which()这是我想出的最好的。这为您提供了布尔向量中第一个真值的索引（0 = false，-1 = true）。使用 movemask 在 AVX 上可以做得更好一些。我不知道最好的NEON。

// vector of signed ints
typedef int vNi __attribute__((vector_size(VLEN*sizeof(int))));
// vector of bytes, same number of elements, 1/4 the size
typedef unsigned char vNb __attribute__((vector_size(VLEN*sizeof(unsigned char))));
// scalar type the same size as the byte vector
using sNb = std::conditional_t<VLEN == 4, uint32_t, uint64_t>;
static int which(vNi x) {
    vNb cidx = __builtin_convertvector(x, vNb);
    return __builtin_ctzll((sNb)cidx) / 8u;
}

Answer 1

正如 chtz 所评论的，最通用和典型的方法是使用另一个掩码来收集索引：

Vec8s indices = { 0,1,2,3,4,5,6,7};
Vec8s max_idx = indices;
Vec8f max_abs = abs(load8(ptr)); 

for (auto i = 8; i + 8 <= vec_length; i+=8) { 
    Vec8s data = abs(load8(ptr[i]));
    auto mask = is_greater(data, max_abs);
    max_idx = bitselect(mask, indices, max_idx);
    max_abs = max(max_abs, data);    
    indices = indices + 8;
}

另一种选择是交错值和索引：

auto data = load8s(ptr) & 0x7fffffff; // can load data as int32_t
auto idx = vec8s{0,1,2,3,4,5,6,7};
auto lo = zip_lo(idx, data);
auto hi = zip_hi(idx, data);

for (int i = 8; i + 8 <= size; i+=8) {
    idx = idx + 8;
    auto d1 = load8s(ptr + i) & 0x7fffffff;
    auto lo1 = zip_lo(idx, d1);
    auto hi1 = zip_hi(idx, d1);
    lo = max_u64(lo, lo1);
    hi = max_u64(hi, hi1);
}

如果输入范围足够小，可以将输入左移，同时将索引中的一些位附加到同一字的 LSB 位，则此方法尤其有利可图。

即使在这种情况下，我们也可以重新利用浮点数中的 1 位，从而节省一半的位/索引选择操作。

auto data0 = load8u(ptr) << 1; // take abs by shifting left 
auto data1 = (load8u(ptr + 8) << 1) + 1;  // encode odd index to data
auto mx = max_u32(data0, data1);  // the LSB contains one bit of index

看起来可以用作double存储，因为甚至 SSE2 也支持_mm_max_pd（需要注意 Inf/Nan 处理，当重新解释为 64 位双精度的高位部分时，它不再编码为 Inf/Nan）。