add: scalar kernel engine with fusion

1 files changed, 287 insertions, 0 deletions

diff --git a/src/core/kernel.c b/src/core/kernel.c
new file mode 100644
index 0000000..fab3ea0
--- /dev/null
+++ b/src/core/kernel.c
@@ -0,0 +1,287 @@
+#include "core/kernel.h"
+
+#include <assert.h>
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+static char *dup_string(const char *s)
+{
+	size_t n = strlen(s) + 1;
+	char *p = malloc(n);
+	assert(p != NULL);
+	memcpy(p, s, n);
+
+	return p;
+}
+
+static size_t *dup_targets(const size_t *targets, size_t count)
+{
+	size_t *p = malloc(count * sizeof(size_t));
+	assert(p != NULL || count == 0);
+
+	if (count > 0)
+		memcpy(p, targets, count * sizeof(size_t));
+
+	return p;
+}
+
+static bool starts_with(const char *s, const char *prefix)
+{
+	return strncmp(s, prefix, strlen(prefix)) == 0;
+}
+
+static enum PsiGateType detect_gate_type(const char *name, struct PsiMatrix matrix)
+{
+	static const char *diagonal[] = { "Z", "S", "T", "Sdg", "Tdg", "Rz", "P", "U1", "CZ", "CP", "CRz" };
+	for (size_t i = 0; i < sizeof(diagonal) / sizeof(diagonal[0]); i++)
+		if (starts_with(name, diagonal[i]))
+			return PSI_GATE_TYPE_DIAGONAL;
+
+	static const char *controlled[] = { "CNOT", "CZ", "SWAP", "CRx", "CRy", "CRz", "CP", "CCNOT", "CSWAP" };
+	for (size_t i = 0; i < sizeof(controlled) / sizeof(controlled[0]); i++)
+		if (starts_with(name, controlled[i]))
+			return PSI_GATE_TYPE_CONTROLLED;
+
+	if (matrix.rows == 2 && matrix.cols == 2)
+	{
+		bool is_diag = fabs(matrix.data[1].real) < 1e-10
+			&& fabs(matrix.data[1].imaginary) < 1e-10
+			&& fabs(matrix.data[2].real) < 1e-10
+			&& fabs(matrix.data[2].imaginary) < 1e-10;
+		if (is_diag)
+			return PSI_GATE_TYPE_DIAGONAL;
+	}
+
+	return PSI_GATE_TYPE_NON_DIAGONAL;
+}
+
+struct PsiKernel psi_new_kernel(const char *name, struct PsiMatrix matrix, const size_t *targets, size_t target_count)
+{
+	struct PsiKernel kernel;
+	kernel.matrix = matrix;
+	kernel.targets = dup_targets(targets, target_count);
+	kernel.target_count = target_count;
+	kernel.name = dup_string(name);
+	kernel.gate_type = detect_gate_type(name, matrix);
+
+	return kernel;
+}
+
+struct PsiKernel psi_clone_kernel(struct PsiKernel kernel)
+{
+	struct PsiKernel copy;
+	copy.matrix = psi_clone_matrix(kernel.matrix);
+	copy.targets = dup_targets(kernel.targets, kernel.target_count);
+	copy.target_count = kernel.target_count;
+	copy.name = dup_string(kernel.name);
+	copy.gate_type = kernel.gate_type;
+
+	return copy;
+}
+
+void psi_free_kernel(struct PsiKernel *kernel)
+{
+	psi_free_matrix(&kernel->matrix);
+	free(kernel->targets);
+	free(kernel->name);
+	kernel->targets = NULL;
+	kernel->name = NULL;
+	kernel->target_count = 0;
+}
+
+static bool targets_equal(struct PsiKernel a, struct PsiKernel b)
+{
+	if (a.target_count != b.target_count)
+		return false;
+
+	for (size_t i = 0; i < a.target_count; i++)
+		if (a.targets[i] != b.targets[i])
+			return false;
+
+	return true;
+}
+
+bool psi_kernels_share_qubits(struct PsiKernel a, struct PsiKernel b)
+{
+	for (size_t i = 0; i < a.target_count; i++)
+		for (size_t j = 0; j < b.target_count; j++)
+			if (a.targets[i] == b.targets[j])
+				return true;
+
+	return false;
+}
+
+bool psi_kernels_commute(struct PsiKernel a, struct PsiKernel b)
+{
+	if (!psi_kernels_share_qubits(a, b))
+		return true;
+
+	if (a.gate_type == PSI_GATE_TYPE_DIAGONAL && b.gate_type == PSI_GATE_TYPE_DIAGONAL && targets_equal(a, b))
+		return true;
+
+	return false;
+}
+
+bool psi_kernels_can_fuse(struct PsiKernel a, struct PsiKernel b)
+{
+	if (a.target_count != 1 || b.target_count != 1)
+		return false;
+
+	return a.targets[0] == b.targets[0];
+}
+
+bool psi_fuse_kernels(struct PsiKernel a, struct PsiKernel b, struct PsiKernel *out)
+{
+	if (!psi_kernels_can_fuse(a, b))
+		return false;
+
+	enum PsiGateType new_type = a.gate_type == PSI_GATE_TYPE_DIAGONAL && b.gate_type == PSI_GATE_TYPE_DIAGONAL
+		? PSI_GATE_TYPE_DIAGONAL
+		: PSI_GATE_TYPE_NON_DIAGONAL;
+
+	size_t name_len = strlen(a.name) + strlen(b.name) + 2;
+	char *fused_name = malloc(name_len);
+	assert(fused_name != NULL);
+	snprintf(fused_name, name_len, "%s+%s", a.name, b.name);
+
+	out->matrix = psi_dot_matrix(b.matrix, a.matrix);
+	out->targets = dup_targets(a.targets, a.target_count);
+	out->target_count = a.target_count;
+	out->name = fused_name;
+	out->gate_type = new_type;
+
+	return true;
+}
+
+static struct PsiComplex *apply_kernel(const struct PsiComplex *state, struct PsiKernel kernel, size_t num_qubits)
+{
+	size_t dim = (size_t)1 << num_qubits;
+	size_t g = kernel.target_count;
+	size_t gate_dim = (size_t)1 << g;
+
+	size_t *target_bits = malloc(g * sizeof(size_t));
+	assert(target_bits != NULL || g == 0);
+	for (size_t k = 0; k < g; k++)
+		target_bits[k] = num_qubits - 1 - kernel.targets[k];
+
+	size_t non_target_mask = dim - 1;
+	for (size_t k = 0; k < g; k++)
+		non_target_mask &= ~((size_t)1 << target_bits[k]);
+
+	struct PsiComplex *new_state = malloc(dim * sizeof(struct PsiComplex));
+	assert(new_state != NULL);
+
+	for (size_t i = 0; i < dim; i++)
+	{
+		size_t target_idx = 0;
+		for (size_t k = 0; k < g; k++)
+			if ((i >> target_bits[k]) & 1)
+				target_idx |= (size_t)1 << (g - 1 - k);
+
+		struct PsiComplex sum = psi_new_complex(0.0, 0.0);
+		for (size_t j = 0; j < gate_dim; j++)
+		{
+			struct PsiComplex gate_elem = kernel.matrix.data[target_idx * gate_dim + j];
+			if (fabs(gate_elem.real) < 1e-15 && fabs(gate_elem.imaginary) < 1e-15)
+				continue;
+
+			size_t source_idx = i & non_target_mask;
+			for (size_t k = 0; k < g; k++)
+				if ((j >> (g - 1 - k)) & 1)
+					source_idx |= (size_t)1 << target_bits[k];
+
+			sum = psi_add_complex(sum, psi_mul_complex(gate_elem, state[source_idx]));
+		}
+
+		new_state[i] = sum;
+	}
+
+	free(target_bits);
+	return new_state;
+}
+
+struct PsiKernelBatch psi_new_kernel_batch(size_t num_qubits)
+{
+	struct PsiKernelBatch batch;
+	batch.kernels = NULL;
+	batch.count = 0;
+	batch.capacity = 0;
+	batch.num_qubits = num_qubits;
+
+	return batch;
+}
+
+void psi_free_kernel_batch(struct PsiKernelBatch *batch)
+{
+	for (size_t i = 0; i < batch->count; i++)
+		psi_free_kernel(&batch->kernels[i]);
+
+	free(batch->kernels);
+	batch->kernels = NULL;
+	batch->count = 0;
+	batch->capacity = 0;
+}
+
+void psi_add_kernel(struct PsiKernelBatch *batch, struct PsiKernel kernel)
+{
+	if (batch->count == batch->capacity)
+	{
+		size_t new_capacity = batch->capacity == 0 ? 8 : batch->capacity * 2;
+		batch->kernels = realloc(batch->kernels, new_capacity * sizeof(struct PsiKernel));
+		assert(batch->kernels != NULL);
+		batch->capacity = new_capacity;
+	}
+
+	batch->kernels[batch->count++] = kernel;
+}
+
+void psi_optimize_kernel_batch(struct PsiKernelBatch *batch)
+{
+	if (batch->count < 2)
+		return;
+
+	size_t original = batch->count;
+	struct PsiKernel *out = malloc(original * sizeof(struct PsiKernel));
+	assert(out != NULL);
+	size_t out_count = 0;
+
+	size_t i = 0;
+	while (i < batch->count)
+	{
+		if (i + 1 < batch->count)
+		{
+			struct PsiKernel fused;
+			if (psi_fuse_kernels(batch->kernels[i], batch->kernels[i + 1], &fused))
+			{
+				psi_free_kernel(&batch->kernels[i]);
+				psi_free_kernel(&batch->kernels[i + 1]);
+				out[out_count++] = fused;
+				i += 2;
+				continue;
+			}
+		}
+
+		out[out_count++] = batch->kernels[i];
+		i += 1;
+	}
+
+	free(batch->kernels);
+	batch->kernels = out;
+	batch->count = out_count;
+	batch->capacity = original;
+}
+
+void psi_execute_kernel_batch(struct PsiKernelBatch batch, struct PsiVector *state)
+{
+	size_t dim = (size_t)1 << batch.num_qubits;
+	assert(state->size == dim);
+
+	for (size_t i = 0; i < batch.count; i++)
+	{
+		struct PsiComplex *next = apply_kernel(state->data, batch.kernels[i], batch.num_qubits);
+		free(state->data);
+		state->data = next;
+	}
+}