如何在Eigen中利用AVX2指令实现布尔数组对浮点数组的条件操作

背景 #

我使用Eigen的固定长度数组Eigen::Array<float, 8, 1>（通过typedef定义为f32x8）编写模板函数，实现同时支持标量float和AVX2 ymm向量的运算。例如：

template<typename T> inline T add(const T &a, const T &b) { return a + b; }

当用f32x8调用该函数时，编译器会生成单条vaddps指令，完全满足SIMD加速需求。

问题 #

但当我尝试用typedef Eigen::Array<bool, 8, 1> boolx8进行条件操作时：

f32x8 a, b, c;
boolx8 m1 = a < b;
boolx8 m2 = b < c;
boolx8 m3 = m1 && m2;

Eigen会生成标量循环代码（逐个处理元素），而非预期的AVX2指令。原因是bool类型占1字节，boolx8总大小为8字节，和AVX2寄存器中位掩码形式的比较结果不兼容。

解决方案 #

查看Eigen的AVX2架构实现代码（Eigen/src/Core/arch/AVX2/PacketMath.h第584行）发现，pcmp_le这类比较操作的模板特化返回的是Packet8f（即f32x8）——这说明Eigen本身是复用浮点数组类型来存储SIMD比较结果的，而非单独的布尔数组类型。

正确的做法是直接用f32x8作为布尔掩码的存储类型：

1. 比较操作直接返回浮点掩码：浮点向量间的比较（如<）会自动返回f32x8，其中每个元素是全1（对应true）或全0（对应false）的float值，完全匹配vcmpgtps等AVX2比较指令的输出格式。
2. 用逐元素乘积替代逻辑与：掩码间的逻辑与操作，可使用Eigen的cwiseProduct()方法实现，对应AVX2的vandps指令：

f32x8 a, b, c;
f32x8 m1 = a < b;
f32x8 m2 = b < c;
f32x8 m3 = m1.cwiseProduct(m2); // 等价于逻辑与，生成vandps指令

这样修改后，编译器会生成纯AVX2指令，不再出现标量循环。

额外说明 #

当需要将掩码应用到浮点向量进行条件选择时，直接使用Eigen的select()方法即可，它能识别这种浮点掩码类型并生成vblendvps等高效指令：

f32x8 result = m3.select(a, b); // 掩码元素为true时取a的对应元素，否则取b的元素

内容的提问来源于stack exchange，提问作者Martijn Courteaux