From 2ef371ae0d159a6dfd643805f3e5e5fdec37ab88 Mon Sep 17 00:00:00 2001 From: Matthew Treinish Date: Mon, 9 Sep 2024 09:09:24 -0400 Subject: [PATCH] Fully port Split2QUnitaries to rust (#13025) * 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 * Update pass logic with changes from #13095 Some of the logic inside the Split2QUnitaries pass was updated in a recently merged PR. This commit makes those changes so the rust implementation matches the current state of the previous python version. * Use op_nodes() instead of topological_op_nodes() * Use Fn trait instead of FnMut for callback We don't need the callback to be mutable currently so relax the trait to just be `Fn` instead of `FnMut`. If we have a need for a mutable environment callback in the future we can change this easily enough without any issues. * Avoid extra edge operations in replace_on_incoming_qubits * Rename function --- crates/accelerate/src/lib.rs | 1 + crates/accelerate/src/split_2q_unitaries.rs | 75 +++++++++++++++++++ crates/accelerate/src/two_qubit_decompose.rs | 12 +-- crates/circuit/src/dag_circuit.rs | 71 +++++++++++++++++- crates/circuit/src/imports.rs | 4 + crates/pyext/src/lib.rs | 8 +- qiskit/__init__.py | 1 + .../passes/optimization/split_2q_unitaries.py | 44 +---------- 8 files changed, 164 insertions(+), 52 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 916aeeb8351c..9e391e0a9030 100644 --- a/crates/accelerate/src/lib.rs +++ b/crates/accelerate/src/lib.rs @@ -35,6 +35,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..83ab94bf301d --- /dev/null +++ b/crates/accelerate/src/split_2q_unitaries.rs @@ -0,0 +1,75 @@ +// 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.op_nodes(false).collect(); + for node in nodes { + if let NodeType::Operation(inst) = &dag.dag[node] { + let qubits = dag.get_qargs(inst.qubits).to_vec(); + let matrix = inst.op.matrix(inst.params_view()); + // 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" { + 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_node_with_1q_ops(py, node, insert_fn)?; + dag.add_global_phase(py, &Param::Float(decomp.global_phase))?; + } + // TODO: also look into splitting on Specialization::Swap and just + // swap the virtual qubits. Doing this we will need to update the + // permutation like in ElidePermutations + } + } + 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 2c745b199e9e..92ad4724682f 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 5b9769b4d19e..0b5a43c1eb4b 100644 --- a/crates/circuit/src/dag_circuit.rs +++ b/crates/circuit/src/dag_circuit.rs @@ -5254,7 +5254,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(_))) })) @@ -6285,6 +6285,75 @@ impl DAGCircuit { } } + /// Replace a node with individual operations from a provided callback + /// function on each qubit of that node. + #[allow(unused_variables)] + pub fn replace_node_with_1q_ops( + &mut self, + py: Python, // Unused if cache_pygates isn't enabled + node: NodeIndex, + insert: F, + ) -> PyResult<()> + where + F: Fn(&Wire) -> PyResult, + { + let mut edge_list: Vec<(NodeIndex, NodeIndex, Wire)> = Vec::with_capacity(2); + for (source, in_weight) in self + .dag + .edges_directed(node, Incoming) + .map(|x| (x.source(), x.weight())) + { + for (target, out_weight) in self + .dag + .edges_directed(node, Outgoing) + .map(|x| (x.target(), x.weight())) + { + if in_weight == out_weight { + edge_list.push((source, target, in_weight.clone())); + } + } + } + for (source, target, weight) in edge_list { + let new_op = insert(&weight)?; + self.increment_op(new_op.operation.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() { + 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: 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, + #[cfg(feature = "cache_pygates")] + py_op: py_op, + }; + let new_index = self.dag.add_node(NodeType::Operation(inst)); + self.dag.add_edge(source, new_index, weight.clone()); + self.dag.add_edge(new_index, target, weight); + } + + match self.dag.remove_node(node) { + Some(NodeType::Operation(packed)) => { + let op_name = packed.op.name(); + self.decrement_op(op_name); + } + _ => panic!("Must be called with valid operation 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 96b5e84a9cbe..03bd0202dae4 100644 --- a/crates/pyext/src/lib.rs +++ b/crates/pyext/src/lib.rs @@ -22,9 +22,10 @@ use qiskit_accelerate::{ optimize_1q_gates::optimize_1q_gates, pauli_exp_val::pauli_expval, remove_diagonal_gates_before_measure::remove_diagonal_gates_before_measure, 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)] @@ -65,6 +66,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 f731465beecb..29e2b8ba5c11 100644 --- a/qiskit/__init__.py +++ b/qiskit/__init__.py @@ -92,6 +92,7 @@ sys.modules["qiskit._accelerate.commutation_checker"] = _accelerate.commutation_checker sys.modules["qiskit._accelerate.commutation_analysis"] = _accelerate.commutation_analysis sys.modules["qiskit._accelerate.synthesis.linear_phase"] = _accelerate.synthesis.linear_phase +sys.modules["qiskit._accelerate.split_2q_unitaries"] = _accelerate.split_2q_unitaries sys.modules["qiskit._accelerate.gate_direction"] = _accelerate.gate_direction sys.modules["qiskit._accelerate.inverse_cancellation"] = _accelerate.inverse_cancellation sys.modules["qiskit._accelerate.check_map"] = _accelerate.check_map diff --git a/qiskit/transpiler/passes/optimization/split_2q_unitaries.py b/qiskit/transpiler/passes/optimization/split_2q_unitaries.py index 1ca53c6c31d2..f6958a00a4c1 100644 --- a/qiskit/transpiler/passes/optimization/split_2q_unitaries.py +++ b/qiskit/transpiler/passes/optimization/split_2q_unitaries.py @@ -13,11 +13,8 @@ """Splits each two-qubit gate in the `dag` into two single-qubit gates, if possible without error.""" from qiskit.transpiler.basepasses import TransformationPass -from qiskit.circuit.quantumcircuitdata import CircuitInstruction from qiskit.dagcircuit.dagcircuit import DAGCircuit -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): @@ -39,42 +36,5 @@ def __init__(self, fidelity: float = 1.0 - 1e-16): def run(self, dag: DAGCircuit) -> DAGCircuit: """Run the Split2QUnitaries pass on `dag`.""" - - for node in dag.topological_op_nodes(): - # 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 not (len(node.qargs) == 2 and node.op.name == "unitary"): - continue - - decomp = TwoQubitWeylDecomposition(node.matrix, 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