Port ElidePermutations transpiler pass to Rust (Qiskit#13094)

* initial commit

* Rust docstring improvements

* lint

* restoring elided comment

* explicitlt setting dtype=int for permutation gates

* another attempt

* fmt after merge

* update after merge + comment from code review

* switching to apply_operation_back

* Comments from code review

* fix using params

---------

Co-authored-by: Kevin Hartman <kevin@hart.mn>

Author: 2 people

crates/accelerate/src/elide_permutations.rs

@@ -0,0 +1,109 @@
+// 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 numpy::PyReadonlyArray1;
+use pyo3::prelude::*;
+
+use qiskit_circuit::dag_circuit::{DAGCircuit, NodeType};
+use qiskit_circuit::operations::{Operation, Param};
+use qiskit_circuit::Qubit;
+
+/// Run the ElidePermutations pass on `dag`.
+/// Args:
+///     dag (DAGCircuit): the DAG to be optimized.
+/// Returns:
+///     An `Option`: the value of `None` indicates that no optimization was
+///     performed and the original `dag` should be used, otherwise it's a
+///     tuple consisting of the optimized DAG and the induced qubit permutation.
+#[pyfunction]
+fn run(py: Python, dag: &mut DAGCircuit) -> PyResult<Option<(DAGCircuit, Vec<usize>)>> {
+    let permutation_gate_names = ["swap".to_string(), "permutation".to_string()];
+    let op_counts = dag.count_ops(py, false)?;
+    if !permutation_gate_names
+        .iter()
+        .any(|name| op_counts.contains_key(name))
+    {
+        return Ok(None);
+    }
+    let mut mapping: Vec<usize> = (0..dag.num_qubits()).collect();
+
+    // note that DAGCircuit::copy_empty_like clones the interners
+    let mut new_dag = dag.copy_empty_like(py, "alike")?;
+    for node_index in dag.topological_op_nodes()? {
+        if let NodeType::Operation(inst) = &dag.dag()[node_index] {
+            match (inst.op.name(), inst.condition()) {
+                ("swap", None) => {
+                    let qargs = dag.get_qargs(inst.qubits);
+                    let index0 = qargs[0].0 as usize;
+                    let index1 = qargs[1].0 as usize;
+                    mapping.swap(index0, index1);
+                }
+                ("permutation", None) => {
+                    if let Param::Obj(ref pyobj) = inst.params.as_ref().unwrap()[0] {
+                        let pyarray: PyReadonlyArray1<i32> = pyobj.extract(py)?;
+                        let pattern = pyarray.as_array();
+
+                        let qindices: Vec<usize> = dag
+                            .get_qargs(inst.qubits)
+                            .iter()
+                            .map(|q| q.0 as usize)
+                            .collect();
+
+                        let remapped_qindices: Vec<usize> = (0..qindices.len())
+                            .map(|i| pattern[i])
+                            .map(|i| qindices[i as usize])
+                            .collect();
+
+                        qindices
+                            .iter()
+                            .zip(remapped_qindices.iter())
+                            .for_each(|(old, new)| {
+                                mapping[*old] = *new;
+                            });
+                    } else {
+                        unreachable!();
+                    }
+                }
+                _ => {
+                    // 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| q.0 as usize)
+                        .map(|q| mapping[q])
+                        .map(|q| Qubit(q.try_into().unwrap()))
+                        .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")]
+                        None,
+                    )?;
+                }
+            }
+        } else {
+            unreachable!();
+        }
+    }
+    Ok(Some((new_dag, mapping)))
+}
+
+pub fn elide_permutations(m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_wrapped(wrap_pyfunction!(run))?;
+    Ok(())
+}

crates/accelerate/src/lib.rs

@@ -24,6 +24,7 @@ pub mod commutation_checker;
 pub mod convert_2q_block_matrix;
 pub mod dense_layout;
 pub mod edge_collections;
+pub mod elide_permutations;
 pub mod equivalence;
 pub mod error_map;
 pub mod euler_one_qubit_decomposer;

crates/pyext/src/lib.rs

