[Oxidize BasisTranslator] Add rust-native compose_transforms()

crates/accelerate/src/basis/basis_translator/compose_transforms.rs

@@ -0,0 +1,192 @@
+// 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 hashbrown::{HashMap, HashSet};
+use pyo3::prelude::*;
+use qiskit_circuit::circuit_instruction::OperationFromPython;
+use qiskit_circuit::imports::{GATE, PARAMETER_VECTOR, QUANTUM_REGISTER};
+use qiskit_circuit::parameter_table::ParameterUuid;
+use qiskit_circuit::Qubit;
+use qiskit_circuit::{
+    circuit_data::CircuitData,
+    dag_circuit::{DAGCircuit, NodeType},
+    operations::{Operation, Param},
+};
+use smallvec::SmallVec;
+
+use crate::equivalence::CircuitFromPython;
+
+// Custom types
+pub type GateIdentifier = (String, u32);
+pub type BasisTransformIn = (SmallVec<[Param; 3]>, CircuitFromPython);
+pub type BasisTransformOut = (SmallVec<[Param; 3]>, DAGCircuit);
+
+#[pyfunction(name = "compose_transforms")]
+pub(super) fn py_compose_transforms(
+    py: Python,
+    basis_transforms: Vec<(GateIdentifier, BasisTransformIn)>,
+    source_basis: HashSet<GateIdentifier>,
+    source_dag: &DAGCircuit,
+) -> PyResult<HashMap<GateIdentifier, BasisTransformOut>> {
+    compose_transforms(py, &basis_transforms, &source_basis, source_dag)
+}
+
+pub(super) fn compose_transforms<'a>(
+    py: Python,
+    basis_transforms: &'a [(GateIdentifier, BasisTransformIn)],
+    source_basis: &'a HashSet<GateIdentifier>,
+    source_dag: &'a DAGCircuit,
+) -> PyResult<HashMap<GateIdentifier, BasisTransformOut>> {
+    let mut gate_param_counts: HashMap<GateIdentifier, usize> = HashMap::default();
+    get_gates_num_params(source_dag, &mut gate_param_counts)?;
+    let mut mapped_instructions: HashMap<GateIdentifier, BasisTransformOut> = HashMap::new();
+
+    for (gate_name, gate_num_qubits) in source_basis.iter().cloned() {
+        let num_params = gate_param_counts[&(gate_name.clone(), gate_num_qubits)];
+
+        let placeholder_params: SmallVec<[Param; 3]> = PARAMETER_VECTOR
+            .get_bound(py)
+            .call1((&gate_name, num_params))?
+            .extract()?;
+
+        let mut dag = DAGCircuit::new(py)?;
+        // Create the mock gate and add to the circuit, use Python for this.
+        let qubits = QUANTUM_REGISTER.get_bound(py).call1((gate_num_qubits,))?;
+        dag.add_qreg(py, &qubits)?;
+
+        let gate = GATE.get_bound(py).call1((
+            &gate_name,
+            gate_num_qubits,
+            placeholder_params
+                .iter()
+                .map(|x| x.clone_ref(py))
+                .collect::<SmallVec<[Param; 3]>>(),
+        ))?;
+        let gate_obj: OperationFromPython = gate.extract()?;
+        let qubits: Vec<Qubit> = (0..dag.num_qubits() as u32).map(Qubit).collect();
+        dag.apply_operation_back(
+            py,
+            gate_obj.operation,
+            &qubits,
+            &[],
+            if gate_obj.params.is_empty() {
+                None
+            } else {
+                Some(gate_obj.params)
+            },
+            gate_obj.extra_attrs,
+            #[cfg(feature = "cache_pygates")]
+            Some(gate.into()),
+        )?;
+        mapped_instructions.insert((gate_name, gate_num_qubits), (placeholder_params, dag));
+
+        for ((gate_name, gate_num_qubits), (equiv_params, equiv)) in basis_transforms {
+            for (_, dag) in &mut mapped_instructions.values_mut() {
+                let nodes_to_replace = dag
+                    .op_nodes(true)
+                    .filter_map(|node| {
+                        if let Some(NodeType::Operation(op)) = dag.dag().node_weight(node) {
+                            if (gate_name.as_str(), *gate_num_qubits)
+                                == (op.op.name(), op.op.num_qubits())
+                            {
+                                Some((
+                                    node,
+                                    op.params_view()
+                                        .iter()
+                                        .map(|x| x.clone_ref(py))
+                                        .collect::<SmallVec<[Param; 3]>>(),
+                                ))
+                            } else {
+                                None
+                            }
+                        } else {
+                            None
+                        }
+                    })
+                    .collect::<Vec<_>>();
+                for (node, params) in nodes_to_replace {
+                    let param_mapping: HashMap<ParameterUuid, Param> = equiv_params
+                        .iter()
+                        .map(|x| ParameterUuid::from_parameter(x.to_object(py).bind(py)))
+                        .zip(params)
+                        .map(|(uuid, param)| -> PyResult<(ParameterUuid, Param)> {
+                            Ok((uuid?, param.clone_ref(py)))
+                        })
+                        .collect::<PyResult<_>>()?;
+                    let mut replacement = equiv.clone();
+                    replacement
+                        .0
+                        .assign_parameters_from_mapping(py, param_mapping)?;
+                    let replace_dag: DAGCircuit =
+                        DAGCircuit::from_circuit_data(py, replacement.0, true)?;
+                    let op_node = dag.get_node(py, node)?;
+                    dag.py_substitute_node_with_dag(
+                        py,
+                        op_node.bind(py),
+                        &replace_dag,
+                        None,
+                        true,
+                    )?;
+                }
+            }
+        }
+    }
+    Ok(mapped_instructions)
+}
+
+/// `DAGCircuit` variant.
+///
+/// Gets the identifier of a gate instance (name, number of qubits) mapped to the
+/// number of parameters it contains currently.
+fn get_gates_num_params(
+    dag: &DAGCircuit,
+    example_gates: &mut HashMap<GateIdentifier, usize>,
+) -> PyResult<()> {
+    for node in dag.op_nodes(true) {
+        if let Some(NodeType::Operation(op)) = dag.dag().node_weight(node) {
+            example_gates.insert(
+                (op.op.name().to_string(), op.op.num_qubits()),
+                op.params_view().len(),
+            );
+            if op.op.control_flow() {
+                let blocks = op.op.blocks();
+                for block in blocks {
+                    get_gates_num_params_circuit(&block, example_gates)?;
+                }
+            }
+        }
+    }
+    Ok(())
+}
+
+/// `CircuitData` variant.
+///
+/// Gets the identifier of a gate instance (name, number of qubits) mapped to the
+/// number of parameters it contains currently.
+fn get_gates_num_params_circuit(
+    circuit: &CircuitData,
+    example_gates: &mut HashMap<GateIdentifier, usize>,
+) -> PyResult<()> {
+    for inst in circuit.iter() {
+        example_gates.insert(
+            (inst.op.name().to_string(), inst.op.num_qubits()),
+            inst.params_view().len(),
+        );
+        if inst.op.control_flow() {
+            let blocks = inst.op.blocks();
+            for block in blocks {
+                get_gates_num_params_circuit(&block, example_gates)?;
+            }
+        }
+    }
+    Ok(())
+}

