Oxidize CheckGateDirection (#13042)

* First working implementation

* A control flow fix, lint, doc and some more polishing

* Addressing review comments and further simplifying the code

* Address more review comments

* Use explicit temp storage for mapping to avoid E0716 in Rust 1.70

* Allow checking 2Q non-control flow instructions

crates/accelerate/src/gate_direction.rs

@@ -0,0 +1,150 @@
+// 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 crate::nlayout::PhysicalQubit;
+use crate::target_transpiler::Target;
+use hashbrown::HashSet;
+use pyo3::prelude::*;
+use qiskit_circuit::imports;
+use qiskit_circuit::operations::OperationRef;
+use qiskit_circuit::{
+    dag_circuit::{DAGCircuit, NodeType},
+    operations::Operation,
+    packed_instruction::PackedInstruction,
+    Qubit,
+};
+use smallvec::smallvec;
+
+/// Check if the two-qubit gates follow the right direction with respect to the coupling map.
+///
+/// Args:
+///     dag: the DAGCircuit to analyze
+///
+///     coupling_edges: set of edge pairs representing a directed coupling map, against which gate directionality is checked
+///
+/// Returns:
+///     true iff all two-qubit gates comply with the coupling constraints
+#[pyfunction]
+#[pyo3(name = "check_gate_direction_coupling")]
+fn py_check_with_coupling_map(
+    py: Python,
+    dag: &DAGCircuit,
+    coupling_edges: HashSet<[Qubit; 2]>,
+) -> PyResult<bool> {
+    let coupling_map_check =
+        |_: &PackedInstruction, op_args: &[Qubit]| -> bool { coupling_edges.contains(op_args) };
+
+    check_gate_direction(py, dag, &coupling_map_check, None)
+}
+
+/// Check if the two-qubit gates follow the right direction with respect to instructions supported in the given target.
+///
+/// Args:
+///     dag: the DAGCircuit to analyze
+///
+///     target: the Target against which gate directionality compliance is checked
+///
+/// Returns:
+///     true iff all two-qubit gates comply with the target's coupling constraints
+#[pyfunction]
+#[pyo3(name = "check_gate_direction_target")]
+fn py_check_with_target(py: Python, dag: &DAGCircuit, target: &Target) -> PyResult<bool> {
+    let target_check = |inst: &PackedInstruction, op_args: &[Qubit]| -> bool {
+        let qargs = smallvec![
+            PhysicalQubit::new(op_args[0].0),
+            PhysicalQubit::new(op_args[1].0)
+        ];
+
+        target.instruction_supported(inst.op.name(), Some(&qargs))
+    };
+
+    check_gate_direction(py, dag, &target_check, None)
+}
+
+// The main routine for checking gate directionality.
+//
+// gate_complies: a function returning true iff the two-qubit gate direction complies with directionality constraints
+//
+// qubit_mapping: used for mapping the index of a given qubit within an instruction qargs vector to the corresponding qubit index of the
+//  original DAGCircuit the pass was called with. This mapping is required since control flow blocks are represented by nested DAGCircuit
+//  objects whose instruction qubit indices are relative to the parent DAGCircuit they reside in, thus when we recurse into nested DAGs, we need
+//  to carry the mapping context relative to the original DAG.
+//  When qubit_mapping is None, the identity mapping is assumed
+fn check_gate_direction<T>(
+    py: Python,
+    dag: &DAGCircuit,
+    gate_complies: &T,
+    qubit_mapping: Option<&[Qubit]>,
+) -> PyResult<bool>
+where
+    T: Fn(&PackedInstruction, &[Qubit]) -> bool,
+{
+    for node in dag.op_nodes(false) {
+        let NodeType::Operation(packed_inst) = &dag.dag[node] else {
+            panic!("PackedInstruction is expected");
+        };
+
+        let inst_qargs = dag.get_qargs(packed_inst.qubits);
+
+        if let OperationRef::Instruction(py_inst) = packed_inst.op.view() {
+            if py_inst.control_flow() {
+                let circuit_to_dag = imports::CIRCUIT_TO_DAG.get_bound(py); // TODO: Take out of the recursion
+                let py_inst = py_inst.instruction.bind(py);
+
+                for block in py_inst.getattr("blocks")?.iter()? {
+                    let inner_dag: DAGCircuit = circuit_to_dag.call1((block?,))?.extract()?;
+
+                    let block_ok = if let Some(mapping) = qubit_mapping {
+                        let mapping = inst_qargs // Create a temp mapping for the recursive call
+                            .iter()
+                            .map(|q| mapping[q.0 as usize])
+                            .collect::<Vec<Qubit>>();
+
+                        check_gate_direction(py, &inner_dag, gate_complies, Some(&mapping))?
+                    } else {
+                        check_gate_direction(py, &inner_dag, gate_complies, Some(inst_qargs))?
+                    };
+
+                    if !block_ok {
+                        return Ok(false);
+                    }
+                }
+                continue;
+            }
+        }
+
+        if inst_qargs.len() == 2
+            && !match qubit_mapping {
+                // Check gate direction based either on a given custom mapping or the identity mapping
+                Some(mapping) => gate_complies(
+                    packed_inst,
+                    &[
+                        mapping[inst_qargs[0].0 as usize],
+                        mapping[inst_qargs[1].0 as usize],
+                    ],
+                ),
+                None => gate_complies(packed_inst, inst_qargs),
+            }
+        {
+            return Ok(false);
+        }
+    }
+
+    Ok(true)
+}
+
+#[pymodule]
+pub fn gate_direction(m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_wrapped(wrap_pyfunction!(py_check_with_coupling_map))?;
+    m.add_wrapped(wrap_pyfunction!(py_check_with_target))?;
+    Ok(())
+}

crates/accelerate/src/lib.rs

@@ -24,6 +24,7 @@ pub mod edge_collections;
 pub mod error_map;
 pub mod euler_one_qubit_decomposer;
 pub mod filter_op_nodes;
+pub mod gate_direction;
 pub mod inverse_cancellation;
 pub mod isometry;
 pub mod nlayout;

crates/accelerate/src/target_transpiler/mod.rs

@@ -50,7 +50,7 @@ mod exceptions {
 }
 
 // Custom types
-type Qargs = SmallVec<[PhysicalQubit; 2]>;
+pub type Qargs = SmallVec<[PhysicalQubit; 2]>;
 type GateMap = IndexMap<String, PropsMap, RandomState>;
 type PropsMap = NullableIndexMap<Qargs, Option<InstructionProperties>>;
 type GateMapState = Vec<(String, Vec<(Option<Qargs>, Option<InstructionProperties>)>)>;

crates/circuit/src/dag_circuit.rs

@@ -3975,19 +3975,12 @@ def _format(operand):
     }
 
     /// Get list of 2 qubit operations. Ignore directives like snapshot and barrier.
