Skip to content

Commit

Permalink
Leverage native UnitaryGate from rust
Browse files Browse the repository at this point in the history
This commit builds off of the native rust representation of a
UnitaryGate added in #13759 and uses the native representation
everywhere we were using UnitaryGate in rust via python previously:
the quantum_volume() function, consolidate blocks, split2qunitaries,
and unitary synthesis.

One future item is consolidate blocks can be updated to use nalgebra
types internally instead of ndarray as for the 1 and 2q cases we know
the fixed size of the array ahead of time. However the block
consolidation code is built using ndarray currently and later synthesis
code also works in ndarray so there isn't any real benefit yet, and we'd
just add unecessary conversions and allocations. However, once #13649
merges this will change and it would make more sense to add the unitary
gate with a Matrix4. But this can be handled separately after this
merges.
  • Loading branch information
mtreinish committed Feb 5, 2025
1 parent 936ec8e commit 8bdc469
Show file tree
Hide file tree
Showing 8 changed files with 221 additions and 148 deletions.
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/accelerate/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ ndarray_einsum_beta = "0.7"
once_cell = "1.20.2"
rustiq-core = "0.0.10"
bytemuck.workspace = true
nalgebra.workspace = true

[dependencies.smallvec]
workspace = true
Expand Down
56 changes: 16 additions & 40 deletions crates/accelerate/src/circuit_library/quantum_volume.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,19 +15,15 @@ use pyo3::prelude::*;
use pyo3::types::PyDict;

use crate::getenv_use_multiple_threads;
use faer_ext::{IntoFaerComplex, IntoNdarrayComplex};
use ndarray::prelude::*;
use num_complex::Complex64;
use numpy::IntoPyArray;
use nalgebra::Matrix4;
use num_complex::{Complex64, ComplexFloat};
use rand::prelude::*;
use rand_distr::StandardNormal;
use rand_pcg::Pcg64Mcg;
use rayon::prelude::*;

