add: structure-aware kernel batch

1 files changed, 301 insertions, 0 deletions

diff --git a/src/core/kernel.c b/src/core/kernel.c
index fab3ea0..0a92ee1 100644
--- a/src/core/kernel.c
+++ b/src/core/kernel.c
@@ -285,3 +285,304 @@ void psi_execute_kernel_batch(struct PsiKernelBatch batch, struct PsiVector *sta
 		state->data = next;
 	}
 }
+
+static void push_kernel(struct PsiKernel **kernels, size_t *count, size_t *capacity, struct PsiKernel kernel)
+{
+	if (*count == *capacity)
+	{
+		size_t new_capacity = *capacity == 0 ? 8 : *capacity * 2;
+		*kernels = realloc(*kernels, new_capacity * sizeof(struct PsiKernel));
+		assert(*kernels != NULL);
+		*capacity = new_capacity;
+	}
+
+	(*kernels)[(*count)++] = kernel;
+}
+
+static bool layer_can_add(struct PsiExecutionLayer layer, struct PsiKernel kernel)
+{
+	for (size_t i = 0; i < layer.count; i++)
+		if (psi_kernels_share_qubits(layer.kernels[i], kernel))
+			return false;
+
+	return true;
+}
+
+static void free_layer(struct PsiExecutionLayer *layer)
+{
+	for (size_t i = 0; i < layer->count; i++)
+		psi_free_kernel(&layer->kernels[i]);
+
+	free(layer->kernels);
+	layer->kernels = NULL;
+	layer->count = 0;
+	layer->capacity = 0;
+}
+
+struct PsiStructureAwareBatch psi_new_structure_aware_batch(size_t num_qubits)
+{
+	struct PsiStructureAwareBatch batch;
+	batch.kernels = NULL;
+	batch.count = 0;
+	batch.capacity = 0;
+	batch.layers = NULL;
+	batch.layer_count = 0;
+	batch.layer_capacity = 0;
+	batch.num_qubits = num_qubits;
+	batch.optimised = false;
+
+	return batch;
+}
+
+static void clear_layers(struct PsiStructureAwareBatch *batch)
+{
+	for (size_t i = 0; i < batch->layer_count; i++)
+		free_layer(&batch->layers[i]);
+
+	free(batch->layers);
+	batch->layers = NULL;
+	batch->layer_count = 0;
+	batch->layer_capacity = 0;
+}
+
+void psi_free_structure_aware_batch(struct PsiStructureAwareBatch *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;
+	clear_layers(batch);
+}
+
+void psi_add_structure_aware_kernel(struct PsiStructureAwareBatch *batch, struct PsiKernel kernel)
+{
+	push_kernel(&batch->kernels, &batch->count, &batch->capacity, kernel);
+	batch->optimised = false;
+}
+
+static struct PsiKernel remove_kernel_at(struct PsiStructureAwareBatch *batch, size_t index)
+{
+	struct PsiKernel removed = batch->kernels[index];
+	for (size_t i = index; i + 1 < batch->count; i++)
+		batch->kernels[i] = batch->kernels[i + 1];
+
+	batch->count--;
+	return removed;
+}
+
+static void insert_kernel_at(struct PsiStructureAwareBatch *batch, size_t index, 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;
+	}
+
+	for (size_t i = batch->count; i > index; i--)
+		batch->kernels[i] = batch->kernels[i - 1];
+
+	batch->kernels[index] = kernel;
+	batch->count++;
+}
+
+static void reorder_commuting_gates(struct PsiStructureAwareBatch *batch)
+{
+	bool changed = true;
+	size_t iterations = 0;
+	const size_t MAX_ITERATIONS = 100;
+
+	while (changed && iterations < MAX_ITERATIONS)
+	{
+		changed = false;
+		iterations++;
+
+		for (size_t i = 0; i + 1 < batch->count; i++)
+		{
+			struct PsiKernel current = batch->kernels[i];
+			struct PsiKernel next = batch->kernels[i + 1];
+
+			if (current.target_count != 1 || next.target_count != 1
+				|| current.targets[0] == next.targets[0]
+				|| !psi_kernels_commute(current, next))
+				continue;
+
+			for (size_t j = i + 2; j < batch->count; j++)
+			{
+				struct PsiKernel candidate = batch->kernels[j];
+				if (candidate.target_count != 1 || candidate.targets[0] != current.targets[0])
+					continue;
+
+				bool can_move = true;
+				for (size_t k = i + 1; k < j; k++)
+				{
+					struct PsiKernel between = batch->kernels[k];
+					if (psi_kernels_share_qubits(between, current) && !psi_kernels_commute(current, between))
+					{
+						can_move = false;
+						break;
+					}
+				}
+
+				if (can_move && psi_kernels_can_fuse(current, candidate))
+				{
+					struct PsiKernel moved = remove_kernel_at(batch, j);
+					insert_kernel_at(batch, i + 1, moved);
+					changed = true;
+					break;
+				}
+			}
+		}
+	}
+}
+
+static void multi_pass_fusion(struct PsiStructureAwareBatch *batch)
+{
+	bool changed = true;
+	size_t iterations = 0;
+	const size_t MAX_ITERATIONS = 50;
+
+	while (changed && iterations < MAX_ITERATIONS)
+	{
+		changed = false;
+		iterations++;
+
+		struct PsiKernel *new_kernels = NULL;
+		size_t new_count = 0;
+		size_t new_capacity = 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]);
+					push_kernel(&new_kernels, &new_count, &new_capacity, fused);
+					i += 2;
+					changed = true;
+					continue;
+				}
+			}
+
+			push_kernel(&new_kernels, &new_count, &new_capacity, batch->kernels[i]);
+			i++;
+		}
+
+		free(batch->kernels);
+		batch->kernels = new_kernels;
+		batch->count = new_count;
+		batch->capacity = new_capacity;
+	}
+}
+
+static void build_execution_layers(struct PsiStructureAwareBatch *batch)
+{
+	clear_layers(batch);
+
+	for (size_t i = 0; i < batch->count; i++)
+	{
+		struct PsiKernel kernel = batch->kernels[i];
+		bool placed = false;
+
+		for (size_t l = 0; l < batch->layer_count; l++)
+			if (layer_can_add(batch->layers[l], kernel))
+			{
+				struct PsiExecutionLayer *layer = &batch->layers[l];
+				push_kernel(&layer->kernels, &layer->count, &layer->capacity, psi_clone_kernel(kernel));
+				placed = true;
+				break;
+			}
+
+		if (placed)
+			continue;
+
+		if (batch->layer_count == batch->layer_capacity)
+		{
+			size_t new_capacity = batch->layer_capacity == 0 ? 4 : batch->layer_capacity * 2;
+			batch->layers = realloc(batch->layers, new_capacity * sizeof(struct PsiExecutionLayer));
+			assert(batch->layers != NULL);
+			batch->layer_capacity = new_capacity;
+		}
+
+		struct PsiExecutionLayer layer;
+		layer.kernels = NULL;
+		layer.count = 0;
+		layer.capacity = 0;
+		push_kernel(&layer.kernels, &layer.count, &layer.capacity, psi_clone_kernel(kernel));
+		batch->layers[batch->layer_count++] = layer;
+	}
+}
+
+void psi_optimize_structure_aware_batch(struct PsiStructureAwareBatch *batch)
+{
+	if (batch->optimised || batch->count < 2)
+		return;
+
+	reorder_commuting_gates(batch);
+	multi_pass_fusion(batch);
+	build_execution_layers(batch);
+	batch->optimised = true;
+}
+
+void psi_execute_structure_aware_batch(struct PsiStructureAwareBatch 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;
+	}
+}
+
+void psi_execute_structure_aware_batch_layered(struct PsiStructureAwareBatch batch, struct PsiVector *state)
+{
+	size_t dim = (size_t)1 << batch.num_qubits;
+	assert(state->size == dim);
+
+	for (size_t l = 0; l < batch.layer_count; l++)
+		for (size_t k = 0; k < batch.layers[l].count; k++)
+		{
+			struct PsiComplex *next = apply_kernel(state->data, batch.layers[l].kernels[k], batch.num_qubits);
+			free(state->data);
+			state->data = next;
+		}
+}
+
+struct PsiKernelStats psi_structure_aware_batch_stats(struct PsiStructureAwareBatch batch)
+{
+	struct PsiKernelStats stats;
+	stats.total_kernels = batch.count;
+	stats.single_qubit = 0;
+	stats.two_qubit = 0;
+	stats.multi_qubit = 0;
+	stats.diagonal = 0;
+	stats.execution_layers = batch.layer_count;
+
+	for (size_t i = 0; i < batch.count; i++)
+	{
+		struct PsiKernel kernel = batch.kernels[i];
+
+		if (kernel.target_count == 1)
+			stats.single_qubit++;
+		else if (kernel.target_count == 2)
+			stats.two_qubit++;
+		else if (kernel.target_count > 2)
+			stats.multi_qubit++;
+
+		if (kernel.gate_type == PSI_GATE_TYPE_DIAGONAL)
+			stats.diagonal++;
+	}
+
+	return stats;
+}