Port GateDirection to Rust

eliarbel · eliarbel · commit 30bf21bc2782 · 2024-09-02T14:15:36.000+03:00
Move GateDirection._run_coupling_map and GateDirection._run_target to
Rust. Also moving the code to create replacement DAGs to Rust. Both
GateDirection and CheckGateDirection will reside in gate_direction.rs.

Depends on the check-gate-direction branch.
diff --git a/crates/accelerate/src/gate_direction.rs b/crates/accelerate/src/gate_direction.rs
@@ -14,15 +14,69 @@ use crate::nlayout::PhysicalQubit;
 use crate::target_transpiler::{Qargs, Target};
 use pyo3::prelude::*;
 use pyo3::types::{PySet, PyTuple};
-use qiskit_circuit::imports;
-use qiskit_circuit::operations::{OperationRef, PyInstruction};
+use pyo3::intern;
+use qiskit_circuit::{imports, TupleLikeArg};
+use qiskit_circuit::operations::{OperationRef, PyInstruction, StandardGate};
 use qiskit_circuit::{
     dag_circuit::{DAGCircuit, NodeType},
     error::DAGCircuitError,
     operations::Operation,
     packed_instruction::PackedInstruction,
-    Qubit,
+    Qubit, operations::Param
 };
+use crate::target_transpiler::exceptions::TranspilerError;
+
+
+
+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))
+    })
+}
+
+// Map a qubit interned in curr_dag to its corresponding qubit entry interned in orig_dag.
+// Required for checking control flow instructions 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")
+}
+
+
+//#########################################################################
+//              CheckGateDirection analysis pass functions
+//#########################################################################
+
 
 /// Check if the two-qubit gates follow the right direction with respect to the coupling map.
 ///
@@ -40,19 +94,10 @@ fn py_check_with_coupling_map(
     dag: &DAGCircuit,
     coupling_edges: &Bound<PySet>,
 ) -> PyResult<bool> {
-    let coupling_map_check =
-        |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)
-        };
-
-    check_gate_direction(py, dag, &coupling_map_check)
+    check_gate_direction(py, dag, &coupling_direction_checker(&py, dag, coupling_edges))
 }
 
+
 /// Check if the two-qubit gates follow the right direction with respect to instructions supported in the given target.
 ///
 /// Args:
@@ -67,31 +112,42 @@ fn py_check_with_coupling_map(
 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: &[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,
-            ));
+    check_gate_direction(py, dag, &target_direction_checker(&py, dag, target))
+}
 
-            target.instruction_supported(inst.op.name(), Some(&qargs))
+// The main routine for checking gate directionality
+fn check_gate_direction(py: Python, dag: &DAGCircuit, gate_complies: &GateDirectionCheckFn) -> PyResult<bool>
+{
+    for node in dag.op_nodes(false) {
+        let NodeType::Operation(packed_inst) = &dag.dag[node] else {
+            return Err(DAGCircuitError::new_err("PackedInstruction is expected"));
         };
 
-    check_gate_direction(py, dag, &target_check)
+        if let OperationRef::Instruction(py_inst) = packed_inst.op.view() {
+            if py_inst.control_flow() {
+                if !check_gate_direction_control_flow(py, py_inst, gate_complies)? {
+                    return Ok(false);
+                } else {
+                    continue;
+                }
+            }
+        }
+
+        let op_args = dag.get_inst_qubits(packed_inst.qubits);
+        if op_args.len() == 2 && !gate_complies(dag, packed_inst, op_args) {
+            return Ok(false);
+        }
+    }
+
+    Ok(true)
 }
 
 // Handle a control flow instruction, namely check recursively into its circuit blocks
-fn check_gate_direction_control_flow<T>(
+fn check_gate_direction_control_flow(
     py: Python,
     py_inst: &PyInstruction,
-    gate_complies: &T,
+    gate_complies: &GateDirectionCheckFn,
 ) -> PyResult<bool>
-where
-    T: Fn(&DAGCircuit, &PackedInstruction, &[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);
@@ -110,50 +166,125 @@ where
     Ok(true)
 }
 
-// The main routine for checking gate directionality
-fn check_gate_direction<T>(py: Python, dag: &DAGCircuit, gate_complies: &T) -> PyResult<bool>
-where
-    T: Fn(&DAGCircuit, &PackedInstruction, &[Qubit]) -> bool,
-{
+//#########################################################################
+//              GateDirection transformation pass functions
+//#########################################################################
+
+///
+///
+///
+///
+///
+#[pyfunction]
+#[pyo3(name = "fix_gate_direction_coupling")]
+fn py_fix_with_coupling_map(py: Python, dag: &DAGCircuit, coupling_edges: &Bound<PySet>) -> PyResult<DAGCircuit> {
+    fix_gate_direction(py, dag, &coupling_direction_checker(&py, dag, coupling_edges))
+}
+
+
+fn fix_gate_direction(py: Python, dag: &DAGCircuit, gate_complies: &GateDirectionCheckFn) -> PyResult<DAGCircuit> {
     for node in dag.op_nodes(false) {
         let NodeType::Operation(packed_inst) = &dag.dag[node] else {
             return Err(DAGCircuitError::new_err("PackedInstruction is expected"));
         };
 
         if let OperationRef::Instruction(py_inst) = packed_inst.op.view() {
             if py_inst.control_flow() {
-                if !check_gate_direction_control_flow(py, py_inst, gate_complies)? {
-                    return Ok(false);
-                } else {
-                    continue;
-                }
+                todo!("direction fix control flow blocks");
             }
         }
 
         let op_args = dag.get_inst_qubits(packed_inst.qubits);
-        if op_args.len() == 2 && !gate_complies(dag, packed_inst, op_args) {
-            return Ok(false);
+        if op_args.len() != 2 {continue;}
+
+        if !gate_complies(dag, packed_inst, op_args) {
+            if !gate_complies(dag, 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(true)
+    Ok(dag.clone()) // TODO: avoid cloning
 }
 
-// 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")
+
+//###################################################
+// 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, None)
+}
+
+fn parameterized_std_gate(py: Python, gate: StandardGate, param: Param) -> PyResult<Py<PyAny>> {
+    gate.create_py_op(py, Some(&[param]), 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(())
 }
diff --git a/crates/accelerate/src/target_transpiler/mod.rs b/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}
diff --git a/crates/circuit/src/dag_circuit.rs b/crates/circuit/src/dag_circuit.rs
@@ -985,7 +985,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."));
         }
@@ -1661,7 +1661,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>,
diff --git a/qiskit/transpiler/passes/utils/gate_direction.py b/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)
