[add]: Kernel+Kernel Batching

1 files changed, 222 insertions, 0 deletions

diff --git a/libpsi-core/src/core/kernel.rs b/libpsi-core/src/core/kernel.rs
new file mode 100644
index 0000000..3d75425
--- /dev/null
+++ b/libpsi-core/src/core/kernel.rs
@@ -0,0 +1,222 @@
+use crate::{complex, Complex, Matrix};
+use rayon::prelude::*;
+
+#[derive(Clone)]
+pub struct Kernel {
+    pub matrix: Matrix<Complex<f64>>,
+    pub targets: Vec<usize>,
+    pub name: String,
+}
+
+impl Kernel {
+    pub fn new(name: &str, matrix: Matrix<Complex<f64>>, targets: Vec<usize>) -> Self {
+        Self {
+            matrix,
+            targets,
+            name: name.to_string(),
+        }
+    }
+
+    pub fn num_qubits(&self) -> usize {
+        self.targets.len()
+    }
+
+    pub fn can_fuse_with(&self, other: &Kernel) -> bool {
+        if self.targets.len() != 1 || other.targets.len() != 1 {
+            return false;
+        }
+        self.targets[0] == other.targets[0]
+    }
+
+    pub fn fuse(&self, other: &Kernel) -> Option<Kernel> {
+        if !self.can_fuse_with(other) {
+            return None;
+        }
+        let fused_matrix = other.matrix.dot(&self.matrix)?;
+        Some(Kernel {
+            matrix: fused_matrix,
+            targets: self.targets.clone(),
+            name: format!("{}+{}", self.name, other.name),
+        })
+    }
+}
+
+pub struct KernelBatch {
+    kernels: Vec<Kernel>,
+    num_qubits: usize,
+}
+
+impl KernelBatch {
+    pub fn new(num_qubits: usize) -> Self {
+        Self {
+            kernels: Vec::new(),
+            num_qubits,
+        }
+    }
+
+    pub fn add(&mut self, kernel: Kernel) {
+        self.kernels.push(kernel);
+    }
+
+    pub fn len(&self) -> usize {
+        self.kernels.len()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.kernels.is_empty()
+    }
+
+    pub fn kernels(&self) -> &[Kernel] {
+        &self.kernels
+    }
+
+    pub fn optimize(&mut self) {
+        if self.kernels.len() < 2 {
+            return;
+        }
+
+        let mut optimized: Vec<Kernel> = Vec::with_capacity(self.kernels.len());
+        let mut i = 0;
+
+        while i < self.kernels.len() {
+            let current = &self.kernels[i];
+
+            if i + 1 < self.kernels.len() {
+                let next = &self.kernels[i + 1];
+                if let Some(fused) = current.fuse(next) {
+                    optimized.push(fused);
+                    i += 2;
+                    continue;
+                }
+            }
+
+            optimized.push(current.clone());
+            i += 1;
+        }
+
+        self.kernels = optimized;
+    }
+
+    pub fn execute(&self, state: &mut Vec<Complex<f64>>) {
+        for kernel in &self.kernels {
+            *state = apply_kernel(state, kernel, self.num_qubits);
+        }
+    }
+
+    pub fn execute_parallel(&self, state: &mut Vec<Complex<f64>>) {
+        for kernel in &self.kernels {
+            *state = apply_kernel_parallel(state, kernel, self.num_qubits);
+        }
+    }
+}
+
+fn apply_kernel(state: &[Complex<f64>], kernel: &Kernel, num_qubits: usize) -> Vec<Complex<f64>> {
+    let dim = 1 << num_qubits;
+    let g = kernel.targets.len();
+    let gate_dim = 1 << g;
+
+    let target_bits: Vec<usize> = kernel.targets.iter().map(|&t| num_qubits - 1 - t).collect();
+
+    let mut non_target_mask: usize = (1 << num_qubits) - 1;
+    for &pos in &target_bits {
+        non_target_mask &= !(1 << pos);
+    }
+
+    let mut new_state = vec![complex!(0.0, 0.0); dim];
+
+    for i in 0..dim {
+        let mut target_idx = 0usize;
+        for (k, &pos) in target_bits.iter().enumerate() {
+            if (i >> pos) & 1 == 1 {
+                target_idx |= 1 << (g - 1 - k);
+            }
+        }
+
+        let mut sum = complex!(0.0, 0.0);
+
+        for j in 0..gate_dim {
+            let gate_elem = kernel.matrix.data[target_idx * gate_dim + j];
+
+            if gate_elem.real.abs() < 1e-15 && gate_elem.imaginary.abs() < 1e-15 {
+                continue;
+            }
+
+            let mut source_idx = i & non_target_mask;
+            for (k, &pos) in target_bits.iter().enumerate() {
+                if (j >> (g - 1 - k)) & 1 == 1 {
+                    source_idx |= 1 << pos;
+                }
+            }
+
+            sum = sum + gate_elem * state[source_idx];
+        }
+
+        new_state[i] = sum;
+    }
+
+    new_state
+}
+
+fn apply_kernel_parallel(
+    state: &[Complex<f64>],
+    kernel: &Kernel,
+    num_qubits: usize,
+) -> Vec<Complex<f64>> {
+    let dim = 1 << num_qubits;
+    let g = kernel.targets.len();
+    let gate_dim = 1 << g;
+
+    let target_bits: Vec<usize> = kernel.targets.iter().map(|&t| num_qubits - 1 - t).collect();
+
+    let mut non_target_mask: usize = (1 << num_qubits) - 1;
+    for &pos in &target_bits {
+        non_target_mask &= !(1 << pos);
+    }
+
+    (0..dim)
+        .into_par_iter()
+        .map(|i| {
+            let mut target_idx = 0usize;
+            for (k, &pos) in target_bits.iter().enumerate() {
+                if (i >> pos) & 1 == 1 {
+                    target_idx |= 1 << (g - 1 - k);
+                }
+            }
+
+            let mut sum = complex!(0.0, 0.0);
+
+            for j in 0..gate_dim {
+                let gate_elem = kernel.matrix.data[target_idx * gate_dim + j];
+
+                if gate_elem.real.abs() < 1e-15 && gate_elem.imaginary.abs() < 1e-15 {
+                    continue;
+                }
+
+                let mut source_idx = i & non_target_mask;
+                for (k, &pos) in target_bits.iter().enumerate() {
+                    if (j >> (g - 1 - k)) & 1 == 1 {
+                        source_idx |= 1 << pos;
+                    }
+                }
+
+                sum = sum + gate_elem * state[source_idx];
+            }
+
+            sum
+        })
+        .collect()
+}
+
+pub struct KernelBuilder {
+    num_qubits: usize,
+}
+
+impl KernelBuilder {
+    pub fn new(num_qubits: usize) -> Self {
+        Self { num_qubits }
+    }
+
+    pub fn num_qubits(&self) -> usize {
+        self.num_qubits
+    }
+}