microsoft/qdk
Publicmirrored fromhttps://github.com/microsoft/qdkAvailable
source/noisy_simulator/src/kernel.rs
119lines · modecode
| 1 | // Copyright (c) Microsoft Corporation. |
| 2 | // Licensed under the MIT License. |
| 3 | |
| 4 | //! This module contains the `apply_kernel` function used by the `DensityMatrixSimualtor` |
| 5 | //! and the `TrajectorySimulator`. |
| 6 | |
| 7 | use crate::{ComplexVector, Error, SquareMatrix}; |
| 8 | use nalgebra::Complex; |
| 9 | |
| 10 | /// This function extracts the relevant entries from the `state_vector` into its own vector. |
| 11 | /// Then it applies the `operation_matrix` to this extracted entries. |
| 12 | /// Finally it stores the results back into the state vector. |
| 13 | /// |
| 14 | /// Performance note: Because `nalgebra` stores its matrices in column major form, we apply |
| 15 | /// `gemv_tr` to avoid incurring cache misses. That is why we transpose all Kraus operators |
| 16 | /// when they enter the simulator: |
| 17 | /// `gemv(1, matrix, vec, 0)` is equivalent to `gemv_tr(1, matrix_tr, vec, 0)`, |
| 18 | /// but the later has much better performance. |
| 19 | /// |
| 20 | /// Errors: If the `operation_matrix` doesn't have the right dimension for the number of target `qubits`, |
| 21 | /// this function will return `Error::MatrixVecDimensionMismatch`. |
| 22 | pub fn apply_kernel( |
| 23 | state: &mut ComplexVector, |
| 24 | operation_matrix: &SquareMatrix, |
| 25 | qubits: &[usize], |
| 26 | ) -> Result<(), Error> { |
| 27 | // Construct a mask that has 1s at locations given by the target `qubits` ids. |
| 28 | let mask = make_mask(state, qubits); |
| 29 | |
| 30 | // Number of elements in small matrix-vector multiplications (dimension of gate matrix). |
| 31 | let num_elements: usize = 1 << qubits.len(); |
| 32 | let (nrows, ncols) = operation_matrix.shape(); |
| 33 | |
| 34 | if num_elements != ncols { |
| 35 | return Err(Error::MatrixVecDimensionMismatch { |
| 36 | nrows, |
| 37 | ncols, |
| 38 | vec_dim: num_elements, |
| 39 | }); |
| 40 | } |
| 41 | |
| 42 | // Compute all index offsets of the entries to load. |
| 43 | // E.g., for a 2-qubit gate acting on qubits [k1,k2], |
| 44 | // index offsets are [0, 2^k1, 2^k2, 2^(k1+k2)] |
| 45 | let mut index_offsets: Vec<usize> = std::vec::Vec::with_capacity(num_elements); |
| 46 | for k in 0..num_elements { |
| 47 | let mut j = 0; |
| 48 | let mut k_ = k; |
| 49 | let mut idx = 0; |
| 50 | while k_ != 0 { |
| 51 | idx |= (k_ & 1) << qubits[j]; |
| 52 | k_ >>= 1; |
| 53 | j += 1; |
| 54 | } |
| 55 | index_offsets.push(idx); |
| 56 | } |
| 57 | |
| 58 | // Main loop. |
| 59 | let mut extracted_entries = ComplexVector::zeros(num_elements); |
| 60 | let mut new_entries = ComplexVector::zeros(num_elements); |
| 61 | for s in 0..state.len() { |
| 62 | if (s & mask) == 0 { |
| 63 | // Extract relevant entries into a vector to make the gate application easier. |
| 64 | (0..index_offsets.len()).for_each(|k| { |
| 65 | // SAFETY: extracted_entries has size index_offset.len(), so that get_unchecked is safe. |
| 66 | // state.get_unchecked(idx) is safe because: |
| 67 | // 1. s has total_qubits_in_system bits |
| 68 | // 2. index_offset has total_qubits_in_system bits |
| 69 | // 3. Therefore, idx = s | index_offset also has total_qubits_in_system bits. |
| 70 | // 4. idx < (1 << total_qubits_in_system) = state.len() because |
| 71 | // it has (total_qubits_in_system + 1) bits. |
| 72 | // 5. Therefore state.get_unchecked(idx) is safe. |
| 73 | let idx = s | index_offsets[k]; |
| 74 | unsafe { |
| 75 | *extracted_entries.get_unchecked_mut(k) = *state.get_unchecked(idx); |
| 76 | } |
| 77 | }); |
| 78 | |
| 79 | // Apply the gate. |
| 80 | new_entries.gemv_tr( |
| 81 | Complex::ONE, |
| 82 | operation_matrix, |
| 83 | &extracted_entries, |
| 84 | Complex::ZERO, |
| 85 | ); |
| 86 | |
| 87 | // Store accumulated result back into the state vector. |
| 88 | (0..index_offsets.len()).for_each(|k| { |
| 89 | // SAFETY: new_entries has size index_offset.len(), so that get_unchecked is safe. |
| 90 | // state.get_unchecked(idx) is safe because: |
| 91 | // 1. s has total_qubits_in_system bits |
| 92 | // 2. index_offset has total_qubits_in_system bits |
| 93 | // 3. Therefore, idx = s | index_offset also has total_qubits_in_system bits. |
| 94 | // 4. idx < (1 << total_qubits_in_system) = state.len() because |
| 95 | // it has (total_qubits_in_system + 1) bits. |
| 96 | // 5. Therefore state.get_unchecked(idx) is safe. |
| 97 | let idx = s | index_offsets[k]; |
| 98 | unsafe { |
| 99 | *state.get_unchecked_mut(idx) = *new_entries.get_unchecked(k); |
| 100 | } |
| 101 | }); |
| 102 | } |
| 103 | } |
| 104 | |
| 105 | Ok(()) |
| 106 | } |
| 107 | |
| 108 | /// Construct a mask that has 1s at locations given by the target `qubits` ids. |
| 109 | fn make_mask(state: &ComplexVector, qubits: &[usize]) -> usize { |
| 110 | // Number of elements in the density matrix. |
| 111 | let num_elements = state.len(); |
| 112 | let mut mask: usize = 0; |
| 113 | for id in qubits { |
| 114 | let id_mask: usize = 1 << id; |
| 115 | assert!(id_mask < num_elements, "invalid qubit id: {id}"); |
| 116 | mask |= id_mask; |
| 117 | } |
| 118 | mask |
| 119 | } |
| 120 | |