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

@@ -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;

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<NodeIndex> = 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<OperationFromPython> {
+                    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<PyModule>) -> PyResult<()> {
+    m.add_wrapped(wrap_pyfunction!(split_2q_unitaries))?;
+    Ok(())
+}

crates/accelerate/src/two_qubit_decompose.rs

@@ -340,7 +340,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,
@@ -409,13 +409,13 @@ pub struct TwoQubitWeylDecomposition {
     #[pyo3(get)]
     c: f64,
     #[pyo3(get)]
-    global_phase: f64,
+    pub global_phase: f64,
     K1l: Array2<Complex64>,
     K2l: Array2<Complex64>,
     K1r: Array2<Complex64>,
     K2r: Array2<Complex64>,
     #[pyo3(get)]
-    specialization: Specialization,
+    pub specialization: Specialization,
     default_euler_basis: EulerBasis,
     #[pyo3(get)]
     requested_fidelity: Option<f64>,
@@ -475,7 +475,7 @@ impl TwoQubitWeylDecomposition {
 
     /// Instantiate a new TwoQubitWeylDecomposition with rust native
     /// data structures
-    fn new_inner(
+    pub fn new_inner(
         unitary_matrix: ArrayView2<Complex64>,
 
         fidelity: Option<f64>,
@@ -1020,13 +1020,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()
     }
 

crates/circuit/src/dag_circuit.rs

@@ -5262,7 +5262,7 @@ impl DAGCircuit {
         Ok(nodes.into_iter())
     }
 
-    fn topological_op_nodes(&self) -> PyResult<impl Iterator<Item = NodeIndex> + '_> {
+    pub fn topological_op_nodes(&self) -> PyResult<impl Iterator<Item = NodeIndex> + '_> {
         Ok(self.topological_nodes()?.filter(|node: &NodeIndex| {
             matches!(self.dag.node_weight(*node), Some(NodeType::Operation(_)))
         }))
@@ -6131,6 +6131,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,
@@ -6193,6 +6197,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<F>(
+        &mut self,
+        py: Python, // Unused if cache_pygates isn't enabled
+        node: NodeIndex,
+        mut insert: F,
+    ) -> PyResult<()>
+    where
+        F: FnMut(&Wire) -> PyResult<OperationFromPython>,
+    {
+        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(_) => {

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

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<PyModule>) -> 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")?;

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
 

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