1 files changed, 65 insertions, 25 deletions

diff --git a/src/core/kernel.c b/src/core/kernel.c
index 377982f..063d47f 100644
--- a/src/core/kernel.c
+++ b/src/core/kernel.c
@@ -6,6 +6,8 @@
 #include <stdlib.h>
 #include <string.h>
 
+#include "parallel.h"
+
 static char* dup_string(const char* s)
 {
     size_t n = strlen(s) + 1;
@@ -159,8 +161,49 @@ bool psi_fuse_kernels(struct PsiKernel a, struct PsiKernel b, struct PsiKernel*
     return true;
 }
 
-static struct PsiComplex* apply_kernel(const struct PsiComplex* state, struct PsiKernel kernel,
-                                       size_t num_qubits)
+struct PsiKernelApply
+{
+    const struct PsiComplex* state;
+    struct PsiComplex* new_state;
+    struct PsiMatrix matrix;
+    const size_t* target_bits;
+    size_t g;
+    size_t gate_dim;
+    size_t non_target_mask;
+};
+
+static void kernel_apply_range(size_t start, size_t end, void* vctx)
+{
+    struct PsiKernelApply* c = vctx;
+
+    for (size_t i = start; i < end; i++)
+    {
+        size_t target_idx = 0;
+        for (size_t k = 0; k < c->g; k++)
+            if ((i >> c->target_bits[k]) & 1)
+                target_idx |= (size_t)1 << (c->g - 1 - k);
+
+        struct PsiComplex sum = psi_new_complex(0.0, 0.0);
+        for (size_t j = 0; j < c->gate_dim; j++)
+        {
+            struct PsiComplex gate_elem = c->matrix.data[target_idx * c->gate_dim + j];
+            if (fabs(gate_elem.real) < 1e-15 && fabs(gate_elem.imaginary) < 1e-15)
+                continue;
+
+            size_t source_idx = i & c->non_target_mask;
+            for (size_t k = 0; k < c->g; k++)
+                if ((j >> (c->g - 1 - k)) & 1)
+                    source_idx |= (size_t)1 << c->target_bits[k];
+
+            sum = psi_add_complex(sum, psi_mul_complex(gate_elem, c->state[source_idx]));
+        }
+
+        c->new_state[i] = sum;
+    }
+}
+
+static struct PsiComplex* run_kernel(const struct PsiComplex* state, struct PsiKernel kernel,
+                                     size_t num_qubits, bool parallel)
 {
     size_t dim = (size_t)1 << num_qubits;
     size_t g = kernel.target_count;
@@ -178,35 +221,25 @@ static struct PsiComplex* apply_kernel(const struct PsiComplex* state, struct Ps
     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 PsiKernelApply ctx = {
+        state, new_state, kernel.matrix, target_bits, g, gate_dim, non_target_mask,
+    };
 
-        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;
-    }
+    if (parallel)
+        psi_parallel_for(dim, kernel_apply_range, &ctx);
+    else
+        kernel_apply_range(0, dim, &ctx);
 
     free(target_bits);
     return new_state;
 }
 
+static struct PsiComplex* apply_kernel(const struct PsiComplex* state, struct PsiKernel kernel,
+                                       size_t num_qubits)
+{
+    return run_kernel(state, kernel, num_qubits, false);
+}
+
 struct PsiKernelBatch psi_new_kernel_batch(size_t num_qubits)
 {
     struct PsiKernelBatch batch;
@@ -298,6 +331,13 @@ void psi_apply_kernel(struct PsiVector* state, struct PsiKernel kernel, size_t n
     state->data = next;
 }
 
+void psi_apply_kernel_parallel(struct PsiVector* state, struct PsiKernel kernel, size_t num_qubits)
+{
+    struct PsiComplex* next = run_kernel(state->data, kernel, num_qubits, true);
+    free(state->data);
+    state->data = next;
+}
+
 static void push_kernel(struct PsiKernel** kernels, size_t* count, size_t* capacity,
                         struct PsiKernel kernel)
 {