Oxidize GateDirection pass

crates/accelerate/src/gate_direction.rs

@@ -10,19 +10,27 @@
 // copyright notice, and modified files need to carry a notice indicating
 // that they have been altered from the originals.
 
+use pyo3::prelude::*;
+use pyo3::intern;
+use pyo3::types::PyTuple;
+use hashbrown::HashSet;
+use smallvec::smallvec;
 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,
+    Qubit, TupleLikeArg,
+    operations::{OperationRef, StandardGate, Param, Operation},
+    circuit_instruction::ExtraInstructionAttributes,
 };
-use smallvec::smallvec;
+use crate::target_transpiler::exceptions::TranspilerError;
+
+//#########################################################################
+//              CheckGateDirection analysis pass functions
+//#########################################################################
+
 
 /// Check if the two-qubit gates follow the right direction with respect to the coupling map.
 ///
@@ -46,6 +54,7 @@ fn py_check_with_coupling_map(
     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:
@@ -142,8 +151,162 @@ where
     Ok(true)
 }
 
+//#########################################################################
+//              GateDirection transformation pass functions
+//#########################################################################
+
+
+// type GateDirectionCheckFn<'a> = Box<dyn Fn(&DAGCircuit, &PackedInstruction, &[Qubit]) -> bool + 'a>;
+
+// // Return a closure function that checks whether the direction of a given gate complies with the given coupling map. This is used in the
+// // pass functions below
+// fn coupling_direction_checker<'a>(py: &'a Python, dag: &'a DAGCircuit, coupling_edges: &'a Bound<PySet>) -> GateDirectionCheckFn<'a> {
+//     Box::new(move |curr_dag: &DAGCircuit, _: &PackedInstruction, op_args: &[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)
+//         })
+// }
+
+// // Return a closure function that checks whether the direction of a given gate complies with the given target. This is used in the
+// // pass functions below
+// fn target_direction_checker<'a>(py: &'a Python, dag: &'a DAGCircuit, target: PyRef<'a, Target>) -> GateDirectionCheckFn<'a> {
+//     Box::new(move |curr_dag: &DAGCircuit, inst: &PackedInstruction, op_args: &[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))
+//     })
+// }
+
+///
+///
+///
+///
+///
+#[pyfunction]
+#[pyo3(name = "fix_gate_direction_coupling")]
+fn py_fix_with_coupling_map(py: Python, dag: &DAGCircuit, coupling_edges: HashSet<[Qubit; 2]>) -> PyResult<DAGCircuit> {
+    let coupling_map_check =
+    |_: &PackedInstruction, op_args: &[Qubit]| -> bool { coupling_edges.contains(op_args) };
+
+
+    fix_gate_direction(py, dag, coupling_map_check)
+}
+
+fn fix_gate_direction<T>(py: Python, dag: &DAGCircuit, gate_complies: T)  -> PyResult<DAGCircuit>
+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");
+        };
+
+        if let OperationRef::Instruction(py_inst) = packed_inst.op.view() {
+            if py_inst.control_flow() {
+                todo!("direction fix control flow blocks");
+            }
+        }
+
+        let op_args = dag.get_qargs(packed_inst.qubits);
+        if op_args.len() != 2 {continue;}
+
+        if !gate_complies(packed_inst, op_args) {
+            if !gate_complies(packed_inst, &[op_args[1], op_args[0]]) {
+                return Err(TranspilerError::new_err(format!("The circuit requires a connection between physical qubits {:?}", op_args)));
+            }
+
+            if let OperationRef::Standard(std_gate) = packed_inst.op.view() {
+                match std_gate {
+                    StandardGate::CXGate |
+                    StandardGate::CZGate |
+                    StandardGate::ECRGate |
+                    StandardGate::SwapGate |
+                    StandardGate::RZXGate |
+                    StandardGate::RXXGate |
+                    StandardGate::RYYGate => todo!("Direction fix for {:?}", std_gate),
+                    StandardGate::RZZGate => println!("PARAMs: {:?}", packed_inst.params),
+                    _ => continue,
+                }
+            }
+        }
+    }
+
+    Ok(dag.clone()) // TODO: avoid cloning
+}
+
+//###################################################
+// Utility functions to build the replacement dags
+// TODO: replace once we have fully Rust-friendly versions of QuantumRegister, DAGCircuit and ParemeterExpression
+
+fn create_qreg<'py>(py: Python<'py>, size: u32) -> PyResult<Bound<'py, PyAny>> {
+    imports::QUANTUM_REGISTER.get_bound(py).call1((size,))
+}
+
+fn qreg_bit<'py>(py: Python, qreg: &Bound<'py, PyAny>, index: u32) -> PyResult<Bound<'py, PyAny>> {
+    qreg.call_method1(intern!(py, "__getitem__"), (index,))
+}
+
+fn std_gate(py: Python, gate: StandardGate) -> PyResult<Py<PyAny>> {
+    gate.create_py_op(py, None, &ExtraInstructionAttributes::new(None, None, None, None))
+}
+
+fn parameterized_std_gate(py: Python, gate: StandardGate, param: Param) -> PyResult<Py<PyAny>> {
+    gate.create_py_op(py, Some(&[param]), &ExtraInstructionAttributes::new(None, None, None, None))
+}
+
+fn apply_op_back(py: Python, dag: &mut DAGCircuit, op: &Py<PyAny>, qargs: &Vec<&Bound<PyAny>>) -> PyResult<()> {
+    dag.py_apply_operation_back(py,
+        op.bind(py).clone(),
+        Some( TupleLikeArg::extract_bound( &PyTuple::new_bound(py, qargs))? ),
+        None,
+        false)?;
+
+    Ok(())
+}
+
+// fn build_dag(py: Python) -> PyResult<DAGCircuit> {
+//     let qreg = create_qreg(py, 2)?;
+//     let new_dag = &mut DAGCircuit::new(py)?;
+//     new_dag.add_qreg(py, &qreg)?;
+
+//     let (q0, q1) = (qreg_bit(py, &qreg, 0)?, qreg_bit(py, &qreg, 0)?);
+
+//     apply_standard_gate_back(py, new_dag, StandardGate::HGate, &vec![&q0])?;
+//     apply_standard_gate_back(py, new_dag, StandardGate::CXGate, &vec![&q0, &q1])?;
+
+//     Ok( new_dag.clone() ) // TODO: Get rid of the clone
+// }
+
+fn cx_replacement_dag(py: Python) -> PyResult<DAGCircuit> {
+    let qreg = create_qreg(py, 2)?;
+    let new_dag = &mut DAGCircuit::new(py)?;
+    new_dag.add_qreg(py, &qreg)?;
+
+    let (q0, q1) = (qreg_bit(py, &qreg, 0)?, qreg_bit(py, &qreg, 0)?);
+    apply_op_back(py, new_dag, &std_gate(py, StandardGate::HGate)?, &vec![&q0])?;
+    apply_op_back(py, new_dag, &std_gate(py, StandardGate::HGate)?, &vec![&q1])?;
+    apply_op_back(py, new_dag, &std_gate(py, StandardGate::HGate)?, &vec![&q1, &q0])?;
+    apply_op_back(py, new_dag, &std_gate(py, StandardGate::HGate)?, &vec![&q0])?;
+    apply_op_back(py, new_dag, &std_gate(py, StandardGate::HGate)?, &vec![&q1])?;
+
+    Ok( new_dag.clone() ) // TODO: Get rid of the clone
+}
+
+
+#[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))?;
+    m.add_wrapped(wrap_pyfunction!(py_fix_with_coupling_map))?;
     Ok(())
 }