-    fn two_qubit_ops(&self, py: Python) -> PyResult<Vec<Py<PyAny>>> {
+    #[pyo3(name = "two_qubit_ops")]
+    pub fn py_two_qubit_ops(&self, py: Python) -> PyResult<Vec<Py<PyAny>>> {
         let mut nodes = Vec::new();
-        for (node, weight) in self.dag.node_references() {
-            if let NodeType::Operation(ref packed) = weight {
-                if packed.op.directive() {
-                    continue;
-                }
-
-                let qargs = self.qargs_interner.get(packed.qubits);
-                if qargs.len() == 2 {
-                    nodes.push(self.unpack_into(py, node, weight)?);
-                }
-            }
+        for node in self.two_qubit_ops() {
+            let weight = self.dag.node_weight(node).expect("NodeIndex in graph");
+            nodes.push(self.unpack_into(py, node, weight)?);
         }
         Ok(nodes)
     }
@@ -5756,6 +5749,19 @@ impl DAGCircuit {
         }
     }
 
+    /// Return an iterator of 2 qubit operations. Ignore directives like snapshot and barrier.
+    pub fn two_qubit_ops(&self) -> impl Iterator<Item = NodeIndex> + '_ {
+        Box::new(self.op_nodes(false).filter(|index| {
+            let weight = self.dag.node_weight(*index).expect("NodeIndex in graph");
+            if let NodeType::Operation(ref packed) = weight {
+                let qargs = self.qargs_interner.get(packed.qubits);
+                qargs.len() == 2
+            } else {
+                false
+            }
+        }))
+    }
+
     // Filter any nodes that don't match a given predicate function
     pub fn filter_op_nodes<F>(&mut self, mut predicate: F)
     where

