diff options
Diffstat (limited to 'src/maths')
| -rw-r--r-- | src/maths/simd.rs | 510 |
1 files changed, 0 insertions, 510 deletions
diff --git a/src/maths/simd.rs b/src/maths/simd.rs deleted file mode 100644 index de0370c..0000000 --- a/src/maths/simd.rs +++ /dev/null @@ -1,510 +0,0 @@ -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()) -} |