crates/accelerate/src/basis/basis_translator/mod.rs

@@ -0,0 +1,21 @@
+// 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 pyo3::prelude::*;
+
+mod compose_transforms;
+
+#[pymodule]
+pub fn basis_translator(m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_wrapped(wrap_pyfunction!(compose_transforms::py_compose_transforms))?;
+    Ok(())
+}

crates/accelerate/src/basis/mod.rs

@@ -0,0 +1,21 @@
+// 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 pyo3::{prelude::*, wrap_pymodule};
+
+pub mod basis_translator;
+
+#[pymodule]
+pub fn basis(m: &Bound<PyModule>) -> PyResult<()> {
+    m.add_wrapped(wrap_pymodule!(basis_translator::basis_translator))?;
+    Ok(())
+}

crates/accelerate/src/equivalence.rs

@@ -130,14 +130,14 @@ pub struct Equivalence {
     #[pyo3(get)]
     pub params: SmallVec<[Param; 3]>,
     #[pyo3(get)]
-    pub circuit: CircuitRep,
+    pub circuit: CircuitFromPython,
 }
 
 #[pymethods]
 impl Equivalence {
     #[new]
     #[pyo3(signature = (params, circuit))]
-    fn new(params: SmallVec<[Param; 3]>, circuit: CircuitRep) -> Self {
+    fn new(params: SmallVec<[Param; 3]>, circuit: CircuitFromPython) -> Self {
         Self { circuit, params }
     }
 
@@ -295,15 +295,17 @@ impl<'py> FromPyObject<'py> for GateOper {
     }
 }
 
-/// Representation of QuantumCircuit by using an instance of `CircuitData`.]
+/// Used to extract an instance of [CircuitData] from a `QuantumCircuit`.
+/// It also ensures seamless conversion back to `QuantumCircuit` once sent
+/// back to Python.
 ///
 /// TODO: Remove this implementation once the `EquivalenceLibrary` is no longer
 /// called from Python, or once the API is able to seamlessly accept instances
-/// of `CircuitData`.
+/// of [CircuitData].
 #[derive(Debug, Clone)]
-pub struct CircuitRep(pub CircuitData);
+pub struct CircuitFromPython(pub CircuitData);
 
-impl FromPyObject<'_> for CircuitRep {
+impl FromPyObject<'_> for CircuitFromPython {
     fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult<Self> {
         if ob.is_instance(QUANTUM_CIRCUIT.get_bound(ob.py()))? {
             let data: Bound<PyAny> = ob.getattr("_data")?;
@@ -318,7 +320,7 @@ impl FromPyObject<'_> for CircuitRep {
     }
 }
 
-impl IntoPy<PyObject> for CircuitRep {
+impl IntoPy<PyObject> for CircuitFromPython {
     fn into_py(self, py: Python<'_>) -> PyObject {
         QUANTUM_CIRCUIT
             .get_bound(py)
@@ -328,7 +330,7 @@ impl IntoPy<PyObject> for CircuitRep {
     }
 }
 
-impl ToPyObject for CircuitRep {
+impl ToPyObject for CircuitFromPython {
     fn to_object(&self, py: Python<'_>) -> PyObject {
         self.clone().into_py(py)
     }
@@ -395,7 +397,7 @@ impl EquivalenceLibrary {
         &mut self,
         py: Python,
         gate: GateOper,
-        equivalent_circuit: CircuitRep,
+        equivalent_circuit: CircuitFromPython,
     ) -> PyResult<()> {
         self.add_equivalence(py, &gate.operation, &gate.params, equivalent_circuit)
     }
@@ -425,7 +427,12 @@ impl EquivalenceLibrary {
     ///     entry (List['QuantumCircuit']) : A list of QuantumCircuits, each
     ///         equivalently implementing the given Gate.
     #[pyo3(name = "set_entry")]
-    fn py_set_entry(&mut self, py: Python, gate: GateOper, entry: Vec<CircuitRep>) -> PyResult<()> {
+    fn py_set_entry(
+        &mut self,
+        py: Python,
+        gate: GateOper,
+        entry: Vec<CircuitFromPython>,
+    ) -> PyResult<()> {
         self.set_entry(py, &gate.operation, &gate.params, entry)
     }
 
@@ -567,7 +574,7 @@ impl EquivalenceLibrary {
         py: Python,
         gate: &PackedOperation,
         params: &[Param],
-        equivalent_circuit: CircuitRep,
+        equivalent_circuit: CircuitFromPython,
     ) -> PyResult<()> {
         raise_if_shape_mismatch(gate, &equivalent_circuit)?;
         raise_if_param_mismatch(py, params, equivalent_circuit.0.unsorted_parameters(py)?)?;
@@ -614,7 +621,7 @@ impl EquivalenceLibrary {
         py: Python,
         gate: &PackedOperation,
         params: &[Param],
-        entry: Vec<CircuitRep>,
+        entry: Vec<CircuitFromPython>,
     ) -> PyResult<()> {
         for equiv in entry.iter() {
             raise_if_shape_mismatch(gate, equiv)?;
@@ -714,7 +721,7 @@ fn raise_if_param_mismatch(
     Ok(())
 }
 
-fn raise_if_shape_mismatch(gate: &PackedOperation, circuit: &CircuitRep) -> PyResult<()> {
+fn raise_if_shape_mismatch(gate: &PackedOperation, circuit: &CircuitFromPython) -> PyResult<()> {
     let op_ref = gate.view();
     if op_ref.num_qubits() != circuit.0.num_qubits() as u32
         || op_ref.num_clbits() != circuit.0.num_clbits() as u32
@@ -732,7 +739,11 @@ fn raise_if_shape_mismatch(gate: &PackedOperation, circuit: &CircuitRep) -> PyRe
     Ok(())
 }
 
-fn rebind_equiv(py: Python, equiv: Equivalence, query_params: &[Param]) -> PyResult<CircuitRep> {
+fn rebind_equiv(
+    py: Python,
+    equiv: Equivalence,
+    query_params: &[Param],
+) -> PyResult<CircuitFromPython> {
     let (equiv_params, mut equiv_circuit) = (equiv.params, equiv.circuit);
     let param_mapping: PyResult<IndexMap<ParameterUuid, &Param>> = equiv_params
         .iter()

crates/accelerate/src/lib.rs

@@ -15,6 +15,7 @@ use std::env;
 use pyo3::import_exception;
 
 pub mod barrier_before_final_measurement;
+pub mod basis;
 pub mod check_map;
 pub mod circuit_library;
 pub mod commutation_analysis;

crates/circuit/src/dag_circuit.rs

@@ -988,7 +988,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."));
         }
@@ -1671,7 +1671,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>,
@@ -2860,8 +2860,8 @@ def _format(operand):
     ///
     /// Raises:
     ///     DAGCircuitError: if met with unexpected predecessor/successors
-    #[pyo3(signature = (node, input_dag, wires=None, propagate_condition=true))]
-    fn substitute_node_with_dag(
+    #[pyo3(name = "substitute_node_with_dag", signature = (node, input_dag, wires=None, propagate_condition=true))]
+    pub fn py_substitute_node_with_dag(
         &mut self,
         py: Python,
         node: &Bound<PyAny>,