crates/circuit/src/lib.rs

@@ -32,7 +32,7 @@ use pyo3::prelude::*;
 use pyo3::types::{PySequence, PyTuple};
 
 pub type BitType = u32;
-#[derive(Copy, Clone, Debug, Hash, Ord, PartialOrd, Eq, PartialEq)]
+#[derive(Copy, Clone, Debug, Hash, Ord, PartialOrd, Eq, PartialEq, FromPyObject)]
 pub struct Qubit(pub BitType);
 #[derive(Copy, Clone, Debug, Hash, Ord, PartialOrd, Eq, PartialEq)]
 pub struct Clbit(pub BitType);

crates/pyext/src/lib.rs

@@ -17,9 +17,9 @@ use qiskit_accelerate::{
     commutation_analysis::commutation_analysis, commutation_checker::commutation_checker,
     convert_2q_block_matrix::convert_2q_block_matrix, dense_layout::dense_layout,
     error_map::error_map, euler_one_qubit_decomposer::euler_one_qubit_decomposer,
-    filter_op_nodes::filter_op_nodes_mod, inverse_cancellation::inverse_cancellation_mod,
-    isometry::isometry, nlayout::nlayout, optimize_1q_gates::optimize_1q_gates,
-    pauli_exp_val::pauli_expval,
+    filter_op_nodes::filter_op_nodes_mod, gate_direction::gate_direction,
+    inverse_cancellation::inverse_cancellation_mod, isometry::isometry, nlayout::nlayout,
+    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,
@@ -72,6 +72,7 @@ fn _accelerate(m: &Bound<PyModule>) -> PyResult<()> {
     add_submodule(m, uc_gate, "uc_gate")?;
     add_submodule(m, utils, "utils")?;
     add_submodule(m, vf2_layout, "vf2_layout")?;
+    add_submodule(m, gate_direction, "gate_direction")?;
     add_submodule(m, commutation_checker, "commutation_checker")?;
     add_submodule(m, commutation_analysis, "commutation_analysis")?;
     Ok(())

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.gate_direction"] = _accelerate.gate_direction
 sys.modules["qiskit._accelerate.inverse_cancellation"] = _accelerate.inverse_cancellation
 sys.modules["qiskit._accelerate.check_map"] = _accelerate.check_map
 sys.modules["qiskit._accelerate.filter_op_nodes"] = _accelerate.filter_op_nodes

qiskit/transpiler/passes/utils/check_gate_direction.py

@@ -12,9 +12,11 @@
 
 """Check if the gates follow the right direction with respect to the coupling map."""
 
-from qiskit.circuit.controlflow import CONTROL_FLOW_OP_NAMES
-from qiskit.converters import circuit_to_dag
 from qiskit.transpiler.basepasses import AnalysisPass
+from qiskit._accelerate.gate_direction import (
+    check_gate_direction_coupling,
+    check_gate_direction_target,
+)
 
 
 class CheckGateDirection(AnalysisPass):
@@ -34,42 +36,6 @@ def __init__(self, coupling_map, target=None):
         self.coupling_map = coupling_map
         self.target = target
 
-    def _coupling_map_visit(self, dag, wire_map, edges=None):
-        if edges is None:
-            edges = self.coupling_map.get_edges()
-        # Don't include directives to avoid things like barrier, which are assumed always supported.
-        for node in dag.op_nodes(include_directives=False):
-            if node.name in CONTROL_FLOW_OP_NAMES:
-                for block in node.op.blocks:
-                    inner_wire_map = {
-                        inner: wire_map[outer] for outer, inner in zip(node.qargs, block.qubits)
-                    }
-
-                    if not self._coupling_map_visit(circuit_to_dag(block), inner_wire_map, edges):
-                        return False
-            elif (
-                len(node.qargs) == 2
-                and (wire_map[node.qargs[0]], wire_map[node.qargs[1]]) not in edges
-            ):
-                return False
-        return True
-
-    def _target_visit(self, dag, wire_map):
-        # Don't include directives to avoid things like barrier, which are assumed always supported.
-        for node in dag.op_nodes(include_directives=False):
-            if node.name in CONTROL_FLOW_OP_NAMES:
-                for block in node.op.blocks:
-                    inner_wire_map = {
-                        inner: wire_map[outer] for outer, inner in zip(node.qargs, block.qubits)
-                    }
-                    if not self._target_visit(circuit_to_dag(block), inner_wire_map):
-                        return False
-            elif len(node.qargs) == 2 and not self.target.instruction_supported(
-                node.name, (wire_map[node.qargs[0]], wire_map[node.qargs[1]])
-            ):
-                return False
-        return True
-
     def run(self, dag):
         """Run the CheckGateDirection pass on `dag`.
 
@@ -79,9 +45,8 @@ def run(self, dag):
         Args:
             dag (DAGCircuit): DAG to check.
         """
-        wire_map = {bit: i for i, bit in enumerate(dag.qubits)}
         self.property_set["is_direction_mapped"] = (
-            self._coupling_map_visit(dag, wire_map)
-            if self.target is None
-            else self._target_visit(dag, wire_map)
+            check_gate_direction_target(dag, self.target)
+            if self.target
+            else check_gate_direction_coupling(dag, set(self.coupling_map.get_edges()))
         )

test/python/transpiler/test_check_gate_direction.py

@@ -72,13 +72,13 @@ def test_true_direction(self):
 
     def test_true_direction_in_same_layer(self):
         """Two CXs distance_qubits 1 to each other, in the same layer
-        qr0:--(+)--
+        qr0:---.--
                |
-        qr1:---.---
+        qr1:--(+)---
 
-        qr2:--(+)--
+        qr2:---.---
                |
-        qr3:---.---
+        qr3:--(+)--
 
         CouplingMap map: [0]->[1]->[2]->[3]
         """
@@ -96,9 +96,9 @@ def test_true_direction_in_same_layer(self):
 
     def test_wrongly_mapped(self):
         """Needs [0]-[1] in a [0]--[2]--[1]
-        qr0:--(+)--
+        qr0:---.---
                |
-        qr1:---.---
+        qr1:--(+)--
 
         CouplingMap map: [0]->[2]->[1]
         """
@@ -115,11 +115,11 @@ def test_wrongly_mapped(self):
 
     def test_true_direction_undirected(self):
         """Mapped but with wrong direction
-        qr0:--(+)-[H]--.--
+        qr0:---.--[H]-(+)-
                |       |
-        qr1:---.-------|--
+        qr1:--(+)------|--
                        |
-        qr2:----------(+)-
+        qr2:-----------.--
 
         CouplingMap map: [1]<-[0]->[2]
         """
@@ -138,13 +138,13 @@ def test_true_direction_undirected(self):
 
     def test_false_direction_in_same_layer_undirected(self):
         """Two CXs in the same layer, but one is wrongly directed
-        qr0:--(+)--
+        qr0:---.---
                |
-        qr1:---.---
+        qr1:--(+)--
 
-        qr2:---.---
+        qr2:--(+)--
                |
-        qr3:--(+)--
+        qr3:---.---
 
         CouplingMap map: [0]->[1]->[2]->[3]
         """