First working implementation

crates/accelerate/src/gate_direction.rs

@@ -0,0 +1,126 @@
+use pyo3::prelude::*;
+use pyo3::types::{PySet, PyTuple};
+use crate::nlayout::PhysicalQubit;
+use crate::target_transpiler::{Qargs, Target};
+use qiskit_circuit::imports;
+use qiskit_circuit::operations::{OperationRef, PyInstruction};
+use qiskit_circuit::{
+    dag_circuit::{DAGCircuit, NodeType},
+    operations::Operation,
+    packed_instruction::PackedInstruction,
+    Qubit,
+};
+
+// Handle control flow instruction, namely recurse into its circuit blocks in the analysis
+fn check_gate_direction_control_flow<T>(
+    py: Python,
+    py_inst: &PyInstruction,
+    gate_complies: &T,
+) -> PyResult<bool>
+where
+    T: Fn(&DAGCircuit, &PackedInstruction, &Vec<Qubit>) -> bool,
+{
+    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);
+
+    let raw_blocks = py_inst.getattr("blocks")?;
+    let blocks = raw_blocks.downcast::<PyTuple>()?;
+
+    for block in blocks.iter() {
+        let inner_dag: DAGCircuit = circuit_to_dag.call1((block,))?.extract()?;
+
+        if !check_gate_direction(py, &inner_dag, gate_complies)? {
+            return Ok(false);
+        }
+    }
+
+    Ok(true)
+}
+
+// The main analysis pass routine.
+fn check_gate_direction<T>(py: Python, dag: &DAGCircuit, gate_complies: &T) -> PyResult<bool>
+where
+    T: Fn(&DAGCircuit, &PackedInstruction, &Vec<Qubit>) -> bool,
+{
+    for node in dag.op_nodes(false) {
+        let Some(NodeType::Operation(packed_inst)) = dag.dag.node_weight(node) else {
+            return Ok(false);
+        }; // TODO: proper error
+
+        if let OperationRef::Instruction(py_inst) = packed_inst.op.view() {
+            if py_inst.control_flow()
+                && !check_gate_direction_control_flow(py, py_inst, gate_complies)?
+            {
+                return Ok(false);
+            }
+        } else {
+            let op_args = dag.get_qubits(packed_inst.qubits);
+            if op_args.len() == 2 && !gate_complies(dag, packed_inst, op_args) {
+                return Ok(false);
+            }
+        }
+    }
+
+    Ok(true)
+}
+
+// Map a qubit interned in curr_dag to its corresponding qubit entry interned in orig_dag.
+// Required for checking control flow instruction which are represented in blocks (circuits)
+// and converted to DAGCircuit with possibly different qargs than the original one.
+fn map_qubit(py: Python, orig_dag: &DAGCircuit, curr_dag: &DAGCircuit, qubit: Qubit) -> Qubit {
+    let qubit = curr_dag
+        .qubits
+        .get(qubit)
+        .expect("Qubit in curr_dag")
+        .bind(py);
+    orig_dag.qubits.find(qubit).expect("Qubit in orig_dag")
+}
+
+#[pyfunction]
+#[pyo3(name = "check_gate_direction_coupling")]
+fn py_check_with_coupling_map(
+    py: Python,
+    dag: &DAGCircuit,
+    coupling_edges: &Bound<PySet>,
+) -> PyResult<bool> {
+    let coupling_map_check =
+        |curr_dag: &DAGCircuit, _: &PackedInstruction, op_args: &Vec<Qubit>| -> bool {
+            coupling_edges
+                .contains((
+                    map_qubit(py, dag, curr_dag, op_args[0]).0,
+                    map_qubit(py, dag, curr_dag, op_args[1]).0,
+                ))
+                .unwrap_or(false)
+        };
+
+    check_gate_direction(py, dag, &coupling_map_check)
+}
+
+#[pyfunction]
+#[pyo3(name = "check_gate_direction_target")]
+fn py_check_with_target(py: Python, dag: &DAGCircuit, target: &Bound<Target>) -> PyResult<bool> {
+    let target = target.borrow();
+
+    let target_check =
+        |curr_dag: &DAGCircuit, inst: &PackedInstruction, op_args: &Vec<Qubit>| -> bool {
+            let mut qargs = Qargs::new();
+
+            qargs.push(PhysicalQubit::new(
+                map_qubit(py, dag, curr_dag, op_args[0]).0,
+            ));
+            qargs.push(PhysicalQubit::new(
+                map_qubit(py, dag, curr_dag, op_args[1]).0,
+            ));
+
+            target.instruction_supported(inst.op.name(), Some(&qargs))
+        };
+
+    check_gate_direction(py, dag, &target_check)
+}
+
+#[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

