跳转至

Rust使用SIMD进行加速计算

依赖库

该库封装了simd相关指令集, 并对寄存器的使用进行了抽象, 方便我们使用

packed_simd = { version = "0.3.4", package = "packed_simd_2" }

Demo

    // {element_type}x{number_of_lanes} == Simd<[element_type; number_of_lanes]>
    // i32x4; u32x4....等等


    use packed_simd::*;

    // 1. 纵向操作: 默认情况下, 对打包向量的操作是 "垂直" 的, 可以对齐进行纵向操作
    let v = vec![1, 2, 3, 4];
    let a = i32x4::new(v[0], v[1], v[2], v[3]);
    let b = i32x4::new(5, 6, 7, 8);
    let c = a + b;
    println!("{:?}", a);  // i32x4(1, 2, 3, 4)

    // 纵向 与常量相加
    let mut a = i32x4::new(1, 2, 3, 4);
    a += 1;  // 每个元素 增加1
    println!("{:?}", a);  // i32x4(2, 3, 4, 5)


    // 2. 横向操作
    let a = u32x4::new(2, 5, 8, 8);
    let a = a.wrapping_sum();  // 给元素相加
    println!("{:?}", a);  // 23

    // 3. 从slice 创建向量, 使用from_slice_unaligned
    // 注意, slice 的len 必须大于等于 4 才可以创建成功
    let x = &[1, 2, 3,];  // 出错, 必须要大于4个
    let x = &[1, 2, 3, 4];  // i32x4(1, 2, 3, 4)
    let x = &[1, 2, 3, 4, 5];  // i32x4(1, 2, 3, 4)
    i32x4::from_slice_unaligned(x);


    // 4. 纵向垂直操作都是很快速的, 做一个垂直相加, 然后横向相加 对vec 求和的函数
    fn reduce(mut v: Vec<i32>) -> i32{

        // 不能整除的需要从最后补齐
        let y = v.len() % 4;
        if y != 0 {
            for _ in 0..4-y {
                v.push(0);
            }
        }

        let mut sum = i32x4::splat(0);  // 创建一个都为 0 的向量

        for mut i in 0..v.len()/8 {
            i *= 4;
            sum += i32x4::from_slice_unaligned(&v[i..]);  // 纵向相加
        }
        sum.wrapping_sum()  // 横向相加
    }

    let v = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11];
    let r = reduce(v);
    println!("{:?}", r);

    // 5. 读取, 替换 向量中的指定位置的元素
    let mut a = u32x4::new(5, 8, 9, 0);
    println!("{:?}", a.extract(2));  // 读取下标为 2 的元素  9
    let b = a.replace(2, 14);  // 返回一个新的向量, 其中第 2 个元素被替换为 14
    println!("{:?}", b);  // u32x4(5, 8, 14, 0)

    // 6. 从向量中 写入数据到 slice中
    let mut a = u32x4::new(12, 18, 24, 64);
    let mut v = vec![1, 2, 3, 4, 5];
    a.write_to_slice_aligned(&mut v);
    println!("{:?}", v);  // [12, 18, 24, 64, 5]

    // 7. 使用掩码 "选择" 指定要操作的元素, a和b 的元素  如果 v[1]
    // m.select(a, b); 选择 掩码为真的 a向量中的元素, 否则返回 b向量中的元素
    // 掩码选择 可以选用不同 "宽度" 的向量类型, 下面这个例子使用的 m16x4而不是m32x4, 通常情况下, 适用于备操作的向量的元素宽度相同的掩码元素宽度性能更好!
    let a = u32x4::new(2, 7, 3, 4);
    let b = u32x4::new(1, 7, 4, 5);
    let m = m16x4::new(true, true, false, false);  // 选择操作 0号 1号下标的  a向量中的元素, 其他的元素使用的是b向量中的
    let r = m.select(a, b);
    println!("{:?}", r);  // u32x4(2, 7, 4, 5)//

    // 8. 利用掩码 对 向量中的前 n 个元素 "选择操作", 这里对前两个元素进行 +1 并返回一个新的向量
    let a = i32x4::new(1, 3, 5, 8);
    let m = m32x4::new(true, true, false, false);
    let a = m.select(a + 1, a);
    println!("{:?}", a);  // i32x4(2, 4, 5, 8)

    let a = i32x4::new(0, 1, 2, 3);
    let b = i32x4::new(1, 2, 2, 3);



    // simd 简单测试一下性能
    let mut rng =rand::thread_rng();

    let mut v = Vec::new();
    for i in 0..1000_0000 {
        v.push(rng.gen::<f64>())
    }


    let old_len = v.len();
    let y = v.len() % 8;
    if y != 0 {
        for _ in 0..8-y {
            v.push(0.0);
        }
    }

    let mut i = 0;
    let t = std::time::Instant::now();
    for _ in 0..v.len() / 8 {
        let mut temp = f64x8::from_slice_unaligned(&v[i..]);
        temp += rng.gen::<f64>();
        temp -= rng.gen::<f64>();
        temp *= rng.gen::<f64>();
        temp /= rng.gen::<f64>();
        temp.write_to_slice_unaligned(&mut v[i..]);
        i += 8;
    }
    v.truncate(old_len);
    println!("{:?}", t.elapsed());

    // 不使用simd时间测试
    let mut v = Vec::new();
    for i in 0..1000_0000 {
        v.push(rng.gen::<f64>())
    }
    let t = std::time::Instant::now();
    for i in 0..v.len() {
        v[i] += rng.gen::<f64>();
        v[i] -= rng.gen::<f64>();
        v[i] *= rng.gen::<f64>();
        v[i] /= rng.gen::<f64>();
    };
    println!("{:?}", t.elapsed());

一些记录

在使用SIMD的时候, 尽量避免内存拷贝, 并对寄存器内的值进行多次计算, 否则效率还不如 for循环 (逃