feat: simd + testing + formatting

1 files changed, 324 insertions, 0 deletions

diff --git a/src/maths/simd.c b/src/maths/simd.c
new file mode 100644
index 0000000..9e83dcc
--- /dev/null
+++ b/src/maths/simd.c
@@ -0,0 +1,324 @@
+#include "maths/simd.h"
+
+#include <assert.h>
+#include <stdlib.h>
+
+#if defined(__x86_64__) || defined(__i386__)
+#include <immintrin.h>
+#elif defined(__aarch64__)
+#include <arm_neon.h>
+#endif
+
+enum PsiSimdCapability psi_detect_simd(void)
+{
+#if defined(__x86_64__) || defined(__i386__)
+	if (__builtin_cpu_supports("avx512f") && __builtin_cpu_supports("avx512dq"))
+		return PSI_SIMD_AVX512;
+	if (__builtin_cpu_supports("avx2") && __builtin_cpu_supports("fma"))
+		return PSI_SIMD_AVX2;
+
+	return PSI_SIMD_NONE;
+#elif defined(__aarch64__)
+	return PSI_SIMD_NEON;
+#else
+	return PSI_SIMD_NONE;
+#endif
+}
+
+const char* psi_simd_name(enum PsiSimdCapability cap)
+{
+	switch (cap)
+	{
+		case PSI_SIMD_NONE: return "Scalar";
+		case PSI_SIMD_AVX2: return "AVX2+FMA";
+		case PSI_SIMD_AVX512: return "AVX-512";
+		case PSI_SIMD_NEON: return "NEON";
+	}
+
+	return "Scalar";
+}
+
+static void apply_pair(struct PsiComplex* state, size_t i, size_t j, struct PsiComplex g00,
+                       struct PsiComplex g01, struct PsiComplex g10, struct PsiComplex g11)
+{
+	struct PsiComplex s0 = state[i];
+	struct PsiComplex s1 = state[j];
+
+	state[i] = psi_new_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);
+	state[j] = psi_new_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);
+}
+
+static void apply_scalar(struct PsiComplex* state, const struct PsiComplex gate[2][2],
+                         size_t target, size_t num_qubits)
+{
+	size_t target_bit = num_qubits - 1 - target;
+	size_t step = (size_t)1 << target_bit;
+	size_t dim = (size_t)1 << num_qubits;
+
+	for (size_t i = 0; i < dim; i++)
+	{
+		if ((i >> target_bit) & 1)
+			continue;
+
+		apply_pair(state, i, i | step, gate[0][0], gate[0][1], gate[1][0], gate[1][1]);
+	}
+}
+
+static size_t (*build_pairs(size_t dim, size_t target_bit, size_t step, size_t* out_count))[2]
+{
+	size_t (*pairs)[2] = malloc((dim / 2) * sizeof(*pairs));
+	assert(pairs != NULL || dim == 0);
+
+	size_t n = 0;
+	for (size_t i = 0; i < dim; i++)
+		if (((i >> target_bit) & 1) == 0)
+		{
+			pairs[n][0] = i;
+			pairs[n][1] = i | step;
+			n++;
+		}
+
+	*out_count = n;
+	return pairs;
+}
+
+#if defined(__aarch64__)
+static void apply_neon(struct PsiComplex* state, const struct PsiComplex gate[2][2], size_t target,
+                       size_t num_qubits)
+{
+	size_t target_bit = num_qubits - 1 - target;
+	size_t step = (size_t)1 << target_bit;
+	size_t dim = (size_t)1 << num_qubits;
+
+	struct PsiComplex g00 = gate[0][0];
+	struct PsiComplex g01 = gate[0][1];
+	struct PsiComplex g10 = gate[1][0];
+	struct PsiComplex g11 = gate[1][1];
+
+	size_t np;
+	size_t (*pairs)[2] = build_pairs(dim, target_bit, step, &np);
+	size_t chunks = np / 2;
+
+	float64x2_t g00_re = vdupq_n_f64(g00.real);
+	float64x2_t g00_im = vdupq_n_f64(g00.imaginary);
+	float64x2_t g01_re = vdupq_n_f64(g01.real);
+	float64x2_t g01_im = vdupq_n_f64(g01.imaginary);
+	float64x2_t g10_re = vdupq_n_f64(g10.real);
+	float64x2_t g10_im = vdupq_n_f64(g10.imaginary);
+	float64x2_t g11_re = vdupq_n_f64(g11.real);
+	float64x2_t g11_im = vdupq_n_f64(g11.imaginary);
+
+	for (size_t chunk = 0; chunk < chunks; chunk++)
+	{
+		size_t i0 = pairs[chunk * 2][0];
+		size_t j0 = pairs[chunk * 2][1];
+		size_t i1 = pairs[chunk * 2 + 1][0];
+		size_t j1 = pairs[chunk * 2 + 1][1];
+
+		double s0r[2] = { state[i0].real, state[i1].real };
+		double s0i[2] = { state[i0].imaginary, state[i1].imaginary };
+		double s1r[2] = { state[j0].real, state[j1].real };
+		double s1i[2] = { state[j0].imaginary, state[j1].imaginary };
+
+		float64x2_t s0_re = vld1q_f64(s0r);
+		float64x2_t s0_im = vld1q_f64(s0i);
+		float64x2_t s1_re = vld1q_f64(s1r);
+		float64x2_t s1_im = vld1q_f64(s1i);
+
+		float64x2_t 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));
+		float64x2_t 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));
+		float64x2_t 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));
+		float64x2_t 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] = psi_new_complex(vgetq_lane_f64(new0_re, 0), vgetq_lane_f64(new0_im, 0));
+		state[j0] = psi_new_complex(vgetq_lane_f64(new1_re, 0), vgetq_lane_f64(new1_im, 0));
+		state[i1] = psi_new_complex(vgetq_lane_f64(new0_re, 1), vgetq_lane_f64(new0_im, 1));
+		state[j1] = psi_new_complex(vgetq_lane_f64(new1_re, 1), vgetq_lane_f64(new1_im, 1));
+	}
+
+	for (size_t p = chunks * 2; p < np; p++)
+		apply_pair(state, pairs[p][0], pairs[p][1], g00, g01, g10, g11);
+
+	free(pairs);
+}
+#endif
+
+#if defined(__x86_64__) || defined(__i386__)
+__attribute__((target("avx2,fma"))) static void apply_avx2(struct PsiComplex* state,
+                                                           const struct PsiComplex gate[2][2],
+                                                           size_t target, size_t num_qubits)
+{
+	size_t target_bit = num_qubits - 1 - target;
+	size_t step = (size_t)1 << target_bit;
+	size_t dim = (size_t)1 << num_qubits;
+
+	struct PsiComplex g00 = gate[0][0];
+	struct PsiComplex g01 = gate[0][1];
+	struct PsiComplex g10 = gate[1][0];
+	struct PsiComplex g11 = gate[1][1];
+
+	size_t np;
+	size_t (*pairs)[2] = build_pairs(dim, target_bit, step, &np);
+	size_t chunks = np / 2;
+
+	for (size_t chunk = 0; chunk < chunks; chunk++)
+	{
+		size_t i0 = pairs[chunk * 2][0];
+		size_t j0 = pairs[chunk * 2][1];
+		size_t i1 = pairs[chunk * 2 + 1][0];
+		size_t j1 = pairs[chunk * 2 + 1][1];
+
+		__m256d s0_re =
+				_mm256_set_pd(state[j1].real, state[i1].real, state[j0].real, state[i0].real);
+		__m256d s0_im = _mm256_set_pd(state[j1].imaginary, state[i1].imaginary, state[j0].imaginary,
+		                              state[i0].imaginary);
+
+		__m256d g_re_0 = _mm256_set_pd(g01.real, g00.real, g01.real, g00.real);
+		__m256d g_im_0 = _mm256_set_pd(g01.imaginary, g00.imaginary, g01.imaginary, g00.imaginary);
+		__m256d g_re_1 = _mm256_set_pd(g11.real, g10.real, g11.real, g10.real);
+		__m256d g_im_1 = _mm256_set_pd(g11.imaginary, g10.imaginary, g11.imaginary, g10.imaginary);
+
+		__m256d prod0_re = _mm256_fmsub_pd(s0_re, g_re_0, _mm256_mul_pd(s0_im, g_im_0));
+		__m256d prod0_im = _mm256_fmadd_pd(s0_re, g_im_0, _mm256_mul_pd(s0_im, g_re_0));
+		__m256d prod1_re = _mm256_fmsub_pd(s0_re, g_re_1, _mm256_mul_pd(s0_im, g_im_1));
+		__m256d prod1_im = _mm256_fmadd_pd(s0_re, g_im_1, _mm256_mul_pd(s0_im, g_re_1));
+
+		double res0_re[4];
+		double res0_im[4];
+		double res1_re[4];
+		double res1_im[4];
+
+		_mm256_storeu_pd(res0_re, prod0_re);
+		_mm256_storeu_pd(res0_im, prod0_im);
+		_mm256_storeu_pd(res1_re, prod1_re);
+		_mm256_storeu_pd(res1_im, prod1_im);
+
+		state[i0] = psi_new_complex(res0_re[0] + res0_re[1], res0_im[0] + res0_im[1]);
+		state[j0] = psi_new_complex(res1_re[0] + res1_re[1], res1_im[0] + res1_im[1]);
+		state[i1] = psi_new_complex(res0_re[2] + res0_re[3], res0_im[2] + res0_im[3]);
+		state[j1] = psi_new_complex(res1_re[2] + res1_re[3], res1_im[2] + res1_im[3]);
+	}
+
+	for (size_t p = chunks * 2; p < np; p++)
+		apply_pair(state, pairs[p][0], pairs[p][1], g00, g01, g10, g11);
+
+	free(pairs);
+}
+
+__attribute__((target("avx512f,avx512dq"))) static void
+apply_avx512(struct PsiComplex* state, const struct PsiComplex gate[2][2], size_t target,
+             size_t num_qubits)
+{
+	size_t target_bit = num_qubits - 1 - target;
+	size_t step = (size_t)1 << target_bit;
+	size_t dim = (size_t)1 << num_qubits;
+
+	struct PsiComplex g00 = gate[0][0];
+	struct PsiComplex g01 = gate[0][1];
+	struct PsiComplex g10 = gate[1][0];
+	struct PsiComplex g11 = gate[1][1];
+
+	size_t np;
+	size_t (*pairs)[2] = build_pairs(dim, target_bit, step, &np);
+	size_t chunks = np / 4;
+
+	for (size_t chunk = 0; chunk < chunks; chunk++)
+	{
+		size_t base = chunk * 4;
+		size_t i0 = pairs[base][0], j0 = pairs[base][1];
+		size_t i1 = pairs[base + 1][0], j1 = pairs[base + 1][1];
+		size_t i2 = pairs[base + 2][0], j2 = pairs[base + 2][1];
+		size_t i3 = pairs[base + 3][0], j3 = pairs[base + 3][1];
+
+		__m512d 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);
+		__m512d 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);
+
+		__m512d g_re_0 = _mm512_set_pd(g01.real, g00.real, g01.real, g00.real, g01.real, g00.real,
+		                               g01.real, g00.real);
+		__m512d g_im_0 = _mm512_set_pd(g01.imaginary, g00.imaginary, g01.imaginary, g00.imaginary,
+		                               g01.imaginary, g00.imaginary, g01.imaginary, g00.imaginary);
+		__m512d g_re_1 = _mm512_set_pd(g11.real, g10.real, g11.real, g10.real, g11.real, g10.real,
+		                               g11.real, g10.real);
+		__m512d g_im_1 = _mm512_set_pd(g11.imaginary, g10.imaginary, g11.imaginary, g10.imaginary,
+		                               g11.imaginary, g10.imaginary, g11.imaginary, g10.imaginary);
+
+		__m512d prod0_re = _mm512_fmsub_pd(s0_re, g_re_0, _mm512_mul_pd(s0_im, g_im_0));
+		__m512d prod0_im = _mm512_fmadd_pd(s0_re, g_im_0, _mm512_mul_pd(s0_im, g_re_0));
+		__m512d prod1_re = _mm512_fmsub_pd(s0_re, g_re_1, _mm512_mul_pd(s0_im, g_im_1));
+		__m512d prod1_im = _mm512_fmadd_pd(s0_re, g_im_1, _mm512_mul_pd(s0_im, g_re_1));
+
+		double res0_re[8];
+		double res0_im[8];
+		double res1_re[8];
+		double res1_im[8];
+
+		_mm512_storeu_pd(res0_re, prod0_re);
+		_mm512_storeu_pd(res0_im, prod0_im);
+		_mm512_storeu_pd(res1_re, prod1_re);
+		_mm512_storeu_pd(res1_im, prod1_im);
+
+		state[i0] = psi_new_complex(res0_re[0] + res0_re[1], res0_im[0] + res0_im[1]);
+		state[j0] = psi_new_complex(res1_re[0] + res1_re[1], res1_im[0] + res1_im[1]);
+		state[i1] = psi_new_complex(res0_re[2] + res0_re[3], res0_im[2] + res0_im[3]);
+		state[j1] = psi_new_complex(res1_re[2] + res1_re[3], res1_im[2] + res1_im[3]);
+		state[i2] = psi_new_complex(res0_re[4] + res0_re[5], res0_im[4] + res0_im[5]);
+		state[j2] = psi_new_complex(res1_re[4] + res1_re[5], res1_im[4] + res1_im[5]);
+		state[i3] = psi_new_complex(res0_re[6] + res0_re[7], res0_im[6] + res0_im[7]);
+		state[j3] = psi_new_complex(res1_re[6] + res1_re[7], res1_im[6] + res1_im[7]);
+	}
+
+	for (size_t p = chunks * 4; p < np; p++)
+		apply_pair(state, pairs[p][0], pairs[p][1], g00, g01, g10, g11);
+
+	free(pairs);
+}
+#endif
+
+void psi_apply_single_qubit_gate_simd(struct PsiComplex* state, const struct PsiComplex gate[2][2],
+                                      size_t target, size_t num_qubits)
+{
+	enum PsiSimdCapability cap = psi_detect_simd();
+
+#if defined(__x86_64__) || defined(__i386__)
+	if (cap == PSI_SIMD_AVX512)
+	{
+		apply_avx512(state, gate, target, num_qubits);
+		return;
+	}
+	if (cap == PSI_SIMD_AVX2)
+	{
+		apply_avx2(state, gate, target, num_qubits);
+		return;
+	}
+#elif defined(__aarch64__)
+	if (cap == PSI_SIMD_NEON)
+	{
+		apply_neon(state, gate, target, num_qubits);
+		return;
+	}
+#endif
+
+	(void)cap;
+	apply_scalar(state, gate, target, num_qubits);
+}
+
+void psi_apply_single_qubit_gate_simd_parallel(struct PsiComplex* state,
+                                               const struct PsiComplex gate[2][2], size_t target,
+                                               size_t num_qubits)
+{
+	apply_scalar(state, gate, target, num_qubits);
+}