microsoft/qdk

Public

mirrored fromhttps://github.com/microsoft/qdkAvailable

CodeCommitsIssuesPull requestsActionsInsightsSecurity
iadavis/pipeline-issue-debugging

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

source/pip/src/qir_simulation/gpu_full_state.rs

309lines · modecode

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4use crate::qir_simulation::{NoiseConfig, QirInstruction, QirInstructionId, unbind_noise_config};
5use pyo3::{
6 IntoPyObjectExt, PyResult,
7 exceptions::{PyOSError, PyRuntimeError, PyValueError},
8 prelude::*,
9 pyclass, pymethods,
10 types::{PyDict, PyList},
11};
12use qdk_simulators::gpu_context;
13use qdk_simulators::shader_types::Op;
14
15use 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]
29pub 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]
35pub 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
95type NativeGpuContext = gpu_context::GpuContext;
96#[derive(Debug)]
97#[pyclass(module = "qsharp._native")]
98pub struct GpuContext {
99 native_context: Mutex<NativeGpuContext>,
100 last_set_result_count: usize, // Needed to format results
101}
102
103#[pymethods]
104impl 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
238fn 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