From 8bdc469139bcf47970922f045b64626e0473c4c0 Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Thu, 30 Jan 2025 12:19:45 -0500 Subject: [PATCH] Leverage native UnitaryGate from rust 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. --- Cargo.lock | 1 + crates/accelerate/Cargo.toml | 1 + .../src/circuit_library/quantum_volume.rs | 56 ++----- crates/accelerate/src/consolidate_blocks.rs | 89 ++++++---- crates/accelerate/src/split_2q_unitaries.rs | 49 +++--- crates/accelerate/src/two_qubit_decompose.rs | 17 ++ crates/accelerate/src/unitary_synthesis.rs | 2 +- crates/circuit/src/dag_circuit.rs | 154 ++++++++++++------ 8 files changed, 221 insertions(+), 148 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 40f083bc5469..4fe2aee82e39 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1234,6 +1234,7 @@ dependencies = [ "hashbrown 0.14.5", "indexmap", "itertools 0.13.0", + "nalgebra", "ndarray", "ndarray_einsum_beta", "num-bigint", diff --git a/crates/accelerate/Cargo.toml b/crates/accelerate/Cargo.toml index b17098930b5c..8058139c4f2c 100644 --- a/crates/accelerate/Cargo.toml +++ b/crates/accelerate/Cargo.toml @@ -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 diff --git a/crates/accelerate/src/circuit_library/quantum_volume.rs b/crates/accelerate/src/circuit_library/quantum_volume.rs index 463c123eabb7..04b4fa3d4a0c 100644 --- a/crates/accelerate/src/circuit_library/quantum_volume.rs +++ b/crates/accelerate/src/circuit_library/quantum_volume.rs @@ -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}; @@ -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> { +fn random_unitaries(seed: u64, size: usize) -> impl Iterator> { let mut rng = Pcg64Mcg::seed_from_u64(seed); (0..size).map(move |_| { - let raw_numbers: [[Complex64; 4]; 4] = [ + let mat: Matrix4 = [ [ random_complex(&mut rng), random_complex(&mut rng), @@ -79,23 +75,11 @@ fn random_unitaries(seed: u64, size: usize) -> impl Iterator), + let mut build_instruction = |(unitary_index, unitary_array): (usize, Matrix4), rng: &mut Pcg64Mcg| -> PyResult { 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![], )) @@ -156,7 +132,7 @@ pub fn quantum_volume( .take(num_unitaries) .collect(); - let unitaries: Vec> = if getenv_use_multiple_threads() && num_unitaries > 200 + let unitaries: Vec> = if getenv_use_multiple_threads() && num_unitaries > 200 { seed_vec .par_chunks(per_thread) diff --git a/crates/accelerate/src/consolidate_blocks.rs b/crates/accelerate/src/consolidate_blocks.rs index 30022158e2e5..a4909aef6b5f 100644 --- a/crates/accelerate/src/consolidate_blocks.rs +++ b/crates/accelerate/src/consolidate_blocks.rs @@ -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}; @@ -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; } } @@ -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, @@ -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, @@ -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]; @@ -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 = + 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 = 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, diff --git a/crates/accelerate/src/split_2q_unitaries.rs b/crates/accelerate/src/split_2q_unitaries.rs index 8d54e8875c42..366a333c7a78 100644 --- a/crates/accelerate/src/split_2q_unitaries.rs +++ b/crates/accelerate/src/split_2q_unitaries.rs @@ -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}; @@ -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 @@ -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 { + 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 = [ + [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 = [ + [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."); diff --git a/crates/accelerate/src/two_qubit_decompose.rs b/crates/accelerate/src/two_qubit_decompose.rs index de20430d25bf..5bfc952c548a 100644 --- a/crates/accelerate/src/two_qubit_decompose.rs +++ b/crates/accelerate/src/two_qubit_decompose.rs @@ -531,6 +531,23 @@ impl TwoQubitWeylDecomposition { pub fn c(&self) -> f64 { self.c } + + pub fn k1l_view(&self) -> ArrayView2 { + self.K1l.view() + } + + pub fn k2l_view(&self) -> ArrayView2 { + self.K2l.view() + } + + pub fn k1r_view(&self) -> ArrayView2 { + self.K1r.view() + } + + pub fn k2r_view(&self) -> ArrayView2 { + self.K2r.view() + } + fn weyl_gate( &self, simplify: bool, diff --git a/crates/accelerate/src/unitary_synthesis.rs b/crates/accelerate/src/unitary_synthesis.rs index 2721c58c9237..16478f87f749 100644 --- a/crates/accelerate/src/unitary_synthesis.rs +++ b/crates/accelerate/src/unitary_synthesis.rs @@ -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)?; diff --git a/crates/circuit/src/dag_circuit.rs b/crates/circuit/src/dag_circuit.rs index 42b334e88207..5531aa7a8524 100644 --- a/crates/circuit/src/dag_circuit.rs +++ b/crates/circuit/src/dag_circuit.rs @@ -2864,10 +2864,13 @@ def _format(operand): } let block_ids: Vec<_> = node_block.iter().map(|n| n.node.unwrap()).collect(); - let new_node = self.replace_block_with_py_op( - py, + let py_op = op.extract::()?; + + let new_node = self.replace_block( &block_ids, - op, + py_op.operation, + py_op.params, + py_op.extra_attrs, cycle_check, &qubit_pos_map, &clbit_pos_map, @@ -6343,7 +6346,7 @@ impl DAGCircuit { insert: F, ) -> PyResult<()> where - F: Fn(&Wire) -> PyResult, + F: Fn(&Wire) -> (PackedOperation, SmallVec<[Param; 3]>), { let mut edge_list: Vec<(NodeIndex, NodeIndex, Wire)> = Vec::with_capacity(2); for (source, in_weight) in self @@ -6362,15 +6365,15 @@ impl DAGCircuit { } } for (source, target, weight) in edge_list { - let new_op = insert(&weight)?; - self.increment_op(new_op.operation.name()); + let (new_op, params) = insert(&weight); + self.increment_op(new_op.name()); let qubits = if let Wire::Qubit(qubit) = weight { vec![qubit] } else { panic!("This method only works if the gate being replaced has no classical incident wires") }; #[cfg(feature = "cache_pygates")] - let py_op = match new_op.operation.view() { + let py_op = match new_op.view() { OperationRef::StandardGate(_) | OperationRef::StandardInstruction(_) | OperationRef::Unitary(_) => OnceLock::new(), @@ -6381,11 +6384,11 @@ impl DAGCircuit { OperationRef::Operation(op) => OnceLock::from(op.operation.clone_ref(py)), }; let inst = PackedInstruction { - op: new_op.operation, + op: new_op, qubits: self.qargs_interner.insert_owned(qubits), clbits: self.cargs_interner.get_default(), - params: (!new_op.params.is_empty()).then(|| Box::new(new_op.params)), - extra_attrs: new_op.extra_attrs, + params: (!params.is_empty()).then(|| Box::new(params)), + extra_attrs: ExtraInstructionAttributes::default(), #[cfg(feature = "cache_pygates")] py_op, }; @@ -6850,12 +6853,14 @@ impl DAGCircuit { Self::from_circuit(py, circ, copy_op, None, None) } - /// Replace a block of node indices with a new python operation - pub fn replace_block_with_py_op( + #[allow(clippy::too_many_arguments)] + /// Replace a block of node indices with a new packed operation + pub fn replace_block( &mut self, - py: Python, block_ids: &[NodeIndex], - op: Bound, + op: PackedOperation, + params: SmallVec<[Param; 3]>, + extra_attrs: ExtraInstructionAttributes, cycle_check: bool, qubit_pos_map: &HashMap, clbit_pos_map: &HashMap, @@ -6872,41 +6877,47 @@ impl DAGCircuit { block_cargs.extend(self.cargs_interner.get(packed.clbits)); if let Some(condition) = packed.condition() { - block_cargs.extend( - self.clbits.map_bits( - self.control_flow_module - .condition_resources(condition.bind(py))? - .clbits - .bind(py), - )?, - ); + Python::with_gil(|py| -> PyResult<()> { + block_cargs.extend( + self.clbits.map_bits( + self.control_flow_module + .condition_resources(condition.bind(py))? + .clbits + .bind(py), + )?, + ); + Ok(()) + })?; continue; } // Add classical bits from SwitchCaseOp, if applicable. if let OperationRef::Instruction(op) = packed.op.view() { if op.name() == "switch_case" { - let op_bound = op.instruction.bind(py); - let target = op_bound.getattr(intern!(py, "target"))?; - if target.is_instance(imports::CLBIT.get_bound(py))? { - block_cargs.insert(self.clbits.find(&target).unwrap()); - } else if target - .is_instance(imports::CLASSICAL_REGISTER.get_bound(py))? - { - block_cargs.extend( - self.clbits - .map_bits(target.extract::>>()?)?, - ); - } else { - block_cargs.extend( - self.clbits.map_bits( - self.control_flow_module - .node_resources(&target)? - .clbits - .bind(py), - )?, - ); - } + Python::with_gil(|py| -> PyResult<()> { + let op_bound = op.instruction.bind(py); + let target = op_bound.getattr(intern!(py, "target"))?; + if target.is_instance(imports::CLBIT.get_bound(py))? { + block_cargs.insert(self.clbits.find(&target).unwrap()); + } else if target + .is_instance(imports::CLASSICAL_REGISTER.get_bound(py))? + { + block_cargs.extend( + self.clbits + .map_bits(target.extract::>>()?)?, + ); + } else { + block_cargs.extend( + self.clbits.map_bits( + self.control_flow_module + .node_resources(&target)? + .clbits + .bind(py), + )?, + ); + } + Ok(()) + })?; } } } @@ -6935,25 +6946,23 @@ impl DAGCircuit { .collect(); block_cargs.sort_by_key(|c| clbit_pos_map[c]); - let py_op = op.extract::()?; - - if py_op.operation.num_qubits() as usize != block_qargs.len() { + if op.num_qubits() as usize != block_qargs.len() { return Err(DAGCircuitError::new_err(format!( - "Number of qubits in the replacement operation ({}) is not equal to the number of qubits in the block ({})!", py_op.operation.num_qubits(), block_qargs.len() + "Number of qubits in the replacement operation ({}) is not equal to the number of qubits in the block ({})!", op.num_qubits(), block_qargs.len() ))); } - let op_name = py_op.operation.name().to_string(); + let op_name = op.name().to_string(); let qubits = self.qargs_interner.insert_owned(block_qargs); let clbits = self.cargs_interner.insert_owned(block_cargs); let weight = NodeType::Operation(PackedInstruction { - op: py_op.operation, + op, qubits, clbits, - params: (!py_op.params.is_empty()).then(|| Box::new(py_op.params)), - extra_attrs: py_op.extra_attrs, + params: (!params.is_empty()).then(|| Box::new(params)), + extra_attrs, #[cfg(feature = "cache_pygates")] - py_op: op.unbind().into(), + py_op: OnceLock::new(), }); let new_node = self @@ -6972,6 +6981,47 @@ impl DAGCircuit { Ok(new_node) } + /// Substitute an operation in a node with a new one. The wire counts must match and the same + /// argument order will be used. + pub fn substitute_op( + &mut self, + node_index: NodeIndex, + new_op: PackedOperation, + params: SmallVec<[Param; 3]>, + extra_attrs: ExtraInstructionAttributes, + ) -> PyResult<()> { + let old_packed = self.dag[node_index].unwrap_operation(); + let op_name = old_packed.op.name().to_string(); + + if old_packed.op.num_qubits() != new_op.num_qubits() + || old_packed.op.num_clbits() != new_op.num_clbits() + { + return Err(DAGCircuitError::new_err( + format!( + "Cannot replace node of width ({} qubits, {} clbits) with operation of mismatched width ({} qubits, {} clbits)", + old_packed.op.num_qubits(), old_packed.op.num_clbits(), new_op.num_qubits(), new_op.num_clbits() + ))); + } + let new_op_name = new_op.name().to_string(); + let new_weight = NodeType::Operation(PackedInstruction { + op: new_op, + qubits: old_packed.qubits, + clbits: old_packed.clbits, + params: (!params.is_empty()).then(|| params.into()), + extra_attrs, + #[cfg(feature = "cache_pygates")] + py_op: OnceLock::new(), + }); + if let Some(weight) = self.dag.node_weight_mut(node_index) { + *weight = new_weight; + } + + // Update self.op_names + self.decrement_op(op_name.as_str()); + self.increment_op(new_op_name.as_str()); + Ok(()) + } + /// Substitute a give node in the dag with a new operation from python pub fn substitute_node_with_py_op( &mut self,