@@ -39,6 +39,7 @@ fn _accelerate(m: &Bound<PyModule>) -> PyResult<()> {
     add_submodule(m, ::qiskit_accelerate::dense_layout::dense_layout, "dense_layout")?;
     add_submodule(m, ::qiskit_accelerate::equivalence::equivalence, "equivalence")?;
     add_submodule(m, ::qiskit_accelerate::error_map::error_map, "error_map")?;
+    add_submodule(m, ::qiskit_accelerate::elide_permutations::elide_permutations, "elide_permutations")?;
     add_submodule(m, ::qiskit_accelerate::euler_one_qubit_decomposer::euler_one_qubit_decomposer, "euler_one_qubit_decomposer")?;
     add_submodule(m, ::qiskit_accelerate::filter_op_nodes::filter_op_nodes_mod, "filter_op_nodes")?;
     add_submodule(m, ::qiskit_accelerate::gate_direction::gate_direction, "gate_direction")?;

qiskit/__init__.py

@@ -80,6 +80,7 @@
 sys.modules["qiskit._accelerate.sparse_pauli_op"] = _accelerate.sparse_pauli_op
 sys.modules["qiskit._accelerate.star_prerouting"] = _accelerate.star_prerouting
 sys.modules["qiskit._accelerate.stochastic_swap"] = _accelerate.stochastic_swap
+sys.modules["qiskit._accelerate.elide_permutations"] = _accelerate.elide_permutations
 sys.modules["qiskit._accelerate.target"] = _accelerate.target
 sys.modules["qiskit._accelerate.two_qubit_decompose"] = _accelerate.two_qubit_decompose
 sys.modules["qiskit._accelerate.unitary_synthesis"] = _accelerate.unitary_synthesis

qiskit/circuit/library/generalized_gates/permutation.py

@@ -141,7 +141,7 @@ def __init__(
             raise CircuitError(
                 "Permutation pattern must be some ordering of 0..num_qubits-1 in a list."
             )
-        pattern = np.array(pattern)
+        pattern = np.array(pattern, dtype=np.int32)
 
         super().__init__(name="permutation", num_qubits=num_qubits, params=[pattern])
 

qiskit/transpiler/passes/optimization/elide_permutations.py

@@ -15,10 +15,9 @@
 
 import logging
 
-from qiskit.circuit.library.standard_gates import SwapGate
-from qiskit.circuit.library.generalized_gates import PermutationGate
 from qiskit.transpiler.basepasses import TransformationPass
 from qiskit.transpiler.layout import Layout
+from qiskit._accelerate import elide_permutations as elide_permutations_rs
 
 logger = logging.getLogger(__name__)
 
@@ -67,38 +66,16 @@ def run(self, dag):
             )
             return dag
 
-        op_count = dag.count_ops(recurse=False)
-        if op_count.get("swap", 0) == 0 and op_count.get("permutation", 0) == 0:
+        result = elide_permutations_rs.run(dag)
+
+        # If the pass did not do anything, the result is None
+        if result is None:
             return dag
 
-        new_dag = dag.copy_empty_like()
-        qubit_mapping = list(range(len(dag.qubits)))
-
-        def _apply_mapping(qargs):
-            return tuple(dag.qubits[qubit_mapping[dag.find_bit(qubit).index]] for qubit in qargs)
-
-        for node in dag.topological_op_nodes():
-            if not isinstance(node.op, (SwapGate, PermutationGate)):
-                new_dag.apply_operation_back(
-                    node.op, _apply_mapping(node.qargs), node.cargs, check=False
-                )
-            elif getattr(node.op, "condition", None) is not None:
-                new_dag.apply_operation_back(
-                    node.op, _apply_mapping(node.qargs), node.cargs, check=False
-                )
-            elif isinstance(node.op, SwapGate):
-                index_0 = dag.find_bit(node.qargs[0]).index
-                index_1 = dag.find_bit(node.qargs[1]).index
-                qubit_mapping[index_1], qubit_mapping[index_0] = (
-                    qubit_mapping[index_0],
-                    qubit_mapping[index_1],
-                )
-            elif isinstance(node.op, PermutationGate):
-                starting_indices = [qubit_mapping[dag.find_bit(qarg).index] for qarg in node.qargs]
-                pattern = node.op.params[0]
-                updated_indices = [starting_indices[idx] for idx in pattern]
-                for i, j in zip(starting_indices, updated_indices):
-                    qubit_mapping[i] = j
+        # Otherwise, the result is a pair consisting of the rewritten DAGCircuit and the
+        # qubit mapping.
+        (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:

releasenotes/notes/port-elide-permutations-ed91c3d9cef2fec6.yaml

@@ -0,0 +1,5 @@
+---
+features_transpiler:
+  - |
+    Port most of the logic of the transpiler pass :class:`~.ElidePermutations`
+    to Rust.