use qiskit_circuit::circuit_data::CircuitData;
use qiskit_circuit::imports::UNITARY_GATE;
use qiskit_circuit::operations::Param;
use qiskit_circuit::operations::PyGate;
use qiskit_circuit::operations::{ArrayType, Param, UnitaryGate};
use qiskit_circuit::packed_instruction::PackedOperation;
use qiskit_circuit::{Clbit, Qubit};
use smallvec::{smallvec, SmallVec};
Expand All @@ -50,11 +46,11 @@ fn random_complex(rng: &mut Pcg64Mcg) -> Complex64 {
//
// https://github.com/scipy/scipy/blob/v1.14.1/scipy/stats/_multivariate.py#L4224-L4256
#[inline]
fn random_unitaries(seed: u64, size: usize) -> impl Iterator<Item = Array2<Complex64>> {
fn random_unitaries(seed: u64, size: usize) -> impl Iterator<Item = Matrix4<Complex64>> {
let mut rng = Pcg64Mcg::seed_from_u64(seed);

(0..size).map(move |_| {
let raw_numbers: [[Complex64; 4]; 4] = [
let mat: Matrix4<Complex64> = [
[
random_complex(&mut rng),
random_complex(&mut rng),
Expand All @@ -79,23 +75,11 @@ fn random_unitaries(seed: u64, size: usize) -> impl Iterator<Item = Array2<Compl
random_complex(&mut rng),
random_complex(&mut rng),
],
];

let qr = aview2(&raw_numbers).into_faer_complex().qr();
let r = qr.compute_r();
let diag: [Complex64; 4] = [
r[(0, 0)].to_num_complex() / r[(0, 0)].abs(),
r[(1, 1)].to_num_complex() / r[(1, 1)].abs(),
r[(2, 2)].to_num_complex() / r[(2, 2)].abs(),
r[(3, 3)].to_num_complex() / r[(3, 3)].abs(),
];
let mut q = qr.compute_q().as_ref().into_ndarray_complex().to_owned();
q.axis_iter_mut(Axis(0)).for_each(|mut row| {
row.iter_mut()
.enumerate()
.for_each(|(index, val)| *val *= diag[index])
});
q
]
.into();
let (q, r) = mat.qr().unpack();
let diag = r.map_diagonal(|x| x / x.abs());
q.map_with_location(|i, _j, val| val * diag[i])
})
}

Expand All @@ -115,29 +99,21 @@ pub fn quantum_volume(

let kwargs = PyDict::new(py);
kwargs.set_item(intern!(py, "num_qubits"), 2)?;
let mut build_instruction = |(unitary_index, unitary_array): (usize, Array2<Complex64>),
let mut build_instruction = |(unitary_index, unitary_array): (usize, Matrix4<Complex64>),
rng: &mut Pcg64Mcg|
-> PyResult<Instruction> {
let layer_index = unitary_index % width;
if layer_index == 0 {
permutation.shuffle(rng);
}
let unitary = unitary_array.into_pyarray(py);

let unitary_gate = UNITARY_GATE
.get_bound(py)
.call((unitary.clone(), py.None(), false), Some(&kwargs))?;
let instruction = PyGate {
qubits: 2,
clbits: 0,
params: 1,
op_name: "unitary".to_string(),
gate: unitary_gate.unbind(),
let unitary_gate = UnitaryGate {
array: ArrayType::TwoQ(unitary_array),
};
let qubit = layer_index * 2;
Ok((
PackedOperation::from_gate(Box::new(instruction)),
smallvec![Param::Obj(unitary.into_any().unbind())],
PackedOperation::from_unitary(Box::new(unitary_gate)),
smallvec![],
vec![permutation[qubit], permutation[qubit + 1]],
vec![],
))
Expand All @@ -156,7 +132,7 @@ pub fn quantum_volume(
.take(num_unitaries)
.collect();

let unitaries: Vec<Array2<Complex64>> = if getenv_use_multiple_threads() && num_unitaries > 200
let unitaries: Vec<Matrix4<Complex64>> = if getenv_use_multiple_threads() && num_unitaries > 200
{
seed_vec
.par_chunks(per_thread)
Expand Down
89 changes: 53 additions & 36 deletions crates/accelerate/src/consolidate_blocks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,18 +11,22 @@
// that they have been altered from the originals.

use hashbrown::{HashMap, HashSet};
use nalgebra::Matrix2;
use ndarray::{aview2, Array2};
use num_complex::Complex64;
use numpy::{IntoPyArray, PyReadonlyArray2};
use numpy::PyReadonlyArray2;
use pyo3::intern;
use pyo3::prelude::*;
use rustworkx_core::petgraph::stable_graph::NodeIndex;
use smallvec::smallvec;

use qiskit_circuit::circuit_data::CircuitData;
use qiskit_circuit::circuit_instruction::ExtraInstructionAttributes;
use qiskit_circuit::dag_circuit::DAGCircuit;
use qiskit_circuit::gate_matrix::{ONE_QUBIT_IDENTITY, TWO_QUBIT_IDENTITY};
use qiskit_circuit::imports::{QI_OPERATOR, QUANTUM_CIRCUIT, UNITARY_GATE};
use qiskit_circuit::operations::{Operation, Param};
use qiskit_circuit::imports::{QI_OPERATOR, QUANTUM_CIRCUIT};
use qiskit_circuit::operations::{ArrayType, Operation, Param, UnitaryGate};
use qiskit_circuit::packed_instruction::PackedOperation;
use qiskit_circuit::Qubit;

use crate::convert_2q_block_matrix::{blocks_to_matrix, get_matrix_from_inst};
Expand Down Expand Up @@ -112,11 +116,17 @@ pub(crate) fn consolidate_blocks(
Ok(mat) => mat,
Err(_) => continue,
};
let array = matrix.into_pyarray(py);
let unitary_gate = UNITARY_GATE
.get_bound(py)
.call1((array, py.None(), false))?;
dag.substitute_node_with_py_op(py, inst_node, &unitary_gate, false)?;
// TODO: Use Matrix2/ArrayType::OneQ when we're using nalgebra
// for consolidation
let unitary_gate = UnitaryGate {
array: ArrayType::NDArray(matrix),
};
dag.substitute_op(
inst_node,
PackedOperation::from_unitary(Box::new(unitary_gate)),
smallvec![],
ExtraInstructionAttributes::default(),
)?;
continue;
}
}
Expand Down Expand Up @@ -180,16 +190,16 @@ pub(crate) fn consolidate_blocks(
dag.remove_op_node(node);
}
} else {
let unitary_gate = UNITARY_GATE.get_bound(py).call1((
array.as_ref().into_pyobject(py)?,
py.None(),
false,
))?;
let matrix = array.as_array().to_owned();
let unitary_gate = UnitaryGate {
array: ArrayType::NDArray(matrix),
};
let clbit_pos_map = HashMap::new();
dag.replace_block_with_py_op(
py,
dag.replace_block(
&block,
unitary_gate,
PackedOperation::from_unitary(Box::new(unitary_gate)),
smallvec![],
ExtraInstructionAttributes::default(),
false,
&block_index_map,
&clbit_pos_map,
Expand All @@ -213,21 +223,22 @@ pub(crate) fn consolidate_blocks(
dag.remove_op_node(node);
}
} else {
let array = matrix.into_pyarray(py);
let unitary_gate =
UNITARY_GATE
.get_bound(py)
.call1((array, py.None(), false))?;
// TODO: Use Matrix4/ArrayType::TwoQ when we're using nalgebra
// for consolidation
let unitary_gate = UnitaryGate {
array: ArrayType::NDArray(matrix),
};
let qubit_pos_map = block_index_map
.into_iter()
.enumerate()
.map(|(idx, qubit)| (qubit, idx))
.collect();
let clbit_pos_map = HashMap::new();
dag.replace_block_with_py_op(
py,
dag.replace_block(
&block,
unitary_gate,
PackedOperation::from_unitary(Box::new(unitary_gate)),
smallvec![],
ExtraInstructionAttributes::default(),
false,
&qubit_pos_map,
&clbit_pos_map,
Expand Down Expand Up @@ -258,11 +269,15 @@ pub(crate) fn consolidate_blocks(
Ok(mat) => mat,
Err(_) => continue,
};
let array = matrix.into_pyarray(py);
let unitary_gate = UNITARY_GATE
.get_bound(py)
.call1((array, py.None(), false))?;
dag.substitute_node_with_py_op(py, first_inst_node, &unitary_gate, false)?;
let unitary_gate = UnitaryGate {
array: ArrayType::NDArray(matrix),
};
dag.substitute_op(
first_inst_node,
PackedOperation::from_unitary(Box::new(unitary_gate)),
smallvec![],
ExtraInstructionAttributes::default(),
)?;
continue;
}
let qubit = first_qubits[0];
Expand Down Expand Up @@ -293,17 +308,19 @@ pub(crate) fn consolidate_blocks(
dag.remove_op_node(node);
}
} else {
let array = aview2(&matrix).to_owned().into_pyarray(py);
let unitary_gate = UNITARY_GATE
.get_bound(py)
.call1((array, py.None(), false))?;
let array: Matrix2<Complex64> =
Matrix2::from_row_iterator(matrix.into_iter().flat_map(|x| x.into_iter()));
let unitary_gate = UnitaryGate {
array: ArrayType::OneQ(array),
};
let mut block_index_map: HashMap<Qubit, usize> = HashMap::with_capacity(1);
block_index_map.insert(qubit, 0);
let clbit_pos_map = HashMap::new();
dag.replace_block_with_py_op(
py,
dag.replace_block(
&run,
unitary_gate,
PackedOperation::from_unitary(Box::new(unitary_gate)),
smallvec![],
ExtraInstructionAttributes::default(),
false,
&block_index_map,
&clbit_pos_map,
Expand Down
49 changes: 30 additions & 19 deletions crates/accelerate/src/split_2q_unitaries.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,15 +10,15 @@
// copyright notice, and modified files need to carry a notice indicating
// that they have been altered from the originals.

use pyo3::intern;
use nalgebra::Matrix2;
use num_complex::Complex64;
use pyo3::prelude::*;
use pyo3::types::PyDict;
use rustworkx_core::petgraph::stable_graph::NodeIndex;
use smallvec::{smallvec, SmallVec};

use qiskit_circuit::circuit_instruction::OperationFromPython;
use qiskit_circuit::dag_circuit::{DAGCircuit, NodeType, Wire};
use qiskit_circuit::imports::UNITARY_GATE;
use qiskit_circuit::operations::{Operation, Param};
use qiskit_circuit::operations::{ArrayType, Operation, OperationRef, Param, UnitaryGate};
use qiskit_circuit::packed_instruction::PackedOperation;

use crate::two_qubit_decompose::{Specialization, TwoQubitWeylDecomposition};

Expand All @@ -39,7 +39,7 @@ pub fn split_2q_unitaries(
// We only attempt to split UnitaryGate objects, but this could be extended in future
// -- however we need to ensure that we can compile the resulting single-qubit unitaries
// to the supported basis gate set.
if qubits.len() != 2 || inst.op.name() != "unitary" {
if qubits.len() != 2 || !matches!(inst.op.view(), OperationRef::Unitary(_)) {
continue;
}
let matrix = inst
Expand All @@ -52,22 +52,33 @@ pub fn split_2q_unitaries(
None,
)?;
if matches!(decomp.specialization, Specialization::IdEquiv) {
let k1r_arr = decomp.K1r(py);
let k1l_arr = decomp.K1l(py);
let kwargs = PyDict::new(py);
kwargs.set_item(intern!(py, "num_qubits"), 1)?;
let k1r_gate = UNITARY_GATE
.get_bound(py)
.call((k1r_arr, py.None(), false), Some(&kwargs))?;
let k1l_gate = UNITARY_GATE
.get_bound(py)
.call((k1l_arr, py.None(), false), Some(&kwargs))?;
let insert_fn = |edge: &Wire| -> PyResult<OperationFromPython> {
let k1r_arr = decomp.k1r_view();
let k1l_arr = decomp.k1l_view();

let insert_fn = |edge: &Wire| -> (PackedOperation, SmallVec<[Param; 3]>) {
if let Wire::Qubit(qubit) = edge {
if *qubit == qubits[0] {
k1r_gate.extract()
let mat: Matrix2<Complex64> = [
[k1r_arr[[0, 0]], k1r_arr[[0, 1]]],
[k1r_arr[[1, 0]], k1r_arr[[1, 1]]],
]
.into();
let k1r_gate = Box::new(UnitaryGate {
array: ArrayType::OneQ(mat),
});
(PackedOperation::from_unitary(k1r_gate), smallvec![])
} else {
k1l_gate.extract()
let mat: Matrix2<Complex64> = [
[k1l_arr[[0, 0]], k1l_arr[[0, 1]]],
[k1l_arr[[1, 0]], k1l_arr[[1, 1]]],
]
.into();

let k1l_gate = Box::new(UnitaryGate {
array: ArrayType::OneQ(mat),
});

(PackedOperation::from_unitary(k1l_gate), smallvec![])
}
} else {
unreachable!("This will only be called on ops with no classical wires.");
Expand Down
17 changes: 17 additions & 0 deletions crates/accelerate/src/two_qubit_decompose.rs
Original file line number Diff line number Diff line change
Expand Up @@ -531,6 +531,23 @@ impl TwoQubitWeylDecomposition {
pub fn c(&self) -> f64 {
self.c
}

pub fn k1l_view(&self) -> ArrayView2<Complex64> {
self.K1l.view()
}

pub fn k2l_view(&self) -> ArrayView2<Complex64> {
self.K2l.view()
}

pub fn k1r_view(&self) -> ArrayView2<Complex64> {
self.K1r.view()
}

pub fn k2r_view(&self) -> ArrayView2<Complex64> {
self.K2r.view()
}

fn weyl_gate(
&self,
simplify: bool,
Expand Down
2 changes: 1 addition & 1 deletion crates/accelerate/src/unitary_synthesis.rs
Original file line number Diff line number Diff line change
Expand Up @@ -289,7 +289,7 @@ fn py_run_main_loop(
py_op: new_node.unbind().into(),
};
}
if !(packed_instr.op.name() == "unitary"
if !(matches!(packed_instr.op.view(), OperationRef::Unitary(_))
&& packed_instr.op.num_qubits() >= min_qubits as u32)
{
out_dag.push_back(py, packed_instr)?;
Expand Down
Loading

0 comments on commit 8bdc469

Please sign in to comment.