Update Split2QUnitaries to handle SWAP unitaries as well

crates/accelerate/src/split_2q_unitaries.rs

@@ -9,6 +9,8 @@
 // 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 std::f64::consts::PI;
+const PI4: f64 = PI / 4.;
 
 use pyo3::intern;
 use pyo3::prelude::*;
@@ -19,20 +21,64 @@ 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::packed_instruction::PackedInstruction;
+use qiskit_circuit::Qubit;
 
 use crate::two_qubit_decompose::{Specialization, TwoQubitWeylDecomposition};
 
+fn create_k1_gates<'a>(decomp: &'a TwoQubitWeylDecomposition, py: Python<'a>) -> PyResult<(Bound<'a,PyAny>, Bound<'a,PyAny>)> {
+    let k1r_arr = decomp.K1r(py);
+    let k1l_arr = decomp.K1l(py);
+    let kwargs = PyDict::new_bound(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))?;
+    return Ok((k1r_gate, k1l_gate));
+}    
+
+fn add_new_op(new_dag: &mut DAGCircuit, new_op: OperationFromPython, qargs: Vec<Qubit>, mapping: &Vec<usize>, py: Python) -> PyResult<()> {
+    let inst = PackedInstruction {
+        op: new_op.operation,
+        qubits: new_dag.qargs_interner.insert_owned(qargs),
+        clbits: new_dag.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: std::sync::OnceLock::new(),
+    };
+    let qargs = new_dag.get_qargs(inst.qubits);
+    let mapped_qargs: Vec<Qubit> = qargs
+    .iter()
+    .map(|q| Qubit::new(mapping[q.index()]))
+    .collect();
+    new_dag.apply_operation_back(
+        py,
+        inst.op.clone(),
+        &mapped_qargs,
+        &[],
+        inst.params.as_deref().cloned(),
+        inst.extra_attrs.clone(),
+        #[cfg(feature = "cache_pygates")]
+        None,
+    )?;
+    Ok(())
+}
+
 #[pyfunction]
 pub fn split_2q_unitaries(
     py: Python,
     dag: &mut DAGCircuit,
     requested_fidelity: f64,
-) -> PyResult<()> {
+) -> PyResult<Option<(DAGCircuit, Vec<usize>)>> {
     if !dag.get_op_counts().contains_key("unitary") {
-        return Ok(());
+        return Ok(None);
     }
     let nodes: Vec<NodeIndex> = dag.op_nodes(false).collect();
-
+    let mut has_swaps = false;
     for node in nodes {
         if let NodeType::Operation(inst) = &dag.dag()[node] {
             let qubits = dag.get_qargs(inst.qubits).to_vec();
@@ -51,17 +97,11 @@ pub fn split_2q_unitaries(
                 Some(requested_fidelity),
                 None,
             )?;
+            if matches!(decomp.specialization, Specialization::SWAPEquiv) {
+                has_swaps = true;
+            }
             if matches!(decomp.specialization, Specialization::IdEquiv) {
-                let k1r_arr = decomp.K1r(py);
-                let k1l_arr = decomp.K1l(py);
-                let kwargs = PyDict::new_bound(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 (k1r_gate, k1l_gate) = create_k1_gates(&decomp, py)?;
                 let insert_fn = |edge: &Wire| -> PyResult<OperationFromPython> {
                     if let Wire::Qubit(qubit) = edge {
                         if *qubit == qubits[0] {
@@ -76,15 +116,69 @@ pub fn split_2q_unitaries(
                 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(())
+    if !has_swaps {
+        return Ok(None);
+    }
+    // We have swap-like unitaries, so we create a new DAG in a manner similar to
+    // The Elide Permutations pass, while also splitting the unitaries to 1-qubit gates
+    let mut mapping: Vec<usize> = (0..dag.num_qubits()).collect();
+    let mut new_dag = dag.copy_empty_like(py, "alike")?;
+    for node in dag.topological_op_nodes()? {
+        if let NodeType::Operation(inst) = &dag.dag()[node] {
+            let qubits = dag.get_qargs(inst.qubits).to_vec();
+            let mut is_swap = false;
+            if qubits.len() == 2 && inst.op.name() == "unitary" {
+                let matrix = inst
+                    .op
+                    .matrix(inst.params_view())
+                    .expect("'unitary' gates should always have a matrix form");
+                let decomp = TwoQubitWeylDecomposition::new_inner(
+                    matrix.view(),
+                    Some(requested_fidelity),
+                    None,
+                )?;
+                if matches!(decomp.specialization, Specialization::SWAPEquiv) {
+                    // perform the virtual swap
+                    is_swap = true;
+                    let qargs = dag.get_qargs(inst.qubits);
+                    let index0 = qargs[0].index();
+                    let index1 = qargs[1].index();
+                    mapping.swap(index0, index1);                   
+                    // now add the two 1-qubit gates
+                    let (k1r_gate, k1l_gate) = create_k1_gates(&decomp, py)?;
+                    add_new_op(&mut new_dag, k1r_gate.extract()?, vec![qubits[0]], &mapping, py)?;
+                    add_new_op(&mut new_dag, k1l_gate.extract()?, vec![qubits[1]], &mapping, py)?;
+                    new_dag.add_global_phase(py, &Param::Float(decomp.global_phase+PI4))?;
+                }
+            }
+            if !is_swap{
+                // General instruction
+                let qargs = dag.get_qargs(inst.qubits);
+                let cargs = dag.get_cargs(inst.clbits);
+                let mapped_qargs: Vec<Qubit> = qargs
+                    .iter()
+                    .map(|q| Qubit::new(mapping[q.index()]))
+                    .collect();
+
+                new_dag.apply_operation_back(
+                    py,
+                    inst.op.clone(),
+                    &mapped_qargs,
+                    cargs,
+                    inst.params.as_deref().cloned(),
+                    inst.extra_attrs.clone(),
+                    #[cfg(feature = "cache_pygates")]
+                    inst.py_op.get().map(|x| x.clone_ref(py)),
+                )?;
+            }
+        }
+    }
+    return Ok(Some((new_dag, mapping)));
 }
 
 pub fn split_2q_unitaries_mod(m: &Bound<PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(split_2q_unitaries))?;
     Ok(())
-}
+}

qiskit/transpiler/passes/optimization/split_2q_unitaries.py

@@ -13,6 +13,7 @@
 """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.transpiler.layout import Layout
 from qiskit.dagcircuit.dagcircuit import DAGCircuit
 from qiskit._accelerate.split_2q_unitaries import split_2q_unitaries
 
@@ -36,5 +37,22 @@ def __init__(self, fidelity: float = 1.0 - 1e-16):
 
     def run(self, dag: DAGCircuit) -> DAGCircuit:
         """Run the Split2QUnitaries pass on `dag`."""
-        split_2q_unitaries(dag, self.requested_fidelity)
-        return dag
+        result = split_2q_unitaries(dag, self.requested_fidelity)
+        if result is None:
+            return dag
+
+        (new_dag, qubit_mapping) = result
+        input_qubit_mapping = {qubit: index for index, qubit in enumerate(dag.qubits)}
+        self.property_set["original_layout"] = Layout(input_qubit_mapping)
+        if self.property_set["original_qubit_indices"] is None:
+            self.property_set["original_qubit_indices"] = input_qubit_mapping
+
+        new_layout = Layout({dag.qubits[out]: idx for idx, out in enumerate(qubit_mapping)})
+        if current_layout := self.property_set["virtual_permutation_layout"]:
+            self.property_set["virtual_permutation_layout"] = new_layout.compose(
+                current_layout.inverse(dag.qubits, dag.qubits), dag.qubits
+            )
+        else:
+            self.property_set["virtual_permutation_layout"] = new_layout
+        return new_dag
+

releasenotes/notes/update-split-2q-unitaries-with-swap-557a1252e3208257.yaml

@@ -0,0 +1,8 @@
+---
+features_transpiler:
+  - |
+    The :class:`.Split2QUnitaries` transpiler pass has been upgraded to
+    handle the case where the unitary in consideration can be written
+    as a SWAP gate and two 1-qubit gates. In this case, it splits the
+    unitary and also applies virtual swapping similar to what is done in
+    :class:`.ElidePermutations`.

test/python/transpiler/test_split_2q_unitaries.py

@@ -266,3 +266,49 @@ def to_matrix(self):
         no_split = Split2QUnitaries()(qc)
 
         self.assertDictEqual({"mygate": 1}, no_split.count_ops())
+
+    def test_2q_swap(self):
+        """Test that a 2q unitary matching a swap gate is correctly processed."""
+        qc = QuantumCircuit(2)
+        qc.swap(0,1)
+        qc.global_phase += 1.2345
+        qc_split = QuantumCircuit(2)
+        qc_split.append(UnitaryGate(Operator(qc)), [0, 1])
+
+        pm = PassManager()
+        pm.append(Collect2qBlocks())
+        pm.append(ConsolidateBlocks())
+        pm.append(Split2QUnitaries())
+        res = pm.run(qc_split)
+        res_op = Operator.from_circuit(res, final_layout=res.layout.final_layout)
+        expected_op = Operator(qc_split)
+
+        self.assertTrue(expected_op.equiv(res_op))
+        self.assertTrue(
+            matrix_equal(expected_op.data, res_op.data, ignore_phase=False)
+        )
+
+    def test_2q_swap_with_1_qubit_gates(self):
+        """Test that a 2q unitary matching a swap gate with 1-qubit gates before and after is correctly processed."""
+        qc = QuantumCircuit(2)
+        qc.h(0)
+        qc.x(1)
+        qc.swap(0,1)
+        qc.sx(0)
+        qc.sdg(1)
+        qc.global_phase += 1.2345
+        qc_split = QuantumCircuit(2)
+        qc_split.append(UnitaryGate(Operator(qc)), [0, 1])
+
+        pm = PassManager()
+        pm.append(Collect2qBlocks())
+        pm.append(ConsolidateBlocks())
+        pm.append(Split2QUnitaries())
+        res = pm.run(qc_split)
+        res_op = Operator.from_circuit(res, final_layout=res.layout.final_layout)
+        expected_op = Operator(qc_split)
+
+        self.assertTrue(expected_op.equiv(res_op))
+        self.assertTrue(
+            matrix_equal(expected_op.data, res_op.data, ignore_phase=False)
+        )