[Oxidize BasisTranslator] Add rust-native compose_transforms()

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

@@ -0,0 +1,209 @@
+// 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::{exceptions::PyTypeError, prelude::*};
+use qiskit_circuit::circuit_instruction::OperationFromPython;
+use qiskit_circuit::imports::{GATE, PARAMETER_VECTOR, QUANTUM_CIRCUIT, 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;
+
+// Custom types
+pub type GateIdentifier = (String, u32);
+pub type BasisTransformIn = (SmallVec<[Param; 3]>, CircuitFromPython);
+pub type BasisTransformOut = (SmallVec<[Param; 3]>, DAGCircuit);
+// TODO: Remove these and use the version from `EquivalenceLibrary`
+
+/// Representation of QuantumCircuit which the original circuit object + an
+/// instance of `CircuitData`.
+#[derive(Debug, Clone)]
+pub struct CircuitFromPython(pub CircuitData);
+
+impl FromPyObject<'_> for CircuitFromPython {
+    fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult<Self> {
+        if ob.is_instance(QUANTUM_CIRCUIT.get_bound(ob.py()))? {
+            let data: CircuitData = ob.getattr("_data")?.extract()?;
+            Ok(Self(data))
+        } else {
+            Err(PyTypeError::new_err(
+                "Provided object was not an instance of QuantumCircuit",
+            ))
+        }
+    }
+}
+
+#[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 example_gates: HashMap<GateIdentifier, usize> = HashMap::default();
+    get_gates_num_params(source_dag, &mut example_gates)?;
+    let mut mapped_instructions: HashMap<GateIdentifier, BasisTransformOut> = HashMap::new();
+
+    for (gate_name, gate_num_qubits) in source_basis.iter().cloned() {
+        let num_params = example_gates[&(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/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>,

crates/circuit/src/imports.rs

@@ -72,6 +72,8 @@ pub static CLASSICAL_REGISTER: ImportOnceCell =
     ImportOnceCell::new("qiskit.circuit.classicalregister", "ClassicalRegister");
 pub static PARAMETER_EXPRESSION: ImportOnceCell =
     ImportOnceCell::new("qiskit.circuit.parameterexpression", "ParameterExpression");
+pub static PARAMETER_VECTOR: ImportOnceCell =
+    ImportOnceCell::new("qiskit.circuit.parametervector", "ParameterVector");
 pub static QUANTUM_CIRCUIT: ImportOnceCell =
     ImportOnceCell::new("qiskit.circuit.quantumcircuit", "QuantumCircuit");
 pub static SINGLETON_GATE: ImportOnceCell =

crates/pyext/src/lib.rs

@@ -29,6 +29,7 @@ where
 #[pymodule]
 fn _accelerate(m: &Bound<PyModule>) -> PyResult<()> {
     add_submodule(m, ::qiskit_accelerate::barrier_before_final_measurement::barrier_before_final_measurements_mod, "barrier_before_final_measurement")?;
+    add_submodule(m, ::qiskit_accelerate::basis::basis, "basis")?;
     add_submodule(m, ::qiskit_accelerate::check_map::check_map_mod, "check_map")?;
     add_submodule(m, ::qiskit_accelerate::circuit_library::circuit_library, "circuit_library")?;
     add_submodule(m, ::qiskit_accelerate::commutation_analysis::commutation_analysis, "commutation_analysis")?;

qiskit/__init__.py

@@ -53,6 +53,8 @@
 # and not have to rely on attribute access.  No action needed for top-level extension packages.
 sys.modules["qiskit._accelerate.circuit"] = _accelerate.circuit
 sys.modules["qiskit._accelerate.circuit_library"] = _accelerate.circuit_library
+sys.modules["qiskit._accelerate.basis"] = _accelerate.basis
+sys.modules["qiskit._accelerate.basis.basis_translator"] = _accelerate.basis.basis_translator
 sys.modules["qiskit._accelerate.converters"] = _accelerate.converters
 sys.modules["qiskit._accelerate.convert_2q_block_matrix"] = _accelerate.convert_2q_block_matrix
 sys.modules["qiskit._accelerate.dense_layout"] = _accelerate.dense_layout