Compiler Explorer

Source code

// If you use `main()`, declare it as `pub` to see it in the output:
// pub fn main() { ... }
mod simd {
#[cfg(target_arch = "x86")]
use std::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use std::arch::x86_64::*;

pub type f64x4 = __m256d;
pub const f64x4_SIZE: usize = 4;

#[inline]
pub fn f4zeros() -> f64x4 {
    unsafe { _mm256_setzero_pd() }
}

#[inline]
pub fn f4init(f1: f64, f2: f64, f3: f64, f4: f64) -> f64x4 {
    unsafe { _mm256_set_pd(f1, f2, f3, f4) }
}

#[inline]
pub fn f4add(a: f64x4, b: f64x4) -> f64x4 {
    unsafe { _mm256_add_pd(a, b) }
}

#[inline]
pub fn f4mul(a: f64x4, b: f64x4) -> f64x4 {
    unsafe { _mm256_mul_pd(a, b) }
}

#[inline]
pub fn f4div(a: f64x4, b: f64x4) -> f64x4 {
    unsafe { _mm256_div_pd(a, b) }
}

#[inline]
pub fn lowest(a: f64x4) -> f64 {
    unsafe { _mm256_cvtsd_f64(a) }
}

#[inline]
pub fn shift(a: f64x4, i: i8) -> f64x4 {
    match i {
        1 => unsafe { _mm256_permute4x64_pd(a, 0b_00_11_10_01) },
        2 => unsafe { _mm256_permute4x64_pd(a, 0b_01_00_11_10) },
        3 => unsafe { _mm256_permute4x64_pd(a, 0b_10_01_00_11) },
        _ => panic!("shift must be an integer between 1 and 3"),
    }
}

#[inline]
pub fn get(a: f64x4, i: i8) -> f64 {
    let b = match i {
        0 => shift(a, 3),
        1 => shift(a, 2),
        2 => shift(a, 1),
        3 => a,
        _ => panic!("index out of bounds for f64x4"),
    };
    lowest(b)
}

#[inline]
pub fn vec_sum(a: f64x4) -> f64 {
    let sum1 = f4add(shift(a, 1), a);
    let sum2 = f4add(shift(sum1, 2), sum1);

lowest(sum2)
}

#[inline]
pub fn f4mvp(m: &[f64x4; f64x4_SIZE], v: f64x4) -> f64x4 {
    let mv0 = f4mul(m[0], v);
    let mv1 = f4mul(m[1], v);
    let mv2 = f4mul(m[2], v);
    let mv3 = f4mul(m[3], v);

// { mv0_0+mv0_1, mv1_0+mv1_1, mv0_2+mv0_3, mv1_2+mv1_3 }
    let temp01 = unsafe { _mm256_hadd_pd(mv1, mv0) };

// { mv2_0+mv2_1, mv3_0+mv3_1, mv2_2+mv2_3, mv3_2+mv3_3 }
    let temp23 = unsafe { _mm256_hadd_pd(mv3, mv2) };

// { mv0_2+mv0_3, mv1_2+mv1_3, mv2_0+mv2_1, mv3_0+mv3_1 }
    let swapped = unsafe {_mm256_permute2f128_pd(temp01, temp23, 0b_00_00_00_11)};
    
    // { mv0_0+mv0_1, mv1_0+mv1_1, mv2_2+mv2_3, mv3_2+mv3_3 }
    let blended = unsafe {_mm256_blend_pd(temp01, temp23, 0b_00_00_00_11)};

f4add(swapped, blended)
}

}

const N: usize = 100000;
const SIZE: usize = 4;
const BETA: f64 = 0.33;
const GAMMA: f64 = 1.06;

pub fn with_simd(position: simd::f64x4) -> simd::f64x4 {
    let lm: [simd::f64x4; SIZE] = [
        simd::f4init(GAMMA, -GAMMA * BETA, 0.0, 0.0),
        simd::f4init(-GAMMA * BETA, GAMMA, 0.0, 0.0),
        simd::f4init(0.0, 0.0, 1.0, 0.0),
        simd::f4init(0.0, 0.0, 0.0, 1.0),
    ];
    
    simd::f4mvp(&lm, position)
}

pub fn without_simd(position: &[f64; SIZE]) -> [f64; SIZE] {
    let lm: [[f64; SIZE]; SIZE] = [
        [GAMMA, -GAMMA * BETA, 0.0, 0.0],
        [-GAMMA * BETA, GAMMA, 0.0, 0.0],
        [0.0, 0.0, 1.0, 0.0],
        [0.0, 0.0, 0.0, 1.0],
    ];
    let mut result = [0.0; SIZE];
    for j in 0..SIZE {
        for k in 0..SIZE {
            result[j] += lm[j][k] * position[k];
        }
    }
    result
}

fn main() {
    let position1 = simd::f4init(0.5, -0.5, 0.0, 1.0);
    let position2 = [0.5, -0.5, 0.0, 1.0];

let result1 = with_simd(position1);
    
    let result2 = without_simd(&position2);
}