寻求基于SIMD库与现代X86-64指令集的向量最大元素索引高效实现

问题背景 #

我使用Agner的C++向量类库实现了函数f，用于查找32位无符号字符向量中最大值的第一个出现索引：

int f(const Vec32uc &a) {
  Vec32uc b{horizontal_max(a)};
  return horizontal_find_first(a == b);
}

通过-O3 -march=alderlake编译后生成的汇编代码较长，对比早期的intrinsic实现方案，我认为当前实现效率有待提升，希望找到基于SIMD库和现代X86-64指令集的更高效实现方式。

优化方案（基于AVX2指令集） #

针对Alder Lake支持的AVX2指令集，我们可以从减少水平归约指令数和合并冗余操作两个方向优化：

优化思路 #

1. 利用AVX2的向量max和移位操作，快速完成全局最大值的水平归约，减少中间指令
2. 直接将最大值广播到整个YMM寄存器，与原向量生成比较掩码
3. 将掩码转换为整数后，通过位操作快速定位第一个匹配的索引

优化后的代码实现 #

简化版（基于Agner库内置函数）

#include "vectorclass.h"

int f_opt(const Vec32uc &a) {
    // 直接获取全局最大值
    uint8_t max_val = horizontal_max(a);
    // 广播最大值并生成相等掩码
    Vec32uc eq_mask = a == Vec32uc(max_val);
    // 查找第一个匹配的索引
    return horizontal_find_first(eq_mask);
}

Agner库的horizontal_max在AVX2环境下已做优化，但手动提取最大值后再广播，能避免中间YMM寄存器的冗余操作，让编译器生成更紧凑的代码。

极致优化版（手动控制AVX2指令流）

#include "vectorclass.h"

int f_opt_avx2(const Vec32uc &a) {
    Vec32uc tmp = a;
    // 跨128位lane归约最大值
    tmp = max(tmp, tmp >> 16);
    // 逐步8/4/2/1位归约
    tmp = max(tmp, tmp >> 8);
    tmp = max(tmp, tmp >> 4);
    tmp = max(tmp, tmp >> 2);
    tmp = max(tmp, tmp >> 1);
    // 提取全局最大值
    uint8_t max_val = tmp[0];
    
    // 生成相等掩码并转换为整数
    Vec32uc eq_mask = a == Vec32uc(max_val);
    int mask = eq_mask.movemask();
    // 用内置位操作找第一个置位索引，无匹配返回-1
    return mask ? __builtin_ctz(mask) : -1;
}

这段代码通过向量移位操作替代多次拆分合并指令，进一步压缩了水平归约的指令数量，编译后会生成更精简的AVX2指令序列。

编译后的汇编优势 #

优化后的实现相比原代码：

• 减少了vextracti128、vpunpckhqdq等冗余指令
• 水平归约步骤的指令数大幅减少
• 掩码处理和索引查找的流程更直接

内容的提问来源于stack exchange，提问作者Paul Jurczak