microsoft/qdk
Publicmirrored from https://github.com/microsoft/qdkAvailable
source/noisy_simulator/src/kernel.rs
119lines · modeblame
312af694orpuente-MS2 years ago | 1 | // Copyright (c) Microsoft Corporation. |
| 2 | // Licensed under the MIT License. | |
| 3 | | |
3a78aa85orpuente-MS2 years ago | 4 | //! This module contains the `apply_kernel` function used by the `DensityMatrixSimualtor` |
312af694orpuente-MS2 years ago | 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 | /// | |
84a8c59aorpuente-MS2 years ago | 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 | /// | |
312af694orpuente-MS2 years ago | 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. | |
84a8c59aorpuente-MS2 years ago | 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. | |
312af694orpuente-MS2 years ago | 73 | let idx = s | index_offsets[k]; |
84a8c59aorpuente-MS2 years ago | 74 | unsafe { |
| 75 | *extracted_entries.get_unchecked_mut(k) = *state.get_unchecked(idx); | |
| 76 | } | |
| 77 | }); | |
312af694orpuente-MS2 years ago | 78 | |
| 79 | // Apply the gate. | |
84a8c59aorpuente-MS2 years ago | 80 | new_entries.gemv_tr( |
| 81 | Complex::ONE, | |
| 82 | operation_matrix, | |
| 83 | &extracted_entries, | |
| 84 | Complex::ZERO, | |
| 85 | ); | |
312af694orpuente-MS2 years ago | 86 | |
| 87 | // Store accumulated result back into the state vector. | |
84a8c59aorpuente-MS2 years ago | 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. | |
312af694orpuente-MS2 years ago | 97 | let idx = s | index_offsets[k]; |
84a8c59aorpuente-MS2 years ago | 98 | unsafe { |
| 99 | *state.get_unchecked_mut(idx) = *new_entries.get_unchecked(k); | |
| 100 | } | |
| 101 | }); | |
312af694orpuente-MS2 years ago | 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 | } |