crates/accelerate/src/target_transpiler/mod.rs

@@ -43,7 +43,7 @@ use instruction_properties::InstructionProperties;
 
 use self::exceptions::TranspilerError;
 
-mod exceptions {
+pub mod exceptions {
     use pyo3::import_exception_bound;
     import_exception_bound! {qiskit.exceptions, QiskitError}
     import_exception_bound! {qiskit.transpiler.exceptions, TranspilerError}

crates/circuit/src/dag_circuit.rs

@@ -987,7 +987,7 @@ def _format(operand):
     }
 
     /// Add all wires in a quantum register.
-    fn add_qreg(&mut self, py: Python, qreg: &Bound<PyAny>) -> PyResult<()> {
+    pub fn add_qreg(&mut self, py: Python, qreg: &Bound<PyAny>) -> PyResult<()> {
         if !qreg.is_instance(imports::QUANTUM_REGISTER.get_bound(py))? {
             return Err(DAGCircuitError::new_err("not a QuantumRegister instance."));
         }
@@ -1670,7 +1670,7 @@ def _format(operand):
     /// Raises:
     ///     DAGCircuitError: if a leaf node is connected to multiple outputs
     #[pyo3(name = "apply_operation_back", signature = (op, qargs=None, cargs=None, *, check=true))]
-    fn py_apply_operation_back(
+    pub fn py_apply_operation_back(
         &mut self,
         py: Python,
         op: Bound<PyAny>,

crates/circuit/src/lib.rs

@@ -17,7 +17,7 @@ pub mod converters;
 pub mod dag_circuit;
 pub mod dag_node;
 mod dot_utils;
-mod error;
+pub mod error;
 pub mod gate_matrix;
 pub mod imports;
 mod interner;

qiskit/transpiler/passes/utils/gate_direction.py

@@ -35,6 +35,8 @@
     SwapGate,
 )
 
+from qiskit._accelerate.gate_direction import fix_gate_direction_coupling
+
 
 def _swap_node_qargs(node):
     return DAGOpNode(node.op, node.qargs[::-1], node.cargs)
@@ -345,5 +347,6 @@ def run(self, dag):
         """
         layout_map = {bit: i for i, bit in enumerate(dag.qubits)}
         if self.target is None:
-            return self._run_coupling_map(dag, layout_map)
+            return fix_gate_direction_coupling(dag, set(self.coupling_map.get_edges()))
+            # return self._run_coupling_map(dag, layout_map)
         return self._run_target(dag, layout_map)