[add]: simd

2 files changed, 512 insertions, 0 deletions

diff --git a/libpsi-core/src/maths/mod.rs b/libpsi-core/src/maths/mod.rs
index e716181..85f4872 100644
--- a/libpsi-core/src/maths/mod.rs
+++ b/libpsi-core/src/maths/mod.rs
@@ -2,6 +2,7 @@ pub mod complex;
 pub mod format;
 pub mod matrix;
 pub mod numeric;
+pub mod simd;
 pub mod vector;
 pub mod vector_ops;
 
@@ -9,4 +10,5 @@ pub use complex::*;
 pub use format::*;
 pub use matrix::*;
 pub use numeric::*;
+pub use simd::*;
 pub use vector::*;
diff --git a/libpsi-core/src/maths/simd.rs b/libpsi-core/src/maths/simd.rs
new file mode 100644
index 0000000..de0370c
--- /dev/null
+++ b/libpsi-core/src/maths/simd.rs
@@ -0,0 +1,510 @@
+use crate::{complex, Complex};
+
+#[cfg(target_arch = "x86_64")]
+use std::arch::x86_64::*;
+
+#[cfg(target_arch = "aarch64")]
+use std::arch::aarch64::*;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum SimdCapability {
+    None,
+    #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
+    Avx2,
+    #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
+    Avx512,
+    #[cfg(target_arch = "aarch64")]
+    Neon,
+}
+
+impl SimdCapability {
+    pub fn detect() -> Self {
+        #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
+        {
+            if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512dq") {
+                return SimdCapability::Avx512;
+            }
+            if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
+                return SimdCapability::Avx2;
+            }
+        }
+
+        #[cfg(target_arch = "aarch64")]
+        {
+            return SimdCapability::Neon;
+        }
+
+        #[allow(unreachable_code)]
+        SimdCapability::None
+    }
+
+    pub fn name(&self) -> &'static str {
+        match self {
+            SimdCapability::None => "Scalar",
+            #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
+            SimdCapability::Avx2 => "AVX2+FMA",
+            #[cfg(any(target_arch = "x86_64", target_arch = "x86"))]
+            SimdCapability::Avx512 => "AVX-512",
+            #[cfg(target_arch = "aarch64")]
+            SimdCapability::Neon => "NEON",
+        }
+    }
+}
+
+pub fn apply_single_qubit_gate_simd(
+    state: &mut [Complex<f64>],
+    gate: &[[Complex<f64>; 2]; 2],
+    target: usize,
+    num_qubits: usize,
+) {
+    let capability = SimdCapability::detect();
+
+    match capability {
+        #[cfg(target_arch = "x86_64")]
+        SimdCapability::Avx2 => unsafe {
+            apply_single_qubit_avx2(state, gate, target, num_qubits);
+        },
+        #[cfg(target_arch = "x86_64")]
+        SimdCapability::Avx512 => unsafe {
+            apply_single_qubit_avx512(state, gate, target, num_qubits);
+        },
+        #[cfg(target_arch = "aarch64")]
+        SimdCapability::Neon => unsafe {
+            apply_single_qubit_neon(state, gate, target, num_qubits);
+        },
+        _ => {
+            apply_single_qubit_scalar(state, gate, target, num_qubits);
+        }
+    }
+}
+
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx2", enable = "fma")]
+unsafe fn apply_single_qubit_avx2(
+    state: &mut [Complex<f64>],
+    gate: &[[Complex<f64>; 2]; 2],
+    target: usize,
+    num_qubits: usize,
+) {
+    let target_bit = num_qubits - 1 - target;
+    let step = 1 << target_bit;
+    let dim = 1 << num_qubits;
+
+    let g00 = gate[0][0];
+    let g01 = gate[0][1];
+    let g10 = gate[1][0];
+    let g11 = gate[1][1];
+
+    let pairs: Vec<(usize, usize)> = (0..dim)
+        .filter(|&i| (i >> target_bit) & 1 == 0)
+        .map(|i| (i, i | step))
+        .collect();
+
+    let chunks = pairs.len() / 2;
+
+    for chunk_idx in 0..chunks {
+        let (i0, j0) = pairs[chunk_idx * 2];
+        let (i1, j1) = pairs[chunk_idx * 2 + 1];
+
+        let s0_re = _mm256_set_pd(
+            state[j1].real,
+            state[i1].real,
+            state[j0].real,
+            state[i0].real,
+        );
+        let s0_im = _mm256_set_pd(
+            state[j1].imaginary,
+            state[i1].imaginary,
+            state[j0].imaginary,
+            state[i0].imaginary,
+        );
+
+        let g_re_0 = _mm256_set_pd(g01.real, g00.real, g01.real, g00.real);
+        let g_im_0 = _mm256_set_pd(g01.imaginary, g00.imaginary, g01.imaginary, g00.imaginary);
+        let g_re_1 = _mm256_set_pd(g11.real, g10.real, g11.real, g10.real);
+        let g_im_1 = _mm256_set_pd(g11.imaginary, g10.imaginary, g11.imaginary, g10.imaginary);
+
+        let prod0_re = _mm256_fmsub_pd(s0_re, g_re_0, _mm256_mul_pd(s0_im, g_im_0));
+        let prod0_im = _mm256_fmadd_pd(s0_re, g_im_0, _mm256_mul_pd(s0_im, g_re_0));
+
+        let prod1_re = _mm256_fmsub_pd(s0_re, g_re_1, _mm256_mul_pd(s0_im, g_im_1));
+        let prod1_im = _mm256_fmadd_pd(s0_re, g_im_1, _mm256_mul_pd(s0_im, g_re_1));
+
+        let mut res0_re = [0.0f64; 4];
+        let mut res0_im = [0.0f64; 4];
+        let mut res1_re = [0.0f64; 4];
+        let mut res1_im = [0.0f64; 4];
+
+        _mm256_storeu_pd(res0_re.as_mut_ptr(), prod0_re);
+        _mm256_storeu_pd(res0_im.as_mut_ptr(), prod0_im);
+        _mm256_storeu_pd(res1_re.as_mut_ptr(), prod1_re);
+        _mm256_storeu_pd(res1_im.as_mut_ptr(), prod1_im);
+
+        state[i0] = complex!(res0_re[0] + res0_re[1], res0_im[0] + res0_im[1]);
+        state[j0] = complex!(res1_re[0] + res1_re[1], res1_im[0] + res1_im[1]);
+        state[i1] = complex!(res0_re[2] + res0_re[3], res0_im[2] + res0_im[3]);
+        state[j1] = complex!(res1_re[2] + res1_re[3], res1_im[2] + res1_im[3]);
+    }
+
+    for &(i, j) in pairs.iter().skip(chunks * 2) {
+        let s0 = state[i];
+        let s1 = state[j];
+
+        let new0 = complex!(
+            s0.real * g00.real - s0.imaginary * g00.imaginary + s1.real * g01.real
+                - s1.imaginary * g01.imaginary,
+            s0.real * g00.imaginary
+                + s0.imaginary * g00.real
+                + s1.real * g01.imaginary
+                + s1.imaginary * g01.real
+        );
+
+        let new1 = complex!(
+            s0.real * g10.real - s0.imaginary * g10.imaginary + s1.real * g11.real
+                - s1.imaginary * g11.imaginary,
+            s0.real * g10.imaginary
+                + s0.imaginary * g10.real
+                + s1.real * g11.imaginary
+                + s1.imaginary * g11.real
+        );
+
+        state[i] = new0;
+        state[j] = new1;
+    }
+}
+
+#[cfg(target_arch = "x86_64")]
+#[target_feature(enable = "avx512f", enable = "avx512dq")]
+unsafe fn apply_single_qubit_avx512(
+    state: &mut [Complex<f64>],
+    gate: &[[Complex<f64>; 2]; 2],
+    target: usize,
+    num_qubits: usize,
+) {
+    let target_bit = num_qubits - 1 - target;
+    let step = 1 << target_bit;
+    let dim = 1 << num_qubits;
+
+    let g00 = gate[0][0];
+    let g01 = gate[0][1];
+    let g10 = gate[1][0];
+    let g11 = gate[1][1];
+
+    let pairs: Vec<(usize, usize)> = (0..dim)
+        .filter(|&i| (i >> target_bit) & 1 == 0)
+        .map(|i| (i, i | step))
+        .collect();
+
+    let chunks = pairs.len() / 4;
+
+    for chunk_idx in 0..chunks {
+        let base = chunk_idx * 4;
+        let (i0, j0) = pairs[base];
+        let (i1, j1) = pairs[base + 1];
+        let (i2, j2) = pairs[base + 2];
+        let (i3, j3) = pairs[base + 3];
+
+        let s0_re = _mm512_set_pd(
+            state[j3].real,
+            state[i3].real,
+            state[j2].real,
+            state[i2].real,
+            state[j1].real,
+            state[i1].real,
+            state[j0].real,
+            state[i0].real,
+        );
+        let s0_im = _mm512_set_pd(
+            state[j3].imaginary,
+            state[i3].imaginary,
+            state[j2].imaginary,
+            state[i2].imaginary,
+            state[j1].imaginary,
+            state[i1].imaginary,
+            state[j0].imaginary,
+            state[i0].imaginary,
+        );
+
+        let g_re_0 = _mm512_set_pd(
+            g01.real, g00.real, g01.real, g00.real, g01.real, g00.real, g01.real, g00.real,
+        );
+        let g_im_0 = _mm512_set_pd(
+            g01.imaginary,
+            g00.imaginary,
+            g01.imaginary,
+            g00.imaginary,
+            g01.imaginary,
+            g00.imaginary,
+            g01.imaginary,
+            g00.imaginary,
+        );
+        let g_re_1 = _mm512_set_pd(
+            g11.real, g10.real, g11.real, g10.real, g11.real, g10.real, g11.real, g10.real,
+        );
+        let g_im_1 = _mm512_set_pd(
+            g11.imaginary,
+            g10.imaginary,
+            g11.imaginary,
+            g10.imaginary,
+            g11.imaginary,
+            g10.imaginary,
+            g11.imaginary,
+            g10.imaginary,
+        );
+
+        let prod0_re = _mm512_fmsub_pd(s0_re, g_re_0, _mm512_mul_pd(s0_im, g_im_0));
+        let prod0_im = _mm512_fmadd_pd(s0_re, g_im_0, _mm512_mul_pd(s0_im, g_re_0));
+        let prod1_re = _mm512_fmsub_pd(s0_re, g_re_1, _mm512_mul_pd(s0_im, g_im_1));
+        let prod1_im = _mm512_fmadd_pd(s0_re, g_im_1, _mm512_mul_pd(s0_im, g_re_1));
+
+        let mut res0_re = [0.0f64; 8];
+        let mut res0_im = [0.0f64; 8];
+        let mut res1_re = [0.0f64; 8];
+        let mut res1_im = [0.0f64; 8];
+
+        _mm512_storeu_pd(res0_re.as_mut_ptr(), prod0_re);
+        _mm512_storeu_pd(res0_im.as_mut_ptr(), prod0_im);
+        _mm512_storeu_pd(res1_re.as_mut_ptr(), prod1_re);
+        _mm512_storeu_pd(res1_im.as_mut_ptr(), prod1_im);
+
+        state[i0] = complex!(res0_re[0] + res0_re[1], res0_im[0] + res0_im[1]);
+        state[j0] = complex!(res1_re[0] + res1_re[1], res1_im[0] + res1_im[1]);
+        state[i1] = complex!(res0_re[2] + res0_re[3], res0_im[2] + res0_im[3]);
+        state[j1] = complex!(res1_re[2] + res1_re[3], res1_im[2] + res1_im[3]);
+        state[i2] = complex!(res0_re[4] + res0_re[5], res0_im[4] + res0_im[5]);
+        state[j2] = complex!(res1_re[4] + res1_re[5], res1_im[4] + res1_im[5]);
+        state[i3] = complex!(res0_re[6] + res0_re[7], res0_im[6] + res0_im[7]);
+        state[j3] = complex!(res1_re[6] + res1_re[7], res1_im[6] + res1_im[7]);
+    }
+
+    for &(i, j) in pairs.iter().skip(chunks * 4) {
+        let s0 = state[i];
+        let s1 = state[j];
+
+        let new0 = complex!(
+            s0.real * g00.real - s0.imaginary * g00.imaginary + s1.real * g01.real
+                - s1.imaginary * g01.imaginary,
+            s0.real * g00.imaginary
+                + s0.imaginary * g00.real
+                + s1.real * g01.imaginary
+                + s1.imaginary * g01.real
+        );
+
+        let new1 = complex!(
+            s0.real * g10.real - s0.imaginary * g10.imaginary + s1.real * g11.real
+                - s1.imaginary * g11.imaginary,
+            s0.real * g10.imaginary
+                + s0.imaginary * g10.real
+                + s1.real * g11.imaginary
+                + s1.imaginary * g11.real
+        );
+
+        state[i] = new0;
+        state[j] = new1;
+    }
+}
+
+#[cfg(target_arch = "aarch64")]
+unsafe fn apply_single_qubit_neon(
+    state: &mut [Complex<f64>],
+    gate: &[[Complex<f64>; 2]; 2],
+    target: usize,
+    num_qubits: usize,
+) {
+    let target_bit = num_qubits - 1 - target;
+    let step = 1 << target_bit;
+    let dim = 1 << num_qubits;
+
+    let g00 = gate[0][0];
+    let g01 = gate[0][1];
+    let g10 = gate[1][0];
+    let g11 = gate[1][1];
+
+    let pairs: Vec<(usize, usize)> = (0..dim)
+        .filter(|&i| (i >> target_bit) & 1 == 0)
+        .map(|i| (i, i | step))
+        .collect();
+
+    let chunks = pairs.len() / 2;
+
+    for chunk_idx in 0..chunks {
+        let (i0, j0) = pairs[chunk_idx * 2];
+        let (i1, j1) = pairs[chunk_idx * 2 + 1];
+
+        let s0_0 = state[i0];
+        let s1_0 = state[j0];
+        let s0_1 = state[i1];
+        let s1_1 = state[j1];
+
+        let s0_re = vld1q_f64([s0_0.real, s0_1.real].as_ptr());
+        let s0_im = vld1q_f64([s0_0.imaginary, s0_1.imaginary].as_ptr());
+        let s1_re = vld1q_f64([s1_0.real, s1_1.real].as_ptr());
+        let s1_im = vld1q_f64([s1_0.imaginary, s1_1.imaginary].as_ptr());
+
+        let g00_re = vdupq_n_f64(g00.real);
+        let g00_im = vdupq_n_f64(g00.imaginary);
+        let g01_re = vdupq_n_f64(g01.real);
+        let g01_im = vdupq_n_f64(g01.imaginary);
+        let g10_re = vdupq_n_f64(g10.real);
+        let g10_im = vdupq_n_f64(g10.imaginary);
+        let g11_re = vdupq_n_f64(g11.real);
+        let g11_im = vdupq_n_f64(g11.imaginary);
+
+        let new0_re = vaddq_f64(
+            vfmsq_f64(vmulq_f64(s0_re, g00_re), s0_im, g00_im),
+            vfmsq_f64(vmulq_f64(s1_re, g01_re), s1_im, g01_im),
+        );
+        let new0_im = vaddq_f64(
+            vfmaq_f64(vmulq_f64(s0_re, g00_im), s0_im, g00_re),
+            vfmaq_f64(vmulq_f64(s1_re, g01_im), s1_im, g01_re),
+        );
+
+        let new1_re = vaddq_f64(
+            vfmsq_f64(vmulq_f64(s0_re, g10_re), s0_im, g10_im),
+            vfmsq_f64(vmulq_f64(s1_re, g11_re), s1_im, g11_im),
+        );
+        let new1_im = vaddq_f64(
+            vfmaq_f64(vmulq_f64(s0_re, g10_im), s0_im, g10_re),
+            vfmaq_f64(vmulq_f64(s1_re, g11_im), s1_im, g11_re),
+        );
+
+        state[i0] = complex!(vgetq_lane_f64(new0_re, 0), vgetq_lane_f64(new0_im, 0));
+        state[j0] = complex!(vgetq_lane_f64(new1_re, 0), vgetq_lane_f64(new1_im, 0));
+        state[i1] = complex!(vgetq_lane_f64(new0_re, 1), vgetq_lane_f64(new0_im, 1));
+        state[j1] = complex!(vgetq_lane_f64(new1_re, 1), vgetq_lane_f64(new1_im, 1));
+    }
+
+    for &(i, j) in pairs.iter().skip(chunks * 2) {
+        let s0 = state[i];
+        let s1 = state[j];
+
+        let new0 = complex!(
+            s0.real * g00.real - s0.imaginary * g00.imaginary + s1.real * g01.real
+                - s1.imaginary * g01.imaginary,
+            s0.real * g00.imaginary
+                + s0.imaginary * g00.real
+                + s1.real * g01.imaginary
+                + s1.imaginary * g01.real
+        );
+
+        let new1 = complex!(
+            s0.real * g10.real - s0.imaginary * g10.imaginary + s1.real * g11.real
+                - s1.imaginary * g11.imaginary,
+            s0.real * g10.imaginary
+                + s0.imaginary * g10.real
+                + s1.real * g11.imaginary
+                + s1.imaginary * g11.real
+        );
+
+        state[i] = new0;
+        state[j] = new1;
+    }
+}
+
+fn apply_single_qubit_scalar(
+    state: &mut [Complex<f64>],
+    gate: &[[Complex<f64>; 2]; 2],
+    target: usize,
+    num_qubits: usize,
+) {
+    let target_bit = num_qubits - 1 - target;
+    let step = 1 << target_bit;
+    let dim = 1 << num_qubits;
+
+    let g00 = gate[0][0];
+    let g01 = gate[0][1];
+    let g10 = gate[1][0];
+    let g11 = gate[1][1];
+
+    for i in 0..dim {
+        if (i >> target_bit) & 1 == 1 {
+            continue;
+        }
+
+        let j = i | step;
+        let s0 = state[i];
+        let s1 = state[j];
+
+        let new0 = complex!(
+            s0.real * g00.real - s0.imaginary * g00.imaginary + s1.real * g01.real
+                - s1.imaginary * g01.imaginary,
+            s0.real * g00.imaginary
+                + s0.imaginary * g00.real
+                + s1.real * g01.imaginary
+                + s1.imaginary * g01.real
+        );
+
+        let new1 = complex!(
+            s0.real * g10.real - s0.imaginary * g10.imaginary + s1.real * g11.real
+                - s1.imaginary * g11.imaginary,
+            s0.real * g10.imaginary
+                + s0.imaginary * g10.real
+                + s1.real * g11.imaginary
+                + s1.imaginary * g11.real
+        );
+
+        state[i] = new0;
+        state[j] = new1;
+    }
+}
+
+pub fn apply_single_qubit_gate_simd_parallel(
+    state: &mut [Complex<f64>],
+    gate: &[[Complex<f64>; 2]; 2],
+    target: usize,
+    num_qubits: usize,
+) {
+    use rayon::prelude::*;
+
+    let target_bit = num_qubits - 1 - target;
+    let step = 1 << target_bit;
+    let dim = 1 << num_qubits;
+
+    let g00 = gate[0][0];
+    let g01 = gate[0][1];
+    let g10 = gate[1][0];
+    let g11 = gate[1][1];
+
+    let pairs: Vec<(usize, usize)> = (0..dim)
+        .filter(|&i| (i >> target_bit) & 1 == 0)
+        .map(|i| (i, i | step))
+        .collect();
+
+    let results: Vec<(usize, usize, Complex<f64>, Complex<f64>)> = pairs
+        .par_iter()
+        .map(|&(i, j)| {
+            let s0 = state[i];
+            let s1 = state[j];
+
+            let new0 = complex!(
+                s0.real * g00.real - s0.imaginary * g00.imaginary + s1.real * g01.real
+                    - s1.imaginary * g01.imaginary,
+                s0.real * g00.imaginary
+                    + s0.imaginary * g00.real
+                    + s1.real * g01.imaginary
+                    + s1.imaginary * g01.real
+            );
+
+            let new1 = complex!(
+                s0.real * g10.real - s0.imaginary * g10.imaginary + s1.real * g11.real
+                    - s1.imaginary * g11.imaginary,
+                s0.real * g10.imaginary
+                    + s0.imaginary * g10.real
+                    + s1.real * g11.imaginary
+                    + s1.imaginary * g11.real
+            );
+
+            (i, j, new0, new1)
+        })
+        .collect();
+
+    for (i, j, new0, new1) in results {
+        state[i] = new0;
+        state[j] = new1;
+    }
+}
+
+pub fn get_simd_info() -> String {
+    let cap = SimdCapability::detect();
+    format!("SIMD: {}", cap.name())
+}