aboutsummaryrefslogtreecommitdiff
path: root/libpsi-core/src/maths/simd.rs
diff options
context:
space:
mode:
Diffstat (limited to 'libpsi-core/src/maths/simd.rs')
-rw-r--r--libpsi-core/src/maths/simd.rs510
1 files changed, 510 insertions, 0 deletions
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())
+}