如何在Rust中对模乘法实现SIMD自动向量化？

问题根源 #

原mod_mul64实现依赖u128乘法后直接取模，编译器无法自动向量化的核心原因：

• 主流SIMD指令集对128位整数的乘法/取模原生支持有限，难以批量并行处理；
• 直接的%操作依赖除法运算，而除法是编译器难以向量化的复杂指令。

解决方案：改用Barrett约减实现模乘法 #

Barrett约减通过预计算常数，将模运算转换为乘法、移位和减法操作——这些都是可被编译器高效向量化的整数运算。

修改后的完整代码 #

#[inline(always)]
fn mod_mul64(a: u64, b: u64, modulus: u64, barrett_k: u64) -> u64 {
    // Barrett约减核心逻辑：(a*b) mod m = (a*b) - floor((a*b*m')/2^64) * m，其中m' = floor(2^128 / m)
    let product = a as u128 * b as u128;
    let t = product as u64;
    let high = (product >> 64) as u64;
    
    // 计算floor((a*b*m')/2^64)，barrett_k为预计算的floor(2^128 / modulus)
    let mut q = high.wrapping_mul(barrett_k);
    let q_high = (q as u128 * modulus as u128) >> 64;
    q = q.wrapping_add(q_high);
    
    let mut result = t.wrapping_sub(q.wrapping_mul(modulus));
    // 处理结果可能超出模数的情况
    if result >= modulus {
        result -= modulus;
    }
    result
}

pub fn mul(a: &mut [u64], b: &[u64], modulo: u64) {
    // 预计算Barrett常数，仅执行一次避免重复开销
    let barrett_k = (u128::MAX / modulo as u128 + 1) as u64;
    
    for _ in (0..1000).step_by(4) {
        // 拆分计算消除数据依赖，提升向量化效率
        let temp = [
            mod_mul64(b[0], a[7], modulo, barrett_k),
            mod_mul64(b[1], a[6], modulo, barrett_k),
            mod_mul64(b[2], a[5], modulo, barrett_k),
            mod_mul64(b[3], a[4], modulo, barrett_k),
        ];
        let temp2 = [
            mod_mul64(b[4], a[3], modulo, barrett_k),
            mod_mul64(b[5], a[2], modulo, barrett_k),
            mod_mul64(b[6], a[1], modulo, barrett_k),
            mod_mul64(b[7], a[0], modulo, barrett_k),
        ];
        a[0..4].copy_from_slice(&temp);
        a[4..8].copy_from_slice(&temp2);
    }
}

#[allow(unused)]
pub fn main() {
    let mut a = [0u64; 8];
    let b = [1u64; 8];
    let modulo = 1000000007u64;
    mul(&mut a, &b, modulo);
    println!("a: {:?}", a);
}

关键优化点说明 #

1. Barrett常数预计算：在mul函数开头一次性计算barrett_k，避免每次模乘重复计算，降低额外开销；
2. 消除数据依赖：原代码中a[4]依赖a[3]的新值、a[7]依赖a[0]的新值，形成循环内的数据依赖，阻碍向量化。通过临时数组拆分计算，打破这种依赖，让编译器可以并行处理所有8次模乘；
3. 可选无分支优化：若想进一步消除最后的条件判断，可改用位运算替代：

   result = result.wrapping_sub(modulus) & ((modulus as i64 - 1 - result as i64) >> 63);
   

   此写法可读性稍低，但能彻底移除分支，进一步提升SIMD兼容性。

编译优化要求 #

确保开启以下编译选项：

• 优化等级：-O3或-O2
• 目标CPU指定：-C target-cpu=native（让编译器针对当前CPU生成最优SIMD指令）

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