优化三重数组元素对交集popcount乘积的嵌套循环性能

原函数代码 #

#define N 48
// N = 47 同样适用
int f(
const uint16_t * __restrict A,
const uint16_t * __restrict B,
const uint16_t * __restrict C) {
    int E = 0;
    for (int r = 0; r < N; ++r) {
        for (int s = r; s < N; ++s) {
            int pa = __builtin_popcount(A[r] & A[s]);
            int pb = __builtin_popcount(B[r] & B[s]);
            int pc = __builtin_popcount(C[r] & C[s]);
            E += pa*pb*pc;
        }
    }
    return E;
}

已尝试的优化方法 #

• 对A、B、C数组做线性时间预处理，仅保留pa*pb*pc非零的三元组，但因位分布均匀，几乎无过滤效果
• 通过分块减少缓存缺失
• 将数组重新打包为uint64_t类型，重构popcount逻辑以同时处理4个输入，无明显效果

核心优化方向（基于AVX2指令集） #

1. 数学公式重写：将双重循环转化为位组合统计 #

原函数的核心计算pa*pb*pc可通过数学展开彻底重构：

• pa是A[r]与A[s]的共同置位位数，等价于对每个位k，统计A[r]和A[s]第k位同时为1的次数
• 三重乘积pa*pb*pc可展开为所有位组合(k,m,n)的统计之和，最终推导得：

  E = ΣₖΣₘΣₙ [ S(k,m,n) * (S(k,m,n) + 1) / 2 ]
  

  其中S(k,m,n)是满足A[r]第k位为1、B[r]第m位为1、C[r]第n位为1的r的数量（0≤r<N）

这个转换将原O(N²)的计算（48×49/2=1176次迭代）变为固定的16×16×16=4096次简单计算，完全规避了嵌套循环的开销。

2. 预处理位掩码加速S(k,m,n)计算 #

利用uint64_t（可容纳48个r的位信息）预处理每个数组的位掩码：

• 为A数组创建maskA[16]，maskA[k]的第r位为1当且仅当A[r]的第k位为1
• 同理生成maskB[16]和maskC[16]
• 对每个(k,m,n)，S(k,m,n) = __builtin_popcountll(maskA[k] & maskB[m] & maskC[n])

预处理仅需O(N×16)次操作（48×16=768次），后续的4096次popcount和乘法操作均为CPU单周期或低延迟指令，远快于原循环中的多次位与、popcount和三重乘法。

3. AVX2向量化加速掩码运算 #

针对AVX2指令集，可进一步并行化掩码计算：

• 将maskA/maskB/maskC的16个uint64_t元素加载到AVX2寄存器（256位寄存器可容纳4个uint64_t）
• 使用_mm256_and_si256指令并行计算多组掩码的按位与，一次处理4组(k,m,n)组合
• 结合_mm256_popcnt_epi64（部分CPU支持，或用_mm256_movemask_epi8辅助计算）并行获取多个popcount结果，大幅减少循环次数

4. 编译器优化补充 #

在现有编译选项基础上，可添加：

• -mavx2显式启用AVX2指令集（部分编译器需显式指定）
• 若自动向量化干扰手动优化，可临时添加-fno-tree-vectorize，建议优先依赖编译器自动向量化

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