microsoft/qdk
Publicmirrored fromhttps://github.com/microsoft/qdkAvailable
source/pip/src/qir_simulation/gpu_full_state.rs
309lines · modecode
| 1 | // Copyright (c) Microsoft Corporation. |
| 2 | // Licensed under the MIT License. |
| 3 | |
| 4 | use crate::qir_simulation::{NoiseConfig, QirInstruction, QirInstructionId, unbind_noise_config}; |
| 5 | use pyo3::{ |
| 6 | IntoPyObjectExt, PyResult, |
| 7 | exceptions::{PyOSError, PyRuntimeError, PyValueError}, |
| 8 | prelude::*, |
| 9 | pyclass, pymethods, |
| 10 | types::{PyDict, PyList}, |
| 11 | }; |
| 12 | use qdk_simulators::gpu_context; |
| 13 | use qdk_simulators::shader_types::Op; |
| 14 | |
| 15 | use std::sync::Mutex; |
| 16 | |
| 17 | /// Checks if a compatible GPU adapter is available on the system. |
| 18 | /// |
| 19 | /// This function attempts to request a GPU adapter to determine if GPU-accelerated |
| 20 | /// quantum simulation is supported. It's useful for capability detection before |
| 21 | /// attempting to run GPU-based simulations. |
| 22 | /// |
| 23 | /// # Errors |
| 24 | /// |
| 25 | /// Raises `OSError` if: |
| 26 | /// - No compatible GPU is found |
| 27 | /// - GPU drivers are missing or not functioning properly |
| 28 | #[pyfunction] |
| 29 | pub fn try_create_gpu_adapter() -> PyResult<String> { |
| 30 | let name = qdk_simulators::try_create_gpu_adapter().map_err(PyOSError::new_err)?; |
| 31 | Ok(name) |
| 32 | } |
| 33 | |
| 34 | #[pyfunction] |
| 35 | pub fn run_parallel_shots<'py>( |
| 36 | py: Python<'py>, |
| 37 | input: &Bound<'py, PyList>, |
| 38 | shots: i32, |
| 39 | qubit_count: i32, |
| 40 | result_count: i32, |
| 41 | noise_config: Option<&Bound<'py, NoiseConfig>>, |
| 42 | seed: Option<u32>, |
| 43 | ) -> PyResult<Py<PyAny>> { |
| 44 | // First convert the Python objects to Rust types |
| 45 | let mut ops: Vec<Op> = Vec::with_capacity(input.len()); |
| 46 | for intr in input { |
| 47 | // Error if the instruction can't be converted |
| 48 | let item: QirInstruction = intr |
| 49 | .extract() |
| 50 | .map_err(|e| PyValueError::new_err(format!("expected QirInstruction: {e}")))?; |
| 51 | // However some ops can't be mapped (e.g. OutputRecording), so skip those |
| 52 | if let Some(op) = map_instruction(&item) { |
| 53 | ops.push(op); |
| 54 | } |
| 55 | } |
| 56 | |
| 57 | let noise = noise_config.map(|noise_config| unbind_noise_config(py, noise_config)); |
| 58 | |
| 59 | let rng_seed = seed.unwrap_or(0xfeed_face); |
| 60 | |
| 61 | let sim_results = |
| 62 | qdk_simulators::run_shots_sync(qubit_count, result_count, &ops, &noise, shots, rng_seed, 0) |
| 63 | .map_err(PyRuntimeError::new_err)?; |
| 64 | |
| 65 | // Collect and format the results into a Python list of strings |
| 66 | let result_count: usize = result_count |
| 67 | .try_into() |
| 68 | .map_err(|e| PyValueError::new_err(format!("invalid result count {result_count}: {e}")))?; |
| 69 | |
| 70 | // Turn each shot's results into a string, with '0' for 0, '1' for 1, and 'L' for lost qubits |
| 71 | // The results are a flat list of u32, with each shot's results in sequence + one error code, |
| 72 | // so we need to chunk them up accordingly |
| 73 | let str_results = sim_results |
| 74 | .shot_results |
| 75 | .iter() |
| 76 | .map(|shot_results| { |
| 77 | let mut bitstring = String::with_capacity(result_count); |
| 78 | for res in shot_results { |
| 79 | let char = match res { |
| 80 | 0 => '0', |
| 81 | 1 => '1', |
| 82 | _ => 'L', // lost qubit |
| 83 | }; |
| 84 | bitstring.push(char); |
| 85 | } |
| 86 | bitstring |
| 87 | }) |
| 88 | .collect::<Vec<String>>(); |
| 89 | |
| 90 | PyList::new(py, str_results) |
| 91 | .map_err(|e| PyValueError::new_err(format!("failed to create Python list: {e}")))? |
| 92 | .into_py_any(py) |
| 93 | } |
| 94 | |
| 95 | type NativeGpuContext = gpu_context::GpuContext; |
| 96 | #[derive(Debug)] |
| 97 | #[pyclass(module = "qsharp._native")] |
| 98 | pub struct GpuContext { |
| 99 | native_context: Mutex<NativeGpuContext>, |
| 100 | last_set_result_count: usize, // Needed to format results |
| 101 | } |
| 102 | |
| 103 | #[pymethods] |
| 104 | impl GpuContext { |
| 105 | #[new] |
| 106 | fn new() -> PyResult<Self> { |
| 107 | Ok(GpuContext { |
| 108 | native_context: Mutex::new(NativeGpuContext::default()), |
| 109 | last_set_result_count: 0, |
| 110 | }) |
| 111 | } |
| 112 | |
| 113 | fn load_noise_tables(&mut self, dir_path: &str) -> PyResult<Vec<(u32, String, u32)>> { |
| 114 | let mut gpu_context = self |
| 115 | .native_context |
| 116 | .lock() |
| 117 | .map_err(|_| PyRuntimeError::new_err("Unable to obtain lock on the GPU context"))?; |
| 118 | |
| 119 | gpu_context.clear_correlated_noise_tables(); |
| 120 | for entry in std::fs::read_dir(dir_path)? { |
| 121 | let entry = entry?; |
| 122 | let path = entry.path(); |
| 123 | let is_file = path.is_file(); |
| 124 | // let ends_with_csv = path.extension().map_or(false, |ext| ext == "csv"); |
| 125 | let ends_with_csv = path.extension() == Some("csv".as_ref()); |
| 126 | |
| 127 | if is_file && ends_with_csv { |
| 128 | let contents = std::fs::read_to_string(&path)?; |
| 129 | let filename = path |
| 130 | .file_stem() |
| 131 | .expect("file should have a name") |
| 132 | .to_str() |
| 133 | .expect("file name should be a valid unicode string"); |
| 134 | gpu_context.add_correlated_noise_table(filename, &contents); |
| 135 | } |
| 136 | } |
| 137 | Ok(gpu_context.get_correlated_noise_tables()) |
| 138 | } |
| 139 | |
| 140 | fn get_noise_table_ids(&self) -> PyResult<Vec<(u32, String, u32)>> { |
| 141 | self.native_context |
| 142 | .lock() |
| 143 | .map_err(|_| PyRuntimeError::new_err("Unable to obtain lock on the GPU context")) |
| 144 | .map(|context| Ok(context.get_correlated_noise_tables()))? |
| 145 | } |
| 146 | |
| 147 | fn set_program( |
| 148 | &mut self, |
| 149 | input: &Bound<'_, PyList>, |
| 150 | qubit_count: i32, |
| 151 | result_count: i32, |
| 152 | ) -> PyResult<()> { |
| 153 | let mut ops: Vec<Op> = Vec::with_capacity(input.len()); |
| 154 | for intr in input { |
| 155 | // Error if the instruction can't be converted |
| 156 | let item: QirInstruction = intr |
| 157 | .extract() |
| 158 | .map_err(|e| PyValueError::new_err(format!("expected QirInstruction: {e}")))?; |
| 159 | // However some ops can't be mapped (e.g. OutputRecording), so skip those |
| 160 | if let Some(op) = map_instruction(&item) { |
| 161 | ops.push(op); |
| 162 | } |
| 163 | } |
| 164 | self.native_context |
| 165 | .lock() |
| 166 | .map_err(|_| PyRuntimeError::new_err("Unable to obtain lock on the GPU context"))? |
| 167 | .set_program(&ops, qubit_count, result_count); |
| 168 | |
| 169 | // Save the result count for formatting later |
| 170 | self.last_set_result_count = result_count.try_into().map_err(|e| { |
| 171 | PyValueError::new_err(format!("invalid result count {result_count}: {e}")) |
| 172 | })?; |
| 173 | Ok(()) |
| 174 | } |
| 175 | |
| 176 | fn set_noise<'py>( |
| 177 | &mut self, |
| 178 | py: Python<'py>, |
| 179 | noise_config: &Bound<'py, NoiseConfig>, |
| 180 | ) -> PyResult<()> { |
| 181 | let noise = unbind_noise_config(py, noise_config); |
| 182 | self.native_context |
| 183 | .lock() |
| 184 | .map_err(|_| PyRuntimeError::new_err("Unable to obtain lock on the GPU context"))? |
| 185 | .set_noise_config(noise); |
| 186 | |
| 187 | Ok(()) |
| 188 | } |
| 189 | |
| 190 | fn run_shots(&self, py: Python<'_>, shot_count: i32, seed: u32) -> PyResult<Py<PyAny>> { |
| 191 | let mut gpu_context = self |
| 192 | .native_context |
| 193 | .lock() |
| 194 | .map_err(|_| PyRuntimeError::new_err("Unable to obtain lock on the GPU context"))?; |
| 195 | |
| 196 | let results = gpu_context |
| 197 | .run_shots_sync(shot_count, seed, 0) |
| 198 | .map_err(|_| PyRuntimeError::new_err("Unable to obtain lock on the GPU context"))?; |
| 199 | |
| 200 | let str_results = results |
| 201 | .shot_results |
| 202 | .iter() |
| 203 | .map(|shot_results| { |
| 204 | let mut bitstring = String::with_capacity(self.last_set_result_count); |
| 205 | for res in shot_results { |
| 206 | let char = match res { |
| 207 | 0 => '0', |
| 208 | 1 => '1', |
| 209 | _ => 'L', // lost qubit |
| 210 | }; |
| 211 | bitstring.push(char); |
| 212 | } |
| 213 | bitstring |
| 214 | }) |
| 215 | .collect::<Vec<String>>(); |
| 216 | |
| 217 | let dict = PyDict::new(py); |
| 218 | |
| 219 | dict.set_item("shot_results", PyList::new(py, str_results)?) |
| 220 | .map_err(|e| PyValueError::new_err(format!("failed to set results in dict: {e}")))?; |
| 221 | dict.set_item( |
| 222 | "shot_result_codes", |
| 223 | PyList::new(py, results.shot_result_codes)?, |
| 224 | ) |
| 225 | .map_err(|e| PyValueError::new_err(format!("failed to set result codes in dict: {e}")))?; |
| 226 | |
| 227 | if let Some(diagnostics) = results.diagnostics { |
| 228 | // DiagnosticsData doesn't implement Serialize, so use Debug formatting |
| 229 | dict.set_item("diagnostics", format!("{diagnostics:?}")) |
| 230 | .map_err(|e| { |
| 231 | PyValueError::new_err(format!("failed to set diagnostics in dict: {e}")) |
| 232 | })?; |
| 233 | } |
| 234 | dict.into_py_any(py) |
| 235 | } |
| 236 | } |
| 237 | |
| 238 | fn map_instruction(qir_inst: &QirInstruction) -> Option<Op> { |
| 239 | let op = match qir_inst { |
| 240 | QirInstruction::OneQubitGate(id, qubit) => match id { |
| 241 | QirInstructionId::I => Op::new_id_gate(*qubit), |
| 242 | QirInstructionId::Move => Op::new_move_gate(*qubit), |
| 243 | QirInstructionId::H => Op::new_h_gate(*qubit), |
| 244 | QirInstructionId::X => Op::new_x_gate(*qubit), |
| 245 | QirInstructionId::Y => Op::new_y_gate(*qubit), |
| 246 | QirInstructionId::Z => Op::new_z_gate(*qubit), |
| 247 | QirInstructionId::S => Op::new_s_gate(*qubit), |
| 248 | QirInstructionId::SAdj => Op::new_s_adj_gate(*qubit), |
| 249 | QirInstructionId::SX => Op::new_sx_gate(*qubit), |
| 250 | QirInstructionId::SXAdj => Op::new_sx_adj_gate(*qubit), |
| 251 | QirInstructionId::T => Op::new_t_gate(*qubit), |
| 252 | QirInstructionId::TAdj => Op::new_t_adj_gate(*qubit), |
| 253 | _ => { |
| 254 | panic!("unsupported one-qubit gate: {id:?} on qubit {qubit}"); |
| 255 | } |
| 256 | }, |
| 257 | QirInstruction::TwoQubitGate(id, control, target) => match id { |
| 258 | QirInstructionId::M | QirInstructionId::MZ | QirInstructionId::MResetZ => { |
| 259 | // TODO: These should be distinct in the simulator |
| 260 | Op::new_mresetz_gate(*control, *target) |
| 261 | } |
| 262 | QirInstructionId::CX | QirInstructionId::CNOT => Op::new_cx_gate(*control, *target), |
| 263 | QirInstructionId::CY => Op::new_cy_gate(*control, *target), |
| 264 | QirInstructionId::CZ => Op::new_cz_gate(*control, *target), |
| 265 | QirInstructionId::SWAP => Op::new_swap_gate(*control, *target), |
| 266 | _ => { |
| 267 | panic!("unsupported two-qubit gate: {id:?} on qubits {control}, {target}"); |
| 268 | } |
| 269 | }, |
| 270 | QirInstruction::OneQubitRotationGate(id, angle, qubit) => { |
| 271 | #[allow(clippy::cast_possible_truncation)] |
| 272 | let angle = *angle as f32; |
| 273 | match id { |
| 274 | QirInstructionId::RX => Op::new_rx_gate(angle, *qubit), |
| 275 | QirInstructionId::RY => Op::new_ry_gate(angle, *qubit), |
| 276 | QirInstructionId::RZ => Op::new_rz_gate(angle, *qubit), |
| 277 | _ => { |
| 278 | panic!("unsupported one-qubit rotation gate: {id:?} on qubit {qubit}"); |
| 279 | } |
| 280 | } |
| 281 | } |
| 282 | QirInstruction::TwoQubitRotationGate(id, angle, qubit1, qubit2) => { |
| 283 | #[allow(clippy::cast_possible_truncation)] |
| 284 | let angle = *angle as f32; |
| 285 | match id { |
| 286 | QirInstructionId::RXX => Op::new_rxx_gate(angle, *qubit1, *qubit2), |
| 287 | QirInstructionId::RYY => Op::new_ryy_gate(angle, *qubit1, *qubit2), |
| 288 | QirInstructionId::RZZ => Op::new_rzz_gate(angle, *qubit1, *qubit2), |
| 289 | _ => { |
| 290 | panic!( |
| 291 | "unsupported two-qubit rotation gate: {id:?} on qubits {qubit1}, {qubit2}" |
| 292 | ); |
| 293 | } |
| 294 | } |
| 295 | } |
| 296 | QirInstruction::ThreeQubitGate(QirInstructionId::CCX, c1, c2, target) => { |
| 297 | unimplemented!("{c1}, {c2}, {target}") //Op::new_ccx_gate(*c1, *c2, *target), |
| 298 | } |
| 299 | QirInstruction::OutputRecording(_, _, _) => { |
| 300 | // Ignore for now |
| 301 | return None; |
| 302 | } |
| 303 | QirInstruction::ThreeQubitGate(..) => panic!("unsupported instruction: {qir_inst:?}"), |
| 304 | QirInstruction::CorrelatedNoise(_, table_id, qubit_args) => { |
| 305 | Op::new_correlated_noise_gate(*table_id, qubit_args) |
| 306 | } |
| 307 | }; |
| 308 | Some(op) |
| 309 | } |
| 310 | |