@@ -36,6 +36,7 @@ pub mod two_qubit_decompose;
 pub mod uc_gate;
 pub mod utils;
 pub mod vf2_layout;
+pub mod gate_direction;
 
 mod rayon_ext;
 #[cfg(test)]

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

@@ -22,7 +22,7 @@ use crate::dag_node::{DAGInNode, DAGNode, DAGOpNode, DAGOutNode};
 use crate::dot_utils::build_dot;
 use crate::error::DAGCircuitError;
 use crate::imports;
-use crate::interner::{IndexedInterner, Interner};
+use crate::interner::{Index, IndexedInterner, Interner};
 use crate::operations::{Operation, OperationRef, Param, PyInstruction};
 use crate::packed_instruction::PackedInstruction;
 use crate::rustworkx_core_vnext::isomorphism;
@@ -3958,19 +3958,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_cache.intern(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)
     }
@@ -5798,6 +5791,19 @@ impl DAGCircuit {
         }
     }
 
+    /// Return an iterator of 2 qubit operations. Ignore directives like snapshot and barrier.
+    pub fn two_qubit_ops<'a>(&'a self) -> Box<dyn Iterator<Item = NodeIndex> + 'a> {
+        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_cache.intern(packed.qubits);
+                return qargs.len() == 2;
+            } else {
+                false
+            }
+        }))
+    }
+
     pub fn op_nodes_by_py_type<'a>(
         &'a self,
         op: &'a Bound<PyType>,
@@ -6167,6 +6173,11 @@ impl DAGCircuit {
         }
         Ok(())
     }
+
+    /// Get the qubits interned in the given index
+    pub fn get_qubits(&self, index: Index) -> &Vec<Qubit> {
+        self.qargs_cache.intern(index)
+    }
 }
 
 /// Add to global phase. Global phase can only be Float or ParameterExpression so this

crates/pyext/src/lib.rs

@@ -20,7 +20,7 @@ use qiskit_accelerate::{
     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,
+    utils::utils, vf2_layout::vf2_layout, gate_direction::gate_direction
 };
 
 #[inline(always)]
@@ -60,5 +60,6 @@ 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")?;
     Ok(())
 }

qiskit/__init__.py

@@ -87,6 +87,7 @@
 sys.modules["qiskit._accelerate.synthesis.linear"] = _accelerate.synthesis.linear
 sys.modules["qiskit._accelerate.synthesis.clifford"] = _accelerate.synthesis.clifford
 sys.modules["qiskit._accelerate.synthesis.linear_phase"] = _accelerate.synthesis.linear_phase
+sys.modules["qiskit._accelerate.gate_direction"] = _accelerate.gate_direction
 
 from qiskit.exceptions import QiskitError, MissingOptionalLibraryError
 

qiskit/transpiler/passes/utils/check_gate_direction.py

@@ -12,9 +12,8 @@
 
 """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 +33,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 +42,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)
-        )
+        if self.target:
+            self.property_set["is_direction_mapped"] = check_gate_direction_target(dag, self.target)
+        else:
+            coupling_edges = {x for x in self.coupling_map.get_edges()}
+            self.property_set["is_direction_mapped"] = check_gate_direction_coupling(dag, coupling_edges)