GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/qdk

Public

mirrored from https://github.com/microsoft/qdkAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
v1.28.0

Branches

Tags

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

Clone

HTTPS

Download ZIP

qdk/source/pip/src/qir_simulation

source/pip/src/qir_simulation/gpu_full_state.rs

452lines · modecode

RawDownload
Latest commit unavailable.
1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4use crate::qir_simulation::{
5 NoiseConfig, QirInstruction, QirInstructionId, adaptive_program_from_pydict,
6 unbind_noise_config,
7};
8use pyo3::{
9 IntoPyObjectExt, PyResult,
10 exceptions::{PyOSError, PyRuntimeError, PyValueError},
11 prelude::*,
12 pyclass, pymethods,
13 types::{PyDict, PyList},
14};
15use qdk_simulators::gpu_context;
16use qdk_simulators::shader_types::Op;
17
18use std::sync::Mutex;
19
20/// Checks if a compatible GPU adapter is available on the system.
21///
22/// This function attempts to request a GPU adapter to determine if GPU-accelerated
23/// quantum simulation is supported. It's useful for capability detection before
24/// attempting to run GPU-based simulations.
25///
26/// # Errors
27///
28/// Raises `OSError` if:
29/// - No compatible GPU is found
30/// - GPU drivers are missing or not functioning properly
31#[pyfunction]
32pub fn try_create_gpu_adapter() -> PyResult<String> {
33 let name = qdk_simulators::try_create_gpu_adapter().map_err(PyOSError::new_err)?;
34 Ok(name)
35}
36
37#[pyfunction]
38pub fn run_parallel_shots<'py>(
39 py: Python<'py>,
40 input: &Bound<'py, PyList>,
41 qubit_count: i32,
42 result_count: i32,
43 shots: i32,
44 noise_config: Option<&Bound<'py, NoiseConfig>>,
45 seed: Option<u32>,
46) -> PyResult<Py<PyAny>> {
47 // First convert the Python objects to Rust types
48 let mut ops: Vec<Op> = Vec::with_capacity(input.len());
49 for intr in input {
50 // Error if the instruction can't be converted
51 let item: QirInstruction = intr
52 .extract()
53 .map_err(|e| PyValueError::new_err(format!("expected QirInstruction: {e}")))?;
54 // However some ops can't be mapped (e.g. OutputRecording), so skip those
55 if let Some(op) = map_instruction(&item) {
56 ops.push(op);
57 }
58 }
59
60 let noise = noise_config.map(|noise_config| unbind_noise_config(py, noise_config));
61
62 let rng_seed = seed.unwrap_or(0xfeed_face);
63
64 let sim_results =
65 qdk_simulators::run_shots_sync(qubit_count, result_count, &ops, &noise, shots, rng_seed, 0)
66 .map_err(PyRuntimeError::new_err)?;
67
68 // Collect and format the results into a Python list of strings
69 let result_count: usize = result_count
70 .try_into()
71 .map_err(|e| PyValueError::new_err(format!("invalid result count {result_count}: {e}")))?;
72
73 // Turn each shot's results into a string, with '0' for 0, '1' for 1, and 'L' for lost qubits
74 // The results are a flat list of u32, with each shot's results in sequence + one error code,
75 // so we need to chunk them up accordingly
76 let str_results = sim_results
77 .shot_results
78 .iter()
79 .map(|shot_results| {
80 let mut bitstring = String::with_capacity(result_count);
81 for res in shot_results {
82 let char = match res {
83 0 => '0',
84 1 => '1',
85 _ => 'L', // lost qubit
86 };
87 bitstring.push(char);
88 }
89 bitstring
90 })
91 .collect::<Vec<String>>();
92
93 PyList::new(py, str_results)
94 .map_err(|e| PyValueError::new_err(format!("failed to create Python list: {e}")))?
95 .into_py_any(py)
96}
97
98type NativeGpuContext = gpu_context::GpuContext;
99#[derive(Debug)]
100#[pyclass(module = "qsharp._native")]
101pub struct GpuContext {
102 native_context: Mutex<NativeGpuContext>,
103 last_set_result_count: usize, // Needed to format results
104}
105
106#[pymethods]
107impl GpuContext {
108 #[new]
109 fn new() -> PyResult<Self> {
110 Ok(GpuContext {
111 native_context: Mutex::new(NativeGpuContext::default()),
112 last_set_result_count: 0,
113 })
114 }
115
116 fn load_noise_tables(&mut self, dir_path: &str) -> PyResult<Vec<(u32, String, u32)>> {
117 let mut gpu_context = self
118 .native_context
119 .lock()
120 .map_err(|_| PyRuntimeError::new_err("Unable to obtain lock on the GPU context"))?;
121
122 gpu_context.clear_correlated_noise_tables();
123 for entry in std::fs::read_dir(dir_path)? {
124 let entry = entry?;
125 let path = entry.path();
126 let is_file = path.is_file();
127 // let ends_with_csv = path.extension().map_or(false, |ext| ext == "csv");
128 let ends_with_csv = path.extension() == Some("csv".as_ref());
129
130 if is_file && ends_with_csv {
131 let contents = std::fs::read_to_string(&path)?;
132 let filename = path
133 .file_stem()
134 .expect("file should have a name")
135 .to_str()
136 .expect("file name should be a valid unicode string");
137 gpu_context.add_correlated_noise_table(filename, &contents);
138 }
139 }
140 Ok(gpu_context.get_correlated_noise_tables())
141 }
142
143 fn get_noise_table_ids(&self) -> PyResult<Vec<(u32, String, u32)>> {
144 self.native_context
145 .lock()
146 .map_err(|_| PyRuntimeError::new_err("Unable to obtain lock on the GPU context"))
147 .map(|context| Ok(context.get_correlated_noise_tables()))?
148 }
149
150 fn set_program(
151 &mut self,
152 input: &Bound<'_, PyList>,
153 qubit_count: i32,
154 result_count: i32,
155 ) -> PyResult<()> {
156 let mut gpu_context = self
157 .native_context
158 .lock()
159 .map_err(|_| PyRuntimeError::new_err("Unable to obtain lock on the GPU context"))?;
160
161 gpu_context.switch_to_base();
162
163 let mut ops: Vec<Op> = Vec::with_capacity(input.len());
164 for intr in input {
165 // Error if the instruction can't be converted
166 let item: QirInstruction = intr
167 .extract()
168 .map_err(|e| PyValueError::new_err(format!("expected QirInstruction: {e}")))?;
169 // However some ops can't be mapped (e.g. OutputRecording), so skip those
170 if let Some(op) = map_instruction(&item) {
171 ops.push(op);
172 }
173 }
174 gpu_context.set_program(&ops, qubit_count, result_count);
175
176 // Save the result count for formatting later
177 self.last_set_result_count = result_count.try_into().map_err(|e| {
178 PyValueError::new_err(format!("invalid result count {result_count}: {e}"))
179 })?;
180 Ok(())
181 }
182
183 fn set_noise<'py>(
184 &mut self,
185 py: Python<'py>,
186 noise_config: &Bound<'py, NoiseConfig>,
187 ) -> PyResult<()> {
188 let noise = unbind_noise_config(py, noise_config);
189 self.native_context
190 .lock()
191 .map_err(|_| PyRuntimeError::new_err("Unable to obtain lock on the GPU context"))?
192 .set_noise_config(noise);
193
194 Ok(())
195 }
196
197 fn run_shots(&self, py: Python<'_>, shot_count: i32, seed: u32) -> PyResult<Py<PyAny>> {
198 let mut gpu_context = self
199 .native_context
200 .lock()
201 .map_err(|_| PyRuntimeError::new_err("Unable to obtain lock on the GPU context"))?;
202
203 if gpu_context.is_adaptive() {
204 return Err(PyRuntimeError::new_err(
205 "Context should be non-adaptive. Try setting a base profile program first with `.set_program()`",
206 ));
207 }
208
209 let results = gpu_context
210 .run_shots_sync(shot_count, seed, 0)
211 .map_err(|_| PyRuntimeError::new_err("Unable to obtain lock on the GPU context"))?;
212
213 let str_results = results
214 .shot_results
215 .iter()
216 .map(|shot_results| {
217 let mut bitstring = String::with_capacity(self.last_set_result_count);
218 for res in shot_results {
219 let char = match res {
220 0 => '0',
221 1 => '1',
222 _ => 'L', // lost qubit
223 };
224 bitstring.push(char);
225 }
226 bitstring
227 })
228 .collect::<Vec<String>>();
229
230 let dict = PyDict::new(py);
231
232 dict.set_item("shot_results", PyList::new(py, str_results)?)
233 .map_err(|e| PyValueError::new_err(format!("failed to set results in dict: {e}")))?;
234 dict.set_item(
235 "shot_result_codes",
236 PyList::new(py, results.shot_result_codes)?,
237 )
238 .map_err(|e| PyValueError::new_err(format!("failed to set result codes in dict: {e}")))?;
239
240 if let Some(diagnostics) = results.diagnostics {
241 // DiagnosticsData doesn't implement Serialize, so use Debug formatting
242 dict.set_item("diagnostics", format!("{diagnostics:?}"))
243 .map_err(|e| {
244 PyValueError::new_err(format!("failed to set diagnostics in dict: {e}"))
245 })?;
246 }
247 dict.into_py_any(py)
248 }
249
250 fn set_adaptive_program(&mut self, program: &Bound<'_, PyDict>) -> PyResult<()> {
251 let mut gpu_context = self
252 .native_context
253 .lock()
254 .map_err(|_| PyRuntimeError::new_err("Unable to obtain lock on the GPU context"))?;
255
256 gpu_context.swith_to_adaptive();
257
258 let adaptive_program = adaptive_program_from_pydict(program)?;
259 let num_results = adaptive_program.num_results;
260
261 gpu_context
262 .set_adaptive_program(adaptive_program)
263 .map_err(PyValueError::new_err)?;
264
265 // Save the result count for formatting later
266 self.last_set_result_count = num_results.try_into().map_err(|e| {
267 PyValueError::new_err(format!("invalid result count {num_results}: {e}"))
268 })?;
269
270 Ok(())
271 }
272
273 fn run_adaptive_shots(
274 &self,
275 py: Python<'_>,
276 shot_count: i32,
277 seed: u32,
278 ) -> PyResult<Py<PyAny>> {
279 let mut gpu_context = self
280 .native_context
281 .lock()
282 .map_err(|_| PyRuntimeError::new_err("Unable to obtain lock on the GPU context"))?;
283
284 if !gpu_context.is_adaptive() {
285 return Err(PyRuntimeError::new_err(
286 "Context should be adaptive. Try setting an adaptive program first with `.set_adaptive_program()`",
287 ));
288 }
289
290 let results = gpu_context
291 .run_adaptive_shots_sync(shot_count, seed, 0)
292 .map_err(PyRuntimeError::new_err)?;
293
294 Self::format_results(py, results, self.last_set_result_count)
295 }
296}
297
298impl GpuContext {
299 fn format_results(
300 py: Python<'_>,
301 results: gpu_context::RunResults,
302 result_count: usize,
303 ) -> PyResult<Py<PyAny>> {
304 let str_results = results
305 .shot_results
306 .iter()
307 .map(|shot_results| {
308 let mut bitstring = String::with_capacity(result_count);
309 for res in shot_results {
310 let char = match res {
311 0 => '0',
312 1 => '1',
313 _ => 'L',
314 };
315 bitstring.push(char);
316 }
317 bitstring
318 })
319 .collect::<Vec<String>>();
320
321 let dict = PyDict::new(py);
322 dict.set_item("shot_results", PyList::new(py, str_results)?)
323 .map_err(|e| PyValueError::new_err(format!("failed to set results in dict: {e}")))?;
324 dict.set_item(
325 "shot_result_codes",
326 PyList::new(py, results.shot_result_codes)?,
327 )
328 .map_err(|e| PyValueError::new_err(format!("failed to set result codes in dict: {e}")))?;
329
330 if let Some(diagnostics) = results.diagnostics {
331 dict.set_item("diagnostics", format!("{diagnostics:?}"))
332 .map_err(|e| {
333 PyValueError::new_err(format!("failed to set diagnostics in dict: {e}"))
334 })?;
335 }
336 dict.into_py_any(py)
337 }
338}
339
340fn map_instruction(qir_inst: &QirInstruction) -> Option<Op> {
341 let op = match qir_inst {
342 QirInstruction::OneQubitGate(id, qubit) => match id {
343 QirInstructionId::I => Op::new_id_gate(*qubit),
344 QirInstructionId::Move => Op::new_move_gate(*qubit),
345 QirInstructionId::H => Op::new_h_gate(*qubit),
346 QirInstructionId::X => Op::new_x_gate(*qubit),
347 QirInstructionId::Y => Op::new_y_gate(*qubit),
348 QirInstructionId::Z => Op::new_z_gate(*qubit),
349 QirInstructionId::S => Op::new_s_gate(*qubit),
350 QirInstructionId::SAdj => Op::new_s_adj_gate(*qubit),
351 QirInstructionId::SX => Op::new_sx_gate(*qubit),
352 QirInstructionId::SXAdj => Op::new_sx_adj_gate(*qubit),
353 QirInstructionId::T => Op::new_t_gate(*qubit),
354 QirInstructionId::TAdj => Op::new_t_adj_gate(*qubit),
355 QirInstructionId::RESET => Op::new_resetz_gate(*qubit),
356 _ => {
357 panic!("unsupported one-qubit gate: {id:?} on qubit {qubit}");
358 }
359 },
360 QirInstruction::TwoQubitGate(id, control, target) => match id {
361 QirInstructionId::M | QirInstructionId::MZ => Op::new_mz_gate(*control, *target),
362 QirInstructionId::MResetZ => Op::new_mresetz_gate(*control, *target),
363 QirInstructionId::CX | QirInstructionId::CNOT => Op::new_cx_gate(*control, *target),
364 QirInstructionId::CY => Op::new_cy_gate(*control, *target),
365 QirInstructionId::CZ => Op::new_cz_gate(*control, *target),
366 QirInstructionId::SWAP => Op::new_swap_gate(*control, *target),
367 _ => {
368 panic!("unsupported two-qubit gate: {id:?} on qubits {control}, {target}");
369 }
370 },
371 QirInstruction::OneQubitRotationGate(id, angle, qubit) => {
372 #[allow(clippy::cast_possible_truncation)]
373 let angle = *angle as f32;
374 match id {
375 QirInstructionId::RX => Op::new_rx_gate(angle, *qubit),
376 QirInstructionId::RY => Op::new_ry_gate(angle, *qubit),
377 QirInstructionId::RZ => Op::new_rz_gate(angle, *qubit),
378 _ => {
379 panic!("unsupported one-qubit rotation gate: {id:?} on qubit {qubit}");
380 }
381 }
382 }
383 QirInstruction::TwoQubitRotationGate(id, angle, qubit1, qubit2) => {
384 #[allow(clippy::cast_possible_truncation)]
385 let angle = *angle as f32;
386 match id {
387 QirInstructionId::RXX => Op::new_rxx_gate(angle, *qubit1, *qubit2),
388 QirInstructionId::RYY => Op::new_ryy_gate(angle, *qubit1, *qubit2),
389 QirInstructionId::RZZ => Op::new_rzz_gate(angle, *qubit1, *qubit2),
390 _ => {
391 panic!(
392 "unsupported two-qubit rotation gate: {id:?} on qubits {qubit1}, {qubit2}"
393 );
394 }
395 }
396 }
397 QirInstruction::ThreeQubitGate(QirInstructionId::CCX, c1, c2, target) => {
398 unimplemented!("{c1}, {c2}, {target}") //Op::new_ccx_gate(*c1, *c2, *target),
399 }
400 QirInstruction::OutputRecording(_, _, _) => {
401 // Ignore for now
402 return None;
403 }
404 QirInstruction::ThreeQubitGate(..) => panic!("unsupported instruction: {qir_inst:?}"),
405 QirInstruction::CorrelatedNoise(_, table_id, qubit_args) => {
406 Op::new_correlated_noise_gate(*table_id, qubit_args)
407 }
408 };
409 Some(op)
410}
411
412#[pyfunction]
413pub fn run_adaptive_parallel_shots<'py>(
414 py: Python<'py>,
415 input: &Bound<'py, PyDict>,
416 shots: i32,
417 noise_config: Option<&Bound<'py, NoiseConfig>>,
418 seed: Option<u32>,
419) -> PyResult<Py<PyAny>> {
420 let noise = noise_config.map(|noise_config| unbind_noise_config(py, noise_config));
421 let rng_seed = seed.unwrap_or(0xfeed_face);
422 let program = adaptive_program_from_pydict(input)?;
423 let result_count: usize = program.num_results as usize;
424 let sim_results = qdk_simulators::run_adaptive_shots_sync(program, &noise, shots, rng_seed, 0)
425 .map_err(PyRuntimeError::new_err)?;
426
427 // Collect and format the results into a Python list of strings
428
429 // Turn each shot's results into a string, with '0' for 0, '1' for 1, and 'L' for lost qubits
430 // The results are a flat list of u32, with each shot's results in sequence + one error code,
431 // so we need to chunk them up accordingly
432 let str_results = sim_results
433 .shot_results
434 .iter()
435 .map(|shot_results| {
436 let mut bitstring = String::with_capacity(result_count);
437 for res in shot_results {
438 let char = match res {
439 0 => '0',
440 1 => '1',
441 _ => 'L', // lost qubit
442 };
443 bitstring.push(char);
444 }
445 bitstring
446 })
447 .collect::<Vec<String>>();
448
449 PyList::new(py, str_results)
450 .map_err(|e| PyValueError::new_err(format!("failed to create Python list: {e}")))?
451 .into_py_any(py)
452}
453