【发布时间】：2022-01-24 23:07:12
【问题描述】：
我想看看是否可以编写一些可以高效编译的通用 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，绝对值使用 (v)andps 变为 t & 0x7fffffff，即清除符号位。对于 NEON，使用 vneg + fmaxnm 完成。查找和广播水平最大值的块变成了置换和最大值指令的有效序列。 gcc 能够使用 NEON fabs 作为绝对值。
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 寄存器进行测试。
除了架构特定的内在函数之外，是否有一些代码可以通过 gcc 获得 ptest 以及通过 clang/gcc 和 NEON 获得 vmaxvq？
我尝试了一些其他方法，例如if (x[0] || x[1] || ... x[7])，但效果更差。
更新
我创建了一个updated example，它显示了两种不同的实现，包括原始方法和 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 上，将其更改为使用 > 而不是 < 会导致编译器生成非常糟糕的标量代码。我猜是有异常或异常的东西。
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 上，如果编译器可以使用水平 max 指令然后广播结果可能会更好。令我印象深刻的是，如果在 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。为什么？因为使用vpcmpeqd ymm0, ymm0, ymm0 会生成-1s 的向量，这比广播1s 的向量要快。
这是我想出的最好的which()。这为您提供了布尔向量中第一个真值的索引（0 = 假，-1 = 真）。使用 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;
}
【问题讨论】：
-
为什么要在循环内而不是在最后搜索向量的水平最大值？即，还保留一个 8 大小的索引寄存器，其中包含所有
k + 8*i元素的最大值的索引（对于k=0..8）。 -
您也可以使用 movemask+popcount 指令进行任何操作，以便更快地找到位置（并进行检查）。我想这可能会更快，但这可能需要 VLEN 更大。
-
最后我想到了水平最大值，但没有找到一种方法来跟踪哪个元素在哪个
i。但是向量索引寄存器，它可以工作，并且可以避免if(any(cmp))（我仍然想优化）。变成max = cmp ? t : max;和imax = cmp ? i : imax;，最后找到max的水平最大值和imax的对应元素。 -
我还没有找到不使用特定于体系结构的内在函数来获取移动掩码的方法。 movemask + ctz 给出了我发现的最好的
which()，它对any()很有效。 -
您的编辑看起来应该作为答案发布，而不是问题的一部分。