From 0689b736d9f5d49427f07fb6950e5416aedf861d Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Fri, 23 Aug 2024 09:57:49 -0400 Subject: [PATCH] Fully port Split2QUnitaries to rust This commit builds off of #13013 and the other data model in Rust infrastructure and migrates the InverseCancellation 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 for creating `UnitaryGate` instances and `ParameterExpression` 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 use loops 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 for different gates so we should be able to increase the throughput of the pass by leveraging multithreading to handle each gate 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. Part of #12208 --- crates/accelerate/src/lib.rs | 1 + crates/accelerate/src/split_2q_unitaries.rs | 74 +++++++++++++++++++ crates/accelerate/src/two_qubit_decompose.rs | 12 +-- crates/circuit/src/dag_circuit.rs | 60 ++++++++++++++- crates/circuit/src/imports.rs | 4 + crates/pyext/src/lib.rs | 8 +- qiskit/__init__.py | 1 + .../passes/optimization/split_2q_unitaries.py | 42 +---------- 8 files changed, 152 insertions(+), 50 deletions(-) create mode 100644 crates/accelerate/src/split_2q_unitaries.rs diff --git a/crates/accelerate/src/lib.rs b/crates/accelerate/src/lib.rs index fe9a9213ae8d..70b52100886d 100644 --- a/crates/accelerate/src/lib.rs +++ b/crates/accelerate/src/lib.rs @@ -28,6 +28,7 @@ pub mod results; pub mod sabre; pub mod sampled_exp_val; pub mod sparse_pauli_op; +pub mod split_2q_unitaries; pub mod star_prerouting; pub mod stochastic_swap; pub mod synthesis; diff --git a/crates/accelerate/src/split_2q_unitaries.rs b/crates/accelerate/src/split_2q_unitaries.rs new file mode 100644 index 000000000000..944e05c3a6b5 --- /dev/null +++ b/crates/accelerate/src/split_2q_unitaries.rs @@ -0,0 +1,74 @@ +// This code is part of Qiskit. +// +// (C) Copyright IBM 2024 +// +// This code is licensed under the Apache License, Version 2.0. You may +// obtain a copy of this license in the LICENSE.txt file in the root directory +// of this source tree or at http://www.apache.org/licenses/LICENSE-2.0. +// +// Any modifications or derivative works of this code must retain this +// copyright notice, and modified files need to carry a notice indicating +// that they have been altered from the originals. + +use pyo3::prelude::*; +use rustworkx_core::petgraph::stable_graph::NodeIndex; + +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 crate::two_qubit_decompose::{Specialization, TwoQubitWeylDecomposition}; + +#[pyfunction] +pub fn split_2q_unitaries( + py: Python, + dag: &mut DAGCircuit, + requested_fidelity: f64, +) -> PyResult<()> { + let nodes: Vec = dag.topological_op_nodes()?.collect(); + for node in nodes { + if let NodeType::Operation(inst) = &dag.dag[node] { + let qubits = dag.get_qubits(inst.qubits).to_vec(); + let matrix = inst.op.matrix(inst.params_view()); + if !dag.get_clbits(inst.clbits).is_empty() + || qubits.len() != 2 + || matrix.is_none() + || inst.is_parameterized() + || inst.condition().is_some() + { + continue; + } + let decomp = TwoQubitWeylDecomposition::new_inner( + matrix.unwrap().view(), + Some(requested_fidelity), + None, + )?; + if matches!(decomp.specialization, Specialization::IdEquiv) { + let k1r_arr = decomp.K1r(py); + let k1l_arr = decomp.K1l(py); + let k1r_gate = UNITARY_GATE.get_bound(py).call1((k1r_arr,))?; + let k1l_gate = UNITARY_GATE.get_bound(py).call1((k1l_arr,))?; + let insert_fn = |edge: &Wire| -> PyResult { + if let Wire::Qubit(qubit) = edge { + if *qubit == qubits[0] { + k1r_gate.extract() + } else { + k1l_gate.extract() + } + } else { + unreachable!("This will only be called on ops with no classical wires."); + } + }; + dag.replace_on_incoming_qubits(py, node, insert_fn)?; + dag.add_global_phase(py, &Param::Float(decomp.global_phase))?; + } + } + } + Ok(()) +} + +pub fn split_2q_unitaries_mod(m: &Bound) -> PyResult<()> { + m.add_wrapped(wrap_pyfunction!(split_2q_unitaries))?; + Ok(()) +} diff --git a/crates/accelerate/src/two_qubit_decompose.rs b/crates/accelerate/src/two_qubit_decompose.rs index 77ebb3f9433c..0ca5deb868e9 100644 --- a/crates/accelerate/src/two_qubit_decompose.rs +++ b/crates/accelerate/src/two_qubit_decompose.rs @@ -341,7 +341,7 @@ const DEFAULT_FIDELITY: f64 = 1.0 - 1.0e-9; #[derive(Clone, Debug, Copy)] #[pyclass(module = "qiskit._accelerate.two_qubit_decompose")] -enum Specialization { +pub enum Specialization { General, IdEquiv, SWAPEquiv, @@ -410,13 +410,13 @@ pub struct TwoQubitWeylDecomposition { #[pyo3(get)] c: f64, #[pyo3(get)] - global_phase: f64, + pub global_phase: f64, K1l: Array2, K2l: Array2, K1r: Array2, K2r: Array2, #[pyo3(get)] - specialization: Specialization, + pub specialization: Specialization, default_euler_basis: EulerBasis, #[pyo3(get)] requested_fidelity: Option, @@ -476,7 +476,7 @@ impl TwoQubitWeylDecomposition { /// Instantiate a new TwoQubitWeylDecomposition with rust native /// data structures - fn new_inner( + pub fn new_inner( unitary_matrix: ArrayView2, fidelity: Option, @@ -1021,13 +1021,13 @@ impl TwoQubitWeylDecomposition { #[allow(non_snake_case)] #[getter] - fn K1l(&self, py: Python) -> PyObject { + pub fn K1l(&self, py: Python) -> PyObject { self.K1l.to_pyarray_bound(py).into() } #[allow(non_snake_case)] #[getter] - fn K1r(&self, py: Python) -> PyObject { + pub fn K1r(&self, py: Python) -> PyObject { self.K1r.to_pyarray_bound(py).into() } diff --git a/crates/circuit/src/dag_circuit.rs b/crates/circuit/src/dag_circuit.rs index 3c97bf34945a..9afb3ded7844 100644 --- a/crates/circuit/src/dag_circuit.rs +++ b/crates/circuit/src/dag_circuit.rs @@ -5276,7 +5276,7 @@ impl DAGCircuit { Ok(nodes.into_iter()) } - fn topological_op_nodes(&self) -> PyResult + '_> { + pub fn topological_op_nodes(&self) -> PyResult + '_> { Ok(self.topological_nodes()?.filter(|node: &NodeIndex| { matches!(self.dag.node_weight(*node), Some(NodeType::Operation(_))) })) @@ -6140,6 +6140,10 @@ impl DAGCircuit { self.qargs_cache.intern(index) } + pub fn get_clbits(&self, index: Index) -> &[Clbit] { + self.cargs_cache.intern(index) + } + /// Insert a new 1q standard gate on incoming qubit pub fn insert_1q_on_incoming_qubit( &mut self, @@ -6202,6 +6206,60 @@ impl DAGCircuit { } } + /// Insert an op given by callback on each individual qubit into a node + + #[allow(unused_variables)] + pub fn replace_on_incoming_qubits( + &mut self, + py: Python, // Unused if cache_pygates isn't enabled + node: NodeIndex, + mut insert: F, + ) -> PyResult<()> + where + F: FnMut(&Wire) -> PyResult, + { + let edges: Vec<(NodeIndex, EdgeIndex, Wire)> = self + .dag + .edges_directed(node, Incoming) + .map(|edge| (edge.source(), edge.id(), edge.weight().clone())) + .collect(); + for (edge_source, edge_index, edge_weight) in edges { + let new_op = insert(&edge_weight)?; + self.increment_op(new_op.operation.name()); + let qubits = if let Wire::Qubit(qubit) = edge_weight { + vec![qubit] + } else { + panic!("This method only works if the gate being replaced") + }; + #[cfg(feature = "cache_pygates")] + let py_op = match new_op.operation.view() { + OperationRef::Standard(_) => OnceCell::new(), + OperationRef::Gate(gate) => OnceCell::from(gate.gate.clone_ref(py)), + OperationRef::Instruction(instruction) => { + OnceCell::from(instruction.instruction.clone_ref(py)) + } + OperationRef::Operation(op) => OnceCell::from(op.operation.clone_ref(py)), + }; + let inst = PackedInstruction { + op: new_op.operation, + qubits: Interner::intern(&mut self.qargs_cache, qubits)?, + clbits: Interner::intern(&mut self.cargs_cache, vec![])?, + params: (!new_op.params.is_empty()).then(|| Box::new(new_op.params)), + extra_attrs: new_op.extra_attrs, + #[cfg(feature = "cache_pygates")] + py_op: py_op, + }; + let new_index = self.dag.add_node(NodeType::Operation(inst)); + let parent_index = edge_source; + self.dag + .add_edge(parent_index, new_index, edge_weight.clone()); + self.dag.add_edge(new_index, node, edge_weight); + self.dag.remove_edge(edge_index); + } + self.remove_op_node(node); + Ok(()) + } + pub fn add_global_phase(&mut self, py: Python, value: &Param) -> PyResult<()> { match value { Param::Obj(_) => { diff --git a/crates/circuit/src/imports.rs b/crates/circuit/src/imports.rs index 588471546273..a8f6d9562559 100644 --- a/crates/circuit/src/imports.rs +++ b/crates/circuit/src/imports.rs @@ -109,6 +109,10 @@ pub static SWITCH_CASE_OP_CHECK: ImportOnceCell = pub static FOR_LOOP_OP_CHECK: ImportOnceCell = ImportOnceCell::new("qiskit.dagcircuit.dagnode", "_for_loop_eq"); pub static UUID: ImportOnceCell = ImportOnceCell::new("uuid", "UUID"); +pub static UNITARY_GATE: ImportOnceCell = ImportOnceCell::new( + "qiskit.circuit.library.generalized_gates.unitary", + "UnitaryGate", +); /// A mapping from the enum variant in crate::operations::StandardGate to the python /// module path and class name to import it. This is used to populate the conversion table diff --git a/crates/pyext/src/lib.rs b/crates/pyext/src/lib.rs index f215d567904a..f7e246efc0b7 100644 --- a/crates/pyext/src/lib.rs +++ b/crates/pyext/src/lib.rs @@ -18,9 +18,10 @@ use qiskit_accelerate::{ inverse_cancellation::inverse_cancellation_mod, isometry::isometry, nlayout::nlayout, optimize_1q_gates::optimize_1q_gates, pauli_exp_val::pauli_expval, results::results, sabre::sabre, sampled_exp_val::sampled_exp_val, sparse_pauli_op::sparse_pauli_op, - star_prerouting::star_prerouting, stochastic_swap::stochastic_swap, synthesis::synthesis, - target_transpiler::target, two_qubit_decompose::two_qubit_decompose, uc_gate::uc_gate, - utils::utils, vf2_layout::vf2_layout, + split_2q_unitaries::split_2q_unitaries_mod, star_prerouting::star_prerouting, + stochastic_swap::stochastic_swap, synthesis::synthesis, target_transpiler::target, + two_qubit_decompose::two_qubit_decompose, uc_gate::uc_gate, utils::utils, + vf2_layout::vf2_layout, }; #[inline(always)] @@ -53,6 +54,7 @@ fn _accelerate(m: &Bound) -> PyResult<()> { add_submodule(m, sabre, "sabre")?; add_submodule(m, sampled_exp_val, "sampled_exp_val")?; add_submodule(m, sparse_pauli_op, "sparse_pauli_op")?; + add_submodule(m, split_2q_unitaries_mod, "split_2q_unitaries")?; add_submodule(m, star_prerouting, "star_prerouting")?; add_submodule(m, stochastic_swap, "stochastic_swap")?; add_submodule(m, target, "target")?; diff --git a/qiskit/__init__.py b/qiskit/__init__.py index f5f63a7b737a..c07eb8fa7336 100644 --- a/qiskit/__init__.py +++ b/qiskit/__init__.py @@ -87,6 +87,7 @@ sys.modules["qiskit._accelerate.synthesis.clifford"] = _accelerate.synthesis.clifford sys.modules["qiskit._accelerate.synthesis.linear_phase"] = _accelerate.synthesis.linear_phase sys.modules["qiskit._accelerate.inverse_cancellation"] = _accelerate.inverse_cancellation +sys.modules["qiskit._accelerate.split_2q_unitaries"] = _accelerate.split_2q_unitaries from qiskit.exceptions import QiskitError, MissingOptionalLibraryError diff --git a/qiskit/transpiler/passes/optimization/split_2q_unitaries.py b/qiskit/transpiler/passes/optimization/split_2q_unitaries.py index ac04043a27fa..97f3a339da5a 100644 --- a/qiskit/transpiler/passes/optimization/split_2q_unitaries.py +++ b/qiskit/transpiler/passes/optimization/split_2q_unitaries.py @@ -18,6 +18,7 @@ from qiskit.dagcircuit.dagnode import DAGOpNode from qiskit.circuit.library.generalized_gates import UnitaryGate from qiskit.synthesis.two_qubit.two_qubit_decompose import TwoQubitWeylDecomposition +from qiskit._accelerate.split_2q_unitaries import split_2q_unitaries class Split2QUnitaries(TransformationPass): @@ -41,44 +42,5 @@ def __init__(self, fidelity: Optional[float] = 1.0 - 1e-16): def run(self, dag: DAGCircuit): """Run the Split2QUnitaries pass on `dag`.""" - for node in dag.topological_op_nodes(): - # skip operations without two-qubits and for which we can not determine a potential 1q split - if ( - len(node.cargs) > 0 - or len(node.qargs) != 2 - or node.matrix is None - or node.is_parameterized() - ): - continue - - decomp = TwoQubitWeylDecomposition(node.op, fidelity=self.requested_fidelity) - if ( - decomp._inner_decomposition.specialization - == TwoQubitWeylDecomposition._specializations.IdEquiv - ): - new_dag = DAGCircuit() - new_dag.add_qubits(node.qargs) - - ur = decomp.K1r - ur_node = DAGOpNode.from_instruction( - CircuitInstruction(UnitaryGate(ur), qubits=(node.qargs[0],)) - ) - - ul = decomp.K1l - ul_node = DAGOpNode.from_instruction( - CircuitInstruction(UnitaryGate(ul), qubits=(node.qargs[1],)) - ) - new_dag._apply_op_node_back(ur_node) - new_dag._apply_op_node_back(ul_node) - new_dag.global_phase = decomp.global_phase - dag.substitute_node_with_dag(node, new_dag) - elif ( - decomp._inner_decomposition.specialization - == TwoQubitWeylDecomposition._specializations.SWAPEquiv - ): - # TODO maybe also look into swap-gate-like gates? Things to consider: - # * As the qubit mapping may change, we'll always need to build a new dag in this pass - # * There may not be many swap-gate-like gates in an arbitrary input circuit - # * Removing swap gates from a user-routed input circuit here is unexpected - pass + split_2q_unitaries(dag, self.requested_fidelity) return dag