aboutsummaryrefslogtreecommitdiff
path: root/libpsi-core/src/core/kernel.rs
diff options
context:
space:
mode:
Diffstat (limited to 'libpsi-core/src/core/kernel.rs')
-rw-r--r--libpsi-core/src/core/kernel.rs222
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
+ }
+}