Skip to content

Commit

Permalink
Fully port Optimize1qGatesDecomposition to Rust
Browse files Browse the repository at this point in the history 
This commit builds off of #12550 and the other data model in Rust
infrastructure and migrates the Optimize1qGatesDecomposition pass to
operate fully in Rust. The full path of the transpiler pass now never
leaves rust until it has finished modifying the DAGCircuit. There is
still some python interaction necessary to handle parts of the data
model that are still in Python, mainly calibrations and parameter
expressions (for global phase). But otherwise the entirety of the pass
operates in rust now.

This is just a first pass at the migration here, it moves the pass to be
a single for loop in rust. The next steps here are to look at operating
the pass in parallel. There is no data dependency between the
optimizations being done by the pass so we should be able to the
throughput of the pass by leveraging multithreading to handle each run
in parallel. This commit does not attempt this though, because of the
Python dependency and also the data structures around gates and the
dag aren't really setup for multithreading yet and there likely will
need to be some work to support that (this pass is a good candidate to
work through the bugs on that).

Part of #12208
  • Loading branch information
@mtreinish
mtreinish committed Aug 22, 2024
1 parent 415e9b7 commit 0c8b668
Show file tree
Hide file tree
Showing 7 changed files with 423 additions and 160 deletions.
315 changes: 227 additions & 88 deletions crates/accelerate/src/euler_one_qubit_decomposer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,9 @@
#![allow(clippy::too_many_arguments)]
#![allow(clippy::upper_case_acronyms)]

use hashbrown::HashMap;
use ahash::RandomState;
use hashbrown::{HashMap, HashSet};
use indexmap::IndexSet;
use num_complex::{Complex64, ComplexFloat};
use smallvec::{smallvec, SmallVec};
use std::cmp::Ordering;
Expand All @@ -29,14 +31,19 @@ use pyo3::Python;
use ndarray::prelude::*;
use numpy::PyReadonlyArray2;
use pyo3::pybacked::PyBackedStr;
use rustworkx_core::petgraph::stable_graph::NodeIndex;

use qiskit_circuit::circuit_data::CircuitData;
use qiskit_circuit::dag_circuit::{DAGCircuit, NodeType};
use qiskit_circuit::dag_node::DAGOpNode;
use qiskit_circuit::operations::{Operation, Param, StandardGate};
use qiskit_circuit::slice::{PySequenceIndex, SequenceIndex};
use qiskit_circuit::util::c64;
use qiskit_circuit::Qubit;

use crate::nlayout::PhysicalQubit;
use crate::target_transpiler::Target;

pub const ANGLE_ZERO_EPSILON: f64 = 1e-12;

#[pyclass(module = "qiskit._accelerate.euler_one_qubit_decomposer")]
Expand Down Expand Up @@ -69,6 +76,7 @@ impl OneQubitGateErrorMap {
}
}

#[derive(Debug)]
#[pyclass(sequence)]
pub struct OneQubitGateSequence {
pub gates: Vec<(StandardGate, SmallVec<[f64; 3]>)>,
Expand Down Expand Up @@ -571,7 +579,7 @@ pub fn generate_circuit(
Ok(res)
}

#[derive(Clone, Debug, Copy)]
#[derive(Clone, Debug, Copy, Eq, Hash, PartialEq)]
#[pyclass(module = "qiskit._accelerate.euler_one_qubit_decomposer")]
pub enum EulerBasis {
U321,
Expand Down Expand Up @@ -684,24 +692,6 @@ fn compare_error_fn(
}
}

fn compute_error(
gates: &[(StandardGate, SmallVec<[f64; 3]>)],
error_map: Option<&OneQubitGateErrorMap>,
qubit: usize,
) -> (f64, usize) {
match error_map {
Some(err_map) => {
let num_gates = gates.len();
let gate_fidelities: f64 = gates
.iter()
.map(|gate| 1. - err_map.error_map[qubit].get(gate.0.name()).unwrap_or(&0.))
.product();
(1. - gate_fidelities, num_gates)
}
None => (gates.len() as f64, gates.len()),
}
}

fn compute_error_term(gate: &str, error_map: &OneQubitGateErrorMap, qubit: usize) -> f64 {
1. - error_map.error_map[qubit].get(gate).unwrap_or(&0.)
}
Expand All @@ -724,15 +714,6 @@ fn compute_error_str(
}
}

#[pyfunction]
pub fn compute_error_one_qubit_sequence(
circuit: &OneQubitGateSequence,
qubit: usize,
error_map: Option<&OneQubitGateErrorMap>,
) -> (f64, usize) {
compute_error(&circuit.gates, error_map, qubit)
}

