含if/else逻辑的C++函数基于ARM NEON的SIMD高效向量化方法

针对你遇到的编译器无法自动向量化的问题（分支逻辑+间接数组索引），可以通过ARM NEON的手动向量化来解决，核心思路是用向量指令批量处理数据，用掩码替代分支判断，利用gather指令处理间接索引。

核心优化思路 #

1. 批量加载数据：用NEON向量指令一次性加载多组arr_1的相邻元素，避免循环迭代的开销。
2. 掩码消除分支：通过向量比较指令生成相等/不等的掩码，用掩码将需要置0的乘积结果直接清零，替代原有的if/else分支。
3. Gather处理间接索引：利用AArch64 NEON的gather加载指令，根据向量中的索引批量获取arr_2的元素，解决间接索引无法自动向量化的问题。

完整优化代码 #

#include <arm_neon.h>

int arr_1[1000];
double arr_2[1000];

// 手动向量化的累加函数
double vectorized_sum() {
    double result = 0.0;
    int k = 0;

    // 批量处理4组数据（一次对应4次原val函数调用）
    for (; k <= 96; k += 4) {
        // 加载arr_1的连续元素：a_vec是arr_1[k], arr_1[k+1], arr_1[k+2], arr_1[k+3]
        // b_vec是arr_1[k+1], arr_1[k+2], arr_1[k+3], arr_1[k+4]
        int32x4_t a_vec = vld1q_s32(&arr_1[k]);
        int32x4_t b_vec = vld1q_s32(&arr_1[k+1]);

        // 比较相邻元素是否相等，生成掩码：相等时对应位为0xFFFFFFFF，不等时为0
        uint32x4_t eq_mask = vceqq_s32(a_vec, b_vec);
        // 生成不等掩码：不等时为0xFFFFFFFF，相等时为0
        uint32x4_t ne_mask = vceqzq_s32(eq_mask);

        // 批量加载arr_2中对应索引的元素（gather指令，M1支持）
        float64x2_t arr2_a_low = vld1q_gather_f64(&arr_2[vget_low_s32(a_vec)]);
        float64x2_t arr2_a_high = vld1q_gather_f64(&arr_2[vget_high_s32(a_vec)]);
        float64x4_t arr2_a_vec = vcombine_f64(arr2_a_low, arr2_a_high);

        float64x2_t arr2_b_low = vld1q_gather_f64(&arr_2[vget_low_s32(b_vec)]);
        float64x2_t arr2_b_high = vld1q_gather_f64(&arr_2[vget_high_s32(b_vec)]);
        float64x4_t arr2_b_vec = vcombine_f64(arr2_b_low, arr2_b_high);

        // 计算乘积向量
        float64x4_t prod_vec = vmulq_f64(arr2_a_vec, arr2_b_vec);

        // 生成选择掩码：不等时取1.0，相等时取0.0
        float64x4_t one_vec = vdupq_n_f64(1.0);
        float64x4_t zero_vec = vdupq_n_f64(0.0);
        float64x4_t select_mask = vbslq_f64(ne_mask, one_vec, zero_vec);

        // 用掩码清零相等位置的乘积
        float64x4_t masked_prod = vmulq_f64(prod_vec, select_mask);

        // 累加向量中所有元素到结果
        result += vaddvq_f64(masked_prod);
    }

    // 处理剩余不足4组的元素（本示例中100是4的倍数，此循环不会执行，保留为通用情况）
    for (; k < 100; ++k) {
        int a = arr_1[k];
        int b = arr_1[k+1];
        if (a != b) {
            result += arr_2[a] * arr_2[b];
        }
    }

    return result;
}

int main() {
    double result = vectorized_sum();
    return 0;
}

关键细节说明 #

• Gather加载指令：vld1q_gather_f64是AArch64 NEON的高级指令，支持根据向量中的索引批量从内存加载数据，完美解决arr_2间接索引无法自动向量化的问题，Apple M1基于ARMv8.5-A架构，完全支持该指令。
• 掩码分支消除：通过向量比较和位选择指令，将原有的if/else分支转换成向量运算，避免了分支预测开销，同时满足向量化的要求。
• 向量累加：vaddvq_f64指令直接计算向量中4个double元素的和，比循环累加标量值效率更高。

编译与验证 #

编译时使用clang -O3 -march=armv8.5-a your_code.cpp，-march参数可明确指定M1支持的架构，确保编译器生成最优的NEON指令。

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