__m256i向量置1位与对应位置乘积计算的优化方案问询

一、置位操作优化 #

原问题 #

给定排序且唯一的位位置数组int inds[N]（元素范围[0,255]），需将__m256i向量的对应位设为1，原实现如下：

alignas(32) uint64_t buf[4] = {0};

for (int i = 0; i < N; ++i) {
    int ind = inds[i];
    buf[ind / 64] |= 1ul << (ind % 64);
}
auto r = _mm256_load_si256((__m256i*)buf);

优化方案 #

1. 位运算替代除法取模

CPU对位运算的执行效率远高于除法与取模操作，直接用位运算替换原代码中的除法和取模：

alignas(32) uint64_t buf[4] = {0};

for (int i = 0; i < N; ++i) {
    int ind = inds[i];
    buf[ind >> 6] |= 1ULL << (ind & 63); // 位运算替代除法取模
}
auto r = _mm256_load_si256((__m256i*)buf);

2. 直接在向量寄存器中构造掩码（避免内存中转）

如果N较小，可直接在__m256i寄存器内逐位构造掩码，省去内存缓冲区的读写操作：

__m256i mask = _mm256_setzero_si256();
for (int i = 0; i < N; ++i) {
    int ind = inds[i];
    int lane = ind >> 6; // 0-3，对应__m256i的4个64位 lane
    int bit = ind & 63;
    // 构造仅对应lane位置1的向量掩码
    __m256i bit_mask = _mm256_set_epi64x(
        lane == 3 ? (1ULL << bit) : 0,
        lane == 2 ? (1ULL << bit) : 0,
        lane == 1 ? (1ULL << bit) : 0,
        lane == 0 ? (1ULL << bit) : 0
    );
    mask = _mm256_or_si256(mask, bit_mask);
}

该方法适合N较小的场景，全程在寄存器内完成掩码累积，避免内存访问开销。

二、获取操作优化 #

原问题 #

需计算__m256i掩码中置1位对应double const sizes[256]数组元素的乘积，原实现如下：

inline
double size (__m256i r, double const sizes[256])
{
    alignas(16) uint64_t buf[4];
    _mm256_store_si256((__m256i*)buf, r);

    double s[4] = {1.0, 1.0, 1.0, 1.0};

    for (; buf[0] != 0; buf[0] &= buf[0] - 1)
        s[0] *= sizes[__builtin_ctzl(buf[0]) + 0 * 64];

    for (; buf[1] != 0; buf[1] &= buf[1] - 1)
        s[1] *= sizes[__builtin_ctzl(buf[1]) + 1 * 64];

    for (; buf[2] != 0; buf[2] &= buf[2] - 1)
        s[2] *= sizes[__builtin_ctzl(buf[2]) + 2 * 64];

    for (; buf[3] != 0; buf[3] &= buf[3] - 1)
        s[3] *= sizes[__builtin_ctzl(buf[3]) + 3 * 64];

    return s[0] * s[1] * s[2] * s[3];
}

优化方案 #

1. 跳过内存存储，直接提取寄存器元素

原代码中_mm256_store_si256将向量写入内存再读取的操作完全可以省略，利用AVX2的_mm256_extract_epi64指令直接从__m256i寄存器提取每个64位lane的数值，大幅减少内存开销：

inline double size(__m256i r, double const sizes[256]) {
    double result = 1.0;

    // 直接从寄存器提取每个64位lane，跳过内存存储
    uint64_t lane0 = _mm256_extract_epi64(r, 0);
    while (lane0 != 0) {
        int bit_pos = __builtin_ctzl(lane0);
        result *= sizes[bit_pos];
        lane0 &= lane0 - 1; // 清除最低位的1
    }

    uint64_t lane1 = _mm256_extract_epi64(r, 1);
    while (lane1 != 0) {
        int bit_pos = __builtin_ctzl(lane1);
        result *= sizes[bit_pos + 64];
        lane1 &= lane1 - 1;
    }

    uint64_t lane2 = _mm256_extract_epi64(r, 2);
    while (lane2 != 0) {
        int bit_pos = __builtin_ctzl(lane2);
        result *= sizes[bit_pos + 128];
        lane2 &= lane2 - 1;
    }

    uint64_t lane3 = _mm256_extract_epi64(r, 3);
    while (lane3 != 0) {
        int bit_pos = __builtin_ctzl(lane3);
        result *= sizes[bit_pos + 192];
        lane3 &= lane3 - 1;
    }

    return result;
}

该版本核心操作均在寄存器内完成，省去内存缓冲区的读写，性能提升明显。

2. 批量处理优化（适合置位数量较多的场景）

如果掩码中置位的位数较多，可预将sizes数组按64位块分组，结合SIMD指令批量加载多个元素并相乘。但由于乘法是累积操作，这种方法可能引入浮点精度损失，仅适合对精度要求不高的场景：

• 将每个64位块对应的sizes元素加载到__m256d向量组中，利用_mm256_mul_pd逐元素相乘，再通过水平乘法累积结果，需提前完成分组预处理。

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