#[pyfunction]
pub fn compute_error_list(
circuit: Vec<PyRef<DAGOpNode>>,
Expand Down Expand Up @@ -965,72 +946,231 @@ pub fn params_zxz(unitary: PyReadonlyArray2<Complex64>) -> [f64; 4] {
params_zxz_inner(mat)
}

type OptimizeDecompositionReturn = Option<((f64, usize), (f64, usize), OneQubitGateSequence)>;
fn compute_error_term_from_target(gate: &str, target: &Target, qubit: PhysicalQubit) -> f64 {
1. - target.get_error(gate, &[qubit]).unwrap_or(0.)
}

fn compute_error_from_target_one_qubit_sequence(
circuit: &OneQubitGateSequence,
qubit: PhysicalQubit,
target: Option<&Target>,
) -> (f64, usize) {
match target {
Some(target) => {
let num_gates = circuit.gates.len();
let gate_fidelities: f64 = circuit
.gates
.iter()
.map(|gate| compute_error_term_from_target(gate.0.name(), target, qubit))
.product();
(1. - gate_fidelities, num_gates)
}
None => (circuit.gates.len() as f64, circuit.gates.len()),
}
}

#[pyfunction]
pub fn optimize_1q_gates_decomposition(
runs: Vec<Vec<PyRef<DAGOpNode>>>,
qubits: Vec<usize>,
bases: Vec<Vec<PyBackedStr>>,
simplify: bool,
error_map: Option<&OneQubitGateErrorMap>,
atol: Option<f64>,
) -> Vec<OptimizeDecompositionReturn> {
runs.iter()
.enumerate()
.map(|(index, raw_run)| -> OptimizeDecompositionReturn {
let mut error = match error_map {
Some(_) => 1.,
None => raw_run.len() as f64,
#[pyo3(signature = (dag, *, target=None, basis_gates=None, global_decomposers=None))]
pub(crate) fn optimize_1q_gates_decomposition(
py: Python,
dag: &mut DAGCircuit,
target: Option<&Target>,
basis_gates: Option<HashSet<String>>,
global_decomposers: Option<Vec<String>>,
) -> PyResult<()> {
let runs: Vec<Vec<NodeIndex>> = dag.collect_1q_runs().unwrap().collect();
let dag_qubits = dag.num_qubits();
let mut target_basis_per_qubit: Vec<Option<IndexSet<EulerBasis, RandomState>>> =
vec![None; dag_qubits];
let mut basis_gates_per_qubit: Vec<Option<HashSet<&str>>> = vec![None; dag_qubits];
for raw_run in runs {
let mut error = match target {
Some(_) => 1.,
None => raw_run.len() as f64,
};
let qubit: PhysicalQubit = if let NodeType::Operation(inst) = &dag.dag[raw_run[0]] {
dag.get_qubits(inst.qubits)[0].into()
} else {
unreachable!("nodes in runs will always be op nodes")
};
if !dag.calibrations_empty() {
let mut has_calibration = false;
for node in &raw_run {
if dag.has_calibration_for_index(py, *node)? {
has_calibration = true;
break;
}
}
if has_calibration {
continue;
}
}
if basis_gates_per_qubit[qubit.0 as usize].is_none() {
let basis_gates = match target {
Some(target) => Some(
target
.operation_names_for_qargs(Some(&smallvec![qubit]))
.unwrap(),
),
None => {
let basis = basis_gates.as_ref();
basis.map(|basis| basis.iter().map(|x| x.as_str()).collect())
}
};
let qubit = qubits[index];
let operator = &raw_run
.iter()
.map(|node| {
if let Some(err_map) = error_map {
error *=
compute_error_term(node.instruction.operation.name(), err_map, qubit)
}
node.instruction
.operation
.matrix(&node.instruction.params)
.expect("No matrix defined for operation")
})
.fold(
[
[Complex64::new(1., 0.), Complex64::new(0., 0.)],
[Complex64::new(0., 0.), Complex64::new(1., 0.)],
],
|mut operator, node| {
matmul_1q(&mut operator, node);
operator
basis_gates_per_qubit[qubit.0 as usize] = basis_gates;
}
let basis_gates = &basis_gates_per_qubit[qubit.0 as usize].as_ref();

if target_basis_per_qubit[qubit.0 as usize].is_none() {
let mut target_basis_set: IndexSet<EulerBasis, RandomState> = match target {
Some(_target) => EULER_BASIS_MAP
.iter()
.enumerate()
.filter_map(|(idx, gates)| {
if !gates
.iter()
.all(|gate| basis_gates.as_ref().unwrap().contains(gate))
{
return None;
}
let basis = EULER_BASIS_NAMES[idx];
Some(basis)
})
.collect(),
None => match &global_decomposers {
Some(bases) => bases
.iter()
.map(|basis| EulerBasis::__new__(basis).unwrap())
.collect(),
None => match basis_gates {
Some(gates) => EULER_BASIS_MAP
.iter()
.enumerate()
.filter_map(|(idx, basis_gates)| {
if !gates.iter().all(|gate| basis_gates.as_ref().contains(gate)) {
return None;
}
let basis = EULER_BASIS_NAMES[idx];
Some(basis)
})
.collect(),
None => EULER_BASIS_NAMES.iter().copied().collect(),
},
);
let old_error = if error_map.is_some() {
(1. - error, raw_run.len())
} else {
(error, raw_run.len())
},
};
let target_basis_vec: Vec<EulerBasis> = bases[index]
.iter()
.map(|basis| EulerBasis::__new__(basis).unwrap())
.collect();
unitary_to_gate_sequence_inner(
aview2(operator),
&target_basis_vec,
qubit,
error_map,
simplify,
atol,
)
.map(|out_seq| {
let new_error = compute_error_one_qubit_sequence(&out_seq, qubit, error_map);
(old_error, new_error, out_seq)
if target_basis_set.contains(&EulerBasis::U3)
&& target_basis_set.contains(&EulerBasis::U321)
{
target_basis_set.swap_remove(&EulerBasis::U3);
}
if target_basis_set.contains(&EulerBasis::ZSX)
&& target_basis_set.contains(&EulerBasis::ZSXX)
{
target_basis_set.swap_remove(&EulerBasis::ZSX);
}
target_basis_per_qubit[qubit.0 as usize] = Some(target_basis_set);
}
let target_basis_set = target_basis_per_qubit[qubit.0 as usize].as_ref().unwrap();
let target_basis_vec: Vec<EulerBasis> = target_basis_set.iter().copied().collect();
let operator = raw_run
.iter()
.map(|node_index| {
let node = &dag.dag[*node_index];
if let NodeType::Operation(inst) = node {
if let Some(target) = target {
error *= compute_error_term_from_target(inst.op.name(), target, qubit);
}
inst.op.matrix(inst.params_view()).unwrap()
} else {
unreachable!("Can only have op nodes here")
}
})
})
.collect()
.fold(
[
[Complex64::new(1., 0.), Complex64::new(0., 0.)],
[Complex64::new(0., 0.), Complex64::new(1., 0.)],
],
|mut operator, node| {
matmul_1q(&mut operator, node);
operator
},
);

let old_error = if target.is_some() {
(1. - error, raw_run.len())
} else {
(error, raw_run.len())
};
let sequence = unitary_to_gate_sequence_inner(
aview2(&operator),
&target_basis_vec,
qubit.0 as usize,
None,
true,
None,
);
let sequence = match sequence {
Some(seq) => seq,
None => continue,
};
let new_error = compute_error_from_target_one_qubit_sequence(&sequence, qubit, target);

let mut outside_basis = false;
if let Some(basis) = basis_gates {
for node in &raw_run {
if let NodeType::Operation(inst) = &dag.dag[*node] {
if !basis.contains(inst.op.name()) {
outside_basis = true;
break;
}
}
}
} else {
outside_basis = false;
}
if outside_basis
|| new_error < old_error
|| new_error.0.abs() < 1e-9 && old_error.0.abs() >= 1e-9
{
for gate in sequence.gates {
dag.insert_1q_on_incoming_qubit((gate.0, &gate.1), raw_run[0]);
}
dag.add_global_phase(py, &Param::Float(sequence.global_phase))?;
dag.remove_1q_sequence(&raw_run);
}
}
Ok(())
}

static EULER_BASIS_MAP: [&[&str]; 12] = [
&["u3"],
&["u3", "u2", "u1"],
&["u"],
&["p", "sx"],
&["u1", "rx"],
&["r"],
&["rz", "ry"],
&["rz", "rx"],
&["rz", "rx"],
&["rx", "ry"],
&["rz", "sx", "x"],
&["rz", "sx"],
];
static EULER_BASIS_NAMES: [EulerBasis; 12] = [
EulerBasis::U3,
EulerBasis::U321,
EulerBasis::U,
EulerBasis::PSX,
EulerBasis::U1X,
EulerBasis::RR,
EulerBasis::ZYZ,
EulerBasis::ZXZ,
EulerBasis::XZX,
EulerBasis::XYX,
EulerBasis::ZSXX,
EulerBasis::ZSX,
];

fn matmul_1q(operator: &mut [[Complex64; 2]; 2], other: Array2<Complex64>) {
*operator = [
[
Expand All @@ -1054,7 +1194,6 @@ pub fn euler_one_qubit_decomposer(m: &Bound<PyModule>) -> PyResult<()> {
m.add_wrapped(wrap_pyfunction!(generate_circuit))?;
m.add_wrapped(wrap_pyfunction!(unitary_to_gate_sequence))?;
m.add_wrapped(wrap_pyfunction!(unitary_to_circuit))?;
m.add_wrapped(wrap_pyfunction!(compute_error_one_qubit_sequence))?;
m.add_wrapped(wrap_pyfunction!(compute_error_list))?;
m.add_wrapped(wrap_pyfunction!(optimize_1q_gates_decomposition))?;
m.add_class::<OneQubitGateSequence>()?;
Expand Down
Loading

0 comments on commit 0c8b668

Please sign in to comment.